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