From 160c0cb953c98264e86a89eedfb7de3cfd4521a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C9=B4=E1=B4=87=E1=B4=9B=E1=B4=9B=CA=8F=20=E2=9D=8B?= Date: Fri, 26 Jun 2026 15:36:43 +0000 Subject: [PATCH] refactor(enrichers): deduplicate shared logic into base classes and consolidate utils - Consolidate triplicated utility functions (13+ functions) from flowsint-api/app/utils.py and flowsint-enrichers/utils.py into flowsint-core/utils.py as single source of truth - Add get_root_domain() to flowsint-core/utils.py (was only in enrichers) - Create DehashedBase, HudsonRockBase, AsnmapBase in _shared/ module - Refactor 4 DeHashed enrichers (domain/email/ip/social) to use DehashedBase - Refactor 3 HudsonRock enrichers (email/phone/social) to use HudsonRockBase - Refactor 3 ASN enrichers (domain/ip/organization) to use AsnmapBase - Delete duplicate utils.py files from flowsint-api and flowsint-enrichers - Clean up unused imports in domain/to_root_domain.py Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- flowsint-api/app/utils.py | 313 --------------- flowsint-core/src/flowsint_core/utils.py | 46 +++ .../flowsint_enrichers/_shared/__init__.py | 5 + .../flowsint_enrichers/_shared/asnmap_base.py | 135 +++++++ .../_shared/dehashed_base.py | 172 ++++++++ .../_shared/hudsonrock_base.py | 141 +++++++ .../src/flowsint_enrichers/domain/to_asn.py | 123 +----- .../flowsint_enrichers/domain/to_dehashed.py | 146 +------ .../domain/to_root_domain.py | 4 +- .../flowsint_enrichers/email/to_dehashed.py | 146 +------ .../flowsint_enrichers/email/to_hudsonrock.py | 110 +----- .../src/flowsint_enrichers/ip/to_asn.py | 114 +----- .../src/flowsint_enrichers/ip/to_dehashed.py | 148 +------ .../flowsint_enrichers/organization/to_asn.py | 113 +----- .../flowsint_enrichers/phone/to_hudsonrock.py | 110 +----- .../flowsint_enrichers/social/to_dehashed.py | 146 +------ .../social/to_hudsonrock.py | 115 +----- .../src/flowsint_enrichers/utils.py | 367 ------------------ 18 files changed, 599 insertions(+), 1855 deletions(-) delete mode 100644 flowsint-api/app/utils.py create mode 100644 flowsint-enrichers/src/flowsint_enrichers/_shared/__init__.py create mode 100644 flowsint-enrichers/src/flowsint_enrichers/_shared/asnmap_base.py create mode 100644 flowsint-enrichers/src/flowsint_enrichers/_shared/dehashed_base.py create mode 100644 flowsint-enrichers/src/flowsint_enrichers/_shared/hudsonrock_base.py delete mode 100644 flowsint-enrichers/src/flowsint_enrichers/utils.py diff --git a/flowsint-api/app/utils.py b/flowsint-api/app/utils.py deleted file mode 100644 index 46b0a5e19..000000000 --- a/flowsint-api/app/utils.py +++ /dev/null @@ -1,313 +0,0 @@ -from urllib.parse import urlparse -import phonenumbers -import ipaddress -from phonenumbers import NumberParseException -from pydantic import TypeAdapter, BaseModel -from urllib.parse import urlparse -import re -import ssl -import socket -from typing import Dict, Any, List, Type -import inspect -from typing import Any, Dict, Type -from pydantic import BaseModel, TypeAdapter - - -def is_valid_ip(address: str) -> bool: - try: - ipaddress.ip_address(address) - return True - except ValueError: - return False - - -def is_valid_username(username: str) -> bool: - if not re.match(r"^[a-zA-Z0-9_-]{3,30}$", username): - return False - return True - - -def is_valid_email(email: str) -> bool: - if not re.match(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$", email): - return False - return True - - -def is_valid_domain(url_or_domain: str) -> str: - - try: - parsed = urlparse( - url_or_domain if "://" in url_or_domain else "http://" + url_or_domain - ) - hostname = parsed.hostname or url_or_domain - - if not hostname or "." not in hostname: - return False - - if not re.match(r"^[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$", hostname): - return False - - return True - except Exception as e: - return False - - -def is_root_domain(domain: str) -> bool: - """ - Determine if a domain is a root domain or subdomain. - - Args: - domain: The domain string to check - - Returns: - True if it's a root domain (e.g., example.com), False if it's a subdomain (e.g., sub.example.com) - """ - try: - # Remove protocol if present - if "://" in domain: - parsed = urlparse(domain) - domain = parsed.hostname or domain - - # Split by dots - parts = domain.split(".") - - # Handle common country code TLDs that have 2 parts (e.g., .co.uk, .com.au, .org.uk) - common_cc_tlds = [ - ".co.uk", - ".com.au", - ".org.uk", - ".net.uk", - ".gov.uk", - ".ac.uk", - ".co.nz", - ".com.sg", - ".co.jp", - ".co.kr", - ".com.br", - ".com.mx", - ] - - # Check if the domain ends with a common country code TLD - for cc_tld in common_cc_tlds: - if domain.endswith(cc_tld): - # For country code TLDs, we need exactly 3 parts (e.g., example.co.uk) - return len(parts) == 3 - - # For regular TLDs, a root domain has 2 parts (e.g., example.com) - # A subdomain has 3 or more parts (e.g., sub.example.com, www.sub.example.com) - return len(parts) == 2 - except Exception: - # If we can't parse it, assume it's not a root domain - return False - - -def is_valid_number(phone: str, region: str = "FR") -> bool: - """ - Validates a phone number. Raises InvalidPhoneNumberError if invalid. - - `region` should be ISO 3166-1 alpha-2 country code (e.g., 'FR' for France) - """ - try: - parsed = phonenumbers.parse(phone, region) - if not phonenumbers.is_valid_number(parsed): - return False - except NumberParseException: - return False - - return True - - -def parse_asn(asn: str) -> int: - if not is_valid_asn(asn): - raise ValueError(f"Invalid ASN format: {asn}") - return int(re.sub(r"(?i)^AS", "", asn)) - - -def is_valid_asn(asn: str) -> bool: - if not re.fullmatch(r"(AS)?\d+", asn, re.IGNORECASE): - return False - asn_num = int(re.sub(r"(?i)^AS", "", asn)) - return 0 <= asn_num <= 4294967295 - - -def resolve_type(details: dict, schema_context: dict = None) -> str: - if "anyOf" in details: - types = [] - for option in details["anyOf"]: - if "$ref" in option: - ref = option["$ref"].split("/")[-1] - types.append(ref) - elif option.get("type") == "array": - # Handle array types within anyOf - item_type = resolve_type(option.get("items", {}), schema_context) - types.append(f"{item_type}[]") - else: - types.append(option.get("type", "unknown")) - return " | ".join(types) - - if "type" in details: - if details["type"] == "array": - item_type = resolve_type(details.get("items", {}), schema_context) - return f"{item_type}[]" - return details["type"] - - # Handle $ref in array items or other contexts - if "$ref" in details and schema_context: - ref_path = details["$ref"] - if ref_path.startswith("#/$defs/"): - ref_name = ref_path.split("/")[-1] - return ref_name - - return "any" - - -def extract_input_schema_flow(model: Type[BaseModel]) -> Dict[str, Any]: - adapter = TypeAdapter(model) - schema = adapter.json_schema() - - # Use the main schema properties, not the $defs - type_name = model.__name__ - details = schema - - return { - "class_name": model.__name__, - "name": model.__name__, - "module": model.__module__, - "description": inspect.cleandoc(model.__doc__ or ""), - "outputs": { - "type": type_name, - "properties": [ - {"name": prop, "type": resolve_type(info, schema)} - for prop, info in details.get("properties", {}).items() - ], - }, - "inputs": {"type": "", "properties": []}, - "type": "type", - "category": model.__name__, - } - - -def get_domain_from_ssl(ip: str, port: int = 443) -> str | None: - try: - context = ssl.create_default_context() - with socket.create_connection((ip, port), timeout=3) as sock: - with context.wrap_socket(sock, server_hostname=ip) as ssock: - cert = ssock.getpeercert() - subject = cert.get("subject", []) - for entry in subject: - if entry[0][0] == "commonName": - return entry[0][1] - # Alternative: check subjectAltName - san = cert.get("subjectAltName", []) - for typ, val in san: - if typ == "DNS": - return val - except Exception as e: - print(f"SSL extraction failed for {ip}: {e}") - return None - - -def extract_enricher(enricher: Dict[str, Any]) -> Dict[str, Any]: - nodes = enricher["nodes"] - edges = enricher["edges"] - - input_node = next((node for node in nodes if node["data"]["type"] == "type"), None) - if not input_node: - raise ValueError("No input node found.") - input_output = input_node["data"]["outputs"] - node_lookup = {node["id"]: node for node in nodes} - - enrichers = [] - for edge in edges: - target_id = edge["target"] - source_handle = edge["sourceHandle"] - target_handle = edge["targetHandle"] - - enricher_node = node_lookup.get(target_id) - if enricher_node and enricher_node["data"]["type"] == "enricher": - enrichers.append( - { - "enricher_name": enricher_node["data"]["name"], - "module": enricher_node["data"]["module"], - "input": source_handle, - "output": target_handle, - } - ) - - return { - "input": { - "name": input_node["data"]["name"], - "outputs": input_output, - }, - "enrichers": enrichers, - "enricher_names": [enricher["enricher_name"] for enricher in enrichers], - } - - -def get_label_color(label: str) -> str: - color_map = {"subdomain": "#A5ABB6", "domain": "#68BDF6", "default": "#A5ABB6"} - - return color_map.get(label, color_map["default"]) - - -def flatten(data_dict, prefix=""): - """ - Flattens a dictionary to contain only Neo4j-compatible property values. - Neo4j supports primitive types (string, number, boolean) and arrays of those types. - Args: - data_dict (dict): Dictionary to flatten - Returns: - dict: Flattened dictionary with only Neo4j-compatible values - """ - flattened = {} - if not isinstance(data_dict, dict): - return flattened - for key, value in data_dict.items(): - if value is None: - continue - if isinstance(value, (str, int, float, bool)) or ( - isinstance(value, list) - and all(isinstance(item, (str, int, float, bool)) for item in value) - ): - key = f"{prefix}{key}" - flattened[key] = value - return flattened - - -def get_inline_relationships(nodes: List[Any], edges: List[Any]) -> List[str]: - """ - Get the inline relationships for a list of nodes and edges. - """ - relationships = [] - for edge in edges: - source = next((node for node in nodes if node["id"] == edge["source"]), None) - target = next((node for node in nodes if node["id"] == edge["target"]), None) - if source and target: - relationships.append({"source": source, "edge": edge, "target": target}) - return relationships - - -def to_json_serializable(obj): - """Convert any object to a JSON-serializable format.""" - import json - from pydantic import BaseModel - - try: - # Test if already JSON serializable - json.dumps(obj) - return obj - except (TypeError, ValueError): - # Handle common cases - if isinstance(obj, BaseModel): - # Use mode='json' to ensure all Pydantic types are properly serialized - return ( - obj.model_dump(mode="json") - if hasattr(obj, "model_dump") - else obj.dict() - ) - elif isinstance(obj, (list, tuple)): - return [to_json_serializable(item) for item in obj] - elif isinstance(obj, dict): - return {key: to_json_serializable(value) for key, value in obj.items()} - else: - # Convert anything else to string - return str(obj) diff --git a/flowsint-core/src/flowsint_core/utils.py b/flowsint-core/src/flowsint_core/utils.py index 470cf516f..76cfce078 100644 --- a/flowsint-core/src/flowsint_core/utils.py +++ b/flowsint-core/src/flowsint_core/utils.py @@ -100,6 +100,52 @@ def is_root_domain(domain: str) -> bool: return False +def get_root_domain(domain: str) -> str: + """ + Extract the root domain from a given domain string. + + Args: + domain: The domain string (can be a subdomain or root domain) + + Returns: + The root domain (e.g., "example.com" from "sub.example.com" or "www.sub.example.com") + """ + try: + if "://" in domain: + parsed = urlparse(domain) + domain = parsed.hostname or domain + + parts = domain.split(".") + + common_cc_tlds = [ + ".co.uk", + ".com.au", + ".org.uk", + ".net.uk", + ".gov.uk", + ".ac.uk", + ".co.nz", + ".com.sg", + ".co.jp", + ".co.kr", + ".com.br", + ".com.mx", + ] + + for cc_tld in common_cc_tlds: + if domain.endswith(cc_tld): + if len(parts) >= 3: + return ".".join(parts[-3:]) + return domain + + if len(parts) >= 2: + return ".".join(parts[-2:]) + + return domain + except Exception: + return domain + + def is_valid_number(phone: str, region: str = "FR") -> bool: """ Validates a phone number. Raises InvalidPhoneNumberError if invalid. diff --git a/flowsint-enrichers/src/flowsint_enrichers/_shared/__init__.py b/flowsint-enrichers/src/flowsint_enrichers/_shared/__init__.py new file mode 100644 index 000000000..4fff91115 --- /dev/null +++ b/flowsint-enrichers/src/flowsint_enrichers/_shared/__init__.py @@ -0,0 +1,5 @@ +from .dehashed_base import DehashedBase +from .hudsonrock_base import HudsonRockBase +from .asnmap_base import AsnmapBase + +__all__ = ["DehashedBase", "HudsonRockBase", "AsnmapBase"] diff --git a/flowsint-enrichers/src/flowsint_enrichers/_shared/asnmap_base.py b/flowsint-enrichers/src/flowsint_enrichers/_shared/asnmap_base.py new file mode 100644 index 000000000..0f80ba228 --- /dev/null +++ b/flowsint-enrichers/src/flowsint_enrichers/_shared/asnmap_base.py @@ -0,0 +1,135 @@ +import os +from abc import abstractmethod +from typing import Any, Dict, List, Optional + +from flowsint_core.core.enricher_base import Enricher +from flowsint_core.core.logger import Logger +from flowsint_types.asn import ASN +from tools.network.asnmap import AsnmapTool + + +class AsnmapBase(Enricher): + """Base class for enrichers that look up ASN data via asnmap. + + Subclasses define InputType, name(), category(), key(), + asnmap_type(), input_attr(), and relationship_type(). + """ + + OutputType = ASN + + def __init__( + self, + sketch_id: Optional[str] = None, + scan_id: Optional[str] = None, + vault=None, + params: Optional[Dict[str, Any]] = None, + ): + super().__init__( + sketch_id=sketch_id, + scan_id=scan_id, + params_schema=self.get_params_schema(), + vault=vault, + params=params, + ) + self._input_asn_mapping: List[tuple] = [] + + @classmethod + def required_params(cls) -> bool: + return True + + @classmethod + def get_params_schema(cls) -> List[Dict[str, Any]]: + return [ + { + "name": "PDCP_API_KEY", + "type": "vaultSecret", + "description": "The ProjectDiscovery Cloud Platform API key for asnmap.", + "required": True, + }, + ] + + @classmethod + @abstractmethod + def asnmap_type(cls) -> str: + """Type parameter for asnmap.launch(), e.g. 'domain', 'ip', 'org'.""" + + @classmethod + @abstractmethod + def input_attr(cls) -> str: + """Attribute on InputType that holds the lookup value.""" + + @classmethod + def relationship_type(cls) -> str: + """Graph relationship type between input and ASN. Override if needed.""" + return "BELONGS_TO" + + async def scan(self, data: List) -> List[ASN]: + results: List[ASN] = [] + self._input_asn_mapping = [] + asnmap = AsnmapTool() + attr = self.input_attr() + asnmap_t = self.asnmap_type() + + api_key = self.get_secret("PDCP_API_KEY", os.getenv("PDCP_API_KEY")) + + for item in data: + value = getattr(item, attr) + try: + asn_data = asnmap.launch(value, type=asnmap_t, api_key=api_key) + + if asn_data and "as_number" in asn_data: + asn_string = asn_data["as_number"] + asn_number = int(asn_string.replace("AS", "").replace("as", "")) + + asn = ASN( + number=asn_number, + name=asn_data.get("as_name", ""), + country=asn_data.get("as_country", ""), + description=asn_data.get("as_name", ""), + ) + results.append(asn) + self._input_asn_mapping.append((item, asn)) + + Logger.info( + self.sketch_id, + { + "message": f"[ASNMAP] Found AS{asn.number} ({asn.name}) for {asnmap_t} {value}" + }, + ) + else: + Logger.warn( + self.sketch_id, + { + "message": f"[ASNMAP] No ASN data or missing 'as_number' field for {asnmap_t} {value}" + }, + ) + + except Exception as e: + Logger.error( + self.sketch_id, + {"message": f"Error getting ASN for {asnmap_t} {value}: {e}"}, + ) + continue + + return results + + def postprocess( + self, results: List, input_data: List = None + ) -> List: + if not self._graph_service: + return results + + attr = self.input_attr() + rel_type = self.relationship_type() + + for input_item, asn in self._input_asn_mapping: + self.create_node(input_item) + self.create_node(asn) + self.create_relationship(input_item, asn, rel_type) + + value = getattr(input_item, attr) + self.log_graph_message( + f"{value} -> AS{asn.number} ({asn.name})" + ) + + return results diff --git a/flowsint-enrichers/src/flowsint_enrichers/_shared/dehashed_base.py b/flowsint-enrichers/src/flowsint_enrichers/_shared/dehashed_base.py new file mode 100644 index 000000000..f4efcfc2a --- /dev/null +++ b/flowsint-enrichers/src/flowsint_enrichers/_shared/dehashed_base.py @@ -0,0 +1,172 @@ +import json +import os +from abc import abstractmethod +from typing import Any, Dict, List, Optional + +import requests + +from flowsint_core.core.enricher_base import Enricher +from flowsint_core.core.logger import Logger +from flowsint_types.individual import Individual + + +class DehashedBase(Enricher): + """Base class for all DeHashed enrichers. + + Subclasses only need to define InputType, name(), category(), key(), + and get_query_field() / get_input_value(). + """ + + OutputType = Individual + + def __init__( + self, + sketch_id: Optional[str] = None, + scan_id: Optional[str] = None, + vault=None, + params: Optional[Dict[str, Any]] = None, + ): + super().__init__( + sketch_id=sketch_id, + scan_id=scan_id, + params_schema=self.get_params_schema(), + vault=vault, + params=params, + ) + + @classmethod + def get_params_schema(cls) -> List[Dict[str, Any]]: + return [ + { + "name": "DEHASHED_API_KEY", + "type": "vaultSecret", + "description": "Your Dehashed API key.", + "required": True, + } + ] + + @classmethod + @abstractmethod + def query_field(cls) -> str: + """DeHashed query field name, e.g. 'domain', 'email', 'ip_address', 'username'.""" + + @classmethod + @abstractmethod + def input_attr(cls) -> str: + """Attribute name on the InputType to extract the query value, e.g. 'domain', 'email', 'address', 'value'.""" + + def _log_prefix(self) -> str: + return self.__class__.__name__ + + async def scan(self, data: List) -> List[Individual]: + results: List[Individual] = [] + prefix = self._log_prefix() + + api_key = self.get_secret("DEHASHED_API_KEY", os.getenv("DEHASHED_API_KEY")) + + for item in data: + value = getattr(item, self.input_attr()) + try: + headers = { + "Dehashed-Api-Key": api_key, + "Content-Type": "application/json", + } + raw_data = json.dumps( + {"query": f"{self.query_field()}:{value}"} + ) + + api_request = requests.post( + "https://api.dehashed.com/v2/search", + data=raw_data, + headers=headers, + timeout=30, + ) + + if api_request.status_code != 200: + if api_request.status_code == 401: + Logger.error( + self.sketch_id, + { + "message": f"({prefix}) Enricher failed because API key does not have a valid subscription: '{value}': {api_request.text}" + }, + ) + elif api_request.status_code == 403: + Logger.error( + self.sketch_id, + { + "message": f"({prefix}) Enricher failed because request has malfunctioned (missing API_KEY): '{value}': {api_request.text}" + }, + ) + + try: + response_json = api_request.json() + except Exception as e: + Logger.error( + None, + { + "message": f"({prefix}) Failed to parse JSON for {value}: {e}" + }, + ) + continue + + dehashed_entries = response_json.get("entries", []) + if not dehashed_entries: + Logger.error( + self.sketch_id, + { + "message": f"({prefix}) Enricher failed for '{value}': {api_request.text}" + }, + ) + continue + + api_balance = response_json.get("balance") + if api_balance: + Logger.info( + self.sketch_id, + f"({prefix}) Your remaining API balance is {api_balance}.", + ) + + for entry in dehashed_entries: + entry_name = entry.get("name") + entry_dob = entry.get("dob") + results.append( + Individual( + full_name=entry_name[0] if entry_name else None, + birth_date=entry_dob[0] if entry_dob else None, + email_addresses=entry.get("email") or None, + phone_numbers=entry.get("phone") or None, + social_media_profiles=entry.get("social") or None, + ip_addresses=entry.get("ip_address") or None, + usernames=entry.get("username") or None, + ) + ) + except Exception as e: + Logger.error( + self.sketch_id, + { + "message": f"({prefix}) Exception while querying {value}: {e}" + }, + ) + + return results + + def postprocess(self, results: List, input_data: List = None) -> List: + if not self._graph_service: + return results + + prefix = self._log_prefix() + attr = self.input_attr() + + if input_data and self._graph_service: + for input_item in input_data: + for individual in results: + self.create_node(input_item) + self.create_node(individual) + self.create_relationship( + input_item, individual, "CONNECTION_WITH" + ) + self.log_graph_message( + f"({prefix}) Successfully found individual connections with {getattr(input_item, attr)}. " + ) + + return results diff --git a/flowsint-enrichers/src/flowsint_enrichers/_shared/hudsonrock_base.py b/flowsint-enrichers/src/flowsint_enrichers/_shared/hudsonrock_base.py new file mode 100644 index 000000000..94463eda7 --- /dev/null +++ b/flowsint-enrichers/src/flowsint_enrichers/_shared/hudsonrock_base.py @@ -0,0 +1,141 @@ +from abc import abstractmethod +from typing import List + +import requests + +from flowsint_core.core.enricher_base import Enricher +from flowsint_core.core.logger import Logger +from flowsint_types import Device + + +class HudsonRockBase(Enricher): + """Base class for all HudsonRock enrichers. + + Subclasses define InputType, name(), category(), key(), + api_search_type(), and input_attr(). + """ + + OutputType = Device + + @classmethod + @abstractmethod + def api_search_type(cls) -> str: + """HudsonRock search type: 'email' or 'username'.""" + + @classmethod + @abstractmethod + def input_attr(cls) -> str: + """Attribute on InputType that holds the lookup value.""" + + def _log_prefix(self) -> str: + return self.__class__.__name__ + + @staticmethod + def _parse_ip_addresses(ip_value) -> List[str]: + if isinstance(ip_value, str): + if ip_value.lower() == "not found" or not ip_value.strip(): + return [] + return [ip_value] + if isinstance(ip_value, list): + return ip_value + return [] + + async def scan(self, data: List) -> List[Device]: + results: List[Device] = [] + prefix = self._log_prefix() + search_type = self.api_search_type() + attr = self.input_attr() + + for item in data: + value = getattr(item, attr) + try: + api_request = requests.get( + f"https://cavalier.hudsonrock.com/api/json/v2/osint-tools/search-by-{search_type}?{search_type}={value}", + timeout=30, + ) + + if api_request.status_code != 200: + Logger.error( + self.sketch_id, + { + "message": f"({prefix}) failed for '{value}': {api_request.text}" + }, + ) + continue + + try: + response_json = api_request.json() + except Exception: + Logger.error( + None, + { + "message": f"({prefix}) Failed to parse JSON for {value}: {api_request.text}" + }, + ) + continue + + if response_json["total_user_services"] == 0: + Logger.error( + self.sketch_id, + { + "message": f"({prefix}) did not find anything for '{value}'." + }, + ) + continue + + for stealer_info in response_json["stealers"]: + ip_addresses = self._parse_ip_addresses( + stealer_info.get("ip") + ) + source = f"{stealer_info.get('stealer_family')} (stealer)" + + results.append( + Device( + device_id=stealer_info.get("computer_name"), + type="PC/Laptop", + os=stealer_info.get("operating_system"), + last_seen=stealer_info.get("date_compromised"), + is_desktop=True, + ip_addresses=ip_addresses, + source=source, + ) + ) + except Exception as e: + Logger.error( + self.sketch_id, + { + "message": f"({prefix}) Exception while querying {value}: {e}" + }, + ) + + return results + + def postprocess(self, results: List, original_input: List) -> List: + if not self._graph_service: + return results + + prefix = self._log_prefix() + attr = self.input_attr() + + for device in results: + try: + self.create_node(device) + + for input_item in original_input: + self.create_relationship( + input_item, device, "ASSOCIATED_WITH_DEVICE" + ) + self.log_graph_message( + f"({prefix}) {getattr(input_item, attr)} -> found device '{device.device_id}'" + ) + + except Exception as e: + Logger.error( + self.sketch_id, + { + "message": f"Failed to create graph nodes ({prefix}): {e}" + }, + ) + continue + + return results diff --git a/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py b/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py index fe5dc9925..ab6d9a28f 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py +++ b/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py @@ -1,55 +1,13 @@ -import json -import os -import socket -from typing import Any, Dict, List, Optional, Union -from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers._shared import AsnmapBase from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.domain import Domain -from flowsint_types.asn import ASN -from flowsint_core.utils import is_valid_domain -from flowsint_core.core.logger import Logger -from tools.network.asnmap import AsnmapTool @flowsint_enricher -class DomainToAsnEnricher(Enricher): +class DomainToAsnEnricher(AsnmapBase): """[ASNMAP] Takes a domain and returns its corresponding ASN.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Domain - OutputType = ASN - - def __init__( - self, - sketch_id: Optional[str] = None, - scan_id: Optional[str] = None, - vault=None, - params: Optional[Dict[str, Any]] = None, - ): - super().__init__( - sketch_id=sketch_id, - scan_id=scan_id, - params_schema=self.get_params_schema(), - vault=vault, - params=params, - ) - self.domain_asn_mapping: List[tuple[Domain, ASN]] = [] - - @classmethod - def required_params(cls) -> bool: - return True - - @classmethod - def get_params_schema(cls) -> List[Dict[str, Any]]: - """Declare required parameters for this enricher""" - return [ - { - "name": "PDCP_API_KEY", - "type": "vaultSecret", - "description": "The ProjectDiscovery Cloud Platform API key for asnmap.", - "required": True, - }, - ] @classmethod def name(cls) -> str: @@ -63,76 +21,17 @@ def category(cls) -> str: def key(cls) -> str: return "domain" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - self.domain_asn_mapping = [] - asnmap = AsnmapTool() - - # Retrieve API key from vault or environment - api_key = self.get_secret("PDCP_API_KEY", os.getenv("PDCP_API_KEY")) - - for domain in data: - try: - # Use asnmap tool to get ASN info from domain, passing the API key - asn_data = asnmap.launch(domain.domain, type="domain", api_key=api_key) - - if asn_data and "as_number" in asn_data: - # Parse ASN number from string like "AS16276" to integer 16276 - asn_string = asn_data["as_number"] - asn_number = int(asn_string.replace("AS", "").replace("as", "")) - - # Create ASN object with correct field mapping - asn = ASN( - number=asn_number, - name=asn_data.get("as_name", ""), - country=asn_data.get("as_country", ""), - description=asn_data.get("as_name", ""), - ) - results.append(asn) - self.domain_asn_mapping.append((domain, asn)) - - Logger.info( - self.sketch_id, - { - "message": f"[ASNMAP] Found AS{asn.number} ({asn.name}) for domain {domain.domain}" - }, - ) - else: - Logger.warn( - self.sketch_id, - { - "message": f"[ASNMAP] No ASN data or missing 'as_number' field for domain {domain.domain}" - }, - ) - - except Exception as e: - Logger.error( - self.sketch_id, - {"message": f"Error getting ASN for domain {domain.domain}: {e}"}, - ) - continue - - return results - - def postprocess( - self, results: List[OutputType], input_data: List[InputType] = None - ) -> List[OutputType]: - # Create Neo4j relationships between domains and their corresponding ASNs - if self._graph_service: - for domain, asn in self.domain_asn_mapping: - # Create domain node - self.create_node(domain) - # Create ASN node - self.create_node(asn) - - # Create relationship - self.create_relationship(domain, asn, "HOSTED_IN") + @classmethod + def asnmap_type(cls) -> str: + return "domain" - self.log_graph_message( - f"Domain {domain.domain} is hosted in AS{asn.number} ({asn.name})" - ) + @classmethod + def input_attr(cls) -> str: + return "domain" - return results + @classmethod + def relationship_type(cls) -> str: + return "HOSTED_IN" # Make types available at module level for easy access diff --git a/flowsint-enrichers/src/flowsint_enrichers/domain/to_dehashed.py b/flowsint-enrichers/src/flowsint_enrichers/domain/to_dehashed.py index 80fd8ba8d..ba1c96133 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/domain/to_dehashed.py +++ b/flowsint-enrichers/src/flowsint_enrichers/domain/to_dehashed.py @@ -1,52 +1,13 @@ -import json -import os -import requests - -from typing import Any, Dict, List, Optional -from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers._shared import DehashedBase from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.domain import Domain -from flowsint_types.individual import Individual -from flowsint_core.core.logger import Logger + @flowsint_enricher -class DomainToDehashed(Enricher): +class DomainToDehashed(DehashedBase): """[DeHashed] Get breach intelligence from a domain.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Domain - OutputType = Individual - - def __init__( - self, - sketch_id: Optional[str] = None, - scan_id: Optional[str] = None, - vault=None, - params: Optional[Dict[str, Any]] = None, - ): - super().__init__( - sketch_id=sketch_id, - scan_id=scan_id, - params_schema=self.get_params_schema(), - vault=vault, - params=params, - ) - - # @classmethod - # def required_params(cls) -> bool: - # return True - - @classmethod - def get_params_schema(cls) -> List[Dict[str, Any]]: - """Declare required parameters for this enricher""" - return [ - { - "name": "DEHASHED_API_KEY", # Get your API key from dehashed.com/api - "type": "vaultSecret", - "description": "Your Dehashed API key.", - "required": True, - } - ] @classmethod def name(cls) -> str: @@ -60,102 +21,15 @@ def category(cls) -> str: def key(cls) -> str: return "domain" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - - api_key = self.get_secret("DEHASHED_API_KEY", os.getenv("DEHASHED_API_KEY")) - - for domain in data: - try: - headers = {'Dehashed-Api-Key': api_key, 'Content-Type': 'application/json'} - raw_data = json.dumps({"query": f"domain:{domain.domain}"}) - - api_request = requests.post(f'https://api.dehashed.com/v2/search', data=raw_data, headers=headers, timeout=30) - - if api_request.status_code != 200: - if api_request.status_code == 401: - Logger.error( - self.sketch_id, - { - "message": f"(DomainToDehashed) Enricher failed for the domain because API key does not have a valid subscription: '{domain.domain}': {api_request.text}" - }, - ) - - elif api_request.status_code == 403: - Logger.error( - self.sketch_id, - { - "message": f"(DomainToDehashed) Enricher failed for the domain because request has malfunctioned (missing API_KEY): '{domain.domain}': {api_request.text}" - }, - ) - - try: - response_json = api_request.json() - except Exception as e: - Logger.error(None, {"message": f"(DomainToDehashed) Failed to parse JSON for {domain.domain}: {e}"}) - continue - - dehashed_entries = response_json.get("entries", []) - if not dehashed_entries: - Logger.error( - self.sketch_id, - { - "message": f"(DomainToDehashed) Enricher failed for the domain: '{domain.domain}': {api_request.text}" - }, - ) - continue - - - api_balance = response_json.get("balance") - if api_balance: - Logger.info(self.sketch_id, f'(DomainToDehashed) Your remaining API balance is {api_balance}.') - - for entry in dehashed_entries: - entry_email = entry.get("email") - entry_phone = entry.get("phone") - entry_name = entry.get("name") - entry_socialmedia = entry.get("social") - entry_ip = entry.get("ip_address") - entry_username = entry.get("username") - entry_dob = entry.get("dob") - - results.append( - Individual( - full_name=entry_name[0] if entry_name else None, - birth_date=entry_dob[0] if entry_dob else None, - email_addresses=entry_email if entry_email else None, - phone_numbers=entry_phone if entry_phone else None, - social_media_profiles=entry_socialmedia if entry_socialmedia else None, - ip_addresses=entry_ip if entry_ip else None, - usernames=entry_username if entry_username else None - ) - ) - except Exception as e: - Logger.error(self.sketch_id, {"message": f"(DomainToDehashed) Exception while querying {domain.domain}: {e}"}) - - return results - - def postprocess( - self, results: List[OutputType], input_data: List[InputType] = None - ) -> List[OutputType]: - if not self._graph_service: - return results - - if input_data and self._graph_service: - for domain in input_data: - for individual in results: - self.create_node(domain) - self.create_node(individual) - - # Create relationship - self.create_relationship(domain, individual, "CONNECTION_WITH") - self.log_graph_message( - f"(DomainToDehashed) Successfully found individual connections with {domain.domain}. " - ) + @classmethod + def query_field(cls) -> str: + return "domain" - return results + @classmethod + def input_attr(cls) -> str: + return "domain" # Make types available at module level for easy access InputType = DomainToDehashed.InputType -OutputType = DomainToDehashed.OutputType \ No newline at end of file +OutputType = DomainToDehashed.OutputType diff --git a/flowsint-enrichers/src/flowsint_enrichers/domain/to_root_domain.py b/flowsint-enrichers/src/flowsint_enrichers/domain/to_root_domain.py index eeb236fd0..3d3e56232 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/domain/to_root_domain.py +++ b/flowsint-enrichers/src/flowsint_enrichers/domain/to_root_domain.py @@ -1,5 +1,5 @@ -from typing import List, Union -from flowsint_enrichers.utils import is_valid_domain, get_root_domain +from typing import List +from flowsint_core.utils import get_root_domain from flowsint_core.core.enricher_base import Enricher from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.domain import Domain diff --git a/flowsint-enrichers/src/flowsint_enrichers/email/to_dehashed.py b/flowsint-enrichers/src/flowsint_enrichers/email/to_dehashed.py index 622e8d160..9d1f4a3fa 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/email/to_dehashed.py +++ b/flowsint-enrichers/src/flowsint_enrichers/email/to_dehashed.py @@ -1,52 +1,13 @@ -import json -import os -import requests - -from typing import Any, Dict, List, Optional -from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers._shared import DehashedBase from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.email import Email -from flowsint_types.individual import Individual -from flowsint_core.core.logger import Logger + @flowsint_enricher -class EmailToDehashed(Enricher): +class EmailToDehashed(DehashedBase): """[DeHashed] Get breach intelligence from an email address.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Email - OutputType = Individual - - def __init__( - self, - sketch_id: Optional[str] = None, - scan_id: Optional[str] = None, - vault=None, - params: Optional[Dict[str, Any]] = None, - ): - super().__init__( - sketch_id=sketch_id, - scan_id=scan_id, - params_schema=self.get_params_schema(), - vault=vault, - params=params, - ) - - # @classmethod - # def required_params(cls) -> bool: - # return True - - @classmethod - def get_params_schema(cls) -> List[Dict[str, Any]]: - """Declare required parameters for this enricher""" - return [ - { - "name": "DEHASHED_API_KEY", # Get your API key from dehashed.com/api - "type": "vaultSecret", - "description": "Your Dehashed API key.", - "required": True, - } - ] @classmethod def name(cls) -> str: @@ -60,102 +21,15 @@ def category(cls) -> str: def key(cls) -> str: return "email" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - - api_key = self.get_secret("DEHASHED_API_KEY", os.getenv("DEHASHED_API_KEY")) - - for email in data: - try: - headers = {'Dehashed-Api-Key': api_key, 'Content-Type': 'application/json'} - raw_data = json.dumps({"query": f"email:{email.email}"}) - - api_request = requests.post(f'https://api.dehashed.com/v2/search', data=raw_data, headers=headers, timeout=30) - - if api_request.status_code != 200: - if api_request.status_code == 401: - Logger.error( - self.sketch_id, - { - "message": f"(EmailToDehashed) Enricher failed for the email because API key does not have a valid subscription: '{email.email}': {api_request.text}" - }, - ) - - elif api_request.status_code == 403: - Logger.error( - self.sketch_id, - { - "message": f"(EmailToDehashed) Enricher failed for the email because request has malfunctioned (missing API_KEY): '{email.email}': {api_request.text}" - }, - ) - - try: - response_json = api_request.json() - except Exception as e: - Logger.error(None, {"message": f"(EmailToDehashed) Failed to parse JSON for {email.email}: {e}"}) - continue - - dehashed_entries = response_json.get("entries", []) - if not dehashed_entries: - Logger.error( - self.sketch_id, - { - "message": f"(EmailToDehashed) Enricher failed for the email: '{email.email}': {api_request.text}" - }, - ) - continue - - - api_balance = response_json.get("balance") - if api_balance: - Logger.info(self.sketch_id, f'(EmailToDehashed) Your remaining API balance is {api_balance}.') - - for entry in dehashed_entries: - entry_email = entry.get("email") - entry_phone = entry.get("phone") - entry_name = entry.get("name") - entry_socialmedia = entry.get("social") - entry_ip = entry.get("ip_address") - entry_username = entry.get("username") - entry_dob = entry.get("dob") - - results.append( - Individual( - full_name=entry_name[0] if entry_name else None, - birth_date=entry_dob[0] if entry_dob else None, - email_addresses=entry_email if entry_email else None, - phone_numbers=entry_phone if entry_phone else None, - social_media_profiles=entry_socialmedia if entry_socialmedia else None, - ip_addresses=entry_ip if entry_ip else None, - usernames=entry_username if entry_username else None - ) - ) - except Exception as e: - Logger.error(self.sketch_id, {"message": f"(EmailToDehashed) Exception while querying email {email.email}: {e}"}) - - return results - - def postprocess( - self, results: List[OutputType], input_data: List[InputType] = None - ) -> List[OutputType]: - if not self._graph_service: - return results - - if input_data and self._graph_service: - for email in input_data: - for individual in results: - self.create_node(email) - self.create_node(individual) - - # Create relationship - self.create_relationship(email, individual, "CONNECTION_WITH") - self.log_graph_message( - f"(EmailToDehashed) Successfully found individual connections with {email.email}. " - ) + @classmethod + def query_field(cls) -> str: + return "email" - return results + @classmethod + def input_attr(cls) -> str: + return "email" # Make types available at module level for easy access InputType = EmailToDehashed.InputType -OutputType = EmailToDehashed.OutputType \ No newline at end of file +OutputType = EmailToDehashed.OutputType diff --git a/flowsint-enrichers/src/flowsint_enrichers/email/to_hudsonrock.py b/flowsint-enrichers/src/flowsint_enrichers/email/to_hudsonrock.py index 758d96640..2c34ce56f 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/email/to_hudsonrock.py +++ b/flowsint-enrichers/src/flowsint_enrichers/email/to_hudsonrock.py @@ -1,20 +1,13 @@ -import requests -from typing import List - -from flowsint_core.core.enricher_base import Enricher -from flowsint_core.core.logger import Logger -from flowsint_types import Email, Device - +from flowsint_enrichers._shared import HudsonRockBase from flowsint_enrichers.registry import flowsint_enricher +from flowsint_types import Email @flowsint_enricher -class HudsonRockToEmail(Enricher): +class HudsonRockToEmail(HudsonRockBase): """[HudsonRock] Looks up email/s associated with devices from Infostealer related data using HudsonRock.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Email - OutputType = Device @classmethod def name(cls) -> str: @@ -28,98 +21,15 @@ def category(cls) -> str: def key(cls) -> str: return "email" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - - for email_obj in data: - email_value = email_obj.email - try: - api_request = requests.get(f'https://cavalier.hudsonrock.com/api/json/v2/osint-tools/search-by-email?email={email_value}', timeout=30) - - if api_request.status_code != 200: - Logger.error( - self.sketch_id, - { - "message": f"(HudsonRockToEmail) failed for the email: '{email_value}': {api_request.text}" - }, - ) - continue - - try: - response_json = api_request.json() - except Exception as e: - Logger.error(None, {"message": f"(HudsonRockToEmail) Failed to parse JSON for {email_value}: {api_request.text}"}) - continue - - if response_json["total_user_services"] == 0: - Logger.error( - self.sketch_id, - { - "message": f"(HudsonRockToEmail) did not find anything for the email: '{email_value}'." - }, - ) - continue - - elif response_json["total_user_services"] >= 1: - for stealer_info in response_json["stealers"]: - device_id = stealer_info.get("computer_name") - device_type = "PC/Laptop" - os = stealer_info.get("operating_system") - last_seen = stealer_info.get("date_compromised") - is_desktop = True - - ip = stealer_info.get("ip") - if isinstance(ip, str): - if ip.lower() == "not found" or not ip.strip(): - ip_addresses = [] - else: - ip_addresses = [ip] - elif isinstance(ip, list): - ip_addresses = ip - else: - ip_addresses = [] - - source = f"{stealer_info.get('stealer_family')} (stealer)" - - results.append( - Device( - device_id=device_id, type=device_type, os=os, last_seen=last_seen, is_desktop=is_desktop, - ip_addresses=ip_addresses, source=source - ) - ) - except Exception as e: - Logger.error(self.sketch_id, {"message": f"(HudsonRockToEmail) Exception while querying {email_value}: {e}"}) - - return results - - def postprocess( - self, results: List[OutputType], original_input: List[InputType] - ) -> List[OutputType]: - - if not self._graph_service: - return results - - for device in results: - try: - self.create_node(device) - - for email_obj in original_input: - self.create_relationship( - email_obj, device, "ASSOCIATED_WITH_DEVICE") - self.log_graph_message(f"(HudsonRockToEmail) {email_obj.email} -> found device '{device.device_id}'") - - except Exception as e: - Logger.error( - self.sketch_id, - { - "message": f"Failed to create graph nodes (HudsonRockToEmail): {e}" - }, - ) - continue + @classmethod + def api_search_type(cls) -> str: + return "email" - return results + @classmethod + def input_attr(cls) -> str: + return "email" # Make types available at module level for easy access InputType = HudsonRockToEmail.InputType -OutputType = HudsonRockToEmail.OutputType \ No newline at end of file +OutputType = HudsonRockToEmail.OutputType diff --git a/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py b/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py index 0acb0c2ed..9373cb90e 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py +++ b/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py @@ -1,54 +1,13 @@ -import json -import os -from typing import Any, Dict, List, Optional, Union -from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers._shared import AsnmapBase from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.ip import Ip -from flowsint_types.asn import ASN -from flowsint_core.utils import is_valid_ip -from flowsint_core.core.logger import Logger -from tools.network.asnmap import AsnmapTool @flowsint_enricher -class IpToAsnEnricher(Enricher): +class IpToAsnEnricher(AsnmapBase): """[ASNMAP] Takes an IP address and returns its corresponding ASN.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Ip - OutputType = ASN - - def __init__( - self, - sketch_id: Optional[str] = None, - scan_id: Optional[str] = None, - vault=None, - params: Optional[Dict[str, Any]] = None, - ): - super().__init__( - sketch_id=sketch_id, - scan_id=scan_id, - params_schema=self.get_params_schema(), - vault=vault, - params=params, - ) - self.ip_asn_mapping: List[tuple[Ip, ASN]] = [] - - @classmethod - def required_params(cls) -> bool: - return True - - @classmethod - def get_params_schema(cls) -> List[Dict[str, Any]]: - """Declare required parameters for this enricher""" - return [ - { - "name": "PDCP_API_KEY", - "type": "vaultSecret", - "description": "The ProjectDiscovery Cloud Platform API key for asnmap.", - "required": True, - }, - ] @classmethod def name(cls) -> str: @@ -62,70 +21,13 @@ def category(cls) -> str: def key(cls) -> str: return "address" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - self.ip_asn_mapping = [] - asnmap = AsnmapTool() - - # Retrieve API key from vault or environment - api_key = self.get_secret("PDCP_API_KEY", os.getenv("PDCP_API_KEY")) - - for ip in data: - try: - # Use asnmap tool to get ASN info, passing the API key - asn_data = asnmap.launch(ip.address, type="ip", api_key=api_key) - if asn_data and "as_number" in asn_data: - # Parse ASN number from string like "AS16276" to integer 16276 - asn_string = asn_data["as_number"] - asn_number = int(asn_string.replace("AS", "").replace("as", "")) - # Create ASN object with correct field mapping - asn = ASN( - number=asn_number, - name=asn_data.get("as_name", ""), - country=asn_data.get("as_country", ""), - description=asn_data.get("as_name", ""), - ) - results.append(asn) - self.ip_asn_mapping.append((ip, asn)) - Logger.info( - self.sketch_id, - { - "message": f"[ASNMAP] Found AS{asn.number} ({asn.name}) for IP {ip.address}" - }, - ) - else: - Logger.warn( - self.sketch_id, - { - "message": f"[ASNMAP] No ASN data or missing 'as_number' field for IP {ip.address}. Data keys: {list(asn_data.keys()) if asn_data else 'None'}" - }, - ) - except Exception as e: - Logger.error( - self.sketch_id, - {"message": f"Error getting ASN for IP {ip.address}: {e}"}, - ) - continue - - return results - - def postprocess( - self, results: List[OutputType], input_data: List[InputType] = None - ) -> List[OutputType]: - # Create Neo4j relationships between IPs and their corresponding ASNs - if self._graph_service: - for ip, asn in self.ip_asn_mapping: - # Create IP node - self.create_node(ip) - # Create ASN node - self.create_node(asn) - # Create relationship - self.create_relationship(ip, asn, "BELONGS_TO") - self.log_graph_message( - f"IP {ip.address} belongs to AS{asn.number} ({asn.name})" - ) + @classmethod + def asnmap_type(cls) -> str: + return "ip" - return results + @classmethod + def input_attr(cls) -> str: + return "address" # Make types available at module level for easy access diff --git a/flowsint-enrichers/src/flowsint_enrichers/ip/to_dehashed.py b/flowsint-enrichers/src/flowsint_enrichers/ip/to_dehashed.py index 6612f56ef..cc4b6be72 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/ip/to_dehashed.py +++ b/flowsint-enrichers/src/flowsint_enrichers/ip/to_dehashed.py @@ -1,52 +1,13 @@ -import json -import os -import requests - -from typing import Any, Dict, List, Optional -from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers._shared import DehashedBase from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.ip import Ip -from flowsint_types.individual import Individual -from flowsint_core.core.logger import Logger + @flowsint_enricher -class IpToIntelligence(Enricher): +class IpToIntelligence(DehashedBase): """[DeHashed] Get breach intelligence from an IP address.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Ip - OutputType = Individual - - def __init__( - self, - sketch_id: Optional[str] = None, - scan_id: Optional[str] = None, - vault=None, - params: Optional[Dict[str, Any]] = None, - ): - super().__init__( - sketch_id=sketch_id, - scan_id=scan_id, - params_schema=self.get_params_schema(), - vault=vault, - params=params, - ) - - # @classmethod - # def required_params(cls) -> bool: - # return True - - @classmethod - def get_params_schema(cls) -> List[Dict[str, Any]]: - """Declare required parameters for this enricher""" - return [ - { - "name": "DEHASHED_API_KEY", # Get your API key from dehashed.com/api - "type": "vaultSecret", - "description": "Your Dehashed API key.", - "required": True, - } - ] @classmethod def name(cls) -> str: @@ -60,104 +21,15 @@ def category(cls) -> str: def key(cls) -> str: return "address" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - - api_key = self.get_secret("DEHASHED_API_KEY", os.getenv("DEHASHED_API_KEY")) - - for ip in data: - try: - headers = {'Dehashed-Api-Key': api_key, 'Content-Type': 'application/json'} - raw_data = json.dumps({"query": f"ip_address:{ip.address}"}) - - api_request = requests.post(f'https://api.dehashed.com/v2/search', data=raw_data, headers=headers, timeout=30) - - if api_request.status_code != 200: - if api_request.status_code == 401: - Logger.error( - self.sketch_id, - { - "message": f"(IpToIntelligence) Enricher failed for the IP address because API key does not have a valid subscription: '{ip.address}': {api_request.text}" - }, - ) - - elif api_request.status_code == 403: - Logger.error( - self.sketch_id, - { - "message": f"(IpToIntelligence) Enricher failed for the IP address because request has malfunctioned (missing API_KEY): '{ip.address}': {api_request.text}" - }, - ) - - try: - response_json = api_request.json() - except Exception as e: - Logger.error(None, {"message": f"(IpToIntelligence) Failed to parse JSON for {ip.address}: {e}"}) - continue - - dehashed_entries = response_json.get("entries", []) - if not dehashed_entries: - Logger.error( - self.sketch_id, - { - "message": f"(IpToIntelligence) Enricher failed for the IP address: '{ip.address}': {api_request.text}" - }, - ) - continue - - - api_balance = response_json.get("balance") - if api_balance: - Logger.info(self.sketch_id, f'(EmailToDehashed) Your remaining API balance is {api_balance}.') - - time_took = response_json.get("took") - - for entry in dehashed_entries: - entry_email = entry.get("email") - entry_phone = entry.get("phone") - entry_name = entry.get("name") - entry_socialmedia = entry.get("social") - entry_ip = entry.get("ip_address") - entry_username = entry.get("username") - entry_dob = entry.get("dob") - - results.append( - Individual( - full_name=entry_name[0] if entry_name else None, - birth_date=entry_dob[0] if entry_dob else None, - email_addresses=entry_email if entry_email else None, - phone_numbers=entry_phone if entry_phone else None, - social_media_profiles=entry_socialmedia if entry_socialmedia else None, - ip_addresses=entry_ip if entry_ip else None, - usernames=entry_username if entry_username else None - ) - ) - except Exception as e: - Logger.error(self.sketch_id, {"message": f"(IpToIntelligence) Exception while querying {ip.address}: {e}"}) - - return results - - def postprocess( - self, results: List[OutputType], input_data: List[InputType] = None - ) -> List[OutputType]: - if not self._graph_service: - return results - - if input_data and self._graph_service: - for ip in input_data: - for individual in results: - self.create_node(ip) - self.create_node(individual) - - # Create relationship - self.create_relationship(ip, individual, "CONNECTION_WITH") - self.log_graph_message( - f"(IpToIntelligence) Successfully found individual connections with {ip.address}." - ) + @classmethod + def query_field(cls) -> str: + return "ip_address" - return results + @classmethod + def input_attr(cls) -> str: + return "address" # Make types available at module level for easy access InputType = IpToIntelligence.InputType -OutputType = IpToIntelligence.OutputType \ No newline at end of file +OutputType = IpToIntelligence.OutputType diff --git a/flowsint-enrichers/src/flowsint_enrichers/organization/to_asn.py b/flowsint-enrichers/src/flowsint_enrichers/organization/to_asn.py index ea26784eb..803ce3dde 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/organization/to_asn.py +++ b/flowsint-enrichers/src/flowsint_enrichers/organization/to_asn.py @@ -1,51 +1,13 @@ -import os -from typing import List, Dict, Any, Union, Optional -from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers._shared import AsnmapBase from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.organization import Organization -from flowsint_types.asn import ASN -from flowsint_core.core.logger import Logger -from tools.network.asnmap import AsnmapTool @flowsint_enricher -class OrgToAsnEnricher(Enricher): - """Takes an organization and returns its corresponding ASN.""" +class OrgToAsnEnricher(AsnmapBase): + """[ASNMAP] Takes an organization and returns its corresponding ASN.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Organization - OutputType = ASN - - def __init__( - self, - sketch_id: Optional[str] = None, - scan_id: Optional[str] = None, - vault=None, - params: Optional[Dict[str, Any]] = None, - ): - super().__init__( - sketch_id=sketch_id, - scan_id=scan_id, - params_schema=self.get_params_schema(), - vault=vault, - params=params, - ) - - @classmethod - def required_params(cls) -> bool: - return True - - @classmethod - def get_params_schema(cls) -> List[Dict[str, Any]]: - """Declare required parameters for this enricher""" - return [ - { - "name": "PDCP_API_KEY", - "type": "vaultSecret", - "description": "The ProjectDiscovery Cloud Platform API key for asnmap.", - "required": True, - }, - ] @classmethod def name(cls) -> str: @@ -59,70 +21,13 @@ def category(cls) -> str: def key(cls) -> str: return "name" - async def scan(self, data: List[InputType]) -> List[OutputType]: - """Find ASN information for organizations using asnmap.""" - results: List[OutputType] = [] - asnmap = AsnmapTool() - - # Retrieve API key from vault or environment - api_key = self.get_secret("PDCP_API_KEY", os.getenv("PDCP_API_KEY")) - - for org in data: - try: - # Use asnmap tool to get ASN info, passing the API key - asn_data = asnmap.launch(org.name, type="org", api_key=api_key) - if asn_data and "as_number" in asn_data: - # Parse ASN number from string like "AS16276" to integer 16276 - asn_string = asn_data["as_number"] - asn_number = int(asn_string.replace("AS", "").replace("as", "")) - # Create ASN object with correct field mapping - asn = ASN( - number=asn_number, - name=asn_data.get("as_name", ""), - country=asn_data.get("as_country", ""), - description=asn_data.get("as_name", ""), - ) - results.append(asn) - Logger.info( - self.sketch_id, - { - "message": f"[ASNMAP] Found AS{asn.number} ({asn.name}) for organization {org.name}" - }, - ) - else: - Logger.warn( - self.sketch_id, - { - "message": f"[ASNMAP] No ASN data or missing 'as_number' field for organization {org.name}. Data keys: {list(asn_data.keys()) if asn_data else 'None'}" - }, - ) - except Exception as e: - Logger.error( - self.sketch_id, - {"message": f"Error getting ASN for organization {org.name}: {e}"}, - ) - continue - - return results - - def postprocess(self, results: List[OutputType], original_input: List[InputType]) -> List[OutputType]: - # Create Neo4j relationships between organizations and their corresponding ASNs - for input_org, result_asn in zip(original_input, results): - # Skip if no valid ASN was found - if result_asn.number == 0: - continue - if self._graph_service: - # Create organization node - self.create_node(input_org) - # Create ASN node - self.create_node(result_asn) - # Create relationship - self.create_relationship(input_org, result_asn, "BELONGS_TO") - self.log_graph_message( - f"Found for {input_org.name} -> ASN {result_asn.number}" - ) + @classmethod + def asnmap_type(cls) -> str: + return "org" - return results + @classmethod + def input_attr(cls) -> str: + return "name" # Make types available at module level for easy access diff --git a/flowsint-enrichers/src/flowsint_enrichers/phone/to_hudsonrock.py b/flowsint-enrichers/src/flowsint_enrichers/phone/to_hudsonrock.py index 05e3f11d4..7d72e6dcf 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/phone/to_hudsonrock.py +++ b/flowsint-enrichers/src/flowsint_enrichers/phone/to_hudsonrock.py @@ -1,20 +1,13 @@ -import requests -from typing import List - -from flowsint_core.core.enricher_base import Enricher -from flowsint_core.core.logger import Logger -from flowsint_types import Phone, Device - +from flowsint_enrichers._shared import HudsonRockBase from flowsint_enrichers.registry import flowsint_enricher +from flowsint_types import Phone @flowsint_enricher -class HudsonRockToPhone(Enricher): +class HudsonRockToPhone(HudsonRockBase): """[HudsonRock] Looks up phone number/s associated with devices from Infostealer related data using HudsonRock.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Phone - OutputType = Device @classmethod def name(cls) -> str: @@ -28,98 +21,15 @@ def category(cls) -> str: def key(cls) -> str: return "number" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - - for phone_obj in data: - phonenum_value = phone_obj.number - try: - api_request = requests.get(f'https://cavalier.hudsonrock.com/api/json/v2/osint-tools/search-by-username?username={phonenum_value}', timeout=30) - - if api_request.status_code != 200: - Logger.error( - self.sketch_id, - { - "message": f"(HudsonRockToPhone) failed for the phone number: '{phonenum_value}': {api_request.text}" - }, - ) - continue - - try: - response_json = api_request.json() - except Exception as e: - Logger.error(None, {"message": f"(HudsonRockToPhone) Failed to parse JSON for {phonenum_value}: {api_request.text}"}) - continue - - if response_json["total_user_services"] == 0: - Logger.error( - self.sketch_id, - { - "message": f"(HudsonRockToPhone) did not find anything for the phone number: '{phonenum_value}'." - }, - ) - continue - - elif response_json["total_user_services"] >= 1: - for stealer_info in response_json["stealers"]: - device_id = stealer_info.get("computer_name") - device_type = "PC/Laptop" - os = stealer_info.get("operating_system") - last_seen = stealer_info.get("date_compromised") - is_desktop = True - - ip = stealer_info.get("ip") - if isinstance(ip, str): - if ip.lower() == "not found" or not ip.strip(): - ip_addresses = [] - else: - ip_addresses = [ip] - elif isinstance(ip, list): - ip_addresses = ip - else: - ip_addresses = [] - - source = f"{stealer_info.get('stealer_family')} (stealer)" - - results.append( - Device( - device_id=device_id, type=device_type, os=os, last_seen=last_seen, is_desktop=is_desktop, - ip_addresses=ip_addresses, source=source - ) - ) - except Exception as e: - Logger.error(self.sketch_id, {"message": f"(HudsonRockToPhone) Exception while querying {phonenum_value}: {e}"}) - - return results - - def postprocess( - self, results: List[OutputType], original_input: List[InputType] - ) -> List[OutputType]: - - if not self._graph_service: - return results - - for device in results: - try: - self.create_node(device) - - for phone_obj in original_input: - self.create_relationship( - phone_obj, device, "ASSOCIATED_WITH_DEVICE") - self.log_graph_message(f"(HudsonRockToPhone) {phone_obj.email} -> found device '{device.device_id}'") - - except Exception as e: - Logger.error( - self.sketch_id, - { - "message": f"Failed to create graph nodes (HudsonRockToPhone): {e}" - }, - ) - continue + @classmethod + def api_search_type(cls) -> str: + return "username" - return results + @classmethod + def input_attr(cls) -> str: + return "number" # Make types available at module level for easy access InputType = HudsonRockToPhone.InputType -OutputType = HudsonRockToPhone.OutputType \ No newline at end of file +OutputType = HudsonRockToPhone.OutputType diff --git a/flowsint-enrichers/src/flowsint_enrichers/social/to_dehashed.py b/flowsint-enrichers/src/flowsint_enrichers/social/to_dehashed.py index 006387922..20daf6c3a 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/social/to_dehashed.py +++ b/flowsint-enrichers/src/flowsint_enrichers/social/to_dehashed.py @@ -1,52 +1,13 @@ -import json -import os -import requests - -from typing import Any, Dict, List, Optional -from flowsint_core.core.enricher_base import Enricher +from flowsint_enrichers._shared import DehashedBase from flowsint_enrichers.registry import flowsint_enricher from flowsint_types.username import Username -from flowsint_types.individual import Individual -from flowsint_core.core.logger import Logger + @flowsint_enricher -class UsernameToDehashed(Enricher): +class UsernameToDehashed(DehashedBase): """[DeHashed] Get breach intelligence from a username.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Username - OutputType = Individual - - def __init__( - self, - sketch_id: Optional[str] = None, - scan_id: Optional[str] = None, - vault=None, - params: Optional[Dict[str, Any]] = None, - ): - super().__init__( - sketch_id=sketch_id, - scan_id=scan_id, - params_schema=self.get_params_schema(), - vault=vault, - params=params, - ) - - # @classmethod - # def required_params(cls) -> bool: - # return True - - @classmethod - def get_params_schema(cls) -> List[Dict[str, Any]]: - """Declare required parameters for this enricher""" - return [ - { - "name": "DEHASHED_API_KEY", # Get your API key from dehashed.com/api - "type": "vaultSecret", - "description": "Your Dehashed API key.", - "required": True, - } - ] @classmethod def name(cls) -> str: @@ -60,102 +21,15 @@ def category(cls) -> str: def key(cls) -> str: return "username" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - - api_key = self.get_secret("DEHASHED_API_KEY", os.getenv("DEHASHED_API_KEY")) - - for username in data: - try: - headers = {'Dehashed-Api-Key': api_key, 'Content-Type': 'application/json'} - raw_data = json.dumps({"query": f"username:{username.value}"}) - - api_request = requests.post(f'https://api.dehashed.com/v2/search', data=raw_data, headers=headers, timeout=30) - - if api_request.status_code != 200: - if api_request.status_code == 401: - Logger.error( - self.sketch_id, - { - "message": f"(UsernameToDehashed) Enricher failed for the username because API key does not have a valid subscription: '{username.value}': {api_request.text}" - }, - ) - - elif api_request.status_code == 403: - Logger.error( - self.sketch_id, - { - "message": f"(UsernameToDehashed) Enricher failed for the username because request has malfunctioned (missing API_KEY): '{username.value}': {api_request.text}" - }, - ) - - try: - response_json = api_request.json() - except Exception as e: - Logger.error(None, {"message": f"(UsernameToDehashed) Failed to parse JSON for {username.value}: {e}"}) - continue - - dehashed_entries = response_json.get("entries", []) - if not dehashed_entries: - Logger.error( - self.sketch_id, - { - "message": f"(UsernameToDehashed) failed for the username: '{username.value}': {api_request.text}" - }, - ) - continue - - - api_balance = response_json.get("balance") - if api_balance: - Logger.info(self.sketch_id, f'(UsernameToDehashed) Your remaining API balance is {api_balance}.') - - for entry in dehashed_entries: - entry_email = entry.get("email") - entry_phone = entry.get("phone") - entry_name = entry.get("name") - entry_socialmedia = entry.get("social") - entry_ip = entry.get("ip_address") - entry_username = entry.get("username") - entry_dob = entry.get("dob") - - results.append( - Individual( - full_name=entry_name[0] if entry_name else None, - birth_date=entry_dob[0] if entry_dob else None, - email_addresses=entry_email if entry_email else None, - phone_numbers=entry_phone if entry_phone else None, - social_media_profiles=entry_socialmedia if entry_socialmedia else None, - ip_addresses=entry_ip if entry_ip else None, - usernames=entry_username if entry_username else None - ) - ) - except Exception as e: - Logger.error(self.sketch_id, {"message": f"(UsernameToDehashed) Exception while querying username {username.value}: {e}"}) - - return results - - def postprocess( - self, results: List[OutputType], input_data: List[InputType] = None - ) -> List[OutputType]: - if not self._graph_service: - return results - - if input_data and self._graph_service: - for username in input_data: - for individual in results: - self.create_node(username) - self.create_node(individual) - - # Create relationship - self.create_relationship(username, individual, "CONNECTION_WITH") - self.log_graph_message( - f"(UsernameToDehashed) Successfully found individual connections with {username.value}. " - ) + @classmethod + def query_field(cls) -> str: + return "username" - return results + @classmethod + def input_attr(cls) -> str: + return "value" # Make types available at module level for easy access InputType = UsernameToDehashed.InputType -OutputType = UsernameToDehashed.OutputType \ No newline at end of file +OutputType = UsernameToDehashed.OutputType diff --git a/flowsint-enrichers/src/flowsint_enrichers/social/to_hudsonrock.py b/flowsint-enrichers/src/flowsint_enrichers/social/to_hudsonrock.py index 459a78210..dbda727eb 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/social/to_hudsonrock.py +++ b/flowsint-enrichers/src/flowsint_enrichers/social/to_hudsonrock.py @@ -1,24 +1,13 @@ -import subprocess -from pathlib import Path -from typing import List - -import requests -import json - -from flowsint_core.core.enricher_base import Enricher -from flowsint_core.core.logger import Logger -from flowsint_types import Device, Username - +from flowsint_enrichers._shared import HudsonRockBase from flowsint_enrichers.registry import flowsint_enricher +from flowsint_types import Username @flowsint_enricher -class HudsonRockToUsername(Enricher): +class HudsonRockToUsername(HudsonRockBase): """[HudsonRock] Looks up username/s associated with devices from Infostealer related data using HudsonRock.""" - # Define types as class attributes - base class handles schema generation automatically InputType = Username - OutputType = Device @classmethod def name(cls) -> str: @@ -32,99 +21,15 @@ def category(cls) -> str: def key(cls) -> str: return "username" - async def scan(self, data: List[InputType]) -> List[OutputType]: - results: List[OutputType] = [] - - for username in data: - try: - api_request = requests.get(f'https://cavalier.hudsonrock.com/api/json/v2/osint-tools/search-by-username?username={username.value}', timeout=30) - - if api_request.status_code != 200: - Logger.error( - self.sketch_id, - { - "message": f"(HudsonRockToUsername) failed for the username '{username.value}': {api_request.text}" - }, - ) - continue - response_json = api_request.json() - - if response_json["total_user_services"] == 0: - Logger.error( - self.sketch_id, - { - "message": f"(HudsonRockToUsername) did not find anything for the username '{username.value}'." - }, - ) - continue - - elif response_json["total_user_services"] >= 1: - for stealer_info in response_json["stealers"]: - device_id = stealer_info.get("computer_name") - device_type = "PC/Laptop" - os = stealer_info.get("operating_system") - last_seen = stealer_info.get("date_compromised") - is_desktop = True - - ip = stealer_info.get("ip") - if isinstance(ip, str): - if ip.lower() == "not found" or not ip.strip(): - ip_addresses = [] - else: - ip_addresses = [ip] - elif isinstance(ip, list): - ip_addresses = ip - else: - ip_addresses = [] - - source = f"{stealer_info.get('stealer_family')} (stealer)" - - results.append( - Device( - device_id=device_id, type=device_type, os=os, last_seen=last_seen, is_desktop=is_desktop, - ip_addresses=ip_addresses, source=source - ) - ) - except Exception as e: - Logger.error( - self.sketch_id, - { - "message": f"Failed to run HudsonRockToUsername for {username.value} (error): {e}" - }, - ) - - return results - - def postprocess( - self, results: List[OutputType], original_input: List[InputType] - ) -> List[OutputType]: - - if not self._graph_service: - return results - - for device in results: - try: - self.create_node(device) - - for username in original_input: - self.create_relationship( - username, device, "ASSOCIATED_WITH_DEVICE" - ) - self.log_graph_message( - f"{username.value} -> found device '{device.device_id}'" - ) - except Exception as e: - Logger.error( - self.sketch_id, - { - "message": f"Failed to create graph nodes (HudsonRockToUsername): {e}" - }, - ) - continue + @classmethod + def api_search_type(cls) -> str: + return "username" - return results + @classmethod + def input_attr(cls) -> str: + return "value" # Make types available at module level for easy access InputType = HudsonRockToUsername.InputType -OutputType = HudsonRockToUsername.OutputType \ No newline at end of file +OutputType = HudsonRockToUsername.OutputType diff --git a/flowsint-enrichers/src/flowsint_enrichers/utils.py b/flowsint-enrichers/src/flowsint_enrichers/utils.py deleted file mode 100644 index a7d01c821..000000000 --- a/flowsint-enrichers/src/flowsint_enrichers/utils.py +++ /dev/null @@ -1,367 +0,0 @@ -from urllib.parse import urlparse -import phonenumbers -import ipaddress -from phonenumbers import NumberParseException -from pydantic import TypeAdapter, BaseModel -from urllib.parse import urlparse -import re -import ssl -import socket -from typing import Dict, Any, List, Type -from pydantic import BaseModel -import inspect -from typing import Any, Dict, Type -from pydantic import BaseModel, TypeAdapter - - -def is_valid_ip(address: str) -> bool: - try: - ipaddress.ip_address(address) - return True - except ValueError: - return False - - -def is_valid_username(username: str) -> bool: - if not re.match(r"^[a-zA-Z0-9_-]{3,30}$", username): - return False - return True - - -def is_valid_email(email: str) -> bool: - if not re.match(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$", email): - return False - return True - - -def is_valid_domain(url_or_domain: str) -> str: - - try: - parsed = urlparse( - url_or_domain if "://" in url_or_domain else "http://" + url_or_domain - ) - hostname = parsed.hostname or url_or_domain - - if not hostname or "." not in hostname: - return False - - if not re.match(r"^[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$", hostname): - return False - - return True - except Exception as e: - return False - - -def is_root_domain(domain: str) -> bool: - """ - Determine if a domain is a root domain or subdomain. - - Args: - domain: The domain string to check - - Returns: - True if it's a root domain (e.g., example.com), False if it's a subdomain (e.g., sub.example.com) - """ - try: - # Remove protocol if present - if "://" in domain: - parsed = urlparse(domain) - domain = parsed.hostname or domain - - # Split by dots - parts = domain.split(".") - - # Handle common country code TLDs that have 2 parts (e.g., .co.uk, .com.au, .org.uk) - common_cc_tlds = [ - ".co.uk", - ".com.au", - ".org.uk", - ".net.uk", - ".gov.uk", - ".ac.uk", - ".co.nz", - ".com.sg", - ".co.jp", - ".co.kr", - ".com.br", - ".com.mx", - ] - - # Check if the domain ends with a common country code TLD - for cc_tld in common_cc_tlds: - if domain.endswith(cc_tld): - # For country code TLDs, we need exactly 3 parts (e.g., example.co.uk) - return len(parts) == 3 - - # For regular TLDs, a root domain has 2 parts (e.g., example.com) - # A subdomain has 3 or more parts (e.g., sub.example.com, www.sub.example.com) - return len(parts) == 2 - except Exception: - # If we can't parse it, assume it's not a root domain - return False - - -def get_root_domain(domain: str) -> str: - """ - Extract the root domain from a given domain string. - - Args: - domain: The domain string (can be a subdomain or root domain) - - Returns: - The root domain (e.g., "example.com" from "sub.example.com" or "www.sub.example.com") - """ - try: - # Remove protocol if present - if "://" in domain: - parsed = urlparse(domain) - domain = parsed.hostname or domain - - # Split by dots - parts = domain.split(".") - - # Handle common country code TLDs that have 2 parts (e.g., .co.uk, .com.au, .org.uk) - common_cc_tlds = [ - ".co.uk", - ".com.au", - ".org.uk", - ".net.uk", - ".gov.uk", - ".ac.uk", - ".co.nz", - ".com.sg", - ".co.jp", - ".co.kr", - ".com.br", - ".com.mx", - ] - - # Check if the domain ends with a common country code TLD - for cc_tld in common_cc_tlds: - if domain.endswith(cc_tld): - # For country code TLDs, take the last 3 parts (e.g., example.co.uk) - if len(parts) >= 3: - return ".".join(parts[-3:]) - return domain - - # For regular TLDs, take the last 2 parts (e.g., example.com) - if len(parts) >= 2: - return ".".join(parts[-2:]) - - return domain - except Exception: - # If we can't parse it, return the original domain - return domain - - -def is_valid_number(phone: str, region: str = "FR") -> bool: - """ - Validates a phone number. Raises InvalidPhoneNumberError if invalid. - - `region` should be ISO 3166-1 alpha-2 country code (e.g., 'FR' for France) - """ - try: - parsed = phonenumbers.parse(phone, region) - if not phonenumbers.is_valid_number(parsed): - return False - except NumberParseException: - return False - - return True - - -def parse_asn(asn: str) -> int: - if not is_valid_asn(asn): - raise ValueError(f"Invalid ASN format: {asn}") - return int(re.sub(r"(?i)^AS", "", asn)) - - -def is_valid_asn(asn: str) -> bool: - if not re.fullmatch(r"(AS)?\d+", asn, re.IGNORECASE): - return False - asn_num = int(re.sub(r"(?i)^AS", "", asn)) - return 0 <= asn_num <= 4294967295 - - -def resolve_type(details: dict, schema_context: dict = None) -> str: - if "anyOf" in details: - types = [] - for option in details["anyOf"]: - if "$ref" in option: - ref = option["$ref"].split("/")[-1] - types.append(ref) - elif option.get("type") == "array": - # Handle array types within anyOf - item_type = resolve_type(option.get("items", {}), schema_context) - types.append(f"{item_type}[]") - else: - types.append(option.get("type", "unknown")) - return " | ".join(types) - - if "type" in details: - if details["type"] == "array": - item_type = resolve_type(details.get("items", {}), schema_context) - return f"{item_type}[]" - return details["type"] - - # Handle $ref in array items or other contexts - if "$ref" in details and schema_context: - ref_path = details["$ref"] - if ref_path.startswith("#/$defs/"): - ref_name = ref_path.split("/")[-1] - return ref_name - - return "any" - - -def extract_input_schema_flow(model: Type[BaseModel]) -> Dict[str, Any]: - adapter = TypeAdapter(model) - schema = adapter.json_schema() - - # Use the main schema properties, not the $defs - type_name = model.__name__ - details = schema - - return { - "class_name": model.__name__, - "name": model.__name__, - "module": model.__module__, - "description": inspect.cleandoc(model.__doc__ or ""), - "outputs": { - "type": type_name, - "properties": [ - {"name": prop, "type": resolve_type(info, schema)} - for prop, info in details.get("properties", {}).items() - ], - }, - "inputs": {"type": "", "properties": []}, - "type": "type", - "category": model.__name__, - } - - -def get_domain_from_ssl(ip: str, port: int = 443) -> str | None: - try: - context = ssl.create_default_context() - with socket.create_connection((ip, port), timeout=3) as sock: - with context.wrap_socket(sock, server_hostname=ip) as ssock: - cert = ssock.getpeercert() - subject = cert.get("subject", []) - for entry in subject: - if entry[0][0] == "commonName": - return entry[0][1] - # Alternative: check subjectAltName - san = cert.get("subjectAltName", []) - for typ, val in san: - if typ == "DNS": - return val - except Exception as e: - print(f"SSL extraction failed for {ip}: {e}") - return None - - -def extract_enricher(enricher: Dict[str, Any]) -> Dict[str, Any]: - nodes = enricher["nodes"] - edges = enricher["edges"] - - input_node = next((node for node in nodes if node["data"]["type"] == "type"), None) - if not input_node: - raise ValueError("No input node found.") - input_output = input_node["data"]["outputs"] - node_lookup = {node["id"]: node for node in nodes} - - enrichers = [] - for edge in edges: - target_id = edge["target"] - source_handle = edge["sourceHandle"] - target_handle = edge["targetHandle"] - - enricher_node = node_lookup.get(target_id) - if enricher_node and enricher_node["data"]["type"] == "enricher": - enrichers.append( - { - "enricher_name": enricher_node["data"]["name"], - "module": enricher_node["data"]["module"], - "input": source_handle, - "output": target_handle, - } - ) - - return { - "input": { - "name": input_node["data"]["name"], - "outputs": input_output, - }, - "enrichers": enrichers, - "enricher_names": [enricher["enricher_name"] for enricher in enrichers], - } - - -def get_label_color(label: str) -> str: - color_map = {"subdomain": "#A5ABB6", "domain": "#68BDF6", "default": "#A5ABB6"} - - return color_map.get(label, color_map["default"]) - - -def flatten(data_dict, prefix=""): - """ - Flattens a dictionary to contain only Neo4j-compatible property values. - Neo4j supports primitive types (string, number, boolean) and arrays of those types. - Args: - data_dict (dict): Dictionary to flatten - Returns: - dict: Flattened dictionary with only Neo4j-compatible values - """ - flattened = {} - if not isinstance(data_dict, dict): - return flattened - for key, value in data_dict.items(): - if value is None: - continue - if isinstance(value, (str, int, float, bool)) or ( - isinstance(value, list) - and all(isinstance(item, (str, int, float, bool)) for item in value) - ): - key = f"{prefix}{key}" - flattened[key] = value - return flattened - - -def get_inline_relationships(nodes: List[Any], edges: List[Any]) -> List[str]: - """ - Get the inline relationships for a list of nodes and edges. - """ - relationships = [] - for edge in edges: - source = next((node for node in nodes if node["id"] == edge["source"]), None) - target = next((node for node in nodes if node["id"] == edge["target"]), None) - if source and target: - relationships.append({"source": source, "edge": edge, "target": target}) - return relationships - - -def to_json_serializable(obj): - """Convert any object to a JSON-serializable format.""" - import json - from pydantic import BaseModel - - try: - # Test if already JSON serializable - json.dumps(obj) - return obj - except (TypeError, ValueError): - # Handle common cases - if isinstance(obj, BaseModel): - # Use mode='json' to ensure all Pydantic types are properly serialized - return ( - obj.model_dump(mode="json") - if hasattr(obj, "model_dump") - else obj.dict() - ) - elif isinstance(obj, (list, tuple)): - return [to_json_serializable(item) for item in obj] - elif isinstance(obj, dict): - return {key: to_json_serializable(value) for key, value in obj.items()} - else: - # Convert anything else to string - return str(obj)