from fastmcp import FastMCP

from tool import *

# 使用绝对路径，确保无论从哪个目录运行都正确
PROJECT_ROOT = Path(__file__).parent.parent
# 确保 logs 目录存在
os.makedirs(PROJECT_ROOT / "logs", exist_ok=True)
logger.add(str(PROJECT_ROOT / "logs" / "dirmcp_mcp.log"), rotation="00:00", enqueue=True)
logger.info("日志初始化完成")

app = FastMCP(name="dirmcp",
              instructions="""
              调用dirsearch，扫描指定网址的目录
              """
              )

@app.tool()
def start_scan(
        url: str,
        cookies: str = None,
        extensions: list[str] | None = None,
        wordlist: str | None = None,
        threads: int | None = None,
        timeout: int | None = None,
        exclude_codes: list[str] | None = None,
        include_codes: list[str] | None = None,
        recursive: bool = False,
        recursion_depth: int | None = None,
) -> dict[str, Any]:
    """
    使用dirsearch扫描目标URL的目录和文件。

    args:
        url（必填）: 要扫描的目标 URL（例如：http://example.com）
        cookies: 指定Cookie
        extensions: 要扫描的文件扩展名（例如：["php", "html"]）
        wordlist: 自定义字典文件路径
        threads: 并发线程数（默认：25）
        timeout: 最大运行时间（秒）
        exclude_codes: 要排除的 HTTP 状态码（例如：["404", "403"]）
        include_codes: 要包含的 HTTP 状态码（例如：["200", "301"]）
        recursive: 启用递归扫描
        recursion_depth: 递归扫描的最大深度
    return:
        success：是否成功
        error：具体错误
        results：扫描结果
    """
    logger.info(f"url: {url}")
    if not url:
        return {"success": False, "error": "URL is required"}

    params = {
        "url": url,
        "cookies": cookies,
        "extensions": extensions,
        "wordlist": wordlist,
        "threads": threads,
        "timeout": timeout,
        "exclude_codes": exclude_codes,
        "include_codes": include_codes,
        "recursive": recursive,
        "recursion_depth": recursion_depth,
    }

    result = run_dirsearch(**params)

    if result["success"]:
        logger.info(f"scan completed, found {len(result.get('results', {}).get('results', []))} paths")
        return {
            "success": True,
            "results": result["results"]
        }
    else:
        logger.error(f"scan failed: {result.get('error', 'unknown error')}")
        return {
            "success": False,
            "error": result.get("error", "unknown error")
        }

def main():
    app.run()
    logger.info("日志结束")

if __name__ == "__main__":
    app.run()
    logger.info("日志结束")
