Bjorn/db_utils/stats.py

# db_utils/stats.py
# Statistics tracking and display operations

import time
import sqlite3
from typing import Dict
import logging

from logger import Logger

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


class StatsOps:
    """Statistics tracking and display operations"""
    
    def __init__(self, base):
        self.base = base
    
    def create_tables(self):
        """Create stats table"""
        self.base.execute("""
            CREATE TABLE IF NOT EXISTS stats (
              id INTEGER PRIMARY KEY CHECK (id=1),
              total_open_ports       INTEGER DEFAULT 0,
              alive_hosts_count      INTEGER DEFAULT 0,
              all_known_hosts_count  INTEGER DEFAULT 0,
              vulnerabilities_count  INTEGER DEFAULT 0,
              actions_count          INTEGER DEFAULT 0,
              zombie_count           INTEGER DEFAULT 0
            );
        """)
        logger.debug("Stats table created/verified")
    
    def ensure_stats_initialized(self):
        """Ensure the singleton row in `stats` exists"""
        row = self.base.query("SELECT 1 FROM stats WHERE id=1")
        if not row:
            self.base.execute("INSERT INTO stats(id) VALUES(1);")
    
    # =========================================================================
    # STATS OPERATIONS
    # =========================================================================
    
    def get_livestats(self) -> Dict[str, int]:
        """Return the live counters maintained in the `stats` singleton row"""
        row = self.base.query("""
            SELECT total_open_ports, alive_hosts_count, all_known_hosts_count, vulnerabilities_count
            FROM stats WHERE id=1
        """)[0]
        return {
            "total_open_ports": int(row["total_open_ports"]),
            "alive_hosts_count": int(row["alive_hosts_count"]),
            "all_known_hosts_count": int(row["all_known_hosts_count"]),
            "vulnerabilities_count": int(row["vulnerabilities_count"]),
        }
    
    def update_livestats(self, total_open_ports: int, alive_hosts_count: int,
                        all_known_hosts_count: int, vulnerabilities_count: int):
        """Update the live stats counters (touch in-place)"""
        self.base.invalidate_stats_cache()
        self.base.execute("""
            UPDATE stats
            SET total_open_ports = ?,
                alive_hosts_count = ?,
                all_known_hosts_count = ?,
                vulnerabilities_count = ?
            WHERE id = 1;
        """, (int(total_open_ports), int(alive_hosts_count),
              int(all_known_hosts_count), int(vulnerabilities_count)))
    
    def get_stats(self) -> Dict[str, int]:
        """Compatibility alias to retrieve stats; ensures the singleton row exists"""
        row = self.base.query("SELECT total_open_ports, alive_hosts_count, all_known_hosts_count, vulnerabilities_count FROM stats WHERE id=1;")
        if not row:
            self.ensure_stats_initialized()
            row = self.base.query("SELECT total_open_ports, alive_hosts_count, all_known_hosts_count, vulnerabilities_count FROM stats WHERE id=1;")
        r = row[0]
        return {
            "total_open_ports": int(r["total_open_ports"]),
            "alive_hosts_count": int(r["alive_hosts_count"]),
            "all_known_hosts_count": int(r["all_known_hosts_count"]),
            "vulnerabilities_count": int(r["vulnerabilities_count"]),
        }
    
    def set_stats(self, total_open_ports: int, alive_hosts_count: int, 
                  all_known_hosts_count: int, vulnerabilities_count: int):
        """Compatibility alias that forwards to update_livestats"""
        self.update_livestats(total_open_ports, alive_hosts_count, all_known_hosts_count, vulnerabilities_count)
    
    def get_display_stats(self) -> Dict[str, int]:
        """
        Cached bundle of counters for quick UI refresh using stats table.
        """
        now = time.time()
        
        # Serve from cache when valid
        if self.base._stats_cache['data'] and (now - self.base._stats_cache['timestamp']) < self.base._cache_ttl:
            return self.base._stats_cache['data'].copy()
        
        # Compute fresh counters
        with self.base._lock:
            try:
                # Use stats table for pre-calculated values
                result = self.base.query_one("""
                    SELECT 
                        s.total_open_ports,
                        s.alive_hosts_count,
                        s.all_known_hosts_count,
                        s.vulnerabilities_count,
                        COALESCE(s.actions_count, 
                            (SELECT COUNT(*) FROM actions WHERE b_enabled = 1)
                        ) as actions_count,
                        COALESCE(s.zombie_count, 0) as zombie_count,
                        (SELECT COUNT(*) FROM creds) as creds
                    FROM stats s
                    WHERE s.id = 1
                """)
                
                if result:
                    stats = {
                        'alive_hosts_count': int(result['alive_hosts_count'] or 0),
                        'all_known_hosts_count': int(result['all_known_hosts_count'] or 0),
                        'total_open_ports': int(result['total_open_ports'] or 0),
                        'vulnerabilities_count': int(result['vulnerabilities_count'] or 0),
                        'credentials_count': int(result['creds'] or 0),
                        'actions_count': int(result['actions_count'] or 0),
                        'zombie_count': int(result['zombie_count'] or 0)
                    }
                else:
                    # Fallback if no stats row
                    stats = {
                        'alive_hosts_count': 0,
                        'all_known_hosts_count': 0,
                        'total_open_ports': 0,
                        'vulnerabilities_count': 0,
                        'credentials_count': 0,
                        'actions_count': 0,
                        'zombie_count': 0
                    }
                
                # Update cache
                self.base._stats_cache = {'data': stats, 'timestamp': now}
                return stats
                
            except Exception:
                return {
                    'alive_hosts_count': 0,
                    'all_known_hosts_count': 0,
                    'total_open_ports': 0,
                    'vulnerabilities_count': 0,
                    'credentials_count': 0,
                    'actions_count': 0,
                    'zombie_count': 0
                }