# Install dependencies
just install

# Start the server
just start

# Or with debug logging
just debug

"""YourTool description."""

from __future__ import annotations
from typing import Any

from tools.base import BaseTool
from execution import engine
from validation import validate_required, validate_timeout
from models import ToolError
from responses import success_response, error_response


class YourTool(BaseTool):
    @property
    def name(self) -> str:
        return "yourtool"  # CLI name of the tool on the Kali machine

    @property
    def description(self) -> str:
        return "Human-readable description shown to the AI agent."

    @property
    def default_timeout(self) -> int:
        return 300  # seconds

    def input_schema(self) -> dict[str, Any]:
        """JSON Schema for tool parameters. Becomes the tool's input contract."""
        return {
            "type": "object",
            "properties": {
                "target": {
                    "type": "string",
                    "description": "What to scan",
                },
                "extra_args": {
                    "type": "string",
                    "description": "Additional CLI arguments (e.g. '-v --output json')",
                },
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in seconds",
                    "default": 300,
                },
            },
            "required": ["target"],
        }

    def validate(self, arguments: dict[str, Any]) -> None:
        """Validate inputs before building command. Raise ValueError on bad input."""
        validate_required(arguments, "target")
        if "timeout" in arguments:
            validate_timeout(arguments["timeout"])

    def build_command(self, arguments: dict[str, Any]) -> list[str]:
        """Convert validated arguments into a command list (no shell injection)."""
        cmd = ["yourtool", arguments["target"]]
        if "extra_args" in arguments:
            cmd.extend(arguments["extra_args"].split())
        return cmd

    async def execute(self, arguments: dict[str, Any]) -> dict[str, Any]:
        """Run the tool via the execution engine."""
        try:
            self.validate(arguments)
        except ValueError as e:
            return error_response(ToolError(error="Validation error", details=str(e)))

        result = await engine.execute(
            command=self.build_command(arguments),
            tool=self.name,
            timeout=arguments.get("timeout", self.default_timeout),
        )
        return success_response(result)

from tools.your_tool import YourTool  # add this import

ALL_TOOLS = [
    # ... existing tools ...
    YourTool(),  # add this line
]

cd mcp-server
python3 -c "from tools.your_tool import YourTool; t = YourTool(); print(t.name, t.description)"
just test  # run smoke tests

┌─────────────────────────────────────────────────┐
│  AI Agent (Claude, GPT, Gemini, etc.)           │
└────────────────────┬────────────────────────────┘
                     │ MCP (SSE)
┌────────────────────▼────────────────────────────┐
│  FastMCP Server (server.py)                     │
│  ┌──────────────┐  ┌─────────────────────────┐  │
│  │ ToolRegistry │  │ Health Check             │  │
│  └──────┬───────┘  └─────────────────────────┘  │
│         │                                        │
│  ┌──────▼──────────────────────────────────┐     │
│  │ Tools (20 native + generic_command)     │     │
│  │  validate → build_command → execute     │     │
│  └──────┬──────────────────────────────────┘     │
│         │                                        │
│  ┌──────▼──────────────────────────────────┐     │
│  │ ExecutionEngine (async subprocess)      │     │
│  │  semaphore → run → log → return         │     │
│  └──────┬──────────────────────────────────┘     │
│         │                                        │
│  ┌──────▼──────────────────────────────────┐     │
│  │ Structured JSON Response                │     │
│  │  stdout, stderr, exit_code, timing      │     │
│  └─────────────────────────────────────────┘     │
└─────────────────────────────────────────────────┘

just install    — Install Python dependencies
just start      — Start the server (foreground)
just debug      — Start with debug logging
just test       — Run smoke tests
just health     — Check if server is running
just logs       — Tail server logs
just exec-logs  — Tail execution audit logs
just tools      — List all registered tools
just clean      — Clear logs and artifacts

mcp-server/
├── server.py              ← Entry point — FastMCP SSE server
├── config.py              ← Environment-based configuration
├── models.py              ← ExecutionResult, ToolError dataclasses
├── execution.py           ← Async subprocess execution engine
├── validation.py          ← Input validators (IP, domain, URL, etc.)
├── logging_utils.py       ← Structured JSON logging
├── responses.py           ← MCP response builders
├── security.py            ← Command sanitization
├── registry.py            ← Tool name → instance registry
├── requirements.txt       ← Python dependencies
├── test_server.py         ← Smoke tests
├── tools/
│   ├── __init__.py        ← Auto-imports all tools
│   ├── base.py            ← BaseTool abstract class
│   ├── generic_command.py ← Escape hatch
│   ├── nmap.py            ├── naabu.py
│   ├── httpx.py           ├── nuclei.py
│   ├── ffuf.py            ├── katana.py
│   ├── subfinder.py       ├── amass.py
│   ├── sqlmap.py          ├── commix.py
│   ├── wpscan.py          ├── whatweb.py
│   ├── arjun.py           ├── enum4linux.py
│   ├── netexec.py         ├── crackmapexec.py
│   ├── bloodhound.py      ├── theharvester.py
│   └── spiderfoot.py
├── utils/
│   └── process.py         ← Kill process tree helper
├── logs/                  ← Runtime logs
└── artifacts/             ← Command output artifacts