working build, basically feature complete, could use some bug testing

main.py

 import sys
+from typing import Optional
 import secrets
 import os
 import random
@@ -156,93 +157,40 @@ class State:
         file.close()
 
 
-def caesar_cipher(inp, offset):
-    pass
-
 def generate_password(length, charset):
     s = []
     for i in range(0, length):
         s.append(random.choice(list(charset)))
     return ''.join(s)
 
-def aes_256(inp):
-    pass
-
 def handle_args(db: str, password: str):
     state = State(db, password)
     return state
 
-def main():
-    filepath = ""
-    if len(sys.argv) < 2:
-        print(f"No database path supplied. Usage: {sys.argv[0]} <database path> <password>")
-        return
-    else:
-        filepath = sys.argv[1]
-    
-    if len(sys.argv) < 3:
-        print(f"No password supplied. Usage: {sys.argv[0]} <database path> <password>");
-        return
-
-    password = sys.argv[2]
-
-    try:
-        state = handle_args(filepath, password)
-    except FileNotFoundError:
-        print(f"Error while opening database file, {filepath} does not exist or cannot be opened.")
-        ask = input(f"Create {filepath}? (Y/N)")
-        ask = ask.strip()
-        ask = ask.lower()
-        if ask == "y":
-            file = open(filepath, "xb")
-            salt = secrets.token_bytes(16)
-            key = password_encrypt("super secret hihi >_>".encode(), password)
-            file.write(key)
-            file.write('\n'.encode())
-            file.close()
-            return main()
-        else:
-            return
-    except Exception as e:
-        print(f"Error parsing database {filepath}: {e}")
-        return
-
-    # TODO: make this serialized (somehow... :[)
-    cfg = {
-        "hints": True
-    }
-    print("Welcome to SmashPass. Type 'h' or 'help' for help")
-
-    while True:
-        command_input = input(">")
-        command_input = command_input.strip()
-
-        args = command_input.split(" ")
-        if len(args) == 0:
-            continue
-
+def handle_command(state: State, args: [str], cfg) -> Optional[bool]:
     command = args[0]
-
     if command == "q" or command == "quit" or command == "exit":
         if len(args) == 2:
             if args[1] == "f" or args[1] == "force" or args[1] == "nosave":
-                    if confirm("Are you sure you want to exit without saving?", "y", "n"):
-                        return
+                if len(args) > 2 and args[2] == "y":
+                    return True
+                elif confirm("Are you sure you want to exit without saving?", "y", "n"):
+                    return True
         else:
-                break
+            return False
     if command == "cfg" or command == "config" or command == "settings":
         if len(args) < 2:
             print(f"Settings: hints: {cfg['hints']}, ...")
-                continue
+            return None
 
         s = args[1].split("=")
         if len(s) < 2:
             print(f"Incorrect format for modifying program parameters. Usage: {USAGES['cfg']}")
-                continue
+            return None
 
         if s[0] not in cfg:
             print(f"Unknown cfg parameter '{s[0]}'. Type 'cfg' for a list of config parameters.")
-                continue
+            return None
         try:
             cfg[s[0]] = s[1]
         except:
@@ -257,15 +205,13 @@ def main():
             cmd = args[1]
             if not cmd in USAGES:
                 print("Unknown command, you may have mistyped it.")
-                    continue
+                return None
             print(USAGES[cmd])
     elif command == "echo":
         if len(args) < 2:
             print(f"Too few arguments for {command}, usage: {USAGES['echo']}") 
-                continue
+            return None
         print(" ".join(args[1:]))
-        elif command == "dbg":
-            print(vars(state))
     elif command == "save" or command == "s":
         try:
             state.save()
@@ -275,7 +221,7 @@ def main():
     elif command == "drop" or command == "delete" or command == "del" or command == "rm":
         if len(args) < 2:
             print(f"Too few arguments for command {command}, usage: {USAGES['drop']}")
-                continue
+            return None
         
         label = args[1].strip().lower()
         if label in state.to_remove:
@@ -294,13 +240,13 @@ def main():
     elif command == "undo" or command == "un":
         if len(state.to_remove) == 0:
             print("Already on latest change")
-                continue
+            return None
         ltst = state.to_remove.pop()
         print(f"Undone latest deletion: {ltst}")
     elif command == "restore" or command == "res" or command == "r":
         if len(args) < 2:
             print(f"Too few arguments for command {command}. Usage: {USAGES['restore']}")
-                continue
+            return None
 
         args[1].strip()
         og_label = args[1].lower()
@@ -327,7 +273,7 @@ def main():
     elif command == "rename" or command == "rn":
         if len(args) < 3:
             print(f"Too few args for {command}. Usage: {USAGES['rename']}.")
-                continue
+            return None
 
         if args[1] in state.passwords:
             state.passwords[args[2]] = state.passwords[args[1]]
@@ -346,12 +292,12 @@ def main():
     elif command == "new": # Form: new <length> [ <label> ]
         if len(args) < 2:
             print(f"Too few arguments for {command}, usage: {USAGES['new']}") 
-                continue
+            return None
         try:
             length = int(args[1])
         except:
             print(f"Error parsing length as number, usage: {USAGES['new']}")
-                continue
+            return None
 
         if len(args) >= 3:
             label = args[2]
@@ -370,7 +316,7 @@ def main():
                 if cfg['hints'] == True:
                     print("(Hint: use the 'restore' command to undo this action)")
             else:
-                    continue
+                return None
 
         charset = "all"
 
@@ -403,8 +349,67 @@ def main():
                 print("(Hint: type 'list <key>' to decrypt your chosen password")
         elif len(args) == 2:
             print(state.get_password(args[1]).decode('utf-8'))
+    return None
+
+def main():
+    filepath = ""
+    if len(sys.argv) < 2:
+        print(f"No database path supplied. Usage: {sys.argv[0]} <database path> <password>")
+        return
+    else:
+        filepath = sys.argv[1]
+    
+    if len(sys.argv) < 3:
+        print(f"No password supplied. Usage: {sys.argv[0]} <database path> <password>");
+        return
+
+    password = sys.argv[2]
+
+    try:
+        state = handle_args(filepath, password)
+    except FileNotFoundError:
+        print(f"Error while opening database file, {filepath} does not exist or cannot be opened.")
+        ask = input(f"Create {filepath}? (Y/N)")
+        ask = ask.strip()
+        ask = ask.lower()
+        if ask == "y":
+            file = open(filepath, "xb")
+            salt = secrets.token_bytes(16)
+            key = password_encrypt("super secret hihi >_>".encode(), password)
+            file.write(key)
+            file.write('\n'.encode())
+            file.close()
+            return main()
+        else:
+            return
+    except Exception as e:
+        print(f"Error parsing database {filepath}: {e}")
+        return
+
+    # TODO: make this serialized (somehow... :[)
+    cfg = {
+        "hints": True
+    }
+    print("Welcome to SmashPass. Type 'h' or 'help' for help")
 
+    while True:
+        command_input = input(">")
+        command_input = command_input.strip()
 
+        args = command_input.split(" ")
+        if len(args) == 0:
+            continue
+
+        should_immediately_exit = None
+        try:
+            should_immediately_exit = handle_command(state, args, cfg)
+        except e:
+            print(f"ERROR: handling the inputted command: {e}")
+        if should_immediately_exit is not None:
+            if should_immediately_exit:
+                return
+            else:
+                break
 
     state.save()
     print("Thanks for using!")