import datetime import hashlib import importlib import os import sqlite3 from contextlib import contextmanager from urllib.parse import urlparse from config import USE_POSTGRES, DATABASE_URL import config user_db_file_location = "database_file/users.db" note_db_file_location = "database_file/notes.db" image_db_file_location = "database_file/images.db" def _sha256_hex(value: str) -> str: return hashlib.sha256(value.encode()).hexdigest() def _is_postgres() -> bool: return bool(getattr(config, "USE_POSTGRES", False)) def _postgres_connect_kwargs() -> dict: database_url = getattr(config, "DATABASE_URL", "") or os.environ.get("DATABASE_URL", "") if not database_url: raise RuntimeError( "Postgres er aktivert (ENVIRONMENT=production/staging), men DATABASE_URL er ikke satt. " "Sett DATABASE_URL f.eks. til: postgresql://postgres:postgres@db:5432/postgres" ) parsed = urlparse(database_url) return { "host": parsed.hostname, "port": parsed.port or 5432, "dbname": (parsed.path or "/").lstrip("/") or "postgres", "user": parsed.username or "postgres", "password": parsed.password or "postgres", } def _pg_init_schema(conn) -> None: with conn.cursor() as cur: cur.execute( """ CREATE TABLE IF NOT EXISTS users ( id TEXT PRIMARY KEY, pw TEXT NOT NULL ); """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS notes ( user_id TEXT NOT NULL, timestamp TEXT NOT NULL, note TEXT NOT NULL, note_id TEXT PRIMARY KEY ); """ ) cur.execute( """ CREATE TABLE IF NOT EXISTS images ( uid TEXT PRIMARY KEY, owner TEXT NOT NULL, name TEXT NOT NULL, timestamp TEXT NOT NULL ); """ ) #if admin dont exist, create one with default password "adminadmin" cur.execute( """ INSERT INTO users (id, pw) VALUES (%s, %s) ON CONFLICT (id) DO NOTHING; """, ("ADMIN", _sha256_hex("adminadmin")), ) conn.commit() @contextmanager def _pg_conn(): psycopg2 = importlib.import_module("psycopg2") conn = psycopg2.connect(**_postgres_connect_kwargs()) try: _pg_init_schema(conn) yield conn finally: conn.close() @contextmanager def _sqlite_conn(db_path: str): conn = sqlite3.connect(db_path) try: yield conn finally: conn.close() @contextmanager def _conn_for(operation: str): # operation: 'users' | 'notes' | 'images' if _is_postgres(): with _pg_conn() as conn: yield conn return if operation == "users": with _sqlite_conn(user_db_file_location) as conn: yield conn elif operation == "notes": with _sqlite_conn(note_db_file_location) as conn: yield conn elif operation == "images": with _sqlite_conn(image_db_file_location) as conn: yield conn else: raise ValueError(f"Unknown operation: {operation}") def list_users(): with _conn_for("users") as conn: cur = conn.cursor() cur.execute("SELECT id FROM users;") rows = cur.fetchall() return [x[0] for x in rows] def verify(id, pw): with _conn_for("users") as conn: cur = conn.cursor() if _is_postgres(): cur.execute("SELECT pw FROM users WHERE id = %s;", (id,)) else: cur.execute("SELECT pw FROM users WHERE id = ?;", (id,)) row = cur.fetchone() if not row: return False return row[0] == _sha256_hex(pw) def delete_user_from_db(id): if _is_postgres(): with _conn_for("users") as conn: cur = conn.cursor() cur.execute("DELETE FROM users WHERE id = %s;", (id,)) cur.execute("DELETE FROM notes WHERE user_id = %s;", (id,)) cur.execute("DELETE FROM images WHERE owner = %s;", (id,)) conn.commit() return with _conn_for("users") as conn: cur = conn.cursor() cur.execute("DELETE FROM users WHERE id = ?;", (id,)) conn.commit() with _conn_for("notes") as conn: cur = conn.cursor() cur.execute("DELETE FROM notes WHERE user = ?;", (id,)) conn.commit() with _conn_for("images") as conn: cur = conn.cursor() cur.execute("DELETE FROM images WHERE owner = ?;", (id,)) conn.commit() def add_user(id, pw): with _conn_for("users") as conn: cur = conn.cursor() if _is_postgres(): cur.execute("INSERT INTO users values(%s, %s)", (id.upper(), _sha256_hex(pw))) else: cur.execute("INSERT INTO users values(?, ?)", (id.upper(), _sha256_hex(pw))) conn.commit() def read_note_from_db(id): with _conn_for("notes") as conn: cur = conn.cursor() if _is_postgres(): cur.execute( "SELECT note_id, timestamp, note FROM notes WHERE user_id = %s;", (id.upper(),), ) else: cur.execute( "SELECT note_id, timestamp, note FROM notes WHERE user = ?;", (id.upper(),), ) return cur.fetchall() def match_user_id_with_note_id(note_id): # Given the note id, confirm if the current user is the owner of the note which is being operated. with _conn_for("notes") as conn: cur = conn.cursor() if _is_postgres(): cur.execute("SELECT user_id FROM notes WHERE note_id = %s;", (note_id,)) else: cur.execute("SELECT user FROM notes WHERE note_id = ?;", (note_id,)) row = cur.fetchone() return row[0] if row else None def write_note_into_db(id, note_to_write): with _conn_for("notes") as conn: cur = conn.cursor() current_timestamp = str(datetime.datetime.now()) note_id = hashlib.sha1((id.upper() + current_timestamp).encode()).hexdigest() if _is_postgres(): cur.execute( "INSERT INTO notes values(%s, %s, %s, %s)", (id.upper(), current_timestamp, note_to_write, note_id), ) else: cur.execute( "INSERT INTO notes values(?, ?, ?, ?)", (id.upper(), current_timestamp, note_to_write, note_id), ) conn.commit() def delete_note_from_db(note_id): with _conn_for("notes") as conn: cur = conn.cursor() if _is_postgres(): cur.execute("DELETE FROM notes WHERE note_id = %s;", (note_id,)) else: cur.execute("DELETE FROM notes WHERE note_id = ?;", (note_id,)) conn.commit() def image_upload_record(uid, owner, image_name, timestamp): with _conn_for("images") as conn: cur = conn.cursor() if _is_postgres(): cur.execute( "INSERT INTO images VALUES (%s, %s, %s, %s)", (uid, owner, image_name, timestamp), ) else: cur.execute( "INSERT INTO images VALUES (?, ?, ?, ?)", (uid, owner, image_name, timestamp), ) conn.commit() def list_images_for_user(owner): with _conn_for("images") as conn: cur = conn.cursor() if _is_postgres(): cur.execute( "SELECT uid, timestamp, name FROM images WHERE owner = %s", (owner,), ) else: cur.execute( "SELECT uid, timestamp, name FROM images WHERE owner = ?", (owner,), ) return cur.fetchall() def match_user_id_with_image_uid(image_uid): # Given the note id, confirm if the current user is the owner of the note which is being operated. with _conn_for("images") as conn: cur = conn.cursor() if _is_postgres(): cur.execute("SELECT owner FROM images WHERE uid = %s;", (image_uid,)) else: cur.execute("SELECT owner FROM images WHERE uid = ?;", (image_uid,)) row = cur.fetchone() return row[0] if row else None def delete_image_from_db(image_uid): with _conn_for("images") as conn: cur = conn.cursor() if _is_postgres(): cur.execute("DELETE FROM images WHERE uid = %s;", (image_uid,)) else: cur.execute("DELETE FROM images WHERE uid = ?;", (image_uid,)) conn.commit() if __name__ == "__main__": print(list_users())