aiptx 2.0.2__py3-none-any.whl

aipt_v2/database/repository.py

@@ -0,0 +1,384 @@
+"""
+AIPT Repository - Database access layer
+Provides CRUD operations for all models.
+
+Usage:
+    repo = Repository("sqlite:///aipt.db")
+    project = repo.create_project("Test", "192.168.1.0/24")
+    session = repo.create_session(project.id)
+    repo.add_finding(project.id, session.id, "port", "80/tcp", "HTTP server")
+"""
+
+from datetime import datetime
+from typing import Optional, List
+from contextlib import contextmanager
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, Session as DBSession
+
+from .models import (
+    Base, Project, Session, Finding, Task,
+    SeverityLevel, TaskStatus, PhaseType
+)
+
+
+class Repository:
+    """
+    Database repository for AIPT.
+    Handles all database operations.
+    """
+
+    def __init__(self, db_url: str = "sqlite:///~/.aipt/aipt.db"):
+        """
+        Initialize repository.
+
+        Args:
+            db_url: Database URL (SQLite or PostgreSQL)
+                    SQLite: sqlite:///path/to/db.sqlite
+                    PostgreSQL: postgresql://user:pass@host:port/db
+        """
+        import os
+
+        # Expand ~ in path
+        if db_url.startswith("sqlite:///~"):
+            db_url = db_url.replace("~", os.path.expanduser("~"))
+            # Ensure directory exists
+            db_path = db_url.replace("sqlite:///", "")
+            os.makedirs(os.path.dirname(db_path), exist_ok=True)
+
+        self.engine = create_engine(db_url, echo=False)
+        self.SessionLocal = sessionmaker(bind=self.engine, expire_on_commit=False)
+
+        # Create tables
+        Base.metadata.create_all(self.engine)
+
+    @contextmanager
+    def _get_db(self):
+        """Get database session with automatic cleanup"""
+        session = self.SessionLocal()
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    # ============== Project Operations ==============
+
+    def create_project(
+        self,
+        name: str,
+        target: str,
+        description: Optional[str] = None,
+        scope: Optional[List[str]] = None,
+        config: Optional[dict] = None,
+    ) -> Project:
+        """Create a new project"""
+        with self._get_db() as session:
+            project = Project(
+                name=name,
+                target=target,
+                description=description,
+                scope=scope or [target],
+                config=config or {},
+            )
+            session.add(project)
+            session.flush()
+            session.refresh(project)
+            return project
+
+    def get_project(self, project_id: int) -> Optional[Project]:
+        """Get project by ID"""
+        with self._get_db() as session:
+            return session.query(Project).filter(Project.id == project_id).first()
+
+    def get_project_by_target(self, target: str) -> Optional[Project]:
+        """Get project by target"""
+        with self._get_db() as session:
+            return session.query(Project).filter(Project.target == target).first()
+
+    def list_projects(self, status: Optional[str] = None) -> List[Project]:
+        """List all projects"""
+        with self._get_db() as session:
+            query = session.query(Project)
+            if status:
+                query = query.filter(Project.status == status)
+            return query.order_by(Project.created_at.desc()).all()
+
+    def update_project(
+        self,
+        project_id: int,
+        **kwargs
+    ) -> Optional[Project]:
+        """Update project fields"""
+        with self._get_db() as session:
+            project = session.query(Project).filter(Project.id == project_id).first()
+            if project:
+                for key, value in kwargs.items():
+                    if hasattr(project, key):
+                        setattr(project, key, value)
+                session.flush()
+                session.refresh(project)
+            return project
+
+    def delete_project(self, project_id: int) -> bool:
+        """Delete project and all related data"""
+        with self._get_db() as session:
+            project = session.query(Project).filter(Project.id == project_id).first()
+            if project:
+                session.delete(project)
+                return True
+            return False
+
+    # ============== Session Operations ==============
+
+    def create_session(
+        self,
+        project_id: int,
+        name: Optional[str] = None,
+        phase: PhaseType = PhaseType.RECON,
+        max_iterations: int = 100,
+    ) -> Session:
+        """Create a new session"""
+        with self._get_db() as db_session:
+            session = Session(
+                project_id=project_id,
+                name=name or f"Session {datetime.now().strftime('%Y%m%d_%H%M%S')}",
+                phase=phase,
+                max_iterations=max_iterations,
+            )
+            db_session.add(session)
+            db_session.flush()
+            db_session.refresh(session)
+            return session
+
+    def get_session_by_id(self, session_id: int) -> Optional[Session]:
+        """Get session by ID"""
+        with self._get_db() as db_session:
+            return db_session.query(Session).filter(Session.id == session_id).first()
+
+    def get_active_session(self, project_id: int) -> Optional[Session]:
+        """Get active session for project"""
+        with self._get_db() as db_session:
+            return db_session.query(Session).filter(
+                Session.project_id == project_id,
+                Session.status == "running"
+            ).first()
+
+    def list_sessions(self, project_id: int) -> List[Session]:
+        """List all sessions for a project"""
+        with self._get_db() as db_session:
+            return db_session.query(Session).filter(
+                Session.project_id == project_id
+            ).order_by(Session.started_at.desc()).all()
+
+    def update_session(
+        self,
+        session_id: int,
+        **kwargs
+    ) -> Optional[Session]:
+        """Update session fields"""
+        with self._get_db() as db_session:
+            session = db_session.query(Session).filter(Session.id == session_id).first()
+            if session:
+                for key, value in kwargs.items():
+                    if hasattr(session, key):
+                        setattr(session, key, value)
+                db_session.flush()
+                db_session.refresh(session)
+            return session
+
+    def save_session_state(
+        self,
+        session_id: int,
+        state: dict,
+        memory_summary: Optional[str] = None,
+    ) -> None:
+        """Save session state for resume"""
+        self.update_session(
+            session_id,
+            state=state,
+            memory_summary=memory_summary,
+        )
+
+    def complete_session(self, session_id: int) -> None:
+        """Mark session as completed"""
+        self.update_session(
+            session_id,
+            status="completed",
+            ended_at=datetime.now(),
+        )
+
+    # ============== Finding Operations ==============
+
+    def add_finding(
+        self,
+        project_id: int,
+        session_id: Optional[int],
+        type: str,
+        value: str,
+        description: Optional[str] = None,
+        severity: str = "info",
+        phase: Optional[str] = None,
+        tool: Optional[str] = None,
+        raw_output: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> Finding:
+        """Add a new finding"""
+        with self._get_db() as db_session:
+            # Check for duplicate
+            existing = db_session.query(Finding).filter(
+                Finding.project_id == project_id,
+                Finding.type == type,
+                Finding.value == value,
+            ).first()
+
+            if existing:
+                # Update existing finding
+                if metadata:
+                    existing.extra_data.update(metadata)
+                return existing
+
+            # Create new finding
+            finding = Finding(
+                project_id=project_id,
+                session_id=session_id,
+                type=type,
+                value=value,
+                description=description,
+                severity=SeverityLevel(severity) if severity else SeverityLevel.INFO,
+                phase=PhaseType(phase) if phase else None,
+                tool=tool,
+                raw_output=raw_output,
+                extra_data=metadata or {},
+            )
+            db_session.add(finding)
+            db_session.flush()
+            db_session.refresh(finding)
+            return finding
+
+    def get_findings(
+        self,
+        project_id: int,
+        type: Optional[str] = None,
+        severity: Optional[str] = None,
+        phase: Optional[str] = None,
+        verified_only: bool = False,
+    ) -> List[Finding]:
+        """Get findings with optional filters"""
+        with self._get_db() as db_session:
+            query = db_session.query(Finding).filter(Finding.project_id == project_id)
+
+            if type:
+                query = query.filter(Finding.type == type)
+            if severity:
+                query = query.filter(Finding.severity == SeverityLevel(severity))
+            if phase:
+                query = query.filter(Finding.phase == PhaseType(phase))
+            if verified_only:
+                query = query.filter(Finding.verified == True)
+
+            return query.order_by(Finding.discovered_at.desc()).all()
+
+    def verify_finding(self, finding_id: int, verified: bool = True, notes: Optional[str] = None) -> None:
+        """Mark finding as verified"""
+        with self._get_db() as db_session:
+            finding = db_session.query(Finding).filter(Finding.id == finding_id).first()
+            if finding:
+                finding.verified = verified
+                if notes:
+                    finding.notes = notes
+
+    def mark_false_positive(self, finding_id: int, notes: Optional[str] = None) -> None:
+        """Mark finding as false positive"""
+        with self._get_db() as db_session:
+            finding = db_session.query(Finding).filter(Finding.id == finding_id).first()
+            if finding:
+                finding.false_positive = True
+                if notes:
+                    finding.notes = notes
+
+    def get_findings_summary(self, project_id: int) -> dict:
+        """Get summary of findings by severity"""
+        with self._get_db() as db_session:
+            findings = db_session.query(Finding).filter(
+                Finding.project_id == project_id,
+                Finding.false_positive == False,
+            ).all()
+
+            summary = {
+                "total": len(findings),
+                "by_severity": {},
+                "by_type": {},
+                "verified": 0,
+            }
+
+            for f in findings:
+                # By severity
+                sev = f.severity.value if f.severity else "info"
+                summary["by_severity"][sev] = summary["by_severity"].get(sev, 0) + 1
+
+                # By type
+                summary["by_type"][f.type] = summary["by_type"].get(f.type, 0) + 1
+
+                # Verified count
+                if f.verified:
+                    summary["verified"] += 1
+
+            return summary
+
+    # ============== Task Operations ==============
+
+    def add_task(
+        self,
+        session_id: int,
+        task_id: str,
+        description: str,
+        phase: str,
+        status: str = "to-do",
+    ) -> Task:
+        """Add a new task"""
+        with self._get_db() as db_session:
+            task = Task(
+                session_id=session_id,
+                task_id=task_id,
+                description=description,
+                phase=PhaseType(phase),
+                status=TaskStatus(status),
+            )
+            db_session.add(task)
+            db_session.flush()
+            db_session.refresh(task)
+            return task
+
+    def update_task(
+        self,
+        task_id: int,
+        status: Optional[str] = None,
+        findings_count: Optional[int] = None,
+        notes: Optional[str] = None,
+    ) -> Optional[Task]:
+        """Update task status"""
+        with self._get_db() as db_session:
+            task = db_session.query(Task).filter(Task.id == task_id).first()
+            if task:
+                if status:
+                    task.status = TaskStatus(status)
+                    if status == "in-progress":
+                        task.started_at = datetime.now()
+                    elif status == "completed":
+                        task.completed_at = datetime.now()
+                if findings_count is not None:
+                    task.findings_count = findings_count
+                if notes:
+                    task.notes = notes
+            return task
+
+    def get_tasks(self, session_id: int) -> List[Task]:
+        """Get all tasks for a session"""
+        with self._get_db() as db_session:
+            return db_session.query(Task).filter(
+                Task.session_id == session_id
+            ).order_by(Task.created_at).all()

aipt_v2/docker/__init__.py

@@ -0,0 +1,23 @@
+"""
+AIPT Docker Module
+
+Container management for secure execution:
+- Docker sandbox for isolated command execution
+- Pre-configured images for security tools
+- Resource limits and network isolation
+- Container lifecycle management
+"""
+from __future__ import annotations
+
+from .sandbox import DockerSandbox, SandboxConfig, SandboxResult
+from .manager import ContainerManager, SecurityImage
+from .builder import ImageBuilder
+
+__all__ = [
+    "DockerSandbox",
+    "SandboxConfig",
+    "SandboxResult",
+    "ContainerManager",
+    "SecurityImage",
+    "ImageBuilder",
+]

aipt_v2/docker/builder.py

@@ -0,0 +1,260 @@
+"""
+AIPT Image Builder - Custom Docker image building for security tools
+
+Provides utilities to build custom images with:
+- Multiple security tools pre-installed
+- Custom configurations
+- Optimized for pentest workflows
+"""
+from __future__ import annotations
+
+import subprocess
+import tempfile
+from typing import Optional, List, Dict
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class ImageSpec:
+    """Specification for a custom Docker image"""
+    name: str
+    tag: str = "latest"
+    base_image: str = "kalilinux/kali-rolling"
+    packages: List[str] = field(default_factory=list)
+    pip_packages: List[str] = field(default_factory=list)
+    go_packages: List[str] = field(default_factory=list)
+    custom_commands: List[str] = field(default_factory=list)
+    environment: Dict[str, str] = field(default_factory=dict)
+    workdir: str = "/workspace"
+
+
+class ImageBuilder:
+    """
+    Build custom Docker images for security testing.
+
+    Example:
+        builder = ImageBuilder()
+        spec = ImageSpec(
+            name="aipt-recon",
+            packages=["nmap", "masscan", "subfinder"],
+            pip_packages=["httpx", "dnspython"],
+        )
+        builder.build(spec)
+    """
+
+    # Pre-defined tool sets
+    TOOL_SETS = {
+        "recon": {
+            "packages": ["nmap", "masscan", "dnsutils", "whois", "curl", "wget"],
+            "pip_packages": ["httpx", "dnspython", "shodan"],
+            "go_packages": [
+                "github.com/projectdiscovery/subfinder/v2/cmd/subfinder@latest",
+                "github.com/projectdiscovery/httpx/cmd/httpx@latest",
+                "github.com/tomnomnom/assetfinder@latest",
+            ],
+        },
+        "enum": {
+            "packages": ["nikto", "dirb", "gobuster", "ffuf"],
+            "pip_packages": ["wappalyzer", "whatweb"],
+            "go_packages": [
+                "github.com/projectdiscovery/nuclei/v3/cmd/nuclei@latest",
+                "github.com/ffuf/ffuf/v2@latest",
+            ],
+        },
+        "exploit": {
+            "packages": ["sqlmap", "hydra", "john", "hashcat"],
+            "pip_packages": ["impacket", "pwntools"],
+        },
+        "post": {
+            "packages": ["netcat-openbsd", "socat", "python3-pip"],
+            "pip_packages": ["linpeas", "pspy"],
+        },
+    }
+
+    def __init__(self, registry: Optional[str] = None):
+        """
+        Initialize image builder.
+
+        Args:
+            registry: Optional Docker registry (e.g., "ghcr.io/myorg")
+        """
+        self.registry = registry
+
+    def generate_dockerfile(self, spec: ImageSpec) -> str:
+        """Generate Dockerfile content from spec"""
+        lines = [
+            f"FROM {spec.base_image}",
+            "",
+            "# Update and install base packages",
+            "RUN apt-get update && apt-get install -y \\",
+            "    curl wget git python3 python3-pip golang-go \\",
+            "    && rm -rf /var/lib/apt/lists/*",
+            "",
+        ]
+
+        # Install system packages
+        if spec.packages:
+            lines.extend([
+                "# Install security tools",
+                "RUN apt-get update && apt-get install -y \\",
+                "    " + " \\\n    ".join(spec.packages) + " \\",
+                "    && rm -rf /var/lib/apt/lists/*",
+                "",
+            ])
+
+        # Install pip packages
+        if spec.pip_packages:
+            lines.extend([
+                "# Install Python packages",
+                "RUN pip3 install --no-cache-dir \\",
+                "    " + " \\\n    ".join(spec.pip_packages),
+                "",
+            ])
+
+        # Install Go packages
+        if spec.go_packages:
+            lines.extend([
+                "# Install Go tools",
+                "ENV GOPATH=/go",
+                "ENV PATH=$PATH:/go/bin",
+            ])
+            for pkg in spec.go_packages:
+                lines.append(f"RUN go install {pkg}")
+            lines.append("")
+
+        # Custom commands
+        if spec.custom_commands:
+            lines.extend([
+                "# Custom commands",
+            ])
+            for cmd in spec.custom_commands:
+                lines.append(f"RUN {cmd}")
+            lines.append("")
+
+        # Environment variables
+        if spec.environment:
+            lines.append("# Environment variables")
+            for key, value in spec.environment.items():
+                lines.append(f"ENV {key}={value}")
+            lines.append("")
+
+        # Set workdir
+        lines.extend([
+            f"WORKDIR {spec.workdir}",
+            "",
+            "# Default command",
+            'CMD ["/bin/bash"]',
+        ])
+
+        return "\n".join(lines)
+
+    def build(
+        self,
+        spec: ImageSpec,
+        push: bool = False,
+        no_cache: bool = False,
+    ) -> bool:
+        """
+        Build Docker image from spec.
+
+        Args:
+            spec: Image specification
+            push: Push to registry after build
+            no_cache: Build without cache
+
+        Returns:
+            True if successful
+        """
+        dockerfile_content = self.generate_dockerfile(spec)
+
+        # Create temp directory with Dockerfile
+        with tempfile.TemporaryDirectory() as tmpdir:
+            dockerfile_path = Path(tmpdir) / "Dockerfile"
+            dockerfile_path.write_text(dockerfile_content)
+
+            # Build image
+            image_name = f"{spec.name}:{spec.tag}"
+            if self.registry:
+                image_name = f"{self.registry}/{image_name}"
+
+            cmd = ["docker", "build", "-t", image_name, "."]
+            if no_cache:
+                cmd.append("--no-cache")
+
+            try:
+                result = subprocess.run(
+                    cmd,
+                    cwd=tmpdir,
+                    capture_output=True,
+                    text=True,
+                    timeout=1800,  # 30 min timeout for builds
+                )
+
+                if result.returncode != 0:
+                    print(f"Build failed: {result.stderr}")
+                    return False
+
+                if push and self.registry:
+                    return self.push(image_name)
+
+                return True
+
+            except subprocess.TimeoutExpired:
+                print("Build timed out")
+                return False
+            except Exception as e:
+                print(f"Build error: {e}")
+                return False
+
+    def push(self, image_name: str) -> bool:
+        """Push image to registry"""
+        try:
+            result = subprocess.run(
+                ["docker", "push", image_name],
+                capture_output=True,
+                text=True,
+                timeout=600,
+            )
+            return result.returncode == 0
+        except Exception:
+            return False
+
+    def build_preset(
+        self,
+        preset: str,
+        name: Optional[str] = None,
+        **kwargs
+    ) -> bool:
+        """
+        Build image from preset tool set.
+
+        Args:
+            preset: One of "recon", "enum", "exploit", "post"
+            name: Image name (defaults to "aipt-{preset}")
+            **kwargs: Additional ImageSpec parameters
+
+        Returns:
+            True if successful
+        """
+        if preset not in self.TOOL_SETS:
+            raise ValueError(f"Unknown preset: {preset}. Choose from: {list(self.TOOL_SETS.keys())}")
+
+        tools = self.TOOL_SETS[preset]
+
+        spec = ImageSpec(
+            name=name or f"aipt-{preset}",
+            packages=tools.get("packages", []),
+            pip_packages=tools.get("pip_packages", []),
+            go_packages=tools.get("go_packages", []),
+            **kwargs
+        )
+
+        return self.build(spec)
+
+    def build_all_presets(self, **kwargs) -> Dict[str, bool]:
+        """Build all preset images"""
+        results = {}
+        for preset in self.TOOL_SETS:
+            results[preset] = self.build_preset(preset, **kwargs)
+        return results