#!/usr/bin/env python import asyncio import json, random, re, math, time, traceback from websockets.exceptions import ConnectionClosed, ConnectionClosedOK from websockets.asyncio.server import serve from db import Database userdb = Database("users") userdb.createTable("NAMES",{"id":"int","name":"string","creation_date":"int","last_login":"int","last_namechange":"int"}) userdb.createTable("PASSWORDS",{"id":"int","sum":"string"}) def makeUserID() -> int: count = userdb.rowAmount("NAMES") digits = max(4,1+int(math.log10(10+count))) uid=0 while uid==0 or userdb.hasPK("NAMES","id",uid) or userdb.hasPK("PASSWORDS","id",uid): uid = random.randint(0,(10**digits)-1) & 0x7FFFFFFFFFFFFFFF return uid def isValidUsername(name) -> bool: if len(name)<3 or len(name)>16: return False pattern = re.compile("^[a-z][a-z0-9-_]*") match = pattern.match(name) if match==None: return False return match.start()==0 and match.end()==len(name) def isValidChecksum(cs) -> bool: if len(cs)!=32: return False pattern = re.compile("^[0-9a-f]*") match = pattern.match(cs) if match==None: return False return match.start()==0 and match.end()==len(cs) def response(msg,**kwargs) -> str: return json.dumps({"type":"response","response":msg,**kwargs}) def clientError(msg,**kwargs) -> str: return json.dumps({"type":"error","error":msg,"server":False,**kwargs}) def serverError(msg,**kwargs) -> str: print(f"Server error: {msg}") return json.dumps({"type":"error","error":msg,"server":True,**kwargs}) sessionUsers = {} sessionCreatedAccount = [] def forgetSession(sid): if sid in sessionUsers: del sessionUsers[sid] if sid in sessionCreatedAccount: sessionCreatedAccount.remove(sid) async def handleCommand(ws,cmd): global sessionUsers print(f"recieved command: {cmd} (from {id(ws)})") if not "type" in cmd: return await ws.send(clientError("Malformed command")) uid = None if id(ws) in sessionUsers: uid = sessionUsers[id(ws)] match cmd["type"]: case "ping": return await ws.send(response("Pong!")) # accounts case "createuser": if uid != None: return await ws.send(clientError("You are already logged in.")) uid = makeUserID() name = cmd["name"] name = name.lower() if not isValidUsername(name): return await ws.send(clientError("Username is not valid.")) if userdb.findRow("NAMES","name = ?",(name,))!=None: return await ws.send(clientError(f"The username '{name}' is already taken.")) psum = cmd["pass"] psum = psum.lower() if not isValidChecksum(psum): return await ws.send(clientError("Recieved password has not been correctly checksummed.")) if id(ws) in sessionCreatedAccount: raise sessionCreatedAccount.append(id(ws)) create = time.time() userdb.insertRow("NAMES",{"id":uid,"name":cmd["name"],"creation_date":create,"last_login":create,"last_namechange":create}) userdb.insertRow("PASSWORDS",{"id":uid,"sum":psum}) return await ws.send(response({"uid":uid})) case "login": if uid != None: return await ws.send(clientError("You are already logged in.")) name = cmd["name"] psum = cmd["pass"] if not isValidUsername(name): return await ws.send(clientError("Username is not valid.")) res = userdb.findRow("NAMES","name = ?",(name,)) if res==None: return await ws.send(clientError("Invalid username or password.")) uid = res["id"] if not isValidChecksum(psum): return await ws.send(clientError("Recieved password has not been correctly checksummed.")) dbsum = userdb.rowFromPK("PASSWORDS","id",uid)["sum"] if psum.lower() != dbsum.lower(): return await ws.send(clientError("Invalid username or password.")) print(f"session {id(ws)} has logged in as {name} ({uid})") sessionUsers[id(ws)] = uid userdb.updateFromPK("NAMES","id",uid,{"last_login":time.time()}) return await ws.send(response("success")) case "logout": if uid == None: return await ws.send(clientError("You are already logged out.")) del sessionUsers[id(ws)] return await ws.send(response("success")) case "status": if uid == None: return await ws.send(response({"login":False})) name = userdb.rowFromPK("NAMES","id",uid)["name"] return await ws.send(response({"login":True,"uid":uid,"name":name})) case "changename": if uid == None: return await ws.send(clientError("You are not logged in.")) name = cmd["name"] if not isValidUsername(name): return await ws.send(clientError("Username is not valid.")) if userdb.findRow("NAMES","name = ?",(name,))!=None: return await ws.send(clientError(f"The username '{name}' is already taken.")) lastchange = userdb.rowFromPK("NAMES","id",uid)["last_namechange"] if lastchange == None: lastchange = 0 if (lastchange+7*24*60*60)>time.time(): return await ws.send(clientError("You can only change your username after 7 days.")) userdb.updateFromPK("NAMES","id",uid,{"name":name,"last_namechange":time.time()}) return await ws.send(response("success")) case "changepass": if uid == None: return await ws.send(clientError("You are not logged in.")) oldpass = cmd["oldpass"].lower() newpass = cmd["newpass"].lower() if not (isValidChecksum(oldpass) and isValidChecksum(newpass)): return await ws.send(clientError("Recieved password has not been correctly checksummed.")) dbpass = userdb.rowFromPK("PASSWORDS","id",uid)["sum"].lower() if oldpass != dbpass: return await ws.send(clientError("Invalid username or password.")) userdb.updateFromPK("PASSWORDS","id",uid,{"sum":newpass}) return await ws.send(response("success")) case _: return await ws.send(clientError(f"Command '{cmd["type"]}' does not exist.")) async def serveClient(ws): print("serve client",id(ws)) await ws.send(json.dumps({"type":"connected"})) while True: try: cmdtxt = await ws.recv() if type(cmdtxt)==bytes: cmdtxt=cmdtxt.decode("utf8") if type(cmdtxt)!=str: continue cmd = json.loads(cmdtxt) try: await handleCommand(ws,cmd) except Exception as e: traceback.print_exception(e) await ws.send(serverError(f"Cannot process command ({cmd["type"]})")) except json.decoder.JSONDecodeError: await ws.send(serverError("Cannot process JSON.")) except ConnectionClosedOK: print(f"connection {id(ws)} has closed") forgetSession(id(ws)) break except ConnectionClosed as e: print(f"connection {id(ws)} has closed with a non-OK exit code:") traceback.print_exception(e) break except Exception as e: traceback.print_exception(e) await ws.send(serverError("Cannot parse command.")) async def main(): async with serve(serveClient, "", 30341) as server: await server.serve_forever() if __name__ == "__main__": print("now serving") asyncio.run(main())