save qr code to specific dir, improve help, add tests

- use metavar for files and dirs in help
- support several recursive dirs in saveqr
- add saveqr and debug tests

.vscode/settings.json

@@ -3,5 +3,18 @@
         "."
     ],
     "python.testing.unittestEnabled": false,
-    "python.testing.pytestEnabled": true
+    "python.testing.pytestEnabled": true,
+    "cSpell.words": [
+        "devbox",
+        "HOTP",
+        "otpauth",
+        "pipenv",
+        "proto",
+        "protobuf",
+        "protoc",
+        "pytest",
+        "qrcode",
+        "TOTP",
+        "venv"
+    ]
 }

README.md

@@ -19,7 +19,23 @@ The secret and otp values can be printed and exported to json or csv. The QR cod
 3. Save the captured QR codes in a text file. Save each QR code on a new line. (The captured QR codes look like `otpauth-migration://offline?data=...`)
 4. Call this script with the file as input:
 
-        python extract_otp_secret_keys.py -p example_export.txt
+        python extract_otp_secret_keys.py example_export.txt
+
+## Program help: arguments and options
+
+<pre>usage: extract_otp_secret_keys.py [-h] [--json FILE] [--csv FILE] [--printqr] [--saveqr DIR] [--verbose] [--quiet] infile
+
+positional arguments:
+  infile                file or - for stdin (default: -) with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored
+
+options:
+  -h, --help            show this help message and exit
+  --json FILE, -j FILE  export to json file
+  --csv FILE, -c FILE   export to csv file
+  --printqr, -p         print QR code(s) as text to the terminal (requires qrcode module)
+  --saveqr DIR, -s DIR  save QR code(s) as images to the given folder (requires qrcode module)
+  --verbose, -v         verbose output
+  --quiet, -q           no stdout output</pre>
 
 ## Dependencies
 

extract_otp_secret_keys.code-workspace

@@ -6,5 +6,20 @@
 	],
 	"settings": {
 		"python.testing.pytestEnabled": true
+	},
+	"launch": {
+		"version": "0.2.0",
+		"configurations": [
+            {
+                "name": "Python: extract_otp_secret_keys.py",
+                "type": "python",
+                "request": "launch",
+                "program": "extract_otp_secret_keys.py",
+                "args": [
+                    "example_export.txt"
+                ],
+                "console": "integratedTerminal"
+            },
+		]
 	}
-}
+}

extract_otp_secret_keys.py

@@ -5,7 +5,7 @@
 # 2. Read QR codes with QR code reader (e.g. with a second device)
 # 3. Save the captured QR codes in a text file. Save each QR code on a new line. (The captured QR codes look like "otpauth-migration://offline?data=...")
 # 4. Call this script with the file as input:
-#    python extract_otp_secret_keys.py -p example_export.txt
+#    python extract_otp_secret_keys.py example_export.txt
 #
 # Requirement:
 # The protobuf package of Google for proto3 is required for running this script.
@@ -48,7 +48,7 @@ import sys
 import csv
 import json
 from urllib.parse import parse_qs, urlencode, urlparse, quote
-from os import path, mkdir
+from os import path, makedirs
 from re import compile as rcompile
 import protobuf_generated_python.google_auth_pb2
 
@@ -60,7 +60,7 @@ def sys_main():
 def main(sys_args):
     global verbose, quiet
     args = parse_args(sys_args)
-    verbose = args.verbose
+    verbose = args.verbose if args.verbose else 0
     quiet = args.quiet
 
     otps = extract_otps(args)
@@ -70,13 +70,13 @@ def main(sys_args):
 
 def parse_args(sys_args):
     arg_parser = argparse.ArgumentParser()
-    arg_parser.add_argument('--verbose', '-v', help='verbose output', action='store_true')
-    arg_parser.add_argument('--quiet', '-q', help='no stdout output', action='store_true')
-    arg_parser.add_argument('--saveqr', '-s', help='save QR code(s) as images to the "qr" subfolder', action='store_true')
-    arg_parser.add_argument('--printqr', '-p', help='print QR code(s) as text to the terminal', action='store_true')
-    arg_parser.add_argument('--json', '-j', help='export to json file')
-    arg_parser.add_argument('--csv', '-c', help='export to csv file')
     arg_parser.add_argument('infile', help='file or - for stdin (default: -) with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored')
+    arg_parser.add_argument('--json', '-j', help='export to json file', metavar=('FILE'))
+    arg_parser.add_argument('--csv', '-c', help='export to csv file', metavar=('FILE'))
+    arg_parser.add_argument('--printqr', '-p', help='print QR code(s) as text to the terminal (requires qrcode module)', action='store_true')
+    arg_parser.add_argument('--saveqr', '-s', help='save QR code(s) as images to the given folder (requires qrcode module)', metavar=('DIR'))
+    arg_parser.add_argument('--verbose', '-v', help='verbose output', action='count')
+    arg_parser.add_argument('--quiet', '-q', help='no stdout output', action='store_true')
     args = arg_parser.parse_args(sys_args)
     if args.verbose and args.quiet:
         print("The arguments --verbose and --quite are mutual exclusive.")
@@ -125,15 +125,21 @@ def extract_otps(args):
 
 
 def get_payload_from_line(line, i, args):
+    global verbose
     if not line.startswith('otpauth-migration://'):
         print('\nWARN: line is not a otpauth-migration:// URL\ninput file: {}\nline "{}"\nProbably a wrong file was given'.format(args.infile, line))
     parsed_url = urlparse(line)
-    params = parse_qs(parsed_url.query)
+    if verbose > 1: print('\nDEBUG: parsed_url={}'.format(parsed_url))
+    params = parse_qs(parsed_url.query, strict_parsing=True)
+    if verbose > 1: print('\nDEBUG: querystring params={}'.format(params))
     if 'data' not in params:
         print('\nERROR: no data query parameter in input URL\ninput file: {}\nline "{}"\nProbably a wrong file was given'.format(args.infile, line))
         sys.exit(1)
-    data_encoded = params['data'][0]
-    data = base64.b64decode(data_encoded)
+    data_base64 = params['data'][0]
+    if verbose > 1: print('\nDEBUG: data_base64={}'.format(data_base64))
+    data_base64_fixed = data_base64.replace(' ', '+')
+    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)
     if verbose:
@@ -170,11 +176,12 @@ def print_otp(otp):
 
 
 def save_qr(otp, args, j):
-    if not (path.exists('qr')): mkdir('qr')
+    dir = args.saveqr
+    if not (path.exists(dir)): makedirs(dir, exist_ok=True)
     pattern = rcompile(r'[\W_]+')
-    file_otp_name = pattern.sub('', otp.name)
-    file_otp_issuer = pattern.sub('', otp.issuer)
-    save_qr_file(args, otp.url, 'qr/{}-{}{}.png'.format(j, file_otp_name, '-' + file_otp_issuer if file_otp_issuer else ''))
+    file_otp_name = pattern.sub('', otp['name'])
+    file_otp_issuer = pattern.sub('', otp['issuer'])
+    save_qr_file(args, otp['url'], '{}/{}-{}{}.png'.format(dir, j, file_otp_name, '-' + file_otp_issuer if file_otp_issuer else ''))
     return file_otp_issuer
 
 

test/test_plus_problem_export.txt

@@ -0,0 +1 @@
+otpauth-migration://offline?data=ClEKFAciUeGF4aS6IDCvMv99ySZ1ekKsEiVTZXJlbml0eUxhYnM6dGVzdDFAc2VyZW5pdHlsYWJzLmNvLnVrGgxTZXJlbml0eUxhYnMgASgBMAIKUQoUkIY8/fbrHZWTb4CBln18lvqt0HcSJVNlcmVuaXR5TGFiczp0ZXN0MkBzZXJlbml0eWxhYnMuY28udWsaDFNlcmVuaXR5TGFicyABKAEwAgpRChScf+1/Ua4d4gCY0W/7fj9VBkM9PBIlU2VyZW5pdHlMYWJzOnRlc3QzQHNlcmVuaXR5bGFicy5jby51axoMU2VyZW5pdHlMYWJzIAEoATACClEKFG6Qu0ryTSFA/l5rmvTIXtNeb5LtEiVTZXJlbml0eUxhYnM6dGVzdDRAc2VyZW5pdHlsYWJzLmNvLnVrGgxTZXJlbml0eUxhYnMgASgBMAIQARgBIAAogtTa1vz/////AQ==

test_extract_otp_secret_keys_pytest.py

@@ -18,12 +18,14 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-from utils import read_csv, read_json, remove_file, read_file_to_str
+from utils import read_csv, read_json, remove_file, remove_dir_with_files, read_file_to_str
+from os import path
+from pytest import raises
 
 import extract_otp_secret_keys
 
 
-def test_extract_csv():
+def test_extract_csv(capsys):
     # Arrange
     cleanup()
 
@@ -36,11 +38,16 @@ def test_extract_csv():
 
     assert actual_csv == expected_csv
 
+    captured = capsys.readouterr()
+
+    assert captured.out == ''
+    assert captured.err == ''
+
     # Clean up
     cleanup()
 
 
-def test_extract_json():
+def test_extract_json(capsys):
     # Arrange
     cleanup()
 
@@ -53,6 +60,11 @@ def test_extract_json():
 
     assert actual_json == expected_json
 
+    captured = capsys.readouterr()
+
+    assert captured.out == ''
+    assert captured.err == ''
+
     # Clean up
     cleanup()
 
@@ -88,6 +100,39 @@ Type:   OTP_TOTP
     assert captured.err == ''
 
 
+def test_extract_not_encoded_plus(capsys):
+    # Act
+    extract_otp_secret_keys.main(['test/test_plus_problem_export.txt'])
+
+    # Assert
+    captured = capsys.readouterr()
+
+    expected_stdout = '''Name:   SerenityLabs:test1@serenitylabs.co.uk
+Secret: A4RFDYMF4GSLUIBQV4ZP67OJEZ2XUQVM
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
+Name:   SerenityLabs:test2@serenitylabs.co.uk
+Secret: SCDDZ7PW5MOZLE3PQCAZM7L4S35K3UDX
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
+Name:   SerenityLabs:test3@serenitylabs.co.uk
+Secret: TR76272RVYO6EAEY2FX7W7R7KUDEGPJ4
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
+Name:   SerenityLabs:test4@serenitylabs.co.uk
+Secret: N2ILWSXSJUQUB7S6NONPJSC62NPG7EXN
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
+'''
+
+    assert captured.out == expected_stdout
+    assert captured.err == ''
+
+
 def test_extract_printqr(capsys):
     # Act
     extract_otp_secret_keys.main(['-p', 'example_export.txt'])
@@ -101,6 +146,28 @@ def test_extract_printqr(capsys):
     assert captured.err == ''
 
 
+def test_extract_saveqr(capsys):
+    # Arrange
+    cleanup()
+
+    # Act
+    extract_otp_secret_keys.main(['-q', '-s', 'testout/qr/', 'example_export.txt'])
+
+    # Assert
+    captured = capsys.readouterr()
+
+    assert captured.out == ''
+    assert captured.err == ''
+
+    assert path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png')
+    assert path.isfile('testout/qr/2-piraspberrypi.png')
+    assert path.isfile('testout/qr/3-piraspberrypi.png')
+    assert path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png')
+
+    # Clean up
+    cleanup()
+
+
 def test_extract_verbose(capsys):
     # Act
     extract_otp_secret_keys.main(['-v', 'example_export.txt'])
@@ -114,6 +181,36 @@ def test_extract_verbose(capsys):
     assert captured.err == ''
 
 
+def test_extract_debug(capsys):
+    # Act
+    extract_otp_secret_keys.main(['-vv', 'example_export.txt'])
+
+    # Assert
+    captured = capsys.readouterr()
+
+    expected_stdout = read_file_to_str('test/print_verbose_output.txt')
+
+    assert len(captured.out) > len(expected_stdout)
+    assert "DEBUG: " in captured.out
+    assert captured.err == ''
+
+
+def test_extract_help(capsys):
+    with raises(SystemExit) as pytest_wrapped_e:
+        # Act
+        extract_otp_secret_keys.main(['-h'])
+
+    # Assert
+    captured = capsys.readouterr()
+
+    assert len(captured.out) > 0
+    assert "-h, --help" in captured.out and "--verbose, -v" in captured.out
+    assert captured.err == ''
+    assert pytest_wrapped_e.type == SystemExit
+    assert pytest_wrapped_e.value.code == 0
+
+
 def cleanup():
     remove_file('test_example_output.csv')
     remove_file('test_example_output.json')
+    remove_dir_with_files('testout/')

test_extract_otp_secret_keys_unittest.py

@@ -21,7 +21,8 @@
 import unittest
 import io
 from contextlib import redirect_stdout
-from utils import read_csv, read_json, remove_file, Capturing, read_file_to_str
+from utils import read_csv, read_json, remove_file, remove_dir_with_files, Capturing, read_file_to_str
+from os import path
 
 import extract_otp_secret_keys
 
@@ -95,6 +96,35 @@ Secret: 7KSQL2JTUDIS5EF65KLMRQIIGY
 Issuer: raspberrypi
 Type:   OTP_TOTP
 
+'''
+        self.assertEqual(actual_output, expected_output)
+
+    def test_extract_not_encoded_plus(self):
+        out = io.StringIO()
+        with redirect_stdout(out):
+            extract_otp_secret_keys.main(['test/test_plus_problem_export.txt'])
+        actual_output = out.getvalue()
+
+        expected_output = '''Name:   SerenityLabs:test1@serenitylabs.co.uk
+Secret: A4RFDYMF4GSLUIBQV4ZP67OJEZ2XUQVM
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
+Name:   SerenityLabs:test2@serenitylabs.co.uk
+Secret: SCDDZ7PW5MOZLE3PQCAZM7L4S35K3UDX
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
+Name:   SerenityLabs:test3@serenitylabs.co.uk
+Secret: TR76272RVYO6EAEY2FX7W7R7KUDEGPJ4
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
+Name:   SerenityLabs:test4@serenitylabs.co.uk
+Secret: N2ILWSXSJUQUB7S6NONPJSC62NPG7EXN
+Issuer: SerenityLabs
+Type:   OTP_TOTP
+
 '''
         self.assertEqual(actual_output, expected_output)
 
@@ -108,6 +138,14 @@ Type:   OTP_TOTP
 
         self.assertEqual(actual_output, expected_output)
 
+    def test_extract_saveqr(self):
+        extract_otp_secret_keys.main(['-q', '-s', 'testout/qr/', 'example_export.txt'])
+
+        self.assertTrue(path.isfile('testout/qr/1-piraspberrypi-raspberrypi.png'))
+        self.assertTrue(path.isfile('testout/qr/2-piraspberrypi.png'))
+        self.assertTrue(path.isfile('testout/qr/3-piraspberrypi.png'))
+        self.assertTrue(path.isfile('testout/qr/4-piraspberrypi-raspberrypi.png'))
+
     def test_extract_verbose(self):
         out = io.StringIO()
         with redirect_stdout(out):
@@ -118,6 +156,30 @@ Type:   OTP_TOTP
 
         self.assertEqual(actual_output, expected_output)
 
+    def test_extract_debug(self):
+        out = io.StringIO()
+        with redirect_stdout(out):
+            extract_otp_secret_keys.main(['-vv', 'example_export.txt'])
+        actual_output = out.getvalue()
+
+        expected_stdout = read_file_to_str('test/print_verbose_output.txt')
+
+        self.assertGreater(len(actual_output), len(expected_stdout))
+        self.assertTrue("DEBUG: " in actual_output)
+
+    def test_extract_help(self):
+        out = io.StringIO()
+        with redirect_stdout(out):
+            try:
+                extract_otp_secret_keys.main(['-h'])
+            except SystemExit:
+                pass
+
+        actual_output = out.getvalue()
+
+        self.assertGreater(len(actual_output), 0)
+        self.assertTrue("-h, --help" in actual_output and "--verbose, -v" in actual_output)
+
     def setUp(self):
         self.cleanup()
 
@@ -127,6 +189,7 @@ Type:   OTP_TOTP
     def cleanup(self):
         remove_file('test_example_output.csv')
         remove_file('test_example_output.json')
+        remove_dir_with_files('testout/')
 
 
 if __name__ == '__main__':

utils.py

@@ -16,6 +16,7 @@
 import csv
 import json
 import os
+import shutil
 from io import StringIO
 import sys
 
@@ -38,8 +39,12 @@ with Capturing() as output:
         sys.stdout = self._stdout
 
 
-def remove_file(filename):
-    if os.path.exists(filename): os.remove(filename)
+def remove_file(file):
+    if os.path.isfile(file): os.remove(file)
+
+
+def remove_dir_with_files(dir):
+    if os.path.exists(dir): shutil.rmtree(dir)
 
 
 def read_csv(filename):