Bjorn/db_utils/vulnerabilities.py

# db_utils/vulnerabilities.py
# Vulnerability tracking and CVE metadata operations

import json
import time
from typing import Any, Dict, List, Optional
import logging

from logger import Logger

logger = Logger(name="db_utils.vulnerabilities", level=logging.DEBUG)


class VulnerabilityOps:
    """Vulnerability tracking and CVE metadata operations"""
    
    def __init__(self, base):
        self.base = base
    
    def create_tables(self):
        """Create vulnerability and CVE metadata tables"""
        # CVE metadata cache (NVD/MITRE/EPSS/KEV + Exploit-DB)
        self.base.execute("""
            CREATE TABLE IF NOT EXISTS cve_meta (
                cve_id           TEXT PRIMARY KEY,
                description      TEXT,
                cvss_json        TEXT,
                references_json  TEXT,
                last_modified    TEXT,
                affected_json    TEXT,
                solution         TEXT,
                exploits_json    TEXT,
                is_kev           INTEGER DEFAULT 0,
                epss             REAL,
                epss_percentile  REAL,
                updated_at       INTEGER
            );
        """)
        self.base.execute("CREATE INDEX IF NOT EXISTS idx_cve_meta_updated ON cve_meta(updated_at);")
        
        # Vulnerabilities table
        self.base.execute("""
            CREATE TABLE IF NOT EXISTS vulnerabilities (
                id              INTEGER PRIMARY KEY AUTOINCREMENT,
                mac_address     TEXT NOT NULL,
                ip              TEXT,
                hostname        TEXT,
                port            INTEGER NOT NULL DEFAULT 0,
                vuln_id         TEXT NOT NULL,
                previous_vulns  TEXT,
                first_seen      TEXT DEFAULT CURRENT_TIMESTAMP,
                last_seen       TEXT DEFAULT CURRENT_TIMESTAMP,
                is_active       INTEGER DEFAULT 1
            );
        """)
        
        # Unique index without COALESCE since port is now NOT NULL
        self.base.execute("""
            DROP INDEX IF EXISTS uq_vuln_identity;
        """)
        
        self.base.execute("""
            CREATE UNIQUE INDEX IF NOT EXISTS uq_vuln_identity
            ON vulnerabilities(mac_address, vuln_id, port);
        """)
        
        # Migration: convert NULL to 0
        self.base.execute("""
            UPDATE vulnerabilities SET port = 0 WHERE port IS NULL;
        """)
        
        # Cleanup real duplicates after migration
        self.base.execute("""
            DELETE FROM vulnerabilities 
            WHERE rowid NOT IN (
                SELECT MIN(rowid) 
                FROM vulnerabilities 
                GROUP BY mac_address, vuln_id, port
            );
        """)
        
        self.base.execute("CREATE INDEX IF NOT EXISTS idx_vuln_active ON vulnerabilities(is_active) WHERE is_active=1;")
        self.base.execute("CREATE INDEX IF NOT EXISTS idx_vuln_mac_port ON vulnerabilities(mac_address, port);")
        
        # Vulnerability history (immutable log)
        self.base.execute("""
            CREATE TABLE IF NOT EXISTS vulnerability_history (
              id          INTEGER PRIMARY KEY AUTOINCREMENT,
              mac_address TEXT NOT NULL,
              ip          TEXT,
              hostname    TEXT,
              port        INTEGER,
              vuln_id     TEXT NOT NULL,
              event       TEXT NOT NULL,
              seen_at     TEXT DEFAULT CURRENT_TIMESTAMP
            );
        """)
        
        logger.debug("Vulnerability tables created/verified")
    
    # =========================================================================
    # CVE METADATA OPERATIONS
    # =========================================================================
    
    def get_cve_meta(self, cve_id: str) -> Optional[Dict[str, Any]]:
        """Get CVE metadata from cache"""
        row = self.base.query_one("SELECT * FROM cve_meta WHERE cve_id=? LIMIT 1;", (cve_id,))
        if not row:
            return None
        # Deserialize JSON fields
        for k in ("cvss_json", "references_json", "affected_json", "exploits_json"):
            if row.get(k):
                try:
                    row[k] = json.loads(row[k])
                except Exception:
                    row[k] = None
        return row
    
    def upsert_cve_meta(self, meta: Dict[str, Any]) -> None:
        """Insert or update CVE metadata"""
        # Serialize JSON fields
        cvss = json.dumps(meta.get("cvss"), ensure_ascii=False) if meta.get("cvss") is not None else None
        refs = json.dumps(meta.get("references"), ensure_ascii=False) if meta.get("references") is not None else None
        aff  = json.dumps(meta.get("affected"), ensure_ascii=False) if meta.get("affected") is not None else None
        exps = json.dumps(meta.get("exploits"), ensure_ascii=False) if meta.get("exploits") is not None else None
        
        self.base.execute("""
            INSERT INTO cve_meta(
                cve_id, description, cvss_json, references_json, last_modified,
                affected_json, solution, exploits_json, is_kev, epss, epss_percentile, updated_at
            ) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)
            ON CONFLICT(cve_id) DO UPDATE SET
                description     = excluded.description,
                cvss_json       = excluded.cvss_json,
                references_json = excluded.references_json,
                last_modified   = excluded.last_modified,
                affected_json   = excluded.affected_json,
                solution        = excluded.solution,
                exploits_json   = excluded.exploits_json,
                is_kev          = excluded.is_kev,
                epss            = excluded.epss,
                epss_percentile = excluded.epss_percentile,
                updated_at      = excluded.updated_at;
        """, (
            meta.get("cve_id"),
            meta.get("description"),
            cvss, refs, meta.get("lastModified"),
            aff, meta.get("solution"), exps,
            1 if meta.get("is_kev") else 0,
            meta.get("epss"),
            meta.get("epss_percentile"),
            int(meta.get("updated_at") or time.time())
        ))
    
    def get_cve_meta_bulk(self, cve_ids: List[str]) -> Dict[str, Dict[str, Any]]:
        """Get multiple CVE metadata entries at once"""
        if not cve_ids:
            return {}
        placeholders = ",".join("?" for _ in cve_ids)
        rows = self.base.query(f"SELECT * FROM cve_meta WHERE cve_id IN ({placeholders});", tuple(cve_ids))
        out = {}
        for r in rows:
            for k in ("cvss_json","references_json","affected_json","exploits_json"):
                if r.get(k):
                    try:
                        r[k] = json.loads(r[k])
                    except Exception:
                        r[k] = None
            out[r["cve_id"]] = r
        return out
    
    # =========================================================================
    # VULNERABILITY CRUD OPERATIONS
    # =========================================================================
    
    def add_vulnerability(self, mac_address: str, vuln_id: str, ip: Optional[str] = None,
                         hostname: Optional[str] = None, port: Optional[int] = None):
        """Insert/reactivate a vulnerability row and record history (NULL-safe on port)"""
        self.base.invalidate_stats_cache()
        p = int(port or 0)
        
        try:
            # Try to update existing row
            updated = self.base.execute(
                """
                UPDATE vulnerabilities
                SET is_active = 1,
                    ip        = COALESCE(?, ip),
                    hostname  = COALESCE(?, hostname),
                    last_seen = CURRENT_TIMESTAMP
                WHERE mac_address = ? AND vuln_id = ? AND COALESCE(port, 0) = ?
                """,
                (ip, hostname, mac_address, vuln_id, p)
            )
            
            if updated and updated > 0:
                # Seen again
                self.base.execute(
                    """
                    INSERT INTO vulnerability_history(mac_address, ip, hostname, port, vuln_id, event)
                    VALUES(?,?,?,?,?,'seen')
                    """,
                    (mac_address, ip, hostname, p, vuln_id)
                )
                return
            
            # Insert new row (port=0 if unknown)
            self.base.execute(
                """
                INSERT INTO vulnerabilities(mac_address, ip, hostname, port, vuln_id, is_active)
                VALUES(?,?,?,?,?,1)
                """,
                (mac_address, ip, hostname, p, vuln_id)
            )
            self.base.execute(
                """
                INSERT INTO vulnerability_history(mac_address, ip, hostname, port, vuln_id, event)
                VALUES(?,?,?,?,?,'new')
                """,
                (mac_address, ip, hostname, p, vuln_id)
            )
            
        except Exception:
            # Fallback if the query fails for exotic reason
            row = self.base.query_one(
                """
                SELECT id FROM vulnerabilities
                WHERE mac_address=? AND vuln_id=? AND COALESCE(port,0)=?
                LIMIT 1
                """,
                (mac_address, vuln_id, p)
            )
            if row:
                self.base.execute(
                    """
                    UPDATE vulnerabilities
                    SET is_active=1,
                        ip=COALESCE(?, ip),
                        hostname=COALESCE(?, hostname),
                        last_seen=CURRENT_TIMESTAMP
                    WHERE id=?
                    """,
                    (ip, hostname, row["id"])
                )
                self.base.execute(
                    """
                    INSERT INTO vulnerability_history(mac_address, ip, hostname, port, vuln_id, event)
                    VALUES(?,?,?,?,?,'seen')
                    """,
                    (mac_address, ip, hostname, p, vuln_id)
                )
            else:
                self.base.execute(
                    """
                    INSERT INTO vulnerabilities(mac_address, ip, hostname, port, vuln_id, is_active)
                    VALUES(?,?,?,?,?,1)
                    """,
                    (mac_address, ip, hostname, p, vuln_id)
                )
                self.base.execute(
                    """
                    INSERT INTO vulnerability_history(mac_address, ip, hostname, port, vuln_id, event)
                    VALUES(?,?,?,?,?,'new')
                    """,
                    (mac_address, ip, hostname, p, vuln_id)
                )
    
    def update_vulnerability_status(self, mac_address: str, current_vulns: List[str]):
        """Update vulnerability presence (new/seen/inactive) and touch timestamps/history"""
        existing = self.base.query(
            "SELECT vuln_id FROM vulnerabilities WHERE mac_address=? AND is_active=1",
            (mac_address,)
        )
        existing_ids = {r['vuln_id'] for r in existing}
        current_set = set(current_vulns)
        
        # Mark inactive
        for vuln_id in (existing_ids - current_set):
            self.base.execute("""
                UPDATE vulnerabilities 
                SET is_active=0, last_seen=CURRENT_TIMESTAMP
                WHERE mac_address=? AND vuln_id=? AND is_active=1
            """, (mac_address, vuln_id))
            
            self.base.execute("""
                INSERT INTO vulnerability_history(mac_address, port, vuln_id, event)
                SELECT mac_address, port, vuln_id, 'inactive'
                FROM vulnerabilities
                WHERE mac_address=? AND vuln_id=? LIMIT 1
            """, (mac_address, vuln_id))
        
        # Add new
        for vuln_id in (current_set - existing_ids):
            self.add_vulnerability(mac_address, vuln_id)
        
        # Seen: refresh last_seen and record history
        for vuln_id in (current_set & existing_ids):
            self.base.execute("""
                UPDATE vulnerabilities
                SET last_seen=CURRENT_TIMESTAMP
                WHERE mac_address=? AND vuln_id=? AND is_active=1
            """, (mac_address, vuln_id))
            
            self.base.execute("""
                INSERT INTO vulnerability_history(mac_address, port, vuln_id, event)
                SELECT mac_address, port, vuln_id, 'seen'
                FROM vulnerabilities
                WHERE mac_address=? AND vuln_id=? LIMIT 1
            """, (mac_address, vuln_id))
    
    def update_vulnerability_status_by_port(self, mac_address: str, port: int, current_vulns: List[str]):
        """Update vulnerability status for a specific port to avoid NULL conflicts"""
        port = int(port) if port is not None else 0
        
        existing = self.base.query(
            "SELECT vuln_id FROM vulnerabilities WHERE mac_address=? AND COALESCE(port, 0)=? AND is_active=1",
            (mac_address, port)
        )
        existing_ids = {r['vuln_id'] for r in existing}
        current_set = set(current_vulns)
        
        # Mark inactive (for this specific port)
        for vuln_id in (existing_ids - current_set):
            self.base.execute("""
                UPDATE vulnerabilities 
                SET is_active=0, last_seen=CURRENT_TIMESTAMP
                WHERE mac_address=? AND vuln_id=? AND COALESCE(port, 0)=? AND is_active=1
            """, (mac_address, vuln_id, port))
            
            self.base.execute("""
                INSERT INTO vulnerability_history(mac_address, ip, hostname, port, vuln_id, event)
                VALUES (?, NULL, NULL, ?, ?, 'inactive')
            """, (mac_address, port, vuln_id))
        
        # Add new (calls your existing method with the port)
        for vuln_id in (current_set - existing_ids):
            self.add_vulnerability(mac_address, vuln_id, port=port)
        
        # Mark as seen (for this specific port)
        for vuln_id in (current_set & existing_ids):
            self.base.execute("""
                UPDATE vulnerabilities
                SET last_seen=CURRENT_TIMESTAMP
                WHERE mac_address=? AND vuln_id=? AND COALESCE(port, 0)=? AND is_active=1
            """, (mac_address, vuln_id, port))
            
            self.base.execute("""
                INSERT INTO vulnerability_history(mac_address, ip, hostname, port, vuln_id, event)
                VALUES (?, NULL, NULL, ?, ?, 'seen')
            """, (mac_address, port, vuln_id))
    
    def save_vulnerabilities(self, mac: str, ip: str, findings: List[Dict]):
        """Separate CPE and CVE, update statuses + record new findings"""
        # Group findings by port to avoid conflicts
        findings_by_port = {}
        for f in findings:
            port = f.get('port', 0)
            if port is None:
                port = 0
            port = int(port) if port != 0 else 0
            
            if port not in findings_by_port:
                findings_by_port[port] = {'cves': set(), 'cpes': set(), 'findings': []}
            
            findings_by_port[port]['findings'].append(f)
            
            vid = str(f.get('vuln_id', ''))
            if vid.upper().startswith('CVE-'):
                findings_by_port[port]['cves'].add(vid)
            elif vid.upper().startswith('CPE:'):
                findings_by_port[port]['cpes'].add(vid.split(':', 1)[1])
            elif vid.lower().startswith('cpe:'):
                findings_by_port[port]['cpes'].add(vid)
        
        # Process CVE by port to avoid conflicts
        all_cve_ids = set()
        for port, data in findings_by_port.items():
            if data['cves']:
                try:
                    self.update_vulnerability_status_by_port(mac, port, sorted(data['cves']))
                    all_cve_ids.update(data['cves'])
                except Exception as e:
                    logger.error(f"Failed to update CVE status for port {port}: {e}")
        
        # Process CPE globally (as before) - delegated to SoftwareOps
        all_cpe_vals = set()
        for port, data in findings_by_port.items():
            all_cpe_vals.update(data['cpes'])
        
        # Note: CPE handling would typically be done by SoftwareOps
        # but we keep the call here for compatibility
        
        logger.debug(f"Processed: {len(all_cve_ids)} CVE across {len(findings_by_port)} ports, {len(all_cpe_vals)} CPE for {mac}")
    
    # =========================================================================
    # VULNERABILITY QUERY OPERATIONS
    # =========================================================================
    
    def get_all_vulns(self) -> List[Dict[str, Any]]:
        """Get all vulnerabilities with host details"""
        return self.base.query("""
            SELECT v.id, v.mac_address, v.ip, v.hostname, v.port, v.vuln_id, v.is_active, v.first_seen, v.last_seen,
                   h.ips AS host_ips, h.hostnames AS host_hostnames, h.ports AS host_ports, h.vendor AS host_vendor
            FROM vulnerabilities v
            LEFT JOIN hosts h ON v.mac_address = h.mac_address
            ORDER BY v.mac_address, v.vuln_id;
        """)
    
    def count_vulnerabilities_alive(self, distinct: bool = False, active_only: bool = True) -> int:
        """Count vulnerabilities for hosts with alive=1"""
        where = ["h.alive = 1"]
        if active_only:
            where.append("v.is_active = 1")
        where_sql = " AND ".join(where)
        
        if distinct:
            sql = f"""
                SELECT COUNT(DISTINCT v.vuln_id) AS c
                FROM vulnerabilities v
                JOIN hosts h ON h.mac_address = v.mac_address
                WHERE {where_sql}
            """
        else:
            sql = f"""
                SELECT COUNT(*) AS c
                FROM vulnerabilities v
                JOIN hosts h ON h.mac_address = v.mac_address
                WHERE {where_sql}
            """
        row = self.base.query(sql)
        return int(row[0]["c"]) if row else 0
    
    def count_distinct_vulnerabilities(self, alive_only: bool = False) -> int:
        """Return the number of distinct vulnerabilities (vuln_id)"""
        if alive_only:
            row = self.base.query("""
                SELECT COUNT(DISTINCT v.vuln_id) AS c
                  FROM vulnerabilities v
                  JOIN hosts h ON h.mac_address = v.mac_address
                 WHERE h.alive = 1
            """)
        else:
            row = self.base.query("SELECT COUNT(DISTINCT vuln_id) AS c FROM vulnerabilities")
        return int(row[0]["c"]) if row else 0
    
    def get_vulnerabilities_for_alive_hosts(self) -> List[str]:
        """Return a list of distinct vuln_id affecting hosts currently marked alive=1"""
        rows = self.base.query("""
            SELECT DISTINCT v.vuln_id
              FROM vulnerabilities v
              JOIN hosts h ON h.mac_address = v.mac_address
             WHERE h.alive = 1
        """)
        return [r["vuln_id"] for r in rows]
    
    def list_vulnerability_history(self, cve_id: str | None = None,
                                   mac: str | None = None, limit: int = 500) -> list[dict]:
        """Return vulnerability history (events) sorted most recent first"""
        where = []
        params: list = []
        if cve_id:
            where.append("vuln_id = ?")
            params.append(cve_id)
        if mac:
            where.append("mac_address = ?")
            params.append(mac)
        where_sql = ("WHERE " + " AND ".join(where)) if where else ""
        params.append(int(limit))
        
        return self.base.query(f"""
            SELECT mac_address, ip, hostname, port, vuln_id, event, seen_at
            FROM vulnerability_history
            {where_sql}
            ORDER BY datetime(seen_at) DESC
            LIMIT ?
        """, tuple(params))
    
    # =========================================================================
    # CLEANUP OPERATIONS
    # =========================================================================
    
    def cleanup_vulnerability_duplicates(self):
        """Clean up vulnerability duplicates"""
        self.base.invalidate_stats_cache()
        
        # Delete entries with port NULL if an entry with port=0 exists
        self.base.execute("""
            DELETE FROM vulnerabilities 
            WHERE port IS NULL 
            AND EXISTS (
                SELECT 1 FROM vulnerabilities v2 
                WHERE v2.mac_address = vulnerabilities.mac_address 
                AND v2.vuln_id = vulnerabilities.vuln_id 
                AND v2.port = 0
            )
        """)
        
        # Update remaining NULL ports to 0
        self.base.execute("""
            UPDATE vulnerabilities SET port = 0 WHERE port IS NULL
        """)
        
        # Delete true duplicates (same mac, vuln_id, port) - keep most recent
        self.base.execute("""
            DELETE FROM vulnerabilities 
            WHERE rowid NOT IN (
                SELECT MAX(rowid) 
                FROM vulnerabilities 
                GROUP BY mac_address, vuln_id, COALESCE(port, 0)
            )
        """)
    
    def fix_vulnerability_history_nulls(self):
        """Fix history entries with problematic NULL values"""
        # Update history where ports are NULL but should be 0
        self.base.execute("""
            UPDATE vulnerability_history 
            SET port = 0 
            WHERE port IS NULL 
            AND EXISTS (
                SELECT 1 FROM vulnerabilities v 
                WHERE v.mac_address = vulnerability_history.mac_address 
                AND v.vuln_id = vulnerability_history.vuln_id 
                AND v.port = 0
            )
        """)
        
        # For cases where we can't determine the port, use 0 by default
        self.base.execute("""
            UPDATE vulnerability_history 
            SET port = 0 
            WHERE port IS NULL
        """)