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

souleyez/storage/timeline_tracker.py

@@ -1,6 +1,7 @@
 """
 Timeline tracking and velocity calculations for deliverables.
 """
+
 from typing import Dict, List, Optional
 from datetime import datetime, timedelta
 from .database import get_db
@@ -8,31 +9,34 @@ from .database import get_db
 
 class TimelineTracker:
     """Track time spent on deliverables and predict completion."""
-    
+
     def __init__(self):
         self.db = get_db()
-    
+
     def start_deliverable(self, deliverable_id: int):
         """Mark deliverable as started (sets started_at timestamp)."""
         self.db.execute(
             "UPDATE deliverables SET started_at = CURRENT_TIMESTAMP, status = 'in_progress' WHERE id = ?",
-            (deliverable_id,)
+            (deliverable_id,),
         )
-        
+
         # Log activity
         try:
             from .team_collaboration import TeamCollaboration
-            deliverable = self.db.execute_one("SELECT engagement_id FROM deliverables WHERE id = ?", (deliverable_id,))
+
+            deliverable = self.db.execute_one(
+                "SELECT engagement_id FROM deliverables WHERE id = ?", (deliverable_id,)
+            )
             if deliverable:
                 tc = TeamCollaboration()
-                tc.log_activity(deliverable_id, deliverable['engagement_id'], 'started')
+                tc.log_activity(deliverable_id, deliverable["engagement_id"], "started")
         except:
             pass  # Fail silently if team collaboration not available
-    
+
     def complete_deliverable(self, deliverable_id: int, actual_hours: float = None):
         """
         Mark deliverable as completed.
-        
+
         Args:
             deliverable_id: Target deliverable
             actual_hours: Optional manual time entry (hours)
@@ -44,72 +48,84 @@ class TimelineTracker:
                        status = 'completed',
                        actual_hours = ?
                    WHERE id = ?""",
-                (actual_hours, deliverable_id)
+                (actual_hours, deliverable_id),
             )
         else:
             # Auto-calculate from started_at
             deliverable = self.db.execute_one(
                 "SELECT started_at, engagement_id FROM deliverables WHERE id = ?",
-                (deliverable_id,)
+                (deliverable_id,),
             )
-            
-            engagement_id = deliverable.get('engagement_id') if deliverable else None
-            
-            if deliverable and deliverable.get('started_at'):
-                started = datetime.fromisoformat(deliverable['started_at'].replace('Z', '+00:00'))
+
+            engagement_id = deliverable.get("engagement_id") if deliverable else None
+
+            if deliverable and deliverable.get("started_at"):
+                started = datetime.fromisoformat(
+                    deliverable["started_at"].replace("Z", "+00:00")
+                )
                 completed = datetime.now()
                 hours = (completed - started).total_seconds() / 3600
-                
+
                 self.db.execute(
                     """UPDATE deliverables 
                        SET completed_at = CURRENT_TIMESTAMP,
                            status = 'completed',
                            actual_hours = ?
                        WHERE id = ?""",
-                    (hours, deliverable_id)
+                    (hours, deliverable_id),
                 )
             else:
                 # No started_at, just mark complete
                 self.db.execute(
                     "UPDATE deliverables SET completed_at = CURRENT_TIMESTAMP, status = 'completed' WHERE id = ?",
-                    (deliverable_id,)
+                    (deliverable_id,),
                 )
-        
+
         # Log activity
         try:
             from .team_collaboration import TeamCollaboration
-            deliverable = self.db.execute_one("SELECT engagement_id FROM deliverables WHERE id = ?", (deliverable_id,))
+
+            deliverable = self.db.execute_one(
+                "SELECT engagement_id FROM deliverables WHERE id = ?", (deliverable_id,)
+            )
             if deliverable:
                 tc = TeamCollaboration()
-                tc.log_activity(deliverable_id, deliverable['engagement_id'], 'completed')
+                tc.log_activity(
+                    deliverable_id, deliverable["engagement_id"], "completed"
+                )
         except:
             pass  # Fail silently
-    
+
     def set_blocker(self, deliverable_id: int, blocker: str):
         """Set blocker text for a deliverable."""
         self.db.execute(
             "UPDATE deliverables SET blocker = ?, status = 'pending' WHERE id = ?",
-            (blocker, deliverable_id)
+            (blocker, deliverable_id),
         )
-    
+
     def clear_blocker(self, deliverable_id: int):
         """Clear blocker for a deliverable."""
         self.db.execute(
-            "UPDATE deliverables SET blocker = NULL WHERE id = ?",
-            (deliverable_id,)
+            "UPDATE deliverables SET blocker = NULL WHERE id = ?", (deliverable_id,)
         )
-    
+
     def get_phase_breakdown(self, engagement_id: int) -> Dict:
         """
         Get time breakdown by PTES phase.
-        
+
         Returns:
             Dict with phase stats (count, completed, hours, etc.)
         """
-        phases = ['reconnaissance', 'enumeration', 'exploitation', 'post_exploitation', 'techniques']
-        
+        phases = [
+            "reconnaissance",
+            "enumeration",
+            "exploitation",
+            "post_exploitation",
+            "techniques",
+        ]
+
         breakdown = {}
-        
+
         for phase in phases:
             deliverables = self.db.execute(
                 """SELECT COUNT(*) as total,
@@ -118,25 +134,29 @@ class TimelineTracker:
                           SUM(actual_hours) as actual_hours
                    FROM deliverables 
                    WHERE engagement_id = ? AND category = ?""",
-                (engagement_id, phase)
+                (engagement_id, phase),
             )
-            
+
             if deliverables:
                 stats = deliverables[0]
                 breakdown[phase] = {
-                    'total': stats['total'] or 0,
-                    'completed': stats['completed'] or 0,
-                    'estimated_hours': stats['estimated_hours'] or 0,
-                    'actual_hours': stats['actual_hours'] or 0,
-                    'completion_rate': (stats['completed'] / stats['total'] * 100) if stats['total'] > 0 else 0
+                    "total": stats["total"] or 0,
+                    "completed": stats["completed"] or 0,
+                    "estimated_hours": stats["estimated_hours"] or 0,
+                    "actual_hours": stats["actual_hours"] or 0,
+                    "completion_rate": (
+                        (stats["completed"] / stats["total"] * 100)
+                        if stats["total"] > 0
+                        else 0
+                    ),
                 }
-        
+
         return breakdown
-    
+
     def calculate_velocity(self, engagement_id: int) -> Dict:
         """
         Calculate delivery velocity (deliverables per hour).
-        
+
         Returns:
             Dict with velocity metrics
         """
@@ -148,30 +168,34 @@ class TimelineTracker:
                    AVG(actual_hours) as avg_hours_per_deliverable
                FROM deliverables 
                WHERE engagement_id = ? AND actual_hours > 0""",
-            (engagement_id,)
+            (engagement_id,),
         )
-        
-        if not stats or not stats['total_hours']:
+
+        if not stats or not stats["total_hours"]:
             return {
-                'velocity': 0,
-                'avg_hours_per_deliverable': 0,
-                'total_hours': 0,
-                'completed_deliverables': 0
+                "velocity": 0,
+                "avg_hours_per_deliverable": 0,
+                "total_hours": 0,
+                "completed_deliverables": 0,
             }
-        
-        velocity = stats['completed_deliverables'] / stats['total_hours'] if stats['total_hours'] > 0 else 0
-        
+
+        velocity = (
+            stats["completed_deliverables"] / stats["total_hours"]
+            if stats["total_hours"] > 0
+            else 0
+        )
+
         return {
-            'velocity': velocity,  # deliverables per hour
-            'avg_hours_per_deliverable': stats['avg_hours_per_deliverable'] or 0,
-            'total_hours': stats['total_hours'] or 0,
-            'completed_deliverables': stats['completed_deliverables'] or 0
+            "velocity": velocity,  # deliverables per hour
+            "avg_hours_per_deliverable": stats["avg_hours_per_deliverable"] or 0,
+            "total_hours": stats["total_hours"] or 0,
+            "completed_deliverables": stats["completed_deliverables"] or 0,
         }
-    
+
     def project_completion(self, engagement_id: int) -> Dict:
         """
         Project completion date based on current velocity.
-        
+
         Returns:
             Dict with projection data
         """
@@ -183,43 +207,43 @@ class TimelineTracker:
                    SUM(CASE WHEN status != 'completed' THEN 1 ELSE 0 END) as remaining
                FROM deliverables 
                WHERE engagement_id = ?""",
-            (engagement_id,)
+            (engagement_id,),
         )
-        
-        if not totals or totals['remaining'] == 0:
+
+        if not totals or totals["remaining"] == 0:
             return {
-                'status': 'complete',
-                'remaining_deliverables': 0,
-                'projected_hours': 0,
-                'projected_date': None
+                "status": "complete",
+                "remaining_deliverables": 0,
+                "projected_hours": 0,
+                "projected_date": None,
             }
-        
+
         velocity = self.calculate_velocity(engagement_id)
-        
-        if velocity['velocity'] == 0:
+
+        if velocity["velocity"] == 0:
             # No historical data, use average estimate
-            avg_estimate = velocity['avg_hours_per_deliverable']
+            avg_estimate = velocity["avg_hours_per_deliverable"]
             if avg_estimate == 0:
                 avg_estimate = 2.0  # Default assumption: 2 hours per deliverable
-            
-            projected_hours = totals['remaining'] * avg_estimate
+
+            projected_hours = totals["remaining"] * avg_estimate
         else:
             # Use velocity to project
-            projected_hours = totals['remaining'] / velocity['velocity']
-        
+            projected_hours = totals["remaining"] / velocity["velocity"]
+
         # Calculate projected completion date (assuming 8-hour work days)
         work_days = projected_hours / 8
         projected_date = datetime.now() + timedelta(days=work_days)
-        
+
         return {
-            'status': 'in_progress',
-            'remaining_deliverables': totals['remaining'],
-            'projected_hours': round(projected_hours, 1),
-            'projected_days': round(work_days, 1),
-            'projected_date': projected_date.strftime('%Y-%m-%d'),
-            'velocity': round(velocity['velocity'], 2)
+            "status": "in_progress",
+            "remaining_deliverables": totals["remaining"],
+            "projected_hours": round(projected_hours, 1),
+            "projected_days": round(work_days, 1),
+            "projected_date": projected_date.strftime("%Y-%m-%d"),
+            "velocity": round(velocity["velocity"], 2),
         }
-    
+
     def get_blockers(self, engagement_id: int) -> List[Dict]:
         """Get all deliverables with blockers."""
         return self.db.execute(
@@ -234,9 +258,9 @@ class TimelineTracker:
                        WHEN 'low' THEN 4
                        ELSE 5
                    END""",
-            (engagement_id,)
+            (engagement_id,),
         )
-    
+
     def get_in_progress(self, engagement_id: int) -> List[Dict]:
         """Get all in-progress deliverables with time stats."""
         return self.db.execute(
@@ -244,9 +268,9 @@ class TimelineTracker:
                FROM deliverables 
                WHERE engagement_id = ? AND status = 'in_progress'
                ORDER BY started_at DESC""",
-            (engagement_id,)
+            (engagement_id,),
         )
-    
+
     def get_timeline_summary(self, engagement_id: int) -> Dict:
         """Get comprehensive timeline summary."""
         phase_breakdown = self.get_phase_breakdown(engagement_id)
@@ -254,15 +278,15 @@ class TimelineTracker:
         projection = self.project_completion(engagement_id)
         blockers = self.get_blockers(engagement_id)
         in_progress = self.get_in_progress(engagement_id)
-        
+
         return {
-            'phase_breakdown': phase_breakdown,
-            'velocity': velocity,
-            'projection': projection,
-            'blockers': blockers,
-            'in_progress': in_progress
+            "phase_breakdown": phase_breakdown,
+            "velocity": velocity,
+            "projection": projection,
+            "blockers": blockers,
+            "in_progress": in_progress,
         }
-    
+
     def update_engagement_hours(self, engagement_id: int):
         """Update total hours on engagement based on deliverables."""
         stats = self.db.execute_one(
@@ -271,13 +295,17 @@ class TimelineTracker:
                    SUM(actual_hours) as total_actual
                FROM deliverables 
                WHERE engagement_id = ?""",
-            (engagement_id,)
+            (engagement_id,),
         )
-        
+
         if stats:
             self.db.execute(
                 """UPDATE engagements 
                    SET estimated_hours = ?, actual_hours = ?
                    WHERE id = ?""",
-                (stats['total_estimated'] or 0, stats['total_actual'] or 0, engagement_id)
+                (
+                    stats["total_estimated"] or 0,
+                    stats["total_actual"] or 0,
+                    engagement_id,
+                ),
             )

souleyez/storage/wazuh_vulns.py

@@ -34,7 +34,7 @@ class WazuhVulnsManager:
         detection_time: datetime = None,
         published_date: str = None,
         reference_urls: List[str] = None,
-        raw_data: dict = None
+        raw_data: dict = None,
     ) -> int:
         """
         Insert or update a Wazuh vulnerability.
@@ -54,41 +54,40 @@ class WazuhVulnsManager:
                 AND COALESCE(package_name, '') = ?
         """
         existing = self.db.execute_one(
-            query,
-            (engagement_id, agent_id, cve_id, package_name or '')
+            query, (engagement_id, agent_id, cve_id, package_name or "")
         )
 
         data = {
-            'engagement_id': engagement_id,
-            'agent_id': agent_id,
-            'cve_id': cve_id,
-            'package_name': package_name,
-            'agent_name': agent_name,
-            'agent_ip': agent_ip,
-            'name': name,
-            'severity': severity,
-            'cvss_score': cvss_score,
-            'cvss_version': cvss_version,
-            'package_version': package_version,
-            'package_architecture': package_architecture,
-            'detection_time': detection_time,
-            'published_date': published_date,
-            'reference_urls': json.dumps(reference_urls) if reference_urls else None,
-            'raw_data': json.dumps(raw_data) if raw_data else None,
-            'synced_at': datetime.now().isoformat()
+            "engagement_id": engagement_id,
+            "agent_id": agent_id,
+            "cve_id": cve_id,
+            "package_name": package_name,
+            "agent_name": agent_name,
+            "agent_ip": agent_ip,
+            "name": name,
+            "severity": severity,
+            "cvss_score": cvss_score,
+            "cvss_version": cvss_version,
+            "package_version": package_version,
+            "package_architecture": package_architecture,
+            "detection_time": detection_time,
+            "published_date": published_date,
+            "reference_urls": json.dumps(reference_urls) if reference_urls else None,
+            "raw_data": json.dumps(raw_data) if raw_data else None,
+            "synced_at": datetime.now().isoformat(),
         }
 
         if existing:
             # Update existing
-            vuln_id = existing['id']
-            set_clause = ', '.join([f"{k} = ?" for k in data.keys()])
+            vuln_id = existing["id"]
+            set_clause = ", ".join([f"{k} = ?" for k in data.keys()])
             update_query = f"UPDATE wazuh_vulnerabilities SET {set_clause} WHERE id = ?"
             params = list(data.values()) + [vuln_id]
             self.db.execute(update_query, tuple(params))
             return vuln_id
         else:
             # Insert new
-            return self.db.insert('wazuh_vulnerabilities', data)
+            return self.db.insert("wazuh_vulnerabilities", data)
 
     def get_vulnerability(self, vuln_id: int) -> Optional[Dict[str, Any]]:
         """Get a vulnerability by ID."""
@@ -108,7 +107,7 @@ class WazuhVulnsManager:
         cve_id: str = None,
         status: str = None,
         verified_only: bool = False,
-        limit: int = None
+        limit: int = None,
     ) -> List[Dict[str, Any]]:
         """
         List vulnerabilities with optional filters.
@@ -193,26 +192,21 @@ class WazuhVulnsManager:
         results = self.db.execute(query, (engagement_id,))
 
         summary = {
-            'total': 0,
-            'verified': 0,
-            'by_severity': {
-                'Critical': 0,
-                'High': 0,
-                'Medium': 0,
-                'Low': 0
-            }
+            "total": 0,
+            "verified": 0,
+            "by_severity": {"Critical": 0, "High": 0, "Medium": 0, "Low": 0},
         }
 
         for row in results:
-            severity = row.get('severity', 'Low')
-            count = row.get('count', 0)
-            verified = row.get('verified_count', 0)
+            severity = row.get("severity", "Low")
+            count = row.get("count", 0)
+            verified = row.get("verified_count", 0)
 
-            summary['total'] += count
-            summary['verified'] += verified
+            summary["total"] += count
+            summary["verified"] += verified
 
-            if severity in summary['by_severity']:
-                summary['by_severity'][severity] = count
+            if severity in summary["by_severity"]:
+                summary["by_severity"][severity] = count
 
         return summary
 
@@ -225,7 +219,7 @@ class WazuhVulnsManager:
             ORDER BY cve_id
         """
         results = self.db.execute(query, (engagement_id,))
-        return [row['cve_id'] for row in results if row.get('cve_id')]
+        return [row["cve_id"] for row in results if row.get("cve_id")]
 
     def get_unique_agents(self, engagement_id: int) -> List[Dict[str, str]]:
         """Get list of unique agents with vuln counts."""
@@ -247,7 +241,7 @@ class WazuhVulnsManager:
         vuln_id: int,
         status: str,
         verified_by_scan: bool = None,
-        matched_finding_id: int = None
+        matched_finding_id: int = None,
     ) -> bool:
         """
         Update vulnerability status.
@@ -261,15 +255,15 @@ class WazuhVulnsManager:
         Returns:
             True if update succeeded
         """
-        updates = {'status': status}
+        updates = {"status": status}
 
         if verified_by_scan is not None:
-            updates['verified_by_scan'] = 1 if verified_by_scan else 0
+            updates["verified_by_scan"] = 1 if verified_by_scan else 0
 
         if matched_finding_id is not None:
-            updates['matched_finding_id'] = matched_finding_id
+            updates["matched_finding_id"] = matched_finding_id
 
-        set_clause = ', '.join([f"{k} = ?" for k in updates.keys()])
+        set_clause = ", ".join([f"{k} = ?" for k in updates.keys()])
         query = f"UPDATE wazuh_vulnerabilities SET {set_clause} WHERE id = ?"
         params = list(updates.values()) + [vuln_id]
 
@@ -305,7 +299,7 @@ class WazuhVulnsManager:
             WHERE engagement_id = ? AND agent_ip = ? AND host_id = ?
         """
         result = self.db.execute_one(count_query, (engagement_id, agent_ip, host_id))
-        return result.get('count', 0) if result else 0
+        return result.get("count", 0) if result else 0
 
     def get_unmapped(self, engagement_id: int) -> List[Dict[str, Any]]:
         """Get vulnerabilities not mapped to any host."""
@@ -322,7 +316,7 @@ class WazuhVulnsManager:
         try:
             self.db.execute(
                 "DELETE FROM wazuh_vulnerabilities WHERE engagement_id = ? AND agent_id = ?",
-                (engagement_id, agent_id)
+                (engagement_id, agent_id),
             )
             return True
         except Exception:
@@ -333,7 +327,7 @@ class WazuhVulnsManager:
         try:
             self.db.execute(
                 "DELETE FROM wazuh_vulnerabilities WHERE engagement_id = ?",
-                (engagement_id,)
+                (engagement_id,),
             )
             return True
         except Exception:
@@ -350,27 +344,31 @@ class WazuhVulnsManager:
         self,
         engagement_id: int,
         count: int,
-        status: str = 'success',
-        errors: List[str] = None
+        status: str = "success",
+        errors: List[str] = None,
     ) -> None:
         """Update sync status after a sync operation."""
         existing = self.get_sync_status(engagement_id)
 
         data = {
-            'engagement_id': engagement_id,
-            'last_sync_at': datetime.now().isoformat(),
-            'last_sync_count': count,
-            'last_sync_status': status,
-            'last_sync_errors': json.dumps(errors) if errors else None
+            "engagement_id": engagement_id,
+            "last_sync_at": datetime.now().isoformat(),
+            "last_sync_count": count,
+            "last_sync_status": status,
+            "last_sync_errors": json.dumps(errors) if errors else None,
         }
 
         if existing:
-            set_clause = ', '.join([f"{k} = ?" for k in data.keys() if k != 'engagement_id'])
+            set_clause = ", ".join(
+                [f"{k} = ?" for k in data.keys() if k != "engagement_id"]
+            )
             query = f"UPDATE wazuh_vuln_sync SET {set_clause} WHERE engagement_id = ?"
-            params = [v for k, v in data.items() if k != 'engagement_id'] + [engagement_id]
+            params = [v for k, v in data.items() if k != "engagement_id"] + [
+                engagement_id
+            ]
             self.db.execute(query, tuple(params))
         else:
-            self.db.insert('wazuh_vuln_sync', data)
+            self.db.insert("wazuh_vuln_sync", data)
 
     def is_stale(self, engagement_id: int, max_age_hours: int = 1) -> bool:
         """
@@ -385,10 +383,10 @@ class WazuhVulnsManager:
         """
         sync_status = self.get_sync_status(engagement_id)
 
-        if not sync_status or not sync_status.get('last_sync_at'):
+        if not sync_status or not sync_status.get("last_sync_at"):
             return True
 
-        last_sync = sync_status['last_sync_at']
+        last_sync = sync_status["last_sync_at"]
         if isinstance(last_sync, str):
             try:
                 last_sync = datetime.fromisoformat(last_sync)
@@ -402,16 +400,16 @@ class WazuhVulnsManager:
         """Deserialize JSON fields."""
         result = dict(row)
 
-        if result.get('reference_urls'):
+        if result.get("reference_urls"):
             try:
-                result['reference_urls'] = json.loads(result['reference_urls'])
+                result["reference_urls"] = json.loads(result["reference_urls"])
             except (json.JSONDecodeError, TypeError):
-                result['reference_urls'] = []
+                result["reference_urls"] = []
 
-        if result.get('raw_data'):
+        if result.get("raw_data"):
             try:
-                result['raw_data'] = json.loads(result['raw_data'])
+                result["raw_data"] = json.loads(result["raw_data"])
             except (json.JSONDecodeError, TypeError):
-                result['raw_data'] = {}
+                result["raw_data"] = {}
 
         return result