diff --git a/flowsint-core/src/flowsint_core/templates/loader/yaml_loader.py b/flowsint-core/src/flowsint_core/templates/loader/yaml_loader.py new file mode 100644 index 00000000..325539f9 --- /dev/null +++ b/flowsint-core/src/flowsint_core/templates/loader/yaml_loader.py @@ -0,0 +1,58 @@ +import re +from typing import Any + +import yaml +from flowsint_types import TYPE_REGISTRY + +from flowsint_core.templates.types import Template + +TEMPLATE_RE = re.compile(r"\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}") + + +class YamlLoader: + @staticmethod + def load_enricher_yaml(filename: str) -> dict[str, Any] | yaml.YAMLError: + with open(filename) as stream: + try: + return yaml.safe_load(stream) + except yaml.YAMLError as exc: + return exc + + @staticmethod + def parse_yaml_to_template(raw: dict[str, Any]) -> Template: + if not isinstance(raw, dict): + return + input: dict | None = raw.get("input", None) + + if not input or not isinstance(input, dict): + raise TypeError("Missing 'input' property in the yaml.") + + input_type: str | None = input.get("type", None) + + if not input_type: + raise TypeError("Missing 'input_type' property in the yaml.") + + DetectedType = TYPE_REGISTRY.get(input_type) + + if not DetectedType: + raise TypeError(f"Type '{input_type}' not found in registry.") + + return Template(**raw) + + @staticmethod + def get_template_from_file(filename: str) -> Template | None: + template_dict = YamlLoader.load_enricher_yaml(filename) + if not isinstance(template_dict, dict): + return + + return YamlLoader.parse_yaml_to_template(template_dict) + + @staticmethod + def render_template(template: str, values: dict[str, str]) -> str: + def replace(match: re.Match) -> str: + key = match.group(1) + if key not in values: + raise ValueError(f"Missing variable: {key}") + return values[key] + + return TEMPLATE_RE.sub(replace, template) diff --git a/flowsint-core/src/flowsint_core/templates/types.py b/flowsint-core/src/flowsint_core/templates/types.py new file mode 100644 index 00000000..bfeaab30 --- /dev/null +++ b/flowsint-core/src/flowsint_core/templates/types.py @@ -0,0 +1,85 @@ +from typing import Literal, Optional + +from pydantic import BaseModel, Field + + +class TemplateInput(BaseModel): + type: str = Field(..., description="Flowsint Type to the template takes as input") + key: str = Field( + default="nodeLabel", + description="Flowsint Type to the template takes as input", + ) + + +class TemplateHttpRequestHeader(BaseModel): + name: str = Field(min_length=1, max_length=256, pattern=r"^[A-Za-z0-9\-]+$") + value: str = Field(min_length=1, max_length=4096) + + class Config: + extra = "forbid" + frozen = True + + +class TemplateHttpRequestParams(BaseModel): + key: str = Field(min_length=1, max_length=256, pattern=r"^[A-Za-z0-9\-]+$") + value: str = Field(min_length=1, max_length=4096) + + class Config: + extra = "forbid" + frozen = True + + +class TemplateHttpRequest(BaseModel): + method: Literal["GET", "POST"] + url: str + headers: dict = Field(default_factory=dict) + params: dict = Field(default_factory=dict) + body: Optional[str] = None + + class Config: + extra = "forbid" + frozen = True + + +class TemplateHttpResponseMapping(BaseModel): + key: str = Field( + min_length=1, + max_length=256, + pattern=r"^[A-Za-z0-9\-]+$", + description="The key (from the response format) to map.", + ) + value: str = Field( + min_length=1, + max_length=4096, + description="The key of the field you want to feed (of the expected FlowsintType).", + ) + + class Config: + extra = "forbid" + frozen = True + + +class TemplateHttpResponse(BaseModel): + expect: Literal["json", "xml", "text"] + map: dict = Field(default_factory=dict) + + class Config: + extra = "forbid" + frozen = True + + +class Template(BaseModel): + name: str = Field(..., description="Name of the template") + category: str = Field(..., description="Category of the template") + version: float = Field(..., description="Version of the template") + input: TemplateInput = Field( + ..., + description="Input format of the template, with key to use (default to nodeLabel)", + ) + request: TemplateHttpRequest = Field( + ..., description="Request model for the HTTP request to be made." + ) + + response: TemplateHttpResponse = Field( + ..., description="Response model for the HTTP response to expect." + ) diff --git a/flowsint-core/tests/templates/example.yaml b/flowsint-core/tests/templates/example.yaml new file mode 100644 index 00000000..44e2e70e --- /dev/null +++ b/flowsint-core/tests/templates/example.yaml @@ -0,0 +1,24 @@ +name: shodan_ip_lookup +category: Ip +version: 1.0 +type: enricher +input: + type: Ip + key: address +request: + method: GET + url: http://ip-api.com/json/{{address}} + headers: + Authorization: "Bearer {{secrets.SHODAN_API_KEY}}" + +output: + type: Ip + +response: + expect: json + map: + lat: Ip.latitute + lon: Ip.longitude + country: Ip.country + city: Ip.city + isp: Ip.isp diff --git a/flowsint-core/tests/templates/loader.py b/flowsint-core/tests/templates/loader.py new file mode 100644 index 00000000..e25ba6d8 --- /dev/null +++ b/flowsint-core/tests/templates/loader.py @@ -0,0 +1,23 @@ +import json + +from flowsint_core.templates.loader.yaml_loader import YamlLoader +from flowsint_core.templates.types import Template + + +def test_yaml_loader(): + file = YamlLoader.get_template_from_file( + "/Users/eliottmorcillo/Documents/Perso/projets/reconurge/flowsint/flowsint-core/tests/templates/example.yaml" + ) + assert isinstance(file, Template) + assert file.name == "shodan_ip_lookup" + assert file.category == "Ip" + + print(json.dumps(file.model_dump(), indent=2, ensure_ascii=False)) + + +def test_render_template(): + url_template = "http://ip-api.com/json/{{address}}" + + url = YamlLoader.render_template(url_template, {"address": "8.8.8.8"}) + + assert url == "http://ip-api.com/json/8.8.8.8"