diff --git a/extract_otp_secret_keys.py b/extract_otp_secret_keys.py index d904d15..ef34958 100644 --- a/extract_otp_secret_keys.py +++ b/extract_otp_secret_keys.py @@ -94,38 +94,42 @@ def extract_otps(args): otps = [] i = j = 0 - for line in (line.strip() for line in fileinput.input(args.infile)): - if verbose: print(line) - if line.startswith('#') or line == '': continue - i += 1 - payload = get_payload_from_line(line, i, args) - - # pylint: disable=no-member - for raw_otp in payload.otp_parameters: - j += 1 - if verbose: print('\n{}. Secret Key'.format(j)) - secret = convert_secret_from_bytes_to_base32_str(raw_otp.secret) - otp_type_enum = get_enum_name_by_number(raw_otp, 'type') - otp_type = get_otp_type_str_from_code(raw_otp.type) - otp_url = build_otp_url(secret, raw_otp) - otp = { - "name": raw_otp.name, - "secret": secret, - "issuer": raw_otp.issuer, - "type": otp_type, - "counter": raw_otp.counter if raw_otp.type == 1 else None, - "url": otp_url - } - if not quiet: - print_otp(otp) - if args.printqr: - print_qr(args, otp_url) - if args.saveqr: - save_qr(otp, args, j) - if not quiet: - print() - - otps.append(otp) + finput = fileinput.input(args.infile) + try: + for line in (line.strip() for line in finput): + if verbose: print(line) + if line.startswith('#') or line == '': continue + i += 1 + payload = get_payload_from_line(line, i, args) + + # pylint: disable=no-member + for raw_otp in payload.otp_parameters: + j += 1 + if verbose: print('\n{}. Secret Key'.format(j)) + secret = convert_secret_from_bytes_to_base32_str(raw_otp.secret) + otp_type_enum = get_enum_name_by_number(raw_otp, 'type') + otp_type = get_otp_type_str_from_code(raw_otp.type) + otp_url = build_otp_url(secret, raw_otp) + otp = { + "name": raw_otp.name, + "secret": secret, + "issuer": raw_otp.issuer, + "type": otp_type, + "counter": raw_otp.counter if raw_otp.type == 1 else None, + "url": otp_url + } + if not quiet: + print_otp(otp) + if args.printqr: + print_qr(args, otp_url) + if args.saveqr: + save_qr(otp, args, j) + if not quiet: + print() + + otps.append(otp) + finally: + finput.close() return otps @@ -146,7 +150,12 @@ def get_payload_from_line(line, i, args): if verbose > 1: print('\nDEBUG: data_base64_fixed={}'.format(data_base64)) data = base64.b64decode(data_base64_fixed, validate=True) payload = protobuf_generated_python.google_auth_pb2.MigrationPayload() - payload.ParseFromString(data) + try: + payload.ParseFromString(data) + except: + print('\nERROR: Cannot decode otpauth-migration migration payload.') + print('data={}'.format(data_base64)) + exit(1); if verbose: print('\n{}. Payload Line'.format(i), payload, sep='\n') diff --git a/test/test_export_wrong_content.txt b/test/test_export_wrong_content.txt new file mode 100644 index 0000000..a24d3d5 --- /dev/null +++ b/test/test_export_wrong_content.txt @@ -0,0 +1 @@ +Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua. diff --git a/test/test_export_wrong_data.txt b/test/test_export_wrong_data.txt new file mode 100644 index 0000000..8c0cf14 --- /dev/null +++ b/test/test_export_wrong_data.txt @@ -0,0 +1 @@ +otpauth-migration://offline?data=XXXX diff --git a/test/test_export_wrong_prefix.txt b/test/test_export_wrong_prefix.txt new file mode 100644 index 0000000..afd8dbd --- /dev/null +++ b/test/test_export_wrong_prefix.txt @@ -0,0 +1 @@ +QR-Code:otpauth-migration://offline?data=CjUKEPqlBekzoNEukL7qlsjBCDYSDnBpQHJhc3BiZXJyeXBpGgtyYXNwYmVycnlwaSABKAEwAhABGAEgACjr4JKK%2B%2F%2F%2F%2F%2F8B diff --git a/test_extract_otp_secret_keys_pytest.py b/test_extract_otp_secret_keys_pytest.py index d98a9cc..90d7193 100644 --- a/test_extract_otp_secret_keys_pytest.py +++ b/test_extract_otp_secret_keys_pytest.py @@ -279,6 +279,70 @@ def test_verbose_and_quiet(capsys): assert 'The arguments --verbose and --quiet are mutually exclusive.' in captured.out +def test_wrong_data(capsys): + with raises(SystemExit) as pytest_wrapped_e: + # Act + extract_otp_secret_keys.main(['test/test_export_wrong_data.txt']) + + # Assert + captured = capsys.readouterr() + + expected_stdout = ''' +ERROR: Cannot decode otpauth-migration migration payload. +data=XXXX +''' + + assert captured.out == expected_stdout + assert captured.err == '' + + +def test_wrong_content(capsys): + with raises(SystemExit) as pytest_wrapped_e: + # Act + extract_otp_secret_keys.main(['test/test_export_wrong_content.txt']) + + # Assert + captured = capsys.readouterr() + + expected_stdout = ''' +WARN: line is not a otpauth-migration:// URL +input file: test/test_export_wrong_content.txt +line "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua." +Probably a wrong file was given + +ERROR: no data query parameter in input URL +input file: test/test_export_wrong_content.txt +line "Lorem ipsum dolor sit amet, consetetur sadipscing elitr, sed diam nonumy eirmod tempor invidunt ut labore et dolore magna aliquyam erat, sed diam voluptua." +Probably a wrong file was given +''' + + assert captured.out == expected_stdout + assert captured.err == '' + + +def test_wrong_prefix(capsys): + # Act + extract_otp_secret_keys.main(['test/test_export_wrong_prefix.txt']) + + # Assert + captured = capsys.readouterr() + + expected_stdout = ''' +WARN: line is not a otpauth-migration:// URL +input file: test/test_export_wrong_prefix.txt +line "QR-Code:otpauth-migration://offline?data=CjUKEPqlBekzoNEukL7qlsjBCDYSDnBpQHJhc3BiZXJyeXBpGgtyYXNwYmVycnlwaSABKAEwAhABGAEgACjr4JKK%2B%2F%2F%2F%2F%2F8B" +Probably a wrong file was given +Name: pi@raspberrypi +Secret: 7KSQL2JTUDIS5EF65KLMRQIIGY +Issuer: raspberrypi +Type: totp + +''' + + assert captured.out == expected_stdout + assert captured.err == '' + + def test_add_pre_suffix(capsys): assert extract_otp_secret_keys.add_pre_suffix("name.csv", "totp") == "name.totp.csv" assert extract_otp_secret_keys.add_pre_suffix("name.csv", "") == "name..csv"