prismic/src/database.py

import os, hashlib, sqlite3, logging
from os import getenv as env
import mysql.connector


log = logging.getLogger(__name__)
FETCHALL = 0
FETCHONE = 1


class Database:
    def __init__(self):
        self.connection = None
        self.cursor = None

    def connect_sqlite(self, db_path):
        try:
            self.connection = sqlite3.connect(db_path, check_same_thread=False)
            self.cursor = self.connection.cursor()
            log.info("Connected to SQLite database")
        except sqlite3.Error as e:
            log.error(f"SQLite connection error: {e}")
            raise

    def connect_mysql(self):
        try:
            self.connection = mysql.connector.connect(
                host=env('MYSQL_HOST'),
                user=env('MYSQL_USER'),
                password=env('MYSQL_PASSWORD'),
                database=env('MYSQL_DATABASE')
            )
            self.cursor = self.connection.cursor()
            log.info("Connected to MySQL database")
        except mysql.connector.Error as e:
            log.error(f"MySQL connection error: {e}")
            raise

    def close(self):
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
        log.warning("Database connection closed")

    def execute_query(self, query, params=None, fetch_type=FETCHALL):
        try:
            if params:
                self.cursor.execute(query, params)
            else:
                self.cursor.execute(query)
            self.connection.commit()
            log.debug(f"Executed query: {query}")
            if fetch_type == FETCHALL:
                return self.fetchall()
            elif fetch_type == FETCHONE:
                return self.fetchone()
            elif fetch_type is None:
                log.debug("No fetch type specified, returning None")
                return None
            else:
                log.warning("Invalid fetch type, returning None")
                return None
        except (sqlite3.Error, mysql.connector.Error) as e:
            log.critical(f"Query execution error: {e}")
            raise

    def fetchall(self):
        try:
            result = self.cursor.fetchall()
            log.debug(f"Fetched all results")
            return result
        except (sqlite3.Error, mysql.connector.Error) as e:
            log.error(f"Fetchall error: {e}")
            raise

    def fetchone(self):
        try:
            result = self.cursor.fetchone()
            log.debug(f"Fetched one result")
            return result
        except (sqlite3.Error, mysql.connector.Error) as e:
            log.critical(f"Fetchone error: {e}")
            raise


def first_time_run(db:Database):
    log.info("First time run detected, initializing database")

    # Create users table
    log.info("Creating users table")
    db.execute_query("""
        CREATE TABLE users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT NOT NULL,
            password TEXT NOT NULL,
            about TEXT DEFAULT 'No description',
            permissions TEXT DEFAULT 'user',
            avatar MEDIUMBLOB,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            token TEXT,
            avatar_type TEXT
        )
    """)

    # Create posts table
    log.info("Creating posts table")
    db.execute_query("""
        CREATE TABLE posts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            board_id INTEGER NOT NULL,
            content TEXT NOT NULL,
            reference INTREGER,
            type TEXT DEFAULT 'post',
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE,
            FOREIGN KEY (board_id) REFERENCES boards (id) ON DELETE CASCADE,
            FOREIGN KEY (reference) REFERENCES posts (id) ON DELETE SET NULL
        )
    """)

    # Create boards table
    log.info("Creating boards table")
    db.execute_query("""
        CREATE TABLE boards (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT NOT NULL,
            description TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            owner_id INTEGER NOT NULL,
            FOREIGN KEY (owner_id) REFERENCES users (id)
        )
    """)

    # Create attachments table
    db.execute_query("""
        CREATE TABLE attachments (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            post_id INTEGER NOT NULL,
            file_data MEDIUMBLOB NOT NULL,
            file_name TEXT NOT NULL,
            file_type TEXT NOT NULL,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (post_id) REFERENCES posts (id) ON DELETE CASCADE
        )
    """)

    # Create system user

    if not env('SYSTEM_PASSWORD'):
        log.warning("SYSTEM_PASSWORD is empty, generating random password")
        password = os.urandom(16).hex()
        log.info(f"Generated system user password: {password}")
    else:
        password = env('SYSTEM_PASSWORD')
    
    log.info("Creating system user")
    db.execute_query("""
        INSERT INTO users (username, password, about, permissions)
        VALUES (?, ?, ?, ?)
    """, (
        env('SYSTEM_USER', default='system'),
        env('SYSTEM_PASSWORD', default=hashlib.sha256(password.encode()).hexdigest()),
        env('SYSTEM_ABOUT', default='System User'),
        env('SYSTEM_PERMISSIONS', default='admin')
    ))

    # Create system boards
    boards_names = env('SYSTEM_BOARDS', default='General,Random,System').split(',')
    if "System" not in boards_names:
        boards_names.append("System")

    log.info(f"Creating system boards: {', '.join(boards_names)}")
    for board_name in boards_names:
        db.execute_query("""
            INSERT INTO boards (name, description, owner_id)
            VALUES (?, ?, ?)
        """, (
            board_name,
            f"This is a automatically created board for {board_name}",
            1
        ))

    log.info("First time run completed")