From 127fcfa9645a4e516a4d5ced705686b820bfc9a9 Mon Sep 17 00:00:00 2001
From: Tracker-Friendly <jliwin98@danwin1210.de>
Date: Sun, 31 Mar 2024 12:38:29 +0100
Subject: [PATCH] PKCE support

---
 main       | 35 ++++++++++++++++++++++++++++++++---
 schema.sql |  4 +++-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/main b/main
index ae72f5d..27203c2 100644
--- a/main
+++ b/main
@@ -1,4 +1,6 @@
 #!/usr/bin/python3
+import hashlib
+import base64
 import jwt
 import os
 import sqlite3
@@ -31,6 +33,12 @@ if SECRET_KEY == "supersecretkey" or SECRET_KEY == "placeholder":
 app = Quart(__name__)
 app.config["SECRET_KEY"] = SECRET_KEY
 
+# Hash creation function
+def sha256_base64(s: str) -> str:
+    hashed = hashlib.sha256(s.encode()).digest()
+    encoded = base64.urlsafe_b64encode(hashed).rstrip(b'=').decode()
+    return encoded
+
 # Database functions
 def get_db_connection():
     conn = sqlite3.connect("database.db")
@@ -246,6 +254,8 @@ async def apiauthenticate():
         data = await request.get_json()
         secretKey = data["secretKey"]
         appId = data["appId"]
+        code = data["code"]
+        codemethod = data["codemethod"]
 
         userCookie = get_session(secretKey)
         user = get_user(userCookie["id"])
@@ -282,8 +292,8 @@ async def apiauthenticate():
 
         nextjwt_token = jwt.encode(datatemplate2, SECRET_KEY, algorithm='HS256')
 
-        conn.execute("INSERT INTO logins (appId, secret, nextsecret, code, nextcode, creator, openid, nextopenid) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
-                (str(appId), str(secretkey), str(secrets.token_hex(512)), str(secrets.token_hex(512)), str(secrets.token_hex(512)), int(user["id"]), str(jwt_token), str(nextjwt_token)))
+        conn.execute("INSERT INTO logins (appId, secret, nextsecret, code, nextcode, creator, openid, nextopenid, pkce, pkcemethod) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+                (str(appId), str(secretkey), str(secrets.token_hex(512)), str(secrets.token_hex(512)), str(secrets.token_hex(512)), int(user["id"]), str(jwt_token), str(nextjwt_token), str(code), str(codemethod)))
 
         conn.commit()
         conn.close()
@@ -301,13 +311,32 @@ async def apitokenexchange():
         appId = data["client_id"]
         code = data["code"]
 
+        if "code_verifier" in data:
+            code_verify = data["code_verifier"]
+            verifycode = true
+        else:
+            verifycode = false
+
         conn = get_db_connection()
 
         # Fetch required data in a single query
-        oauth_data = conn.execute("SELECT appId, secret FROM oauth WHERE appId = ?", (str(appId),)).fetchone()
+        oauth_data = conn.execute("SELECT appId, secret, pkce, pkcemethod FROM oauth WHERE appId = ?", (str(appId),)).fetchone()
         if not oauth_data or oauth_data["appId"] != appId or oauth_data["secret"] != secret:
             return {}, 401
 
+        if verifycode:
+            if str(oauth_data["pkce"]) == "none":
+                return 400
+            else:
+                if str(oauth_data["pkcemethod"]) == "S256":
+                    if str(sha256_base64(code_verify)) != str(oauth_data["code"]):
+                        return 403
+                elif str(oauth_data["pkcemethod"]) == "plain":
+                    if str(code_verify) != str(oauth_data["code"]):
+                        return 403
+                else:
+                    return 501
+
         newkey = str(secrets.token_hex(512))
         conn.execute("UPDATE logins SET secret = ?, nextsecret = ? WHERE appId = ? AND secret = ?", (str(newkey), str(secrets.token_hex(512)), str(appId), str(secret)))
 
diff --git a/schema.sql b/schema.sql
index fd22353..0c2f4ec 100644
--- a/schema.sql
+++ b/schema.sql
@@ -31,7 +31,9 @@ CREATE TABLE logins (
     code TEXT NOT NULL,
     nextcode TEXT NOT NULL,
     creator INTEGER NOT NULL,
-    openid TEXT NOT NULL
+    openid TEXT NOT NULL,
+    pkce TEXT NOT NULL,
+    pkcemethod TEXT NOT NULL
 );
 
 CREATE TABLE oauth (