diff --git a/flowsint-core/src/flowsint_core/core/enricher_base.py b/flowsint-core/src/flowsint_core/core/enricher_base.py index 0566954f..ed34a46c 100644 --- a/flowsint-core/src/flowsint_core/core/enricher_base.py +++ b/flowsint-core/src/flowsint_core/core/enricher_base.py @@ -374,7 +374,6 @@ class Enricher(ABC): base_type = self.InputType adapter = TypeAdapter(base_type) - # Trouver le champ primaire marqué par Field(..., primary=True) primary_field = None if issubclass(base_type, BaseModel): for name, field in base_type.model_fields.items(): @@ -402,6 +401,14 @@ class Enricher(ABC): except Exception: continue + if len(cleaned) == 0: + Logger.warn( + self.sketch_id, + { + "message": f"No valid input were provided to enricher '{self.name()}'." + }, + ) + return values return cleaned def postprocess( diff --git a/flowsint-core/src/flowsint_core/core/graph/__init__.py b/flowsint-core/src/flowsint_core/core/graph/__init__.py index 38412b3d..e87db742 100644 --- a/flowsint-core/src/flowsint_core/core/graph/__init__.py +++ b/flowsint-core/src/flowsint_core/core/graph/__init__.py @@ -11,7 +11,7 @@ This module provides all graph-related functionality including: from .connection import Neo4jConnection, neo4j_connection from .repository import Neo4jGraphRepository from .repository_protocol import GraphRepositoryProtocol -from .serializer import GraphSerializer +from .serializer import GraphSerializer, TypeResolver from .service import GraphService, LoggerProtocol, create_graph_service from .types import ( GraphData, @@ -30,6 +30,7 @@ __all__ = [ "GraphRepositoryProtocol", # Serializer "GraphSerializer", + "TypeResolver", # Service "GraphService", "create_graph_service", diff --git a/flowsint-core/src/flowsint_core/core/graph/serializer.py b/flowsint-core/src/flowsint_core/core/graph/serializer.py index 62e0bfd7..9061eb3b 100644 --- a/flowsint-core/src/flowsint_core/core/graph/serializer.py +++ b/flowsint-core/src/flowsint_core/core/graph/serializer.py @@ -5,11 +5,14 @@ This module provides utilities for serializing complex Python objects into Neo4j-compatible primitive types, following the Single Responsibility Principle. """ -from typing import Any, Dict, List +from typing import Any, Callable, Dict, List, Optional, Type -from flowsint_types import TYPE_REGISTRY, FlowsintType +from flowsint_types import FlowsintType from pydantic import BaseModel +# Callable that resolves a type name to a FlowsintType subclass (or None). +TypeResolver = Callable[[str], Optional[Type[FlowsintType]]] + from flowsint_core.utils import flatten, unflatten from .types import GraphEdge, GraphNode, NodeMetadata @@ -53,8 +56,15 @@ class GraphSerializer: return flatten(dict, remove_empty=False) @staticmethod - def parse_flowsint_type(entity: Dict, nodeType: str) -> FlowsintType: - DetectedType = TYPE_REGISTRY.get_lowercase(nodeType) + def parse_flowsint_type( + entity: Dict, + nodeType: str, + type_resolver: Optional[TypeResolver] = None, + ) -> FlowsintType: + if not type_resolver: + from flowsint_core.core.services.type_registry_service import local_type_resolver + type_resolver = local_type_resolver + DetectedType = type_resolver(nodeType) if not DetectedType: raise ValueError(f"Unknown type: {nodeType}") properties = GraphSerializer._clean_empty_values(entity) @@ -65,7 +75,10 @@ class GraphSerializer: return node.nodeProperties @staticmethod - def neo4j_dict_to_graph_node(node_dict: Dict[str, Any]) -> GraphNode: + def neo4j_dict_to_graph_node( + node_dict: Dict[str, Any], + type_resolver: Optional[TypeResolver] = None, + ) -> GraphNode: """Convert a flattened Neo4j node record to a GraphNode instance. Unflattens the data, parses the nodeProperties into the appropriate @@ -88,7 +101,9 @@ class GraphSerializer: "nodeLabel", None ) # remove nodeLabel from original pydantic - entity = GraphSerializer.parse_flowsint_type(node_properties, node_type) + entity = GraphSerializer.parse_flowsint_type( + node_properties, node_type, type_resolver=type_resolver + ) return GraphNode( id=node_id, nodeLabel=nodeLabel, @@ -178,10 +193,15 @@ class GraphSerializer: } @staticmethod - def deserialize_nodes(node_dicts: List[Dict[str, Any]]) -> List[GraphNode]: + def deserialize_nodes( + node_dicts: List[Dict[str, Any]], + type_resolver: Optional[TypeResolver] = None, + ) -> List[GraphNode]: """Convert a list of Neo4j node records to GraphNode instances.""" return [ - GraphSerializer.neo4j_dict_to_graph_node(node_dict) + GraphSerializer.neo4j_dict_to_graph_node( + node_dict, type_resolver=type_resolver + ) for node_dict in node_dicts ] diff --git a/flowsint-core/src/flowsint_core/core/graph/service.py b/flowsint-core/src/flowsint_core/core/graph/service.py index 353615b6..1054b403 100644 --- a/flowsint-core/src/flowsint_core/core/graph/service.py +++ b/flowsint-core/src/flowsint_core/core/graph/service.py @@ -12,7 +12,7 @@ from pydantic import BaseModel from .repository import Neo4jGraphRepository from .repository_protocol import GraphRepositoryProtocol -from .serializer import GraphSerializer +from .serializer import GraphSerializer, TypeResolver from .types import GraphData, GraphDict, GraphNode @@ -41,6 +41,7 @@ class GraphService: repository: GraphRepositoryProtocol, logger: Optional[LoggerProtocol] = None, enable_batching: bool = False, + type_resolver: Optional[TypeResolver] = None, ): """ Initialize the graph service. @@ -50,6 +51,7 @@ class GraphService: repository: Repository instance (required - dependency injection) logger: Optional logger instance enable_batching: Enable batch operations + type_resolver: Optional callable to resolve custom types by name Raises: ValueError: If repository is not provided @@ -63,6 +65,7 @@ class GraphService: self._repository = repository self._logger = logger self._enable_batching = enable_batching + self._type_resolver = type_resolver @property def sketch_id(self) -> str: @@ -138,13 +141,17 @@ class GraphService: def get_sketch_graph(self) -> GraphData: graph_data = self.repository.get_sketch_graph(self.sketch_id) - nodes = GraphSerializer.deserialize_nodes(graph_data.get("nodes", [])) + nodes = GraphSerializer.deserialize_nodes( + graph_data.get("nodes", []), type_resolver=self._type_resolver + ) edges = GraphSerializer.deserialize_edges(graph_data.get("edges", [])) return GraphData(nodes=nodes, edges=edges) def get_nodes_by_ids(self, node_ids: List[str]) -> List[GraphNode]: nodes = self.repository.get_nodes_by_ids(node_ids, self.sketch_id) - return GraphSerializer.deserialize_nodes(nodes) + return GraphSerializer.deserialize_nodes( + nodes, type_resolver=self._type_resolver + ) def get_nodes_by_ids_for_task(self, node_ids: List[str]) -> List[BaseModel]: nodes = self.get_nodes_by_ids(node_ids) @@ -203,7 +210,9 @@ class GraphService: node_id=node_id, sketch_id=self._sketch_id, ) - nodes = GraphSerializer.deserialize_nodes(graph_data.get("nodes", [])) + nodes = GraphSerializer.deserialize_nodes( + graph_data.get("nodes", []), type_resolver=self._type_resolver + ) edges = GraphSerializer.deserialize_edges(graph_data.get("edges", [])) return GraphData(nodes=nodes, edges=edges) @@ -333,6 +342,7 @@ class GraphService: def create_graph_service( sketch_id: str, enable_batching: bool = True, + type_resolver: Optional[TypeResolver] = None, ) -> GraphService: """ Factory function to create a GraphService instance with Neo4j repository. @@ -343,6 +353,7 @@ def create_graph_service( Args: sketch_id: Investigation sketch ID enable_batching: Enable batch operations + type_resolver: Optional callable to resolve custom types by name Returns: Configured GraphService instance @@ -358,4 +369,5 @@ def create_graph_service( repository=repository, logger=Logger, enable_batching=enable_batching, + type_resolver=type_resolver, ) diff --git a/flowsint-core/src/flowsint_core/core/models.py b/flowsint-core/src/flowsint_core/core/models.py index be741fdc..bb7ba5f8 100644 --- a/flowsint-core/src/flowsint_core/core/models.py +++ b/flowsint-core/src/flowsint_core/core/models.py @@ -3,8 +3,8 @@ import uuid from datetime import datetime, timezone from sqlalchemy import ( - Boolean, JSON, + Boolean, Column, DateTime, Float, @@ -358,6 +358,8 @@ class CustomType(Base): nullable=False, ) schema: Mapped[dict] = mapped_column(JSON, nullable=False) + icon: Mapped[str] = mapped_column(String, nullable=True) + color: Mapped[str] = mapped_column(String, nullable=True) status: Mapped[str] = mapped_column(String, server_default="draft", nullable=False) checksum: Mapped[str] = mapped_column(String, nullable=True) description: Mapped[str] = mapped_column(Text, nullable=True) diff --git a/flowsint-core/src/flowsint_core/core/services/__init__.py b/flowsint-core/src/flowsint_core/core/services/__init__.py index 3ae586ad..11510fc2 100644 --- a/flowsint-core/src/flowsint_core/core/services/__init__.py +++ b/flowsint-core/src/flowsint_core/core/services/__init__.py @@ -25,7 +25,7 @@ from .scan_service import ScanService, create_scan_service from .log_service import LogService, create_log_service from .flow_service import FlowService, create_flow_service from .custom_type_service import CustomTypeService, create_custom_type_service -from .type_registry_service import TypeRegistryService, create_type_registry_service +from .type_registry_service import TypeRegistryService, create_type_registry_service, local_type_resolver from .enricher_service import EnricherService, create_enricher_service from .enricher_template_service import EnricherTemplateService, create_enricher_template_service from .template_generator_service import TemplateGeneratorService, create_template_generator_service @@ -68,6 +68,7 @@ __all__ = [ "create_custom_type_service", "TypeRegistryService", "create_type_registry_service", + "local_type_resolver", "EnricherService", "create_enricher_service", "EnricherTemplateService", diff --git a/flowsint-core/src/flowsint_core/core/services/custom_type_service.py b/flowsint-core/src/flowsint_core/core/services/custom_type_service.py index f93f3a9b..e49dd183 100644 --- a/flowsint-core/src/flowsint_core/core/services/custom_type_service.py +++ b/flowsint-core/src/flowsint_core/core/services/custom_type_service.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Session from ..models import CustomType from ..repositories import CustomTypeRepository from .base import BaseService -from .exceptions import NotFoundError, ValidationError, ConflictError +from .exceptions import ConflictError, NotFoundError, ValidationError class CustomTypeService(BaseService): @@ -26,13 +26,13 @@ class CustomTypeService(BaseService): self, user_id: UUID, status: Optional[str] = None ) -> List[CustomType]: if status and status not in ["draft", "published", "archived"]: - raise ValidationError( - "Status must be one of: draft, published, archived" - ) + raise ValidationError("Status must be one of: draft, published, archived") return self._custom_type_repo.get_by_owner(user_id, status=status) def get_by_id(self, custom_type_id: UUID, user_id: UUID) -> CustomType: - custom_type = self._custom_type_repo.get_by_id_and_owner(custom_type_id, user_id) + custom_type = self._custom_type_repo.get_by_id_and_owner( + custom_type_id, user_id + ) if not custom_type: raise NotFoundError("Custom type not found") return custom_type @@ -54,7 +54,9 @@ class CustomTypeService(BaseService): if validate_schema_func: validate_schema_func(json_schema) - checksum = calculate_checksum_func(json_schema) if calculate_checksum_func else None + checksum = ( + calculate_checksum_func(json_schema) if calculate_checksum_func else None + ) existing = self._custom_type_repo.get_by_name_and_owner(name, user_id) if existing: @@ -83,6 +85,8 @@ class CustomTypeService(BaseService): json_schema: Optional[Dict[str, Any]] = None, description: Optional[str] = None, status: Optional[str] = None, + color: Optional[str] = None, + icon: Optional[str] = None, validate_schema_func=None, calculate_checksum_func=None, ) -> CustomType: @@ -107,6 +111,12 @@ class CustomTypeService(BaseService): if status is not None: custom_type.status = status + if icon is not None: + custom_type.icon = icon + + if color is not None: + custom_type.color = color + self._commit() self._refresh(custom_type) diff --git a/flowsint-core/src/flowsint_core/core/services/sketch_service.py b/flowsint-core/src/flowsint_core/core/services/sketch_service.py index 1f4610cd..a898e787 100644 --- a/flowsint-core/src/flowsint_core/core/services/sketch_service.py +++ b/flowsint-core/src/flowsint_core/core/services/sketch_service.py @@ -2,7 +2,9 @@ Sketch service for managing sketches and graph operations. """ -from typing import Any, Dict, List, Optional +from __future__ import annotations + +from typing import Any, Dict, List, Optional, TYPE_CHECKING from uuid import UUID from sqlalchemy.orm import Session @@ -14,6 +16,9 @@ from ..repositories import SketchRepository, InvestigationRepository from .base import BaseService from .exceptions import NotFoundError, PermissionDeniedError, ValidationError, DatabaseError +if TYPE_CHECKING: + from .type_registry_service import TypeRegistryService + class SketchService(BaseService): """ @@ -25,11 +30,13 @@ class SketchService(BaseService): db: Session, sketch_repo: SketchRepository, investigation_repo: InvestigationRepository, + type_registry_service: Optional[TypeRegistryService] = None, **kwargs, ): super().__init__(db, **kwargs) self._sketch_repo = sketch_repo self._investigation_repo = investigation_repo + self._type_registry = type_registry_service def _get_sketch_with_permission( self, sketch_id: UUID, user_id: UUID, actions: List[str] @@ -107,8 +114,11 @@ class SketchService(BaseService): ) -> Dict[str, Any]: sketch = self._get_sketch_with_permission(sketch_id, user_id, ["read"]) + resolver = self._type_registry.build_type_resolver(user_id) if self._type_registry else None graph_service = create_graph_service( - sketch_id=str(sketch_id), enable_batching=False + sketch_id=str(sketch_id), + enable_batching=False, + type_resolver=resolver, ) graph_data = graph_service.get_sketch_graph() @@ -322,7 +332,11 @@ class SketchService(BaseService): self._get_sketch_with_permission(sketch_id, user_id, ["read"]) try: - graph_service = create_graph_service(sketch_id=str(sketch_id)) + resolver = self._type_registry.build_type_resolver(user_id) if self._type_registry else None + graph_service = create_graph_service( + sketch_id=str(sketch_id), + type_resolver=resolver, + ) result = graph_service.get_neighbors(node_id) except Exception as e: print(e) @@ -338,8 +352,11 @@ class SketchService(BaseService): ) -> Dict[str, Any]: sketch = self._get_sketch_with_permission(sketch_id, user_id, ["read"]) + resolver = self._type_registry.build_type_resolver(user_id) if self._type_registry else None graph_service = create_graph_service( - sketch_id=str(sketch_id), enable_batching=False + sketch_id=str(sketch_id), + enable_batching=False, + type_resolver=resolver, ) graph_data = graph_service.get_sketch_graph() @@ -358,9 +375,14 @@ class SketchService(BaseService): def create_sketch_service(db: Session) -> SketchService: - investigation_repo = InvestigationRepository(db) + from ..repositories import CustomTypeRepository + from .type_registry_service import TypeRegistryService + return SketchService( db=db, sketch_repo=SketchRepository(db), - investigation_repo=investigation_repo, + investigation_repo=InvestigationRepository(db), + type_registry_service=TypeRegistryService( + db=db, custom_type_repo=CustomTypeRepository(db) + ), ) diff --git a/flowsint-core/src/flowsint_core/core/services/type_registry_service.py b/flowsint-core/src/flowsint_core/core/services/type_registry_service.py index 16349a3d..270bd51f 100644 --- a/flowsint-core/src/flowsint_core/core/services/type_registry_service.py +++ b/flowsint-core/src/flowsint_core/core/services/type_registry_service.py @@ -5,13 +5,39 @@ Type registry service for managing flowsint types. from typing import Any, Dict, List, Optional, Type from uuid import UUID, uuid4 +from flowsint_types import FlowsintType +from pydantic import BaseModel, TypeAdapter, create_model from sqlalchemy.orm import Session -from pydantic import BaseModel, TypeAdapter +from ..graph.serializer import TypeResolver from ..repositories import CustomTypeRepository from .base import BaseService +def local_type_resolver(type_name: str) -> Type[FlowsintType] | None: + """Resolve a type using only the local TYPE_REGISTRY (no DB). + + Useful as a fallback when no TypeRegistryService is available (tests, CLI, etc.). + """ + from flowsint_types import TYPE_REGISTRY + + return TYPE_REGISTRY.get_lowercase(type_name) + + +def _build_pydantic_model_from_schema(name: str, schema: dict) -> Type[FlowsintType]: + """Build a dynamic Pydantic model from a custom type JSON schema.""" + properties = schema.get("properties", {}) + required = set(schema.get("required", [])) + + fields: Dict[str, Any] = {} + for prop, info in properties.items(): + annotation = Optional[str] if prop not in required else str + default = ... if prop in required else None + fields[prop] = (annotation, default) + + return create_model(name, __base__=FlowsintType, **fields) + + class TypeRegistryService(BaseService): """ Service for type registry operations and schema extraction. @@ -21,6 +47,57 @@ class TypeRegistryService(BaseService): super().__init__(db, **kwargs) self._custom_type_repo = custom_type_repo + def resolve_type(self, type_name: str, user_id: UUID) -> Type[FlowsintType] | None: + """Resolve a type name to a FlowsintType class. + + Checks the local TYPE_REGISTRY first, then falls back to custom types in DB. + """ + from flowsint_types import TYPE_REGISTRY + + model = TYPE_REGISTRY.get_lowercase(type_name) + if model: + return model + + custom_type = self._custom_type_repo.get_published_by_name_and_owner( + name=type_name, owner_id=user_id + ) + if not custom_type: + return None + return _build_pydantic_model_from_schema(custom_type.name, custom_type.schema) + + def build_type_resolver(self, user_id: UUID) -> TypeResolver: + """Return a TypeResolver callable bound to a specific user. + + Usage: + resolver = type_registry_service.build_type_resolver(user_id) + graph_service = create_graph_service(sketch_id, type_resolver=resolver) + """ + + def resolver(type_name: str) -> Type[FlowsintType] | None: + return self.resolve_type(type_name, user_id) + + return resolver + + def get_type(self, user_id: UUID, type_name: str) -> Dict[str, Any] | None: + from flowsint_types.registry import get_type as get_type_from_registry + + model = get_type_from_registry(type_name, case_sensitive=True) + + if model: + return self._extract_input_schema(model, label_key="nodeLabel") + else: + print( + f"Warning: Type {type_name} not found in TYPE_REGISTRY, checking in custom types..." + ) + custom_type = self._custom_type_repo.get_by_name_and_owner( + name=type_name, owner_id=user_id + ) + if not custom_type: + print(f"Warning: Type {type_name} not found.") + return None + print(f"Warning: Type {type_name} found in cutsom types.") + return self._extract_input_schema(custom_type, label_key="nodeLabel") + def get_types_list(self, user_id: UUID) -> List[Dict[str, Any]]: from flowsint_types.registry import get_type @@ -37,7 +114,9 @@ class TypeRegistryService(BaseService): if model: children_schemas.append( - self._extract_input_schema(model, label_key=label_key, icon=icon) + self._extract_input_schema( + model, label_key=label_key, icon=icon + ) ) else: print(f"Warning: Type {type_name} not found in TYPE_REGISTRY") @@ -68,7 +147,8 @@ class TypeRegistryService(BaseService): "type": custom_type.name, "key": custom_type.name.lower(), "label_key": label_key, - "icon": "custom", + "icon": custom_type.icon or "custom", + "color": custom_type.color, "label": custom_type.name, "description": custom_type.description or "", "fields": [ @@ -167,8 +247,8 @@ class TypeRegistryService(BaseService): ("Website", "url", None), ("Ip", "address", None), ("Port", "number", None), - ("DNSRecord", "name", "dns"), - ("SSLCertificate", "subject", "ssl"), + ("DNSRecord", "name", "dnsrecord"), + ("SSLCertificate", "subject", "sslcertificate"), ("WebTracker", "name", "webtracker"), ], }, diff --git a/flowsint-core/src/flowsint_core/core/types.py b/flowsint-core/src/flowsint_core/core/types.py index 2c00ebd1..1b6cb594 100644 --- a/flowsint-core/src/flowsint_core/core/types.py +++ b/flowsint-core/src/flowsint_core/core/types.py @@ -65,7 +65,7 @@ class FlowStep(BaseModel): params: Optional[Dict[str, Any]] = Field( None, description="Parameters for the step", title="Parameters" ) - type: Literal["type", "enricher"] = Field( + type: Literal["type", "enricher", "error"] = Field( ..., description="Type of step - either type transformation or enricher", title="Step Type", diff --git a/flowsint-core/src/flowsint_core/imports/entity_detection.py b/flowsint-core/src/flowsint_core/imports/entity_detection.py index e643fa89..0cc126aa 100644 --- a/flowsint-core/src/flowsint_core/imports/entity_detection.py +++ b/flowsint-core/src/flowsint_core/imports/entity_detection.py @@ -9,6 +9,12 @@ from flowsint_types import TYPE_REGISTRY, FlowsintType def detect_type(value: str) -> Optional[Type[FlowsintType]]: + """Detect entity type from a string value using built-in type patterns. + + Note: This only checks built-in types (TYPE_REGISTRY) since custom types + don't have detect() methods. For custom type resolution by name, use + TypeRegistryService.resolve_type(). + """ for model in TYPE_REGISTRY.all_types().values(): if hasattr(model, "detect") and model.detect(value): return model diff --git a/flowsint-core/src/flowsint_core/imports/file_parser.py b/flowsint-core/src/flowsint_core/imports/file_parser.py index 64abbb15..51c21c0a 100644 --- a/flowsint-core/src/flowsint_core/imports/file_parser.py +++ b/flowsint-core/src/flowsint_core/imports/file_parser.py @@ -7,6 +7,8 @@ Each line represents ONE entity with a single value. from pathlib import Path from typing import BinaryIO, Optional, Union +from flowsint_core.core.graph.serializer import TypeResolver + from .json import parse_json from .txt import parse_txt from .types import FileParseResult @@ -18,6 +20,7 @@ def parse_import_file( file_content: Optional[Union[bytes, BinaryIO]], filename: Optional[str], max_preview_rows: int = 100, + type_resolver: Optional[TypeResolver] = None, ) -> FileParseResult | None: """ Parse an uploaded file and analyze its contents. @@ -40,4 +43,4 @@ def parse_import_file( if file_ext == ".txt": return parse_txt(file_bytes, max_preview_rows) elif file_ext in [".json"]: - return parse_json(file_bytes, max_preview_rows) + return parse_json(file_bytes, max_preview_rows, type_resolver=type_resolver) diff --git a/flowsint-core/src/flowsint_core/imports/import_service.py b/flowsint-core/src/flowsint_core/imports/import_service.py index 3229956c..a52c3ad9 100644 --- a/flowsint-core/src/flowsint_core/imports/import_service.py +++ b/flowsint-core/src/flowsint_core/imports/import_service.py @@ -11,6 +11,7 @@ from typing import Any, Dict, List, Optional from flowsint_types import FlowsintType from flowsint_core.core.graph import GraphSerializer, GraphService +from flowsint_core.core.graph.serializer import TypeResolver @dataclass @@ -45,17 +46,23 @@ class ImportService: - Batch node and edge creation """ - def __init__(self, graph_service: GraphService): + def __init__( + self, + graph_service: GraphService, + type_resolver: Optional[TypeResolver] = None, + ): """ Initialize the import service. Args: graph_service: GraphService instance for database operations + type_resolver: Optional callable to resolve types by name """ self._graph_service = graph_service + self._type_resolver = type_resolver or graph_service._type_resolver - @staticmethod def analyze_file( + self, file_content: bytes, filename: str, max_preview_rows: int = 10000000, @@ -80,6 +87,7 @@ class ImportService: file_content=file_content, filename=filename, max_preview_rows=max_preview_rows, + type_resolver=self._type_resolver, ) def execute_import( @@ -166,6 +174,7 @@ class ImportService: pydantic_obj = GraphSerializer.parse_flowsint_type( entity=entity_data, nodeType=mapping.entity_type, + type_resolver=self._type_resolver, ) pydantic_nodes.append(pydantic_obj) diff --git a/flowsint-core/src/flowsint_core/imports/json/parse_json.py b/flowsint-core/src/flowsint_core/imports/json/parse_json.py index d4ff8583..98cdeddb 100644 --- a/flowsint-core/src/flowsint_core/imports/json/parse_json.py +++ b/flowsint-core/src/flowsint_core/imports/json/parse_json.py @@ -1,7 +1,9 @@ import json -from typing import Dict, List +from typing import Dict, List, Optional -from flowsint_types import TYPE_REGISTRY, Individual +from flowsint_types import Individual + +from flowsint_core.core.graph.serializer import TypeResolver from ..types import Edge, Entity, EntityPreview, FileParseResult from ..utils import create_entity_preview @@ -15,6 +17,7 @@ VALID_EDGE_TO_KEYS = ["to", "target"] def parse_json( file_bytes: bytes, max_preview_rows: int, + type_resolver: Optional[TypeResolver] = None, ) -> FileParseResult: """Parse JSON files.""" try: @@ -33,7 +36,7 @@ def parse_json( raise Exception( f"No valid edges keys found in JSON. Please provide one of : {', '.join(VALID_EDGES_KEYS)}" ) - nodes = _get_nodes(graph.get(node_key)) + nodes = _get_nodes(graph.get(node_key), type_resolver=type_resolver) edges = _get_edges(graph.get(edge_key)) for node in nodes: if node: @@ -67,14 +70,17 @@ def parse_json( raise Exception(f"Invalid JSON: {str(e)}") -def _get_nodes(nodes: List[Dict]) -> List[EntityPreview]: +def _get_nodes( + nodes: List[Dict], + type_resolver: Optional[TypeResolver] = None, +) -> List[EntityPreview]: results = [] for node in nodes: node_id = node.get("id") node_type = node.get("nodeType", node.get("type")) label = node.get("nodeLabel", node.get("label")) node_obj = {"nodeType": node_type, "nodeLabel": label, **node} - preview = _parse_node(node_obj) + preview = _parse_node(node_obj, type_resolver=type_resolver) if preview: results.append( EntityPreview( @@ -87,7 +93,10 @@ def _get_nodes(nodes: List[Dict]) -> List[EntityPreview]: return results -def _parse_node(nodeDict: dict) -> EntityPreview | None: +def _parse_node( + nodeDict: dict, + type_resolver: Optional[TypeResolver] = None, +) -> EntityPreview | None: type = nodeDict.get("nodeType") first_key = next(iter(nodeDict), None) label = ( @@ -99,8 +108,11 @@ def _parse_node(nodeDict: dict) -> EntityPreview | None: return None if not type: return create_entity_preview(label) - # Try to find a matching type in TYPE_REGISTRY - DetectedType = TYPE_REGISTRY.get_lowercase(type.lower()) + # Try to resolve the type via the registry service + if not type_resolver: + from flowsint_core.core.services.type_registry_service import local_type_resolver + type_resolver = local_type_resolver + DetectedType = type_resolver(type.lower()) if not DetectedType: # Type not recognized, fall back to generic return create_entity_preview(label) diff --git a/flowsint-core/src/flowsint_core/templates/loader/yaml_loader.py b/flowsint-core/src/flowsint_core/templates/loader/yaml_loader.py index 0b895e40..297cbb72 100644 --- a/flowsint-core/src/flowsint_core/templates/loader/yaml_loader.py +++ b/flowsint-core/src/flowsint_core/templates/loader/yaml_loader.py @@ -1,11 +1,11 @@ import ipaddress import re -from typing import Any, Set +from typing import Any, Optional, Set from urllib.parse import urlparse import yaml -from flowsint_types import TYPE_REGISTRY +from flowsint_core.core.graph.serializer import TypeResolver from flowsint_core.templates.types import Template # Template variable pattern: {{variable_name}} or {{secrets.NAME}} @@ -116,7 +116,10 @@ class YamlLoader: return exc @staticmethod - def parse_yaml_to_template(raw: dict[str, Any]) -> Template: + def parse_yaml_to_template( + raw: dict[str, Any], + type_resolver: Optional[TypeResolver] = None, + ) -> Template: if not isinstance(raw, dict): raise ValueError("Template must be a YAML dictionary") @@ -130,7 +133,10 @@ class YamlLoader: if not input_type: raise ValueError("Missing 'input.type' property in the yaml.") - DetectedType = TYPE_REGISTRY.get(input_type) + if not type_resolver: + from flowsint_core.core.services.type_registry_service import local_type_resolver + type_resolver = local_type_resolver + DetectedType = type_resolver(input_type) request = raw.get("request", {}) method = request.get("method", "GET") @@ -146,12 +152,17 @@ class YamlLoader: return Template(**raw) @staticmethod - def get_template_from_file(filename: str) -> Template | None: + def get_template_from_file( + filename: str, + type_resolver: Optional[TypeResolver] = None, + ) -> Template | None: template_dict = YamlLoader.load_enricher_yaml(filename) if not isinstance(template_dict, dict): return None - return YamlLoader.parse_yaml_to_template(template_dict) + return YamlLoader.parse_yaml_to_template( + template_dict, type_resolver=type_resolver + ) @staticmethod def render_template(