From 99432461125ef6f9c3c99bf3c6afde172f05926e Mon Sep 17 00:00:00 2001 From: Ontolocy Date: Sun, 12 Oct 2025 17:18:44 +0100 Subject: [PATCH] build out tool support --- .github/workflows/ci.yml | 2 +- docker-compose.yml | 17 -- setup.cfg | 2 +- src/ontolocy/__init__.py | 2 - src/ontolocy/models/dnsrecord.py | 27 ++- src/ontolocy/models/domainname.py | 4 +- src/ontolocy/tools/capec.py | 2 +- src/ontolocy/tools/cisa_kev.py | 2 +- src/ontolocy/tools/ctid_attack_mappings.py | 2 +- src/ontolocy/tools/cwe.py | 2 +- src/ontolocy/tools/ht_dns.py | 177 +++++++++++++++++++ src/ontolocy/tools/ht_ptrip.py | 184 ++++++++++++++++++++ src/ontolocy/tools/ht_revdns.py | 190 +++++++++++++++++++++ src/ontolocy/tools/ht_subdomain.py | 167 ++++++++++++++++++ src/ontolocy/tools/mitre_attack.py | 2 +- src/ontolocy/tools/nist_csf_1.py | 4 +- src/ontolocy/tools/nist_csf_2.py | 4 +- src/ontolocy/tools/nist_sp80053_v4.py | 3 +- src/ontolocy/tools/nist_sp80053_v5.py | 4 +- src/ontolocy/tools/nvd.py | 2 +- src/ontolocy/tools/ontolocy_enricher.py | 4 +- src/ontolocy/tools/ontolocy_parser.py | 12 +- src/ontolocy/tools/shodan.py | 2 +- tests/conftest.py | 61 ++++--- tests/test_cli.py | 17 +- tests/test_tools/test_capec.py | 8 +- tests/test_tools/test_cisa_kev.py | 23 ++- tests/test_tools/test_ht_dns.py | 138 +++++++++++++++ tests/test_tools/test_ht_revdns.py | 93 ++++++++++ tests/test_tools/test_ht_subdomain.py | 97 +++++++++++ tests/test_tools/test_htptrip.py | 97 +++++++++++ tests/test_tools/test_ontolocy_parser.py | 2 +- 32 files changed, 1237 insertions(+), 116 deletions(-) create mode 100644 src/ontolocy/tools/ht_dns.py create mode 100644 src/ontolocy/tools/ht_ptrip.py create mode 100644 src/ontolocy/tools/ht_revdns.py create mode 100644 src/ontolocy/tools/ht_subdomain.py create mode 100644 tests/test_tools/test_ht_dns.py create mode 100644 tests/test_tools/test_ht_revdns.py create mode 100644 tests/test_tools/test_ht_subdomain.py create mode 100644 tests/test_tools/test_htptrip.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0fc3911..0d6b331 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,7 +26,7 @@ jobs: fail-fast: true max-parallel: 2 matrix: - python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"] + python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"] steps: - uses: actions/checkout@v3 diff --git a/docker-compose.yml b/docker-compose.yml index 79f055b..0b31c90 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,5 +1,3 @@ -version: "3" - services: neo4j: image: neo4j:5.26-community @@ -29,18 +27,3 @@ services: retries: 5 start_period: 20s timeout: 10s - testmemgraph: - image: memgraph/memgraph:latest - ports: - - 9687:7687 - - 7444:7444 - command: ["--log-level=TRACE"] - healthcheck: - test: ["CMD-SHELL", "echo 'RETURN 0;' | mgconsole || exit 1"] - interval: 10s - timeout: 5s - retries: 3 - start_period: 0s - environment: - - MEMGRAPH_USER=${TEST_MEMGRAPH_USER} - - MEMGRAPH_PASSWORD=${TEST_MEMGRAPH_PASSWORD} diff --git a/setup.cfg b/setup.cfg index c24e51e..aa06c79 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = ontolocy -version = 0.9.0 +version = 0.9.1 author = Ontolocy description = A Python package for modeling cybersecurity data in a Neo4j graph database. long_description = file: README.md diff --git a/src/ontolocy/__init__.py b/src/ontolocy/__init__.py index 2f636ca..3a71054 100644 --- a/src/ontolocy/__init__.py +++ b/src/ontolocy/__init__.py @@ -40,7 +40,6 @@ from ontolocy.models.detection import Detection, DetectionForAttackTechnique from ontolocy.models.dnsrecord import ( DNSRecord, - DNSRecordForDomain, DNSRecordPointsToDomainName, DNSRecordPointsToIPAddress, ) @@ -208,7 +207,6 @@ "Detection", "DetectionForAttackTechnique", "DNSRecord", - "DNSRecordForDomain", "DNSRecordPointsToDomainName", "DNSRecordPointsToIPAddress", "DomainName", diff --git a/src/ontolocy/models/dnsrecord.py b/src/ontolocy/models/dnsrecord.py index e151b18..7593ff9 100644 --- a/src/ontolocy/models/dnsrecord.py +++ b/src/ontolocy/models/dnsrecord.py @@ -1,8 +1,7 @@ from datetime import datetime from typing import ClassVar, Optional -from uuid import UUID -from pydantic import ValidationInfo, field_validator +from pydantic import ValidationInfo, field_validator, Field from ..node import OntolocyNode from ..relationship import OntolocyRelationship @@ -16,11 +15,14 @@ class DNSRecord(OntolocyNode): type: str name: str content: str + record_class: str = "IN" - unique_id: Optional[UUID] = None + ttl: Optional[int] = None + + unique_id: Optional[str] = None @field_validator("unique_id") - def generate_dnsrecord_uuid(cls, v: Optional[UUID], info: ValidationInfo) -> UUID: + def generate_dnsrecord_uuid(cls, v: Optional[str], info: ValidationInfo) -> str: values = info.data if v is None: @@ -32,7 +34,7 @@ def generate_dnsrecord_uuid(cls, v: Optional[UUID], info: ValidationInfo) -> UUI v = generate_deterministic_uuid(key_values) - return v + return str(v) class DNSRecordPointsToIPAddress(OntolocyRelationship): @@ -41,7 +43,7 @@ class DNSRecordPointsToIPAddress(OntolocyRelationship): source: DNSRecord target: "IPAddressNode" - observation_date: datetime + observation_date: datetime = Field(default_factory=datetime.now) class DNSRecordPointsToDomainName(OntolocyRelationship): @@ -50,22 +52,13 @@ class DNSRecordPointsToDomainName(OntolocyRelationship): source: DNSRecord target: "DomainName" - observation_date: datetime - - -class DNSRecordForDomain(OntolocyRelationship): - __relationshiptype__: ClassVar[str] = "DNS_RECORD_FOR_DOMAIN" - - source: DNSRecord - target: "DomainName" - - observation_date: datetime + observation_date: datetime = Field(default_factory=datetime.now) from .domainname import DomainName # noqa: E402 DNSRecordPointsToDomainName.model_rebuild() -DNSRecordForDomain.model_rebuild() + from .ip import IPAddressNode # noqa: E402 diff --git a/src/ontolocy/models/domainname.py b/src/ontolocy/models/domainname.py index ee3d4ac..bf8ec5f 100644 --- a/src/ontolocy/models/domainname.py +++ b/src/ontolocy/models/domainname.py @@ -2,7 +2,7 @@ from datetime import datetime from typing import ClassVar, Optional -from pydantic import field_validator +from pydantic import field_validator, Field from ..node import OntolocyNode from ..relationship import OntolocyRelationship @@ -32,7 +32,7 @@ class DomainNameHasDNSRecord(OntolocyRelationship): source: DomainName target: "DNSRecord" - observation_date: datetime + observation_date: datetime = Field(default_factory=datetime.now) from .dnsrecord import DNSRecord # noqa: E402 diff --git a/src/ontolocy/tools/capec.py b/src/ontolocy/tools/capec.py index 85b47cf..6a4bd90 100644 --- a/src/ontolocy/tools/capec.py +++ b/src/ontolocy/tools/capec.py @@ -70,7 +70,7 @@ def _load_data(self, raw_data): return data - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: node_dfs = {} rel_dfs = {} diff --git a/src/ontolocy/tools/cisa_kev.py b/src/ontolocy/tools/cisa_kev.py index 12668eb..b02796f 100644 --- a/src/ontolocy/tools/cisa_kev.py +++ b/src/ontolocy/tools/cisa_kev.py @@ -52,7 +52,7 @@ def _detect(self, input_data) -> bool: def _load_data(self, raw_data): return pd.read_csv(StringIO(raw_data)) - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: node_dfs = {} rel_dfs = {} diff --git a/src/ontolocy/tools/ctid_attack_mappings.py b/src/ontolocy/tools/ctid_attack_mappings.py index 0a7dd51..86866ed 100644 --- a/src/ontolocy/tools/ctid_attack_mappings.py +++ b/src/ontolocy/tools/ctid_attack_mappings.py @@ -38,7 +38,7 @@ def _detect(self, input_data) -> bool: def _load_data(self, raw_data): return json.loads(raw_data) - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: node_dfs = {} rel_dfs = {} diff --git a/src/ontolocy/tools/cwe.py b/src/ontolocy/tools/cwe.py index 883a94c..fc21733 100644 --- a/src/ontolocy/tools/cwe.py +++ b/src/ontolocy/tools/cwe.py @@ -61,7 +61,7 @@ def _load_url(self, url): return response.content - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: node_dfs = {} rel_dfs = {} diff --git a/src/ontolocy/tools/ht_dns.py b/src/ontolocy/tools/ht_dns.py new file mode 100644 index 0000000..0cf42bb --- /dev/null +++ b/src/ontolocy/tools/ht_dns.py @@ -0,0 +1,177 @@ +import json +import requests +import pandas as pd + +from ontolocy import ( + DNSRecord, + DomainName, + IPAddressNode, + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, +) + +from .ontolocy_parser import OntolocyParser +from .ontolocy_enricher import OntolocyClient, OntolocyEnricher, SeedTypeEnum + + +class HackerTargetDNSParser(OntolocyParser): + """Parser for HackerTarget DNS lookup data. + + See https://hackertarget.com/dns-lookup/ for more details. + """ + + node_types = [DNSRecord, DomainName, IPAddressNode] + + rel_types = [ + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, + ] + + def _detect(self, input_data) -> bool: + + if not isinstance(input_data, dict): + return False + + for key in input_data.keys(): + if key not in ["A", "AAAA", "CNAME", "MX", "NS", "SOA", "TXT", "PTR"]: + return False + + return True + + def _load_data(self, raw_data): + return json.loads(raw_data) + + def _parse(self, input_data, private_namespace, ctx): + """ + Parse the data. + + Expects ctx to be a dictionary with a 'domain' key for the domain name queried. + + """ + + # create complete individual records + + records = [] + domains = [] + ips = [] + domain_to_dnsrecord_rels = [] + dnsrecord_to_ip_rels = [] + dnsrecord_to_domain_rels = [] + + domain_name = ctx["query"] + + domains.append({"name": domain_name}) + + for record_type, content in input_data.items(): + for entry in content: + record = { + "type": record_type, + "name": domain_name, + "content": entry, + } + + records.append(record) + + record_id = DNSRecord( + type=record_type, + name=domain_name, + content=entry, + ).unique_id + + domain_to_dnsrecord_rels.append( + {"source": domain_name, "target": record_id} + ) + + if record_type in ["A", "AAAA"]: + ips.append({"ip_address": entry}) + dnsrecord_to_ip_rels.append({"source": record_id, "target": entry}) + + elif record_type in ["CNAME", "NS", "PTR"]: + target_domain = entry.rstrip(".") + domains.append({"name": target_domain}) + dnsrecord_to_domain_rels.append( + {"source": record_id, "target": target_domain} + ) + + node_dfs = { + DNSRecord.__primarylabel__: pd.DataFrame.from_records(records) + .drop_duplicates() + .reset_index(drop=True), + DomainName.__primarylabel__: pd.DataFrame.from_records(domains) + .drop_duplicates() + .reset_index(drop=True), + IPAddressNode.__primarylabel__: pd.DataFrame.from_records(ips) + .drop_duplicates() + .reset_index(drop=True), + } + + rel_dfs = { + DomainNameHasDNSRecord.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToIPAddress.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToDomainName.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["target"] + ].copy(), + }, + } + + return node_dfs, rel_dfs + + +class HackerTargetDNSClient(OntolocyClient): + """Lightweight client for querying the HackerTarget DNS lookup API. + + See https://hackertarget.com/dns-lookup/ for more details. + + Query method expects a domain name as input. + """ + + def __init__(self): + super().__init__() + self.parser = HackerTargetDNSParser() + + def _query(self, query: str): + """Query the HackerTarget DNS lookup API. + + Args: + query (str): domain name to lookup + """ + + api_endpoint = "https://api.hackertarget.com/dnslookup/" + + response = requests.get(api_endpoint, params={"q": query, "output": "json"}) + + # raise an exception for bad responses + response.raise_for_status() + + return response.json() + + +class HackerTargetDNSEnricher(OntolocyEnricher): + seed_type = SeedTypeEnum.DOMAIN + + def __init__(self): + super().__init__() + self.client = HackerTargetDNSClient() + + def _generate_single_query(self, seed): + return seed diff --git a/src/ontolocy/tools/ht_ptrip.py b/src/ontolocy/tools/ht_ptrip.py new file mode 100644 index 0000000..7e162d1 --- /dev/null +++ b/src/ontolocy/tools/ht_ptrip.py @@ -0,0 +1,184 @@ +import requests +import re + +import pandas as pd + +from ontolocy import ( + DNSRecord, + DomainName, + IPAddressNode, + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, +) + +from .ontolocy_parser import OntolocyParser +from .ontolocy_enricher import ( + OntolocyClient, + OntolocyEnricher, + SeedTypeEnum, + SEED_MAPPINGS, +) + + +class HackerTargetPtrIPParser(OntolocyParser): + """Parser for HackerTarget reverse DNS lookup by IP. + + This endpoint actively resolves reverse DNS (PTR) records for given IP addresses. + + See https://hackertarget.com/reverse-dns-lookup/ for more details. + """ + + node_types = [DNSRecord, DomainName, IPAddressNode] + + rel_types = [ + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, + ] + + def _detect(self, input_data: str) -> bool: + + # expects new line separated entries of "IP DOMAIN" + for line in input_data.splitlines(): + parts = line.split() + if len(parts) != 2: + return False + + ip_part = parts[0] + + if not re.search(SEED_MAPPINGS["ip"]["pattern"], ip_part): + return False + + return True + + def _parse(self, input_data, private_namespace, ctx): + """ + Parse the data. + + Expects ctx to be a dictionary with a 'domain' key for the domain name queried. + + """ + + records = [] + domains = [] + ips = [] + domain_to_dnsrecord_rels = [] + dnsrecord_to_ip_rels = [] + dnsrecord_to_domain_rels = [] + + for line in input_data.splitlines(): + parts = line.split() + ip_part = parts[0] + domain_part = parts[1] + + record_name = f"{ip_part}.in-addr.arpa." + + ptr_domain = record_name.rstrip(".") + + domains.append({"name": ptr_domain}) + + record = { + "type": "PTR", + "name": record_name, + "content": domain_part, + } + + records.append(record) + + record_id = DNSRecord(**record).unique_id + + ips.append({"ip_address": ip_part}) + + dnsrecord_to_ip_rels.append({"source": record_id, "target": ip_part}) + + domain_to_dnsrecord_rels.append({"source": ptr_domain, "target": record_id}) + + # entries may be wildcards rather than explicit domain names + if re.match(SEED_MAPPINGS["domain"]["pattern"], domain_part): + domains.append({"name": domain_part}) + + dnsrecord_to_domain_rels.append( + {"source": record_id, "target": domain_part} + ) + + node_dfs = { + DNSRecord.__primarylabel__: pd.DataFrame.from_records(records) + .drop_duplicates() + .reset_index(drop=True), + DomainName.__primarylabel__: pd.DataFrame.from_records(domains) + .drop_duplicates() + .reset_index(drop=True), + IPAddressNode.__primarylabel__: pd.DataFrame.from_records(ips) + .drop_duplicates() + .reset_index(drop=True), + } + + rel_dfs = { + DomainNameHasDNSRecord.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToIPAddress.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToDomainName.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["target"] + ].copy(), + }, + } + + return node_dfs, rel_dfs + + +class HackerTargetPtrIPClient(OntolocyClient): + """Lightweight client for querying the HackerTarget DNS lookup API. + + See https://hackertarget.com/dns-lookup/ for more details. + + Query method expects a domain name as input. + """ + + def __init__(self): + super().__init__() + self.parser = HackerTargetPtrIPParser() + + def _query(self, query: str): + """Query the HackerTarget Reverse DNS lookup API. + + Args: + query (str): domain name to lookup + """ + + api_endpoint = "https://api.hackertarget.com/reversedns/" + + response = requests.get(api_endpoint, params={"q": query}) + + # raise an exception for bad responses + response.raise_for_status() + + return response.text + + +class HackerTargetPtrIPEnricher(OntolocyEnricher): + seed_type = SeedTypeEnum.IP + + def __init__(self): + super().__init__() + self.client = HackerTargetPtrIPClient() + + def _generate_single_query(self, seed): + return seed diff --git a/src/ontolocy/tools/ht_revdns.py b/src/ontolocy/tools/ht_revdns.py new file mode 100644 index 0000000..61525d4 --- /dev/null +++ b/src/ontolocy/tools/ht_revdns.py @@ -0,0 +1,190 @@ +import requests +import pandas as pd +import re + +from ontolocy import ( + DNSRecord, + DomainName, + IPAddressNode, + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, +) + +from .ontolocy_parser import OntolocyParser +from .ontolocy_enricher import ( + OntolocyClient, + OntolocyEnricher, + SeedTypeEnum, + SEED_MAPPINGS, +) + + +class HackerTargetReverseDNSParser(OntolocyParser): + """Parser for HackerTarget reverse DNS lookups. + + This parser is for lookups by domain name. + + See https://hackertarget.com/reverse-dns-lookup/ for more details. + """ + + node_types = [DNSRecord, DomainName, IPAddressNode] + + rel_types = [ + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, + ] + + def _detect(self, input_data: str) -> bool: + + print(input_data) + + # expects new line separated entries of "IP DOMAIN" + for line in input_data.splitlines(): + parts = line.split(",") + + print(parts) + if len(parts) != 2: + return False + + ip_part = parts[1] + + if not re.search(SEED_MAPPINGS["ip"]["pattern"], ip_part): + return False + + return True + + def _parse(self, input_data, private_namespace, ctx): + """ + Parse the data. + + Expects ctx to be a dictionary with a 'domain' key for the domain name queried. + + """ + + records = [] + domains = [] + ips = [] + domain_to_dnsrecord_rels = [] + dnsrecord_to_ip_rels = [] + dnsrecord_to_domain_rels = [] + + for line in input_data.splitlines(): + parts = line.split(",") + domain_part = parts[0] + ip_part = parts[1] + + rev_ip = ".".join(ip_part.split(".")[::-1]) + + record_name = f"{rev_ip}.in-addr.arpa." + + ptr_domain = record_name.rstrip(".") + + domains.append({"name": ptr_domain}) + + record = { + "type": "PTR", + "name": record_name, + "content": domain_part, + } + + records.append(record) + + record_id = DNSRecord(**record).unique_id + + # create the relationships + domain_to_dnsrecord_rels.append({"source": ptr_domain, "target": record_id}) + + ips.append({"ip_address": ip_part}) + + # entries may be wildcards rather than explicit domain names + if re.match(SEED_MAPPINGS["domain"]["pattern"], domain_part): + domains.append({"name": domain_part}) + + dnsrecord_to_domain_rels.append( + {"source": record_id, "target": domain_part} + ) + + dnsrecord_to_ip_rels.append({"source": record_id, "target": ip_part}) + + node_dfs = { + DNSRecord.__primarylabel__: pd.DataFrame.from_records(records) + .drop_duplicates() + .reset_index(drop=True), + DomainName.__primarylabel__: pd.DataFrame.from_records(domains) + .drop_duplicates() + .reset_index(drop=True), + IPAddressNode.__primarylabel__: pd.DataFrame.from_records(ips) + .drop_duplicates() + .reset_index(drop=True), + } + + rel_dfs = { + DomainNameHasDNSRecord.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToIPAddress.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToDomainName.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["target"] + ].copy(), + }, + } + + return node_dfs, rel_dfs + + +class HackerTargetReverseDNSClient(OntolocyClient): + """Lightweight client for querying the HackerTarget DNS lookup API. + + See https://hackertarget.com/dns-lookup/ for more details. + + Query method expects a domain name as input. + """ + + def __init__(self): + super().__init__() + self.parser = HackerTargetReverseDNSParser() + + def _query(self, query: str): + """Query the HackerTarget DNS lookup API. + + Args: + query (str): domain name to lookup + """ + + api_endpoint = "https://api.hackertarget.com/reversedns/" + + response = requests.get(api_endpoint, params={"q": query}) + + # raise an exception for bad responses + response.raise_for_status() + + return response.text + + +class HackerTargetReverseDNSEnricher(OntolocyEnricher): + seed_type = SeedTypeEnum.DOMAIN + + def __init__(self): + super().__init__() + self.client = HackerTargetReverseDNSClient() + + def _generate_single_query(self, seed): + return seed diff --git a/src/ontolocy/tools/ht_subdomain.py b/src/ontolocy/tools/ht_subdomain.py new file mode 100644 index 0000000..44a7156 --- /dev/null +++ b/src/ontolocy/tools/ht_subdomain.py @@ -0,0 +1,167 @@ +import requests +import pandas as pd +import re + +from ontolocy import ( + DNSRecord, + DomainName, + IPAddressNode, + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, +) + +from .ontolocy_parser import OntolocyParser +from .ontolocy_enricher import ( + OntolocyClient, + OntolocyEnricher, + SeedTypeEnum, + SEED_MAPPINGS, +) + + +def reverse_ip(ip): + """Reverse the octets of an IPv4 address for PTR record naming.""" + return ".".join(ip.split(".")[::-1]) + + +class HackerTargetSubdomainParser(OntolocyParser): + """ + Parser for HackerTarget subdomain enumeration endpoint. + + See https://hackertarget.com/find-dns-host-records/ + """ + + node_types = [DNSRecord, DomainName, IPAddressNode] + rel_types = [ + DNSRecordPointsToDomainName, + DNSRecordPointsToIPAddress, + DomainNameHasDNSRecord, + ] + + def _detect(self, input_data: str) -> bool: + # expects new line separated entries of "domain,ip" + for line in input_data.splitlines(): + parts = line.split(",") + if len(parts) != 2: + return False + ip_part = parts[1] + if not re.search(SEED_MAPPINGS["ip"]["pattern"], ip_part): + return False + return True + + def _parse(self, input_data, private_namespace, ctx): + """ + Parse the data. + + Expects ctx to be a dictionary with a 'query' key for the domain name queried. + """ + records = [] + domains = [] + ips = [] + domain_to_dnsrecord_rels = [] + dnsrecord_to_ip_rels = [] + dnsrecord_to_domain_rels = [] + + for line in input_data.splitlines(): + parts = line.split(",") + domain_part = parts[0] + ip_part = parts[1] + + # PTR record name: reversed IP + .in-addr.arpa. + ptr_name = f"{reverse_ip(ip_part)}.in-addr.arpa." + + record = { + "type": "PTR", + "name": ptr_name, + "content": domain_part, + } + records.append(record) + record_id = DNSRecord(**record).unique_id + + ips.append({"ip_address": ip_part}) + + dnsrecord_to_ip_rels.append({"source": record_id, "target": ip_part}) + + # Only create DomainName node and relationship if not a wildcard + if re.match(SEED_MAPPINGS["domain"]["pattern"], domain_part): + domains.append({"name": domain_part}) + dnsrecord_to_domain_rels.append( + {"source": record_id, "target": domain_part} + ) + + # Always create DomainName node for PTR domain and relationship + ptr_domain = ptr_name.rstrip(".") + domains.append({"name": ptr_domain}) + domain_to_dnsrecord_rels.append({"source": ptr_domain, "target": record_id}) + + node_dfs = { + DNSRecord.__primarylabel__: pd.DataFrame.from_records(records) + .drop_duplicates() + .reset_index(drop=True), + DomainName.__primarylabel__: pd.DataFrame.from_records(domains) + .drop_duplicates() + .reset_index(drop=True), + IPAddressNode.__primarylabel__: pd.DataFrame.from_records(ips) + .drop_duplicates() + .reset_index(drop=True), + } + + rel_dfs = { + DomainNameHasDNSRecord.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(domain_to_dnsrecord_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToIPAddress.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_ip_rels)[ + ["target"] + ].copy(), + }, + DNSRecordPointsToDomainName.__relationshiptype__: { + "src_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["source"] + ].copy(), + "tgt_df": pd.DataFrame.from_records(dnsrecord_to_domain_rels)[ + ["target"] + ].copy(), + }, + } + + return node_dfs, rel_dfs + + +class HackerTargetSubdomainClient(OntolocyClient): + """Client for querying the HackerTarget subdomain endpoint.""" + + def __init__(self): + super().__init__() + self.parser = HackerTargetSubdomainParser() + + def _query(self, query: str): + """Query the HackerTarget subdomain API. + + Args: + query (str): domain name to lookup + """ + api_endpoint = "https://api.hackertarget.com/hostsearch/" + response = requests.get(api_endpoint, params={"q": query}) + response.raise_for_status() + return response.text + + +class HackerTargetSubdomainEnricher(OntolocyEnricher): + seed_type = SeedTypeEnum.DOMAIN + + def __init__(self): + super().__init__() + self.client = HackerTargetSubdomainClient() + + def _generate_single_query(self, seed): + return seed diff --git a/src/ontolocy/tools/mitre_attack.py b/src/ontolocy/tools/mitre_attack.py index 4f4e634..85e39ac 100644 --- a/src/ontolocy/tools/mitre_attack.py +++ b/src/ontolocy/tools/mitre_attack.py @@ -162,7 +162,7 @@ def _stix_rels_to_df(self, stix_data): return rel_df - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: stix_json = json.loads(input_data) # diff --git a/src/ontolocy/tools/nist_csf_1.py b/src/ontolocy/tools/nist_csf_1.py index 202e788..62dcfc3 100644 --- a/src/ontolocy/tools/nist_csf_1.py +++ b/src/ontolocy/tools/nist_csf_1.py @@ -68,7 +68,9 @@ def _load_url(self, url): return response.content - def _parse(self, input_data: pd.DataFrame, private_namespace=None) -> tuple: + def _parse( + self, input_data: pd.DataFrame, private_namespace=None, ctx=None + ) -> tuple: FRAMEWORK = "NIST CSF" FRAMEWORK_VERSION = "1.1" FRAMEWORK_URL = "https://doi.org/10.6028/NIST.CSWP.04162018" diff --git a/src/ontolocy/tools/nist_csf_2.py b/src/ontolocy/tools/nist_csf_2.py index 9782294..df7f872 100644 --- a/src/ontolocy/tools/nist_csf_2.py +++ b/src/ontolocy/tools/nist_csf_2.py @@ -64,7 +64,9 @@ def _load_url(self, url): return response.content - def _parse(self, input_data: pd.DataFrame, private_namespace=None) -> tuple: + def _parse( + self, input_data: pd.DataFrame, private_namespace=None, ctx=None + ) -> tuple: FRAMEWORK = "NIST CSF" FRAMEWORK_VERSION = "2.0" FRAMEWORK_URL = "https://doi.org/10.6028/NIST.CSWP.29" diff --git a/src/ontolocy/tools/nist_sp80053_v4.py b/src/ontolocy/tools/nist_sp80053_v4.py index 932380b..d644b26 100644 --- a/src/ontolocy/tools/nist_sp80053_v4.py +++ b/src/ontolocy/tools/nist_sp80053_v4.py @@ -19,7 +19,6 @@ class NistSP80053v4Parser(OntolocyParser): rel_types = [ControlHasParentControl] def _detect(self, input_data) -> bool: - print(input_data) columns = [ "FAMILY", @@ -45,7 +44,7 @@ def _detect(self, input_data) -> bool: def _load_data(self, raw_data) -> pd.DataFrame: return pd.read_csv(StringIO(raw_data)) - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: FRAMEWORK = "NIST SP 800-53" FRAMEWORK_VERSION = "Rev. 4" FRAMEWORK_URL = "https://doi.org/10.6028/NIST.SP.800-53r4" diff --git a/src/ontolocy/tools/nist_sp80053_v5.py b/src/ontolocy/tools/nist_sp80053_v5.py index c80f398..c47c6b5 100644 --- a/src/ontolocy/tools/nist_sp80053_v5.py +++ b/src/ontolocy/tools/nist_sp80053_v5.py @@ -53,7 +53,9 @@ def _load_url(self, url): return response.content - def _parse(self, input_data: pd.DataFrame, private_namespace=None) -> tuple: + def _parse( + self, input_data: pd.DataFrame, private_namespace=None, ctx=None + ) -> tuple: FRAMEWORK = "NIST SP 800-53" FRAMEWORK_VERSION = "Rev. 5" FRAMEWORK_URL = "https://doi.org/10.6028/NIST.SP.800-53r5" diff --git a/src/ontolocy/tools/nvd.py b/src/ontolocy/tools/nvd.py index 83c4f3a..d740913 100644 --- a/src/ontolocy/tools/nvd.py +++ b/src/ontolocy/tools/nvd.py @@ -47,7 +47,7 @@ def _detect(self, input_data) -> bool: else: return False - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: node_dfs = {} rel_dfs = {} diff --git a/src/ontolocy/tools/ontolocy_enricher.py b/src/ontolocy/tools/ontolocy_enricher.py index 540bb1d..f6302c4 100644 --- a/src/ontolocy/tools/ontolocy_enricher.py +++ b/src/ontolocy/tools/ontolocy_enricher.py @@ -78,10 +78,10 @@ def query(self, query, populate=True): print(f"Failed after {self.retries} retries.") raise - self.parser.parse_data(result, populate=populate) + self.parser.parse_data(result, populate=populate, ctx={"query": query}) @abstractmethod - def _query(self, query): + def _query(self, query) -> tuple: raise NotImplementedError("Query not implemented") diff --git a/src/ontolocy/tools/ontolocy_parser.py b/src/ontolocy/tools/ontolocy_parser.py index f1c749c..408aa4a 100644 --- a/src/ontolocy/tools/ontolocy_parser.py +++ b/src/ontolocy/tools/ontolocy_parser.py @@ -60,7 +60,7 @@ def __init__( else: self.data_origin = None - def _process_data(self, input_data, private_namespace=None): + def _process_data(self, input_data, private_namespace=None, ctx=None): """ Takes input data in the form of a tuple of dicts: Nodes: indexed by node label - with corresponding df for all nodes to merge @@ -73,7 +73,9 @@ def _process_data(self, input_data, private_namespace=None): self.data_inputs.append(input_data) - node_oriented_dfs, rel_input_dfs = self._parse(input_data, private_namespace) + node_oriented_dfs, rel_input_dfs = self._parse( + input_data, private_namespace, ctx + ) # update node entries for label, df in node_oriented_dfs.items(): @@ -212,7 +214,7 @@ def populate( def detect(self, input_data) -> bool: return self._detect(input_data) - def parse_data(self, input_data, populate=True): + def parse_data(self, input_data, populate=True, ctx=None): if self.detect(input_data) is False: raise ValueError( "Detection suggests input data is not valid for this parser" @@ -224,7 +226,7 @@ def parse_data(self, input_data, populate=True): else: private_namespace = self.private_namespace - self._process_data(input_data, private_namespace) + self._process_data(input_data, private_namespace, ctx) if populate is True: self.populate() @@ -277,7 +279,7 @@ def _detect(self, input_data) -> bool: raise NotImplementedError @abstractmethod - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: """Returns a tuple where the first entry is a dictionary of node dataframes indexed by label The second entry is a dictionary of relationship dataframes indexed by relationship type Each relationship entry includes the following: diff --git a/src/ontolocy/tools/shodan.py b/src/ontolocy/tools/shodan.py index 10219f6..7bb9377 100644 --- a/src/ontolocy/tools/shodan.py +++ b/src/ontolocy/tools/shodan.py @@ -162,7 +162,7 @@ def _generate_df(self, raw_data): return pd.DataFrame.from_records(all_results) - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace, ctx) -> tuple: node_dfs = {} rel_dfs = {} diff --git a/tests/conftest.py b/tests/conftest.py index 32b4a3f..3c86276 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,7 +7,7 @@ import pytest from dotenv import load_dotenv from neontology import GraphConnection, init_neontology -from neontology.graphengines import MemgraphConfig, Neo4jConfig +from neontology.graphengines import Neo4jConfig logger = logging.getLogger(__name__) @@ -26,17 +26,6 @@ }, id="neo4j-engine", ), - pytest.param( - { - "graph_config_vars": { - "uri": "TEST_MEMGRAPH_URI", - "username": "TEST_MEMGRAPH_USER", - "password": "TEST_MEMGRAPH_PASSWORD", - }, - "graph_engine": "MEMGRAPH", - }, - id="memgraph-engine", - ), ], ) def get_graph_config(request, tmp_path_factory) -> tuple: @@ -44,7 +33,6 @@ def get_graph_config(request, tmp_path_factory) -> tuple: graph_engines = { "NEO4J": Neo4jConfig, - "MEMGRAPH": MemgraphConfig, } graph_config_vars = request.param["graph_config_vars"] @@ -53,8 +41,11 @@ def get_graph_config(request, tmp_path_factory) -> tuple: # build config using environment variables for key, value in graph_config_vars.items(): + graph_config[key] = os.getenv(value) - assert graph_config[key] is not None + assert ( + graph_config[key] is not None + ), f"Environment variable {value} is not set." graph_engine = request.param["graph_engine"] @@ -63,14 +54,18 @@ def get_graph_config(request, tmp_path_factory) -> tuple: return config -@pytest.fixture(scope="session") -def neo4j_db(get_graph_config): +@pytest.fixture( + scope="session", +) +def graph_db(request, tmp_path_factory, get_graph_config): load_dotenv() init_neontology(get_graph_config) gc = GraphConnection() + gc.change_engine(get_graph_config) + # confirm we're starting with an empty database cypher = """ MATCH (n) @@ -79,8 +74,10 @@ def neo4j_db(get_graph_config): node_count = gc.evaluate_query_single(cypher) + # most backends will return 0 + # Grand will return an empty list assert ( - node_count == 0 + not node_count ), f"Looks like there are {node_count} nodes in the database, it should be empty." yield gc @@ -94,17 +91,17 @@ def pytest_collection_modifyitems(config, items): @pytest.fixture(scope="function") -def use_graph(neo4j_db): - yield neo4j_db +def use_graph(request, graph_db): + """Fixture to use the graph database in tests.""" + yield graph_db # at the end of every individual test function, we want to empty the database cypher = """ - MATCH (n) - DETACH DELETE n + MATCH (n) DETACH DELETE n; """ - neo4j_db.evaluate_query_single(cypher) + graph_db.evaluate_query_single(cypher) @pytest.fixture @@ -115,16 +112,28 @@ def cli_runner(get_graph_config): "MemgraphConfig": "MEMGRAPH", } - runner = CliRunner( - env={ + graph_engine = graph_engine_vars[get_graph_config.__class__.__name__] + + cli_env = { + "NEONTOLOGY_ENGINE": graph_engine, + } + + if graph_engine == "NEO4J": + cli_env = { + **cli_env, "NEO4J_URI": get_graph_config.uri, "NEO4J_USERNAME": get_graph_config.username, "NEO4J_PASSWORD": get_graph_config.password, + } + + elif graph_engine == "MEMGRAPH": + cli_env = { + **cli_env, "MEMGRAPH_URI": get_graph_config.uri, "MEMGRAPH_USERNAME": get_graph_config.username, "MEMGRAPH_PASSWORD": get_graph_config.password, - "NEONTOLOGY_ENGINE": graph_engine_vars[get_graph_config.__class__.__name__], } - ) + + runner = CliRunner(env=cli_env) return runner diff --git a/tests/test_cli.py b/tests/test_cli.py index 053fa64..a95e0f3 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4,6 +4,7 @@ from ontolocy.models.actortype import actor_type_taxonomy from ontolocy.models.country import country_codes from ontolocy.models.sector import sectors +from ontolocy import MitreAttackTechnique @pytest.mark.parametrize("prompt_input", ["\n", "no\n"]) @@ -30,12 +31,7 @@ def test_mitre_attack_no_populate(prompt_input, use_graph, cli_runner): assert "Goodbye" in result.output - cypher = """ - MATCH (n) - RETURN COUNT(DISTINCT n) - """ - - assert use_graph.evaluate_query_single(cypher) == 0 + assert MitreAttackTechnique.get_count() == 0 @pytest.mark.webtest @@ -64,14 +60,7 @@ def test_mitre_attack_populate_url(use_graph, cli_runner): assert "populating the graph..." in result.output assert "ingest complete" in result.output - technique_cypher = """ - MATCH (n:MitreAttackTechnique) - WHERE (n.stix_revoked = false OR n.stix_revoked is NULL) - AND (n.attack_deprecated = false or n.attack_deprecated is NULL) - RETURN COUNT(DISTINCT n) - """ - - assert use_graph.evaluate_query_single(technique_cypher) == 625 + assert MitreAttackTechnique.get_count() == 625 def test_populate_all(use_graph, cli_runner): diff --git a/tests/test_tools/test_capec.py b/tests/test_tools/test_capec.py index 2d9dfcf..0836534 100644 --- a/tests/test_tools/test_capec.py +++ b/tests/test_tools/test_capec.py @@ -10,7 +10,7 @@ @pytest.mark.webtest @pytest.mark.parametrize("test_data", test_cases) -def test_detect(test_data): +def test_detect_capec(test_data): parser = CapecParser() response = requests.get(test_data["url"]) @@ -19,7 +19,7 @@ def test_detect(test_data): assert parser.detect(input_data) is True -def test_detect_bad(): +def test_detect_bad_capec(): parser = CapecParser() assert parser.detect("just some text") is False @@ -27,7 +27,7 @@ def test_detect_bad(): @pytest.mark.webtest @pytest.mark.parametrize("test_data", test_cases) -def test_node_parse(test_data): +def test_node_parse_capec(test_data): parser = CapecParser() parser.parse_url(test_data["url"], populate=False) @@ -43,7 +43,7 @@ def test_node_parse(test_data): @pytest.mark.slow @pytest.mark.webtest @pytest.mark.parametrize("test_data", test_cases) -def test_populate(use_graph, test_data): +def test_populate_capec(use_graph, test_data): parser = CapecParser() parser.parse_url(test_data["url"], populate=True) diff --git a/tests/test_tools/test_cisa_kev.py b/tests/test_tools/test_cisa_kev.py index 9a27e9f..fdd11d7 100644 --- a/tests/test_tools/test_cisa_kev.py +++ b/tests/test_tools/test_cisa_kev.py @@ -1,6 +1,7 @@ import pytest from ontolocy.tools import CisaKevParser +from ontolocy import CVE test_data = """cveID,vendorProject,product,vulnerabilityName,dateAdded,shortDescription,requiredAction,dueDate,knownRansomwareCampaignUse,notes,cwes CVE-2025-23006,SonicWall,"SMA1000 Appliances","SonicWall SMA1000 Appliances Deserialization Vulnerability",2025-01-24,"SonicWall SMA1000 Appliance Management Console (AMC) and Central Management Console (CMC) contain a deserialization of untrusted data vulnerability, which can enable a remote, unauthenticated attacker to execute arbitrary OS commands.","Apply mitigations per vendor instructions or discontinue use of the product if mitigations are unavailable.",2025-02-14,Unknown,"https://psirt.global.sonicwall.com/vuln-detail/SNWLID-2025-0002 ; https://nvd.nist.gov/vuln/detail/CVE-2025-23006",CWE-502 @@ -10,7 +11,7 @@ CVE-2025-21334,Microsoft,Windows,"Microsoft Windows Hyper-V NT Kernel Integration VSP Use-After-Free Vulnerability",2025-01-14,"Microsoft Windows Hyper-V NT Kernel Integration VSP contains a use-after-free vulnerability that allows a local attacker to gain SYSTEM privileges.","Apply mitigations per vendor instructions or discontinue use of the product if mitigations are unavailable.",2025-02-04,Unknown,"https://msrc.microsoft.com/update-guide/en-US/vulnerability/CVE-2025-21334 ; https://nvd.nist.gov/vuln/detail/CVE-2025-21334",CWE-416""" -def test_detect(): +def test_kev_detect(): parser = CisaKevParser() input_data = parser._load_data(test_data) @@ -18,13 +19,13 @@ def test_detect(): assert parser.detect(input_data) is True -def test_detect_bad(): +def test_kev_detect_bad(): parser = CisaKevParser() assert parser.detect("just some text") is False -def test_node_parse(): +def test_kev_node_parse(): parser = CisaKevParser() input_data = parser._load_data(test_data) @@ -39,7 +40,7 @@ def test_node_parse(): assert len(org_df.index) == 1 -def test_rel_parse(): +def test_kev_rel_parse(): parser = CisaKevParser() input_data = parser._load_data(test_data) @@ -51,20 +52,18 @@ def test_rel_parse(): assert len(org_to_cve_df.index) == 5 -def test_populate(use_graph): +def test_kev_populate(use_graph): parser = CisaKevParser() input_data = parser._load_data(test_data) parser.parse_data(input_data, populate=True) - cypher = """MATCH (:Organisation)-[:ORGANISATION_REPORTED_EXPLOITATION_OF_CVE]->(n:CVE) RETURN COUNT(DISTINCT n)""" - - assert use_graph.evaluate_query_single(cypher) == 5 + assert CVE.get_count() == 5 @pytest.mark.slow @pytest.mark.webtest -def test_populate_web(use_graph): +def test_kev_populate_web(use_graph): parser = CisaKevParser() parser.parse_url( @@ -72,12 +71,12 @@ def test_populate_web(use_graph): populate=True, ) - cypher = """MATCH (cve:CVE) WHERE cve.cve_id = 'CVE-2019-0708' RETURN COUNT(DISTINCT cve)""" + cypher = """MATCH (cve:CVE) WHERE cve.cve_id = 'CVE-2019-0708' RETURN cve""" # check for a specific CVE (BlueKeep) which should be in the dataset - assert use_graph.evaluate_query_single(cypher) == 1 + assert len(use_graph.evaluate_query(cypher).nodes) == 1 cypher = """MATCH (:Organisation)-[:ORGANISATION_REPORTED_EXPLOITATION_OF_CVE]->(n:CVE) RETURN COUNT(DISTINCT n)""" # the list is always growing, so we can't check for a specific number - assert use_graph.evaluate_query_single(cypher) > 1200 + assert CVE.get_count() > 1200 diff --git a/tests/test_tools/test_ht_dns.py b/tests/test_tools/test_ht_dns.py new file mode 100644 index 0000000..3433360 --- /dev/null +++ b/tests/test_tools/test_ht_dns.py @@ -0,0 +1,138 @@ +import json + +from ontolocy.tools.ht_dns import ( + HackerTargetDNSParser, + HackerTargetDNSEnricher, + HackerTargetDNSClient, +) + + +test_data = """{ + "A": [ + "23.192.228.80", + "23.220.75.232", + "23.215.0.136", + "23.220.75.245", + "23.215.0.138", + "23.192.228.84" + ], + "AAAA": [ + "2600:1408:ec00:36::1736:7f24", + "2600:1406:bc00:53::b81e:94c8", + "2600:1406:bc00:53::b81e:94ce", + "2600:1406:5e00:6::17ce:bc1b", + "2600:1408:ec00:36::1736:7f31", + "2600:1406:5e00:6::17ce:bc12" + ], + "MX": [ + "0 ." + ], + "NS": [ + "b.iana-servers.net.", + "a.iana-servers.net." + ], + "TXT": [ + "v=spf1 -all", + "_k2n1y4vw3qtb4skdx9e7dxt97qrmmq9" + ], + "CNAME": [], + "SOA": [ + "ns.icann.org. noc.dns.icann.org. 2025082219 7200 3600 1209600 3600" + ] +}""" + + +def test_detect_ht_dns(): + parser = HackerTargetDNSParser() + + input_data = parser._load_data(test_data) + + assert parser.detect(input_data) is True + + +def test_detect_ht_dns_bad(): + parser = HackerTargetDNSParser() + + assert parser.detect("incorrect") is False + + +def test_node_parse_ht_dns(): + parser = HackerTargetDNSParser() + + input_data = parser._load_data(test_data) + ctx = {"query": "example.com"} + parser.parse_data(input_data, ctx=ctx, populate=False) + + dnsrecord_df = parser.node_oriented_dfs["DNSRecord"] + + assert len(dnsrecord_df.index) == 18 + + domain_df = parser.node_oriented_dfs["DomainName"] + + assert len(domain_df.index) == 3 + assert domain_df.iloc[0]["name"] == "example.com" + + ip_df = parser.node_oriented_dfs["IPAddress"] + + assert len(ip_df.index) == 12 + + +def test_rel_parse_ht_dns(): + parser = HackerTargetDNSParser() + + input_data = parser._load_data(test_data) + ctx = {"query": "example.com"} + parser.parse_data(input_data, ctx=ctx, populate=False) + + domain_to_dnsrecord_df = parser.rel_input_dfs["DOMAIN_NAME_HAS_DNS_RECORD"][ + "src_df" + ] + assert len(domain_to_dnsrecord_df.index) == 18 + + dnsrecord_to_ip_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_IP_ADDRESS"][ + "src_df" + ] + assert len(dnsrecord_to_ip_df.index) == 12 + + dnsrecord_to_domain_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_DOMAIN_NAME"][ + "src_df" + ] + assert len(dnsrecord_to_domain_df.index) == 2 + + +def test_populate_ht_dns(use_graph): + parser = HackerTargetDNSParser() + + input_data = parser._load_data(test_data) + ctx = {"query": "example.com"} + parser.parse_data(input_data, ctx=ctx, populate=True) + + cypher = """MATCH (d:DomainName)-[:DOMAIN_NAME_HAS_DNS_RECORD]->(r:DNSRecord) + WHERE d.name = "example.com" + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 18 + + cypher = """MATCH (r:DNSRecord)-[:DNS_RECORD_POINTS_TO_IP_ADDRESS]->(i:IPAddress) + WHERE r.name = "example.com" + RETURN i""" + assert len(use_graph.evaluate_query(cypher).nodes) == 12 + + cypher = """MATCH (r:DNSRecord)-[:DNS_RECORD_POINTS_TO_DOMAIN_NAME]->(d:DomainName) + WHERE r.name = "example.com" + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 2 + + +def test_enrich_domain_ht_dns(monkeypatch, use_graph): + def mockreturn(self, query): + return json.loads(test_data) + + monkeypatch.setattr(HackerTargetDNSClient, "_query", mockreturn) + + enricher = HackerTargetDNSEnricher() + enricher.enrich("example.com") + + cypher = """MATCH (d:DomainName)-[:DOMAIN_NAME_HAS_DNS_RECORD]->(r:DNSRecord) + WHERE d.name = "example.com" + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 18 diff --git a/tests/test_tools/test_ht_revdns.py b/tests/test_tools/test_ht_revdns.py new file mode 100644 index 0000000..16a8d0f --- /dev/null +++ b/tests/test_tools/test_ht_revdns.py @@ -0,0 +1,93 @@ +from ontolocy.tools.ht_revdns import ( + HackerTargetReverseDNSParser, + HackerTargetReverseDNSClient, + HackerTargetReverseDNSEnricher, +) + +# Example reverse DNS data returned by the API for a domain lookup +test_data = """007.google.com,92.223.30.94 +01sin.google.com,173.234.14.228 +100-cache-blicnet.google.com,92.241.132.100 +""" + + +def test_detect_ht_revdns(): + parser = HackerTargetReverseDNSParser() + assert parser.detect(test_data) is True + + +def test_detect_ht_revdns_bad(): + parser = HackerTargetReverseDNSParser() + bad_data = "not,a,valid,line" + assert parser.detect(bad_data) is False + + +def test_node_parse_ht_revdns(): + parser = HackerTargetReverseDNSParser() + ctx = {"query": "google.com"} + parser.parse_data(test_data, ctx=ctx, populate=False) + + dnsrecord_df = parser.node_oriented_dfs["DNSRecord"] + domain_df = parser.node_oriented_dfs["DomainName"] + ip_df = parser.node_oriented_dfs["IPAddress"] + + assert len(dnsrecord_df.index) == 3 + assert len(domain_df.index) >= 3 # PTR domains + valid domains + assert len(ip_df.index) == 3 + + +def test_rel_parse_ht_revdns(): + parser = HackerTargetReverseDNSParser() + ctx = {"query": "google.com"} + parser.parse_data(test_data, ctx=ctx, populate=False) + + domain_to_dnsrecord_df = parser.rel_input_dfs["DOMAIN_NAME_HAS_DNS_RECORD"][ + "src_df" + ] + dnsrecord_to_ip_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_IP_ADDRESS"][ + "src_df" + ] + dnsrecord_to_domain_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_DOMAIN_NAME"][ + "src_df" + ] + + assert len(domain_to_dnsrecord_df.index) == 3 + assert len(dnsrecord_to_ip_df.index) == 3 + assert len(dnsrecord_to_domain_df.index) >= 2 # Only valid domains matched + + +def test_populate_ht_revdns(use_graph): + parser = HackerTargetReverseDNSParser() + ctx = {"query": "google.com"} + parser.parse_data(test_data, ctx=ctx, populate=True) + + cypher = """MATCH (d:DomainName)-[:DOMAIN_NAME_HAS_DNS_RECORD]->(r:DNSRecord) + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 3 + + cypher = """MATCH (r:DNSRecord)-[:DNS_RECORD_POINTS_TO_IP_ADDRESS]->(i:IPAddress) + RETURN i""" + assert len(use_graph.evaluate_query(cypher).nodes) == 3 + + +def test_client_query(monkeypatch): + def mockreturn(self, query): + return test_data + + monkeypatch.setattr(HackerTargetReverseDNSClient, "_query", mockreturn) + client = HackerTargetReverseDNSClient() + result = client._query("google.com") + assert "007.google.com" in result + + +def test_enrich_domain_ht_revdns(monkeypatch, use_graph): + def mockreturn(self, query): + return test_data + + monkeypatch.setattr(HackerTargetReverseDNSClient, "_query", mockreturn) + enricher = HackerTargetReverseDNSEnricher() + enricher.enrich("google.com") + + cypher = """MATCH (d:DomainName)-[:DOMAIN_NAME_HAS_DNS_RECORD]->(r:DNSRecord) + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 3 diff --git a/tests/test_tools/test_ht_subdomain.py b/tests/test_tools/test_ht_subdomain.py new file mode 100644 index 0000000..b0ad06d --- /dev/null +++ b/tests/test_tools/test_ht_subdomain.py @@ -0,0 +1,97 @@ +from ontolocy.tools.ht_subdomain import ( + HackerTargetSubdomainParser, + HackerTargetSubdomainClient, + HackerTargetSubdomainEnricher, +) + +# Example output from the API (comma separated domain,ip pairs) +test_data = """google.com,142.251.184.102 +*.evenflow-test.adz.google.com,173.194.65.129 +accounts.google.com,74.125.137.84 +""" + + +def test_detect_ht_subdomain(): + parser = HackerTargetSubdomainParser() + assert parser.detect(test_data) is True + + +def test_detect_ht_subdomain_bad(): + parser = HackerTargetSubdomainParser() + bad_data = "not,a,valid,line" + assert parser.detect(bad_data) is False + + +def test_node_parse_ht_subdomain(): + parser = HackerTargetSubdomainParser() + ctx = {"query": "google.com"} + parser.parse_data(test_data, ctx=ctx, populate=False) + + dnsrecord_df = parser.node_oriented_dfs["DNSRecord"] + domain_df = parser.node_oriented_dfs["DomainName"] + ip_df = parser.node_oriented_dfs["IPAddress"] + + # 3 records, 2 non-wildcard domains, 3 PTR domains, 3 IPs + assert len(dnsrecord_df.index) == 3 + assert "google.com" in domain_df["name"].values + assert "accounts.google.com" in domain_df["name"].values + assert "*.evenflow-test.adz.google.com" not in domain_df["name"].values + assert len(ip_df.index) == 3 + + +def test_rel_parse_ht_subdomain(): + parser = HackerTargetSubdomainParser() + ctx = {"query": "google.com"} + parser.parse_data(test_data, ctx=ctx, populate=False) + + domain_to_dnsrecord_df = parser.rel_input_dfs["DOMAIN_NAME_HAS_DNS_RECORD"][ + "src_df" + ] + dnsrecord_to_ip_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_IP_ADDRESS"][ + "src_df" + ] + dnsrecord_to_domain_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_DOMAIN_NAME"][ + "src_df" + ] + + # 3 PTR domains, 3 records, 2 record->domain relationships (no wildcard) + assert len(domain_to_dnsrecord_df.index) == 3 + assert len(dnsrecord_to_ip_df.index) == 3 + assert len(dnsrecord_to_domain_df.index) == 2 + + +def test_populate_ht_subdomain(use_graph): + parser = HackerTargetSubdomainParser() + ctx = {"query": "google.com"} + parser.parse_data(test_data, ctx=ctx, populate=True) + + cypher = """MATCH (d:DomainName)-[:DOMAIN_NAME_HAS_DNS_RECORD]->(r:DNSRecord) + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 3 + + cypher = """MATCH (r:DNSRecord)-[:DNS_RECORD_POINTS_TO_IP_ADDRESS]->(i:IPAddress) + RETURN i""" + assert len(use_graph.evaluate_query(cypher).nodes) == 3 + + +def test_client_query(monkeypatch): + def mockreturn(self, query): + return test_data + + monkeypatch.setattr(HackerTargetSubdomainClient, "_query", mockreturn) + client = HackerTargetSubdomainClient() + result = client._query("google.com") + assert "google.com" in result + + +def test_enrich_domain_ht_subdomain(monkeypatch, use_graph): + def mockreturn(self, query): + return test_data + + monkeypatch.setattr(HackerTargetSubdomainClient, "_query", mockreturn) + enricher = HackerTargetSubdomainEnricher() + enricher.enrich("google.com") + + cypher = """MATCH (d:DomainName)-[:DOMAIN_NAME_HAS_DNS_RECORD]->(r:DNSRecord) + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 3 diff --git a/tests/test_tools/test_htptrip.py b/tests/test_tools/test_htptrip.py new file mode 100644 index 0000000..86520f5 --- /dev/null +++ b/tests/test_tools/test_htptrip.py @@ -0,0 +1,97 @@ +from ontolocy import DomainName, IPAddressNode + +from ontolocy.tools.ht_ptrip import ( + HackerTargetPtrIPParser, + HackerTargetPtrIPClient, + HackerTargetPtrIPEnricher, +) + +test_data = "1.1.1.1 one.one.one.one" + + +def test_detect_ht_ptrip(): + parser = HackerTargetPtrIPParser() + + input_data = parser._load_data(test_data) + + assert parser.detect(input_data) is True + + +def test_detect_ht_ptrip_bad(): + parser = HackerTargetPtrIPParser() + + assert parser.detect("incorrect") is False + + +def test_node_parse_ht_ptrip(): + + parser = HackerTargetPtrIPParser() + + parser.parse_data(test_data, populate=False) + + dnsrecord_df = parser.node_oriented_dfs["DNSRecord"] + assert len(dnsrecord_df.index) == 1 + + assert dnsrecord_df.iloc[0]["name"] == "1.1.1.1.in-addr.arpa." + + domain_df = parser.node_oriented_dfs["DomainName"] + assert len(domain_df.index) == 2 + + ip_df = parser.node_oriented_dfs["IPAddress"] + assert len(ip_df.index) == 1 + + +def test_rel_parse_ht_ptrip(): + parser = HackerTargetPtrIPParser() + + parser.parse_data(test_data, populate=False) + + domain_to_dnsrecord_df = parser.rel_input_dfs["DOMAIN_NAME_HAS_DNS_RECORD"][ + "src_df" + ] + assert len(domain_to_dnsrecord_df.index) == 1 + + dnsrecord_to_ip_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_IP_ADDRESS"][ + "src_df" + ] + assert len(dnsrecord_to_ip_df.index) == 1 + + dnsrecord_to_domain_df = parser.rel_input_dfs["DNS_RECORD_POINTS_TO_DOMAIN_NAME"][ + "src_df" + ] + assert len(dnsrecord_to_domain_df.index) == 1 + + +def test_populate_ht_ptrip(use_graph): + parser = HackerTargetPtrIPParser() + + parser.parse_data(test_data, populate=True) + + assert DomainName.get_count() == 2 + + cypher = """MATCH (d:DomainName)-[:DOMAIN_NAME_HAS_DNS_RECORD]->(r:DNSRecord) + WHERE d.name = "1.1.1.1.in-addr.arpa" + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 1 + + cypher = """MATCH (r:DNSRecord)-[:DNS_RECORD_POINTS_TO_IP_ADDRESS]->(i:IPAddress) + WHERE i.ip_address = "1.1.1.1" + RETURN i""" + assert len(use_graph.evaluate_query(cypher).nodes) == 1 + + cypher = """MATCH (r:DNSRecord)-[:DNS_RECORD_POINTS_TO_DOMAIN_NAME]->(d:DomainName) + WHERE d.name = "one.one.one.one" + RETURN r""" + assert len(use_graph.evaluate_query(cypher).nodes) == 1 + + +def test_enrich_ip_ht_ptrip(monkeypatch, use_graph): + def mockreturn(self, query): + return test_data + + monkeypatch.setattr(HackerTargetPtrIPClient, "_query", mockreturn) + + enricher = HackerTargetPtrIPEnricher() + enricher.enrich("1.1.1.1") + + assert IPAddressNode.match("1.1.1.1") diff --git a/tests/test_tools/test_ontolocy_parser.py b/tests/test_tools/test_ontolocy_parser.py index da90556..ab26021 100644 --- a/tests/test_tools/test_ontolocy_parser.py +++ b/tests/test_tools/test_ontolocy_parser.py @@ -42,7 +42,7 @@ def _detect(self, input_data) -> bool: return True - def _parse(self, input_data, private_namespace=None) -> tuple: + def _parse(self, input_data, private_namespace=None, ctx=None) -> tuple: data = json.loads(input_data) input_df = pd.DataFrame.from_records(data["SOCKETS"]).rename(