test wrong data and improve error handling

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')
 

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.

test/test_export_wrong_data.txt

@@ -0,0 +1 @@
+otpauth-migration://offline?data=XXXX

test/test_export_wrong_prefix.txt

@@ -0,0 +1 @@
+QR-Code:otpauth-migration://offline?data=CjUKEPqlBekzoNEukL7qlsjBCDYSDnBpQHJhc3BiZXJyeXBpGgtyYXNwYmVycnlwaSABKAEwAhABGAEgACjr4JKK%2B%2F%2F%2F%2F%2F8B

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"