Add validation of endpoints for API/ETL and postprocessor hook

Author: gaya3-zipstack

backend/notification_v2/provider/webhook/webhook.py

@@ -7,6 +7,7 @@
 from backend.celery_service import app as celery_app
 from notification_v2.enums import AuthorizationType
 from notification_v2.provider.notification_provider import NotificationProvider
+from unstract.sdk.adapters.url_validator import URLValidator
 
 logger = logging.getLogger(__name__)
 
@@ -51,6 +52,13 @@ def validate(self):
         """
         if not self.notification.url:
             raise ValueError("Webhook URL is required.")
+
+        # Validate webhook URL for security
+        is_valid, error_message = URLValidator.validate_url(self.notification.url)
+        logger.info(f"Notification url {self.notification_url}")
+        if not is_valid:
+            raise ValueError(f"Webhook URL validation failed: {error_message}")
+
         if not self.payload:
             raise ValueError("Payload is required.")
         return super().validate()

backend/sample.env

@@ -201,7 +201,7 @@ RUNNER_POLLING_INTERVAL_SECONDS=2
 # Examples: 900 (15 min), 1800 (30 min), 3600 (60 min)
 MIN_SCHEDULE_INTERVAL_SECONDS=1800
 
-# WHitelisted adapter URLs to allow user to connect to locally hosted adapters.
+# Whitelisted adapter URLs to allow user to connect to locally hosted adapters.
 # Whitelisting 10.68.0.10 to allow frictionless adapter connection to
 # managed Postgres for VectorDB
 WHITELISTED_ENDPOINTS="10.68.0.10"

prompt-service/sample.env

@@ -64,3 +64,8 @@ ADAPTER_LLMW_STATUS_RETRIES=5
 # Rentroll Service
 RENTROLL_SERVICE_HOST=http://unstract-rentroll-service
 RENTROLL_SERVICE_PORT=5003
+
+# Whitelisted adapter URLs to allow user to connect to locally hosted adapters.
+# Whitelisting 10.68.0.10 to allow frictionless adapter connection to
+# managed Postgres for VectorDB
+WHITELISTED_ENDPOINTS="10.68.0.10"

prompt-service/src/unstract/prompt_service/services/answer_prompt.py

@@ -1,8 +1,5 @@
-import ipaddress
-import socket
 from logging import Logger
 from typing import Any
-from urllib.parse import urlparse
 
 from flask import current_app as app
 
@@ -17,6 +14,7 @@
     repair_json_with_best_structure,
 )
 from unstract.prompt_service.utils.log import publish_log
+from unstract.sdk.adapters.url_validator import URLValidator
 from unstract.sdk.constants import LogLevel
 from unstract.sdk.exceptions import RateLimitError as SdkRateLimitError
 from unstract.sdk.exceptions import SdkError
@@ -26,58 +24,6 @@
 from unstract.sdk.llm import LLM
 
 
-def _is_safe_public_url(url: str) -> bool:
-    """Validate webhook URL for SSRF protection.
-
-    Only allows HTTPS and blocks private/loopback/internal addresses.
-    Resolves all DNS records (A/AAAA) to prevent DNS rebinding attacks.
-    """
-    try:
-        p = urlparse(url)
-        if p.scheme not in ("https",):  # Only allow HTTPS for security
-            return False
-        host = p.hostname or ""
-        # Block obvious local hosts
-        if host in ("localhost",):
-            return False
-
-        addrs: set[str] = set()
-        # If literal IP, validate directly; else resolve all records (A/AAAA)
-        try:
-            ipaddress.ip_address(host)
-            addrs.add(host)
-        except ValueError:
-            try:
-                for family, _type, _proto, _canonname, sockaddr in socket.getaddrinfo(
-                    host, None, type=socket.SOCK_STREAM
-                ):
-                    addr = sockaddr[0]
-                    addrs.add(addr)
-            except Exception:
-                return False
-
-        if not addrs:
-            return False
-
-        # Validate all resolved addresses
-        for addr in addrs:
-            try:
-                ip = ipaddress.ip_address(addr)
-            except ValueError:
-                return False
-            if (
-                ip.is_private
-                or ip.is_loopback
-                or ip.is_link_local
-                or ip.is_reserved
-                or ip.is_multicast
-            ):
-                return False
-        return True
-    except Exception:
-        return False
-
-
 class AnswerPromptService:
     @staticmethod
     def extract_variable(
@@ -342,23 +288,25 @@ def handle_json(
                         app.logger.warning(
                             "Postprocessing webhook enabled but URL missing; skipping."
                         )
-                    elif not _is_safe_public_url(webhook_url):
-                        app.logger.warning(
-                            "Postprocessing webhook URL is not allowed; skipping."
-                        )
                     else:
-                        try:
-                            processed_data, updated_highlight_data = postprocess_data(
-                                parsed_data,
-                                webhook_enabled=True,
-                                webhook_url=webhook_url,
-                                highlight_data=highlight_data,
-                                timeout=60,
-                            )
-                        except Exception as e:
+                        is_valid, error_message = URLValidator.validate_url(webhook_url)
+                        if not is_valid:
                             app.logger.warning(
-                                f"Postprocessing webhook failed: {e}. Using unprocessed data."
+                                f"Postprocessing webhook URL validation failed: {error_message}; skipping."
                             )
+                        else:
+                            try:
+                                processed_data, updated_highlight_data = postprocess_data(
+                                    parsed_data,
+                                    webhook_enabled=True,
+                                    webhook_url=webhook_url,
+                                    highlight_data=highlight_data,
+                                    timeout=60,
+                                )
+                            except Exception as e:
+                                app.logger.warning(
+                                    f"Postprocessing webhook failed: {e}. Using unprocessed data."
+                                )
 
                 structured_output[prompt_key] = processed_data