Try to implement authentication

auth/auth.py

@@ -0,0 +1,56 @@
+from threading import Lock
+import sqlite3
+import contextlib
+
+class Authentication:
+    def __init__(self, database):
+        self.db = database
+        self.create()
+
+    def connect(self):
+        return contextlib.closing(sqlite3.connect(self.db))
+
+    def create(self):
+        with self.connect() as c:
+            c.execute(
+                "CREATE TABLE IF NOT EXISTS users(profileId, token, user)")
+            c.execute(
+                "CREATE TABLE IF NOT EXISTS new(user)")
+
+    def check_token(self, profileId, token):
+        token = token.split(".")[-1]
+        if len(token) != 43:
+            return False
+        with self.connect() as c:
+            res = cur.execute("SELECT token FROM users WHERE profileId=?",
+                (profileId,)).fetchone()
+        if res is not None and res[0] != token:
+            return False
+        return True
+
+    def check_user(self, profileId, token, user):
+        with self.connect() as c:
+            res = c.execute("SELECT user FROM users WHERE profileId=?",
+                (profileId,)).fetchone()
+            if res is not None:
+                return res[0] == user
+
+            # Check if new user can be created
+            res = c.execute("SELECT 1 FROM new WHERE user=?",
+                (user,)).fetchone()
+            if res is None and res[0] != 1:
+                return False
+            c.execute("DELETE FROM new WHERE user=?", (user,))
+
+            # Create new user
+            c.execute("INSERT INTO users VALUES (?, ?, ?)",
+                (profileId, token, user))
+        return True
+
+    def get_user(self, profileId):
+        with self.connect() as c:
+            res = c.execute("SELECT user FROM users WHERE profileId=?",
+                (profileId)).fetchone()
+        if res is None or not isinstance(res[0], str):
+            return None
+        return res[0]

auth/server.py

@@ -8,6 +8,8 @@ from base64 import b64encode
 from os.path import getmtime
 import json
 
+from auth import Authentication
+
 userdata = {}
 _g_connections = {}
 _g_connections_lock = Lock()
@@ -136,14 +138,14 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
         token, profileId = params["sessionId"][0].split(":")[1:]
         serverId = params["serverId"][0]
 
-        with _g_connections_lock:
-            _g_connections[serverId] = profileId
-
-        if True:
-            self.send_ok(b"OK")
-        else:
+        if not self.auth.check_token(profileId, token):
             # Displayed directly to the user
             self.send_ok(b"Bad login")
+            return
+
+        with _g_connections_lock:
+            _g_connections[serverId] = (profileId, token)
+        self.send_ok(b"OK")
 
     def ygg_joinserver(self, body):
         data = json.loads(body.decode())
@@ -151,18 +153,19 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
         profileId = data["selectedProfile"]
         serverId = data["serverId"]
 
-        with _g_connections_lock:
-            _g_connections[serverId] = profileId
-
-        if True:
-            self.send_response(HTTPStatus.NO_CONTENT)
-            self.end_headers()
-        else:
+        if not self.auth.check_token(profileId, token):
             resp = b'{"error":"ForbiddenOperationException"}'
             self.send_response(HTTPStatus.FORBIDDEN)
             self.send_header("Content-Length", len(resp))
             self.end_headers()
             self.wfile.write(resp)
+            return
+
+        with _g_connections_lock:
+            _g_connections[serverId] = (profileId, token)
+
+        self.send_response(HTTPStatus.NO_CONTENT)
+        self.end_headers()
 
     def lgy_checkserver(self, params):
         user = params["user"][0]
@@ -172,8 +175,13 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
             if not serverId in _g_connections:
                 self.send_ok(b"NO")
                 return
+            profileId, token = _g_connections[serverId]
             del _g_connections[serverId]
 
+        if not self.auth.check_user(profileId, token, user):
+            self.send_ok(b"NO")
+            return
+
         self.send_ok(b"YES")
 
     def ygg_checkserver(self, params):
@@ -184,8 +192,14 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
             if not serverId in _g_connections:
                 self.send_response(HTTPStatus.NO_CONTENT)
                 self.end_headers()
+                return
+            profileId, token = _g_connections[serverId]
             del _g_connections[serverId]
-            profileId = _g_connections[serverId]
+
+        if not self.auth.check_user(profileId, token, user):
+            self.send_response(HTTPStatus.NO_CONTENT)
+            self.end_headers()
+            return
 
         data = self.get_profile(profileId)
         if data is None:
@@ -196,31 +210,30 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
         self.send_ok(json.dumps(data).encode())
 
     def get_profile(self, profileId):
-        if profileId not in userdata:
+        user = self.auth.get_user(profileId)
+        if user is None:
             return None
 
-        user = userdata[profileId]
-
         properties = []
 
-        check_skin = self.check_skin(user["name"])
-        check_cape = self.check_cape(user["name"])
+        check_skin = self.check_skin(user)
+        check_cape = self.check_cape(user)
 
         url = self.config["server_url"]
         textures = {}
         textures_stamp = 0
         if check_skin is not None:
             textures_stamp = max(textures_stamp, check_skin)
-            textures["SKIN"] = {"url": url + "/skin/%s.png" % user["name"]}
+            textures["SKIN"] = {"url": url + "/skin/%s.png" % user}
         if check_cape is not None:
             textures_stamp = max(textures_stamp, check_cape)
-            textures["CAPE"] = {"url": url + "/cape/%s.cape.png" % user["name"]}
+            textures["CAPE"] = {"url": url + "/cape/%s.cape.png" % user}
 
         if textures:
             prop = {
                 "timestamp": textures_stamp,
                 "profileId": profileId,
-                "profileName": user["name"],
+                "profileName": user,
                 "textures": textures
             }
             properties.append({
@@ -230,7 +243,7 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
 
         data = {
             "id": profileId,
-            "name": user["name"],
+            "name": user,
             "properties": properties,
             "profileActions": []
         }
@@ -267,8 +280,10 @@ if __name__ == "__main__":
     server_address = ("", 25564)
     config = ConfigParser()
     config.read("config.ini")
+    auth = Authentication("users.db")
     httpd = ThreadingHTTPServer(server_address, HTTPRequestHandler)
     httpd.config = config
+    httpd.auth = auth
     try:
         httpd.serve_forever()
     except KeyboardInterrupt: