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

souleyez/storage/exploit_attempts.py

@@ -8,10 +8,10 @@ from souleyez.storage.database import get_db
 
 
 # Valid status values
-STATUS_NOT_TRIED = 'not_tried'
-STATUS_ATTEMPTED = 'attempted'
-STATUS_FAILED = 'failed'
-STATUS_SUCCESS = 'success'
+STATUS_NOT_TRIED = "not_tried"
+STATUS_ATTEMPTED = "attempted"
+STATUS_FAILED = "failed"
+STATUS_SUCCESS = "success"
 
 VALID_STATUSES = [STATUS_NOT_TRIED, STATUS_ATTEMPTED, STATUS_FAILED, STATUS_SUCCESS]
 
@@ -23,7 +23,7 @@ def record_attempt(
     exploit_title: str,
     status: str,
     service_id: Optional[int] = None,
-    notes: Optional[str] = None
+    notes: Optional[str] = None,
 ) -> int:
     """
     Record an exploit attempt. Creates new record or updates existing one.
@@ -47,15 +47,19 @@ def record_attempt(
     conn = db.get_connection()
 
     # Check if record exists
-    existing = conn.execute("""
+    existing = conn.execute(
+        """
         SELECT id FROM exploit_attempts
         WHERE engagement_id = ? AND host_id = ? AND service_id IS ? AND exploit_identifier = ?
-    """, (engagement_id, host_id, service_id, exploit_identifier)).fetchone()
+    """,
+        (engagement_id, host_id, service_id, exploit_identifier),
+    ).fetchone()
 
     if existing:
         # Update existing record
-        attempt_id = existing['id']
-        conn.execute("""
+        attempt_id = existing["id"]
+        conn.execute(
+            """
             UPDATE exploit_attempts
             SET status = ?,
                 exploit_title = ?,
@@ -63,18 +67,29 @@ def record_attempt(
                 attempted_at = ?,
                 updated_at = CURRENT_TIMESTAMP
             WHERE id = ?
-        """, (status, exploit_title, notes, datetime.now(), attempt_id))
+        """,
+            (status, exploit_title, notes, datetime.now(), attempt_id),
+        )
     else:
         # Create new record
-        cursor = conn.execute("""
+        cursor = conn.execute(
+            """
             INSERT INTO exploit_attempts (
                 engagement_id, host_id, service_id, exploit_identifier,
                 exploit_title, status, attempted_at, notes
             ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
-        """, (
-            engagement_id, host_id, service_id, exploit_identifier,
-            exploit_title, status, datetime.now() if status != STATUS_NOT_TRIED else None, notes
-        ))
+        """,
+            (
+                engagement_id,
+                host_id,
+                service_id,
+                exploit_identifier,
+                exploit_title,
+                status,
+                datetime.now() if status != STATUS_NOT_TRIED else None,
+                notes,
+            ),
+        )
         attempt_id = cursor.lastrowid
 
     conn.commit()
@@ -87,7 +102,7 @@ def get_attempt_status(
     engagement_id: int,
     host_id: int,
     exploit_identifier: str,
-    service_id: Optional[int] = None
+    service_id: Optional[int] = None,
 ) -> Optional[str]:
     """
     Get the status of a specific exploit attempt.
@@ -98,14 +113,17 @@ def get_attempt_status(
     db = get_db()
     conn = db.get_connection()
 
-    result = conn.execute("""
+    result = conn.execute(
+        """
         SELECT status FROM exploit_attempts
         WHERE engagement_id = ? AND host_id = ? AND service_id IS ? AND exploit_identifier = ?
-    """, (engagement_id, host_id, service_id, exploit_identifier)).fetchone()
+    """,
+        (engagement_id, host_id, service_id, exploit_identifier),
+    ).fetchone()
 
     conn.close()
 
-    return result['status'] if result else None
+    return result["status"] if result else None
 
 
 def get_attempts_by_host(host_id: int) -> List[Dict[str, Any]]:
@@ -113,7 +131,8 @@ def get_attempts_by_host(host_id: int) -> List[Dict[str, Any]]:
     db = get_db()
     conn = db.get_connection()
 
-    results = conn.execute("""
+    results = conn.execute(
+        """
         SELECT
             ea.*,
             s.port,
@@ -125,14 +144,18 @@ def get_attempts_by_host(host_id: int) -> List[Dict[str, Any]]:
         LEFT JOIN hosts h ON ea.host_id = h.id
         WHERE ea.host_id = ?
         ORDER BY ea.updated_at DESC
-    """, (host_id,)).fetchall()
+    """,
+        (host_id,),
+    ).fetchall()
 
     conn.close()
 
     return [dict(r) for r in results]
 
 
-def get_attempts_by_engagement(engagement_id: int, status_filter: Optional[str] = None) -> List[Dict[str, Any]]:
+def get_attempts_by_engagement(
+    engagement_id: int, status_filter: Optional[str] = None
+) -> List[Dict[str, Any]]:
     """
     Get all exploit attempts for an engagement, optionally filtered by status.
 
@@ -162,7 +185,9 @@ def get_attempts_by_engagement(engagement_id: int, status_filter: Optional[str]
 
     if status_filter:
         if status_filter not in VALID_STATUSES:
-            raise ValueError(f"Invalid status filter. Must be one of: {', '.join(VALID_STATUSES)}")
+            raise ValueError(
+                f"Invalid status filter. Must be one of: {', '.join(VALID_STATUSES)}"
+            )
         query += " AND ea.status = ?"
         params.append(status_filter)
 
@@ -186,43 +211,43 @@ def get_attempt_stats(engagement_id: int) -> Dict[str, Any]:
     conn = db.get_connection()
 
     # Get status breakdown
-    status_counts = conn.execute("""
+    status_counts = conn.execute(
+        """
         SELECT status, COUNT(*) as count
         FROM exploit_attempts
         WHERE engagement_id = ?
         GROUP BY status
-    """, (engagement_id,)).fetchall()
+    """,
+        (engagement_id,),
+    ).fetchall()
 
-    stats = {
-        'total': 0,
-        'not_tried': 0,
-        'attempted': 0,
-        'failed': 0,
-        'success': 0
-    }
+    stats = {"total": 0, "not_tried": 0, "attempted": 0, "failed": 0, "success": 0}
 
     for row in status_counts:
-        stats['total'] += row['count']
-        stats[row['status']] = row['count']
+        stats["total"] += row["count"]
+        stats[row["status"]] = row["count"]
 
     # Get success rate if attempts have been made
-    total_attempts = stats['attempted'] + stats['failed'] + stats['success']
+    total_attempts = stats["attempted"] + stats["failed"] + stats["success"]
     if total_attempts > 0:
-        stats['success_rate'] = (stats['success'] / total_attempts) * 100
+        stats["success_rate"] = (stats["success"] / total_attempts) * 100
     else:
-        stats['success_rate'] = 0
+        stats["success_rate"] = 0
 
     # Get host coverage
-    host_coverage = conn.execute("""
+    host_coverage = conn.execute(
+        """
         SELECT
             COUNT(DISTINCT host_id) as hosts_with_attempts,
             (SELECT COUNT(*) FROM hosts WHERE engagement_id = ?) as total_hosts
         FROM exploit_attempts
         WHERE engagement_id = ?
-    """, (engagement_id, engagement_id)).fetchone()
+    """,
+        (engagement_id, engagement_id),
+    ).fetchone()
 
-    stats['hosts_with_attempts'] = host_coverage['hosts_with_attempts']
-    stats['total_hosts'] = host_coverage['total_hosts']
+    stats["hosts_with_attempts"] = host_coverage["hosts_with_attempts"]
+    stats["total_hosts"] = host_coverage["total_hosts"]
 
     conn.close()
 
@@ -230,9 +255,7 @@ def get_attempt_stats(engagement_id: int) -> Dict[str, Any]:
 
 
 def update_attempt_status(
-    attempt_id: int,
-    status: str,
-    notes: Optional[str] = None
+    attempt_id: int, status: str, notes: Optional[str] = None
 ) -> bool:
     """
     Update the status of an existing attempt.
@@ -251,14 +274,22 @@ def update_attempt_status(
     db = get_db()
     conn = db.get_connection()
 
-    cursor = conn.execute("""
+    cursor = conn.execute(
+        """
         UPDATE exploit_attempts
         SET status = ?,
             notes = ?,
             attempted_at = ?,
             updated_at = CURRENT_TIMESTAMP
         WHERE id = ?
-    """, (status, notes, datetime.now() if status != STATUS_NOT_TRIED else None, attempt_id))
+    """,
+        (
+            status,
+            notes,
+            datetime.now() if status != STATUS_NOT_TRIED else None,
+            attempt_id,
+        ),
+    )
 
     rows_updated = cursor.rowcount
 
@@ -287,7 +318,9 @@ def delete_attempts_by_engagement(engagement_id: int) -> int:
     db = get_db()
     conn = db.get_connection()
 
-    cursor = conn.execute("DELETE FROM exploit_attempts WHERE engagement_id = ?", (engagement_id,))
+    cursor = conn.execute(
+        "DELETE FROM exploit_attempts WHERE engagement_id = ?", (engagement_id,)
+    )
     rows_deleted = cursor.rowcount
 
     conn.commit()
@@ -327,13 +360,13 @@ def bulk_record_attempts(attempts: List[Dict[str, Any]]) -> int:
     count = 0
     for attempt in attempts:
         record_attempt(
-            engagement_id=attempt['engagement_id'],
-            host_id=attempt['host_id'],
-            exploit_identifier=attempt['exploit_identifier'],
-            exploit_title=attempt['exploit_title'],
-            status=attempt['status'],
-            service_id=attempt.get('service_id'),
-            notes=attempt.get('notes')
+            engagement_id=attempt["engagement_id"],
+            host_id=attempt["host_id"],
+            exploit_identifier=attempt["exploit_identifier"],
+            exploit_title=attempt["exploit_title"],
+            status=attempt["status"],
+            service_id=attempt.get("service_id"),
+            notes=attempt.get("notes"),
         )
         count += 1
 
@@ -352,12 +385,15 @@ def get_untried_exploits_for_host(engagement_id: int, host_id: int) -> List[str]
     db = get_db()
     conn = db.get_connection()
 
-    results = conn.execute("""
+    results = conn.execute(
+        """
         SELECT exploit_identifier
         FROM exploit_attempts
         WHERE engagement_id = ? AND host_id = ? AND status = ?
-    """, (engagement_id, host_id, STATUS_NOT_TRIED)).fetchall()
+    """,
+        (engagement_id, host_id, STATUS_NOT_TRIED),
+    ).fetchall()
 
     conn.close()
 
-    return [r['exploit_identifier'] for r in results]
+    return [r["exploit_identifier"] for r in results]

souleyez/storage/exploits.py

@@ -15,26 +15,36 @@ def add_exploit(
     url: str,
     date_published: str,
     search_term: str,
-    service_id: Optional[int] = None
+    service_id: Optional[int] = None,
 ) -> int:
     """Add a new exploit to the database."""
     db = get_db()
     conn = db.get_connection()
-    
-    cursor = conn.execute("""
+
+    cursor = conn.execute(
+        """
         INSERT INTO exploits (
             engagement_id, service_id, edb_id, title, platform, 
             type, url, date_published, search_term
         ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
-    """, (
-        engagement_id, service_id, edb_id, title, platform,
-        exploit_type, url, date_published, search_term
-    ))
-    
+    """,
+        (
+            engagement_id,
+            service_id,
+            edb_id,
+            title,
+            platform,
+            exploit_type,
+            url,
+            date_published,
+            search_term,
+        ),
+    )
+
     exploit_id = cursor.lastrowid
     conn.commit()
     conn.close()
-    
+
     return exploit_id
 
 
@@ -42,8 +52,9 @@ def get_exploits_by_engagement(engagement_id: int) -> List[Dict[str, Any]]:
     """Get all exploits for an engagement."""
     db = get_db()
     conn = db.get_connection()
-    
-    results = conn.execute("""
+
+    results = conn.execute(
+        """
         SELECT 
             e.*,
             s.port,
@@ -55,10 +66,12 @@ def get_exploits_by_engagement(engagement_id: int) -> List[Dict[str, Any]]:
         LEFT JOIN hosts h ON s.host_id = h.id
         WHERE e.engagement_id = ?
         ORDER BY e.found_at DESC
-    """, (engagement_id,)).fetchall()
-    
+    """,
+        (engagement_id,),
+    ).fetchall()
+
     conn.close()
-    
+
     return [dict(r) for r in results]
 
 
@@ -66,15 +79,18 @@ def get_exploits_by_service(service_id: int) -> List[Dict[str, Any]]:
     """Get all exploits for a specific service."""
     db = get_db()
     conn = db.get_connection()
-    
-    results = conn.execute("""
+
+    results = conn.execute(
+        """
         SELECT * FROM exploits
         WHERE service_id = ?
         ORDER BY found_at DESC
-    """, (service_id,)).fetchall()
-    
+    """,
+        (service_id,),
+    ).fetchall()
+
     conn.close()
-    
+
     return [dict(r) for r in results]
 
 
@@ -84,7 +100,8 @@ def get_exploit_stats(engagement_id: int) -> Dict[str, Any]:
     conn = db.get_connection()
 
     # Get basic stats
-    stats = conn.execute("""
+    stats = conn.execute(
+        """
         SELECT
             COUNT(*) as total,
             COUNT(DISTINCT platform) as platforms,
@@ -92,36 +109,48 @@ def get_exploit_stats(engagement_id: int) -> Dict[str, Any]:
             COUNT(DISTINCT service_id) as services_with_exploits
         FROM exploits
         WHERE engagement_id = ?
-    """, (engagement_id,)).fetchone()
+    """,
+        (engagement_id,),
+    ).fetchone()
 
-    result = dict(stats) if stats else {
-        'total': 0,
-        'platforms': 0,
-        'types': 0,
-        'services_with_exploits': 0
-    }
+    result = (
+        dict(stats)
+        if stats
+        else {"total": 0, "platforms": 0, "types": 0, "services_with_exploits": 0}
+    )
 
     # Get platform breakdown (top 5)
-    platform_breakdown = conn.execute("""
+    platform_breakdown = conn.execute(
+        """
         SELECT platform, COUNT(*) as count
         FROM exploits
         WHERE engagement_id = ?
         GROUP BY platform
         ORDER BY count DESC
         LIMIT 5
-    """, (engagement_id,)).fetchall()
-    result['platform_breakdown'] = [{'platform': row['platform'], 'count': row['count']} for row in platform_breakdown]
+    """,
+        (engagement_id,),
+    ).fetchall()
+    result["platform_breakdown"] = [
+        {"platform": row["platform"], "count": row["count"]}
+        for row in platform_breakdown
+    ]
 
     # Get type breakdown (top 5)
-    type_breakdown = conn.execute("""
+    type_breakdown = conn.execute(
+        """
         SELECT type, COUNT(*) as count
         FROM exploits
         WHERE engagement_id = ?
         GROUP BY type
         ORDER BY count DESC
         LIMIT 5
-    """, (engagement_id,)).fetchall()
-    result['type_breakdown'] = [{'type': row['type'], 'count': row['count']} for row in type_breakdown]
+    """,
+        (engagement_id,),
+    ).fetchall()
+    result["type_breakdown"] = [
+        {"type": row["type"], "count": row["count"]} for row in type_breakdown
+    ]
 
     conn.close()
 
@@ -147,7 +176,9 @@ def delete_exploits_by_engagement(engagement_id: int) -> int:
     db = get_db()
     conn = db.get_connection()
 
-    cursor = conn.execute("DELETE FROM exploits WHERE engagement_id = ?", (engagement_id,))
+    cursor = conn.execute(
+        "DELETE FROM exploits WHERE engagement_id = ?", (engagement_id,)
+    )
     rows_deleted = cursor.rowcount
 
     conn.commit()
@@ -164,7 +195,7 @@ def delete_exploits_by_ids(exploit_ids: List[int]) -> int:
     db = get_db()
     conn = db.get_connection()
 
-    placeholders = ','.join('?' * len(exploit_ids))
+    placeholders = ",".join("?" * len(exploit_ids))
     query = f"DELETE FROM exploits WHERE id IN ({placeholders})"
 
     cursor = conn.execute(query, exploit_ids)
@@ -210,7 +241,7 @@ def search_exploits(engagement_id: int, search_term: str) -> List[Dict[str, Any]
         )
         """
         # Add three parameters (one for each field) for this keyword
-        keyword_pattern = f'%{keyword}%'
+        keyword_pattern = f"%{keyword}%"
         params.extend([keyword_pattern, keyword_pattern, keyword_pattern])
 
     query += "ORDER BY e.found_at DESC"