add allowed domains cli option (#11)

Author: eyurtsev

mcpdoc/cli.py

@@ -42,6 +42,12 @@ class CustomFormatter(
   
   # Using SSE transport with additional HTTP options
   mcpdoc --yaml sample_config.yaml --follow-redirects --timeout 15 --transport sse --host localhost --port 8080
+  
+  # Allow fetching from additional domains. The domains hosting the llms.txt files are always allowed.
+  mcpdoc --yaml sample_config.yaml --allowed-domains https://example.com/ https://another-example.com/
+  
+  # Allow fetching from any domain
+  mcpdoc --yaml sample_config.yaml --allowed-domains '*'
 """
 
 
@@ -74,6 +80,12 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Whether to follow HTTP redirects",
     )
+    parser.add_argument(
+        "--allowed-domains",
+        type=str,
+        nargs="*",
+        help="Additional allowed domains to fetch documentation from. Use '*' to allow all domains",
+    )
     parser.add_argument(
         "--timeout", type=float, default=10.0, help="HTTP request timeout in seconds"
     )
@@ -229,6 +241,7 @@ def main() -> None:
         follow_redirects=args.follow_redirects,
         timeout=args.timeout,
         settings=settings,
+        allowed_domains=args.allowed_domains,
     )
 
     if args.transport == "sse":

mcpdoc/main.py

@@ -40,6 +40,7 @@ def create_server(
     follow_redirects: bool = False,
     timeout: float = 10,
     settings: dict | None = None,
+    allowed_domains: list[str] | None = None,
 ) -> FastMCP:
     """Create the server and generate documentation retrieval tools.
 
@@ -48,6 +49,10 @@ def create_server(
         follow_redirects: Whether to follow HTTP redirects when fetching docs
         timeout: HTTP request timeout in seconds
         settings: Additional settings to pass to FastMCP
+        allowed_domains: Additional domains to allow fetching from.
+            Use ['*'] to allow all domains
+            The domain hosting the llms.txt file is always appended to the list
+            of allowed domains.
 
     Returns:
         A FastMCP server instance configured with documentation tools
@@ -81,7 +86,14 @@ def list_doc_sources() -> str:
         return content
 
     # Parse the domain names in the llms.txt URLs
-    allowed_domains = set(extract_domain(entry["llms_txt"]) for entry in doc_source)
+    domains = set(extract_domain(entry["llms_txt"]) for entry in doc_source)
+
+    # Add additional allowed domains if specified
+    if allowed_domains:
+        if "*" in allowed_domains:
+            domains = {"*"}  # Special marker for allowing all domains
+        else:
+            domains.update(allowed_domains)
 
     @server.tool()
     async def fetch_docs(url: str) -> str:
@@ -99,11 +111,11 @@ async def fetch_docs(url: str) -> str:
             The fetched documentation content converted to markdown, or an error message
             if the request fails or the URL is not from an allowed domain.
         """
-        nonlocal allowed_domains
-        if not any(url.startswith(domain) for domain in allowed_domains):
+        nonlocal domains
+        if "*" not in domains and not any(url.startswith(domain) for domain in domains):
             return (
                 "Error: URL not allowed. Must start with one of the following domains: "
-                + ", ".join(allowed_domains)
+                + ", ".join(domains)
             )
 
         try: