136 lines
4.7 KiB
Python
136 lines
4.7 KiB
Python
# src/core/database/db.py
|
||
import sqlite3
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
from contextlib import contextmanager
|
||
|
||
from src.utils.logger import get_logger
|
||
|
||
logger = get_logger("database")
|
||
|
||
class Database:
|
||
def __init__(self, db_path: str = "/home/lodqa/attack_module_data/security_scanner.db"):
|
||
self.db_path = Path(db_path)
|
||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||
self.conn: Optional[sqlite3.Connection] = None
|
||
self.init_db()
|
||
|
||
def connect(self):
|
||
"""Подключение к базе данных."""
|
||
try:
|
||
self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
|
||
self.conn.row_factory = sqlite3.Row
|
||
logger.info(f"Connected to database at {self.db_path}")
|
||
except sqlite3.Error as e:
|
||
logger.error(f"Failed to connect to database: {e}")
|
||
raise
|
||
|
||
def init_db(self):
|
||
"""Инициализация базы данных и создание таблиц."""
|
||
self.connect()
|
||
cursor = self.conn.cursor()
|
||
|
||
# Создание таблиц
|
||
cursor.executescript("""
|
||
CREATE TABLE IF NOT EXISTS sessions (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS scans (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
session_id INTEGER,
|
||
tool TEXT,
|
||
args TEXT,
|
||
summary TEXT,
|
||
duration INTEGER,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS hosts (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
scan_id INTEGER,
|
||
ip TEXT,
|
||
mac TEXT,
|
||
os TEXT,
|
||
FOREIGN KEY (scan_id) REFERENCES scans(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS ports (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
host_id INTEGER,
|
||
protocol TEXT,
|
||
port_num INTEGER,
|
||
state TEXT,
|
||
FOREIGN KEY (host_id) REFERENCES hosts(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS service (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
port_id INTEGER,
|
||
name TEXT,
|
||
product TEXT,
|
||
extrainfo TEXT,
|
||
ostype TEXT,
|
||
FOREIGN KEY (port_id) REFERENCES ports(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS cpe (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
service_id INTEGER,
|
||
name TEXT,
|
||
FOREIGN KEY (service_id) REFERENCES service(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS modbus_scan_result (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
scan_id INTEGER,
|
||
active_coils TEXT,
|
||
active_discrete_inputs TEXT,
|
||
active_holding_registers TEXT,
|
||
active_input_registers TEXT,
|
||
FOREIGN KEY (scan_id) REFERENCES scans(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS attacks (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
session_id INTEGER,
|
||
tool TEXT,
|
||
args TEXT,
|
||
summary TEXT,
|
||
duration INTEGER,
|
||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||
FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS attack_result (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
attack_id INTEGER,
|
||
summary TEXT,
|
||
FOREIGN KEY (attack_id) REFERENCES attacks(id) ON DELETE CASCADE
|
||
);
|
||
""")
|
||
self.conn.commit()
|
||
logger.info("Database tables initialized")
|
||
|
||
@contextmanager
|
||
def get_cursor(self):
|
||
"""Контекстный менеджер для работы с курсором."""
|
||
cursor = self.conn.cursor()
|
||
try:
|
||
yield cursor
|
||
self.conn.commit()
|
||
except Exception as e:
|
||
self.conn.rollback()
|
||
logger.error(f"Database operation failed: {e}")
|
||
raise
|
||
finally:
|
||
cursor.close()
|
||
|
||
def close(self):
|
||
"""Закрытие соединения с базой данных."""
|
||
if self.conn:
|
||
self.conn.close()
|
||
logger.info("Database connection closed") |