mirror of
https://github.com/reconurge/flowsint.git
synced 2026-07-16 09:02:57 -05:00
72 lines
2.0 KiB
Python
72 lines
2.0 KiB
Python
from typing import List
|
|
from flowsint_core.core.enricher_base import Enricher
|
|
from flowsint_types import Phone
|
|
from flowsint_types.domain import Domain
|
|
from flowsint_types.ip import Ip
|
|
|
|
|
|
class ResolveEnricher(Enricher):
|
|
InputType = List[Domain]
|
|
OutputType = List[Ip]
|
|
|
|
@classmethod
|
|
def name(cls) -> str:
|
|
return "domain_to_ip"
|
|
|
|
@classmethod
|
|
def category(cls) -> str:
|
|
return "Domain"
|
|
|
|
@classmethod
|
|
def key(cls) -> str:
|
|
return "domain"
|
|
|
|
async def scan(self, data: InputType) -> OutputType:
|
|
return []
|
|
|
|
def postprocess(self, results: OutputType, original_input: InputType) -> OutputType:
|
|
return []
|
|
|
|
|
|
# Make types available at module level for easy access
|
|
InputType = ResolveEnricher.InputType
|
|
OutputType = ResolveEnricher.OutputType
|
|
|
|
enricher = ResolveEnricher("sketch_123", "scan_123")
|
|
|
|
|
|
def test_correct_preprocess():
|
|
inputs = [
|
|
Domain(domain="mydomain.com"),
|
|
{"domain": "blog.mydomain2.com"},
|
|
"mydomain3.com",
|
|
"notADomaiN",
|
|
]
|
|
preprocessed = enricher.preprocess(inputs)
|
|
assert len(preprocessed) == 3 # 3 valid domains
|
|
assert preprocessed[0].domain == "mydomain.com"
|
|
assert preprocessed[0].label == "mydomain.com"
|
|
assert preprocessed[0].root == True
|
|
|
|
assert preprocessed[1].domain == "blog.mydomain2.com"
|
|
assert preprocessed[1].label == "blog.mydomain2.com"
|
|
assert preprocessed[1].root == False
|
|
|
|
assert preprocessed[2].domain == "mydomain3.com"
|
|
assert preprocessed[2].label == "mydomain3.com"
|
|
assert preprocessed[2].root == True
|
|
|
|
|
|
def test_incorrect_preprocess():
|
|
inputs = [
|
|
Phone(number="+33634565423"),
|
|
{"name": "JohnDoe"},
|
|
"mydomain.com",
|
|
"notADomaiN",
|
|
]
|
|
preprocessed = enricher.preprocess(inputs)
|
|
assert len(preprocessed) == 1 # 1 valid domain
|
|
assert preprocessed[0].domain == "mydomain.com"
|
|
assert preprocessed[0].root == True
|
|
assert preprocessed[0].label == "mydomain.com"
|