Changes based on suggestions from scito

https://github.com/scito/extract_otp_secret_keys/pull/2#pullrequestreview-555264890

extract_otp_secret_keys.py

@@ -46,19 +46,27 @@ import base64
 import fileinput
 import sys
 from urllib.parse import parse_qs, urlencode, urlparse, quote
-
+from os import path, mkdir
+from re import sub, compile as rcompile
 import generated_python.google_auth_pb2
-import os,re,string
 
 arg_parser = argparse.ArgumentParser()
 arg_parser.add_argument('--verbose', '-v', help='verbose output', action='store_true')
-arg_parser.add_argument('--qr', '-q', help='print QR codes (otpauth://...)', action='store_true')
+arg_parser.add_argument('--qr', '-q', help='generate QR codes (otpauth://...)', action='store_true')
+arg_parser.add_argument('--saveqr', '-s', help='save QR code(s) as images to the "qr" folder', action='store_true')
+arg_parser.add_argument('--printqr', '-p', help='print the QR code(s) as text to the terminal', action='store_true')
 arg_parser.add_argument('infile', help='file or - for stdin (default: -) with "otpauth-migration://..." URLs separated by newlines, lines starting with # are ignored')
 args = arg_parser.parse_args()
 
 verbose = args.verbose
+saveqr = args.saveqr
+printqr = args.printqr
+
 if args.qr:
     from qrcode import QRCode
+    if not(saveqr or printqr):
+        print("Error: If you specify to generate a QR code, you must specify how you want to receive it. Please specify -s or -p. Use -h or --help for more information")
+        exit()
 
 # https://stackoverflow.com/questions/40226049/find-enums-listed-in-python-descriptor-for-protobuf
 def get_enum_name_by_number(parent, field_name):
@@ -69,12 +77,17 @@ def convert_secret_from_bytes_to_base32_str(bytes):
     return str(base64.b32encode(otp.secret), 'utf-8').replace('=', '')
 
 
-def print_qr(data, name):
+def save_qr(data, name):
     qr = QRCode()
     qr.add_data(data)
     img = qr.make_image(fill_color="black", back_color="white")
+    if verbose: print("Saving to {}".format(name))
     img.save(name)
         
+def print_qr(data):
+    qr = QRCode()
+    qr.add_data(data)
+    qr.print_tty()
 
 for line in (line.strip() for line in fileinput.input(args.infile)):
     if verbose: print(line)
@@ -88,9 +101,9 @@ for line in (line.strip() for line in fileinput.input(args.infile)):
     if verbose: print(payload)
 
     # pylint: disable=no-member
-    i = int(0)
+    i = 0
     for otp in payload.otp_parameters:
-        i+=1
+        i += 1
         print('\nName:   {}'.format(otp.name))
         secret = convert_secret_from_bytes_to_base32_str(otp.secret)
         print('Secret: {}'.format(secret))
@@ -102,9 +115,11 @@ for line in (line.strip() for line in fileinput.input(args.infile)):
         otp_url = 'otpauth://{}/{}?'.format('totp' if otp.type == 2 else 'hotp', quote(otp.name)) + urlencode(url_params)
         if args.qr:
             if verbose: print(otp_url)
-            if not(os.path.exists("qr")):os.mkdir("qr")
-            pattern = re.compile('[\W_]+')
-            file_otp_name = pattern.sub('', otp.name)
-            file_otp_issuer = pattern.sub('',otp.issuer)
-            if not(file_otp_issuer):print_qr(otp_url, "qr/{}-{}.png".format(i,file_otp_name))
-            if file_otp_issuer: print_qr(otp_url, "qr/{}-{}-{}.png".format(i,file_otp_name,file_otp_issuer))
+            if saveqr:
+                if not(path.exists("qr")): mkdir("qr")
+                pattern = rcompile('[\W_]+')
+                file_otp_name = pattern.sub('', otp.name)
+                file_otp_issuer = pattern.sub('', otp.issuer)
+                if not(file_otp_issuer): print_qr(otp_url, "qr/{}-{}.png".format(i, file_otp_name))
+                if file_otp_issuer: save_qr(otp_url, "qr/{}-{}-{}.png".format(i,file_otp_name, file_otp_issuer))
+            if printqr: print_qr(otp_url)