From 5a6b24182deead4de9c778136cf77897ead0ffe2 Mon Sep 17 00:00:00 2001 From: rachit367 Date: Tue, 9 Jun 2026 11:34:10 +0530 Subject: [PATCH] feat(enrichers): make domain_to_website extractions opt-out domain_to_website always fetched the page and extracted headers, the full page text (up to 5000 chars), and technologies on every run, which is slow/noisy when a user only wants liveness + title (#90). Adds three optional select params (extract_content, extract_technologies, extract_headers), all defaulting to "true" so existing behavior is unchanged. Setting any to "false" skips that work. Title, description, status_code and active are always captured (cheap). Also de-duplicates the near-identical HTTPS/HTTP branches in scan() into a scheme loop + a _build_website_data helper, and drops two unused imports. This follows the params_schema convention from #60. Tests (requests mocked, no network): params schema, default full extraction, each toggle off, all-heavy-off keeping core fields, and the request-failure -> inactive path. Closes #90 --- .../flowsint_enrichers/domain/to_website.py | 223 ++++++++++-------- .../tests/enrichers/test_domain_to_website.py | 147 ++++++++++++ 2 files changed, 276 insertions(+), 94 deletions(-) create mode 100644 flowsint-enrichers/tests/enrichers/test_domain_to_website.py diff --git a/flowsint-enrichers/src/flowsint_enrichers/domain/to_website.py b/flowsint-enrichers/src/flowsint_enrichers/domain/to_website.py index 8a247164d..56642edff 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/domain/to_website.py +++ b/flowsint-enrichers/src/flowsint_enrichers/domain/to_website.py @@ -1,7 +1,6 @@ -from typing import Dict, List, Union +from typing import Any, Dict, List import requests from bs4 import BeautifulSoup -from flowsint_core.utils import is_valid_domain from flowsint_core.core.enricher_base import Enricher from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.domain import Domain @@ -29,8 +28,65 @@ def category(cls) -> str: def key(cls) -> str: return "domain" - def _extract_page_info(self, html_content: str) -> Dict[str, any]: - """Extract title, description, content and technologies from HTML.""" + @classmethod + def get_params_schema(cls) -> List[Dict[str, Any]]: + """Optional toggles to skip the heavier extractions for faster scans.""" + return [ + { + "name": "extract_content", + "type": "select", + "description": "Extract the page text content (can be slow / large)", + "required": False, + "default": "true", + "options": [ + {"label": "Enabled", "value": "true"}, + {"label": "Disabled", "value": "false"}, + ], + }, + { + "name": "extract_technologies", + "type": "select", + "description": "Detect web technologies from the HTML", + "required": False, + "default": "true", + "options": [ + {"label": "Enabled", "value": "true"}, + {"label": "Disabled", "value": "false"}, + ], + }, + { + "name": "extract_headers", + "type": "select", + "description": "Capture relevant HTTP response headers", + "required": False, + "default": "true", + "options": [ + {"label": "Enabled", "value": "true"}, + {"label": "Disabled", "value": "false"}, + ], + }, + ] + + def _param_enabled(self, name: str) -> bool: + """A select param is enabled unless explicitly set to 'false'.""" + return self.params.get(name, "true") != "false" + + def _extract_headers(self, response: requests.Response) -> Dict[str, str]: + """Pick out the relevant HTTP headers, dropping missing ones.""" + headers = { + "content-type": response.headers.get("content-type"), + "server": response.headers.get("server"), + "x-powered-by": response.headers.get("x-powered-by"), + } + return {k: v for k, v in headers.items() if v} + + def _extract_page_info( + self, + html_content: str, + extract_content: bool = True, + extract_technologies: bool = True, + ) -> Dict[str, Any]: + """Extract title, description, and (optionally) content and technologies.""" soup = BeautifulSoup(html_content, 'html.parser') # Extract title @@ -45,46 +101,70 @@ def _extract_page_info(self, html_content: str) -> Dict[str, any]: if meta_desc and meta_desc.get('content'): description = meta_desc.get('content').strip() + info: Dict[str, Any] = {'title': title, 'description': description} + # Extract text content (remove scripts and styles) - for script in soup(['script', 'style']): - script.decompose() - content = soup.get_text(separator=' ', strip=True) - # Limit content to first 5000 characters - if len(content) > 5000: - content = content[:5000] + "..." + if extract_content: + for script in soup(['script', 'style']): + script.decompose() + content = soup.get_text(separator=' ', strip=True) + # Limit content to first 5000 characters + if len(content) > 5000: + content = content[:5000] + "..." + info['content'] = content # Detect technologies - technologies = [] - - # Check for common frameworks and libraries - if soup.find('meta', attrs={'name': 'generator'}): - generator = soup.find('meta', attrs={'name': 'generator'}).get('content') - if generator: - technologies.append(generator) - - # Check for React - if soup.find('div', id='root') or soup.find('div', id='react-root'): - technologies.append('React') - - # Check for Vue.js - if soup.find(attrs={'data-v-'}): - technologies.append('Vue.js') - - # Check for Angular - if soup.find(attrs={'ng-app'}) or soup.find(attrs={'ng-version'}): - technologies.append('Angular') - - # Check for WordPress - if soup.find('meta', attrs={'name': 'generator', 'content': lambda x: x and 'WordPress' in x}): - technologies.append('WordPress') - - return { - 'title': title, - 'description': description, - 'content': content, - 'technologies': technologies + if extract_technologies: + technologies = [] + + # Check for common frameworks and libraries + if soup.find('meta', attrs={'name': 'generator'}): + generator = soup.find('meta', attrs={'name': 'generator'}).get('content') + if generator: + technologies.append(generator) + + # Check for React + if soup.find('div', id='root') or soup.find('div', id='react-root'): + technologies.append('React') + + # Check for Vue.js + if soup.find(attrs={'data-v-'}): + technologies.append('Vue.js') + + # Check for Angular + if soup.find(attrs={'ng-app'}) or soup.find(attrs={'ng-version'}): + technologies.append('Angular') + + # Check for WordPress + if soup.find('meta', attrs={'name': 'generator', 'content': lambda x: x and 'WordPress' in x}): + technologies.append('WordPress') + + info['technologies'] = technologies + + return info + + def _build_website_data( + self, domain: Domain, url: str, response: requests.Response + ) -> Dict[str, Any]: + """Assemble Website fields from a successful response, honoring params.""" + website_data: Dict[str, Any] = { + 'url': url, + 'domain': domain, + 'active': True, + 'status_code': response.status_code, } + if self._param_enabled("extract_headers"): + website_data['headers'] = self._extract_headers(response) + + page_info = self._extract_page_info( + response.text, + extract_content=self._param_enabled("extract_content"), + extract_technologies=self._param_enabled("extract_technologies"), + ) + website_data.update(page_info) + return website_data + async def scan(self, data: List[InputType]) -> List[OutputType]: results: List[OutputType] = [] for domain in data: @@ -95,65 +175,20 @@ async def scan(self, data: List[InputType]) -> List[OutputType]: 'active': False } - # Try HTTPS first - try: - https_url = f"https://{domain.domain}" - response = requests.get( - https_url, timeout=10, allow_redirects=True - ) - - if response.status_code < 400: - website_data['url'] = https_url - website_data['active'] = True - website_data['status_code'] = response.status_code - - # Extract relevant headers - website_data['headers'] = { - 'content-type': response.headers.get('content-type'), - 'server': response.headers.get('server'), - 'x-powered-by': response.headers.get('x-powered-by') - } - # Remove None values - website_data['headers'] = {k: v for k, v in website_data['headers'].items() if v} - - # Extract page information - page_info = self._extract_page_info(response.text) - website_data.update(page_info) - - results.append(Website(**website_data)) + # Try HTTPS first, then HTTP. + for scheme in ("https", "http"): + url = f"{scheme}://{domain.domain}" + try: + response = requests.get( + url, timeout=10, allow_redirects=True + ) + except requests.RequestException: continue - except requests.RequestException: - pass - - # Try HTTP if HTTPS fails - try: - http_url = f"http://{domain.domain}" - response = requests.get(http_url, timeout=10, allow_redirects=True) if response.status_code < 400: - website_data['url'] = http_url - website_data['active'] = True - website_data['status_code'] = response.status_code - - # Extract relevant headers - website_data['headers'] = { - 'content-type': response.headers.get('content-type'), - 'server': response.headers.get('server'), - 'x-powered-by': response.headers.get('x-powered-by') - } - # Remove None values - website_data['headers'] = {k: v for k, v in website_data['headers'].items() if v} - - # Extract page information - page_info = self._extract_page_info(response.text) - website_data.update(page_info) - - results.append(Website(**website_data)) - continue - except requests.RequestException: - pass + website_data = self._build_website_data(domain, url, response) + break - # If both fail, still add HTTPS URL as default results.append(Website(**website_data)) except Exception as e: diff --git a/flowsint-enrichers/tests/enrichers/test_domain_to_website.py b/flowsint-enrichers/tests/enrichers/test_domain_to_website.py new file mode 100644 index 000000000..72e4d50fc --- /dev/null +++ b/flowsint-enrichers/tests/enrichers/test_domain_to_website.py @@ -0,0 +1,147 @@ +import pytest + +from flowsint_enrichers import ENRICHER_REGISTRY +from flowsint_enrichers.domain.to_website import DomainToWebsiteEnricher +from flowsint_types.domain import Domain + +HTML = """ + + + Example Domain + + + +
hello world content here
+ +""" + + +class _FakeResponse: + def __init__(self, text=HTML, status_code=200, headers=None): + self.text = text + self.status_code = status_code + self.headers = headers or { + "content-type": "text/html", + "server": "nginx", + "x-powered-by": "PHP/8.1", + } + + +def _patch_get(monkeypatch, response): + def fake_get(url, *args, **kwargs): + return response + + monkeypatch.setattr( + "flowsint_enrichers.domain.to_website.requests.get", fake_get + ) + + +def _enricher(params=None): + return DomainToWebsiteEnricher( + sketch_id="s", scan_id="t", graph_service=None, params=params or {} + ) + + +# --------------------------------------------------------------------------- +# Registry / params schema +# --------------------------------------------------------------------------- +def test_registered(): + assert ( + ENRICHER_REGISTRY.get_enricher("domain_to_website", "1", "1").name() + == "domain_to_website" + ) + + +def test_params_schema_exposes_three_toggles(): + names = {p["name"] for p in DomainToWebsiteEnricher.get_params_schema()} + assert names == {"extract_content", "extract_technologies", "extract_headers"} + # all default to enabled -> backwards compatible + assert all( + p["default"] == "true" for p in DomainToWebsiteEnricher.get_params_schema() + ) + + +# --------------------------------------------------------------------------- +# Default behavior (all extractions on) — unchanged +# --------------------------------------------------------------------------- +@pytest.mark.asyncio +async def test_default_extracts_everything(monkeypatch): + _patch_get(monkeypatch, _FakeResponse()) + [site] = await _enricher().scan([Domain(domain="example.com")]) + + assert site.active is True + assert site.title == "Example Domain" + assert site.description == "An example site" + assert site.content and "hello world" in site.content + assert "WordPress 6.0" in site.technologies + assert "React" in site.technologies + assert site.headers.get("server") == "nginx" + + +# --------------------------------------------------------------------------- +# Toggles off — heavy extractions skipped, cheap fields still present +# --------------------------------------------------------------------------- +@pytest.mark.asyncio +async def test_disable_content(monkeypatch): + _patch_get(monkeypatch, _FakeResponse()) + [site] = await _enricher({"extract_content": "false"}).scan( + [Domain(domain="example.com")] + ) + assert site.content in (None, [], "") # not populated + assert site.title == "Example Domain" # cheap field still there + assert site.technologies # still on + + +@pytest.mark.asyncio +async def test_disable_technologies(monkeypatch): + _patch_get(monkeypatch, _FakeResponse()) + [site] = await _enricher({"extract_technologies": "false"}).scan( + [Domain(domain="example.com")] + ) + assert site.technologies == [] + assert site.content # still on + + +@pytest.mark.asyncio +async def test_disable_headers(monkeypatch): + _patch_get(monkeypatch, _FakeResponse()) + [site] = await _enricher({"extract_headers": "false"}).scan( + [Domain(domain="example.com")] + ) + assert site.headers in (None, {}) + assert site.title == "Example Domain" + + +@pytest.mark.asyncio +async def test_all_heavy_disabled_keeps_core_fields(monkeypatch): + _patch_get(monkeypatch, _FakeResponse()) + [site] = await _enricher( + { + "extract_content": "false", + "extract_technologies": "false", + "extract_headers": "false", + } + ).scan([Domain(domain="example.com")]) + + assert site.active is True + assert site.status_code == 200 + assert site.title == "Example Domain" + assert site.description == "An example site" + assert not site.content + assert site.technologies == [] + assert not site.headers + + +@pytest.mark.asyncio +async def test_inactive_when_request_fails(monkeypatch): + import requests as _requests + + def boom(*a, **k): + raise _requests.RequestException("no network") + + monkeypatch.setattr( + "flowsint_enrichers.domain.to_website.requests.get", boom + ) + [site] = await _enricher().scan([Domain(domain="example.com")]) + assert site.active is False + assert str(site.url).startswith("https://example.com")