IDEV-2204: Add DocstringPatcher to patch the specs from DT official techdocs.

Author: jbabac

domaintools/api.py

@@ -5,13 +5,15 @@
 
 import re
 import ssl
+import yaml
 
 from domaintools.constants import (
     Endpoint,
     OutputFormat,
     ENDPOINT_TO_SOURCE_MAP,
     RTTF_PRODUCTS_LIST,
     RTTF_PRODUCTS_CMD_MAPPING,
+    SPECS_MAPPING,
 )
 from domaintools._version import current as version
 from domaintools.results import (
@@ -29,7 +31,11 @@
     filter_by_field,
     DTResultFilter,
 )
-from domaintools.utils import validate_feeds_parameters
+from domaintools.utils import (
+    api_endpoint,
+    auto_patch_docstrings,
+    validate_feeds_parameters,
+)
 
 
 AVAILABLE_KEY_SIGN_HASHES = ["sha1", "sha256"]
@@ -40,6 +46,7 @@ def delimited(items, character="|"):
     return character.join(items) if type(items) in (list, tuple, set) else items
 
 
+@auto_patch_docstrings
 class API(object):
     """Enables interacting with the DomainTools API via Python:
 
@@ -94,8 +101,10 @@ def __init__(
         self.key_sign_hash = key_sign_hash
         self.default_parameters["app_name"] = app_name
         self.default_parameters["app_version"] = app_version
+        self.specs = {}
 
         self._build_api_url(api_url, api_port)
+        self._initialize_specs()
 
         if not https:
             raise Exception(
@@ -104,8 +113,25 @@ def __init__(
         if proxy_url and not isinstance(proxy_url, str):
             raise Exception("Proxy URL must be a string. For example: '127.0.0.1:8888'")
 
+    def _initialize_specs(self):
+        for spec_name, file_path in SPECS_MAPPING.items():
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    spec_content = yaml.safe_load(f)
+                if not spec_content:
+                    raise ValueError("Spec file is empty or invalid.")
+
+                self.specs[spec_name] = spec_content
+
+            except Exception as e:
+                print(f"Error loading {file_path}: {e}")
+
     def _get_ssl_default_context(self, verify_ssl: Union[str, bool]):
-        return ssl.create_default_context(cafile=verify_ssl) if isinstance(verify_ssl, str) else verify_ssl
+        return (
+            ssl.create_default_context(cafile=verify_ssl)
+            if isinstance(verify_ssl, str)
+            else verify_ssl
+        )
 
     def _build_api_url(self, api_url=None, api_port=None):
         """Build the API url based on the given url and port. Defaults to `https://api.domaintools.com`"""
@@ -133,11 +159,18 @@ def _rate_limit(self, product):
             hours = limit_hours and 3600 / float(limit_hours)
             minutes = limit_minutes and 60 / float(limit_minutes)
 
-            self.limits[product["id"]] = {"interval": timedelta(seconds=minutes or hours or default)}
+            self.limits[product["id"]] = {
+                "interval": timedelta(seconds=minutes or hours or default)
+            }
 
     def _results(self, product, path, cls=Results, **kwargs):
         """Returns _results for the specified API path with the specified **kwargs parameters"""
-        if product != "account-information" and self.rate_limit and not self.limits_set and not self.limits:
+        if (
+            product != "account-information"
+            and self.rate_limit
+            and not self.limits_set
+            and not self.limits
+        ):
             always_sign_api_key_previous_value = self.always_sign_api_key
             header_authentication_previous_value = self.header_authentication
             self._rate_limit(product)
@@ -181,7 +214,9 @@ def handle_api_key(self, is_rttf_product, path, parameters):
             else:
                 raise ValueError(
                     "Invalid value '{0}' for 'key_sign_hash'. "
-                    "Values available are {1}".format(self.key_sign_hash, ",".join(AVAILABLE_KEY_SIGN_HASHES))
+                    "Values available are {1}".format(
+                        self.key_sign_hash, ",".join(AVAILABLE_KEY_SIGN_HASHES)
+                    )
                 )
 
             parameters["timestamp"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
@@ -193,7 +228,9 @@ def handle_api_key(self, is_rttf_product, path, parameters):
 
     def account_information(self, **kwargs):
         """Provides a snapshot of your accounts current API usage"""
-        return self._results("account-information", "/v1/account", items_path=("products",), **kwargs)
+        return self._results(
+            "account-information", "/v1/account", items_path=("products",), **kwargs
+        )
 
     def available_api_calls(self):
         """Provides a list of api calls that you can use based on your account information."""
@@ -396,7 +433,9 @@ def reputation(self, query, include_reasons=False, **kwargs):
 
     def reverse_ip(self, domain=None, limit=None, **kwargs):
         """Pass in a domain name."""
-        return self._results("reverse-ip", "/v1/{0}/reverse-ip".format(domain), limit=limit, **kwargs)
+        return self._results(
+            "reverse-ip", "/v1/{0}/reverse-ip".format(domain), limit=limit, **kwargs
+        )
 
     def host_domains(self, ip=None, limit=None, **kwargs):
         """Pass in an IP address."""
@@ -570,8 +609,12 @@ def iris_enrich(self, *domains, **kwargs):
         younger_than_date = kwargs.pop("younger_than_date", {}) or None
         older_than_date = kwargs.pop("older_than_date", {}) or None
         updated_after = kwargs.pop("updated_after", {}) or None
-        include_domains_with_missing_field = kwargs.pop("include_domains_with_missing_field", {}) or None
-        exclude_domains_with_missing_field = kwargs.pop("exclude_domains_with_missing_field", {}) or None
+        include_domains_with_missing_field = (
+            kwargs.pop("include_domains_with_missing_field", {}) or None
+        )
+        exclude_domains_with_missing_field = (
+            kwargs.pop("exclude_domains_with_missing_field", {}) or None
+        )
 
         filtered_results = DTResultFilter(result_set=results).by(
             [
@@ -624,6 +667,7 @@ def iris_enrich_cli(self, domains=None, **kwargs):
             **kwargs,
         )
 
+    @api_endpoint(spec_name="iris", path="/v1/iris-investigate/")
     def iris_investigate(
         self,
         domains=None,
@@ -641,29 +685,6 @@ def iris_investigate(
         **kwargs,
     ):
         """Returns back a list of domains based on the provided filters.
-        The following filters are available beyond what is parameterized as kwargs:
-
-            - ip: Search for domains having this IP.
-            - email: Search for domains with this email in their data.
-            - email_domain: Search for domains where the email address uses this domain.
-            - nameserver_host: Search for domains with this nameserver.
-            - nameserver_domain: Search for domains with a nameserver that has this domain.
-            - nameserver_ip: Search for domains with a nameserver on this IP.
-            - registrar: Search for domains with this registrar.
-            - registrant: Search for domains with this registrant name.
-            - registrant_org: Search for domains with this registrant organization.
-            - mailserver_host: Search for domains with this mailserver.
-            - mailserver_domain: Search for domains with a mailserver that has this domain.
-            - mailserver_ip: Search for domains with a mailserver on this IP.
-            - redirect_domain: Search for domains which redirect to this domain.
-            - ssl_hash: Search for domains which have an SSL certificate with this hash.
-            - ssl_subject: Search for domains which have an SSL certificate with this subject string.
-            - ssl_email: Search for domains which have an SSL certificate with this email in it.
-            - ssl_org: Search for domains which have an SSL certificate with this organization in it.
-            - google_analytics: Search for domains which have this Google Analytics code.
-            - adsense: Search for domains which have this AdSense code.
-            - tld: Filter by TLD. Must be combined with another parameter.
-            - search_hash: Use search hash from Iris to bring back domains.
 
         You can loop over results of your investigation as if it was a native Python list:
 

domaintools/constants.py

@@ -56,3 +56,8 @@ class OutputFormat(Enum):
     "real-time-domain-discovery-feed-(api)": "domaindiscovery",
     "real-time-domain-discovery-feed-(s3)": "domaindiscovery",
 }
+
+SPECS_MAPPING = {
+    "iris": "domaintools/specs/iris-openapi.yaml",
+    # "rttf": "domaintools/specs/feeds-openapi.yaml",
+}

domaintools/docstring_patcher.py

@@ -0,0 +1,164 @@
+import inspect
+import functools
+import textwrap
+
+
+class DocstringPatcher:
+    """
+    Patches docstrings for methods decorated with @api_endpoint.
+    """
+
+    def patch(self, api_instance):
+        method_names = []
+        for attr_name in dir(api_instance):
+            attr = getattr(api_instance, attr_name)
+            # Look for the new decorator's tags
+            if (
+                inspect.ismethod(attr)
+                and hasattr(attr, "_api_spec_name")
+                and hasattr(attr, "_api_path")
+            ):
+                method_names.append(attr_name)
+
+        for attr_name in method_names:
+            original_method = getattr(api_instance, attr_name)
+            original_function = original_method.__func__
+
+            spec_name = original_function._api_spec_name
+            path = original_function._api_path
+
+            spec_to_use = api_instance.specs.get(spec_name)
+            original_doc = inspect.getdoc(original_function) or ""
+
+            all_doc_sections = []
+            if spec_to_use:
+                path_item = spec_to_use.get("paths", {}).get(path, {})
+
+                # Loop over all HTTP methods defined for this path
+                for http_method in ["get", "post", "put", "delete", "patch"]:
+                    if http_method in path_item:
+                        # Generate a doc section for this specific operation
+                        api_doc = self._generate_api_doc_string(spec_to_use, path, http_method)
+                        all_doc_sections.append(api_doc)
+
+            if not all_doc_sections:
+                all_doc_sections.append(
+                    f"\n--- API Details Error ---"
+                    f"\n  (Could not find any operations for path '{path}')"
+                )
+
+            # Combine the original doc with all operation docs
+            new_doc = textwrap.dedent(original_doc) + "\n\n" + "\n\n".join(all_doc_sections)
+
+            @functools.wraps(original_function)
+            def method_wrapper(*args, _orig_meth=original_method, **kwargs):
+                return _orig_meth(*args, **kwargs)
+
+            method_wrapper.__doc__ = new_doc
+            setattr(
+                api_instance,
+                attr_name,
+                method_wrapper.__get__(api_instance, api_instance.__class__),
+            )
+
+    def _generate_api_doc_string(self, spec: dict, path: str, method: str) -> str:
+        """Creates the formatted API docstring section for ONE operation."""
+
+        details = self._get_operation_details(spec, path, method)
+        # Add a clear title for this specific method
+        lines = [f"--- Operation: {method.upper()} {path} ---"]
+
+        # Render Query Params
+        lines.append(f"\n  Summary: {details.get('summary')}")
+        lines.append(f"  Description: {details.get('description')}")
+        lines.append(f"  External Doc: {details.get('external_doc')}")
+        lines.append("\n  Query Parameters:")
+        if not details["query_params"]:
+            lines.append("    (No query parameters)")
+        else:
+            for param in details["query_params"]:
+                lines.append(f"\n    **{param['name']}** ({param['type']})")
+                lines.append(f"      Required:    {param['required']}")
+                lines.append(f"      Description: {param['description']}")
+
+        # Render Request Body
+        lines.append("\n  Request Body:")
+        if not details["request_body"]:
+            lines.append("    (No request body)")
+        else:
+            body = details["request_body"]
+            lines.append(f"\n    **{body['type']}**")
+            lines.append(f"      Required:    {body['required']}")
+            lines.append(f"      Description: {body['description']}")
+
+        return "\n".join(lines)
+
+    def _get_operation_details(self, spec: dict, path: str, method: str) -> dict:
+        details = {"query_params": [], "request_body": None}
+        if not spec:
+            return details
+        try:
+            path_item = spec.get("paths", {}).get(path, {})
+            operation = path_item.get(method.lower(), {})
+            if not operation:
+                return details
+            all_param_defs = path_item.get("parameters", []) + operation.get("parameters", [])
+            details["summary"] = operation.get("summary")
+            details["description"] = operation.get("description")
+            details["external_doc"] = operation.get("externalDocs", {}).get("url", "N/A")
+            resolved_params = []
+            for param_def in all_param_defs:
+                if "$ref" in param_def:
+                    resolved_params.append(self._resolve_ref(spec, param_def["$ref"]))
+                else:
+                    resolved_params.append(param_def)
+            for p in [p for p in resolved_params if p.get("in") == "query"]:
+                details["query_params"].append(
+                    {
+                        "name": p.get("name"),
+                        "required": p.get("required", False),
+                        "description": p.get("description", "N/A"),
+                        "type": self._get_param_type(spec, p.get("schema")),
+                    }
+                )
+            body_def = operation.get("requestBody")
+            if body_def:
+                if "$ref" in body_def:
+                    body_def = self._resolve_ref(spec, body_def["$ref"])
+                content = body_def.get("content", {})
+                media_type = next(iter(content.values()), None)
+                if media_type and "schema" in media_type:
+                    schema = media_type["schema"]
+                    schema_type = self._get_param_type(spec, schema)
+                    if "$ref" in schema:
+                        schema_type = schema["$ref"].split("/")[-1]
+                    details["request_body"] = {
+                        "required": body_def.get("required", False),
+                        "description": body_def.get("description", "N/A"),
+                        "type": schema_type,
+                    }
+            return details
+        except Exception:
+            return details
+
+    def _resolve_ref(self, spec: dict, ref: str):
+        if not spec or not ref.startswith("#/"):
+            return {}
+        parts = ref.split("/")[1:]
+        current_obj = spec
+        for part in parts:
+            if not isinstance(current_obj, dict):
+                return {}
+            current_obj = current_obj.get(part)
+            if current_obj is None:
+                return {}
+        return current_obj
+
+    def _get_param_type(self, spec: dict, schema: dict) -> str:
+        if not schema:
+            return "N/A"
+        schema_ref = schema.get("$ref")
+        if schema_ref:
+            resolved_schema = self._resolve_ref(spec, schema_ref)
+            return resolved_schema.get("type", "N/A")
+        return schema.get("type", "N/A")