289 lines
8.6 KiB
Python
289 lines
8.6 KiB
Python
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())
|