"""
The Model-Context Protocol (MCP) module for Pentest-Tools.com. It uses the ptt API to provide tools, resources and prompts
to LLM agents that implement MCP.
"""

try:
    from mcp.server.fastmcp import Context, FastMCP
    from fastmcp.utilities import json_schema_type
    from mcp.types import SamplingMessage, TextContent
    from fastmcp.server.auth import TokenVerifier, AccessToken
    from fastmcp.server.dependencies import get_access_token
    from mcp.server.auth.settings import AuthSettings

    # TODO: It might be a good idea to use fastmcp2 here
    # from fastmcp import Context, FastMCP
except ImportError:
    # We use this as a guard, in case `pentesttools` was not installed with mcp support
    FastMCP = None


import asyncio
from copy import deepcopy
import json
import os
import sys
import time
import logging
from typing import Any, Literal

import pentesttools as ptt
from pentesttools.api import Finding

from pydantic import BaseModel, Field


class ConfirmDeleteSchema(BaseModel):
    confirm: str = Field(
        ...,
        description="Type 'yes' to confirm deletion, 'no' to cancel",
        pattern="^(yes|no)$",
    )


# JSON schema for the FindSubdomainsParams type
FIND_SUBDOMAINS_SEARCH_METHODS_SCHEMA = {
    "type": "object",
    "description": "Methods to be used for discovery. All are used by default for a deep scan. For a custom scan, you need to specify at least one.",
    "properties": {
        "passive_detection": {
            "type": "boolean",
            "default": True,
            "description": "Search historical subdomains",
        },
        "dns_records": {
            "type": "boolean",
            "default": True,
            "description": "Zone transfer",
        },
        "dns_enumeration": {
            "type": "object",
            "properties": {
                "enabled": {
                    "type": "boolean",
                    "default": True,
                },
                "wordlist_id": {
                    "type": "integer",
                    "default": 1600,
                },
            },
            "description": "DNS bruteforcing using a wordlist",
        },
        "ctr_search": {
            "type": "boolean",
            "default": True,
            "description": "Certificate Transparency Logs",
        },
        "external_api_search": {
            "type": "boolean",
            "default": True,
            "description": "Subdomains from external APIs",
        },
        "bing_search": {
            "type": "boolean",
            "default": True,
            "description": "Bing search",
        },
        "google_search": {
            "type": "boolean",
            "default": True,
            "description": "Google search",
        },
        "html_search": {
            "type": "boolean",
            "default": True,
            "description": "Subdomains extracted from HTML links and HTTP headers",
        },
        "ssl_search": {
            "type": "boolean",
            "default": True,
            "description": "Extract from SSL certificates",
        },
        "revdns_search": {
            "type": "boolean",
            "default": True,
            "description": "Reverse DNS on target IP ranges",
        },
        "alteration_search": {
            "type": "boolean",
            "default": True,
            "description": (
                "Generates permutations and alterations of the subdomain names found so far in the scan. "
                "Disabling this option should speed up scans when there is a large number of subdomains, but might reduce the number of results."
            ),
        },
        "cname_search": {
            "type": "boolean",
            "default": True,
            "description": "Extract from CNAMEs",
        },
    },
}
FindSubdomainsSearchMethods = json_schema_type.json_schema_to_type(FIND_SUBDOMAINS_SEARCH_METHODS_SCHEMA)


# JSON schema for the Discovery type
DISCOVERY_SCHEMA = {
    "type": "array",
    "items": {
        "type": "string",
        "enum": [
            "fingerprint",
            "software_vulnerabilities",
            "check_robots",
            "outdated_js",
            "untrusted_certificates",
            "client_access_policies",
            "http_debug_methods",
            "security_txt",
            "cors_misconfiguration",
            "resource_discovery",
            "sensitive_files",
            "admin_consoles",
            "interesting_files",
            "server_info_disc",
            "server_software",
            "misconfigurations",
            "graphql_endpoint",
            "fuzz_openapi_locations",
        ],
    },
    "description": "Discovery modules to run. If not set, all will run. If empty list, none will run.",
}

ATTACK_PASSIVE_SCHEMA = {
    "type": "array",
    "items": {
        "type": "string",
        "enum": [
            "security_headers",
            "cookie_security",
            "directory_listing",
            "secure_communication",
            "weak_password_submission",
            "error_debug_messages",
            "password_cleartext",
            "cross_domain_source",
            "mixed_content",
            "sensitive_data",
            "login_interfaces",
            "file_upload",
            "path_disclosure",
            "sql_statement_in_request",
            "password_in_response",
            "session_token_in_url",
            "api_endpoint",
        ],
    },
    "description": "Passive attacks to run. If not set, all will run. If empty list, none will run.",
}

# JSON schema for the AttackActive type
ATTACK_ACTIVE_SCHEMA = {
    "type": "array",
    "items": {
        "type": "string",
        "enum": [
            "xss",
            "sqli",
            "lfi",
            "oscmdi",
            "ssrf",
            "open_redirect",
            "broken_authentication",
            "php_code_injection",
            "js_code_injection",
            "ruby_code_injection",
            "python_code_injection",
            "perl_code_injection",
            "log4j_rce",
            "ssti",
            "xxe",
            "viewstate_rce",
            "prototype_pollution",
            "backup_files",
            "request_url_override",
            "http_request_smuggling",
            "session_fixation",
            "idor",
            "jwt",
            "response_header_injection",
        ],
    },
    "description": "Active attacks to run. If not set, all will run. If empty list, none will run.",
}

# JSON schema for the Spider type
SPIDER_SCHEMA = {
    "type": "object",
    "properties": {
        "approach": {
            "type": "string",
            "enum": ["auto", "classic", "spa"],
            "default": "auto",
        },
        "depth": {"type": "integer", "default": 10},
        "exclude_urls": {
            "type": "array",
            "items": {"type": "string"},
            "description": "URLs to exclude from spidering.",
        },
        "include_urls": {
            "type": "array",
            "items": {"type": "string"},
            "description": "URLs to include in spidering.",
        },
        "subdomains_in_scope": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Subdomains to include in scope.",
        },
    },
}

Discovery = json_schema_type.json_schema_to_type(DISCOVERY_SCHEMA)
AttackPassive = json_schema_type.json_schema_to_type(ATTACK_PASSIVE_SCHEMA)
AttackActive = json_schema_type.json_schema_to_type(ATTACK_ACTIVE_SCHEMA)
Spider = json_schema_type.json_schema_to_type(SPIDER_SCHEMA)


NETWORK_SCANNER_PRESET_SCHEMA = {
    "type": "string",
    "enum": ["light", "deep", "custom"],
    "default": "deep",
    "description": (
        "Determines the Network Scanner scan type that will be used. "
        "`light` uses the `version_based` scanning engine on top 100 ports. "
        "`deep` uses the `version_based`, `sniper` and `nuclei` scanning engines on top 5000 ports (TCP) or top 1000 ports (UDP). "
        "`custom` allows you to use custom configurations of scanning engines and ports."
    ),
}
NetworkScannerPreset = json_schema_type.json_schema_to_type(NETWORK_SCANNER_PRESET_SCHEMA)

NETWORK_SCANNER_SCANNING_ENGINES = {
    "type": "array",
    "description": (
        "Scanning engines that will be loaded for a scan. Can be set only if `custom` preset was selected. "
        "It must contain at least one engine."
    ),
    "items": {
        "type": "string",
        "enum": ["version_based", "sniper", "nuclei", "openvas"],
    },
    "example": ["version_based", "sniper", "nuclei", "openvas"],
}
NetworkScannerScanningEngines = json_schema_type.json_schema_to_type(NETWORK_SCANNER_SCANNING_ENGINES)


NETWORK_SCANNER_PORT_RANGE = {
    "type": "object",
    "properties": {
        "start_port": {"type": "integer"},
        "end_port": {"type": "integer"},
    },
    "example": {"start_port": 0, "end_port": 443},
}
NetworkScannerPortRange = json_schema_type.json_schema_to_type(NETWORK_SCANNER_PORT_RANGE)

NETWORK_SCANNER_PORT_LIST = {
    "type": "array",
    "items": {
        "type": "integer",
        "minimum": 1,
        "maximum": 65535,
    },
    "example": [22, 80, 443],
}
NetworkScannerPortList = json_schema_type.json_schema_to_type(NETWORK_SCANNER_PORT_LIST)

NETWORK_SCANNER_ENGINE_OPTIONS = {
    "type": "object",
    "properties": {
        "requests_per_second": {
            "type": "integer",
            "default": 50,
            "description": (
                "Specify a maximum number of requests per second between 1 and 1000. "
                "It is only valid with the `deep` preset or with the `custom` preset if `nuclei` and/or `openvas` are included in `scanning_engines`."
            ),
            "minimum": 1,
            "maximum": 1000,
        },
        "cves_to_scan": {
            "type": "array",
            "default": [],
            "description": ("Specify particular CVEs for scanning. Can be set only if `custom` preset was selected."),
            "items": {"type": "string"},
            "example": ["CVE-2024-23108", "CVE-2024-27198"],
        },
    },
}
NetworkScannerEngineOptions = json_schema_type.json_schema_to_type(NETWORK_SCANNER_ENGINE_OPTIONS)


# TODO: This is a partial schema.
SCAN_PARAMS_SCHEMA = {
    "type": "object",
    "properties": {
        "workspace_id": {
            "type": "integer",
            "nullable": True,
            "writeOnly": True,
            "description": "Workspace where the scan is started. It has to match the workspace ID of the target",
        },
        "vpn_profile_uuid": {
            "type": "string",
            "nullable": True,
            "writeOnly": True,
            "description": "VPN profile to use for the scan. If null, there will be no VPN profile used. If not specified, the profile attached to the workspace will be used.",
        },
        "vpn_scan": {
            "type": "boolean",
            "readOnly": True,
            "description": "Indicates if the scan is performed over VPN. Read-only.",
        },
        "max_scan_time": {
            "type": "integer",
            "minimum": 1,
            "maximum": 1440,
            "default": 1440,
            "writeOnly": True,
            "description": (
                "Maximum number of minutes that the scan should run. "
                "Not supported by: Sniper, tools with short scan duration (like Website Recon or ICMP Ping)."
            ),
        },
    },
}

ScanParams = json_schema_type.json_schema_to_type(SCAN_PARAMS_SCHEMA)

os.environ["PTT_USER_AGENT"] = "ptt-cli-mcp"

API_URL = "https://app.pentest-tools.com/api/v2/"
MCP_METHOD = "stdio"
HEADERS = {
    "User-Agent": "ptt-cli-mcp",
    "Accept": "*/*",
}


class PTTApiKeyVerifier(TokenVerifier):
    """Verifies that the Bearer token is a valid Pentest-Tools API key."""

    async def verify_token(self, token: str) -> AccessToken | None:
        test_headers = {**HEADERS, "Authorization": f"Bearer {token}"}
        try:
            res = await asyncio.to_thread(ptt.api.get_workspaces, headers=test_headers)
            if res.status_code == 200:
                return AccessToken(token=token, client_id="ptt-user", scopes=[])
        except Exception:
            pass
        return None


def get_authentication_headers(context: Context) -> dict[str, str]:
    """A function that gives the headers for the request. We need this to get the PTT-API-KEY header

    Returns the headers necessary for a request, with the authentication added. If the key is not found, return an empty dict.
    """
    global HEADERS

    # For stdio, we get the api key from `ptt.api.api_key`
    if MCP_METHOD == "stdio":
        HEADERS["Authorization"] = f"Bearer {ptt.api.api_key}"
        return HEADERS

    # For http, we get the api key from the auth context (Bearer token = PTT API key)
    elif MCP_METHOD == "http":
        access_token = get_access_token()
        if access_token is not None:
            headers = deepcopy(HEADERS)
            headers["Authorization"] = f"Bearer {access_token.token}"
            return headers

    return {}


def run_tool(
    tool_id: int,
    target: str,
    tool_params: dict,
    scan_params: dict,
    context: Context,
    raw: bool = False,
):
    """The main way of running a scanning tool"""

    # `HEADERS` + authentication
    headers = get_authentication_headers(context)
    if not headers:
        return "The api key is missing."

    res = ptt.api.start_scan(target, tool_id, tool_params, scan_params, headers=headers)
    res_json = res.json()

    scan_id = res_json["data"]["created_id"]
    logging.info("Started scan %i, for tool_id %s", scan_id, tool_id)

    while True:
        time.sleep(2)

        # Get the status of our scan
        status = ptt.api.get_scan_status(scan_id, headers=headers)
        status_name = status.json()["data"]["status_name"]

        if status_name == "finished":
            logging.debug("Scan status: %s" % res_json["data"])
            logging.debug("Getting JSON report")

            if raw:
                res = ptt.api.get_raw_output(scan_id, headers=headers)
                return res.text

            res = ptt.api.get_output(scan_id, headers=headers)
            output_json = json.dumps(res.json())
            return output_json


async def run_tool_long_running(
    tool_id: int,
    target: str,
    tool_params: dict,
    scan_params: dict,
    context: Context,
    raw: bool = False,
):
    """
    Run a scanning tool.

    For light scans: runs to completion with progress tracking.
    For deep/custom scans: starts the scan and returns the scan_id immediately so the
    user can keep interacting with the LLM. Use the `scan` tool with verb="get_output"
    to retrieve results when ready.
    """

    headers = get_authentication_headers(context)

    if not headers:
        return "The api key is missing."

    res = ptt.api.start_scan(target, tool_id, tool_params, scan_params, headers=headers)
    res_json = res.json()

    try:
        scan_id = res_json["data"]["created_id"]
    except KeyError:
        return f"There was a problem starting the scan: {res_json}"

    logging.info("Started scan %i, for tool_id %s", scan_id, tool_id)

    # For deep/custom scans, return the scan_id immediately so the conversation is not blocked.
    # The user can retrieve results later using the `scan` tool with verb="get_output".
    if tool_params.get("scan_type", "light") != "light":
        return json.dumps({
            "scan_id": scan_id,
            "message": (
                f"Scan {scan_id} has been started. "
                "Use the `scan` tool with verb=\"get_output\" and "
                f"scan_id={scan_id} to retrieve results when ready, "
                "or verb=\"get_status\" to check progress."
                "You can always see the current running scans at https://app.pentest-tools.com/scans"
            ),
        })

    # Status groups
    progress_statuses = {"running", "waiting"}
    finished_statuses = {"finished"}
    failed_statuses = {
        "failed to start",
        "stopped",
        "timed out",
        "aborted",
        "VPN connection error",
        "auth failed",
        "connection error",
    }

    while True:
        await asyncio.sleep(2)

        # Get the status of our scan
        status = ptt.api.get_scan_status(scan_id, headers=headers)
        status_data = status.json()["data"]
        status_name = status_data["status_name"]

        # Report progress if running or waiting
        if context is not None and status_name in progress_statuses:
            await context.report_progress(
                float(status_data.get("progress", 0.0)),
                total=100.0,
                message="Scan id "
                + str(scan_id)
                + " | Status: "
                + str(status_name)
                + f" | Progress: {float(status_data.get('progress', 0.0))}",
            )

        if status_name in finished_statuses:
            logging.debug("Scan status: %s" % status_data)
            logging.debug("Getting JSON report")

            if raw:
                res = ptt.api.get_raw_output(scan_id, headers=headers)
                return res.text

            res = ptt.api.get_output(scan_id, headers=headers)
            output_json = json.dumps(res.json())
            return output_json

        elif status_name in failed_statuses:
            # Report failure to the LLM
            if context is not None:
                await context.info(
                    f"Scan {scan_id} failed with status: {status_name}. Status message: {status_data.get('status_message', '')}"
                )
            return json.dumps(
                {
                    "error": f"Scan failed with status: {status_name}",
                    "details": status_data,
                }
            )

        # Otherwise, keep waiting and reporting progress


_sentry_dsn = os.getenv("SENTRY_DSN")
if _sentry_dsn:
    try:
        import sentry_sdk
        from sentry_sdk.integrations.mcp import MCPIntegration

        sentry_sdk.init(
            dsn=_sentry_dsn,
            traces_sample_rate=1.0,
            integrations=[MCPIntegration()],
            send_default_pii=False,
        )
    except ImportError:
        pass  # sentry-sdk not installed, skip


if FastMCP is not None:
    mcp = FastMCP(
        "ptt-http",
        host=os.getenv("MCP_HOST", "127.0.0.1"),
        port=9000,
        token_verifier=PTTApiKeyVerifier(),
        auth=AuthSettings(
            issuer_url="https://app.pentest-tools.com",
            resource_server_url="http://localhost:9000",
        ),
    )

    ## Resources
    # TODO: json data may not be suitable for LLMs, at least for the less powerful ones.
    # Resources are generally the GET requests from API
    @mcp.resource("scans://{garbage}")
    def scans(garbage: str, context: Context) -> str:
        """Get the current scans (all, not just running) as JSON string."""
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_scans(headers=headers)
        return json.dumps(res.json())

    @mcp.resource("targets://{garbage}")
    def targets(garbage: str, context: Context) -> str:
        """Get the current targets available"""
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_targets(headers=headers)
        return json.dumps(res.json())

    @mcp.resource("workspaces://{garbage}")
    def workspaces(garbage: str, context: Context) -> str:
        """Get the current workspaces available"""
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_workspaces(headers=headers)
        return json.dumps(res.json())

    @mcp.resource("http_loggers://{garbage}")
    def http_loggers(garbage: str, context: Context) -> str:
        """Get the current http_loggers available"""
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_http_loggers(headers=headers)
        return json.dumps(res.json())

    @mcp.resource("wordlists://{garbage}")
    def wordlists(garbage: str, context: Context) -> str:
        """Get the current wordlists available"""
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_wordlists(headers=headers)
        return json.dumps(res.json())

    @mcp.resource("findings://{garbage}")
    def findings(garbage: str, context: Context) -> str:
        """Get the current findings available"""
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_findings(headers=headers)
        return json.dumps(res.json())

    @mcp.resource("reports://{garbage}")
    def reports(garbage: str, context: Context) -> str:
        """Get the current reports available"""
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_all_reports(headers=headers)
        return json.dumps(res.json())

    ### Prompts
    # Prompts should match the intended use-cases for the pentesttools MCP
    @mcp.prompt()
    def pentest(target: str) -> str:
        """Create a pentesting report prompt"""
        return f"Please provide a comprehensive pentesting report for this target: {target}, using the tools at your disposal."

    @mcp.prompt()
    def is_up(target: str) -> str:
        """Create a simple target checking prompt"""
        return f"Please see if this target: {target} is up, using the tools at your disposal."

    @mcp.prompt()
    def subdomains(target: str) -> str:
        """Discover all subdomains for the target domain"""
        return (
            f"Please find all the subdomains that you can for this domain: {target}, using the tools at your disposal."
        )

    @mcp.prompt()
    def summary(workspace_id: int) -> str:
        """Create an executive summary for scans in a certain time period from workspace `workspace_id`"""
        return f"Please create an executive summary for the latest scans in the workspace {workspace_id}, using the tools at your disposal."

    @mcp.prompt()
    def vibeseccoding() -> str:
        """Run a vpn website scanner on the target and fix the problems"""

        return """
        You are inside an application. Please run a `pwd` command to check that you have terminal commands enabled.

        Please inspect the docker related files to see how the app is built and run.

        Run the application. In case you don't figure out how, try `docker compose build -d`, or other similar general commands.

        Use `docker network inspect mcp` in order to see what ip docker gave the application.

        Then please run a deep website scan on it and fix the most critical vulnerabilities you find.
        """

    def get_target_name_by_scan_id(scan_id: int, headers) -> str:
        """Get the target name for a given scan ID."""
        scan_res = ptt.api.get_scan_status(scan_id, headers=headers)
        target_id = scan_res.json()["data"]["target_id"]

        target_res = ptt.api.get_target_by_id(target_id, headers=headers)
        return target_res.json()["data"]["name"]


    def get_workspace_id_by_scan_id(scan_id: int, headers) -> int:
        """Get the workspace ID for a given scan ID."""
        scan_res = ptt.api.get_scan_status(scan_id, headers=headers)
        target_id = scan_res.json()["data"]["target_id"]

        target_res = ptt.api.get_target_by_id(target_id, headers=headers)
        return target_res.json()["data"]["workspace_id"]


    def get_findings_from_scans(scans, headers) -> tuple[int, str, list[str], str]:
        finding_ids = []
        target_name = ""
        workspace_id = 0
        # If scans are provided, get findings from each scan
        for scan_id in scans:
            if not target_name:
                try:
                    target_name = get_target_name_by_scan_id(scan_id, headers)
                except Exception:
                    logging.info("Cannot get target from scan_id, skipping this scan")
                    continue

            if not workspace_id:
                try:
                    workspace_id = get_workspace_id_by_scan_id(scan_id, headers)
                except Exception:
                    logging.info("Cannot get the workspace_id from scan_id, skipping")
                    continue

            scan_output = ptt.api.get_output(scan_id=scan_id, headers=headers).json()
            try:
                for finding in scan_output["data"]["output_data"]["findings"]:
                    finding_ids.append(finding["id"])
            except KeyError:
                continue

        if not finding_ids:
            return workspace_id, target_name, finding_ids, "No findings found in the provided scans."
        
        if not target_name:
            return workspace_id, target_name, finding_ids, "No target name could be extracted"

        return workspace_id, target_name, finding_ids, ""

    def get_findings_from_findings(findings, headers):
        if len(findings) == 0:
            return 0, "", [], "No valid finding IDs provided."

        error_message = ""
        target_name = ""
        workspace_id = 0

        first_finding = findings[0]

        finding_res = ptt.api.get_finding_information_by_id(finding_id=first_finding, headers=headers)
        finding_data = finding_res.json().get("data", {})
        if "target_name" in finding_data:
            target_name = finding_data["target_name"]

        target_id = finding_data.get("target_id")
        if target_id:
            try:
                target_res = ptt.api.get_target_by_id(target_id, headers=headers)
                workspace_id = target_res.json()["data"]["workspace_id"]
            except Exception:
                logging.info("Cannot get workspace_id from target_id %s", target_id)

        return workspace_id, target_name, findings, error_message

    def get_findings_from_json_report(report_id, headers):
        finding_ids = []
        error_message = ""
        target_name = ""
        workspace_id = 0

        headers["Accept"] = "application/json"
        report_json = []
        try:
            report_res = ptt.api.download_report(report_id, headers=headers)
            report_json = report_res.json()
        except Exception:
            error_message = "The report could not be downloaded as JSON. Please ensure that the target report is available in JSON format."
        # Extract finding_ids and target_name
        if isinstance(report_json, list) and report_json:
            for item in report_json:
                findings_list = item.get("findings", [])
                finding_ids.extend([f["id"] for f in findings_list if "id" in f])
            if report_json[0].get("findings"):
                target_name = report_json[0]["findings"][0].get("target_name", "")
            else:
                error_message = "Could not extract findings from the report."

        if finding_ids:
            try:
                finding_res = ptt.api.get_finding_information_by_id(finding_id=finding_ids[0], headers=headers)
                target_id = finding_res.json().get("data", {}).get("target_id")
                if target_id:
                    target_res = ptt.api.get_target_by_id(target_id, headers=headers)
                    workspace_id = target_res.json()["data"]["workspace_id"]
            except Exception:
                logging.info("Cannot get workspace_id from report findings")

        return workspace_id, target_name, finding_ids, error_message

    ### Tools
    # These are the actual pentesting tools, with some added management tools, available through `ptt.api`.
    # There are also some meta-tools that the MCP sees as tools as well.
    @mcp.tool()
    async def translate_report(
        report_id: int = None,
        findings: list[int] = None,
        scans: list[int] = None,
        target_language: str = None,
        workspace_id: int = None,  # TODO: We can't get this from the report, but can we get this from the findings themselves?
        context: Context = None,
    ) -> int | str:
        """Translate a report from English to the target language.

        IMPORTANT: This tool can only read and translate reports in JSON format.
        If you want to translate an existing report by report_id, that report must have
        been created with format="json". For PDF, HTML, or other formats, use `scans` or
        `findings` as input instead — the tool will create a new translated JSON report from
        those directly, without needing a pre-existing report.

        Args:
            report_id: The id of a JSON-format report to translate. Only one of `report_id`, `findings`, `scans` should be present.
            findings: The list of finding IDs to translate into a new report. Only one of `report_id`, `findings`, `scans` should be present.
            scans: The list of scan IDs to translate into a new report. Only one of `report_id`, `findings`, `scans` should be present.
            target_language: Target language to translate to, from English.
            workspace_id: The Workspace ID of the new findings that will be imported. The report is not associated with any workspace_id, so we need to pass one explicitly.

        Returns:
            The report_id of the new report.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        # 0. Ensure only one of report_id, findings, or scans is present, and determine which one was chosen
        provided = [x is not None for x in (report_id, findings, scans)]
        if sum(provided) != 1:
            return "You must provide exactly one of: report_id, findings, or scans."

        if report_id is not None:
            chosen_input = "report_id"
        elif findings is not None:
            chosen_input = "findings"
        elif scans is not None:
            chosen_input = "scans"
        else:
            chosen_input = None  # Should not happen due to the check above

        # For report_id, verify it is a JSON report before proceeding
        if chosen_input == "report_id":
            info_res = ptt.api.get_report_information_by_id(report_id, headers=headers)
            report_info = info_res.json().get("data", {})
            report_format = report_info.get("format")
            if report_format != "json":
                return (
                    f"Report {report_id} is in '{report_format}' format. "
                    "translate_report can only read JSON reports. "
                    "Either create the report again with format='json', or pass the "
                    "underlying scan IDs via the `scans` parameter instead."
                )

        # Log the chosen input for debugging
        logging.info(f"Chosen input for translation: {chosen_input}")

        # 1. Get the finding_ids list and target_name based on the chosen input
        finding_ids: list[int] = []
        target_name = ""

        if chosen_input == "report_id":
            workspace_id, target_name, finding_ids, error_message = get_findings_from_json_report(report_id, headers)
        elif chosen_input == "findings":
            workspace_id, target_name, finding_ids, error_message = get_findings_from_findings(findings, headers)
        elif chosen_input == "scans":
            workspace_id, target_name, finding_ids, error_message = get_findings_from_scans(scans, headers)
        else:
            return "Invalid input for translation. Only one of `report_id`, `findings`, `scans` should be present."

        if error_message:
            logging.info(error_message)
            return error_message

        if not finding_ids:
            logging.info("No findings found.")
            return "No findings found."

        logging.info("finding_ids %s", finding_ids)

        # TODO: Translate these fields: risk_description, recommendation, name, vuln_description
        # 3. For each finding, get its details and elicit translation
        translated_findings: list[Finding] = []

        for fid in finding_ids:
            finding_res = ptt.api.get_finding_information_by_id(finding_id=fid, headers=headers)
            finding = finding_res.json().get("data", {})
            logging.info("Translating finding: %s", finding)

            # For now, we only translate the following fields:
            vuln_description = finding.get("vuln_description", "")
            vuln_recommendation = finding.get("recommendation", "")
            test_description = finding.get("test_description", "")
            finding_name = finding.get("name", "")

            # TODO: We have some problems when vuln_evidence should be a table here
            # It's still unclear if the api supports this
            vuln_evidence_obj = finding.get("vuln_evidence", {})
            if isinstance(vuln_evidence_obj, dict) and vuln_evidence_obj.get("type") == "text":
                vuln_evidence = vuln_evidence_obj.get("data", "")
            else:
                vuln_evidence = json.dumps(vuln_evidence_obj)

            logging.info("Will actually translate:")
            logging.info([vuln_description, vuln_recommendation, test_description, finding_name, vuln_evidence])

            # Do MCP Sampling for the translation for each field
            translated = {}
            for field, value in [
                ("vuln_description", vuln_description),
                ("vuln_recommendation", vuln_recommendation),
                ("test_description", test_description),
                ("finding_name", finding_name),
                ("vuln_evidence", vuln_evidence),
            ]:
                prompt = f"Translate the following {field} from English to {target_language}. Do not translate URLs, API endpoints, or vulnerability names unless appropriate.\n\n{value}"

                if value:
                    sampling_result = await context.session.create_message(
                        messages=[
                            SamplingMessage(
                                role="user",
                                content=TextContent(type="text", text=prompt),
                            )
                        ],
                        max_tokens=200,
                    )

                    if sampling_result.content.type == "text":
                        translated[field] = sampling_result.content.text
                    else:
                        translated[field] = value
                else:
                    translated[field] = value

            logging.info("Translated: %s", translated)

            # Prepare the finding to import
            finding_to_import = Finding(
                target_name=target_name,
                source=2,
                name=translated["finding_name"],
                test_description=translated["test_description"],
                vuln_description=translated["vuln_description"],
                vuln_risk_level=finding.get("vuln_risk_level"),
                vuln_recommendation=translated["vuln_recommendation"],
                cvss_v3=finding.get("cvss_v3"),
                vuln_evidence=translated["vuln_evidence"],
            )
            translated_findings.append(finding_to_import)

        # Import the translated finding (as a list)
        import_res = ptt.api.import_findings(
            findings=translated_findings,
            workspace_id=workspace_id,
            headers=headers,
        )
        import_data = import_res.json()

        logging.info("Imported data: %s", import_data)

        translated_finding_ids = []
        for item in import_data.get("data", {}).get("results", []):
            new_finding_id = item.get("id")
            if new_finding_id:
                translated_finding_ids.append(new_finding_id)

        if not translated_finding_ids:
            return "No translated findings were imported."

        logging.info(translated_finding_ids)

        # 4. Create a new report with the translated findings
        create_report_res = ptt.api.create_report(
            format="pdf",
            group_by="target",
            source="findings",
            resources=translated_finding_ids,
            headers=headers,
        )
        create_report_data = create_report_res.json()

        logging.info("Created report: %s", create_report_data)
        new_report_id = create_report_data.get("data", {}).get("report_id")
        if not new_report_id:
            return "Failed to create the translated report."

        return f"Download your new report from https://app.pentest-tools.com/reports/{new_report_id}"

    # Pentesting tools
    @mcp.tool()
    async def run_website_scanner(
        target: str,
        scan_type: Literal["light", "deep", "custom"] = "light",
        attack_active: "AttackActive" = None,
        attack_passive: "AttackPassive" = None,
        discovery: "Discovery" = None,
        spider: "Spider" = None,
        requests_per_second: int = 100,
        workspace_id: int = None,
        vpn_profile_uuid: str = None,
        auth_cookies: str = None,
        auth_headers: list[str] = None,
        auth_login_form_url: str = None,
        auth_login_form_username: str = None,
        auth_login_form_password: str = None,
        max_scan_time: int = 5,
        context: Context = None,
    ) -> str:
        """
        Run a website scanner on the target URL, reporting progress/status updates.

        Args:
            target: The target URL to scan.
            scan_type: The type of scan to perform. Possible values: "light", "deep", "custom". Default is "deep".
            attack_active: List of active attack modules to run.
            attack_passive: List of passive attack modules to run.
            discovery: List of discovery modules to run.
            spider: Spidering options as a dict.
            requests_per_second: Max requests per second (1-10000). Default is 100.
            workspace_id: The Workspace ID for the workspace the scan is going to be running in.
            vpn_profile_uuid: The VPN profile ID for the scan.
            auth_cookies: Cookie header string for authentication. Example: "PHPSESSID=abc123; token=xyz".
            auth_headers: List of HTTP headers for authentication. Example: ["Authorization: Bearer token123", "Cookie: session=abc"].
            auth_login_form_url: URL of the login page for form-based authentication.
            auth_login_form_username: Username for login form authentication.
            auth_login_form_password: Password for login form authentication.
            max_scan_time: The maximum duration in minutes for the scan. Between 5 and 1440 minutes.
        """
        tool_params = {
            "scan_type": scan_type,
            "requests_per_second": requests_per_second,
        }
        if attack_active:
            tool_params["attack_active"] = attack_active
        if attack_passive:
            tool_params["attack_passive"] = attack_passive
        if discovery:
            tool_params["discovery"] = discovery
        if spider:
            tool_params["spider"] = spider

        scan_params: dict[str, Any] = {"max_scan_time": str(max_scan_time)}
        if vpn_profile_uuid:
            scan_params["vpn_profile_uuid"] = vpn_profile_uuid

        if workspace_id:
            scan_params["workspace_id"] = workspace_id

        if auth_cookies:
            scan_params["authentication"] = {
                "type": "cookies",
                "params": {"cookie_header": auth_cookies},
            }
        elif auth_headers:
            scan_params["authentication"] = {
                "type": "headers",
                "params": {"headers": auth_headers},
            }
        elif auth_login_form_url and auth_login_form_username and auth_login_form_password:
            scan_params["authentication"] = {
                "type": "login_form",
                "params": {
                    "login_page_url": auth_login_form_url,
                    "username": auth_login_form_username,
                    "password": auth_login_form_password,
                },
            }
        elif any(
            (
                bool(auth_login_form_url),
                bool(auth_login_form_username),
                bool(auth_login_form_password),
            )
        ):
            return "You passed only one of `auth_login_form_url`, `auth_login_form_username`, `auth_login_form_password`, but you must pass all of them."

        return await run_tool_long_running(
            ptt.api.Tool.WEBSITE_SCANNER,
            target,
            tool_params,
            scan_params,
            context,
        )

    @mcp.tool()
    async def run_subdomain_finder(
        target: str,
        scan_type: Literal["light", "deep", "custom"] = "light",
        web_details: bool = None,
        whois: bool = None,
        unresolved_results: bool = None,
        search_methods: "FindSubdomainsSearchMethods" = None,
        max_scan_time: int = 5,
        context: Context = None,
    ) -> str:
        """
        Discover all subdomains for the target domain. Use to enumerate subdomains for further scanning.
        """
        tool_params = {
            "scan_type": scan_type,
        }

        if web_details is not None:
            tool_params["web_details"] = web_details
        if whois is not None:
            tool_params["whois"] = whois
        if unresolved_results is not None:
            tool_params["unresolved_results"] = unresolved_results
        if search_methods is not None:
            tool_params["search_methods"] = search_methods

        scan_params = {"max_scan_time": str(max_scan_time)}

        return await run_tool_long_running(
            ptt.api.Tool.SUBDOMAIN_FINDER,
            target,
            tool_params,
            scan_params,
            context,
        )

    @mcp.tool()
    async def run_network_scanner(
        target: str,
        preset: Literal["light", "deep", "custom"] = "light",
        scanning_engines: "NetworkScannerScanningEngines" = None,
        engines_options: "NetworkScannerEngineOptions" = None,
        check_alive: bool = None,  # True
        extensive_modules: bool = None,  # False
        protocol_type: Literal["tcp", "udp"] = "tcp",
        port_range: "NetworkScannerPortRange" = None,
        port_top: Literal["OpenVAS default", "10", "100", "1000", "5000"] = None,
        port_list: "NetworkScannerPortList" = None,
        just_scan_ports: bool = None,
        workspace_id: int = None,
        max_scan_time: int = 5,
        context: Context = None,
    ) -> str:
        """Run a network scanner on the target hostname.

        `just_scan_ports` will only scan the selected ports to see which are open, nothing heavy. It's also called "`port_scanner`".

        You can set port top only in custom mode.
        The scanning engines field is required when preset is custom.
        """

        tool_params = {"scan_type": preset}

        if just_scan_ports:
            # Set tool params
            if port_range:
                tool_params["port_range"] = port_range
            if port_top:
                tool_params["port_top"] = port_top
            if port_list:
                tool_params["port_list"] = port_list

            # Set scan params
            scan_params: dict[str, str | int] = {"max_scan_time": str(max_scan_time)}
            if workspace_id is not None:
                scan_params["workspace_id"] = workspace_id

            return await run_tool_long_running(
                ptt.api.Tool.PORT_SCANNER,
                target,
                tool_params,
                scan_params,
                context,
            )

        # Set tool params
        if check_alive is not None:
            tool_params["check_alive"] = check_alive
        if extensive_modules is not None:
            tool_params["extensive_modules"] = extensive_modules
        if protocol_type is not None:
            tool_params["protocol_type"] = protocol_type
        if scanning_engines:
            tool_params["scanning_engines"] = scanning_engines
        if engines_options:
            tool_params["engines_options"] = engines_options
        if port_range:
            tool_params["port_range"] = port_range
        if port_top:
            tool_params["port_top"] = port_top
        if port_list:
            tool_params["port_list"] = port_list

        # Set scan params
        scan_params = {"max_scan_time": str(max_scan_time)}
        if workspace_id is not None:
            scan_params["workspace_id"] = workspace_id

        return await run_tool_long_running(
            ptt.api.Tool.NETWORK_SCANNER,
            target,
            tool_params,
            scan_params,
            context,
        )

    # Administrative tools
    @mcp.tool()
    async def get_scans(
        context: Context,
        workspace_id: int = None,
        target_id: int = None,
        status: Literal[
            "running",
            "finished",
            "failed to start",
            "stopped",
            "timed out",
            "waiting",
            "aborted",
            "VPN connection error",
            "auth failed",
            "connection error",
        ] = None,
        limit: int = None,
        page: int = None,
        start_time: dict = None,
    ) -> str:
        """
        Retrieve scans from Pentest-Tools.com with optional filtering.

        Usage:
        - To get all scans: call with no extra parameters.
        - To filter by workspace: set workspace_id=<workspace id>.
        - To filter by target: set target_id=<target id>.
        - To filter by scan status: set status to one of the allowed values.
        - To limit results: set limit=<max number of scans>.
        - To paginate: set page=<page number>.
        - To filter by start time: set start_time as a dictionary (e.g. {"from": "...", "to": "..."}).

        Args:
            context: The MCP context object.
            workspace_id: (optional) ID of the workspace to show scans from.
            target_id: (optional) ID of the target to show scans from.
            status: (optional) Status of the scans to be shown. Allowed values: running, finished, failed to start, stopped, timed out, waiting, aborted, VPN connection error, auth failed, connection error.
            limit: (optional) Maximum number of scans to return.
            page: (optional) Page number to return.
            start_time: (optional) Dictionary specifying start time interval for scans (e.g. {"from": "...", "to": "..."}).

        Returns:
            JSON string containing the list of scans matching the filters.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"
        res = ptt.api.get_scans(
            workspace_id=workspace_id,
            target_id=target_id,
            status=status,
            limit=limit,
            page=page,
            start_time=start_time,
            headers=headers,
        )
        return json.dumps(res.json())

    @mcp.tool()
    async def scan(
        verb: Literal[
            "get_status",
            "get_raw_output",
            "get_output",
            "start_by_targetid",
            "stop",
            "delete",
            "get_scan_output_by_id",
        ],
        context: Context,
        scan_id: int = None,
        target_id: int = None,
        tool_id: int = None,
        tool_params: dict = None,
    ) -> str:
        """
        Manage scans in Pentest-Tools.com.

        Usage:
        - To get the status of a scan: verb="get_status", scan_id=<scan id>
        - To get the raw output of a scan: verb="get_raw_output", scan_id=<scan id>
        - To get the parsed output of a scan: verb="get_output", scan_id=<scan id>
        - To start a scan by target id: verb="start_by_targetid", target_id=<target id>, tool_id=<tool id>, tool_params=<tool params>
        - To stop a scan: verb="stop", scan_id=<scan id>
        - To delete a scan: verb="delete", scan_id=<scan id>
        - To get a scan output of a scan: verb="get_scan_output_by_id", scan_id=<scan id>

        Only set the parameters required for the verb you use.

        Args:
            verb: The action to perform. Must be one of: get_status, get_raw_output, get_output, start_by_targetid, stop, delete.
            scan_id: The scan ID (required for get_status, get_raw_output, get_output, stop, delete).
            target_id: The target ID (required for start_by_targetid).
            tool_id: The tool ID (required for start_by_targetid).
            tool_params: The tool parameters (required for start_by_targetid).
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        if verb == "get_status":
            if not scan_id:
                return "Missing scan_id for get_status."
            res = ptt.api.get_scan_status(scan_id, headers=headers)
            return json.dumps(res.json())
        elif verb == "get_raw_output":
            if not scan_id:
                return "Missing scan_id for get_raw_output."
            res = ptt.api.get_raw_output(scan_id, headers=headers)
            return res.text
        elif verb == "get_output":
            if not scan_id:
                return "Missing scan_id for get_output."
            res = ptt.api.get_output(scan_id, headers=headers)
            return json.dumps(res.json())
        elif verb == "start_by_targetid":
            if not target_id or not tool_id or tool_params is None:
                return "Missing target_id, tool_id, or tool_params for start_by_targetid."
            res = ptt.api.start_scan_by_targetid(target_id, tool_id, tool_params, headers=headers)
            return json.dumps(res.json())
        elif verb == "stop":
            if not scan_id:
                return "Missing scan_id for stop."
            res = ptt.api.stop_scan(scan_id, headers=headers)
            return json.dumps(res.json())
        elif verb == "delete":
            if not scan_id:
                return "Missing scan_id for delete."

            # Elicit confirmation from the user before deleting
            if context is not None:
                result = await context.elicit(
                    message="Deleting a scan is an irreversible operation. Respond to this message in order to proceed",
                    schema=ConfirmDeleteSchema,
                )
                # Only proceed if user typed 'yes'
                if result.action != "accept":
                    return "Scan deletion cancelled by user."

            res = ptt.api.delete_scan(scan_id, headers=headers)
            return json.dumps(res.json())
        elif verb == "get_scan_output_by_id":
            if not scan_id:
                return "Missing scan_id for get_scan_output_by_id."
            res = ptt.api.get_output(scan_id, headers=headers)
            return json.dumps(res.json())
        else:
            return "Unknown verb: {verb}. Use one of: get_status, get_raw_output, get_output, start_by_targetid, stop, delete."

    # VPNs
    @mcp.tool()
    async def get_vpn_profiles(
        context: Context,
    ) -> str:
        """
        Use this tool to get a list of vpn profiles from Pentest-Tools.com.

        Returns:
            JSON string containing the list of targets matching the filters.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.get_vpn_profiles(
            headers=headers,
        )
        return json.dumps(res.json())

    ### Targets
    @mcp.tool()
    async def get_targets(
        context: Context,
        workspace_id: int = None,
        target_type: Literal["hostname", "ip_address", "url"] = None,
        limit: int = None,
        page: int = None,
    ) -> str:
        """
        Use this tool to get a list of targets from Pentest-Tools.com.

        Args:
            context: The MCP context object.
            workspace_id: (optional) ID of the workspace to show targets from.
            target_type: (optional) The type of target. Possible values: "hostname", "ip_address", "url".
            limit: (optional) The maximum number of targets to return.
            page: (optional) The page number to return.

        Returns:
            JSON string containing the list of targets matching the filters.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"
        res = ptt.api.get_targets(
            workspace_id=workspace_id,
            target_type=target_type,
            limit=limit,
            page=page,
            headers=headers,
        )
        return json.dumps(res.json())

    @mcp.tool()
    async def target(
        verb: Literal["add", "delete"],
        context: Context,
        name: str = None,
        description: str = "",
        workspace_id: int = None,
        target_id: int = None,
    ) -> str:
        """
        Manage targets in Pentest-Tools.com.

        Usage:
        - To add a new target: verb="add", name=<target name>, description=<optional>, workspace_id=<optional>.
        - To delete a target: verb="delete", target_id=<target id>.

        Only set the parameters required for the verb you use.

        Args:
            verb: The action to perform. Must be one of: add, delete.
            context: The MCP context object.
            name: The name of the target to add (required for add).
            description: Description for the target (optional for add).
            workspace_id: Workspace ID to associate with the target (optional for add).
            target_id: The ID of the target to delete (required for delete).
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        if verb == "add":
            if not name:
                return "Missing target name for add."
            res = ptt.api.add_target(name, description, workspace_id, headers=headers)
            return json.dumps(res.json())
        elif verb == "delete":
            if not target_id:
                return "Missing target_id for delete."
            res = ptt.api.delete_target(target_id, headers=headers)
            return json.dumps(res.json())
        else:
            return f"Unknown verb: {verb}. Use one of: add, delete."

    ### Reports
    @mcp.tool()
    async def create_report(
        format: Literal["pdf", "html", "json", "csv", "xlsx", "docx"],
        group_by: Literal["target", "vulnerability"],
        source: Literal["scans", "findings"],
        resources: list[int],
        context: Context,
        filters: dict = None,
        webhook_url: str = None,
    ) -> str:
        """Create a pentesting report, based on the findings already present in the platform.

        Args:
            format: The format in which to create the report. Possible values: pdf, html, json, csv, xlsx, docx.
            group_by: What key is used to create the report. Possible values: target, vulnerability.
            source: This is the source that is used to create the report. Possible values: scans, findings.
            resources: A list of ids, depeding on which source was chosen. If `scans` was chosed, then it's a list of scan_ids. If `findings` was chosed, then it's a list of finding_ids.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"
        res = ptt.api.create_report(format, group_by, source, resources, filters, webhook_url, headers=headers)
        return json.dumps(res.json())

    @mcp.tool()
    async def report(
        verb: Literal["delete", "download", "get_all", "get_report_information_by_id"],
        context: Context,
        report_id: int = None,
    ) -> str:
        """
        Manage reports in Pentest-Tools.com.

        Usage:
        - To delete a report: verb="delete", report_id=<report id>
        - To download a report: verb="download", report_id=<report id>
        - To get all reports: verb="get_all"
        - To get report information: verb="get_report_information_by_id", report_id=<report id>

        Only set the parameters required for the verb you use.

        Args:
            verb: The action to perform. Must be one of: delete, download, get_all.
            context: The MCP context object.
            report_id: The ID of the report (required for delete and download).
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        if verb == "delete":
            if not report_id:
                return "Missing report_id for delete."
            res = ptt.api.delete_report(report_id, headers=headers)
            return json.dumps(res.json())
        elif verb == "download":
            if not report_id:
                return "Missing report_id for download."
            res = ptt.api.download_report(report_id, headers=headers)
            try:
                return res.json()
            except Exception:
                return res.text

        elif verb == "get_all":
            res = ptt.api.get_all_reports(headers=headers)
            return json.dumps(res.json())
        elif verb == "get_report_information_by_id":
            if not report_id:
                return "Missing report_id for get_report_information_by_id."
            res = ptt.api.get_report_information_by_id(report_id, headers=headers)
            return json.dumps(res.json())
        else:
            return f"Unknown verb: {verb}. Use one of: delete, download, get_all."

    ### Loggers
    @mcp.tool()
    async def http_logger(
        verb: Literal["create", "delete", "clear"],
        context: Context,
        label: str = None,
        logger_id: int = None,
        workspace_id: int = None,
    ) -> str:
        """
        Manage HTTP request loggers.

        Usage:
        - To create a new HTTP logger: call with verb="create", label="logger label", workspace_id=<optional workspace id>.
        - To delete an HTTP logger: call with verb="delete", logger_id=<logger id>.
        - To clear all data for a logger: call with verb="clear", logger_id=<logger id>.

        Only set the parameters required for the verb you use.

        Args:
            verb: The action to perform. Must be one of: create, delete, clear.
            label: The label for the logger (required for create).
            logger_id: The ID of the logger (required for delete and clear).
            workspace_id: The workspace to associate with the logger (optional for create).
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        if verb == "create":
            if not label:
                return "Missing label for logger creation."
            res = ptt.api.create_http_logger(label, workspace_id, headers=headers)
        elif verb == "delete":
            if not logger_id:
                return "Missing logger_id for deletion."
            res = ptt.api.delete_logger(logger_id, headers=headers)
        elif verb == "clear":
            if not logger_id:
                return "Missing logger_id for clearing data."
            res = ptt.api.clear_data_for_logger(logger_id, headers=headers)
        else:
            return f"Unknown verb: {verb}. Use one of: create, delete, clear."
        return json.dumps(res.json())

    ### Workspaces
    @mcp.tool()
    async def workspace(
        verb: Literal["create", "delete", "edit", "list", "get"],
        context: Context,
        workspace_id: int = None,
        name: str = None,
        description: str = None,
    ) -> str:
        """
        Everything you need to manage workspaces.

        Usage:
        - To create a workspace: call with verb="create", name="workspace name", description="optional description".
        - To delete a workspace: call with verb="delete", workspace_id=<id>. You will need to get the id beforehand.
        - To edit a workspace: call with verb="edit", workspace_id=<id>, name="new name" (optional), description="new description" (optional).
        - To get all workspaces: call with verb="list".
        - To get a workspace by id: call with verb="get", workspace_id=<id>.

        Only set the parameters required for the verb you use.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        if verb == "create":
            if not name:
                return "Missing workspace name for creation."
            res = ptt.api.create_workspace(name, description or "", headers=headers)
        elif verb == "delete":
            if not workspace_id:
                return "Missing workspace_id for deletion."
            res = ptt.api.delete_workspace(workspace_id, headers=headers)
        elif verb == "edit":
            if not workspace_id:
                return "Missing workspace_id for editing."
            res = ptt.api.edit_workspace(workspace_id, name, description, headers=headers)
        elif verb == "list":
            res = ptt.api.get_workspaces(headers=headers)
        elif verb == "get":
            if not workspace_id:
                return "Missing workspace_id for get."
            res = ptt.api.get_workspace_by_id(workspace_id, headers=headers)
        else:
            return f"Unknown verb: {verb}. Use one of: create, delete, edit, list, get."
        return json.dumps(res.json())

    ### Findings
    @mcp.tool()
    async def findings(
        verb: Literal["get_finding_information_by_id", "get"],
        context: Context,
        finding_id: int = None,
        workspace_id: int = None,
        target_id: int = None,
        task_id: int = None,
        group_duplicates: bool = False,
    ) -> str:
        """
        Manage findings in Pentest-Tools.com.

        Usage:
        - To get all findings: call with verb="get".
        - To get finding information by id: call with verb="get_finding_information_by_id"

        Only set the parameters required for the verb you use.

        Args:
            verb: The action to perform. Must be one of: get_finding_information_by_id, get.
            finding_id: The id of the finding you want more information for.
            workspace_id: (optional) ID of the workspace to show findings from.
            target_id: (optional) ID of the target to show findings from.
            task_id: (optional) ID of the scan to show findings from.
            group_duplicates: (optional) Group duplicates to show only the latest finding from each group. Default is False.
            context: The MCP context object.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        if verb == "get_finding_information_by_id":
            if not finding_id:
                return "Missing finding_id for get_finding_information_by_id."
            res = ptt.api.get_finding_information_by_id(finding_id=finding_id, headers=headers)
            return json.dumps(res.json())
        elif verb == "get":
            res = ptt.api.get_findings(
                workspace_id=workspace_id,
                target_id=target_id,
                task_id=task_id,
                group_duplicates=group_duplicates,
                headers=headers,
            )
            return json.dumps(res.json())
        else:
            return f"Unknown verb: {verb}. Use one of: get_finding_information_by_id, get."

    @mcp.tool()
    def import_findings(
        findings: list[Finding],
        context: Context,
    ):
        """Import multiple findings into Pentest-Tools.com.

        This is the only way to create new findings manually. Use this when you need to add findings that do not come from scans.

        Args:
            findings: A list of findings.
        """
        headers = get_authentication_headers(context)
        if not headers:
            return "The api key was not set"

        res = ptt.api.import_findings(
            findings=findings,
            headers=headers,
        )
        return json.dumps(res.json())


def entrypoint(mcp_method: str = "stdio"):
    """Entrypoint for ptt-mcp"""
    global MCP_METHOD

    logging.basicConfig(
        level=logging.DEBUG,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    if FastMCP is None:
        logging.error(
            "PentestTools has been installed without MCP support. Install with pip install PentestTools[mcp] if you want MCP support."
        )
        return

    logging.info("Starting the PentestTools MCP server")
    if mcp_method == "stdio":
        logging.info("Using the `stdio` mcp type")
        MCP_METHOD = "stdio"

        # Get the api key for stdio
        if not getattr(ptt.api, "api_key", False):
            api_key = os.getenv("PTT_API_KEY", "")
            if not api_key:
                sys.exit(
                    "API key not found. Please provide it using the --key flag or set the PTT_API_KEY environment variable."
                )

            ptt.api.api_key = api_key

        ptt.api.init()  # This call sets the given key

        return mcp.run(transport="stdio")
    elif mcp_method == "http":
        logging.info("Using the `http` mcp type")
        MCP_METHOD = "http"
        return mcp.run(transport="streamable-http")
