Merge pull request #182 from rachit367/feat/domain-to-dns-enricher

feat(enrichers): add domain_to_dns enricher using dnsx

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

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 <domain> -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

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 == []