mirror of
https://github.com/infinition/Bjorn.git
synced 2025-12-13 16:14:57 +00:00
BREAKING CHANGE: Complete refactor of architecture to prepare BJORN V2 release, APIs, assets, and UI, webapp, logics, attacks, a lot of new features...
This commit is contained in:
Fabien POLLY
155
db_utils/stats.py
Normal file
155
db_utils/stats.py
Normal file
@@ -0,0 +1,155 @@
|
||||
# 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
|
||||
}
|
||||
Reference in New Issue
Block a user