working build, added strength checker

main.py

@@ -10,6 +10,9 @@ from cryptography.hazmat.backends  import default_backend as backend
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
+import rsa
+from rsa import RsaKey
+
 DEFAULT_DATABASE_PATH = "db.txt"
 USAGES = {
         "new": 
@@ -42,6 +45,8 @@ CHARSET_MAP = {
         "s": SPECIAL,
 }
 
+SAFEST = 50
+
 def confirm(msg, yes, no):
     i = input(msg + f" ({yes}/{no})")
     i.strip()
@@ -89,6 +94,12 @@ class Password:
         self.key = key
         self.mode = mode
 
+    def __str__(self):
+        return f"{self.key}, {self.mode}, {repr(self.password)}"
+
+    def __repr__(self):
+        return self.__str__()
+
     def decrypt(password, key, mode):
         if mode == "csr":
             s = []
@@ -96,7 +107,12 @@ class Password:
                 s.append(chr(ord(l) - int(key)))
             return "".join(s)
         if mode == "rsa":
-            print("TODO: decrypt RSA")
+            # inSANE SECURITY VULNERABILITY 
+            # fix: figure out how to turn b64 hex byte string into bytes object
+            # DO NOT LET THIS GO INTO PROD!!!!!!!!!!!!!!!!!!!!!!!!
+            # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+            password = eval(password)
+            return rsa.decrypt_data(key.q, password)
 
     def encrypt(self):
         match self.mode:
@@ -111,12 +127,11 @@ class Password:
         return f"{self.key},{self.encrypt()}"
 
     def rsa_encrypt(self):
-        print("Todo! RSA encryption.")
-        return self.password.encode()
+        return rsa.encrypt_data(self.key.p, self.password)
 
     def csr_encrypt(self):
         s = []
-        for l in self.password:
+        for l in str(self.password):
             s.append(chr(ord(l) + int(self.key)))
         return "".join(s)
 
@@ -138,9 +153,13 @@ class State:
             print("Password does not match for this database, and as such it cannot be decrypted.")
             raise Exception
 
+        tail = lines[1]
+        s = tail.split(":")
+        print(s)
+        self.rsa_key = RsaKey(eval(s[0]),eval(s[1]))
+
         self.passwords = {}
         self.to_remove = []
-        self.rsa_key = None
         self.overwritten = []
         self.head = head
 
@@ -152,7 +171,7 @@ class State:
         key = None
 
         for i, line in enumerate(lines):
-            if i == 0:
+            if i < 2:
                 continue
             if line == "":
                 continue
@@ -166,15 +185,15 @@ class State:
 
                 mode = parts[1]
                 key = parts[2]
-                if mode == "rsa":
-                    self.rsa_key = key.encode()
+                if mode == 'rsa':
+                    key = self.rsa_key
                 continue
 
             # Expect: key,label,encrypted_password
             parts = line.split(",", 1)
-            # if len(parts) != 3:
-            #     print(f"Error parsing password on line {i}: {line}\n ^^^^ -> Expected three parts, got {len(parts)}.")
-            #     continue
+            if len(parts) != 2:
+                print(f"Error parsing password on line {i}: {line}\n ^^^^ -> Expected two parts, got {len(parts)}.")
+                continue
 
             if mode is None or key is None:
                 print(f"Error parsing password on line {i}: {line}\n ^^^^ -> Missing mode description above.")
@@ -215,7 +234,11 @@ class State:
 
     def get_password(self, label):
         if label in self.passwords:
-            return self.passwords[label].password
+            p = self.passwords[label]
+            if p.mode == "csr":
+                return p.password[2:-1]
+            else:
+                return p.password.decode()
         else:
             return b"[ERROR: no such label found]"
 
@@ -231,13 +254,16 @@ class State:
 
     def save(self):
         content = self.head + "\n"
+        content += f"{self.rsa_key}" + "\n"
         for l, p in self.passwords.items():
             if p.mode == "csr":
                 content += f"mode:csr:{p.key}\n"
                 content += f"{l},{p.encrypt()}\n"
         if self.rsa_key is not None:
-            content += f"mode:rsa:{self.rsa_key.encode('utf-8')}\n"
-        
+            content += f"mode:rsa:_\n"
+        for l, p in self.passwords.items():
+            if p.mode == "rsa":
+                content += f"{l},{p.encrypt()}\n"
         file = open(self.filepath, 'w')
         file.write(content)
         file.close()
@@ -300,8 +326,12 @@ def main():
 
             password = ask
             file = open(filepath, "xb")
-            key = password_encrypt("super secret hihi >_>".encode(), password)
-            file.write(key)
+            key = rsa.generate_rsa_key();
+            keysha = password_encrypt("super secret hihi >_>".encode(), password)
+            file.write(keysha)
+            file.write('\n'.encode())
+            print(repr(key))
+            file.write(f"{key}".encode())
             file.write('\n'.encode())
             file.close()
             return main()
@@ -452,6 +482,28 @@ def main():
             for over in state.overwritten:
                 print(f"{over}: [...]", end=' ')
             print("")
+
+        elif command == "strength" or command == "strchk" or command == 'howstrong':
+            if len(args) < 2:
+                print(f"Too few arguments for {command}. Usage: {USAGES['strength']}")
+                continue
+
+            pword = args[1]
+
+            if args[1] in state.passwords:
+                pword = str(state.passwords[args[1]].password)
+
+            # 'p' means percent (or coefficient)
+            # 'd' = diversity 
+            lp = len(pword) 
+            dp = len([ele for ele in SPECIAL if (ele in pword)])
+            dp += len([ele for ele in NUMBERS if (ele in pword)])
+            total_safety = lp + dp
+            percent = total_safety / SAFEST
+            if percent > 1.0:
+                print("Passwords exceeds max safety expected, perfect!")
+            else:
+                print(f"[{'*' * int(percent * 15)}{'_' * int((1-percent) * 15)}]")
         elif command == "new": # Form: 0:new 1:<mode> 2:<key> 3:<length> [ 4:<label> ]
             if len(args) < 2:
                 print(f"Too few arguments for {command}, usage: {USAGES['new']}") 
@@ -522,8 +574,8 @@ def main():
                         chars ^= set(c)
 
             password = generate_password(length, chars)
-            print(str(password))
-            state.add_password(label, password, mode, key)
+            print(password)
+            state.add_password(label, password.encode(), mode, key)
         elif command == "list" or command == "ls":
             if len(args) == 1:
                 # if state.passwords.len() >= 10:

rsa.py

+import math
+
+class RsaKey:
+    def __init__(self, p, q):
+        self.p = p
+        self.q = q
+
+    def __str__(self):
+        return f"{self.p}:{self.q}"
+
+    def __repr__(self):
+        return self.__str__()
+
+# Helper functions
+def gcd(a, b):
+    while b:
+        a, b = b, a % b
+    return a
+
+def modinv(a, m):
+    m0, x0, x1 = m, 0, 1
+    while a > 1:
+        q = a // m
+        m, a = a % m, m
+        x0, x1 = x1 - q * x0, x0
+    return x1 + m0 if x1 < 0 else x1
+
+# Generate RSA Key Pair
+def generate_rsa_key():
+    p = 61  # Example prime 1
+    q = 53  # Example prime 2
+    n = p * q
+    phi = (p - 1) * (q - 1)
+
+    e = 65537  # Common choice for e
+    while gcd(e, phi) != 1:
+        e += 2
+
+    d = modinv(e, phi)
+
+    public_key = (e, n)
+    private_key = (d, n)
+    return RsaKey(private_key, public_key)
+
+# Encrypt data
+def encrypt_data(public_key, data):
+    e, n = public_key
+    encrypted_data = b"".join(pow(byte, e, n).to_bytes((n.bit_length() + 7) // 8, byteorder="big") for byte in data)
+    return encrypted_data
+
+# Decrypt data
+def decrypt_data(private_key, encrypted_data):
+    d, n = private_key
+    key_size = (n.bit_length() + 7) // 8
+    decrypted_data = bytes(
+        pow(int.from_bytes(encrypted_data[i:i + key_size], byteorder="big"), d, n)
+        for i in range(0, len(encrypted_data), key_size)
+    )
+    return decrypted_data
+
+if __name__ == "__main__":
+    # Generate RSA keys
+    key = generate_rsa_key()
+    public_key = key.p
+    private_key = key.q
+
+    # Original data
+    original_data = b"Hello RSA!"
+    print(f"Original Data: {original_data}")
+
+    # Encrypt data
+    encrypted_data = encrypt_data(public_key, original_data)
+    print(f"Encrypted Data: {encrypted_data}")
+
+    # Decrypt data
+    decrypted_data = decrypt_data(private_key, encrypted_data)
+    print(f"Decrypted Data: {decrypted_data}")
+
+    # Verify decryption
+    assert original_data == decrypted_data, "Decrypted data does not match original!"