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

souleyez/storage/migrations/_026_add_engagement_scope.py

@@ -9,6 +9,7 @@ Columns added:
 - engagements.scope_enforcement: Enforcement mode (off, warn, block)
 - hosts.scope_status: Host scope status (in_scope, out_of_scope, unknown)
 """
+
 import os
 
 
@@ -16,7 +17,8 @@ def upgrade(conn):
     """Add scope validation tables and columns."""
 
     # Engagement scope definitions table
-    conn.execute("""
+    conn.execute(
+        """
         CREATE TABLE IF NOT EXISTS engagement_scope (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
             engagement_id INTEGER NOT NULL,
@@ -29,10 +31,12 @@ def upgrade(conn):
             FOREIGN KEY (engagement_id) REFERENCES engagements(id) ON DELETE CASCADE,
             UNIQUE(engagement_id, scope_type, value)
         )
-    """)
+    """
+    )
 
     # Scope validation audit log
-    conn.execute("""
+    conn.execute(
+        """
         CREATE TABLE IF NOT EXISTS scope_validation_log (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
             engagement_id INTEGER NOT NULL,
@@ -45,11 +49,14 @@ def upgrade(conn):
             created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
             FOREIGN KEY (engagement_id) REFERENCES engagements(id) ON DELETE CASCADE
         )
-    """)
+    """
+    )
 
     # Add scope_enforcement column to engagements
     try:
-        conn.execute("ALTER TABLE engagements ADD COLUMN scope_enforcement TEXT DEFAULT 'off'")
+        conn.execute(
+            "ALTER TABLE engagements ADD COLUMN scope_enforcement TEXT DEFAULT 'off'"
+        )
     except Exception:
         pass  # Column may already exist
 
@@ -60,16 +67,28 @@ def upgrade(conn):
         pass  # Column may already exist
 
     # Indexes for performance
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_scope_engagement ON engagement_scope(engagement_id)")
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_scope_type ON engagement_scope(scope_type)")
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_scope_log_engagement ON scope_validation_log(engagement_id)")
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_scope_log_result ON scope_validation_log(validation_result)")
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_scope_log_timestamp ON scope_validation_log(created_at DESC)")
-    conn.execute("CREATE INDEX IF NOT EXISTS idx_hosts_scope_status ON hosts(scope_status)")
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_scope_engagement ON engagement_scope(engagement_id)"
+    )
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_scope_type ON engagement_scope(scope_type)"
+    )
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_scope_log_engagement ON scope_validation_log(engagement_id)"
+    )
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_scope_log_result ON scope_validation_log(validation_result)"
+    )
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_scope_log_timestamp ON scope_validation_log(created_at DESC)"
+    )
+    conn.execute(
+        "CREATE INDEX IF NOT EXISTS idx_hosts_scope_status ON hosts(scope_status)"
+    )
 
     conn.commit()
 
-    if not os.environ.get('SOULEYEZ_MIGRATION_SILENT'):
+    if not os.environ.get("SOULEYEZ_MIGRATION_SILENT"):
         print("Migration 026: Engagement scope validation tables created")
 
 

souleyez/storage/migrations/_027_multi_siem_persistence.py

@@ -6,6 +6,7 @@ Changes wazuh_config table to support multiple SIEM configs per engagement.
 - Adds UNIQUE constraint on (engagement_id, siem_type)
 - Allows each engagement to have separate configs for Wazuh, Splunk, etc.
 """
+
 import os
 
 
@@ -17,13 +18,16 @@ def upgrade(conn):
     cursor.execute("PRAGMA table_info(wazuh_config)")
     columns = [col[1] for col in cursor.fetchall()]
 
-    if 'siem_type' not in columns:
-        cursor.execute("ALTER TABLE wazuh_config ADD COLUMN siem_type TEXT DEFAULT 'wazuh'")
-    if 'config_json' not in columns:
+    if "siem_type" not in columns:
+        cursor.execute(
+            "ALTER TABLE wazuh_config ADD COLUMN siem_type TEXT DEFAULT 'wazuh'"
+        )
+    if "config_json" not in columns:
         cursor.execute("ALTER TABLE wazuh_config ADD COLUMN config_json TEXT")
 
     # Create new table with correct constraint
-    cursor.execute("""
+    cursor.execute(
+        """
         CREATE TABLE IF NOT EXISTS siem_config_new (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
             engagement_id INTEGER NOT NULL,
@@ -42,10 +46,12 @@ def upgrade(conn):
             FOREIGN KEY (engagement_id) REFERENCES engagements(id) ON DELETE CASCADE,
             UNIQUE(engagement_id, siem_type)
         )
-    """)
+    """
+    )
 
     # Copy existing data
-    cursor.execute("""
+    cursor.execute(
+        """
         INSERT OR IGNORE INTO siem_config_new (
             id, engagement_id, siem_type, api_url, api_user, api_password,
             indexer_url, indexer_user, indexer_password, verify_ssl, enabled,
@@ -56,19 +62,24 @@ def upgrade(conn):
             indexer_url, indexer_user, indexer_password, verify_ssl, enabled,
             config_json, created_at, updated_at
         FROM wazuh_config
-    """)
+    """
+    )
 
     # Drop old table and rename new one
     cursor.execute("DROP TABLE wazuh_config")
     cursor.execute("ALTER TABLE siem_config_new RENAME TO wazuh_config")
 
     # Recreate index
-    cursor.execute("CREATE INDEX IF NOT EXISTS idx_wazuh_config_engagement ON wazuh_config(engagement_id)")
-    cursor.execute("CREATE INDEX IF NOT EXISTS idx_wazuh_config_siem_type ON wazuh_config(siem_type)")
+    cursor.execute(
+        "CREATE INDEX IF NOT EXISTS idx_wazuh_config_engagement ON wazuh_config(engagement_id)"
+    )
+    cursor.execute(
+        "CREATE INDEX IF NOT EXISTS idx_wazuh_config_siem_type ON wazuh_config(siem_type)"
+    )
 
     conn.commit()
 
-    if not os.environ.get('SOULEYEZ_MIGRATION_SILENT'):
+    if not os.environ.get("SOULEYEZ_MIGRATION_SILENT"):
         print("Migration 027: Multi-SIEM persistence enabled")
 
 
@@ -76,7 +87,8 @@ def downgrade(conn):
     """Revert to single SIEM per engagement (lossy - keeps only first config per engagement)."""
     cursor = conn.cursor()
 
-    cursor.execute("""
+    cursor.execute(
+        """
         CREATE TABLE IF NOT EXISTS wazuh_config_old (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
             engagement_id INTEGER NOT NULL UNIQUE,
@@ -94,10 +106,12 @@ def downgrade(conn):
             updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
             FOREIGN KEY (engagement_id) REFERENCES engagements(id) ON DELETE CASCADE
         )
-    """)
+    """
+    )
 
     # Copy only first config per engagement
-    cursor.execute("""
+    cursor.execute(
+        """
         INSERT OR IGNORE INTO wazuh_config_old (
             engagement_id, api_url, api_user, api_password,
             indexer_url, indexer_user, indexer_password, verify_ssl, enabled,
@@ -109,11 +123,14 @@ def downgrade(conn):
             siem_type, config_json, created_at, updated_at
         FROM wazuh_config
         GROUP BY engagement_id
-    """)
+    """
+    )
 
     cursor.execute("DROP TABLE wazuh_config")
     cursor.execute("ALTER TABLE wazuh_config_old RENAME TO wazuh_config")
-    cursor.execute("CREATE INDEX IF NOT EXISTS idx_wazuh_config_engagement ON wazuh_config(engagement_id)")
+    cursor.execute(
+        "CREATE INDEX IF NOT EXISTS idx_wazuh_config_engagement ON wazuh_config(engagement_id)"
+    )
 
     conn.commit()
     print("Migration 027: Reverted to single SIEM per engagement")

souleyez/storage/migrations/__init__.py

@@ -36,32 +36,32 @@ from . import (
 
 # Migration registry - maps version to module
 MIGRATIONS_REGISTRY = {
-    '001': _001_add_credential_enhancements,
-    '002': _002_add_status_tracking,
-    '003': _003_add_execution_log,
-    '005': _005_screenshots,
-    '006': _006_deliverables,
-    '007': _007_deliverable_templates,
-    '008': _008_add_nuclei_table,
-    '009': _009_add_cme_tables,
-    '010': _010_evidence_linking,
-    '011': _011_timeline_tracking,
-    '012': _012_team_collaboration,
-    '013': _013_add_host_tags,
-    '014': _014_exploit_attempts,
-    '015': _015_add_mac_os_fields,
-    '016': _016_add_domain_field,
-    '017': _017_msf_sessions,
-    '018': _018_add_osint_target,
-    '019': _019_add_engagement_type,
-    '020': _020_add_rbac,
-    '021': _021_wazuh_integration,
-    '022': _022_wazuh_indexer_columns,
-    '023': _023_fix_detection_results_fk,
-    '024': _024_wazuh_vulnerabilities,
-    '025': _025_multi_siem_support,
-    '026': _026_add_engagement_scope,
-    '027': _027_multi_siem_persistence,
+    "001": _001_add_credential_enhancements,
+    "002": _002_add_status_tracking,
+    "003": _003_add_execution_log,
+    "005": _005_screenshots,
+    "006": _006_deliverables,
+    "007": _007_deliverable_templates,
+    "008": _008_add_nuclei_table,
+    "009": _009_add_cme_tables,
+    "010": _010_evidence_linking,
+    "011": _011_timeline_tracking,
+    "012": _012_team_collaboration,
+    "013": _013_add_host_tags,
+    "014": _014_exploit_attempts,
+    "015": _015_add_mac_os_fields,
+    "016": _016_add_domain_field,
+    "017": _017_msf_sessions,
+    "018": _018_add_osint_target,
+    "019": _019_add_engagement_type,
+    "020": _020_add_rbac,
+    "021": _021_wazuh_integration,
+    "022": _022_wazuh_indexer_columns,
+    "023": _023_fix_detection_results_fk,
+    "024": _024_wazuh_vulnerabilities,
+    "025": _025_multi_siem_support,
+    "026": _026_add_engagement_scope,
+    "027": _027_multi_siem_persistence,
 }
 
 

souleyez/storage/migrations/migration_manager.py

@@ -20,14 +20,16 @@ class MigrationManager:
     def _ensure_migrations_table(self):
         """Create migrations tracking table if it doesn't exist."""
         conn = sqlite3.connect(self.db_path)
-        conn.execute("""
+        conn.execute(
+            """
             CREATE TABLE IF NOT EXISTS schema_migrations (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
                 version TEXT UNIQUE NOT NULL,
                 name TEXT NOT NULL,
                 applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
             )
-        """)
+        """
+        )
         conn.commit()
         conn.close()
 
@@ -48,12 +50,8 @@ class MigrationManager:
         for version in get_all_versions():
             if version not in applied:
                 module = MIGRATIONS_REGISTRY[version]
-                name = getattr(module, 'DESCRIPTION', version)
-                pending.append({
-                    'version': version,
-                    'name': name,
-                    'module': module
-                })
+                name = getattr(module, "DESCRIPTION", version)
+                pending.append({"version": version, "name": name, "module": module})
 
         return pending
 
@@ -65,12 +63,13 @@ class MigrationManager:
         try:
             # Set environment variable for migration scripts to check
             import os
-            old_silent = os.environ.get('SOULEYEZ_MIGRATION_SILENT')
+
+            old_silent = os.environ.get("SOULEYEZ_MIGRATION_SILENT")
             if silent:
-                os.environ['SOULEYEZ_MIGRATION_SILENT'] = '1'
+                os.environ["SOULEYEZ_MIGRATION_SILENT"] = "1"
 
             # Get migration module from registry (already imported)
-            migration_module = migration['module']
+            migration_module = migration["module"]
 
             # Execute upgrade function
             conn = sqlite3.connect(self.db_path)
@@ -79,7 +78,7 @@ class MigrationManager:
             # Record migration
             conn.execute(
                 "INSERT INTO schema_migrations (version, name) VALUES (?, ?)",
-                (migration['version'], migration['name'])
+                (migration["version"], migration["name"]),
             )
             conn.commit()
             conn.close()
@@ -87,9 +86,9 @@ class MigrationManager:
             # Restore environment variable
             if silent:
                 if old_silent is None:
-                    os.environ.pop('SOULEYEZ_MIGRATION_SILENT', None)
+                    os.environ.pop("SOULEYEZ_MIGRATION_SILENT", None)
                 else:
-                    os.environ['SOULEYEZ_MIGRATION_SILENT'] = old_silent
+                    os.environ["SOULEYEZ_MIGRATION_SILENT"] = old_silent
 
             if not silent:
                 print(f"[{migration['version']}] ✅ Successfully applied")
@@ -157,7 +156,7 @@ class MigrationManager:
                     continue
 
                 # Check if downgrade function exists
-                if not hasattr(migration_module, 'downgrade'):
+                if not hasattr(migration_module, "downgrade"):
                     print(f"[{version}] ❌ No downgrade() function found")
                     continue
 
@@ -167,8 +166,7 @@ class MigrationManager:
 
                 # Remove from migrations table
                 conn.execute(
-                    "DELETE FROM schema_migrations WHERE version = ?",
-                    (version,)
+                    "DELETE FROM schema_migrations WHERE version = ?", (version,)
                 )
                 conn.commit()
                 conn.close()
@@ -201,11 +199,13 @@ class MigrationManager:
         if applied:
             print("✅ Applied:")
             conn = sqlite3.connect(self.db_path)
-            cursor = conn.execute("""
+            cursor = conn.execute(
+                """
                 SELECT version, name, applied_at
                 FROM schema_migrations
                 ORDER BY version
-            """)
+            """
+            )
             for row in cursor.fetchall():
                 print(f"  [{row[0]}] {row[1]} (applied: {row[2]})")
             conn.close()

souleyez/storage/msf_sessions.py

@@ -25,7 +25,7 @@ def add_msf_session(
     port: Optional[int] = None,
     tunnel_peer: Optional[str] = None,
     opened_at: Optional[datetime] = None,
-    notes: Optional[str] = None
+    notes: Optional[str] = None,
 ) -> int:
     """
     Add or update MSF session record
@@ -52,16 +52,20 @@ def add_msf_session(
     cursor = db.cursor()
 
     # Check if session already exists
-    cursor.execute("""
+    cursor.execute(
+        """
         SELECT id FROM msf_sessions
         WHERE engagement_id = ? AND msf_session_id = ?
-    """, (engagement_id, msf_session_id))
+    """,
+        (engagement_id, msf_session_id),
+    )
 
     existing = cursor.fetchone()
 
     if existing:
         # Update existing session
-        cursor.execute("""
+        cursor.execute(
+            """
             UPDATE msf_sessions
             SET session_type = ?,
                 via_exploit = ?,
@@ -76,34 +80,54 @@ def add_msf_session(
                 is_active = 1,
                 updated_at = CURRENT_TIMESTAMP
             WHERE id = ?
-        """, (
-            session_type, via_exploit, via_payload, platform, arch,
-            username, port, tunnel_peer, opened_at, notes, existing[0]
-        ))
+        """,
+            (
+                session_type,
+                via_exploit,
+                via_payload,
+                platform,
+                arch,
+                username,
+                port,
+                tunnel_peer,
+                opened_at,
+                notes,
+                existing[0],
+            ),
+        )
         return existing[0]
     else:
         # Insert new session
-        cursor.execute("""
+        cursor.execute(
+            """
             INSERT INTO msf_sessions (
                 engagement_id, host_id, msf_session_id,
                 session_type, via_exploit, via_payload,
                 platform, arch, username, port,
                 tunnel_peer, opened_at, notes, is_active
             ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1)
-        """, (
-            engagement_id, host_id, msf_session_id,
-            session_type, via_exploit, via_payload,
-            platform, arch, username, port,
-            tunnel_peer, opened_at, notes
-        ))
+        """,
+            (
+                engagement_id,
+                host_id,
+                msf_session_id,
+                session_type,
+                via_exploit,
+                via_payload,
+                platform,
+                arch,
+                username,
+                port,
+                tunnel_peer,
+                opened_at,
+                notes,
+            ),
+        )
         return cursor.lastrowid
 
 
 def close_msf_session(
-    db,
-    engagement_id: int,
-    msf_session_id: int,
-    close_reason: Optional[str] = None
+    db, engagement_id: int, msf_session_id: int, close_reason: Optional[str] = None
 ) -> bool:
     """
     Mark MSF session as closed
@@ -119,23 +143,23 @@ def close_msf_session(
     """
     cursor = db.cursor()
 
-    cursor.execute("""
+    cursor.execute(
+        """
         UPDATE msf_sessions
         SET is_active = 0,
             closed_at = CURRENT_TIMESTAMP,
             close_reason = ?,
             updated_at = CURRENT_TIMESTAMP
         WHERE engagement_id = ? AND msf_session_id = ?
-    """, (close_reason, engagement_id, msf_session_id))
+    """,
+        (close_reason, engagement_id, msf_session_id),
+    )
 
     return cursor.rowcount > 0
 
 
 def get_msf_sessions(
-    db,
-    engagement_id: int,
-    active_only: bool = True,
-    host_id: Optional[int] = None
+    db, engagement_id: int, active_only: bool = True, host_id: Optional[int] = None
 ) -> List[Dict[str, Any]]:
     """
     Get MSF sessions
@@ -178,9 +202,7 @@ def get_msf_sessions(
 
 
 def get_msf_session_by_id(
-    db,
-    engagement_id: int,
-    msf_session_id: int
+    db, engagement_id: int, msf_session_id: int
 ) -> Optional[Dict[str, Any]]:
     """
     Get MSF session by MSF session ID
@@ -195,7 +217,8 @@ def get_msf_session_by_id(
     """
     cursor = db.cursor()
 
-    cursor.execute("""
+    cursor.execute(
+        """
         SELECT
             s.*,
             h.ip_address as host_ip,
@@ -203,7 +226,9 @@ def get_msf_session_by_id(
         FROM msf_sessions s
         JOIN hosts h ON s.host_id = h.id
         WHERE s.engagement_id = ? AND s.msf_session_id = ?
-    """, (engagement_id, msf_session_id))
+    """,
+        (engagement_id, msf_session_id),
+    )
 
     row = cursor.fetchone()
     if row:
@@ -212,11 +237,7 @@ def get_msf_session_by_id(
     return None
 
 
-def update_session_last_seen(
-    db,
-    engagement_id: int,
-    msf_session_id: int
-) -> bool:
+def update_session_last_seen(db, engagement_id: int, msf_session_id: int) -> bool:
     """
     Update session last_seen timestamp
 
@@ -230,12 +251,15 @@ def update_session_last_seen(
     """
     cursor = db.cursor()
 
-    cursor.execute("""
+    cursor.execute(
+        """
         UPDATE msf_sessions
         SET last_seen = CURRENT_TIMESTAMP,
             updated_at = CURRENT_TIMESTAMP
         WHERE engagement_id = ? AND msf_session_id = ?
-    """, (engagement_id, msf_session_id))
+    """,
+        (engagement_id, msf_session_id),
+    )
 
     return cursor.rowcount > 0
 
@@ -254,71 +278,79 @@ def get_session_stats(db, engagement_id: int) -> Dict[str, Any]:
     cursor = db.cursor()
 
     stats = {
-        'total': 0,
-        'active': 0,
-        'closed': 0,
-        'by_type': {},
-        'by_exploit': {},
-        'compromised_hosts': 0
+        "total": 0,
+        "active": 0,
+        "closed": 0,
+        "by_type": {},
+        "by_exploit": {},
+        "compromised_hosts": 0,
     }
 
     # Total and active sessions
-    cursor.execute("""
+    cursor.execute(
+        """
         SELECT
             COUNT(*) as total,
             SUM(CASE WHEN is_active = 1 THEN 1 ELSE 0 END) as active
         FROM msf_sessions
         WHERE engagement_id = ?
-    """, (engagement_id,))
+    """,
+        (engagement_id,),
+    )
 
     row = cursor.fetchone()
     if row:
-        stats['total'] = row[0] or 0
-        stats['active'] = row[1] or 0
-        stats['closed'] = stats['total'] - stats['active']
+        stats["total"] = row[0] or 0
+        stats["active"] = row[1] or 0
+        stats["closed"] = stats["total"] - stats["active"]
 
     # Sessions by type
-    cursor.execute("""
+    cursor.execute(
+        """
         SELECT session_type, COUNT(*) as count
         FROM msf_sessions
         WHERE engagement_id = ?
         GROUP BY session_type
-    """, (engagement_id,))
+    """,
+        (engagement_id,),
+    )
 
     for row in cursor.fetchall():
-        session_type = row[0] or 'unknown'
-        stats['by_type'][session_type] = row[1]
+        session_type = row[0] or "unknown"
+        stats["by_type"][session_type] = row[1]
 
     # Sessions by exploit
-    cursor.execute("""
+    cursor.execute(
+        """
         SELECT via_exploit, COUNT(*) as count
         FROM msf_sessions
         WHERE engagement_id = ? AND via_exploit IS NOT NULL
         GROUP BY via_exploit
-    """, (engagement_id,))
+    """,
+        (engagement_id,),
+    )
 
     for row in cursor.fetchall():
-        stats['by_exploit'][row[0]] = row[1]
+        stats["by_exploit"][row[0]] = row[1]
 
     # Compromised hosts (hosts with active sessions)
-    cursor.execute("""
+    cursor.execute(
+        """
         SELECT COUNT(DISTINCT host_id)
         FROM msf_sessions
         WHERE engagement_id = ? AND is_active = 1
-    """, (engagement_id,))
+    """,
+        (engagement_id,),
+    )
 
     row = cursor.fetchone()
     if row:
-        stats['compromised_hosts'] = row[0] or 0
+        stats["compromised_hosts"] = row[0] or 0
 
     return stats
 
 
-def delete_msf_session(
-    db,
-    engagement_id: int,
-    msf_session_id: int
-) -> bool:
+def delete_msf_session(db, engagement_id: int, msf_session_id: int) -> bool:
     """
     Delete MSF session record
 
@@ -332,9 +364,12 @@ def delete_msf_session(
     """
     cursor = db.cursor()
 
-    cursor.execute("""
+    cursor.execute(
+        """
         DELETE FROM msf_sessions
         WHERE engagement_id = ? AND msf_session_id = ?
-    """, (engagement_id, msf_session_id))
+    """,
+        (engagement_id, msf_session_id),
+    )
 
     return cursor.rowcount > 0