finalizando trabalho 4

trabalho4/as.py

@@ -1,23 +1,23 @@
 from common import *
 from hashlib import sha256
-import socket
 from AES import AESCipher
+import socket
 import os
 import base64
+import time
 
 HOST = "127.0.0.1"
 PORT = AS_PORT
 
 def process_message(message):
-    print(message)
     message = message.split("||")
-
     if message[0] == "register":
         body = message[1].split("<>")
         username = body[0]
         password = body[1]
         password = sha256(password.encode()).hexdigest()
         key = os.urandom(32)
+        print(f"Registering user [{username}]")
         return register_user(username, password, key)
     elif message[0] == "login":
         body = message[1].split("<>")
@@ -35,8 +35,8 @@ def process_message(message):
         body = message[1].split("<>")
 
         ID_C = body[0]
-        key = get_ktgs()
-        AES = AESCipher(key)
+        Kc = get_user_key(ID_C)
+        AES = AESCipher(Kc)
 
         inner_message = AES.decrypt(body[1])
         inner_message = inner_message.split("<>")
@@ -76,6 +76,18 @@ def login(username, password):
         print(f"Login failed! [{username}]")
         return "failure"
 
+def get_user_key(username):
+    with open("users.data", "r", newline='\n') as f:
+        for line in f:
+            line = line.strip()
+            line = line.replace("\n", "")
+            user, pw, key = line.split("<>")
+            if user == username:
+                key = eval(key)
+                key = base64.b64decode(key)
+                return key
+    return None
+
 def get_ktgs():
     message = f"getkey"
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -95,7 +107,6 @@ def main():
         while True:
             conn, addr = s.accept()
             with conn:
-                print(f"Connected by {addr}")
                 data = conn.recv(4096)
                 data = data.decode()