254 lines
8.0 KiB
Python
254 lines
8.0 KiB
Python
|
import sqlite3
|
||
|
|
import logging
|
||
|
|
import json
|
||
|
|
from datetime import datetime, timedelta
|
||
|
|
from typing import Dict, List, Optional
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
class Database:
|
||
|
|
"""SQLite database handler for Watchdog Docker"""
|
||
|
|
|
||
|
|
def __init__(self, db_path: str):
|
||
|
|
self.db_path = db_path
|
||
|
|
logger.info(f"Database initialized at {db_path}")
|
||
|
|
|
||
|
|
def _get_connection(self):
|
||
|
|
"""Get database connection"""
|
||
|
|
conn = sqlite3.connect(self.db_path)
|
||
|
|
conn.row_factory = sqlite3.Row
|
||
|
|
return conn
|
||
|
|
|
||
|
|
def initialize(self):
|
||
|
|
"""Create database tables if they don't exist"""
|
||
|
|
logger.info("Initializing database tables")
|
||
|
|
|
||
|
|
conn = self._get_connection()
|
||
|
|
cursor = conn.cursor()
|
||
|
|
|
||
|
|
# Events table
|
||
|
|
cursor.execute('''
|
||
|
|
CREATE TABLE IF NOT EXISTS events (
|
||
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
|
|
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||
|
|
type TEXT NOT NULL,
|
||
|
|
interface TEXT,
|
||
|
|
details TEXT NOT NULL,
|
||
|
|
data JSON
|
||
|
|
)
|
||
|
|
''')
|
||
|
|
|
||
|
|
# Create index on timestamp for faster queries
|
||
|
|
cursor.execute('''
|
||
|
|
CREATE INDEX IF NOT EXISTS idx_events_timestamp
|
||
|
|
ON events(timestamp DESC)
|
||
|
|
''')
|
||
|
|
|
||
|
|
# Create index on type for filtering
|
||
|
|
cursor.execute('''
|
||
|
|
CREATE INDEX IF NOT EXISTS idx_events_type
|
||
|
|
ON events(type)
|
||
|
|
''')
|
||
|
|
|
||
|
|
# Known devices table
|
||
|
|
cursor.execute('''
|
||
|
|
CREATE TABLE IF NOT EXISTS known_devices (
|
||
|
|
mac TEXT PRIMARY KEY,
|
||
|
|
name TEXT,
|
||
|
|
first_seen DATETIME DEFAULT CURRENT_TIMESTAMP,
|
||
|
|
last_seen DATETIME DEFAULT CURRENT_TIMESTAMP
|
||
|
|
)
|
||
|
|
''')
|
||
|
|
|
||
|
|
conn.commit()
|
||
|
|
conn.close()
|
||
|
|
logger.info("Database tables initialized successfully")
|
||
|
|
|
||
|
|
def add_event(self, event: Dict):
|
||
|
|
"""Add an event to the database"""
|
||
|
|
conn = self._get_connection()
|
||
|
|
cursor = conn.cursor()
|
||
|
|
|
||
|
|
try:
|
||
|
|
cursor.execute('''
|
||
|
|
INSERT INTO events (type, interface, details, data)
|
||
|
|
VALUES (?, ?, ?, ?)
|
||
|
|
''', (
|
||
|
|
event.get('type'),
|
||
|
|
event.get('interface'),
|
||
|
|
event.get('details'),
|
||
|
|
json.dumps(event)
|
||
|
|
))
|
||
|
|
|
||
|
|
conn.commit()
|
||
|
|
logger.debug(f"Event added: {event.get('type')} - {event.get('details')}")
|
||
|
|
|
||
|
|
# Auto-add new devices to known_devices if applicable
|
||
|
|
if event.get('type') == 'new_device' and not event.get('known'):
|
||
|
|
mac = event.get('mac')
|
||
|
|
hostname = event.get('hostname', 'Unknown')
|
||
|
|
if mac:
|
||
|
|
self.add_known_device(mac, hostname)
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error adding event: {e}", exc_info=True)
|
||
|
|
conn.rollback()
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
def get_recent_events(self, limit: int = 100, event_type: Optional[str] = None,
|
||
|
|
interface: Optional[str] = None) -> List[Dict]:
|
||
|
|
"""Get recent events with optional filtering"""
|
||
|
|
conn = self._get_connection()
|
||
|
|
cursor = conn.cursor()
|
||
|
|
|
||
|
|
query = 'SELECT * FROM events WHERE 1=1'
|
||
|
|
params = []
|
||
|
|
|
||
|
|
if event_type:
|
||
|
|
query += ' AND type = ?'
|
||
|
|
params.append(event_type)
|
||
|
|
|
||
|
|
if interface:
|
||
|
|
query += ' AND interface = ?'
|
||
|
|
params.append(interface)
|
||
|
|
|
||
|
|
query += ' ORDER BY timestamp DESC LIMIT ?'
|
||
|
|
params.append(limit)
|
||
|
|
|
||
|
|
try:
|
||
|
|
cursor.execute(query, params)
|
||
|
|
rows = cursor.fetchall()
|
||
|
|
|
||
|
|
events = []
|
||
|
|
for row in rows:
|
||
|
|
event = {
|
||
|
|
'id': row['id'],
|
||
|
|
'timestamp': row['timestamp'],
|
||
|
|
'type': row['type'],
|
||
|
|
'interface': row['interface'],
|
||
|
|
'details': row['details'],
|
||
|
|
'data': json.loads(row['data']) if row['data'] else {}
|
||
|
|
}
|
||
|
|
events.append(event)
|
||
|
|
|
||
|
|
return events
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error getting events: {e}", exc_info=True)
|
||
|
|
return []
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
def get_statistics(self) -> Dict:
|
||
|
|
"""Get event statistics"""
|
||
|
|
conn = self._get_connection()
|
||
|
|
cursor = conn.cursor()
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Total events
|
||
|
|
cursor.execute('SELECT COUNT(*) as count FROM events')
|
||
|
|
total = cursor.fetchone()['count']
|
||
|
|
|
||
|
|
# Events today
|
||
|
|
today = datetime.now().strftime('%Y-%m-%d')
|
||
|
|
cursor.execute(
|
||
|
|
'SELECT COUNT(*) as count FROM events WHERE DATE(timestamp) = ?',
|
||
|
|
(today,)
|
||
|
|
)
|
||
|
|
today_count = cursor.fetchone()['count']
|
||
|
|
|
||
|
|
# Events last hour
|
||
|
|
one_hour_ago = (datetime.now() - timedelta(hours=1)).strftime('%Y-%m-%d %H:%M:%S')
|
||
|
|
cursor.execute(
|
||
|
|
'SELECT COUNT(*) as count FROM events WHERE timestamp >= ?',
|
||
|
|
(one_hour_ago,)
|
||
|
|
)
|
||
|
|
hour_count = cursor.fetchone()['count']
|
||
|
|
|
||
|
|
# Events by type
|
||
|
|
cursor.execute('''
|
||
|
|
SELECT type, COUNT(*) as count
|
||
|
|
FROM events
|
||
|
|
GROUP BY type
|
||
|
|
''')
|
||
|
|
by_type = {row['type']: row['count'] for row in cursor.fetchall()}
|
||
|
|
|
||
|
|
# Known devices count
|
||
|
|
cursor.execute('SELECT COUNT(*) as count FROM known_devices')
|
||
|
|
known_devices = cursor.fetchone()['count']
|
||
|
|
|
||
|
|
return {
|
||
|
|
'total_events': total,
|
||
|
|
'events_today': today_count,
|
||
|
|
'events_last_hour': hour_count,
|
||
|
|
'events_by_type': by_type,
|
||
|
|
'known_devices': known_devices
|
||
|
|
}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error getting statistics: {e}", exc_info=True)
|
||
|
|
return {
|
||
|
|
'total_events': 0,
|
||
|
|
'events_today': 0,
|
||
|
|
'events_last_hour': 0,
|
||
|
|
'events_by_type': {},
|
||
|
|
'known_devices': 0
|
||
|
|
}
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
def is_known_device(self, mac: str) -> bool:
|
||
|
|
"""Check if a device is known"""
|
||
|
|
conn = self._get_connection()
|
||
|
|
cursor = conn.cursor()
|
||
|
|
|
||
|
|
try:
|
||
|
|
cursor.execute('SELECT mac FROM known_devices WHERE mac = ?', (mac,))
|
||
|
|
result = cursor.fetchone()
|
||
|
|
return result is not None
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error checking known device: {e}", exc_info=True)
|
||
|
|
return False
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
def add_known_device(self, mac: str, name: str = 'Unknown'):
|
||
|
|
"""Add a device to known devices"""
|
||
|
|
conn = self._get_connection()
|
||
|
|
cursor = conn.cursor()
|
||
|
|
|
||
|
|
try:
|
||
|
|
cursor.execute('''
|
||
|
|
INSERT OR REPLACE INTO known_devices (mac, name, last_seen)
|
||
|
|
VALUES (?, ?, CURRENT_TIMESTAMP)
|
||
|
|
''', (mac, name))
|
||
|
|
|
||
|
|
conn.commit()
|
||
|
|
logger.info(f"Device added to known devices: {mac} ({name})")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error adding known device: {e}", exc_info=True)
|
||
|
|
conn.rollback()
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
def cleanup_old_events(self, retention_days: int):
|
||
|
|
"""Delete events older than retention period"""
|
||
|
|
conn = self._get_connection()
|
||
|
|
cursor = conn.cursor()
|
||
|
|
|
||
|
|
try:
|
||
|
|
cutoff_date = (datetime.now() - timedelta(days=retention_days)).strftime('%Y-%m-%d %H:%M:%S')
|
||
|
|
|
||
|
|
cursor.execute('DELETE FROM events WHERE timestamp < ?', (cutoff_date,))
|
||
|
|
deleted = cursor.rowcount
|
||
|
|
|
||
|
|
conn.commit()
|
||
|
|
logger.info(f"Cleaned up {deleted} events older than {retention_days} days")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Error cleaning up old events: {e}", exc_info=True)
|
||
|
|
conn.rollback()
|
||
|
|
finally:
|
||
|
|
conn.close()
|