souleyez 2.43.29__py3-none-any.whl → 3.0.0__py3-none-any.whl

souleyez/storage/database.py

@@ -12,15 +12,13 @@ from souleyez import config
 from souleyez.log_config import get_logger
 
 logger = get_logger(__name__)
-DB_PATH = Path(config.get('database.path', '~/.souleyez/souleyez.db')).expanduser()
+DB_PATH = Path(config.get("database.path", "~/.souleyez/souleyez.db")).expanduser()
 
 
 class Database:
     def __init__(self, db_path: str = None):
         self.db_path = db_path or str(DB_PATH)
-        logger.info("Initializing database", extra={
-            "db_path": self.db_path
-        })
+        logger.info("Initializing database", extra={"db_path": self.db_path})
         self._ensure_db()
 
     def _ensure_db(self):
@@ -41,10 +39,10 @@ class Database:
             conn.execute("PRAGMA journal_mode=WAL")
             conn.execute("PRAGMA busy_timeout=30000")
 
-            logger.info("Database connection established", extra={
-                "wal_mode": True,
-                "foreign_keys": True
-            })
+            logger.info(
+                "Database connection established",
+                extra={"wal_mode": True, "foreign_keys": True},
+            )
 
             # Check if this is an existing database (has tables)
             cursor = conn.execute(
@@ -58,22 +56,28 @@ class Database:
             if is_existing_db:
                 try:
                     from .migrations.migration_manager import MigrationManager
+
                     manager = MigrationManager(self.db_path)
                     pending = manager.get_pending_migrations()
                     if pending:
-                        logger.info("Running pending migrations for existing database", extra={
-                            "pending_count": len(pending)
-                        })
+                        logger.info(
+                            "Running pending migrations for existing database",
+                            extra={"pending_count": len(pending)},
+                        )
                         manager.migrate()
-                        logger.info("Migrations completed successfully", extra={
-                            "count": len(pending)
-                        })
+                        logger.info(
+                            "Migrations completed successfully",
+                            extra={"count": len(pending)},
+                        )
                 except Exception as migration_error:
-                    logger.error("Failed to run migrations on existing database", extra={
-                        "error": str(migration_error),
-                        "error_type": type(migration_error).__name__,
-                        "traceback": traceback.format_exc()
-                    })
+                    logger.error(
+                        "Failed to run migrations on existing database",
+                        extra={
+                            "error": str(migration_error),
+                            "error_type": type(migration_error).__name__,
+                            "traceback": traceback.format_exc(),
+                        },
+                    )
                     raise  # Don't continue if migrations fail for existing DB
 
             # Reconnect for schema loading
@@ -87,19 +91,21 @@ class Database:
             schema_path = Path(__file__).parent / "schema.sql"
 
             if schema_path.exists():
-                with open(schema_path, 'r') as f:
+                with open(schema_path, "r") as f:
                     schema_sql = f.read()
                     conn.executescript(schema_sql)
-                logger.info("Schema loaded from file", extra={
-                    "schema_file": str(schema_path)
-                })
+                logger.info(
+                    "Schema loaded from file", extra={"schema_file": str(schema_path)}
+                )
             else:
                 # If schema file doesn't exist, create essential tables inline
                 # This must match schema.sql structure for migrations to work
-                logger.warning("Schema file not found, using inline schema", extra={
-                    "expected_path": str(schema_path)
-                })
-                conn.executescript("""
+                logger.warning(
+                    "Schema file not found, using inline schema",
+                    extra={"expected_path": str(schema_path)},
+                )
+                conn.executescript(
+                    """
                 CREATE TABLE IF NOT EXISTS engagements (
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
                     name TEXT UNIQUE NOT NULL,
@@ -261,7 +267,8 @@ class Database:
                 CREATE INDEX IF NOT EXISTS idx_credentials_engagement ON credentials(engagement_id);
                 CREATE INDEX IF NOT EXISTS idx_credentials_host ON credentials(host_id);
                 CREATE INDEX IF NOT EXISTS idx_credentials_status ON credentials(status);
-            """)
+            """
+                )
 
             conn.commit()
             conn.close()
@@ -271,29 +278,35 @@ class Database:
             if not is_existing_db:
                 try:
                     from .migrations.migration_manager import MigrationManager
+
                     manager = MigrationManager(self.db_path)
                     pending = manager.get_pending_migrations()
                     if pending:
-                        logger.info("Running pending migrations for fresh database", extra={
-                            "pending_count": len(pending)
-                        })
+                        logger.info(
+                            "Running pending migrations for fresh database",
+                            extra={"pending_count": len(pending)},
+                        )
                         manager.migrate()
-                        logger.info("Migrations completed successfully", extra={
-                            "count": len(pending)
-                        })
+                        logger.info(
+                            "Migrations completed successfully",
+                            extra={"count": len(pending)},
+                        )
                 except Exception as migration_error:
-                    logger.error("Failed to run migrations", extra={
-                        "error": str(migration_error),
-                        "error_type": type(migration_error).__name__,
-                        "traceback": traceback.format_exc()
-                    })
+                    logger.error(
+                        "Failed to run migrations",
+                        extra={
+                            "error": str(migration_error),
+                            "error_type": type(migration_error).__name__,
+                            "traceback": traceback.format_exc(),
+                        },
+                    )
 
         except Exception as e:
-            logger.error("Database initialization failed", extra={
-                "error": str(e),
-                "traceback": traceback.format_exc()
-            })
-            if 'conn' in locals():
+            logger.error(
+                "Database initialization failed",
+                extra={"error": str(e), "traceback": traceback.format_exc()},
+            )
+            if "conn" in locals():
                 try:
                     conn.close()
                 except:
@@ -319,7 +332,13 @@ class Database:
         conn.execute(f"PRAGMA busy_timeout={busy_timeout}")
         return conn
 
-    def execute(self, query: str, params: tuple = None, retries: int = 3, fast_mode: bool = False) -> List[Dict[str, Any]]:
+    def execute(
+        self,
+        query: str,
+        params: tuple = None,
+        retries: int = 3,
+        fast_mode: bool = False,
+    ) -> List[Dict[str, Any]]:
         """
         Execute query and return results with retry logic.
 
@@ -346,13 +365,17 @@ class Database:
 
                 return results
             except sqlite3.OperationalError as e:
-                if "locked" in str(e).lower() and attempt < max_retries - 1 and not fast_mode:
+                if (
+                    "locked" in str(e).lower()
+                    and attempt < max_retries - 1
+                    and not fast_mode
+                ):
                     time.sleep(0.1 * (attempt + 1))  # Exponential backoff
                     continue
                 # For fast_mode, don't retry - just raise immediately
                 raise
             except Exception:
-                if 'conn' in locals():
+                if "conn" in locals():
                     conn.close()
                 raise
 
@@ -363,8 +386,8 @@ class Database:
 
     def insert(self, table: str, data: Dict[str, Any], retries: int = 3) -> int:
         """Insert row and return ID with retry logic."""
-        columns = ', '.join(data.keys())
-        placeholders = ', '.join(['?' for _ in data])
+        columns = ", ".join(data.keys())
+        placeholders = ", ".join(["?" for _ in data])
         query = f"INSERT INTO {table} ({columns}) VALUES ({placeholders})"
 
         for attempt in range(retries):
@@ -383,7 +406,7 @@ class Database:
                     continue
                 raise
             except Exception:
-                if 'conn' in locals():
+                if "conn" in locals():
                     conn.close()
                 raise
 

souleyez/storage/db.py

@@ -17,7 +17,7 @@ from typing import Optional, List, Dict, Any
 from ..utils import ensure_dirs, HISTORY_FILE, read_json
 from .. import config
 
-DB_PATH = Path(config.get('database.path', '~/.souleyez/souleyez.db')).expanduser()
+DB_PATH = Path(config.get("database.path", "~/.souleyez/souleyez.db")).expanduser()
 ensure_dirs()
 
 
@@ -41,7 +41,8 @@ def _get_conn(fast_mode: bool = False):
 def init_db():
     conn = _get_conn()
     cur = conn.cursor()
-    cur.executescript("""
+    cur.executescript(
+        """
     CREATE TABLE IF NOT EXISTS scans (
         id INTEGER PRIMARY KEY AUTOINCREMENT,
         ts TEXT,
@@ -59,13 +60,15 @@ def init_db():
     CREATE INDEX IF NOT EXISTS idx_scans_tool ON scans(tool);
     CREATE INDEX IF NOT EXISTS idx_scans_target ON scans(target);
     CREATE INDEX IF NOT EXISTS idx_scans_ts ON scans(ts);
-    """)
+    """
+    )
     conn.commit()
     conn.close()
-    
+
     # Auto-apply pending migrations on first run or database initialization
     try:
         from .migrations.migration_manager import MigrationManager
+
         manager = MigrationManager(str(DB_PATH))
         pending = manager.get_pending_migrations()
         if pending:
@@ -84,22 +87,25 @@ def insert_scan(entry: Dict[str, Any]) -> int:
     """
     conn = _get_conn()
     cur = conn.cursor()
-    cur.execute("""
+    cur.execute(
+        """
       INSERT INTO scans (ts, tool, target, label, args, log_path, raw_path, summary, per_host, status, rc)
       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-    """, (
-        entry.get("ts"),
-        entry.get("tool"),
-        entry.get("target"),
-        entry.get("label"),
-        json.dumps(entry.get("args") or []),
-        entry.get("log"),
-        entry.get("xml"),
-        json.dumps(entry.get("summary") or {}),
-        json.dumps(entry.get("per_host") or []),
-        entry.get("status") or "pending",
-        entry.get("rc")
-    ))
+    """,
+        (
+            entry.get("ts"),
+            entry.get("tool"),
+            entry.get("target"),
+            entry.get("label"),
+            json.dumps(entry.get("args") or []),
+            entry.get("log"),
+            entry.get("xml"),
+            json.dumps(entry.get("summary") or {}),
+            json.dumps(entry.get("per_host") or []),
+            entry.get("status") or "pending",
+            entry.get("rc"),
+        ),
+    )
     conn.commit()
     rowid = cur.lastrowid
     conn.close()
@@ -121,7 +127,19 @@ def update_scan(scan_id: int, **fields):
             k_db = "raw_path"
         else:
             k_db = k
-        if k_db not in ("ts", "tool", "target", "label", "args", "log_path", "raw_path", "summary", "per_host", "status", "rc"):
+        if k_db not in (
+            "ts",
+            "tool",
+            "target",
+            "label",
+            "args",
+            "log_path",
+            "raw_path",
+            "summary",
+            "per_host",
+            "status",
+            "rc",
+        ):
             continue
         if k_db in ("args", "summary", "per_host"):
             set_clauses.append(f"{k_db} = ?")
@@ -144,27 +162,31 @@ def get_scans(limit: int = 100, tool: Optional[str] = None) -> List[Dict[str, An
     conn = _get_conn()
     cur = conn.cursor()
     if tool:
-        cur.execute("SELECT * FROM scans WHERE tool = ? ORDER BY id DESC LIMIT ?", (tool, limit))
+        cur.execute(
+            "SELECT * FROM scans WHERE tool = ? ORDER BY id DESC LIMIT ?", (tool, limit)
+        )
     else:
         cur.execute("SELECT * FROM scans ORDER BY id DESC LIMIT ?", (limit,))
     rows = cur.fetchall()
     conn.close()
     out = []
     for r in rows:
-        out.append({
-            "id": r["id"],
-            "ts": r["ts"],
-            "tool": r["tool"],
-            "target": r["target"],
-            "label": r["label"],
-            "args": json.loads(r["args"] or "[]"),
-            "log": r["log_path"],
-            "xml": r["raw_path"],
-            "summary": json.loads(r["summary"] or "{}"),
-            "per_host": json.loads(r["per_host"] or "[]"),
-            "status": r["status"],
-            "rc": r["rc"]
-        })
+        out.append(
+            {
+                "id": r["id"],
+                "ts": r["ts"],
+                "tool": r["tool"],
+                "target": r["target"],
+                "label": r["label"],
+                "args": json.loads(r["args"] or "[]"),
+                "log": r["log_path"],
+                "xml": r["raw_path"],
+                "summary": json.loads(r["summary"] or "{}"),
+                "per_host": json.loads(r["per_host"] or "[]"),
+                "status": r["status"],
+                "rc": r["rc"],
+            }
+        )
     return out
 
 
@@ -188,7 +210,7 @@ def get_scan(scan_id: int) -> Optional[Dict[str, Any]]:
         "summary": json.loads(row["summary"] or "{}"),
         "per_host": json.loads(row["per_host"] or "[]"),
         "status": row["status"],
-        "rc": row["rc"]
+        "rc": row["rc"],
     }
 
 
@@ -215,7 +237,7 @@ def import_json_history_to_db():
             "log": e.get("log"),
             "xml": e.get("xml"),
             "summary": e.get("summary") or {},
-            "per_host": (e.get("summary") or {}).get("per_host") or []
+            "per_host": (e.get("summary") or {}).get("per_host") or [],
         }
         insert_scan(entry)
         inserted += 1