Merge pull request #87 from NatLabRockies/pp/safer_redirect

Safer redirect

Author: ppinchuk

elm/version.py

@@ -2,4 +2,4 @@
 ELM version number
 """
 
-__version__ = "0.0.43"
+__version__ = "0.0.44"

elm/web/file_loader.py

@@ -110,7 +110,7 @@ async def fetch(self, source):
             raise
         except Exception as e:
             msg = ("Encountered error of type %r while fetching document from "
-                   "%s:")
+                   "%s :")
             err_type = type(e)
             logger.exception(msg, err_type, source)
             return HTMLDocument(pages=[])

elm/web/search/google.py

@@ -18,6 +18,7 @@
 
 
 logger = logging.getLogger(__name__)
+_SERPER_SEMAPHORE = asyncio.Semaphore(50)  # default limit is 50 q per second
 
 
 class PlaywrightGoogleLinkSearch(PlaywrightSearchEngineLinkSearch):
@@ -309,11 +310,13 @@ async def _search(self, query, num_results=10, raw=False):
 
         payload = {"q": query, "num": num_results}
         headers = {"X-API-KEY": self.api_key}
+        c_kwargs = {"verify": self.verify, "timeout": 120}
 
-        async with httpx.AsyncClient(verify=self.verify) as client:
+        async with httpx.AsyncClient(**c_kwargs) as client, _SERPER_SEMAPHORE:
             response = await client.post(self._URL, headers=headers,
                                          json=payload)
             response.raise_for_status()
+            await asyncio.sleep(1)
 
         results = response.json().get("organic", [])
         return format_search_results(self._SE_NAME, query, results,

elm/web/search/run.py

@@ -404,6 +404,9 @@ async def search_all_se(queries, search_engines=_DEFAULT_SE,
     ELMInputError
         If `search_engines` input is empty.
     """
+    logger.debug("Running %d queries over these search engines: %r\n"
+                 "Queries:\n\t- %s", len(queries), search_engines,
+                 "\n\t- ".join(queries))
     num_urls = num_urls or 3 * len(queries)
     if len(search_engines) < 1:
         msg = f"Must provide at least one search engine! Got {search_engines=}"
@@ -441,7 +444,7 @@ async def load_docs(sources, file_loader):
     docs = await file_loader.fetch_all(*sources)
     logger.debug("Loaded %d docs from %d sources", len(docs), len(sources))
     docs = [doc for doc in docs if not doc.empty]
-    if len(docs)== 1:
+    if len(docs) == 1:
         logger.debug("%d doc is not empty", len(docs))
     else:
         logger.debug("%d docs are not empty", len(docs))
@@ -625,4 +628,4 @@ def _handle_old_ignore_key(url_ignore_substrings, kwargs):
     else:
         url_ignore_substrings = old_ignore_key
 
-    return url_ignore_substrings, kwargs
+    return url_ignore_substrings, kwargs

elm/web/utilities.py

@@ -4,10 +4,13 @@
 import hashlib
 import logging
 import asyncio
+import socket
+import ipaddress
 from pathlib import Path
 from copy import deepcopy
 from random import randint, choice
 from contextlib import asynccontextmanager
+from urllib.parse import urlparse, urljoin
 
 import httpx
 from slugify import slugify
@@ -81,7 +84,13 @@ async def get_redirected_url(url, **kwargs):
         The final URL after following redirects, or the original URL if
         no redirects are found or an error occurs.
     """
+
     kwargs["follow_redirects"] = True
+    event_hooks = kwargs.setdefault("event_hooks", {})
+    response_hooks = list(event_hooks.get("response") or [])
+    response_hooks.append(_check_redirect_safety)
+    event_hooks["response"] = response_hooks
+
     try:
         async with httpx.AsyncClient(**kwargs) as client:
             response = await client.head(url)
@@ -90,6 +99,53 @@ async def get_redirected_url(url, **kwargs):
         return url
 
 
+async def _check_redirect_safety(response):
+    """Validate each redirect target before following it."""
+    if not response.is_redirect:
+        return
+
+    redirect_url = response.headers.get("location")
+    if not redirect_url:
+        return
+
+    if not redirect_url.startswith(("http://", "https://")):
+        redirect_url = urljoin(str(response.url), redirect_url)
+
+    if not _is_safe_url(redirect_url):
+        raise ValueError(f"Redirect target is not allowed: {redirect_url}")
+
+
+def _is_safe_url(url):
+    """Return whether a URL resolves to a globally routable address."""
+    try:
+        parsed = urlparse(url)
+        hostname = parsed.hostname
+
+        if not hostname:
+            return False
+
+        try:
+            ip = ipaddress.ip_address(hostname)
+        except ValueError:
+
+            try:
+                ip_str = socket.gethostbyname(hostname)
+                ip = ipaddress.ip_address(ip_str)
+            except (socket.gaierror, socket.herror):
+                return False
+
+        return ip.is_global and not (
+            ip.is_private
+            or ip.is_loopback
+            or ip.is_link_local
+            or ip.is_reserved
+            or ip.is_multicast
+            or ip.is_unspecified
+        )
+    except Exception:
+        return False
+
+
 def clean_search_query(query):
     """Check if the first character is a digit and remove it if so.
 

tests/web/test_web_utilities.py

@@ -3,11 +3,13 @@
 from pathlib import Path
 
 import pytest
+import httpx
 
 from elm.web.document import HTMLDocument
 from elm.web.utilities import (
     clean_search_query,
     compute_fn_from_url,
+    get_redirected_url,
     write_url_doc_to_file,
 )
 
@@ -90,5 +92,27 @@ def test_write_url_doc_to_file(tmp_path):
     assert out_fp.name == "examplecom20test.txt"
 
 
+@pytest.mark.asyncio
+async def test_get_redirected_url_preserves_response_hooks():
+    """Test redirect lookup with an existing response hook."""
+
+    seen_urls = []
+
+    async def response_hook(response):
+        seen_urls.append(str(response.url))
+
+    def handler(request):
+        return httpx.Response(200, request=request)
+
+    out = await get_redirected_url(
+        "https://www.example.com",
+        transport=httpx.MockTransport(handler),
+        event_hooks={"response": [response_hook]},
+    )
+
+    assert out == "https://www.example.com"
+    assert seen_urls == ["https://www.example.com"]
+
+
 if __name__ == "__main__":
     pytest.main(["-q", "--show-capture=all", Path(__file__), "-rapP"])