diff --git a/database/db_commands.py b/database/db_commands.py index ca43acd..a2b66ba 100644 --- a/database/db_commands.py +++ b/database/db_commands.py @@ -8,9 +8,10 @@ import random def register_user(username, email, password, account_type="private_user"): + email = email.strip().lower() with get_db() as conn: existing = conn.execute( - "SELECT user_id FROM users WHERE username = ? OR email = ?", + "SELECT user_id FROM users WHERE username = ? OR LOWER(email) = LOWER(?)", (username, email) ).fetchone() if existing: @@ -61,6 +62,49 @@ def create_token_for_user(user_id): return token print("hello") +def create_password_reset_request(email): + user = get_user_by_email(email) + if not user: + return None + reset_token = secrets.token_hex(64) + with get_db() as conn: + conn.execute(""" + INSERT INTO password_resets ( + token, + user_id, + created_at + ) + VALUES (?, ?, ?) + """, (reset_token, user["user_id"], datetime.now())) + return reset_token + + +def get_password_reset_request(token): + with get_db() as conn: + row = conn.execute( + "SELECT * FROM password_resets WHERE token = ? LIMIT 1", + (token,) + ).fetchone() + return dict(row) if row else None + + +def reset_password(token, password): + reset_request = get_password_reset_request(token) + if not reset_request: + return False + password_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt(14)) + with get_db() as conn: + conn.execute( + "UPDATE users SET password_hash = ? WHERE user_id = ?", + (password_hash, reset_request["user_id"]) + ) + conn.execute( + "DELETE FROM password_resets WHERE token = ?", + (token,) + ) + return True + + def delete_token(token): with get_db() as conn: conn.execute(""" @@ -68,6 +112,14 @@ def delete_token(token): WHERE token = ? """, (token,)) +def get_user_by_email(email): + with get_db() as conn: + row = conn.execute( + "SELECT * FROM users WHERE LOWER(email) = LOWER(?) LIMIT 1", + (email,) + ).fetchone() + return dict(row) if row else None + def add_species(english, latin, body, category, extinction_risk, image_id): with get_db() as conn: try: diff --git a/database/db_connection.py b/database/db_connection.py index 1becf89..b0b9f51 100644 --- a/database/db_connection.py +++ b/database/db_connection.py @@ -1,11 +1,51 @@ +import os import sqlite3 from contextlib import contextmanager DATABASE_PATH = "database/database.db" +DATABASE_URL = os.getenv("DATABASE_URL") +USE_POSTGRES = DATABASE_URL is not None -db = sqlite3.connect(DATABASE_PATH, check_same_thread=False) -db.row_factory = sqlite3.Row -db.execute(""" +if USE_POSTGRES: + import psycopg2 + import psycopg2.extras + + class DBCursor: + def __init__(self, cursor): + self._cursor = cursor + + def execute(self, sql, params=None): + if params is None: + params = () + sql = sql.replace("?", "%s") + return self._cursor.execute(sql, params) + + def executemany(self, sql, params=None): + if params is None: + params = () + sql = sql.replace("?", "%s") + return self._cursor.executemany(sql, params) + + def __getattr__(self, name): + return getattr(self._cursor, name) + + @contextmanager + def get_db(): + conn = psycopg2.connect(DATABASE_URL, cursor_factory=psycopg2.extras.RealDictCursor) + cursor = conn.cursor() + try: + yield DBCursor(cursor) + conn.commit() + except Exception: + conn.rollback() + raise + finally: + cursor.close() + conn.close() +else: + 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, username TEXT UNIQUE NOT NULL, @@ -15,22 +55,29 @@ CREATE TABLE IF NOT EXISTS users ( created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) -db.execute(""" + db.execute(""" CREATE TABLE IF NOT EXISTS tokens ( token TEXT PRIMARY KEY, user_id INTEGER ) """) -db.commit() + db.execute(""" +CREATE TABLE IF NOT EXISTS password_resets ( + token TEXT PRIMARY KEY, + user_id INTEGER NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP +) +""") + db.commit() -@contextmanager -def get_db(): - conn = db.cursor() - try: - yield conn - db.commit() - except Exception: - db.rollback() - raise - finally: - conn.close() \ No newline at end of file + @contextmanager + def get_db(): + conn = db.cursor() + try: + yield conn + db.commit() + except Exception: + db.rollback() + raise + finally: + conn.close() \ No newline at end of file diff --git a/database/db_init.py b/database/db_init.py index 65b6aba..1764976 100644 --- a/database/db_init.py +++ b/database/db_init.py @@ -1,4 +1,4 @@ -from database.db_connection import get_db, get_db +from database.db_connection import get_db from database.db_commands import register_user from datetime import datetime @@ -127,6 +127,7 @@ def init_database(): photoid TEXT ) """) + init_accounts_database() print("Database initialised.") seed_species() @@ -210,6 +211,19 @@ def init_accounts_database(): created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) + conn.execute(""" + CREATE TABLE IF NOT EXISTS tokens ( + token TEXT PRIMARY KEY, + user_id INTEGER + ) + """) + conn.execute(""" + CREATE TABLE IF NOT EXISTS password_resets ( + token TEXT PRIMARY KEY, + user_id INTEGER NOT NULL, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP + ) + """) print("Accounts database initialised.") diff --git a/main.py b/main.py index b7b0e8d..a7c7bb5 100644 --- a/main.py +++ b/main.py @@ -35,6 +35,7 @@ def import_pages(route, app): app = Flask(__name__, static_folder="static", static_url_path="") app.secret_key = "9279dc6c-d2bc-4f89-a1b3-fc1438a6910c" +init_database() @app.before_request def load_user(): diff --git a/requirements.txt b/requirements.txt index aef7c50..2fffba8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ itsdangerous==2.2.0 Jinja2==3.1.6 MarkupSafe==3.0.3 Werkzeug==3.1.5 +psycopg2-binary==2.9.10 diff --git a/routes/accounts/accounts.py b/routes/accounts/accounts.py index 4c77edc..84f8121 100644 --- a/routes/accounts/accounts.py +++ b/routes/accounts/accounts.py @@ -1,5 +1,11 @@ from flask import render_template, Blueprint, request, redirect, url_for, Response, g -from database.db_commands import register_user, login_user +from database.db_commands import ( + register_user, + login_user, + create_password_reset_request, + get_password_reset_request, + reset_password, +) page = Blueprint("login", __name__, url_prefix="") @@ -24,6 +30,58 @@ def login(): return resp +@page.route("/password-reset", methods=["GET", "POST"]) +def password_reset_request(): + if g.user: + return redirect(url_for("login.dashboard")) + message = None + reset_link = None + if request.method == "POST": + email = request.form.get("email", "").strip().lower() + token = create_password_reset_request(email) + message = ( + "If an account with that email exists, a password reset link has been created." + ) + if token: + reset_link = url_for("login.password_reset", token=token, _external=True) + return render_template( + "accounts/password_reset_request.jinja", + message=message, + reset_link=reset_link, + ) + + +@page.route("/password-reset/", methods=["GET", "POST"]) +def password_reset(token): + if g.user: + return redirect(url_for("login.dashboard")) + error = None + success = None + reset_request = get_password_reset_request(token) + if not reset_request: + error = "This reset link is invalid or has expired." + if request.method == "POST" and reset_request: + password = request.form.get("password", "") + password_confirm = request.form.get("password_confirm", "") + if password != password_confirm: + error = "Passwords do not match." + elif len(password) < 8: + error = "Password must be at least 8 characters." + else: + if reset_password(token, password): + success = "Your password has been reset. Please log in." + reset_request = None + else: + error = "Unable to reset your password. Please request a new link." + return render_template( + "accounts/password_reset.jinja", + error=error, + success=success, + token=token, + reset_request=reset_request, + ) + + @page.route("/register", methods=["GET", "POST"]) def register(): if g.user: diff --git a/templates/accounts/login.jinja b/templates/accounts/login.jinja index 1e2e2b6..0bfe5e4 100644 --- a/templates/accounts/login.jinja +++ b/templates/accounts/login.jinja @@ -21,6 +21,7 @@ + {% endblock %} \ No newline at end of file diff --git a/templates/accounts/password_reset.jinja b/templates/accounts/password_reset.jinja new file mode 100644 index 0000000..f4896ea --- /dev/null +++ b/templates/accounts/password_reset.jinja @@ -0,0 +1,35 @@ +{% extends "base.jinja" %} +{% block title %} Set New Password {% endblock %} +{% block head %} +{% endblock %} +{% block content %} +{% from 'button.jinja' import button %} +
+
+ {% if error %} +

{{ error }}

+ {% endif %} + {% if success %} +

{{ success }}

+

Return to login

+ {% elif reset_request %} +
+
+ + +
+
+ + +
+ +
+ {% else %} +

Please request a new password reset link.

+

Request another reset link

+ {% endif %} +
+
+{% endblock %} diff --git a/templates/accounts/password_reset_request.jinja b/templates/accounts/password_reset_request.jinja new file mode 100644 index 0000000..267c769 --- /dev/null +++ b/templates/accounts/password_reset_request.jinja @@ -0,0 +1,26 @@ +{% extends "base.jinja" %} +{% block title %} Password Reset {% endblock %} +{% block head %} +{% endblock %} +{% block content %} +{% from 'button.jinja' import button %} +
+
+ {% if message %} +

{{ message }}

+ {% endif %} + {% if reset_link %} +

+ Reset link: {{ reset_link }} +

+ {% endif %} +
+ + +
+ +
+
+{% endblock %}