build out tool support

.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

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}

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

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",

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
 

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

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 = {}
 

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 = {}
 

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 = {}
 

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 = {}
 

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