database.py

"""
Enterprise Honeypot Database Layer
SQLite-backed persistent storage with async-safe operations.
"""

import sqlite3
import threading
import time
import json
import os
from contextlib import contextmanager
from config import DB_PATH, DATA_DIR, RETENTION_DAYS


class HoneypotDB:
    """Asyncio-safe SQLite database for honeypot event storage.

    Uses a single connection with WAL mode (safe for concurrent readers)
    and a threading lock to serialize writes. This prevents cursor
    interleaving across coroutines sharing the same event loop thread
    and protects against executor threads (e.g. GeoIP worker).
    """

    _instance = None
    _init_lock = threading.Lock()

    def __new__(cls):
        with cls._init_lock:
            if cls._instance is None:
                cls._instance = super().__new__(cls)
                cls._instance._initialized = False
            return cls._instance

    def __init__(self):
        if self._initialized:
            return
        self._initialized = True
        self._db_lock = threading.Lock()
        self._conn = None
        os.makedirs(DATA_DIR, exist_ok=True)
        self._init_schema()

    def _get_conn(self):
        if self._conn is None:
            self._conn = sqlite3.connect(DB_PATH, timeout=30,
                                         check_same_thread=False)
            self._conn.row_factory = sqlite3.Row
            self._conn.execute("PRAGMA journal_mode=WAL")
            self._conn.execute("PRAGMA synchronous=NORMAL")
            self._conn.execute("PRAGMA foreign_keys=ON")
            self._conn.execute("PRAGMA busy_timeout=5000")
        return self._conn

    @contextmanager
    def _cursor(self):
        with self._db_lock:
            conn = self._get_conn()
            cur = conn.cursor()
            try:
                yield cur
                conn.commit()
            except Exception:
                conn.rollback()
                raise
            finally:
                cur.close()

    def _init_schema(self):
        with self._cursor() as cur:
            cur.executescript("""
                CREATE TABLE IF NOT EXISTS events (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    service TEXT NOT NULL,
                    event_type TEXT NOT NULL,
                    src_ip TEXT NOT NULL,
                    src_port INTEGER,
                    dst_port INTEGER,
                    protocol TEXT DEFAULT 'tcp',
                    payload TEXT,
                    metadata TEXT,
                    session_id TEXT,
                    threat_score INTEGER DEFAULT 0
                );

                CREATE TABLE IF NOT EXISTS connections (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT UNIQUE NOT NULL,
                    service TEXT NOT NULL,
                    src_ip TEXT NOT NULL,
                    src_port INTEGER,
                    dst_port INTEGER,
                    started_at REAL NOT NULL,
                    ended_at REAL,
                    bytes_received INTEGER DEFAULT 0,
                    bytes_sent INTEGER DEFAULT 0,
                    protocol_data TEXT
                );

                CREATE TABLE IF NOT EXISTS credentials (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    service TEXT NOT NULL,
                    src_ip TEXT NOT NULL,
                    username TEXT,
                    password TEXT,
                    success INTEGER DEFAULT 0,
                    session_id TEXT
                );

                CREATE TABLE IF NOT EXISTS attackers (
                    ip TEXT PRIMARY KEY,
                    first_seen REAL NOT NULL,
                    last_seen REAL NOT NULL,
                    total_connections INTEGER DEFAULT 1,
                    total_events INTEGER DEFAULT 0,
                    services_targeted TEXT DEFAULT '[]',
                    threat_score INTEGER DEFAULT 0,
                    country TEXT,
                    country_code TEXT,
                    city TEXT,
                    latitude REAL,
                    longitude REAL,
                    isp TEXT,
                    org TEXT,
                    asn TEXT,
                    is_banned INTEGER DEFAULT 0,
                    ban_expires REAL
                );

                CREATE TABLE IF NOT EXISTS geo_cache (
                    ip TEXT PRIMARY KEY,
                    data TEXT NOT NULL,
                    cached_at REAL NOT NULL
                );

                CREATE TABLE IF NOT EXISTS dashboard_users (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    username TEXT UNIQUE NOT NULL,
                    password_hash TEXT NOT NULL,
                    salt TEXT NOT NULL,
                    created_at REAL NOT NULL,
                    last_login REAL,
                    is_active INTEGER DEFAULT 1,
                    failed_attempts INTEGER DEFAULT 0,
                    locked_until REAL
                );

                CREATE TABLE IF NOT EXISTS sessions (
                    token TEXT PRIMARY KEY,
                    user_id INTEGER NOT NULL,
                    created_at REAL NOT NULL,
                    expires_at REAL NOT NULL,
                    ip_address TEXT,
                    FOREIGN KEY (user_id) REFERENCES dashboard_users(id)
                );

                CREATE TABLE IF NOT EXISTS alerts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp REAL NOT NULL,
                    alert_type TEXT NOT NULL,
                    severity TEXT NOT NULL,
                    source_ip TEXT,
                    service TEXT,
                    message TEXT NOT NULL,
                    acknowledged INTEGER DEFAULT 0
                );

                CREATE INDEX IF NOT EXISTS idx_events_timestamp ON events(timestamp);
                CREATE INDEX IF NOT EXISTS idx_events_service ON events(service);
                CREATE INDEX IF NOT EXISTS idx_events_src_ip ON events(src_ip);
                CREATE INDEX IF NOT EXISTS idx_events_type ON events(event_type);
                CREATE INDEX IF NOT EXISTS idx_connections_session ON connections(session_id);
                CREATE INDEX IF NOT EXISTS idx_connections_src_ip ON connections(src_ip);
                CREATE INDEX IF NOT EXISTS idx_credentials_src_ip ON credentials(src_ip);
                CREATE INDEX IF NOT EXISTS idx_credentials_service ON credentials(service);
                CREATE INDEX IF NOT EXISTS idx_attackers_threat ON attackers(threat_score);
                CREATE INDEX IF NOT EXISTS idx_alerts_timestamp ON alerts(timestamp);
                CREATE INDEX IF NOT EXISTS idx_geo_cache_time ON geo_cache(cached_at);
            """)

    # ─── Event Recording ────────────────────────────────────────────────

    def record_event(self, service, event_type, src_ip, src_port=None,
                     dst_port=None, protocol="tcp", payload=None,
                     metadata=None, session_id=None, threat_score=0):
        ts = time.time()
        with self._cursor() as cur:
            cur.execute("""
                INSERT INTO events (timestamp, service, event_type, src_ip,
                    src_port, dst_port, protocol, payload, metadata,
                    session_id, threat_score)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (ts, service, event_type, src_ip, src_port, dst_port,
                  protocol, payload, json.dumps(metadata) if metadata else None,
                  session_id, threat_score))
            self._update_attacker(cur, src_ip, service, ts, threat_score)
            return cur.lastrowid

    def _update_attacker(self, cur, ip, service, ts, threat_score):
        cur.execute("SELECT * FROM attackers WHERE ip = ?", (ip,))
        row = cur.fetchone()
        if row:
            services = json.loads(row["services_targeted"])
            if service not in services:
                services.append(service)
            new_score = min(100, row["threat_score"] + threat_score)
            cur.execute("""
                UPDATE attackers SET last_seen=?, total_events=total_events+1,
                    services_targeted=?, threat_score=?
                WHERE ip=?
            """, (ts, json.dumps(services), new_score, ip))
        else:
            cur.execute("""
                INSERT INTO attackers (ip, first_seen, last_seen, total_events,
                    services_targeted, threat_score)
                VALUES (?, ?, ?, 1, ?, ?)
            """, (ip, ts, ts, json.dumps([service]), min(100, threat_score)))

    # ─── Connection Tracking ────────────────────────────────────────────

    def open_connection(self, session_id, service, src_ip, src_port, dst_port):
        ts = time.time()
        with self._cursor() as cur:
            cur.execute("""
                INSERT OR IGNORE INTO connections
                    (session_id, service, src_ip, src_port, dst_port, started_at)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (session_id, service, src_ip, src_port, dst_port, ts))
            cur.execute("""
                UPDATE attackers SET total_connections = total_connections + 1,
                    last_seen = ? WHERE ip = ?
            """, (ts, src_ip))

    def close_connection(self, session_id, bytes_recv=0, bytes_sent=0, proto_data=None):
        with self._cursor() as cur:
            cur.execute("""
                UPDATE connections SET ended_at=?, bytes_received=?,
                    bytes_sent=?, protocol_data=?
                WHERE session_id=?
            """, (time.time(), bytes_recv, bytes_sent,
                  json.dumps(proto_data) if proto_data else None, session_id))

    # ─── Credential Capture ─────────────────────────────────────────────

    def record_credential(self, service, src_ip, username, password,
                          success=False, session_id=None):
        with self._cursor() as cur:
            cur.execute("""
                INSERT INTO credentials (timestamp, service, src_ip,
                    username, password, success, session_id)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (time.time(), service, src_ip, username, password,
                  1 if success else 0, session_id))

    # ─── Geo Cache ──────────────────────────────────────────────────────

    def get_geo_cache(self, ip):
        with self._cursor() as cur:
            cur.execute("SELECT data, cached_at FROM geo_cache WHERE ip=?", (ip,))
            row = cur.fetchone()
            if row:
                return json.loads(row["data"]), row["cached_at"]
        return None, None

    def set_geo_cache(self, ip, data):
        with self._cursor() as cur:
            cur.execute("""
                INSERT OR REPLACE INTO geo_cache (ip, data, cached_at)
                VALUES (?, ?, ?)
            """, (ip, json.dumps(data), time.time()))

    def update_attacker_geo(self, ip, geo_data):
        with self._cursor() as cur:
            cur.execute("""
                UPDATE attackers SET country=?, country_code=?, city=?,
                    latitude=?, longitude=?, isp=?, org=?, asn=?
                WHERE ip=?
            """, (geo_data.get("country"), geo_data.get("countryCode"),
                  geo_data.get("city"), geo_data.get("lat"),
                  geo_data.get("lon"), geo_data.get("isp"),
                  geo_data.get("org"), geo_data.get("as"), ip))

    # ─── Alerts ─────────────────────────────────────────────────────────

    def create_alert(self, alert_type, severity, message, source_ip=None, service=None):
        with self._cursor() as cur:
            cur.execute("""
                INSERT INTO alerts (timestamp, alert_type, severity,
                    source_ip, service, message)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (time.time(), alert_type, severity, source_ip, service, message))

    # ─── Dashboard Queries ──────────────────────────────────────────────

    def get_events(self, limit=100, offset=0, service=None, event_type=None,
                   src_ip=None, since=None):
        query = "SELECT * FROM events WHERE 1=1"
        params = []
        if service:
            query += " AND service=?"
            params.append(service)
        if event_type:
            query += " AND event_type=?"
            params.append(event_type)
        if src_ip:
            query += " AND src_ip=?"
            params.append(src_ip)
        if since:
            query += " AND timestamp>=?"
            params.append(since)
        query += " ORDER BY timestamp DESC LIMIT ? OFFSET ?"
        params.extend([limit, offset])
        with self._cursor() as cur:
            cur.execute(query, params)
            results = []
            for r in cur.fetchall():
                d = dict(r)
                if isinstance(d.get("metadata"), str):
                    try:
                        d["metadata"] = json.loads(d["metadata"])
                    except (json.JSONDecodeError, TypeError):
                        pass
                results.append(d)
            return results

    def get_stats(self, hours=24):
        since = time.time() - (hours * 3600)
        with self._cursor() as cur:
            stats = {}
            cur.execute("SELECT COUNT(*) as c FROM events WHERE timestamp>=?", (since,))
            stats["total_events"] = cur.fetchone()["c"]
            cur.execute("SELECT COUNT(DISTINCT src_ip) as c FROM events WHERE timestamp>=?", (since,))
            stats["unique_ips"] = cur.fetchone()["c"]
            cur.execute("SELECT COUNT(*) as c FROM credentials WHERE timestamp>=?", (since,))
            stats["credential_attempts"] = cur.fetchone()["c"]
            cur.execute("SELECT COUNT(*) as c FROM connections WHERE started_at>=?", (since,))
            stats["total_connections"] = cur.fetchone()["c"]
            cur.execute("""
                SELECT service, COUNT(*) as c FROM events
                WHERE timestamp>=? GROUP BY service ORDER BY c DESC
            """, (since,))
            stats["by_service"] = {r["service"]: r["c"] for r in cur.fetchall()}
            cur.execute("""
                SELECT event_type, COUNT(*) as c FROM events
                WHERE timestamp>=? GROUP BY event_type ORDER BY c DESC
            """, (since,))
            stats["by_type"] = {r["event_type"]: r["c"] for r in cur.fetchall()}
            cur.execute("""
                SELECT src_ip, COUNT(*) as c FROM events
                WHERE timestamp>=? GROUP BY src_ip ORDER BY c DESC LIMIT 20
            """, (since,))
            stats["top_attackers"] = [{"ip": r["src_ip"], "count": r["c"]} for r in cur.fetchall()]
            return stats

    def get_timeline(self, hours=24, buckets=48):
        since = time.time() - (hours * 3600)
        bucket_size = (hours * 3600) / buckets
        with self._cursor() as cur:
            cur.execute("""
                SELECT CAST((timestamp - ?) / ? AS INTEGER) as bucket,
                       service, COUNT(*) as c
                FROM events WHERE timestamp >= ?
                GROUP BY bucket, service ORDER BY bucket
            """, (since, bucket_size, since))
            timeline = {}
            for r in cur.fetchall():
                b = r["bucket"]
                if b not in timeline:
                    timeline[b] = {"timestamp": since + (b * bucket_size), "services": {}}
                timeline[b]["services"][r["service"]] = r["c"]
            return [timeline[k] for k in sorted(timeline.keys())]

    def get_attackers(self, limit=50, offset=0, min_score=0, order_by="threat_score"):
        allowed_orders = {"threat_score", "last_seen", "total_events", "total_connections"}
        if order_by not in allowed_orders:
            order_by = "threat_score"
        with self._cursor() as cur:
            cur.execute(f"""
                SELECT * FROM attackers WHERE threat_score >= ?
                ORDER BY {order_by} DESC LIMIT ? OFFSET ?
            """, (min_score, limit, offset))
            results = []
            for r in cur.fetchall():
                d = dict(r)
                # Parse JSON string fields for the API consumer
                if isinstance(d.get("services_targeted"), str):
                    try:
                        d["services_targeted"] = json.loads(d["services_targeted"])
                    except (json.JSONDecodeError, TypeError):
                        d["services_targeted"] = []
                results.append(d)
            return results

    def get_geo_stats(self):
        with self._cursor() as cur:
            cur.execute("""
                SELECT country_code, country, COUNT(*) as attacker_count,
                       SUM(total_events) as total_events,
                       AVG(threat_score) as avg_threat
                FROM attackers WHERE country_code IS NOT NULL
                GROUP BY country_code ORDER BY total_events DESC
            """)
            return [dict(r) for r in cur.fetchall()]

    def get_attack_map_data(self, hours=24):
        since = time.time() - (hours * 3600)
        with self._cursor() as cur:
            # Exclude private IPs (lat=0,lon=0) since they have no real location
            cur.execute("""
                SELECT a.ip, a.latitude, a.longitude, a.country_code, a.city,
                       a.threat_score, COUNT(e.id) as recent_events
                FROM attackers a
                JOIN events e ON e.src_ip = a.ip AND e.timestamp >= ?
                WHERE a.latitude IS NOT NULL
                  AND NOT (a.latitude = 0 AND a.longitude = 0)
                GROUP BY a.ip
                ORDER BY recent_events DESC LIMIT 500
            """, (since,))
            return [dict(r) for r in cur.fetchall()]

    def get_top_credentials(self, limit=50):
        with self._cursor() as cur:
            cur.execute("""
                SELECT username, password, COUNT(*) as attempts,
                       COUNT(DISTINCT src_ip) as unique_sources
                FROM credentials
                GROUP BY username, password
                ORDER BY attempts DESC LIMIT ?
            """, (limit,))
            return [dict(r) for r in cur.fetchall()]

    def get_alert_count(self, acknowledged=False):
        with self._cursor() as cur:
            cur.execute("SELECT COUNT(*) as c FROM alerts WHERE acknowledged=?",
                        (1 if acknowledged else 0,))
            return cur.fetchone()["c"]

    def get_alerts(self, limit=50, unack_only=True):
        with self._cursor() as cur:
            query = "SELECT * FROM alerts"
            if unack_only:
                query += " WHERE acknowledged=0"
            query += " ORDER BY timestamp DESC LIMIT ?"
            cur.execute(query, (limit,))
            return [dict(r) for r in cur.fetchall()]

    def acknowledge_alert(self, alert_id):
        with self._cursor() as cur:
            cur.execute("UPDATE alerts SET acknowledged=1 WHERE id=?", (alert_id,))

    # ─── Maintenance ────────────────────────────────────────────────────

    def cleanup_old_data(self):
        cutoff = time.time() - (RETENTION_DAYS * 86400)
        with self._cursor() as cur:
            cur.execute("DELETE FROM events WHERE timestamp < ?", (cutoff,))
            cur.execute("DELETE FROM connections WHERE started_at < ?", (cutoff,))
            cur.execute("DELETE FROM credentials WHERE timestamp < ?", (cutoff,))
            cur.execute("DELETE FROM alerts WHERE timestamp < ?", (cutoff,))
            cur.execute("DELETE FROM geo_cache WHERE cached_at < ?",
                        (time.time() - 86400,))
            cur.execute("VACUUM")

    # ─── Dashboard User Management ─────────────────────────────────────

    def get_user(self, username):
        with self._cursor() as cur:
            cur.execute("SELECT * FROM dashboard_users WHERE username=?", (username,))
            row = cur.fetchone()
            return dict(row) if row else None

    def create_user(self, username, password_hash, salt):
        with self._cursor() as cur:
            cur.execute("""
                INSERT OR IGNORE INTO dashboard_users
                    (username, password_hash, salt, created_at)
                VALUES (?, ?, ?, ?)
            """, (username, password_hash, salt, time.time()))

    def update_user_login(self, username):
        with self._cursor() as cur:
            cur.execute("""
                UPDATE dashboard_users SET last_login=?, failed_attempts=0,
                    locked_until=NULL WHERE username=?
            """, (time.time(), username))

    def increment_failed_login(self, username, lock_until=None):
        with self._cursor() as cur:
            if lock_until:
                cur.execute("""
                    UPDATE dashboard_users SET failed_attempts=failed_attempts+1,
                        locked_until=? WHERE username=?
                """, (lock_until, username))
            else:
                cur.execute("""
                    UPDATE dashboard_users SET failed_attempts=failed_attempts+1
                    WHERE username=?
                """, (username,))

    def create_session(self, token, user_id, expires_at, ip_address):
        with self._cursor() as cur:
            cur.execute("""
                INSERT INTO sessions (token, user_id, created_at, expires_at, ip_address)
                VALUES (?, ?, ?, ?, ?)
            """, (token, user_id, time.time(), expires_at, ip_address))

    def get_session(self, token):
        with self._cursor() as cur:
            cur.execute("""
                SELECT s.*, u.username FROM sessions s
                JOIN dashboard_users u ON s.user_id = u.id
                WHERE s.token=? AND s.expires_at > ?
            """, (token, time.time()))
            row = cur.fetchone()
            return dict(row) if row else None

    def delete_session(self, token):
        with self._cursor() as cur:
            cur.execute("DELETE FROM sessions WHERE token=?", (token,))

    def cleanup_sessions(self):
        with self._cursor() as cur:
            cur.execute("DELETE FROM sessions WHERE expires_at < ?", (time.time(),))

    # ─── Ban Management ─────────────────────────────────────────────────

    def ban_ip(self, ip, duration):
        with self._cursor() as cur:
            cur.execute("""
                UPDATE attackers SET is_banned=1, ban_expires=?
                WHERE ip=?
            """, (time.time() + duration, ip))

    def is_banned(self, ip):
        with self._cursor() as cur:
            cur.execute("""
                SELECT is_banned, ban_expires FROM attackers WHERE ip=?
            """, (ip,))
            row = cur.fetchone()
            if row and row["is_banned"]:
                if row["ban_expires"] and row["ban_expires"] < time.time():
                    cur.execute("UPDATE attackers SET is_banned=0 WHERE ip=?", (ip,))
                    return False
                return True
        return False