safety checks

raps.py

@@ -3,9 +3,10 @@
 import random
 import secrets
 import os
+import sys
 import hashlib
 
-SECRET_FILE = 'secret'
+DEFAULT_SECRET_PATH = 'secret'
 
 def create_secret():
     return secrets.token_hex(128)
@@ -27,28 +28,44 @@ def generate_seed(secret):
     m.update(str(month_timestamp()).encode())
     return m.digest()
 
+def generate_secret(secret_path):
+    print('Generating a new secret... ', end='')
+    secret = create_secret()
+    with open(secret_path, 'w') as f:
+        f.write(secret)
+    print('done. Send this to the other party:')
+    print(secret)
+
 def main():
     import argparse
-    parser = argparse.ArgumentParser(prog='Remote Admin Password Solution', description='Generate rotating passwords based on a shared secret')
+    parser = argparse.ArgumentParser(
+                prog='Remote Admin Password Solution',
+                description='Generate rotating passwords based on a shared secret'
+    )
     parser.add_argument('--secret', action='store', help='path to secret file')
     parser.add_argument('--new-secret', action='store_true', help='generate a new secret')
 
     args = parser.parse_args()
 
-    secret_path = args.secret if args.secret else SECRET_FILE
+    secret_path = args.secret if args.secret else DEFAULT_SECRET_PATH
 
     secret = ''
-    if not os.path.isfile(secret_path) or args.new_secret:
-        print('Generating a new secret... ', end='')
-        secret = create_secret()
-        with open(secret_path, 'w') as f:
-            f.write(secret)
-        print('done. Send this to the other party:')
-        print(secret)
-
+    if args.new_secret:
+        if os.path.isfile(secret_path):
+            choice = ''
+            while choice.lower() != 'y' and choice.lower() != 'n':
+                choice = input(f'Secret {secret_path} exists! Overwrite? [y/n]: ')
+            if choice.lower() == 'y':
+                generate_secret(secret_path)
         else:
-        secret = open(SECRET_FILE, 'r').read().strip()
-        print('Secret loaded.')
+            generate_secret(secret_path)
+
+    try:
+        secret = open(secret_path, 'r').read().strip()
+    except FileNotFoundError:
+        print(f'ERROR: Secret file {secret_path} could not be found.')
+        sys.exit(1)
+    print(f'Secret {secret_path} loaded.')
 
     random.seed(generate_seed(secret))