From 9d59b8882b99694e6c0335fe30df4cf738fbbac6 Mon Sep 17 00:00:00 2001
From: magmaus3 <magmaus3@disroot.org>
Date: Fri, 7 Jul 2023 19:20:32 +0200
Subject: [PATCH] Add more security

still insecure but slightly less :3
---
 customiwmserver/database.py | 28 +++++++++++++++++++++++-----
 customiwmserver/main.py     |  7 +++++--
 2 files changed, 28 insertions(+), 7 deletions(-)

diff --git a/customiwmserver/database.py b/customiwmserver/database.py
index ca31e6d..4a22017 100644
--- a/customiwmserver/database.py
+++ b/customiwmserver/database.py
@@ -1,11 +1,13 @@
 from datetime import datetime
 from typing import Literal
 from pymongo import MongoClient
+from itsdangerous.serializer import Serializer
+from itsdangerous import TimestampSigner
+from os import environ, urandom
+import json
 
-
+# Database
 client = MongoClient("mongodb://root:catboys@mongo:27017")
-
-
 db = client["IWM_CustomServer_DB"]
 
 user_collection = db.users
@@ -15,6 +17,12 @@ reports_collection = db.reports
 general_collection = db.general
 admin_log_collection = db.admin_log
 
+# Auth token stuff
+
+SECRET_KEY = environ.get("SECRET_KEY", urandom(32).hex())
+# s = TimestampSigner(SECRET_KEY)
+s = TimestampSigner(SECRET_KEY)
+
 
 def LogAdminAction(
     action_type: str, action_data: dict, UserID: int = None, success: bool = True
@@ -28,6 +36,13 @@ def LogAdminAction(
         }
     )
 
+def login_auth_check(username: str, password: str):
+    # FIXME: This function currently DOES NOT perform any authentication.
+    # This means that ANYONE knowing the username could perform actions as the user.
+    query = user_collection.find_one({"Username": username})
+    if not query:
+        return False, "noauth"
+    return True, query
 
 def auth_check(Authorization) -> (tuple[Literal[False], Literal["noauth"]] | tuple[Literal[True], dict]):
     """Checks credentials.
@@ -37,12 +52,15 @@ def auth_check(Authorization) -> (tuple[Literal[False], Literal["noauth"]] | tup
         - False if wrong username or password 
         - True, [dict] if correct
     """
-    # FIXME: This function currently DOES NOT perform any authentication.
+    # FIXME (not relevant here anymore): This function currently DOES NOT perform any authentication.
     # This means that ANYONE knowing the username could perform actions as the user.
     if Authorization is None:
         return False, "noauth"
-    username, password = Authorization.split(":")
 
+    # In this case I assume that the server already authenticated the user,
+    # and signed the token. 
+    username = s.unsign(Authorization).decode()
+    print("DBG[auth_check]:", username)
     query = user_collection.find_one({"Username": username})
     if not query:
         return False, "noauth"
diff --git a/customiwmserver/main.py b/customiwmserver/main.py
index d9c8b17..7b1d760 100644
--- a/customiwmserver/main.py
+++ b/customiwmserver/main.py
@@ -39,11 +39,14 @@ async def http_exception_handler(request, exc):
 async def login(username: str = Form(), password: str = Form(), version: str = Form()):
     """User login"""
     hook.execute_hooks("player_login", username=username)
-    auth = db.auth_check(username + ":" + password)
+    auth = db.login_auth_check(username, password) 
+    # auth = db.auth_check(username + ":" + password)
     if not auth[0]:
         raise HTTPException(403, detail="Wrong username or password.")
     else:
-        return {"token": username + ":" + password, "userId": auth[1]["ID"]}
+        token = db.s.sign(username)
+        return {"token": token, "userId": auth[1]["ID"]}
+        # return {"token": username + ":" + password, "userId": auth[1]["ID"]}
 
 
 @app.put("/api/v1/user")