not working trab 4

trabalho4/as.py

@@ -3,6 +3,7 @@ from hashlib import sha256
 import socket
 from AES import AESCipher
 import os
+import base64
 
 HOST = "127.0.0.1"
 PORT = AS_PORT
@@ -16,7 +17,7 @@ def process_message(message):
         username = body[0]
         password = body[1]
         password = sha256(password.encode()).hexdigest()
-        key = message[3]
+        key = os.urandom(32)
         return register_user(username, password, key)
     elif message[0] == "login":
         body = message[1].split("<>")
@@ -34,16 +35,17 @@ def process_message(message):
         body = message[1].split("<>")
 
         ID_C = body[0]
-        key = get_key(ID_C)
+        key = get_ktgs()
         AES = AESCipher(key)
-        K_tgs = C_K_tgs
-        M2_AES = AESCipher(K_tgs)
 
         inner_message = AES.decrypt(body[1])
         inner_message = inner_message.split("<>")
         T_R = inner_message[1]
         N1 = inner_message[2]
 
+        # generate m2
+        K_tgs = get_ktgs()
+        M2_AES = AESCipher(K_tgs)
         K_c_tgs = os.urandom(16)
         T_c_tgs = f"{ID_C}<>{T_R}<>{K_c_tgs}"
         T_c_tgs = M2_AES.encrypt(T_c_tgs)
@@ -51,14 +53,13 @@ def process_message(message):
         M2 = f"{K_c_tgs}<>{N1}"
         M2 = AES.encrypt(M2)
         M2 = f"{M2}<>{T_c_tgs}"
-        print('sending m2', M2)
         return M2
     else:
         print("Invalid message")
 
 def register_user(username, password, key):
     with open("users.data", "a") as f:
-        f.write(f"{username}<>{password}<>{key.encode()}\n")
+        f.write(f"{username}<>{password}<>{base64.b64encode(key)}\n")
     return f"success<>{key}"
 
 def login(username, password):
@@ -69,21 +70,23 @@ def login(username, password):
             user, pw, key = line.split("<>")
             if user == username and pw == password:
                 key = eval(key)
-                print(f"Login succeeded!")
+                key = base64.b64decode(key)
+                print(f"Login succeeded [{username}]")
                 return f"success<>{key}"
-        print("Login failed!")
+        print(f"Login failed! [{username}]")
         return "failure"
 
-def get_key(username):
-    return C_User_Key
+def get_ktgs():
+    message = f"getkey"
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.connect((HOST, TGS_PORT))
+        s.sendall(message.encode())
 
-    with open("users.data", "r", newline='\n') as f:
-        for line in f:
-            line = line.strip()
-            line = line.replace("\n", "")
-            user, _, key = line.split("<>")
-            if user == username:
-                return eval(key)
+        data = s.recv(4096)
+        data = data.decode()
+        key = eval(data)
+        key = base64.b64decode(key)
+        return key
 
 def main():
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: