Merge pull request #172 from melihdaskiran/fix/zip-and-vault-bugs

Fix zip alignment mismatch, phone validation return value, and vault …
This commit is contained in:
dextmorgn
2026-06-05 08:24:39 +02:00
committed by GitHub
10 changed files with 57 additions and 33 deletions

View File

@@ -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):

View File

@@ -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")

View File

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

View File

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

View File

@@ -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(

View File

@@ -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")

View File

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

View File

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

View File

@@ -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"})

View File

@@ -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):