not working trab 4

trabalho4/tgs.py

@@ -3,14 +3,19 @@ from common import *
 import socket
 from AES import AESCipher
 import os
+import base64
 
 HOST = "127.0.0.1"
 PORT = TGS_PORT
 
+TGS_KEY = os.urandom(32)
+
 def process_message(message):
     print(message)
     message = message.split("||")
 
+    if message[0] == "getkey":
+        return f"{base64.b64encode(TGS_KEY)}"
     if message[0] == "request":
         """
             T_c_tgs = {ID_C + T_R + K_c_tgs}K_tgs
@@ -20,13 +25,12 @@ def process_message(message):
             T_c_s = {ID_C + T_A + K_c_s}K_s
         """
         body = message[1].split("<>")
-        K_tgs = C_K_tgs
+        K_tgs = TGS_KEY
 
         M3 = body[0]
         M3 = M3.split("<>")
         M3_inner = M3[0]
         T_c_tgs = body[1]
-        print('T_c_tgs', type(T_c_tgs), T_c_tgs)
 
         M3_AES = AESCipher(K_tgs)
         T_c_tgs = M3_AES.decrypt(T_c_tgs)
@@ -46,7 +50,7 @@ def process_message(message):
         T_A = 600 if T_R > 600 else T_R
         T_c_s = f"{ID_C}<>{T_A}<>{K_c_s}"
 
-        K_s = C_K_s
+        K_s = get_ks()
         T_c_s_AES = AESCipher(K_s)
         T_c_s = T_c_s_AES.encrypt(T_c_s)
 
@@ -57,6 +61,18 @@ def process_message(message):
         M4 = f"{M4_inner}<>{T_c_s}"
         return M4
 
+def get_ks():
+    message = f"getkey"
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.connect((HOST, SERVICE_PORT))
+        s.sendall(message.encode())
+
+        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:
         s.bind((HOST, PORT))
@@ -74,4 +90,8 @@ def main():
                 conn.sendall(response.encode())
 
 if __name__ == "__main__":
+    if not os.path.exists("tgs_db.data"):
+        with open("tgs_db.data", "w") as f:
+            pass
+
     main()