Update app.py

app.py

@@ -1,28 +1,38 @@
-"""High-Performance Stealth Web Scraper & Browser Automation MCP Server
-
-All-in-one file designed for Hugging Face Spaces (Port 7860).
+"""
+High-Performance Stealth Web Scraper & Browser Automation MCP Server
+Configured for Hugging Face Spaces with Cloud Host-Header validation bypass.
 """
 
-import asyncio
-import base64
-import json
-import logging
 import os
 import sys
-from contextlib import asynccontextmanager
+import json
+import asyncio
+import logging
+import base64
 from typing import Any, AsyncGenerator
+from contextlib import asynccontextmanager
 
 from bs4 import BeautifulSoup
-from duckduckgo_search import DDGS
 import httpx
+import trafilatura
 from lxml_html_clean import Cleaner
 import markdownify
+from duckduckgo_search import DDGS
+from patchright.async_api import async_playwright, Browser, BrowserContext, Page
+
 from mcp.server.fastmcp import FastMCP
-from patchright.async_api import Browser, BrowserContext, Page, async_playwright
+try:
+    from mcp.server.transport_security import TransportSecuritySettings
+    security_settings = TransportSecuritySettings(
+        enable_dns_rebinding_protection=False,
+        allowed_hosts=["*"]
+    )
+except ImportError:
+    security_settings = None
+
 from starlette.applications import Starlette
+from starlette.routing import Route, Mount
 from starlette.responses import JSONResponse
-from starlette.routing import Mount, Route
-import trafilatura
 import uvicorn
 
 # ==========================================
@@ -31,7 +41,7 @@ import uvicorn
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-    handlers=[logging.StreamHandler(sys.stdout)],
+    handlers=[logging.StreamHandler(sys.stdout)]
 )
 logger = logging.getLogger("mcp_server")
 
@@ -46,10 +56,7 @@ BROWSER_TIMEOUT_MS = 30000
 # ==========================================
 # 2. Browser Pool
 # ==========================================
-
-
 class StealthBrowserPool:
-
     def __init__(self):
         self._playwright = None
         self._browser: Browser | None = None
@@ -62,7 +69,7 @@ class StealthBrowserPool:
             if not self._playwright:
                 self._playwright = await async_playwright().start()
             if not self._browser or not self._browser.is_connected():
-                logger.info("Starting Chromium Stealth instance...")
+                logger.info("Initializing Chromium Stealth instance...")
                 self._browser = await self._playwright.chromium.launch(
                     headless=True,
                     args=[
@@ -72,8 +79,8 @@ class StealthBrowserPool:
                         "--disable-gpu",
                         "--disable-blink-features=AutomationControlled",
                         "--no-first-run",
-                        "--window-size=1920,1080",
-                    ],
+                        "--window-size=1920,1080"
+                    ]
                 )
                 self._request_counter = 0
             return self._browser
@@ -104,7 +111,7 @@ class StealthBrowserPool:
                     "Chrome/131.0.0.0 Safari/537.36"
                 ),
                 locale="en-US",
-                timezone_id="America/New_York",
+                timezone_id="America/New_York"
             )
             await context.add_init_script("""
                 Object.defineProperty(navigator, 'webdriver', { get: () => undefined });
@@ -130,35 +137,25 @@ class StealthBrowserPool:
             self._semaphore.release()
             await self._check_recycle()
 
-
 browser_pool = StealthBrowserPool()
 
 # ==========================================
 # 3. Content Extraction Pipeline
 # ==========================================
-
-
 def clean_html(raw_html: str) -> str:
     cleaner = Cleaner(
-        scripts=True,
-        javascript=True,
-        comments=True,
-        style=True,
-        links=False,
-        meta=False,
-        page_structure=False,
-        safe_attrs_only=False,
+        scripts=True, javascript=True, comments=True, style=True,
+        links=False, meta=False, page_structure=False, safe_attrs_only=False
     )
     return cleaner.clean_html(raw_html)
 
-
 def extract_content(html: str, url: str = "") -> dict[str, Any]:
     extracted = trafilatura.extract(
         html,
         url=url,
         include_links=True,
         include_images=True,
-        output_format="markdown",
+        output_format="markdown"
     )
     metadata = trafilatura.extract_metadata(html) or {}
     meta_dict = {
@@ -171,25 +168,17 @@ def extract_content(html: str, url: str = "") -> dict[str, Any]:
     if not extracted or len(extracted.strip()) < 100:
         cleaned = clean_html(html)
         soup = BeautifulSoup(cleaned, "html.parser")
-        for tag in soup(
-            ["nav", "footer", "aside", "header", "script", "style", "noscript"]
-        ):
+        for tag in soup(["nav", "footer", "aside", "header", "script", "style", "noscript"]):
             tag.decompose()
         if not meta_dict["title"] and soup.title:
             meta_dict["title"] = soup.title.string or ""
-        extracted = markdownify.markdownify(
-            str(soup), heading_style="ATX", strip=["svg"]
-        )
+        extracted = markdownify.markdownify(str(soup), heading_style="ATX", strip=["svg"])
 
     if len(extracted) > MAX_CONTENT_CHARS:
-        extracted = (
-            extracted[:MAX_CONTENT_CHARS]
-            + f"\n\n... [Truncated: reached {MAX_CONTENT_CHARS} limit]"
-        )
+        extracted = extracted[:MAX_CONTENT_CHARS] + f"\n\n... [Truncated: reached {MAX_CONTENT_CHARS} limit]"
 
     return {"content": extracted.strip(), "metadata": meta_dict}
 
-
 def extract_json_ld(html: str) -> list[dict]:
     soup = BeautifulSoup(html, "html.parser")
     schemas = []
@@ -202,27 +191,19 @@ def extract_json_ld(html: str) -> list[dict]:
             continue
     return schemas
 
-
 # ==========================================
 # 4. Hybrid Scraping Core
 # ==========================================
-
 FAST_HEADERS = {
     "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
     "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
     "Accept-Language": "en-US,en;q=0.9",
 }
 
-
 async def dismiss_popups(page: Page):
     selectors = [
-        'button[id*="cookie" i]',
-        'button[class*="cookie" i]',
-        'button[aria-label*="accept" i]',
-        'button:has-text("Accept all")',
-        'button:has-text("Accept")',
-        'button:has-text("I agree")',
-        'button:has-text("Got it")',
+        'button[id*="cookie" i]', 'button[class*="cookie" i]', 'button[aria-label*="accept" i]',
+        'button:has-text("Accept all")', 'button:has-text("Accept")', 'button:has-text("I agree")', 'button:has-text("Got it")'
     ]
     for sel in selectors:
         try:
@@ -234,32 +215,14 @@ async def dismiss_popups(page: Page):
         except Exception:
             continue
 
-
-async def scrape_pipeline(
-    url: str,
-    force_browser: bool = False,
-    auto_scroll: bool = True,
-    wait_for_selector: str | None = None,
-) -> dict:
+async def scrape_pipeline(url: str, force_browser: bool = False, auto_scroll: bool = True, wait_for_selector: str | None = None) -> dict:
     if not force_browser:
         try:
-            async with httpx.AsyncClient(
-                http2=True,
-                timeout=HTTP_TIMEOUT_SEC,
-                follow_redirects=True,
-                headers=FAST_HEADERS,
-            ) as client:
+            async with httpx.AsyncClient(http2=True, timeout=HTTP_TIMEOUT_SEC, follow_redirects=True, headers=FAST_HEADERS) as client:
                 resp = await client.get(url)
                 if resp.status_code == 200:
                     html = resp.text
-                    if not any(
-                        s in html.lower()
-                        for s in [
-                            "cf-challenge",
-                            "ray-id",
-                            "just a moment...",
-                        ]
-                    ):
+                    if not any(s in html.lower() for s in ["cf-challenge", "ray-id", "just a moment..."]):
                         parsed = extract_content(html, url=url)
                         if len(parsed["content"]) >= 150:
                             return {
@@ -267,7 +230,7 @@ async def scrape_pipeline(
                                 "status": 200,
                                 "engine": "fast-http",
                                 "content": parsed["content"],
-                                "metadata": parsed["metadata"],
+                                "metadata": parsed["metadata"]
                             }
         except Exception:
             pass
@@ -275,16 +238,11 @@ async def scrape_pipeline(
     async with browser_pool.get_page() as page:
         await page.goto(url, wait_until="domcontentloaded")
         try:
-            cf_frame = await page.query_selector(
-                "iframe[src*='challenges.cloudflare.com']"
-            )
+            cf_frame = await page.query_selector("iframe[src*='challenges.cloudflare.com']")
             if cf_frame:
                 box = await cf_frame.bounding_box()
                 if box:
-                    await page.mouse.click(
-                        box["x"] + box["width"] / 2,
-                        box["y"] + box["height"] / 2,
-                    )
+                    await page.mouse.click(box["x"] + box["width"] / 2, box["y"] + box["height"] / 2)
                     await asyncio.sleep(2.0)
         except Exception:
             pass
@@ -293,9 +251,7 @@ async def scrape_pipeline(
 
         if wait_for_selector:
             try:
-                await page.wait_for_selector(
-                    wait_for_selector, state="visible", timeout=8000
-                )
+                await page.wait_for_selector(wait_for_selector, state="visible", timeout=8000)
             except Exception:
                 pass
 
@@ -318,32 +274,30 @@ async def scrape_pipeline(
         "status": 200,
         "engine": "patchright-stealth",
         "content": parsed["content"],
-        "metadata": parsed["metadata"],
+        "metadata": parsed["metadata"]
     }
 
-
 # ==========================================
 # 5. MCP Server & Tool Definitions
 # ==========================================
-mcp = FastMCP("WebScraper")
-
-
-@mcp.tool()
-async def scrape_url(
-    url: str,
-    force_browser: bool = False,
-    auto_scroll: bool = True,
-    wait_for_selector: str | None = None,
-) -> dict:
-    """Scrapes cleaned Markdown content and metadata from any URL."""
-    return await scrape_pipeline(
-        url, force_browser, auto_scroll, wait_for_selector
-    )
-
+# Initialize FastMCP with transport_security disabled for public cloud proxying
+if security_settings:
+    mcp = FastMCP("WebScraper", host=HOST, transport_security=security_settings)
+else:
+    mcp = FastMCP("WebScraper", host=HOST)
+
+@mcp.tool(
+    name="scrape_url",
+    description="Scrapes cleaned Markdown content and metadata from any URL."
+)
+async def scrape_url(url: str, force_browser: bool = False, auto_scroll: bool = True, wait_for_selector: str | None = None) -> dict:
+    return await scrape_pipeline(url, force_browser, auto_scroll, wait_for_selector)
 
-@mcp.tool()
+@mcp.tool(
+    name="search_and_scrape",
+    description="Searches DuckDuckGo and concurrently extracts content from top results."
+)
 async def search_and_scrape(query: str, max_results: int = 3) -> dict:
-    """Searches DuckDuckGo and concurrently extracts content from top results."""
     max_results = min(max(1, max_results), 5)
     loop = asyncio.get_running_loop()
 
@@ -355,61 +309,50 @@ async def search_and_scrape(query: str, max_results: int = 3) -> dict:
     if not results:
         return {"query": query, "results": []}
 
-    tasks = [
-        scrape_pipeline(res["href"], force_browser=False, auto_scroll=False)
-        for res in results
-    ]
+    tasks = [scrape_pipeline(res["href"], force_browser=False, auto_scroll=False) for res in results]
     scraped_payloads = await asyncio.gather(*tasks, return_exceptions=True)
 
     enriched = []
     for meta, payload in zip(results, scraped_payloads):
         if isinstance(payload, dict):
-            enriched.append(
-                {
-                    "title": meta.get("title", ""),
-                    "url": meta.get("href", ""),
-                    "snippet": meta.get("body", ""),
-                    "content": payload.get("content", ""),
-                }
-            )
+            enriched.append({
+                "title": meta.get("title", ""),
+                "url": meta.get("href", ""),
+                "snippet": meta.get("body", ""),
+                "content": payload.get("content", "")
+            })
         else:
-            enriched.append(
-                {
-                    "title": meta.get("title", ""),
-                    "url": meta.get("href", ""),
-                    "snippet": meta.get("body", ""),
-                    "content": f"[Error: {str(payload)}]",
-                }
-            )
+            enriched.append({
+                "title": meta.get("title", ""),
+                "url": meta.get("href", ""),
+                "snippet": meta.get("body", ""),
+                "content": f"[Error: {str(payload)}]"
+            })
     return {"query": query, "results": enriched}
 
-
-@mcp.tool()
-async def take_screenshot(
-    url: str, full_page: bool = True, wait_seconds: float = 1.0
-) -> dict:
-    """Takes a full-page or viewport screenshot of a webpage."""
+@mcp.tool(
+    name="take_screenshot",
+    description="Takes a full-page or viewport screenshot of a webpage as Base64 PNG."
+)
+async def take_screenshot(url: str, full_page: bool = True, wait_seconds: float = 1.0) -> dict:
     async with browser_pool.get_page() as page:
         await page.goto(url, wait_until="networkidle")
         if wait_seconds > 0:
             await asyncio.sleep(min(wait_seconds, 10.0))
-        screenshot_bytes = await page.screenshot(
-            full_page=full_page, type="png"
-        )
+        screenshot_bytes = await page.screenshot(full_page=full_page, type="png")
         b64 = base64.b64encode(screenshot_bytes).decode("utf-8")
         return {
             "url": page.url,
             "title": await page.title(),
             "format": "image/png;base64",
-            "base64_image": b64,
+            "base64_image": b64
         }
 
-
-@mcp.tool()
-async def interact_page(
-    url: str, actions: list[dict[str, Any]], extract_markdown: bool = True
-) -> dict:
-    """Executes a list of browser actions (click, type, press, wait, evaluate)."""
+@mcp.tool(
+    name="interact_page",
+    description="Executes a list of browser actions (click, type, press, wait, evaluate)."
+)
+async def interact_page(url: str, actions: list[dict[str, Any]], extract_markdown: bool = True) -> dict:
     async with browser_pool.get_page() as page:
         await page.goto(url, wait_until="domcontentloaded")
         logs = []
@@ -450,12 +393,11 @@ async def interact_page(
         res["metadata"] = parsed["metadata"]
     return res
 
-
-@mcp.tool()
-async def extract_structured_data(
-    url: str, css_selectors: dict[str, str] | None = None
-) -> dict:
-    """Parses structured JSON-LD schemas and arbitrary CSS selector targets."""
+@mcp.tool(
+    name="extract_structured_data",
+    description="Parses structured JSON-LD schemas and arbitrary CSS selector targets."
+)
+async def extract_structured_data(url: str, css_selectors: dict[str, str] | None = None) -> dict:
     async with browser_pool.get_page() as page:
         await page.goto(url, wait_until="domcontentloaded")
         html = await page.content()
@@ -470,31 +412,36 @@ async def extract_structured_data(
     return {
         "url": url,
         "custom_fields": custom_data,
-        "json_ld_schemas": extract_json_ld(html),
+        "json_ld_schemas": extract_json_ld(html)
     }
 
-
 # ==========================================
-# 6. Starlette SSE Router & Server Runner
+# 6. Starlette SSE Router & Tool Inspection
 # ==========================================
-
-
 async def health_check(request):
-    return JSONResponse(
-        {
-            "status": "healthy",
-            "mcp_sse_endpoint": "/sse",
-            "transport": "Server-Sent Events (SSE)",
-        }
-    )
-
+    return JSONResponse({
+        "status": "healthy",
+        "mcp_sse_endpoint": "/sse",
+        "tools_endpoint": "/tools"
+    })
+
+async def list_tools_endpoint(request):
+    """Direct HTTP endpoint exposing all registered MCP tools and schemas."""
+    tools_list = []
+    for tool_name, tool_obj in mcp._tool_manager._tools.items():
+        tools_list.append({
+            "name": tool_name,
+            "description": getattr(tool_obj, "description", "") or "",
+            "parameters": getattr(tool_obj, "parameters", {}) or {}
+        })
+    return JSONResponse({"tools": tools_list})
 
-# Mount the MCP SSE application with health endpoints
 starlette_app = Starlette(
     routes=[
         Route("/", endpoint=health_check),
         Route("/healthz", endpoint=health_check),
-        Mount("/", app=mcp.sse_app()),
+        Route("/tools", endpoint=list_tools_endpoint),
+        Mount("/", app=mcp.sse_app())
     ]
 )