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

souleyez/ai/action_mapper.py

@@ -12,7 +12,7 @@ logger = logging.getLogger(__name__)
 class ActionMapper:
     """
     Maps AI recommendations to executable commands.
-    
+
     Handles translation of natural language actions into
     concrete shell commands.
     """
@@ -20,12 +20,15 @@ class ActionMapper:
     def __init__(self):
         """Initialize action mapper."""
         from ..storage.credentials import CredentialsManager
+
         self.creds_mgr = CredentialsManager()
 
-    def map_to_command(self, recommendation: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def map_to_command(
+        self, recommendation: Dict[str, Any], engagement_id: Optional[int] = None
+    ) -> Optional[str]:
         """
         Map AI recommendation to executable command.
-        
+
         Args:
             recommendation: Dict with keys:
                 - action: Human-readable action
@@ -34,150 +37,203 @@ class ActionMapper:
                 - expected_outcome: Expected result
                 - risk_level: Risk level
             engagement_id: Optional engagement ID for credential lookup
-                
+
         Returns:
             Shell command string or None if can't map
         """
-        action = recommendation.get('action', '').lower()
-        target = recommendation.get('target', '')
-        rationale = recommendation.get('rationale', '').lower()
-        
+        action = recommendation.get("action", "").lower()
+        target = recommendation.get("target", "")
+        rationale = recommendation.get("rationale", "").lower()
+
         # Extract IP and port from target
-        ip_match = re.search(r'\b(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\b', target)
-        port_match = re.search(r':(\d+)', target)
-        
+        ip_match = re.search(r"\b(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\b", target)
+        port_match = re.search(r":(\d+)", target)
+
         ip = ip_match.group(1) if ip_match else None
         port = port_match.group(1) if port_match else None
-        
+
         if not ip:
             logger.warning(f"Could not extract IP from target: {target}")
             return None
-        
+
         # PostgreSQL credential testing (check before MySQL)
-        if ('postgres' in action or 'postgresql' in action or 'psql' in action) and \
-           ('test' in action or 'login' in action or 'credential' in action):
-            return self._map_postgresql_test(ip, port or '5432', recommendation, engagement_id)
-        
+        if ("postgres" in action or "postgresql" in action or "psql" in action) and (
+            "test" in action or "login" in action or "credential" in action
+        ):
+            return self._map_postgresql_test(
+                ip, port or "5432", recommendation, engagement_id
+            )
+
         # MySQL credential testing (check before SSH)
-        if 'mysql' in action and ('test' in action or 'login' in action or 'credential' in action):
-            return self._map_mysql_test(ip, port or '3306', recommendation, engagement_id)
-        
+        if "mysql" in action and (
+            "test" in action or "login" in action or "credential" in action
+        ):
+            return self._map_mysql_test(
+                ip, port or "3306", recommendation, engagement_id
+            )
+
         # MySQL enumeration (expanded patterns - check before SSH)
-        if 'mysql' in action and ('enum' in action or 'database' in action or 'data' in action):
-            return self._map_mysql_enum(ip, port or '3306', recommendation, engagement_id)
-        
+        if "mysql" in action and (
+            "enum" in action or "database" in action or "data" in action
+        ):
+            return self._map_mysql_enum(
+                ip, port or "3306", recommendation, engagement_id
+            )
+
         # SMB credential testing
-        if ('smb' in action or 'cifs' in action or 'smb' in target.lower() or
-            ('authenticate' in action or 'login' in action or 'credential' in action) and
-            ('445' in target or '139' in target or 'smb' in rationale)):
-            return self._map_smb_test(ip, port or '445', recommendation, engagement_id)
-        
+        if (
+            "smb" in action
+            or "cifs" in action
+            or "smb" in target.lower()
+            or ("authenticate" in action or "login" in action or "credential" in action)
+            and ("445" in target or "139" in target or "smb" in rationale)
+        ):
+            return self._map_smb_test(ip, port or "445", recommendation, engagement_id)
+
         # RDP credential testing
-        if ('rdp' in action or 'rdp' in target.lower() or 'remote desktop' in action or
-            ('authenticate' in action or 'login' in action or 'credential' in action) and
-            ('3389' in target or 'rdp' in rationale)):
-            return self._map_rdp_test(ip, port or '3389', recommendation, engagement_id)
-        
+        if (
+            "rdp" in action
+            or "rdp" in target.lower()
+            or "remote desktop" in action
+            or ("authenticate" in action or "login" in action or "credential" in action)
+            and ("3389" in target or "rdp" in rationale)
+        ):
+            return self._map_rdp_test(ip, port or "3389", recommendation, engagement_id)
+
         # FTP credential testing
-        if ('ftp' in action or 'ftp' in target.lower() or
-            ('authenticate' in action or 'login' in action or 'credential' in action) and
-            ('21' in target or 'ftp' in rationale)):
-            return self._map_ftp_test(ip, port or '21', recommendation, engagement_id)
-        
+        if (
+            "ftp" in action
+            or "ftp" in target.lower()
+            or ("authenticate" in action or "login" in action or "credential" in action)
+            and ("21" in target or "ftp" in rationale)
+        ):
+            return self._map_ftp_test(ip, port or "21", recommendation, engagement_id)
+
         # SSH credential testing (expanded patterns)
-        if ('ssh' in action or 'ssh' in target.lower() or 
-            ('authenticate' in action or 'login' in action or 'credential' in action) and 
-            ('22' in target or 'ssh' in rationale or not port)):
-            return self._map_ssh_test(ip, port or '22', recommendation, engagement_id)
-        
+        if (
+            "ssh" in action
+            or "ssh" in target.lower()
+            or ("authenticate" in action or "login" in action or "credential" in action)
+            and ("22" in target or "ssh" in rationale or not port)
+        ):
+            return self._map_ssh_test(ip, port or "22", recommendation, engagement_id)
+
         # HTTP/Web enumeration (BEFORE generic nmap scan)
-        if ('http' in action or 'web' in action or 'directory' in action or 'path' in action or
-            'http' in target.lower() or 'web' in target.lower() or
-            ('enumerate' in action or 'enum' in action) and 
-            ('http' in rationale or 'web' in rationale or port in ['80', '443', '8080', '8443'])):
+        if (
+            "http" in action
+            or "web" in action
+            or "directory" in action
+            or "path" in action
+            or "http" in target.lower()
+            or "web" in target.lower()
+            or ("enumerate" in action or "enum" in action)
+            and (
+                "http" in rationale
+                or "web" in rationale
+                or port in ["80", "443", "8080", "8443"]
+            )
+        ):
             return self._map_http_enum(ip, port, recommendation)
-        
+
         # Nmap port scan (keep as fallback, but more specific)
-        if ('nmap' in action or 'port scan' in action or 
-            ('scan' in action and 'port' in action) or
-            ('discover' in action and 'port' in action)):
+        if (
+            "nmap" in action
+            or "port scan" in action
+            or ("scan" in action and "port" in action)
+            or ("discover" in action and "port" in action)
+        ):
             return self._map_nmap_scan(ip, recommendation)
-        
+
         # Service enumeration
-        if 'enumerate' in action and 'service' in action:
+        if "enumerate" in action and "service" in action:
             return self._map_service_enum(ip, port, recommendation)
-        
+
         logger.warning(f"No command mapping found for action: {action}")
         return None
 
-    def _find_credentials(self, engagement_id: Optional[int], service: str, ip: Optional[str] = None) -> Optional[tuple]:
+    def _find_credentials(
+        self, engagement_id: Optional[int], service: str, ip: Optional[str] = None
+    ) -> Optional[tuple]:
         """
         Find valid credentials from database.
-        
+
         Args:
             engagement_id: Engagement ID
             service: Service name (ssh, mysql, etc.)
             ip: Optional IP to filter by
-            
+
         Returns:
             Tuple of (username, password) or None
         """
         if not engagement_id:
             return None
-        
+
         try:
             creds = self.creds_mgr.list_credentials(engagement_id)
-            
+
             # Look for valid credentials for this service
             for cred in creds:
-                cred_service = cred.get('service', '').lower()
-                cred_status = cred.get('status', 'untested')
-                
+                cred_service = cred.get("service", "").lower()
+                cred_status = cred.get("status", "untested")
+
                 if service.lower() in cred_service or cred_service in service.lower():
                     # Prefer valid credentials, but fall back to untested
-                    if cred_status == 'valid':
-                        return (cred.get('username'), cred.get('password'))
-            
+                    if cred_status == "valid":
+                        return (cred.get("username"), cred.get("password"))
+
             # If no valid creds, try untested for same service
             for cred in creds:
-                cred_service = cred.get('service', '').lower()
-                cred_status = cred.get('status', 'untested')
-                
-                if (service.lower() in cred_service or cred_service in service.lower()) and cred_status == 'untested':
-                    return (cred.get('username'), cred.get('password'))
-            
+                cred_service = cred.get("service", "").lower()
+                cred_status = cred.get("status", "untested")
+
+                if (
+                    service.lower() in cred_service or cred_service in service.lower()
+                ) and cred_status == "untested":
+                    return (cred.get("username"), cred.get("password"))
+
             # If still nothing and service is mysql, try SSH credentials (common on same host)
-            if service.lower() == 'mysql':
+            if service.lower() == "mysql":
                 for cred in creds:
-                    if cred.get('service', '').lower() == 'ssh' and cred.get('status') == 'valid':
-                        logger.info("Using SSH credentials for MySQL (common username on same host)")
-                        return (cred.get('username'), cred.get('password'))
-                    
+                    if (
+                        cred.get("service", "").lower() == "ssh"
+                        and cred.get("status") == "valid"
+                    ):
+                        logger.info(
+                            "Using SSH credentials for MySQL (common username on same host)"
+                        )
+                        return (cred.get("username"), cred.get("password"))
+
         except Exception as e:
             logger.error(f"Failed to lookup credentials: {e}")
-        
+
         return None
 
-    def _map_ssh_test(self, ip: str, port: str, rec: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def _map_ssh_test(
+        self,
+        ip: str,
+        port: str,
+        rec: Dict[str, Any],
+        engagement_id: Optional[int] = None,
+    ) -> Optional[str]:
         """Map SSH credential testing to command."""
         # Try to extract credentials from action or rationale
-        action = rec.get('action', '')
-        rationale = rec.get('rationale', '')
+        action = rec.get("action", "")
+        rationale = rec.get("rationale", "")
         text = f"{action} {rationale}".lower()
-        
+
         # Look for credential patterns
-        cred_match = re.search(r'(\w+):(\w+)', text)
+        cred_match = re.search(r"(\w+):(\w+)", text)
         if cred_match:
             username = cred_match.group(1)
             password = cred_match.group(2)
         else:
             # Try to find credentials from database
-            creds = self._find_credentials(engagement_id, 'ssh', ip)
+            creds = self._find_credentials(engagement_id, "ssh", ip)
             if not creds:
                 logger.warning("SSH test requested but no credentials found")
                 return None
             username, password = creds
-        
+
         return (
             f"sshpass -p '{password}' ssh -o ConnectTimeout=5 "
             f"-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null "
@@ -185,20 +241,26 @@ class ActionMapper:
             f"-p {port} {username}@{ip} 'echo SSH_SUCCESS && whoami && id'"
         )
 
-    def _map_ftp_test(self, ip: str, port: str, rec: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def _map_ftp_test(
+        self,
+        ip: str,
+        port: str,
+        rec: Dict[str, Any],
+        engagement_id: Optional[int] = None,
+    ) -> Optional[str]:
         """Map FTP credential testing to command."""
-        action = rec.get('action', '')
-        rationale = rec.get('rationale', '')
+        action = rec.get("action", "")
+        rationale = rec.get("rationale", "")
         text = f"{action} {rationale}".lower()
 
         # Look for credential patterns
-        cred_match = re.search(r'(\w+):(\w+)', text)
+        cred_match = re.search(r"(\w+):(\w+)", text)
         if cred_match:
             username = cred_match.group(1)
             password = cred_match.group(2)
         else:
             # Try to find credentials from database
-            creds = self._find_credentials(engagement_id, 'ftp', ip)
+            creds = self._find_credentials(engagement_id, "ftp", ip)
             if not creds:
                 logger.warning("FTP test requested but no credentials found")
                 return None
@@ -210,20 +272,26 @@ class ActionMapper:
             f"echo 'FTP_FAILED'"
         )
 
-    def _map_smb_test(self, ip: str, port: str, rec: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def _map_smb_test(
+        self,
+        ip: str,
+        port: str,
+        rec: Dict[str, Any],
+        engagement_id: Optional[int] = None,
+    ) -> Optional[str]:
         """Map SMB credential testing to command."""
-        action = rec.get('action', '')
-        rationale = rec.get('rationale', '')
+        action = rec.get("action", "")
+        rationale = rec.get("rationale", "")
         text = f"{action} {rationale}".lower()
 
         # Look for credential patterns
-        cred_match = re.search(r'(\w+):(\w+)', text)
+        cred_match = re.search(r"(\w+):(\w+)", text)
         if cred_match:
             username = cred_match.group(1)
             password = cred_match.group(2)
         else:
             # Try to find credentials from database
-            creds = self._find_credentials(engagement_id, 'smb', ip)
+            creds = self._find_credentials(engagement_id, "smb", ip)
             if not creds:
                 logger.warning("SMB test requested but no credentials found")
                 return None
@@ -235,20 +303,26 @@ class ActionMapper:
             f"-N 2>&1 || echo 'SMB_FAILED'"
         )
 
-    def _map_rdp_test(self, ip: str, port: str, rec: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def _map_rdp_test(
+        self,
+        ip: str,
+        port: str,
+        rec: Dict[str, Any],
+        engagement_id: Optional[int] = None,
+    ) -> Optional[str]:
         """Map RDP credential testing to command."""
-        action = rec.get('action', '')
-        rationale = rec.get('rationale', '')
+        action = rec.get("action", "")
+        rationale = rec.get("rationale", "")
         text = f"{action} {rationale}".lower()
 
         # Look for credential patterns
-        cred_match = re.search(r'(\w+):(\w+)', text)
+        cred_match = re.search(r"(\w+):(\w+)", text)
         if cred_match:
             username = cred_match.group(1)
             password = cred_match.group(2)
         else:
             # Try to find credentials from database
-            creds = self._find_credentials(engagement_id, 'rdp', ip)
+            creds = self._find_credentials(engagement_id, "rdp", ip)
             if not creds:
                 logger.warning("RDP test requested but no credentials found")
                 return None
@@ -260,73 +334,91 @@ class ActionMapper:
             f"/cert-ignore +auth-only /sec:nla 2>&1 || echo 'RDP_FAILED'"
         )
 
-    def _map_mysql_test(self, ip: str, port: str, rec: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def _map_mysql_test(
+        self,
+        ip: str,
+        port: str,
+        rec: Dict[str, Any],
+        engagement_id: Optional[int] = None,
+    ) -> Optional[str]:
         """Map MySQL credential testing to command."""
-        action = rec.get('action', '')
-        rationale = rec.get('rationale', '')
+        action = rec.get("action", "")
+        rationale = rec.get("rationale", "")
         text = f"{action} {rationale}".lower()
-        
+
         # Look for credential patterns
-        cred_match = re.search(r'(\w+):(\w+)', text)
+        cred_match = re.search(r"(\w+):(\w+)", text)
         if cred_match:
             username = cred_match.group(1)
             password = cred_match.group(2)
         else:
             # Try to find credentials from database
-            creds = self._find_credentials(engagement_id, 'mysql', ip)
+            creds = self._find_credentials(engagement_id, "mysql", ip)
             if not creds:
                 logger.warning("MySQL test requested but no credentials found")
                 return None
             username, password = creds
-        
+
         return (
             f"mysql -h {ip} -P {port} -u {username} -p'{password}' "
             f"--skip-ssl "
             f"-e 'SELECT version();' 2>&1"
         )
 
-    def _map_mysql_enum(self, ip: str, port: str, rec: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def _map_mysql_enum(
+        self,
+        ip: str,
+        port: str,
+        rec: Dict[str, Any],
+        engagement_id: Optional[int] = None,
+    ) -> Optional[str]:
         """Map MySQL enumeration to command."""
-        action = rec.get('action', '')
-        rationale = rec.get('rationale', '')
+        action = rec.get("action", "")
+        rationale = rec.get("rationale", "")
         text = f"{action} {rationale}".lower()
-        
+
         # Look for credentials
-        cred_match = re.search(r'(\w+):(\w+)', text)
+        cred_match = re.search(r"(\w+):(\w+)", text)
         if cred_match:
             username = cred_match.group(1)
             password = cred_match.group(2)
         else:
             # Try to find credentials from database
-            creds = self._find_credentials(engagement_id, 'mysql', ip)
+            creds = self._find_credentials(engagement_id, "mysql", ip)
             if not creds:
                 logger.warning("MySQL enumeration requested but no credentials found")
                 return None
             username, password = creds
-        
+
         return (
             f"mysql -h {ip} -P {port} -u {username} -p'{password}' "
             f"--skip-ssl "
             f"-e 'SHOW DATABASES; SELECT user,host FROM mysql.user;' 2>&1"
         )
 
-    def _map_postgresql_test(self, ip: str, port: str, rec: Dict[str, Any], engagement_id: Optional[int] = None) -> Optional[str]:
+    def _map_postgresql_test(
+        self,
+        ip: str,
+        port: str,
+        rec: Dict[str, Any],
+        engagement_id: Optional[int] = None,
+    ) -> Optional[str]:
         """Map PostgreSQL credential testing to command."""
-        action = rec.get('action', '')
-        rationale = rec.get('rationale', '')
+        action = rec.get("action", "")
+        rationale = rec.get("rationale", "")
         text = f"{action} {rationale}".lower()
 
         # Look for credential patterns
-        cred_match = re.search(r'(\w+):(\w+)', text)
+        cred_match = re.search(r"(\w+):(\w+)", text)
         if cred_match:
             username = cred_match.group(1)
             password = cred_match.group(2)
         else:
             # Try to find credentials from database
-            creds = self._find_credentials(engagement_id, 'postgresql', ip)
+            creds = self._find_credentials(engagement_id, "postgresql", ip)
             if not creds:
                 # Also try 'postgres' service name
-                creds = self._find_credentials(engagement_id, 'postgres', ip)
+                creds = self._find_credentials(engagement_id, "postgres", ip)
             if not creds:
                 logger.warning("PostgreSQL test requested but no credentials found")
                 return None
@@ -339,12 +431,12 @@ class ActionMapper:
 
     def _map_nmap_scan(self, ip: str, rec: Dict[str, Any]) -> str:
         """Map nmap scanning to command."""
-        action = rec.get('action', '').lower()
-        
+        action = rec.get("action", "").lower()
+
         # Quick scan by default
-        if 'full' in action or 'all' in action:
+        if "full" in action or "all" in action:
             return f"nmap -sV -p- {ip}"
-        elif 'quick' in action or 'fast' in action:
+        elif "quick" in action or "fast" in action:
             return f"nmap -F {ip}"
         else:
             # Default: top 1000 ports with version detection
@@ -353,51 +445,58 @@ class ActionMapper:
     def _map_http_enum(self, ip: str, port: Optional[str], rec: Dict[str, Any]) -> str:
         """
         Map HTTP/web enumeration to appropriate tools.
-        
+
         Uses gobuster for web content discovery and vulnerability identification.
         """
-        action = rec.get('action', '').lower()
-        target_info = rec.get('target', '').lower()
-        
+        action = rec.get("action", "").lower()
+        target_info = rec.get("target", "").lower()
+
         # Determine port
         if not port:
             # Try to extract from target or default based on protocol
-            if 'https' in target_info or '443' in target_info:
-                port = '443'
+            if "https" in target_info or "443" in target_info:
+                port = "443"
             else:
-                port = '80'
-        
+                port = "80"
+
         # Build URL
-        protocol = 'https' if port in ['443', '8443'] else 'http'
+        protocol = "https" if port in ["443", "8443"] else "http"
         url = f"{protocol}://{ip}:{port}"
-        
+
         # Prefer gobuster for directory/path enumeration
-        if 'directory' in action or 'path' in action or 'dir' in action or 'file' in action:
-            ssl_flag = ' -k' if protocol == 'https' else ''
-            return f"gobuster dir -u {url} -w /usr/share/wordlists/dirb/common.txt -t 10{ssl_flag}"
-        
+        if (
+            "directory" in action
+            or "path" in action
+            or "dir" in action
+            or "file" in action
+        ):
+            ssl_flag = " -k" if protocol == "https" else ""
+            return f"gobuster dir -u {url} -w data/wordlists/web_dirs_common.txt -t 10{ssl_flag}"
+
         # Use gobuster for web vulnerability and content discovery
-        elif 'vulnerab' in action or 'vuln' in action or 'scan' in action:
-            ssl_flag = ' -k' if protocol == 'https' else ''
-            return f"gobuster dir -u {url} -w /usr/share/wordlists/dirb/common.txt -t 10{ssl_flag}"
-        
+        elif "vulnerab" in action or "vuln" in action or "scan" in action:
+            ssl_flag = " -k" if protocol == "https" else ""
+            return f"gobuster dir -u {url} -w data/wordlists/web_dirs_common.txt -t 10{ssl_flag}"
+
         # Default to gobuster for general HTTP enumeration
         else:
-            ssl_flag = ' -k' if protocol == 'https' else ''
-            return f"gobuster dir -u {url} -w /usr/share/wordlists/dirb/common.txt -t 10{ssl_flag}"
+            ssl_flag = " -k" if protocol == "https" else ""
+            return f"gobuster dir -u {url} -w data/wordlists/web_dirs_common.txt -t 10{ssl_flag}"
 
-    def _map_service_enum(self, ip: str, port: Optional[str], rec: Dict[str, Any]) -> Optional[str]:
+    def _map_service_enum(
+        self, ip: str, port: Optional[str], rec: Dict[str, Any]
+    ) -> Optional[str]:
         """Map service enumeration to command."""
         if not port:
             return f"nmap -sV {ip}"
-        
-        service = rec.get('target', '').lower()
-        
-        if 'http' in service or port in ['80', '443', '8080']:
+
+        service = rec.get("target", "").lower()
+
+        if "http" in service or port in ["80", "443", "8080"]:
             return f"curl -I http://{ip}:{port}/ 2>&1"
-        elif 'ssh' in service or port == '22':
+        elif "ssh" in service or port == "22":
             return f"ssh -V {ip} 2>&1 | head -1"
-        elif 'mysql' in service or port == '3306':
+        elif "mysql" in service or port == "3306":
             return f"nmap -sV -p {port} {ip}"
         else:
             return f"nc -zv {ip} {port} 2>&1"