118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
from common import *
|
|
from hashlib import sha256
|
|
from AES import AESCipher
|
|
import socket
|
|
import os
|
|
import base64
|
|
import time
|
|
|
|
HOST = "127.0.0.1"
|
|
PORT = AS_PORT
|
|
|
|
def process_message(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("<>")
|
|
username = body[0]
|
|
password = body[1]
|
|
password = sha256(password.encode()).hexdigest()
|
|
return login(username, password)
|
|
elif message[0] == "request":
|
|
"""
|
|
M1 = [ID_C + {ID_S + T_R + N1}Kc]
|
|
|
|
T_c_tgs = {ID_C + T_R + K_c_tgs}K_tgs
|
|
M2 = [{K_c_tgs + N_1}Kc + T_c_tgs]
|
|
"""
|
|
body = message[1].split("<>")
|
|
|
|
ID_C = body[0]
|
|
Kc = get_user_key(ID_C)
|
|
AES = AESCipher(Kc)
|
|
|
|
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)
|
|
|
|
M2 = f"{K_c_tgs}<>{N1}"
|
|
M2 = AES.encrypt(M2)
|
|
M2 = f"{M2}<>{T_c_tgs}"
|
|
return M2
|
|
else:
|
|
print("Invalid message")
|
|
|
|
def register_user(username, password, key):
|
|
with open("users.data", "a") as f:
|
|
f.write(f"{username}<>{password}<>{base64.b64encode(key)}\n")
|
|
return f"success<>{key}"
|
|
|
|
def login(username, password):
|
|
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 and pw == password:
|
|
key = eval(key)
|
|
key = base64.b64decode(key)
|
|
print(f"Login succeeded [{username}]")
|
|
return f"success<>{key}"
|
|
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:
|
|
s.connect((HOST, TGS_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))
|
|
s.listen()
|
|
while True:
|
|
conn, addr = s.accept()
|
|
with conn:
|
|
data = conn.recv(4096)
|
|
data = data.decode()
|
|
|
|
response = process_message(data)
|
|
|
|
conn.sendall(response.encode())
|
|
|
|
if __name__ == "__main__":
|
|
main() |