zen-ai-pentest 2.0.0__py3-none-any.whl

core/input_validator.py

@@ -0,0 +1,353 @@
+"""
+Input Validation & Sanitization
+Prevents injection attacks and ensures data integrity
+"""
+
+import html
+import logging
+import re
+import subprocess
+from dataclasses import dataclass
+from typing import Any, List, Optional, Pattern
+from urllib.parse import urlparse
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ValidationRule:
+    """Validation rule configuration"""
+
+    pattern: Pattern
+    max_length: int = 255
+    allow_empty: bool = False
+    error_message: str = "Invalid input"
+
+
+class InputValidator:
+    """
+    Centralized input validation for all user inputs
+    """
+
+    # Domain regex (RFC compliant)
+    DOMAIN_PATTERN = re.compile(
+        r"^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)*"
+        r"[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?$"
+    )
+
+    # IP address regex
+    IP_PATTERN = re.compile(
+        r"^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}"
+        r"(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$"
+    )
+
+    # Email regex (simplified)
+    EMAIL_PATTERN = re.compile(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$")
+
+    # Safe filename regex
+    FILENAME_PATTERN = re.compile(r"^[\w\-. ]+$")
+
+    # Command injection pattern
+    DANGEROUS_CHARS = re.compile(r'[;&|`$(){}[\]\\\'"<>]')
+
+    # Path traversal pattern
+    PATH_TRAVERSAL = re.compile(r"\.\.|^/|\\")
+
+    def __init__(self):
+        self.rules: dict[str, ValidationRule] = {
+            "domain": ValidationRule(
+                pattern=self.DOMAIN_PATTERN,
+                max_length=253,
+                error_message="Invalid domain format",
+            ),
+            "ip": ValidationRule(
+                pattern=self.IP_PATTERN,
+                max_length=15,
+                error_message="Invalid IP address",
+            ),
+            "email": ValidationRule(
+                pattern=self.EMAIL_PATTERN,
+                max_length=254,
+                error_message="Invalid email format",
+            ),
+            "filename": ValidationRule(
+                pattern=self.FILENAME_PATTERN,
+                max_length=255,
+                error_message="Invalid filename (special chars not allowed)",
+            ),
+        }
+
+    def validate_domain(self, domain: str) -> Optional[str]:
+        """Validate and sanitize domain name"""
+        if not domain:
+            return None
+
+        # Strip whitespace and lowercase
+        domain = domain.strip().lower()
+
+        # Check length
+        if len(domain) > 253:
+            logger.warning(f"Domain too long: {domain[:50]}...")
+            return None
+
+        # Validate pattern
+        if not self.DOMAIN_PATTERN.match(domain):
+            logger.warning(f"Invalid domain format: {domain}")
+            return None
+
+        # Check for dangerous characters
+        if self.DANGEROUS_CHARS.search(domain):
+            logger.warning(f"Domain contains dangerous chars: {domain}")
+            return None
+
+        return domain
+
+    def validate_ip(self, ip: str) -> Optional[str]:
+        """Validate IP address"""
+        if not ip:
+            return None
+
+        ip = ip.strip()
+
+        if not self.IP_PATTERN.match(ip):
+            logger.warning(f"Invalid IP format: {ip}")
+            return None
+
+        return ip
+
+    def validate_email(self, email: str) -> Optional[str]:
+        """Validate email address"""
+        if not email:
+            return None
+
+        email = email.strip().lower()
+
+        if len(email) > 254:
+            logger.warning(f"Email too long: {email[:50]}...")
+            return None
+
+        if not self.EMAIL_PATTERN.match(email):
+            logger.warning(f"Invalid email format: {email}")
+            return None
+
+        return email
+
+    def validate_url(
+        self, url: str, allowed_schemes: List[str] = None
+    ) -> Optional[str]:
+        """Validate and sanitize URL"""
+        if not url:
+            return None
+
+        allowed_schemes = allowed_schemes or ["http", "https"]
+
+        try:
+            parsed = urlparse(url.strip())
+
+            # Check scheme
+            if parsed.scheme not in allowed_schemes:
+                logger.warning(f"URL scheme not allowed: {parsed.scheme}")
+                return None
+
+            # Validate hostname
+            if not parsed.hostname:
+                logger.warning(f"URL missing hostname: {url}")
+                return None
+
+            if not self.validate_domain(parsed.hostname):
+                logger.warning(f"Invalid URL hostname: {parsed.hostname}")
+                return None
+
+            # Reconstruct safe URL
+            safe_url = f"{parsed.scheme}://{parsed.hostname}"
+            if parsed.port:
+                safe_url += f":{parsed.port}"
+            if parsed.path:
+                # Sanitize path
+                safe_path = self.sanitize_path(parsed.path)
+                safe_url += safe_path
+
+            return safe_url
+
+        except Exception as e:
+            logger.warning(f"URL validation error: {e}")
+            return None
+
+    def validate_filename(self, filename: str) -> Optional[str]:
+        """Validate filename to prevent path traversal"""
+        if not filename:
+            return None
+
+        # Strip path components
+        filename = filename.split("/")[-1].split("\\")[-1]
+
+        if not self.FILENAME_PATTERN.match(filename):
+            logger.warning(f"Invalid filename: {filename}")
+            return None
+
+        # Check for path traversal
+        if self.PATH_TRAVERSAL.search(filename):
+            logger.warning(f"Path traversal detected: {filename}")
+            return None
+
+        return filename
+
+    def sanitize_path(self, path: str) -> str:
+        """Sanitize file path"""
+        if not path:
+            return "/"
+
+        # Normalize
+        path = path.replace("\\", "/")
+
+        # Remove path traversal
+        parts = []
+        for part in path.split("/"):
+            if part == ".." or part == ".":
+                continue
+            if part:
+                parts.append(part)
+
+        return "/" + "/".join(parts)
+
+    def sanitize_for_shell(self, value: str) -> str:
+        """
+        Sanitize string for safe use in shell commands.
+        Uses shell=False approach internally.
+        """
+        if not value:
+            return ""
+
+        # Remove all shell metacharacters
+        sanitized = self.DANGEROUS_CHARS.sub("", value)
+
+        return sanitized
+
+    def escape_html(self, text: str) -> str:
+        """Escape HTML entities to prevent XSS"""
+        return html.escape(str(text), quote=True)
+
+    def sanitize_llm_output(self, output: str, allowed_tags: List[str] = None) -> str:
+        """
+        Sanitize LLM output before storage/display.
+        Removes potentially dangerous content.
+        """
+        if not output:
+            return ""
+
+        # Remove null bytes
+        output = output.replace("\x00", "")
+
+        # Remove control characters except newlines and tabs
+        output = "".join(
+            char
+            for char in output
+            if char == "\n" or char == "\t" or (ord(char) >= 32 and ord(char) < 127)
+        )
+
+        # Escape HTML
+        output = self.escape_html(output)
+
+        return output
+
+
+class SecureSubprocess:
+    """
+    Secure subprocess execution wrapper.
+    Prevents shell injection by using shell=False and argument lists.
+    """
+
+    @staticmethod
+    def run(
+        command: List[str],
+        cwd: Optional[str] = None,
+        env: Optional[dict] = None,
+        timeout: int = 300,
+        capture_output: bool = True,
+        check: bool = False,
+    ) -> subprocess.CompletedProcess:
+        """
+        Execute command securely with shell=False.
+
+        Args:
+            command: List of command arguments (NEVER use shell=True)
+            cwd: Working directory
+            env: Environment variables
+            timeout: Timeout in seconds
+            capture_output: Capture stdout/stderr
+            check: Raise exception on non-zero exit
+
+        Returns:
+            CompletedProcess instance
+
+        Raises:
+            subprocess.SubprocessError: On execution failure
+        """
+        if not command:
+            raise ValueError("Command cannot be empty")
+
+        # Validate all arguments
+        validator = InputValidator()
+        for arg in command:
+            if validator.DANGEROUS_CHARS.search(str(arg)):
+                logger.error(f"Dangerous character in command argument: {arg}")
+                raise ValueError(f"Invalid command argument: {arg}")
+
+        logger.debug(f"Executing: {' '.join(command)}")
+
+        try:
+            return subprocess.run(
+                command,
+                cwd=cwd,
+                env=env,
+                timeout=timeout,
+                capture_output=capture_output,
+                text=True,
+                check=check,
+                shell=False,  # NEVER use shell=True
+            )
+        except subprocess.TimeoutExpired as e:
+            logger.error(f"Command timeout after {timeout}s: {command[0]}")
+            raise
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Command failed: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Subprocess error: {e}")
+            raise
+
+    @staticmethod
+    def validate_nuclei_args(args: List[str]) -> bool:
+        """
+        Validate Nuclei arguments for security.
+        Blocks dangerous flags.
+        """
+        blocked_flags = [
+            "-shell",
+            "-exec",
+            "-command",
+            "-cmd",
+            "--shell",
+            "--exec",
+            "--command",
+            "--cmd",
+        ]
+
+        for arg in args:
+            if any(arg.startswith(blocked) for blocked in blocked_flags):
+                logger.error(f"Blocked dangerous Nuclei flag: {arg}")
+                return False
+
+        return True
+
+
+# Global validator instance
+_validator = None
+
+
+def get_validator() -> InputValidator:
+    """Get global validator instance"""
+    global _validator
+    if _validator is None:
+        _validator = InputValidator()
+    return _validator

core/models.py

@@ -0,0 +1,288 @@
+"""
+Pydantic Models for Type Safety and Validation
+All configuration, API requests/responses use these models
+"""
+
+import re
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+# Enums
+class Severity(str, Enum):
+    CRITICAL = "critical"
+    HIGH = "high"
+    MEDIUM = "medium"
+    LOW = "low"
+    INFO = "info"
+
+
+class ScanStatus(str, Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class BackendType(str, Enum):
+    DUCKDUCKGO = "duckduckgo"
+    OPENROUTER = "openrouter"
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+
+
+# Base Models
+class TimestampedModel(BaseModel):
+    """Base model with timestamps"""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+    updated_at: Optional[datetime] = None
+
+
+class APIKeyConfig(BaseModel):
+    """Secure API key configuration"""
+
+    model_config = ConfigDict(extra="forbid")
+
+    openrouter_key: Optional[str] = Field(None, pattern=r"^sk-or-[a-zA-Z0-9]{20,}$")
+    openai_key: Optional[str] = Field(None, pattern=r"^sk-[a-zA-Z0-9]{20,}$")
+    anthropic_key: Optional[str] = Field(None, pattern=r"^sk-ant-[a-zA-Z0-9]{20,}$")
+    github_token: Optional[str] = Field(None, min_length=20)
+    shodan_key: Optional[str] = None
+
+    @field_validator("*")
+    @classmethod
+    def mask_keys(cls, v: Optional[str]) -> Optional[str]:
+        """Mask API keys in logs"""
+        if v and len(v) > 10:
+            return v  # Return full value, masking happens in repr
+        return v
+
+    def get_key(self, provider: BackendType) -> Optional[str]:
+        """Get API key for provider"""
+        mapping = {
+            BackendType.OPENROUTER: self.openrouter_key,
+            BackendType.OPENAI: self.openai_key,
+            BackendType.ANTHROPIC: self.anthropic_key,
+        }
+        return mapping.get(provider)
+
+
+class ScanConfig(BaseModel):
+    """Scan configuration with validation"""
+
+    model_config = ConfigDict(extra="forbid")
+
+    target: str = Field(..., min_length=1, max_length=253)
+    scan_type: Literal["quick", "full", "stealth"] = "quick"
+    ports: List[int] = Field(default_factory=lambda: [80, 443])
+    templates: List[str] = Field(default_factory=list)
+    timeout: int = Field(300, ge=10, le=3600)
+    concurrent: int = Field(5, ge=1, le=50)
+    follow_redirects: bool = True
+
+    @field_validator("target")
+    @classmethod
+    def validate_target(cls, v: str) -> str:
+        """Validate target is domain or IP"""
+        v = v.strip().lower()
+
+        # Check for dangerous characters
+        if re.search(r'[;&|`$(){}[\]\\\'"<>]', v):
+            raise ValueError("Target contains dangerous characters")
+
+        # Simple domain validation
+        if not re.match(r"^[a-z0-9][a-z0-9.-]*[a-z0-9]$", v):
+            # Could be IP
+            if not re.match(r"^(\d{1,3}\.){3}\d{1,3}$", v):
+                raise ValueError("Invalid target format")
+
+        return v
+
+    @field_validator("ports")
+    @classmethod
+    def validate_ports(cls, v: List[int]) -> List[int]:
+        """Validate port numbers"""
+        for port in v:
+            if not 1 <= port <= 65535:
+                raise ValueError(f"Invalid port: {port}")
+        return v
+
+
+class Finding(BaseModel):
+    """Security finding/vulnerability"""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    id: Optional[str] = None
+    title: str = Field(..., min_length=1, max_length=500)
+    description: str = Field(..., min_length=1)
+    severity: Severity
+    cvss_score: Optional[float] = Field(None, ge=0, le=10)
+    host: str
+    port: Optional[int] = Field(None, ge=1, le=65535)
+    service: Optional[str] = None
+    evidence: Optional[str] = None
+    remediation: Optional[str] = None
+    references: List[str] = Field(default_factory=list)
+    cve_ids: List[str] = Field(default_factory=list)
+    tags: List[str] = Field(default_factory=list)
+    confidence: Literal["confirmed", "likely", "possible"] = "possible"
+
+    @field_validator("cve_ids")
+    @classmethod
+    def validate_cve_format(cls, v: List[str]) -> List[str]:
+        """Validate CVE ID format"""
+        for cve in v:
+            if not re.match(r"^CVE-\d{4}-\d{4,}$", cve, re.IGNORECASE):
+                raise ValueError(f"Invalid CVE format: {cve}")
+        return [cve.upper() for cve in v]
+
+
+class ScanResult(BaseModel):
+    """Complete scan result"""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    scan_id: str
+    target: str
+    status: ScanStatus
+    started_at: datetime
+    completed_at: Optional[datetime] = None
+    findings: List[Finding] = Field(default_factory=list)
+    stats: Dict[str, Any] = Field(default_factory=dict)
+    error_message: Optional[str] = None
+
+    @property
+    def duration_seconds(self) -> Optional[float]:
+        """Calculate scan duration"""
+        if self.completed_at:
+            return (self.completed_at - self.started_at).total_seconds()
+        return None
+
+    @property
+    def severity_counts(self) -> Dict[str, int]:
+        """Count findings by severity"""
+        counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0}
+        for finding in self.findings:
+            counts[finding.severity.value] += 1
+        return counts
+
+
+class LLMRequest(BaseModel):
+    """LLM request with validation"""
+
+    model_config = ConfigDict(extra="forbid")
+
+    prompt: str = Field(..., min_length=1, max_length=10000)
+    system_prompt: Optional[str] = Field(None, max_length=5000)
+    temperature: float = Field(0.7, ge=0, le=2)
+    max_tokens: Optional[int] = Field(None, ge=1, le=32000)
+    backend: Optional[BackendType] = None
+
+    @field_validator("prompt")
+    @classmethod
+    def sanitize_prompt(cls, v: str) -> str:
+        """Basic prompt sanitization"""
+        # Remove null bytes
+        v = v.replace("\x00", "")
+        # Remove control chars except newlines/tabs
+        v = "".join(c for c in v if c == "\n" or c == "\t" or ord(c) >= 32)
+        return v.strip()
+
+
+class LLMResponse(BaseModel):
+    """LLM response model"""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    content: str
+    backend: BackendType
+    model: Optional[str] = None
+    tokens_used: Optional[int] = None
+    latency_ms: Optional[float] = None
+    cached: bool = False
+    error: Optional[str] = None
+
+    @property
+    def success(self) -> bool:
+        return self.error is None
+
+
+class SubdomainInfo(BaseModel):
+    """Subdomain information"""
+
+    name: str
+    ip_addresses: List[str] = Field(default_factory=list)
+    technologies: List[str] = Field(default_factory=list)
+    ports: List[int] = Field(default_factory=list)
+    is_alive: bool = False
+
+
+class DomainRecon(BaseModel):
+    """Domain reconnaissance results"""
+
+    domain: str
+    registrar: Optional[str] = None
+    creation_date: Optional[datetime] = None
+    expiration_date: Optional[datetime] = None
+    name_servers: List[str] = Field(default_factory=list)
+    subdomains: List[SubdomainInfo] = Field(default_factory=list)
+    emails: List[str] = Field(default_factory=list)
+    technologies: List[str] = Field(default_factory=list)
+
+
+class HealthStatus(BaseModel):
+    """System health status"""
+
+    status: Literal["healthy", "degraded", "unhealthy"]
+    version: str
+    uptime_seconds: float
+    checks: Dict[str, bool] = Field(default_factory=dict)
+    backends: Dict[str, str] = Field(default_factory=dict)
+
+
+class PaginatedResponse(BaseModel):
+    """Paginated API response"""
+
+    items: List[Any]
+    total: int
+    page: int
+    per_page: int
+    pages: int
+
+    @property
+    def has_next(self) -> bool:
+        return self.page < self.pages
+
+    @property
+    def has_prev(self) -> bool:
+        return self.page > 1
+
+
+class ReportConfig(BaseModel):
+    """Report generation configuration"""
+
+    model_config = ConfigDict(extra="forbid")
+
+    title: str = Field(..., min_length=1, max_length=200)
+    client_name: str = Field(..., min_length=1, max_length=200)
+    format: Literal["markdown", "html", "pdf", "json"] = "markdown"
+    template: str = "technical"
+    include_evidence: bool = True
+    include_remediation: bool = True
+    severity_filter: Optional[List[Severity]] = None
+
+    @field_validator("template")
+    @classmethod
+    def validate_template(cls, v: str) -> str:
+        allowed = ["executive", "technical", "detailed"]
+        if v not in allowed:
+            raise ValueError(f"Template must be one of: {allowed}")
+        return v