raps.py

@@ -1,80 +1,51 @@
 #!/usr/bin/env python3
 
-import argparse
-import datetime
-import hashlib
-import os
 import random
 import secrets
-import string
+import os
-import sys
+import hashlib
 
-DEFAULT_SECRET_PATH = 'secret'
+SECRET_FILE = 'secret'
 
 def create_secret():
-    """returns 128 bytes of randomness as hex bytes"""
     return secrets.token_hex(128)
 
 def generate_password():
-    """returns a random 24 char password"""
+    import string
     alphabet = string.ascii_letters + string.digits + '!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}'
     return ''.join(random.choice(alphabet) for i in range(24))
 
 def month_timestamp():
-    """returns the unix timestamp of the first of the current month"""
+    import datetime
     today = datetime.datetime.now()
     start_of_month = datetime.datetime(today.year, today.month, 1, 0, 0, 0, 0)
     return int(start_of_month.timestamp())
 
 def generate_seed(secret):
-    """returns an rng seed based on the current date"""
+    m = hashlib.sha256()
-    hash_object = hashlib.sha256()
+    m.update(bytes.fromhex(secret))
-    hash_object.update(bytes.fromhex(secret))
+    m.update(str(month_timestamp()).encode())
-    hash_object.update(str(month_timestamp()).encode())
+    return m.digest()
-    return hash_object.digest()
-
-def generate_secret(secret_path):
-    """generates a secret and writes it as a file"""
-    print('Generating a new secret... ', end='')
-    secret = create_secret()
-    with open(secret_path, 'w') as file:
-        file.write(secret)
-    print('done. Send this to the other party:')
-    print(secret)
 
 def main():
-    parser = argparse.ArgumentParser(
+    import argparse
-                prog='Remote Admin Password Solution',
+    parser = argparse.ArgumentParser(prog='Remote Admin Password Solution', description='Generate rotating passwords based on a shared secret')
-                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()
 
-    if args.secret:
-        secret_path = os.path.realpath(args.secret)
-    else:
-        secret_path = os.path.realpath(DEFAULT_SECRET_PATH)
-
     secret = ''
-    if args.new_secret:
+    if not os.path.isfile(SECRET_FILE) or args.new_secret:
-        if os.path.isfile(secret_path):
+        print('Generating a new secret... ', end='')
-            choice = ''
+        secret = create_secret()
-            while choice.lower() != 'y' and choice.lower() != 'n':
+        with open(SECRET_FILE, 'w') as f:
-                choice = input(f'Secret {secret_path} exists! Overwrite? [y/n]: ')
+            f.write(secret)
-            if choice.lower() == 'y':
+        print('done. Send this to the other party:')
-                generate_secret(secret_path)
+        print(secret)
+
     else:
-            generate_secret(secret_path)
+        secret = open(SECRET_FILE, 'r').read().strip()
-
+        print('Secret loaded.')
-    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))