diff --git a/src/mvt/android/cmd_check_backup.py b/src/mvt/android/cmd_check_backup.py index 066cd8aa..b4878b00 100644 --- a/src/mvt/android/cmd_check_backup.py +++ b/src/mvt/android/cmd_check_backup.py @@ -27,6 +27,8 @@ from .modules.backup import BACKUP_MODULES log = logging.getLogger(__name__) +MAX_BACKUP_PASSWORD_ATTEMPTS = 3 + class InvalidAndroidBackup(Exception): pass @@ -72,6 +74,15 @@ class CmdAndroidCheckBackup(Command): self.__tar: Optional[tarfile.TarFile] = None self.__files: List[str] = [] + def _get_backup_password(self) -> str: + password = prompt_or_load_android_backup_password(log, self.module_options) + if not password: + if self.sub_command: + raise InvalidAndroidBackup("No backup password provided") + log.critical("No backup password provided.") + sys.exit(1) + return password + def from_ab(self, ab_file_bytes: bytes) -> None: self.__type = "ab" header = parse_ab_header(ab_file_bytes) @@ -84,28 +95,40 @@ class CmdAndroidCheckBackup(Command): sys.exit(1) password = None + # A password typed at the prompt is asked for again if it is wrong. + # One given on the command line or in the environment is not. + max_attempts = 1 if header["encryption"] != "none": - password = prompt_or_load_android_backup_password(log, self.module_options) - if not password: + if not self.module_options.get("backup_password", None): + max_attempts = MAX_BACKUP_PASSWORD_ATTEMPTS + password = self._get_backup_password() + attempt = 1 + while True: + try: + tardata = parse_backup_file(ab_file_bytes, password=password) + break + except InvalidBackupPassword: + if attempt < max_attempts: + attempt += 1 + log.error( + "Invalid backup password, try again (attempt %d of %d)", + attempt, + max_attempts, + ) + password = self._get_backup_password() + continue if self.sub_command: - raise InvalidAndroidBackup("No backup password provided") - log.critical("No backup password provided.") + raise InvalidAndroidBackup("Invalid backup password") + log.critical("Invalid backup password") + sys.exit(1) + except AndroidBackupParsingError as exc: + if self.sub_command: + raise InvalidAndroidBackup( + f"Impossible to parse this backup file: {exc}" + ) from exc + log.critical("Impossible to parse this backup file: %s", exc) + log.critical("Please use Android Backup Extractor (ABE) instead") sys.exit(1) - try: - tardata = parse_backup_file(ab_file_bytes, password=password) - except InvalidBackupPassword: - if self.sub_command: - raise InvalidAndroidBackup("Invalid backup password") - log.critical("Invalid backup password") - sys.exit(1) - except AndroidBackupParsingError as exc: - if self.sub_command: - raise InvalidAndroidBackup( - f"Impossible to parse this backup file: {exc}" - ) from exc - log.critical("Impossible to parse this backup file: %s", exc) - log.critical("Please use Android Backup Extractor (ABE) instead") - sys.exit(1) dbytes = io.BytesIO(tardata) try: diff --git a/tests/test_check_android_backup.py b/tests/test_check_android_backup.py index ac0680f0..61ba7a12 100644 --- a/tests/test_check_android_backup.py +++ b/tests/test_check_android_backup.py @@ -59,6 +59,66 @@ class TestCheckAndroidBackupCommand: assert result.exit_code == 1 assert "Invalid backup password" in caplog.text + def test_check_encrypted_backup_prompt_retry(self, mocker): + """A wrong password at the prompt is asked for again""" + prompt_mock = mocker.patch( + "mvt.android.modules.backup.helpers.prompt_password", + side_effect=["wrong_password", TEST_BACKUP_PASSWORD], + ) + runner = CliRunner() + path = os.path.join(get_artifact_folder(), "androidqf_encrypted/backup.ab") + result = runner.invoke(check_backup, [path]) + + assert prompt_mock.call_count == 2 + assert result.exit_code == 0 + + def test_check_encrypted_backup_prompt_gives_up(self, mocker, caplog): + """The prompt gives up after three wrong passwords""" + prompt_mock = mocker.patch( + "mvt.android.modules.backup.helpers.prompt_password", + return_value="wrong_password", + ) + runner = CliRunner() + path = os.path.join(get_artifact_folder(), "androidqf_encrypted/backup.ab") + + with caplog.at_level(logging.CRITICAL): + result = runner.invoke(check_backup, [path]) + + assert prompt_mock.call_count == 3 + assert result.exit_code == 1 + assert "Invalid backup password" in caplog.text + + def test_check_encrypted_backup_prompt_retry_empty(self, mocker, caplog): + """An empty answer on a retry ends the check""" + prompt_mock = mocker.patch( + "mvt.android.modules.backup.helpers.prompt_password", + side_effect=["wrong_password", ""], + ) + runner = CliRunner() + path = os.path.join(get_artifact_folder(), "androidqf_encrypted/backup.ab") + + with caplog.at_level(logging.CRITICAL): + result = runner.invoke(check_backup, [path]) + + assert prompt_mock.call_count == 2 + assert result.exit_code == 1 + assert "No backup password provided" in caplog.text + + def test_check_encrypted_backup_cli_invalid_no_prompt(self, mocker): + """A password given on the command line is not asked for again""" + prompt_mock = mocker.patch( + "mvt.android.modules.backup.helpers.prompt_password", + return_value=TEST_BACKUP_PASSWORD, + ) + runner = CliRunner() + path = os.path.join(get_artifact_folder(), "androidqf_encrypted/backup.ab") + result = runner.invoke( + check_backup, ["--backup-password", "invalid_password", path] + ) + + assert prompt_mock.call_count == 0 + assert result.exit_code == 1 + def test_check_encrypted_backup_env(self, mocker): """Provide password as environment variable""" prompt_mock = mocker.patch(