# database.py # Main database facade - delegates to specialized modules in db_utils/ # Maintains backward compatibility with existing code import os from typing import Any, Dict, Iterable, List, Optional, Tuple from contextlib import contextmanager from threading import RLock import sqlite3 import logging from logger import Logger from db_utils.base import DatabaseBase from db_utils.config import ConfigOps from db_utils.hosts import HostOps from db_utils.actions import ActionOps from db_utils.queue import QueueOps from db_utils.vulnerabilities import VulnerabilityOps from db_utils.software import SoftwareOps from db_utils.credentials import CredentialOps from db_utils.services import ServiceOps from db_utils.scripts import ScriptOps from db_utils.stats import StatsOps from db_utils.backups import BackupOps from db_utils.comments import CommentOps from db_utils.agents import AgentOps from db_utils.studio import StudioOps from db_utils.webenum import WebEnumOps logger = Logger(name="database.py", level=logging.DEBUG) _DEFAULT_DB = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data", "bjorn.db") class BjornDatabase: """ Main database facade that delegates operations to specialized modules. All existing method calls remain unchanged - they're automatically forwarded. """ def __init__(self, db_path: Optional[str] = None): self.db_path = db_path or _DEFAULT_DB os.makedirs(os.path.dirname(self.db_path), exist_ok=True) # Initialize base connection manager self._base = DatabaseBase(self.db_path) # Initialize all operational modules (they share the base connection) self._config = ConfigOps(self._base) self._hosts = HostOps(self._base) self._actions = ActionOps(self._base) self._queue = QueueOps(self._base) self._vulnerabilities = VulnerabilityOps(self._base) self._software = SoftwareOps(self._base) self._credentials = CredentialOps(self._base) self._services = ServiceOps(self._base) self._scripts = ScriptOps(self._base) self._stats = StatsOps(self._base) self._backups = BackupOps(self._base) self._comments = CommentOps(self._base) self._agents = AgentOps(self._base) self._studio = StudioOps(self._base) self._webenum = WebEnumOps(self._base) # Ensure schema is created self.ensure_schema() logger.info(f"BjornDatabase initialized: {self.db_path}") # ========================================================================= # CORE PRIMITIVES - Delegated to base # ========================================================================= @property def _conn(self): """Access to underlying connection""" return self._base._conn @property def _lock(self): """Access to thread lock""" return self._base._lock @property def _cache_ttl(self): return self._base._cache_ttl @property def _stats_cache(self): return self._base._stats_cache @_stats_cache.setter def _stats_cache(self, value): self._base._stats_cache = value def _cursor(self): return self._base._cursor() def transaction(self, immediate: bool = True): return self._base.transaction(immediate) def execute(self, sql: str, params: Iterable[Any] = (), many: bool = False) -> int: return self._base.execute(sql, params, many) def executemany(self, sql: str, seq_of_params: Iterable[Iterable[Any]]) -> int: return self._base.executemany(sql, seq_of_params) def query(self, sql: str, params: Iterable[Any] = ()) -> List[Dict[str, Any]]: return self._base.query(sql, params) def query_one(self, sql: str, params: Iterable[Any] = ()) -> Optional[Dict[str, Any]]: return self._base.query_one(sql, params) def invalidate_stats_cache(self): return self._base.invalidate_stats_cache() # ========================================================================= # SCHEMA INITIALIZATION # ========================================================================= def ensure_schema(self) -> None: """Create all database tables if missing""" logger.info("Ensuring database schema...") # Each module creates its own tables self._config.create_tables() self._actions.create_tables() self._hosts.create_tables() self._services.create_tables() self._queue.create_tables() self._stats.create_tables() self._vulnerabilities.create_tables() self._software.create_tables() self._credentials.create_tables() self._scripts.create_tables() self._backups.create_tables() self._comments.create_tables() self._agents.create_tables() self._studio.create_tables() self._webenum.create_tables() # Initialize stats singleton self._stats.ensure_stats_initialized() logger.info("Database schema ready") # ========================================================================= # METHOD DELEGATION - All existing methods forwarded automatically # ========================================================================= # Config operations def get_config(self) -> Dict[str, Any]: return self._config.get_config() def save_config(self, config: Dict[str, Any]) -> None: return self._config.save_config(config) # Host operations def get_all_hosts(self) -> List[Dict[str, Any]]: return self._hosts.get_all_hosts() def update_host(self, mac_address: str, ips: Optional[str] = None, hostnames: Optional[str] = None, alive: Optional[int] = None, ports: Optional[str] = None, vendor: Optional[str] = None, essid: Optional[str] = None): return self._hosts.update_host(mac_address, ips, hostnames, alive, ports, vendor, essid) def merge_ip_stub_into_real(self, ip: str, real_mac: str, hostname: Optional[str] = None, essid_hint: Optional[str] = None): return self._hosts.merge_ip_stub_into_real(ip, real_mac, hostname, essid_hint) def update_hostname(self, mac_address: str, new_hostname: str): return self._hosts.update_hostname(mac_address, new_hostname) def get_current_hostname(self, mac_address: str) -> Optional[str]: return self._hosts.get_current_hostname(mac_address) def record_hostname_seen(self, mac_address: str, hostname: str): return self._hosts.record_hostname_seen(mac_address, hostname) def list_hostname_history(self, mac_address: str) -> List[Dict[str, Any]]: return self._hosts.list_hostname_history(mac_address) def update_ips_current(self, mac_address: str, current_ips: Iterable[str], cap_prev: int = 200): return self._hosts.update_ips_current(mac_address, current_ips, cap_prev) def update_ports_current(self, mac_address: str, current_ports: Iterable[int], cap_prev: int = 500): return self._hosts.update_ports_current(mac_address, current_ports, cap_prev) def update_essid_current(self, mac_address: str, new_essid: Optional[str], cap_prev: int = 50): return self._hosts.update_essid_current(mac_address, new_essid, cap_prev) # Action operations def sync_actions(self, actions): return self._actions.sync_actions(actions) def list_actions(self): return self._actions.list_actions() def list_studio_actions(self): return self._actions.list_studio_actions() def get_action_by_class(self, b_class: str) -> dict | None: return self._actions.get_action_by_class(b_class) def delete_action(self, b_class: str) -> None: return self._actions.delete_action(b_class) def upsert_simple_action(self, *, b_class: str, b_module: str, **kw) -> None: return self._actions.upsert_simple_action(b_class=b_class, b_module=b_module, **kw) def list_action_cards(self) -> list[dict]: return self._actions.list_action_cards() def get_action_definition(self, b_class: str) -> Optional[Dict[str, Any]]: return self._actions.get_action_definition(b_class) # Queue operations def get_next_queued_action(self) -> Optional[Dict[str, Any]]: return self._queue.get_next_queued_action() def update_queue_status(self, queue_id: int, status: str, error_msg: str = None, result: str = None): return self._queue.update_queue_status(queue_id, status, error_msg, result) def promote_due_scheduled_to_pending(self) -> int: return self._queue.promote_due_scheduled_to_pending() def ensure_scheduled_occurrence(self, action_name: str, next_run_at: str, mac: Optional[str] = "", ip: Optional[str] = "", **kwargs) -> bool: return self._queue.ensure_scheduled_occurrence(action_name, next_run_at, mac, ip, **kwargs) def queue_action(self, action_name: str, mac: str, ip: str, port: int = None, priority: int = 50, trigger: str = None, metadata: Dict = None) -> None: return self._queue.queue_action(action_name, mac, ip, port, priority, trigger, metadata) def queue_action_at(self, action_name: str, mac: Optional[str] = "", ip: Optional[str] = "", **kwargs) -> None: return self._queue.queue_action_at(action_name, mac, ip, **kwargs) def list_action_queue(self, statuses: Optional[Iterable[str]] = None) -> List[Dict[str, Any]]: return self._queue.list_action_queue(statuses) def get_upcoming_actions_summary(self) -> List[Dict[str, Any]]: return self._queue.get_upcoming_actions_summary() def supersede_old_attempts(self, action_name: str, mac_address: str, port: Optional[int] = None, ref_ts: Optional[str] = None) -> int: return self._queue.supersede_old_attempts(action_name, mac_address, port, ref_ts) def list_attempt_history(self, action_name: str, mac_address: str, port: Optional[int] = None, limit: int = 20) -> List[Dict[str, Any]]: return self._queue.list_attempt_history(action_name, mac_address, port, limit) def get_action_status_from_queue(self, action_name: str, mac_address: Optional[str] = None) -> Optional[Dict[str, Any]]: return self._queue.get_action_status_from_queue(action_name, mac_address) def get_last_action_status_from_queue(self, mac_address: str, action_name: str) -> Optional[Dict[str, str]]: return self._queue.get_last_action_status_from_queue(mac_address, action_name) def get_last_action_statuses_for_mac(self, mac_address: str) -> Dict[str, Dict[str, str]]: return self._queue.get_last_action_statuses_for_mac(mac_address) # Vulnerability operations def add_vulnerability(self, mac_address: str, vuln_id: str, ip: Optional[str] = None, hostname: Optional[str] = None, port: Optional[int] = None): return self._vulnerabilities.add_vulnerability(mac_address, vuln_id, ip, hostname, port) def update_vulnerability_status(self, mac_address: str, current_vulns: List[str]): return self._vulnerabilities.update_vulnerability_status(mac_address, current_vulns) def update_vulnerability_status_by_port(self, mac_address: str, port: int, current_vulns: List[str]): return self._vulnerabilities.update_vulnerability_status_by_port(mac_address, port, current_vulns) def get_all_vulns(self) -> List[Dict[str, Any]]: return self._vulnerabilities.get_all_vulns() def save_vulnerabilities(self, mac: str, ip: str, findings: List[Dict]): return self._vulnerabilities.save_vulnerabilities(mac, ip, findings) def cleanup_vulnerability_duplicates(self): return self._vulnerabilities.cleanup_vulnerability_duplicates() def fix_vulnerability_history_nulls(self): return self._vulnerabilities.fix_vulnerability_history_nulls() def count_vulnerabilities_alive(self, distinct: bool = False, active_only: bool = True) -> int: return self._vulnerabilities.count_vulnerabilities_alive(distinct, active_only) def count_distinct_vulnerabilities(self, alive_only: bool = False) -> int: return self._vulnerabilities.count_distinct_vulnerabilities(alive_only) def get_vulnerabilities_for_alive_hosts(self) -> List[str]: return self._vulnerabilities.get_vulnerabilities_for_alive_hosts() def list_vulnerability_history(self, cve_id: str | None = None, mac: str | None = None, limit: int = 500) -> list[dict]: return self._vulnerabilities.list_vulnerability_history(cve_id, mac, limit) # CVE metadata def get_cve_meta(self, cve_id: str) -> Optional[Dict[str, Any]]: return self._vulnerabilities.get_cve_meta(cve_id) def upsert_cve_meta(self, meta: Dict[str, Any]) -> None: return self._vulnerabilities.upsert_cve_meta(meta) def get_cve_meta_bulk(self, cve_ids: List[str]) -> Dict[str, Dict[str, Any]]: return self._vulnerabilities.get_cve_meta_bulk(cve_ids) # Software operations def add_detected_software(self, mac_address: str, cpe: str, ip: Optional[str] = None, hostname: Optional[str] = None, port: Optional[int] = None) -> None: return self._software.add_detected_software(mac_address, cpe, ip, hostname, port) def update_detected_software_status(self, mac_address: str, current_cpes: List[str]) -> None: return self._software.update_detected_software_status(mac_address, current_cpes) def migrate_cpe_from_vulnerabilities(self) -> int: return self._software.migrate_cpe_from_vulnerabilities() # Credential operations def insert_cred(self, service: str, mac: Optional[str] = None, ip: Optional[str] = None, hostname: Optional[str] = None, user: Optional[str] = None, password: Optional[str] = None, port: Optional[int] = None, database: Optional[str] = None, extra: Optional[Dict[str, Any]] = None): return self._credentials.insert_cred(service, mac, ip, hostname, user, password, port, database, extra) def list_creds_grouped(self) -> List[Dict[str, Any]]: return self._credentials.list_creds_grouped() # Service operations def upsert_port_service(self, mac_address: str, ip: Optional[str], port: int, **kwargs): return self._services.upsert_port_service(mac_address, ip, port, **kwargs) def get_services_for_host(self, mac_address: str) -> List[Dict]: return self._services.get_services_for_host(mac_address) def find_hosts_by_service(self, service: str) -> List[Dict]: return self._services.find_hosts_by_service(service) def get_service_for_host_port(self, mac_address: str, port: int, protocol: str = "tcp") -> Optional[Dict]: return self._services.get_service_for_host_port(mac_address, port, protocol) def _rebuild_host_ports(self, mac_address: str): return self._services._rebuild_host_ports(mac_address) # Script operations def add_script(self, name: str, type_: str, path: str, main_file: Optional[str] = None, category: Optional[str] = None, description: Optional[str] = None): return self._scripts.add_script(name, type_, path, main_file, category, description) def list_scripts(self) -> List[Dict[str, Any]]: return self._scripts.list_scripts() def delete_script(self, name: str) -> None: return self._scripts.delete_script(name) # Stats operations def get_livestats(self) -> Dict[str, int]: return self._stats.get_livestats() def update_livestats(self, total_open_ports: int, alive_hosts_count: int, all_known_hosts_count: int, vulnerabilities_count: int): return self._stats.update_livestats(total_open_ports, alive_hosts_count, all_known_hosts_count, vulnerabilities_count) def get_stats(self) -> Dict[str, int]: return self._stats.get_stats() def set_stats(self, total_open_ports: int, alive_hosts_count: int, all_known_hosts_count: int, vulnerabilities_count: int): return self._stats.set_stats(total_open_ports, alive_hosts_count, all_known_hosts_count, vulnerabilities_count) def get_display_stats(self) -> Dict[str, int]: return self._stats.get_display_stats() def ensure_stats_initialized(self): return self._stats.ensure_stats_initialized() # Backup operations def add_backup(self, filename: str, description: str, date: str, type_: str = "User Backup", is_default: bool = False, is_restore: bool = False, is_github: bool = False): return self._backups.add_backup(filename, description, date, type_, is_default, is_restore, is_github) def list_backups(self) -> List[Dict[str, Any]]: return self._backups.list_backups() def delete_backup(self, filename: str) -> None: return self._backups.delete_backup(filename) def clear_default_backup(self) -> None: return self._backups.clear_default_backup() def set_default_backup(self, filename: str) -> None: return self._backups.set_default_backup(filename) # Comment operations def count_comments(self) -> int: return self._comments.count_comments() def insert_comments(self, comments: List[Tuple[str, str, str, str, int]]): return self._comments.insert_comments(comments) def import_comments_from_json(self, json_path: str, lang: Optional[str] = None, default_theme: str = "general", default_weight: int = 1, clear_existing: bool = False) -> int: return self._comments.import_comments_from_json(json_path, lang, default_theme, default_weight, clear_existing) def random_comment_for(self, status: str, lang: str = "en") -> Optional[Dict[str, Any]]: return self._comments.random_comment_for(status, lang) # Agent operations (C2) def save_agent(self, agent_data: dict) -> None: return self._agents.save_agent(agent_data) def save_command(self, agent_id: str, command: str, response: str | None = None, success: bool = False) -> None: return self._agents.save_command(agent_id, command, response, success) def save_telemetry(self, agent_id: str, telemetry: dict) -> None: return self._agents.save_telemetry(agent_id, telemetry) def save_loot(self, loot: dict) -> None: return self._agents.save_loot(loot) def get_agent_history(self, agent_id: str) -> List[dict]: return self._agents.get_agent_history(agent_id) def purge_stale_agents(self, threshold_seconds: int) -> int: return self._agents.purge_stale_agents(threshold_seconds) def get_stale_agents(self, threshold_seconds: int) -> list[dict]: return self._agents.get_stale_agents(threshold_seconds) # Agent key management def get_active_key(self, agent_id: str) -> str | None: return self._agents.get_active_key(agent_id) def list_keys(self, agent_id: str) -> list[dict]: return self._agents.list_keys(agent_id) def save_new_key(self, agent_id: str, key_b64: str) -> int: return self._agents.save_new_key(agent_id, key_b64) def rotate_key(self, agent_id: str, new_key_b64: str) -> int: return self._agents.rotate_key(agent_id, new_key_b64) def revoke_keys(self, agent_id: str) -> int: return self._agents.revoke_keys(agent_id) def verify_client_key(self, agent_id: str, key_b64: str) -> bool: return self._agents.verify_client_key(agent_id, key_b64) def migrate_keys_from_file(self, json_path: str) -> int: return self._agents.migrate_keys_from_file(json_path) # Studio operations def get_studio_actions(self): return self._studio.get_studio_actions() def get_db_actions(self): return self._studio.get_db_actions() def update_studio_action(self, b_class: str, updates: dict): return self._studio.update_studio_action(b_class, updates) def get_studio_edges(self): return self._studio.get_studio_edges() def upsert_studio_edge(self, from_action: str, to_action: str, edge_type: str, metadata: dict = None): return self._studio.upsert_studio_edge(from_action, to_action, edge_type, metadata) def delete_studio_edge(self, edge_id: int): return self._studio.delete_studio_edge(edge_id) def get_studio_hosts(self, include_real: bool = True): return self._studio.get_studio_hosts(include_real) def upsert_studio_host(self, mac_address: str, data: dict): return self._studio.upsert_studio_host(mac_address, data) def delete_studio_host(self, mac: str): return self._studio.delete_studio_host(mac) def save_studio_layout(self, name: str, layout_data: dict, description: str = None): return self._studio.save_studio_layout(name, layout_data, description) def load_studio_layout(self, name: str): return self._studio.load_studio_layout(name) def apply_studio_to_runtime(self): return self._studio.apply_studio_to_runtime() def _replace_actions_studio_with_actions(self, vacuum: bool = False): return self._studio._replace_actions_studio_with_actions(vacuum) def _sync_actions_studio_schema_and_rows(self): return self._studio._sync_actions_studio_schema_and_rows() # WebEnum operations # Add webenum methods if you have any... # ========================================================================= # UTILITY OPERATIONS # ========================================================================= def checkpoint(self, mode: str = "TRUNCATE") -> Tuple[int, int, int]: """Force a WAL checkpoint""" return self._base.checkpoint(mode) def wal_checkpoint(self, mode: str = "TRUNCATE") -> Tuple[int, int, int]: """Alias for checkpoint""" return self.checkpoint(mode) def optimize(self) -> None: """Run PRAGMA optimize""" return self._base.optimize() def vacuum(self) -> None: """Vacuum the database""" return self._base.vacuum() # Internal helper methods used by modules def _table_exists(self, name: str) -> bool: return self._base._table_exists(name) def _column_names(self, table: str) -> List[str]: return self._base._column_names(table) def _ensure_column(self, table: str, column: str, ddl: str) -> None: return self._base._ensure_column(table, column, ddl)