firegex/backend/utils/sqlite.py

import json
import sqlite3
import os
from hashlib import md5

class SQLite():
    def __init__(self, db_name: str, schema:dict = None) -> None:
        self.conn: sqlite3.Connection|None = None
        self.cur = None
        self.db_name = db_name
        self.__backup = None
        self.schema = {} if schema is None else schema
        self.DB_VER = md5(json.dumps(self.schema).encode()).hexdigest()

    def connect(self) -> None:
        try:
            self.conn = sqlite3.connect(self.db_name, check_same_thread = False)
        except Exception:
            path_name = os.path.dirname(self.db_name)
            if not os.path.exists(path_name):
                os.makedirs(path_name)
            with open(self.db_name, 'x'):
                pass
            self.conn = sqlite3.connect(self.db_name, check_same_thread = False)
        def dict_factory(cursor, row):
            d = {}
            for idx, col in enumerate(cursor.description):
                d[col[0]] = row[idx]
            return d
        self.conn.row_factory = dict_factory

    def backup(self):
        with open(self.db_name, "rb") as f:
            self.__backup = f.read()
    
    def restore(self):
        were_active = True if self.conn else False
        self.disconnect()
        if self.__backup:
            with open(self.db_name, "wb") as f:
                f.write(self.__backup)
            self.__backup = None
        if were_active:
            self.connect()
            
    def disconnect(self) -> None:
        if self.conn:
            self.conn.close()
        self.conn = None

    def create_schema(self, tables = {}) -> None:
        if self.conn:
            cur = self.conn.cursor()
            cur.execute("CREATE TABLE IF NOT EXISTS main.keys_values(key VARCHAR(100) PRIMARY KEY, value VARCHAR(100) NOT NULL);")
            for t in tables:
                if t == "QUERY":
                    continue
                cur.execute('CREATE TABLE IF NOT EXISTS main.{}({});'.format(t, ''.join([(c + ' ' + tables[t][c] + ', ') for c in tables[t]])[:-2]))
            if "QUERY" in tables:
                [cur.execute(qry) for qry in tables["QUERY"]]
            cur.close()
    
    def query(self, query, *values):
        return self.queries([(query, *values)])[0]

    def queries(self, queries: list[tuple[str, ...]]):
        return list(self.queries_iter(queries))

    def queries_iter(self, queries: list[tuple[str, ...]]):
        cur = self.conn.cursor()
        try:
            cur.execute("BEGIN")
            for query_data in queries:
                values = []
                str_query = None
                if isinstance(query_data, str):
                    str_query = query_data
                elif (isinstance(query_data, tuple) or isinstance(query_data, list)) and len(query_data) > 0 and isinstance(query_data[0], str):
                    str_query = query_data[0]
                    values = query_data[1:]
                if str_query:
                    cur.execute(str_query, values)
                yield cur.fetchall()
            cur.execute("COMMIT")
        except Exception as e:
            cur.execute("ROLLBACK")
            raise e
        finally:
            cur.close()
            try:
                self.conn.commit()
            except Exception:
                pass
    
    def delete(self):
        self.disconnect()
        os.remove(self.db_name)
    
    def init(self):
        self.connect()
        try:
            if self.get('DB_VERSION') != self.DB_VER:
                raise Exception("DB_VERSION is not correct")
        except Exception:
            self.delete()
            self.connect()
            self.create_schema(self.schema)
            self.put('DB_VERSION', self.DB_VER)

    def get(self, key, default = None):
        q = self.query('SELECT value FROM keys_values WHERE key = ?', key)
        if len(q) == 0:
            return default
        else:
            return q[0]["value"]

    def put(self, key, value):
        if self.get(key) is None:
            self.query('INSERT INTO keys_values (key, value) VALUES (?, ?);', key, str(value))
        else:
            self.query('UPDATE keys_values SET value=? WHERE key = ?;', str(value), key)
    
    def set(self, key, value):
        return self.put(key, value)
        
    def dump(self) -> dict:
        self.connect()
        res = {}
        tables = self.query("SELECT name FROM sqlite_master WHERE type='table';")
        for table in tables:
            tname = table["name"]
            if tname == 'keys_values':
                # Four keys never travel in a backup. `password` and `secret` are this
                # instance's credentials, and `auth_disabled` with `auth_disabled_host`
                # are *whether it asks for them* — a property of where this firegex is
                # deployed, not of the configuration being backed up. A backup that could
                # carry them would be a backup that can turn a firewall's authentication
                # off on restore.
                res[tname] = self.query(
                    f"SELECT * FROM {tname} WHERE key NOT IN "
                    f"('password', 'secret', 'auth_disabled', 'auth_disabled_host');"
                )
            else:
                res[tname] = self.query(f"SELECT * FROM {tname};")
        return res
        
    def load(self, data: dict, dry_run: bool = False):
        """Replace the tables named in `data` with its rows, all in one transaction.

        `dry_run` does every write and then rolls it back: what an import uses to find out
        whether each of its databases *would* load before it changes any of them. Loaded
        one after another for real, a refusal from the third left the first two replaced —
        `keys_values` among them, which is where the password and the signing secret live.

        Every identifier is quoted. Columns are taken from the table itself rather than
        from the backup, so quoting cannot let anything in, and leaving them bare made a
        column called `table` — the firewall's — a syntax error: no backup holding a
        single firewall rule could be imported at all.
        """
        self.connect()
        cur = self.conn.cursor()
        try:
            cur.execute("BEGIN")
            for table_name, rows in data.items():
                # Verify the table exists to prevent SQL injection or schema mismatches
                cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?;", (table_name,))
                if len(cur.fetchall()) == 0:
                    continue
                # Ignore anything that isn't a proper list of row objects
                if not isinstance(rows, list) or not all(isinstance(r, dict) for r in rows):
                    continue
                cur.execute(f'DELETE FROM "{table_name}";')
                if rows:
                    # Column names are identifiers and can't be parameterized, so only
                    # keep the ones that actually exist on the table - this stops a
                    # crafted backup from injecting arbitrary identifiers into the query.
                    cur.execute(f'PRAGMA table_info("{table_name}");')
                    valid_cols = {r["name"] for r in cur.fetchall()}
                    cols = [c for c in rows[0].keys() if c in valid_cols]
                    if not cols:
                        continue
                    placeholders = ", ".join(["?"] * len(cols))
                    col_names = ", ".join(f'"{c}"' for c in cols)
                    query = f'INSERT INTO "{table_name}" ({col_names}) VALUES ({placeholders})'
                    for row in rows:
                        cur.execute(query, [row.get(c) for c in cols])
            cur.execute("ROLLBACK" if dry_run else "COMMIT")
        except Exception as e:
            cur.execute("ROLLBACK")
            raise e
        finally:
            cur.close()
            try:
                self.conn.commit()
            except Exception:
                pass