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

souleyez/core/msf_database.py

@@ -13,6 +13,7 @@ from datetime import datetime
 try:
     import psycopg2
     from psycopg2.extras import RealDictCursor
+
     PSYCOPG2_AVAILABLE = True
 except ImportError:
     PSYCOPG2_AVAILABLE = False
@@ -20,12 +21,13 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 # Tested MSF versions (major.minor)
-TESTED_MSF_VERSIONS = ['6.2', '6.3', '6.4']
+TESTED_MSF_VERSIONS = ["6.2", "6.3", "6.4"]
 KNOWN_SCHEMA_VERSION = 20230313  # MSF schema version we're compatible with
 
 
 class MSFDatabaseSchemaError(Exception):
     """Raised when MSF database schema is incompatible"""
+
     pass
 
 
@@ -39,7 +41,7 @@ class MSFDatabase:
         database: str = "msf",
         username: str = "msf",
         password: str = "",
-        workspace: str = "default"
+        workspace: str = "default",
     ):
         """
         Initialize MSF database connection
@@ -81,7 +83,7 @@ class MSFDatabase:
                 database=self.database,
                 user=self.username,
                 password=self.password,
-                connect_timeout=10  # Prevent hanging on unreachable database
+                connect_timeout=10,  # Prevent hanging on unreachable database
             )
 
             # Get workspace ID
@@ -109,8 +111,7 @@ class MSFDatabase:
         try:
             with self.conn.cursor() as cur:
                 cur.execute(
-                    "SELECT id FROM workspaces WHERE name = %s",
-                    (workspace_name,)
+                    "SELECT id FROM workspaces WHERE name = %s", (workspace_name,)
                 )
                 result = cur.fetchone()
                 return result[0] if result else None
@@ -125,10 +126,14 @@ class MSFDatabase:
             version_info = self._get_msf_version_info()
 
             if version_info:
-                version_str = version_info.get('version', 'unknown')
+                version_str = version_info.get("version", "unknown")
                 # Check if version is tested
-                major_minor = '.'.join(version_str.split('.')[:2]) if '.' in version_str else version_str
-                if major_minor not in TESTED_MSF_VERSIONS and version_str != 'unknown':
+                major_minor = (
+                    ".".join(version_str.split(".")[:2])
+                    if "." in version_str
+                    else version_str
+                )
+                if major_minor not in TESTED_MSF_VERSIONS and version_str != "unknown":
                     logger.warning(
                         f"MSF version {version_str} has not been tested with this integration. "
                         f"Tested versions: {', '.join(TESTED_MSF_VERSIONS)}. "
@@ -152,18 +157,20 @@ class MSFDatabase:
             # or check for schema_migrations to infer version
             with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
                 # Check latest schema migration
-                cur.execute("""
+                cur.execute(
+                    """
                     SELECT version
                     FROM schema_migrations
                     ORDER BY version DESC
                     LIMIT 1
-                """)
+                """
+                )
                 result = cur.fetchone()
 
                 if result:
                     return {
-                        'schema_version': str(result['version']),
-                        'version': 'unknown'  # Can't reliably determine from DB
+                        "schema_version": str(result["version"]),
+                        "version": "unknown",  # Can't reliably determine from DB
                     }
         except Exception as e:
             logger.debug(f"Could not determine MSF version: {e}")
@@ -175,33 +182,48 @@ class MSFDatabase:
         try:
             with self.conn.cursor() as cur:
                 # Check critical tables exist
-                critical_tables = ['hosts', 'services', 'vulns', 'sessions', 'workspaces']
+                critical_tables = [
+                    "hosts",
+                    "services",
+                    "vulns",
+                    "sessions",
+                    "workspaces",
+                ]
 
                 for table in critical_tables:
-                    cur.execute("""
+                    cur.execute(
+                        """
                         SELECT EXISTS (
                             SELECT FROM information_schema.tables
                             WHERE table_name = %s
                         )
-                    """, (table,))
+                    """,
+                        (table,),
+                    )
 
                     exists = cur.fetchone()[0]
                     if not exists:
-                        logger.error(f"Critical table '{table}' not found in MSF database")
+                        logger.error(
+                            f"Critical table '{table}' not found in MSF database"
+                        )
                         return False
 
                 # Validate key columns exist in hosts table
-                cur.execute("""
+                cur.execute(
+                    """
                     SELECT column_name
                     FROM information_schema.columns
                     WHERE table_name = 'hosts'
-                """)
+                """
+                )
                 columns = {row[0] for row in cur.fetchall()}
 
-                required_columns = {'id', 'address', 'workspace_id'}
+                required_columns = {"id", "address", "workspace_id"}
                 if not required_columns.issubset(columns):
                     missing = required_columns - columns
-                    logger.error(f"Required columns missing from 'hosts' table: {missing}")
+                    logger.error(
+                        f"Required columns missing from 'hosts' table: {missing}"
+                    )
                     return False
 
                 return True
@@ -219,7 +241,8 @@ class MSFDatabase:
         """
         try:
             with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
-                cur.execute("""
+                cur.execute(
+                    """
                     SELECT id, name, created_at, updated_at,
                            boundary,
                            description,
@@ -227,7 +250,8 @@ class MSFDatabase:
                            limit_to_network
                     FROM workspaces
                     ORDER BY name
-                """)
+                """
+                )
                 return [dict(row) for row in cur.fetchall()]
         except Exception as e:
             logger.error(f"Failed to list workspaces: {e}")
@@ -250,7 +274,8 @@ class MSFDatabase:
 
         try:
             with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
-                cur.execute("""
+                cur.execute(
+                    """
                     SELECT
                         id,
                         address,
@@ -270,13 +295,17 @@ class MSFDatabase:
                     FROM hosts
                     WHERE workspace_id = %s
                     ORDER BY address
-                """, (ws_id,))
+                """,
+                    (ws_id,),
+                )
                 return [dict(row) for row in cur.fetchall()]
         except Exception as e:
             logger.error(f"Failed to get hosts: {e}")
             return []
 
-    def get_services(self, workspace_id: Optional[int] = None, host_id: Optional[int] = None) -> List[Dict[str, Any]]:
+    def get_services(
+        self, workspace_id: Optional[int] = None, host_id: Optional[int] = None
+    ) -> List[Dict[str, Any]]:
         """
         Get all services from MSF database
 
@@ -343,7 +372,8 @@ class MSFDatabase:
             with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
                 # Try full query first (without v.refs which doesn't exist in some MSF versions)
                 try:
-                    cur.execute("""
+                    cur.execute(
+                        """
                         SELECT
                             v.id,
                             v.host_id,
@@ -361,16 +391,21 @@ class MSFDatabase:
                         LEFT JOIN services s ON v.service_id = s.id
                         WHERE h.workspace_id = %s
                         ORDER BY v.created_at DESC
-                    """, (ws_id,))
+                    """,
+                        (ws_id,),
+                    )
                     results = cur.fetchall()
                     return [dict(row) for row in results]
                 except Exception as schema_error:
                     # Rollback the failed transaction
                     self.conn.rollback()
-                    logger.warning(f"Full vulns query failed, trying minimal query: {schema_error}")
+                    logger.warning(
+                        f"Full vulns query failed, trying minimal query: {schema_error}"
+                    )
 
                     # Try minimal query with only essential columns
-                    cur.execute("""
+                    cur.execute(
+                        """
                         SELECT
                             v.id,
                             v.host_id,
@@ -382,12 +417,15 @@ class MSFDatabase:
                         JOIN hosts h ON v.host_id = h.id
                         WHERE h.workspace_id = %s
                         ORDER BY v.id DESC
-                    """, (ws_id,))
+                    """,
+                        (ws_id,),
+                    )
                     results = cur.fetchall()
                     return [dict(row) for row in results]
         except Exception as e:
             logger.error(f"Failed to get vulnerabilities: {e}")
             import traceback
+
             logger.error(f"Traceback: {traceback.format_exc()}")
             # Rollback transaction on error
             try:
@@ -415,7 +453,8 @@ class MSFDatabase:
             with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
                 # MSF uses a complex credential schema with cores, logins, and privates
                 # We'll join them to get usable credential data
-                cur.execute("""
+                cur.execute(
+                    """
                     SELECT
                         l.id,
                         l.core_id,
@@ -439,7 +478,9 @@ class MSFDatabase:
                     LEFT JOIN hosts h ON s.host_id = h.id
                     WHERE c.workspace_id = %s
                     ORDER BY l.created_at DESC
-                """, (ws_id,))
+                """,
+                    (ws_id,),
+                )
                 return [dict(row) for row in cur.fetchall()]
         except Exception as e:
             logger.debug(f"Failed to get credentials: {e}")
@@ -450,7 +491,9 @@ class MSFDatabase:
                 pass
             return []
 
-    def get_sessions(self, workspace_id: Optional[int] = None, active_only: bool = True) -> List[Dict[str, Any]]:
+    def get_sessions(
+        self, workspace_id: Optional[int] = None, active_only: bool = True
+    ) -> List[Dict[str, Any]]:
         """
         Get all sessions from MSF database
 
@@ -503,7 +546,9 @@ class MSFDatabase:
                 except Exception as schema_error:
                     # Rollback the failed transaction
                     self.conn.rollback()
-                    logger.warning(f"Full sessions query failed, trying minimal query: {schema_error}")
+                    logger.warning(
+                        f"Full sessions query failed, trying minimal query: {schema_error}"
+                    )
 
                     # Try minimal query with only essential columns
                     query = """
@@ -532,6 +577,7 @@ class MSFDatabase:
         except Exception as e:
             logger.error(f"Failed to get sessions: {e}")
             import traceback
+
             logger.error(f"Traceback: {traceback.format_exc()}")
             # Rollback transaction on error
             try:
@@ -552,7 +598,8 @@ class MSFDatabase:
         """
         try:
             with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
-                cur.execute("""
+                cur.execute(
+                    """
                     SELECT
                         id,
                         session_id,
@@ -565,13 +612,17 @@ class MSFDatabase:
                     FROM session_events
                     WHERE session_id = %s
                     ORDER BY created_at
-                """, (session_id,))
+                """,
+                    (session_id,),
+                )
                 return [dict(row) for row in cur.fetchall()]
         except Exception as e:
             logger.error(f"Failed to get session events: {e}")
             return []
 
-    def get_exploit_attempts(self, workspace_id: Optional[int] = None) -> List[Dict[str, Any]]:
+    def get_exploit_attempts(
+        self, workspace_id: Optional[int] = None
+    ) -> List[Dict[str, Any]]:
         """
         Get exploit attempt data from MSF database
 
@@ -591,7 +642,8 @@ class MSFDatabase:
 
         try:
             with self.conn.cursor(cursor_factory=RealDictCursor) as cur:
-                cur.execute("""
+                cur.execute(
+                    """
                     SELECT
                         v.id,
                         v.host_id,
@@ -609,13 +661,17 @@ class MSFDatabase:
                     WHERE h.workspace_id = %s
                       AND v.exploited_at IS NOT NULL
                     ORDER BY v.exploited_at DESC
-                """, (ws_id,))
+                """,
+                    (ws_id,),
+                )
                 return [dict(row) for row in cur.fetchall()]
         except Exception as e:
             logger.error(f"Failed to get exploit attempts: {e}")
             return []
 
-    def get_database_stats(self, workspace_id: Optional[int] = None, max_retries: int = 3) -> Dict[str, int]:
+    def get_database_stats(
+        self, workspace_id: Optional[int] = None, max_retries: int = 3
+    ) -> Dict[str, int]:
         """
         Get statistics about MSF database contents
 
@@ -632,16 +688,17 @@ class MSFDatabase:
             return {}
 
         stats = {
-            'hosts': 0,
-            'services': 0,
-            'vulns': 0,
-            'creds': 0,
-            'active_sessions': 0,
-            'total_sessions': 0
+            "hosts": 0,
+            "services": 0,
+            "vulns": 0,
+            "creds": 0,
+            "active_sessions": 0,
+            "total_sessions": 0,
         }
 
         # Retry logic for database locks
         import time
+
         for attempt in range(max_retries):
             try:
                 with self.conn.cursor() as cur:
@@ -649,35 +706,41 @@ class MSFDatabase:
                     try:
                         cur.execute(
                             "SELECT COUNT(*) FROM hosts WHERE workspace_id = %s",
-                            (ws_id,)
+                            (ws_id,),
                         )
-                        stats['hosts'] = cur.fetchone()[0]
+                        stats["hosts"] = cur.fetchone()[0]
                     except Exception as e:
                         logger.warning(f"Failed to count hosts: {e}")
                         self.conn.rollback()  # Clear aborted transaction state
 
                     # Count services
                     try:
-                        cur.execute("""
+                        cur.execute(
+                            """
                             SELECT COUNT(*)
                             FROM services s
                             JOIN hosts h ON s.host_id = h.id
                             WHERE h.workspace_id = %s
-                        """, (ws_id,))
-                        stats['services'] = cur.fetchone()[0]
+                        """,
+                            (ws_id,),
+                        )
+                        stats["services"] = cur.fetchone()[0]
                     except Exception as e:
                         logger.warning(f"Failed to count services: {e}")
                         self.conn.rollback()  # Clear aborted transaction state
 
                     # Count vulns
                     try:
-                        cur.execute("""
+                        cur.execute(
+                            """
                             SELECT COUNT(*)
                             FROM vulns v
                             JOIN hosts h ON v.host_id = h.id
                             WHERE h.workspace_id = %s
-                        """, (ws_id,))
-                        stats['vulns'] = cur.fetchone()[0]
+                        """,
+                            (ws_id,),
+                        )
+                        stats["vulns"] = cur.fetchone()[0]
                     except Exception as e:
                         logger.warning(f"Failed to count vulns: {e}")
                         self.conn.rollback()  # Clear aborted transaction state
@@ -685,39 +748,50 @@ class MSFDatabase:
                     # Count creds - credential schema varies by MSF version
                     # Some versions don't have credential tables at all
                     try:
-                        cur.execute("""
+                        cur.execute(
+                            """
                             SELECT COUNT(*)
                             FROM cores
                             WHERE workspace_id = %s
-                        """, (ws_id,))
-                        stats['creds'] = cur.fetchone()[0]
+                        """,
+                            (ws_id,),
+                        )
+                        stats["creds"] = cur.fetchone()[0]
                     except Exception as e:
-                        logger.debug(f"Failed to count credentials (table may not exist in this MSF version): {e}")
+                        logger.debug(
+                            f"Failed to count credentials (table may not exist in this MSF version): {e}"
+                        )
                         self.conn.rollback()  # Clear aborted transaction state
 
                     # Count active sessions
                     try:
-                        cur.execute("""
+                        cur.execute(
+                            """
                             SELECT COUNT(*)
                             FROM sessions s
                             JOIN hosts h ON s.host_id = h.id
                             WHERE h.workspace_id = %s
                               AND s.closed_at IS NULL
-                        """, (ws_id,))
-                        stats['active_sessions'] = cur.fetchone()[0]
+                        """,
+                            (ws_id,),
+                        )
+                        stats["active_sessions"] = cur.fetchone()[0]
                     except Exception as e:
                         logger.warning(f"Failed to count active sessions: {e}")
                         self.conn.rollback()  # Clear aborted transaction state
 
                     # Count total sessions
                     try:
-                        cur.execute("""
+                        cur.execute(
+                            """
                             SELECT COUNT(*)
                             FROM sessions s
                             JOIN hosts h ON s.host_id = h.id
                             WHERE h.workspace_id = %s
-                        """, (ws_id,))
-                        stats['total_sessions'] = cur.fetchone()[0]
+                        """,
+                            (ws_id,),
+                        )
+                        stats["total_sessions"] = cur.fetchone()[0]
                     except Exception as e:
                         logger.warning(f"Failed to count total sessions: {e}")
                         self.conn.rollback()  # Clear aborted transaction state
@@ -728,10 +802,14 @@ class MSFDatabase:
             except Exception as e:
                 # Check if it's a database lock error
                 error_msg = str(e).lower()
-                if 'lock' in error_msg or 'locked' in error_msg:
+                if "lock" in error_msg or "locked" in error_msg:
                     if attempt < max_retries - 1:
-                        wait_time = 0.5 * (attempt + 1)  # Incremental backoff: 0.5s, 1s, 1.5s
-                        logger.warning(f"MSF database is locked (attempt {attempt + 1}/{max_retries}), retrying in {wait_time}s...")
+                        wait_time = 0.5 * (
+                            attempt + 1
+                        )  # Incremental backoff: 0.5s, 1s, 1.5s
+                        logger.warning(
+                            f"MSF database is locked (attempt {attempt + 1}/{max_retries}), retrying in {wait_time}s..."
+                        )
                         time.sleep(wait_time)
                         # Rollback the transaction before retrying
                         try:
@@ -739,11 +817,16 @@ class MSFDatabase:
                         except:
                             pass
                     else:
-                        logger.error(f"MSF database is locked after {max_retries} attempts: {e}")
-                        logger.error("Recommendation: Wait for MSF operations to complete, then try again")
+                        logger.error(
+                            f"MSF database is locked after {max_retries} attempts: {e}"
+                        )
+                        logger.error(
+                            "Recommendation: Wait for MSF operations to complete, then try again"
+                        )
                 else:
                     logger.error(f"Failed to get database stats: {e}")
                     import traceback
+
                     logger.debug(f"Traceback: {traceback.format_exc()}")
                     break  # Don't retry for non-lock errors
 
@@ -764,7 +847,7 @@ def test_msf_database_connection(
     port: int = 5432,
     database: str = "msf",
     username: str = "msf",
-    password: str = ""
+    password: str = "",
 ) -> bool:
     """
     Test MSF database connection