diff --git a/flowsint-api/app/utils.py b/flowsint-api/app/utils.py index 7126bdc..46b0a5e 100644 --- a/flowsint-api/app/utils.py +++ b/flowsint-api/app/utils.py @@ -101,7 +101,7 @@ def is_root_domain(domain: str) -> bool: return False -def is_valid_number(phone: str, region: str = "FR") -> None: +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) @@ -113,6 +113,8 @@ def is_valid_number(phone: str, region: str = "FR") -> None: except NumberParseException: return False + return True + def parse_asn(asn: str) -> int: if not is_valid_asn(asn): diff --git a/flowsint-core/src/flowsint_core/core/vault.py b/flowsint-core/src/flowsint_core/core/vault.py index d842740..c48e959 100644 --- a/flowsint-core/src/flowsint_core/core/vault.py +++ b/flowsint-core/src/flowsint_core/core/vault.py @@ -25,7 +25,7 @@ class Vault(VaultProtocol): if not owner_id: raise ValueError("owner_id is required to use the vault.") self.db = db - self.owner_id = str(owner_id) + self.owner_id = owner_id self.version = "V1" def _get_master_key(self) -> bytes: @@ -52,7 +52,7 @@ class Vault(VaultProtocol): algorithm=hashes.SHA256(), length=32, salt=salt, - info=self.owner_id.encode("utf-8"), + info=str(self.owner_id).encode("utf-8"), ) master_key = self._get_master_key() return hkdf.derive(master_key) @@ -72,7 +72,7 @@ class Vault(VaultProtocol): ciphertext = aesgcm.encrypt( iv, plaintext.encode("utf-8"), - self.owner_id.encode("utf-8"), # AAD = links secret to user_id + str(self.owner_id).encode("utf-8"), # AAD = links secret to user_id ) return { @@ -92,7 +92,7 @@ class Vault(VaultProtocol): plaintext = aesgcm.decrypt( row.get("iv"), row.get("ciphertext"), - self.owner_id.encode("utf-8"), + str(self.owner_id).encode("utf-8"), ) return plaintext.decode("utf-8") diff --git a/flowsint-core/tests/vault/test_vault.py b/flowsint-core/tests/vault/test_vault.py index 5d31f9b..ef12466 100644 --- a/flowsint-core/tests/vault/test_vault.py +++ b/flowsint-core/tests/vault/test_vault.py @@ -45,7 +45,7 @@ class TestVaultInitialization: """Test successful Vault initialization.""" vault = Vault(db=mock_db, owner_id=owner_id) assert vault.db == mock_db - assert vault.owner_id == str(owner_id) + assert vault.owner_id == owner_id assert vault.version == "V1" diff --git a/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py b/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py index b58ff2c..fe5dc99 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py +++ b/flowsint-enrichers/src/flowsint_enrichers/domain/to_asn.py @@ -33,6 +33,7 @@ class DomainToAsnEnricher(Enricher): vault=vault, params=params, ) + self.domain_asn_mapping: List[tuple[Domain, ASN]] = [] @classmethod def required_params(cls) -> bool: @@ -64,6 +65,7 @@ class DomainToAsnEnricher(Enricher): 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 @@ -87,6 +89,7 @@ class DomainToAsnEnricher(Enricher): description=asn_data.get("as_name", ""), ) results.append(asn) + self.domain_asn_mapping.append((domain, asn)) Logger.info( self.sketch_id, @@ -115,8 +118,8 @@ class DomainToAsnEnricher(Enricher): self, results: List[OutputType], input_data: List[InputType] = None ) -> List[OutputType]: # Create Neo4j relationships between domains and their corresponding ASNs - if input_data and self._graph_service: - for domain, asn in zip(input_data, results): + if self._graph_service: + for domain, asn in self.domain_asn_mapping: # Create domain node self.create_node(domain) # Create ASN node diff --git a/flowsint-enrichers/src/flowsint_enrichers/domain/to_ip.py b/flowsint-enrichers/src/flowsint_enrichers/domain/to_ip.py index d583333..4157e2a 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/domain/to_ip.py +++ b/flowsint-enrichers/src/flowsint_enrichers/domain/to_ip.py @@ -14,6 +14,10 @@ class ResolveEnricher(Enricher): InputType = Domain OutputType = Ip + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.domain_ip_mapping: List[tuple[Domain, Ip]] = [] + @classmethod def name(cls) -> str: return "domain_to_ip" @@ -33,10 +37,13 @@ class ResolveEnricher(Enricher): async def scan(self, data: List[InputType]) -> List[OutputType]: results: List[OutputType] = [] + self.domain_ip_mapping = [] for d in data: try: ip = socket.gethostbyname(d.domain) - results.append(Ip(address=ip)) + ip_obj = Ip(address=ip) + results.append(ip_obj) + self.domain_ip_mapping.append((d, ip_obj)) except Exception as e: Logger.info( self.sketch_id, @@ -46,7 +53,9 @@ class ResolveEnricher(Enricher): return results def postprocess(self, results: List[OutputType], original_input: List[InputType]) -> List[OutputType]: - for domain_obj, ip_obj in zip(original_input, results): + for domain_obj, ip_obj in self.domain_ip_mapping: + if not self._graph_service: + continue self.create_node(domain_obj) self.create_node(ip_obj) self.create_relationship( diff --git a/flowsint-enrichers/src/flowsint_enrichers/email/to_gravatar.py b/flowsint-enrichers/src/flowsint_enrichers/email/to_gravatar.py index 8179892..bb0589c 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/email/to_gravatar.py +++ b/flowsint-enrichers/src/flowsint_enrichers/email/to_gravatar.py @@ -15,6 +15,10 @@ class EmailToGravatarEnricher(Enricher): InputType = Email OutputType = Gravatar + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.email_gravatar_mapping: List[tuple[Email, Gravatar]] = [] + @classmethod def name(cls) -> str: return "email_to_gravatar" @@ -29,6 +33,7 @@ class EmailToGravatarEnricher(Enricher): async def scan(self, data: List[InputType]) -> List[OutputType]: results: List[OutputType] = [] + self.email_gravatar_mapping = [] for email in data: try: @@ -72,6 +77,7 @@ class EmailToGravatarEnricher(Enricher): gravatar = Gravatar(**gravatar_data) results.append(gravatar) + self.email_gravatar_mapping.append((email, gravatar)) except Exception as e: Logger.error( @@ -85,13 +91,12 @@ class EmailToGravatarEnricher(Enricher): return results def postprocess(self, results: List[OutputType], original_input: List[InputType]) -> List[OutputType]: - for email_obj, gravatar_obj in zip(original_input, results): + for email_obj, gravatar_obj in self.email_gravatar_mapping: if not self._graph_service: continue # Create email node self.create_node(email_obj) # Create gravatar node - gravatar_key = f"{email_obj.email}_{self.sketch_id}" self.create_node(gravatar_obj) # Create relationship between email and gravatar self.create_relationship(email_obj, gravatar_obj, "HAS_GRAVATAR") diff --git a/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py b/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py index 6044626..0acb0c2 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py +++ b/flowsint-enrichers/src/flowsint_enrichers/ip/to_asn.py @@ -32,6 +32,7 @@ class IpToAsnEnricher(Enricher): vault=vault, params=params, ) + self.ip_asn_mapping: List[tuple[Ip, ASN]] = [] @classmethod def required_params(cls) -> bool: @@ -63,6 +64,7 @@ class IpToAsnEnricher(Enricher): 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 @@ -84,6 +86,7 @@ class IpToAsnEnricher(Enricher): description=asn_data.get("as_name", ""), ) results.append(asn) + self.ip_asn_mapping.append((ip, asn)) Logger.info( self.sketch_id, { @@ -110,8 +113,8 @@ class IpToAsnEnricher(Enricher): self, results: List[OutputType], input_data: List[InputType] = None ) -> List[OutputType]: # Create Neo4j relationships between IPs and their corresponding ASNs - if input_data and self._graph_service: - for ip, asn in zip(input_data, results): + if self._graph_service: + for ip, asn in self.ip_asn_mapping: # Create IP node self.create_node(ip) # Create ASN node diff --git a/flowsint-enrichers/src/flowsint_enrichers/ip/to_fraudscore.py b/flowsint-enrichers/src/flowsint_enrichers/ip/to_fraudscore.py index e8ad32f..25afa6e 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/ip/to_fraudscore.py +++ b/flowsint-enrichers/src/flowsint_enrichers/ip/to_fraudscore.py @@ -33,6 +33,7 @@ class IpToFraudScore(Enricher): vault=vault, params=params, ) + self.ip_risk_mapping: List[tuple[Ip, RiskProfile]] = [] # @classmethod # def required_params(cls) -> bool: @@ -70,6 +71,7 @@ class IpToFraudScore(Enricher): async def scan(self, data: List[InputType]) -> List[OutputType]: results: List[OutputType] = [] + self.ip_risk_mapping = [] api_username = self.get_secret("SCAMLYTICS_USERNAME", os.getenv("SCAMLYTICS_USERNAME")) api_key = self.get_secret("SCAMLYTICS_API_KEY", os.getenv("SCAMLYTICS_API_KEY")) @@ -98,7 +100,7 @@ class IpToFraudScore(Enricher): Logger.error( self.sketch_id, { - "message": f"(IpToFraudScore) Request to Scamlytics failed (status): '{response_json["status"]}'." + "message": f"(IpToFraudScore) Request to Scamlytics failed (status): '{response_json['status']}'." }, ) continue @@ -121,11 +123,11 @@ class IpToFraudScore(Enricher): if proxy_info.get(key): proxy_flags.append(label) - results.append( - RiskProfile( - entity_id=ip.address, entity_type="IP address", overall_risk_score=fraud_score, risk_level=fraud_risk, last_updated=datetime.datetime.utcnow().isoformat(), risk_factors=proxy_flags, source="Scamalytics" - ) - ) + risk_profile = RiskProfile( + entity_id=ip.address, entity_type="IP address", overall_risk_score=fraud_score, risk_level=fraud_risk, last_updated=datetime.datetime.utcnow().isoformat(), risk_factors=proxy_flags, source="Scamalytics" + ) + results.append(risk_profile) + self.ip_risk_mapping.append((ip, risk_profile)) except Exception as e: Logger.error(self.sketch_id, {"message": f"(IpToFraudScore) Exception while querying {ip.address}: {e}"}) @@ -138,16 +140,15 @@ class IpToFraudScore(Enricher): if not self._graph_service: return results - if input_data and self._graph_service: - for ip, risk_profile in zip(input_data, results): - self.create_node(ip) - self.create_node(risk_profile) + for ip, risk_profile in self.ip_risk_mapping: + self.create_node(ip) + self.create_node(risk_profile) - # Create relationship - self.create_relationship(ip, risk_profile, "HAS_RISK_PROFILE") - self.log_graph_message( - f"(IpToFraudScore) IP {ip.address} has risk level of {risk_profile.risk_level}" - ) + # Create relationship + self.create_relationship(ip, risk_profile, "HAS_RISK_PROFILE") + self.log_graph_message( + f"(IpToFraudScore) IP {ip.address} has risk level of {risk_profile.risk_level}" + ) return results diff --git a/flowsint-enrichers/src/flowsint_enrichers/phone/to_infos.py b/flowsint-enrichers/src/flowsint_enrichers/phone/to_infos.py index e9ab824..d904a29 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/phone/to_infos.py +++ b/flowsint-enrichers/src/flowsint_enrichers/phone/to_infos.py @@ -34,9 +34,8 @@ class IgnorantEnricher(Enricher): results: List[OutputType] = [] for phone_obj in data: try: - cleaned_phone = is_valid_number(phone_obj.number) - if cleaned_phone: - result = await self._perform_ignorant_research(cleaned_phone) + if is_valid_number(phone_obj.number): + result = await self._perform_ignorant_research(phone_obj.number) results.append(result) else: results.append({"number": phone_obj.number, "error": "Invalid phone number"}) diff --git a/flowsint-enrichers/src/flowsint_enrichers/utils.py b/flowsint-enrichers/src/flowsint_enrichers/utils.py index 103aa75..a7d01c8 100644 --- a/flowsint-enrichers/src/flowsint_enrichers/utils.py +++ b/flowsint-enrichers/src/flowsint_enrichers/utils.py @@ -155,7 +155,7 @@ def get_root_domain(domain: str) -> str: return domain -def is_valid_number(phone: str, region: str = "FR") -> None: +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) @@ -167,6 +167,8 @@ def is_valid_number(phone: str, region: str = "FR") -> None: except NumberParseException: return False + return True + def parse_asn(asn: str) -> int: if not is_valid_asn(asn):