diff --git a/flowsint-enrichers/src/flowsint_enrichers/domain/to_dns.py b/flowsint-enrichers/src/flowsint_enrichers/domain/to_dns.py new file mode 100644 index 0000000..3709b9f --- /dev/null +++ b/flowsint-enrichers/src/flowsint_enrichers/domain/to_dns.py @@ -0,0 +1,135 @@ +from typing import Any, Dict, List, Optional +from flowsint_core.core.logger import Logger +from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers.registry import flowsint_enricher +from flowsint_types.domain import Domain +from flowsint_types.ip import Ip +from tools.network.dnsx import DnsxTool + + +@flowsint_enricher +class DomainToDnsEnricher(Enricher): + """[DNSX] Resolve a domain's A (IPv4) and AAAA (IPv6) records using dnsx.""" + + InputType = Domain + OutputType = Ip + + def __init__( + self, + sketch_id: Optional[str] = None, + scan_id: Optional[str] = None, + vault=None, + params: Optional[Dict[str, Any]] = None, + **kwargs, + ): + super().__init__( + sketch_id=sketch_id, + scan_id=scan_id, + params_schema=self.get_params_schema(), + vault=vault, + params=params, + **kwargs, + ) + + @classmethod + def name(cls) -> str: + return "domain_to_dns" + + @classmethod + def category(cls) -> str: + return "Domain" + + @classmethod + def key(cls) -> str: + return "domain" + + @classmethod + def get_params_schema(cls) -> List[Dict[str, Any]]: + """Declare parameters for this enricher.""" + return [ + { + "name": "ipv6", + "type": "select", + "description": "Also resolve AAAA (IPv6) records", + "required": False, + "default": "true", + "options": [ + {"label": "Enabled", "value": "true"}, + {"label": "Disabled", "value": "false"}, + ], + }, + { + "name": "PDCP_API_KEY", + "type": "vaultSecret", + "description": "ProjectDiscovery Cloud Platform API key (optional)", + "required": False, + }, + ] + + async def scan(self, data: List[InputType]) -> List[OutputType]: + results: List[OutputType] = [] + + try: + dnsx = DnsxTool() + except Exception as e: + Logger.error( + self.sketch_id, + {"message": f"[DNSX] Failed to initialize dnsx: {e}"}, + ) + return results + + aaaa = self.params.get("ipv6", "true") == "true" + api_key = self.get_secret("PDCP_API_KEY", None) + + for d in data: + try: + ips = dnsx.resolve_domain(d.domain, aaaa=aaaa, api_key=api_key) + except Exception as e: + Logger.error( + self.sketch_id, + {"message": f"[DNSX] Error resolving {d.domain}: {e}"}, + ) + continue + + for ip in ips: + try: + ip_obj = Ip(address=ip) + except Exception: + # dnsx occasionally emits non-address artifacts; skip them. + continue + # Carry the source domain through to postprocess for graph wiring. + setattr(ip_obj, "_source_domain", d.domain) + results.append(ip_obj) + Logger.info( + self.sketch_id, + {"message": f"[DNSX] {d.domain} -> {ip}"}, + ) + + return results + + def postprocess( + self, results: List[OutputType], original_input: List[InputType] = None + ) -> List[OutputType]: + for ip_obj in results: + if not self._graph_service: + continue + source_domain = getattr(ip_obj, "_source_domain", None) + if not source_domain: + continue + + domain_obj = Domain(domain=source_domain) + self.create_node(domain_obj) + self.create_node(ip_obj) + self.create_relationship(domain_obj, ip_obj, "RESOLVES_TO") + self.log_graph_message( + f"IP found for domain {source_domain} -> {ip_obj.address}" + ) + + # Clean up the temporary attribute used to thread context. + delattr(ip_obj, "_source_domain") + + return results + + +InputType = DomainToDnsEnricher.InputType +OutputType = DomainToDnsEnricher.OutputType diff --git a/flowsint-enrichers/src/tools/network/dnsx.py b/flowsint-enrichers/src/tools/network/dnsx.py index d493db0..57d88b3 100644 --- a/flowsint-enrichers/src/tools/network/dnsx.py +++ b/flowsint-enrichers/src/tools/network/dnsx.py @@ -1,5 +1,5 @@ import json -from typing import Any, List +from typing import List from ..dockertool import DockerTool @@ -86,3 +86,72 @@ class DnsxTool(DockerTool): raise RuntimeError( f"Error running dnsx: {str(e)}. Output: {getattr(e, 'output', 'No output')}" ) + + def resolve_domain( + self, domain: str, aaaa: bool = True, api_key: str = None + ) -> List[str]: + """ + Resolve a domain's A (IPv4) and, optionally, AAAA (IPv6) records. + + Runs ``dnsx -d -a [-aaaa] -json -silent`` and parses the JSONL + output, returning the de-duplicated list of resolved IP addresses. + + Args: + domain: Domain name to resolve (e.g., "example.com") + aaaa: Whether to also query AAAA (IPv6) records. Defaults to True. + api_key: Optional ProjectDiscovery Cloud Platform API key + + Returns: + Ordered, de-duplicated list of resolved IP addresses (A then AAAA). + """ + flags = ["-a"] + if aaaa: + flags.append("-aaaa") + flags += ["-json", "-silent"] + + command = f"-d {domain} {' '.join(flags)}" + + env = {} + if api_key: + env["PDCP_API_KEY"] = api_key + + try: + result = super().launch(command, environment=env) + except Exception as e: + raise RuntimeError( + f"Error running dnsx: {str(e)}. Output: {getattr(e, 'output', 'No output')}" + ) + + return self._parse_resolved_ips(result) + + @staticmethod + def _parse_resolved_ips(result: str) -> List[str]: + """ + Parse dnsx JSONL output into an ordered, de-duplicated list of IPs. + + Each line is a JSON object whose ``a``/``aaaa`` keys hold the resolved + IPv4/IPv6 addresses (see retryabledns.DNSData). Malformed lines are + skipped so a single bad record never breaks the whole resolution. + """ + ips: List[str] = [] + seen: set[str] = set() + + if not result or not result.strip(): + return ips + + for line in result.split("\n"): + line = line.strip() + if not line: + continue + try: + record = json.loads(line) + except json.JSONDecodeError: + continue + + for key in ("a", "aaaa"): + for ip in record.get(key, []) or []: + if ip and ip not in seen: + seen.add(ip) + ips.append(ip) + + return ips diff --git a/flowsint-enrichers/tests/enrichers/test_domain_to_dns.py b/flowsint-enrichers/tests/enrichers/test_domain_to_dns.py new file mode 100644 index 0000000..a7c8952 --- /dev/null +++ b/flowsint-enrichers/tests/enrichers/test_domain_to_dns.py @@ -0,0 +1,131 @@ +import pytest + +from flowsint_enrichers import ENRICHER_REGISTRY +from flowsint_enrichers.domain.to_dns import DomainToDnsEnricher +from flowsint_types.ip import Ip +from tools.network.dnsx import DnsxTool + + +# --------------------------------------------------------------------------- +# Registry wiring +# --------------------------------------------------------------------------- +def test_domain_to_dns_is_registered(): + enricher = ENRICHER_REGISTRY.get_enricher("domain_to_dns", "123", "123") + assert enricher.name() == "domain_to_dns" + + +def test_domain_to_dns_metadata(): + assert DomainToDnsEnricher.category() == "Domain" + assert DomainToDnsEnricher.key() == "domain" + assert DomainToDnsEnricher.input_schema()["type"] == "Domain" + assert DomainToDnsEnricher.output_schema()["type"] == "Ip" + + +# --------------------------------------------------------------------------- +# dnsx JSONL parsing (no Docker required) +# --------------------------------------------------------------------------- +def test_parse_resolved_ips_extracts_a_and_aaaa(): + output = ( + '{"host":"example.com","a":["93.184.216.34"],' + '"aaaa":["2606:2800:220:1:248:1893:25c8:1946"]}' + ) + assert DnsxTool._parse_resolved_ips(output) == [ + "93.184.216.34", + "2606:2800:220:1:248:1893:25c8:1946", + ] + + +def test_parse_resolved_ips_dedupes_and_preserves_order(): + output = "\n".join( + [ + '{"host":"a.com","a":["1.1.1.1","1.0.0.1"]}', + '{"host":"a.com","a":["1.1.1.1"],"aaaa":["2606:4700:4700::1111"]}', + ] + ) + assert DnsxTool._parse_resolved_ips(output) == [ + "1.1.1.1", + "1.0.0.1", + "2606:4700:4700::1111", + ] + + +def test_parse_resolved_ips_skips_blank_and_malformed_lines(): + output = "\n".join( + [ + "", + "not-json", + '{"host":"a.com","a":["8.8.8.8"]}', + " ", + ] + ) + assert DnsxTool._parse_resolved_ips(output) == ["8.8.8.8"] + + +def test_parse_resolved_ips_handles_empty_output(): + assert DnsxTool._parse_resolved_ips("") == [] + assert DnsxTool._parse_resolved_ips(" \n ") == [] + + +# --------------------------------------------------------------------------- +# scan() — dnsx tool mocked, no Docker daemon touched +# --------------------------------------------------------------------------- +class _FakeDnsx: + def __init__(self, mapping): + self._mapping = mapping + self.calls = [] + + def resolve_domain(self, domain, aaaa=True, api_key=None): + self.calls.append((domain, aaaa, api_key)) + return self._mapping.get(domain, []) + + +@pytest.mark.asyncio +async def test_scan_returns_ip_objects_tagged_with_source_domain(monkeypatch): + fake = _FakeDnsx({"example.com": ["93.184.216.34", "2606:2800:220:1:248:1893:25c8:1946"]}) + monkeypatch.setattr( + "flowsint_enrichers.domain.to_dns.DnsxTool", lambda: fake + ) + + enricher = DomainToDnsEnricher(sketch_id="s", scan_id="t", graph_service=None) + from flowsint_types.domain import Domain + + results = await enricher.scan([Domain(domain="example.com")]) + + assert {ip.address for ip in results} == { + "93.184.216.34", + "2606:2800:220:1:248:1893:25c8:1946", + } + assert all(isinstance(ip, Ip) for ip in results) + assert all(getattr(ip, "_source_domain") == "example.com" for ip in results) + + +@pytest.mark.asyncio +async def test_scan_ipv6_param_disables_aaaa(monkeypatch): + fake = _FakeDnsx({"example.com": ["93.184.216.34"]}) + monkeypatch.setattr( + "flowsint_enrichers.domain.to_dns.DnsxTool", lambda: fake + ) + + enricher = DomainToDnsEnricher( + sketch_id="s", scan_id="t", graph_service=None, params={"ipv6": "false"} + ) + from flowsint_types.domain import Domain + + await enricher.scan([Domain(domain="example.com")]) + + # aaaa flag forwarded to the tool as False + assert fake.calls == [("example.com", False, None)] + + +@pytest.mark.asyncio +async def test_scan_skips_unresolvable_domain(monkeypatch): + fake = _FakeDnsx({}) # no records for anything + monkeypatch.setattr( + "flowsint_enrichers.domain.to_dns.DnsxTool", lambda: fake + ) + + enricher = DomainToDnsEnricher(sketch_id="s", scan_id="t", graph_service=None) + from flowsint_types.domain import Domain + + results = await enricher.scan([Domain(domain="example.com")]) + assert results == []