From 343520acb80cb797edc2afc0458dfb5ae134a5ba Mon Sep 17 00:00:00 2001
From: scito <info@scito.ch>
Date: Sat, 24 Dec 2022 04:48:12 +0100
Subject: [PATCH] support multiple infiles

---
 README.md                              |  2 +-
 extract_otp_secret_keys.py             | 70 +++++++++++++-------------
 test_extract_otp_secret_keys_pytest.py | 53 ++++++++++++-------
 3 files changed, 70 insertions(+), 55 deletions(-)

diff --git a/README.md b/README.md
index a1bcf58..37068d9 100644
--- a/README.md
+++ b/README.md
@@ -31,7 +31,7 @@ cd extract_otp_secret_keys
 
 ## Program help: arguments and options
 
-<pre>usage: extract_otp_secret_keys.py [-h] [--json FILE] [--csv FILE] [--keepass FILE] [--printqr] [--saveqr DIR] [--verbose | --quiet] infile
+<pre>usage: extract_otp_secret_keys.py [-h] [--json FILE] [--csv FILE] [--keepass FILE] [--printqr] [--saveqr DIR] [--verbose | --quiet] infile [infile ...]
 
 positional arguments:
   infile                   1) file or - for stdin with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored; 2) image file containing a QR code or = for stdin for an image containing a QR code
diff --git a/extract_otp_secret_keys.py b/extract_otp_secret_keys.py
index 5dfa313..02d62be 100644
--- a/extract_otp_secret_keys.py
+++ b/extract_otp_secret_keys.py
@@ -79,7 +79,7 @@ def main(sys_args):
 def parse_args(sys_args):
     formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=52)
     arg_parser = argparse.ArgumentParser(formatter_class=formatter)
-    arg_parser.add_argument('infile', help='1) file or - for stdin with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored; 2) image file containing a QR code or = for stdin for an image containing a QR code')
+    arg_parser.add_argument('infile', help='1) file or - for stdin with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored; 2) image file containing a QR code or = for stdin for an image containing a QR code', nargs='+')
     arg_parser.add_argument('--json', '-j', help='export json file or - for stdout', metavar=('FILE'))
     arg_parser.add_argument('--csv', '-c', help='export csv file or - for stdout', metavar=('FILE'))
     arg_parser.add_argument('--keepass', '-k', help='export totp/hotp csv file(s) for KeePass, - for stdout', metavar=('FILE'))
@@ -101,39 +101,39 @@ def extract_otps(args):
     otps = []
 
     i = j = 0
+    for infile in args.infile:
+        for line in get_lines_from_file(infile):
+            if verbose: print(line)
+            if line.startswith('#') or line == '': continue
+            i += 1
+            payload = get_payload_from_line(line, i, infile)
 
-    for line in get_lines_from_file(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()
 
-        # 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)
+                otps.append(otp)
 
     return otps
 
@@ -209,10 +209,10 @@ def convert_img_to_line(filename):
             abort('\nERROR: Encountered exception "{}".\ninput file: {}'.format(str(e), filename))
 
 
-def get_payload_from_line(line, i, args):
+def get_payload_from_line(line, i, infile):
     global verbose
     if not line.startswith('otpauth-migration://'):
-        eprint( '\nWARN: line is not a otpauth-migration:// URL\ninput file: {}\nline "{}"\nProbably a wrong file was given'.format( args.infile, line))
+        eprint( '\nWARN: line is not a otpauth-migration:// URL\ninput file: {}\nline "{}"\nProbably a wrong file was given'.format(infile, line))
     parsed_url = urlparse(line)
     if verbose > 1: print('\nDEBUG: parsed_url={}'.format(parsed_url))
     try:
@@ -221,7 +221,7 @@ def get_payload_from_line(line, i, args):
         params = []
     if verbose > 1: print('\nDEBUG: querystring params={}'.format(params))
     if 'data' not in params:
-        abort('\nERROR: no data query parameter in input URL\ninput file: {}\nline "{}"\nProbably a wrong file was given'.format( args.infile, line))
+        abort('\nERROR: no data query parameter in input URL\ninput file: {}\nline "{}"\nProbably a wrong file was given'.format(infile, line))
     data_base64 = params['data'][0]
     if verbose > 1: print('\nDEBUG: data_base64={}'.format(data_base64))
     data_base64_fixed = data_base64.replace(' ', '+')
diff --git a/test_extract_otp_secret_keys_pytest.py b/test_extract_otp_secret_keys_pytest.py
index 8d0e561..58a9ec8 100644
--- a/test_extract_otp_secret_keys_pytest.py
+++ b/test_extract_otp_secret_keys_pytest.py
@@ -38,6 +38,21 @@ def test_extract_stdout(capsys):
     assert captured.err == ''
 
 
+def test_extract_multiple_files_and_mixed(capsys):
+    # Act
+    extract_otp_secret_keys.main([
+        'example_export.txt',
+        'test/test_googleauth_export.png',
+        'example_export.txt',
+        'test/test_googleauth_export.png'])
+
+    # Assert
+    captured = capsys.readouterr()
+
+    assert captured.out == EXPECTED_STDOUT_FROM_EXAMPLE_EXPORT + EXPECTED_STDOUT_FROM_EXAMPLE_EXPORT_PNG + EXPECTED_STDOUT_FROM_EXAMPLE_EXPORT + EXPECTED_STDOUT_FROM_EXAMPLE_EXPORT_PNG
+    assert captured.err == ''
+
+
 def test_extract_non_existent_file(capsys):
     # Act
     with raises(SystemExit) as e:
@@ -498,25 +513,7 @@ def test_img_qr_reader_from_file_happy_path(capsys):
     # Assert
     captured = capsys.readouterr()
 
-    expected_stdout =\
-'''Name:    Test1:test1@example1.com
-Secret:  JBSWY3DPEHPK3PXP
-Issuer:  Test1
-Type:    totp
-
-Name:    Test2:test2@example2.com
-Secret:  JBSWY3DPEHPK3PXQ
-Issuer:  Test2
-Type:    totp
-
-Name:    Test3:test3@example3.com
-Secret:  JBSWY3DPEHPK3PXR
-Issuer:  Test3
-Type:    totp
-
-'''
-
-    assert captured.out == expected_stdout
+    assert captured.out == EXPECTED_STDOUT_FROM_EXAMPLE_EXPORT_PNG
     assert captured.err == ''
 
 
@@ -664,3 +661,21 @@ Secret:  7KSQL2JTUDIS5EF65KLMRQIIGY
 Type:    totp
 
 '''
+
+EXPECTED_STDOUT_FROM_EXAMPLE_EXPORT_PNG =\
+'''Name:    Test1:test1@example1.com
+Secret:  JBSWY3DPEHPK3PXP
+Issuer:  Test1
+Type:    totp
+
+Name:    Test2:test2@example2.com
+Secret:  JBSWY3DPEHPK3PXQ
+Issuer:  Test2
+Type:    totp
+
+Name:    Test3:test3@example3.com
+Secret:  JBSWY3DPEHPK3PXR
+Issuer:  Test3
+Type:    totp
+
+'''