v 0.1

2025-06-02 14:53:57 +03:00
parent 91231b6a53
commit d7dc3e4ec4
29 changed files with 1252 additions and 3571 deletions
--- a/src/core/database/database.py
+++ b/src/core/database/database.py
@ -1,121 +1,136 @@
+# src/core/database/db.py
 import sqlite3
-from typing import List, Dict, Any
+from pathlib import Path
+from typing import Optional
+from contextlib import contextmanager
+
 from src.utils.logger import get_logger

-logger = get_logger("sqlite_db")
+logger = get_logger("database")

-class SQLiteDB:
-    def __init__(self, db_path: str):
-        """Инициализация соединения с базой данных."""
-        self.db_path = db_path
-        self._create_tables()
-        logger.info(f"Инициализация базы данных SQLite: {db_path}")
+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 _create_tables(self):
-        """Создание таблиц в базе данных."""
+    def connect(self):
+        """Подключение к базе данных."""
        try:
-            with sqlite3.connect(self.db_path) as conn:
-                cursor = conn.cursor()
-                # Таблица для хостов
-                cursor.execute("""
-                    CREATE TABLE IF NOT EXISTS hosts (
-                        id INTEGER PRIMARY KEY AUTOINCREMENT,
-                        ip TEXT NOT NULL UNIQUE,
-                        state TEXT NOT NULL,
-                        scan_timestamp TEXT NOT NULL
-                    )
-                """)
-                # Таблица для портов
-                cursor.execute("""
-                    CREATE TABLE IF NOT EXISTS ports (
-                        id INTEGER PRIMARY KEY AUTOINCREMENT,
-                        host_id INTEGER NOT NULL,
-                        port INTEGER NOT NULL,
-                        protocol TEXT,
-                        state TEXT,
-                        service TEXT,
-                        version TEXT,
-                        FOREIGN KEY (host_id) REFERENCES hosts(id) ON DELETE CASCADE
-                    )
-                """)
-                conn.commit()
-                logger.info("Таблицы успешно созданы или уже существуют")
+            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"Ошибка при создании таблиц: {e}")
+            logger.error(f"Failed to connect to database: {e}")
            raise

-    def save_scan_results(self, scan_results: Dict[str, Any]):
-        """Сохранение результатов сканирования в базу данных."""
+    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:
-            with sqlite3.connect(self.db_path) as conn:
-                cursor = conn.cursor()
-                # Сохранение хостов
-                for ip, data in scan_results["hosts"].items():
-                    state = data.get("state", {}).get("state", "unknown")
-                    timestamp = data.get("scan_timestamp", "unknown")
-                    cursor.execute(
-                        "INSERT OR REPLACE INTO hosts (ip, state, scan_timestamp) VALUES (?, ?, ?)",
-                        (ip, state, timestamp)
-                    )
-                    host_id = cursor.lastrowid if cursor.lastrowid else cursor.execute(
-                        "SELECT id FROM hosts WHERE ip = ?", (ip,)
-                    ).fetchone()[0]
-                    # Сохранение портов
-                    ports = data.get("ports", [])
-                    for port_data in ports:
-                        cursor.execute(
-                            """
-                            INSERT INTO ports (host_id, port, protocol, state, service, version)
-                            VALUES (?, ?, ?, ?, ?, ?)
-                            """,
-                            (
-                                host_id,
-                                port_data.get("portid"),
-                                port_data.get("protocol"),
-                                port_data.get("state"),
-                                port_data.get("service", {}).get("name"),
-                                port_data.get("service", {}).get("version")
-                            )
-                        )
-                conn.commit()
-                logger.info(f"Результаты сканирования сохранены для {len(scan_results['hosts'])} хостов")
-        except sqlite3.Error as e:
-            logger.error(f"Ошибка при сохранении результатов: {e}")
+            yield cursor
+            self.conn.commit()
+        except Exception as e:
+            self.conn.rollback()
+            logger.error(f"Database operation failed: {e}")
            raise
+        finally:
+            cursor.close()

-    def get_all_hosts(self) -> List[Dict[str, Any]]:
-        """Получение всех хостов из базы данных."""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                cursor = conn.cursor()
-                cursor.execute("SELECT * FROM hosts")
-                hosts = [{"id": row[0], "ip": row[1], "state": row[2], "scan_timestamp": row[3]} for row in cursor.fetchall()]
-                logger.info(f"Извлечено {len(hosts)} хостов из базы данных")
-                return hosts
-        except sqlite3.Error as e:
-            logger.error(f"Ошибка при получении хостов: {e}")
-            return []
-
-    def get_ports_by_host_id(self, host_id: int) -> List[Dict[str, Any]]:
-        """Получение портов для указанного хоста."""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                cursor = conn.cursor()
-                cursor.execute("SELECT * FROM ports WHERE host_id = ?", (host_id,))
-                ports = [
-                    {
-                        "id": row[0],
-                        "host_id": row[1],
-                        "port": row[2],
-                        "protocol": row[3],
-                        "state": row[4],
-                        "service": row[5],
-                        "version": row[6]
-                    }
-                    for row in cursor.fetchall()
-                ]
-                logger.debug(f"Извлечено {len(ports)} портов для хоста с ID {host_id}")
-                return ports
-        except sqlite3.Error as e:
-            logger.error(f"Ошибка при получении портов для хоста {host_id}: {e}")
-            return []
+    def close(self):
+        """Закрытие соединения с базой данных."""
+        if self.conn:
+            self.conn.close()
+            logger.info("Database connection closed")