switch to using a token based authentication system fixing #57

database/db_commands.py

@@ -3,6 +3,8 @@ from database.db_connection import get_db, get_db
 import hashlib
 import os
 import bcrypt
+import secrets
+import random
 
 
 def register_user(username, email, password, account_type="private_user"):
@@ -12,27 +14,59 @@ def register_user(username, email, password, account_type="private_user"):
             (username, email)
         ).fetchone()
         if existing:
-            raise ValueError("A user with that username or email already exists.")
+            return None
+        user_id = random.randint(1000000000, 9999999999)
         password_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt(14))
         conn.execute("""
-            INSERT INTO users (username, email, password_hash, account_type, created_at)
-            VALUES (?, ?, ?, ?, ?)
-        """, (username, email, password_hash, account_type, datetime.now()))
+            INSERT INTO users (user_id, username, email, password_hash, account_type, created_at)
+            VALUES (?, ?, ?, ?, ?, ?)
+        """, (user_id, username, email, password_hash, account_type, datetime.now()))
+        return create_token_for_user(user_id)
 
 
 def login_user(username, password):
+    user = None
     with get_db() as conn:
         row = conn.execute(
-            "SELECT * FROM users WHERE username = ?",
+            "SELECT user_id, password_hash FROM users WHERE username = ?",
             (username,)
         ).fetchone()
         if not row:
-            raise ValueError("Invalid username or password.")
+            return None
         user = dict(row)
         if not bcrypt.checkpw(password.encode(), user["password_hash"]):
-            raise ValueError("Invalid username or password.")
-        return {k: user[k] for k in ("user_id", "username", "email", "account_type")}
+            return None
+    return user and create_token_for_user(user["user_id"])
 
+def get_user_from_token(token):
+    with get_db() as conn:
+        row = conn.execute(
+            "SELECT * FROM users WHERE user_id = (SELECT user_id FROM tokens WHERE token = ? LIMIT 1) LIMIT 1",
+            (token,)
+        ).fetchone()
+        if not row:
+            return None
+        return dict(row)
+
+def create_token_for_user(user_id):
+    with get_db() as conn:
+        token = secrets.token_hex(64)
+        conn.execute("""
+            INSERT INTO tokens (
+                token,
+                user_id
+            )
+            VALUES (?, ?)
+        """, (token, user_id))
+        return token
+    print("hello")
+    
+def delete_token(token):
+    with get_db() as conn:
+        conn.execute("""
+            DELETE FROM tokens
+            WHERE token = ?
+        """, (token,))
 
 def add_species(english, latin, body, category, extinction_risk, image_id):
     with get_db() as conn:

database/db_connection.py

@@ -6,14 +6,20 @@ DATABASE_PATH = "database/database.db"
 db = sqlite3.connect(DATABASE_PATH, check_same_thread=False)
 db.row_factory = sqlite3.Row
 db.execute("""
-    CREATE TABLE IF NOT EXISTS users (
-        user_id INTEGER PRIMARY KEY AUTOINCREMENT,
-        username TEXT UNIQUE NOT NULL,
-        email TEXT UNIQUE NOT NULL,
-        password_hash TEXT NOT NULL,
-        account_type TEXT NOT NULL DEFAULT 'private_user',
-        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-    )
+CREATE TABLE IF NOT EXISTS users (
+    user_id INTEGER PRIMARY KEY,
+    username TEXT UNIQUE NOT NULL,
+    email TEXT UNIQUE NOT NULL,
+    password_hash TEXT NOT NULL,
+    account_type TEXT NOT NULL DEFAULT 'private_user',
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+)
+""")
+db.execute("""
+CREATE TABLE IF NOT EXISTS tokens (
+    token TEXT PRIMARY KEY,
+    user_id INTEGER
+)
 """)
 db.commit()
 
@@ -23,7 +29,8 @@ def get_db():
     try:
         yield conn
         db.commit()
-    except:
+    except Exception:
         db.rollback()
+        raise
     finally:
         conn.close()

main.py

@@ -1,7 +1,8 @@
-from flask import Flask
+from flask import Flask, request, g
 import os
 import importlib.util
 from database.db_init import init_database
+from database.db_commands import get_user_from_token
 
 def import_pages(route, app):
     for root, _, files in os.walk(route):
@@ -31,7 +32,15 @@ def import_pages(route, app):
 
 
 app = Flask(__name__, static_folder='static', static_url_path='')
-app.secret_key = "59474221-13e4-43dd-bc07-46a8d051d456"
+
+@app.before_request
+def load_user():
+    token = request.cookies.get("token")
+    g.user = token and get_user_from_token(token)
+
+@app.context_processor
+def inject_user():
+    return {"user": g.user}
 
 init_database()
 import_pages("routes", app)

routes/accounts/accounts.py

@@ -1,4 +1,4 @@
-from flask import render_template, Blueprint, request, redirect, url_for, session
+from flask import render_template, Blueprint, request, redirect, url_for, Response, g
 from database.db_commands import register_user, login_user
 
 page = Blueprint("login", __name__, url_prefix="")
@@ -6,27 +6,27 @@ page = Blueprint("login", __name__, url_prefix="")
 
 @page.route("/login", methods=["GET", "POST"])
 def login():
-    if "user_id" in session:
+    if g.user:
         return redirect(url_for("login.dashboard"))
     error = None
     if request.method == "POST":
         username = request.form.get("username", "").strip()
         password = request.form.get("password", "")
-        try:
-            user = login_user(username, password)
-            session.clear()
-            session["user_id"] = user["user_id"]
-            session["username"] = user["username"]
-            session["account_type"] = user["account_type"]
-            return redirect(url_for("login.dashboard"))
-        except ValueError as e:
-            error = str(e)
-    return render_template("accounts/login.jinja", error=error)
+        token = login_user(username, password)
+        if not token:
+            error = "invalid username or password"
+        else:
+            resp = redirect(url_for("login.dashboard"))
+            resp.set_cookie("token",token)
+            return resp
+    resp = Response(render_template("accounts/login.jinja", error=error))
+    if error: resp.status = 401
+    return resp
 
 
 @page.route("/register", methods=["GET", "POST"])
 def register():
-    if "user_id" in session:
+    if g.user:
         return redirect(url_for("login.dashboard"))
     error = None
     if request.method == "POST":
@@ -40,22 +40,27 @@ def register():
         elif len(password) < 8:
             error = "Password must be at least 8 characters."
         else:
-            try:
-                register_user(username, email, password, account_type)
-                return redirect(url_for("login.login"))
-            except ValueError as e:
-                error = str(e)
-    return render_template("accounts/user_registration.jinja", error=error)
+            token = register_user(username, email, password, account_type)
+            if not token:
+                error = "A user with that username or email already exists."
+            else:
+                resp = redirect(url_for("login.dashboard"))
+                resp.set_cookie("token",token)
+                return resp
+    resp = Response(render_template("accounts/user_registration.jinja", error=error))
+    if error: resp.status = 401
+    return resp
 
 
 @page.route("/logout")
 def logout():
-    session.clear()
-    return redirect(url_for("login.login"))
+    resp = redirect(url_for("login.login"))
+    resp.delete_cookie("token")
+    return resp
 
 
 @page.route("/dashboard")
 def dashboard():
-    if "user_id" not in session:
+    if not g.user:
         return redirect(url_for("login.login"))
     return render_template("accounts/dashboard.jinja")

templates/accounts/dashboard.jinja

@@ -5,7 +5,7 @@
 {% from 'button.jinja' import button %}
 
 <section id="hero" style="padding: 40px 20px; text-align: center;">
-    <h1>Welcome back, {{ session.username }}.</h1>
+    <h1>Welcome back, {{ user.username }}.</h1>
     <p style="max-width: 700px; margin: 0.5rem auto;">Your conservation dashboard.</p>
 </section>
 
@@ -40,7 +40,7 @@
     {% endif %}
 
     {#Teacher#}
-    {% if session.account_type == "teacher" %}
+    {% if user.account_type == "teacher" %}
     <div class="container-full" style="margin-bottom: 20px;">
         <h2>Classes</h2>
         <p>View your classes or set up a new one for your students.</p>
@@ -51,7 +51,7 @@
     {% endif %}
 
     {#Community Leader#}
-    {% if session.account_type == "community_leader" %}
+    {% if user.account_type == "community_leader" %}
     <div class="container-full" style="margin-bottom: 20px;">
         <h2>Programmes</h2>
         <p>Create a new conservation programme for others to join.</p>