v 0.1

src/core/database/database.py

@@ -0,0 +1,121 @@
+import sqlite3
+from typing import List, Dict, Any
+from src.utils.logger import get_logger
+
+logger = get_logger("sqlite_db")
+
+class SQLiteDB:
+    def __init__(self, db_path: str):
+        """Инициализация соединения с базой данных."""
+        self.db_path = db_path
+        self._create_tables()
+        logger.info(f"Инициализация базы данных SQLite: {db_path}")
+
+    def _create_tables(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("Таблицы успешно созданы или уже существуют")
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка при создании таблиц: {e}")
+            raise
+
+    def save_scan_results(self, scan_results: Dict[str, Any]):
+        """Сохранение результатов сканирования в базу данных."""
+        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}")
+            raise
+
+    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 []