mirror of
https://github.com/reconurge/flowsint.git
synced 2026-07-21 12:50:12 -05:00
Merge pull request #35 from reconurge/refactor/best-pratices
refactor(core): follow best practice principles
This commit is contained in:
dextmorgn
GitHub
@@ -1,56 +1,177 @@
|
||||
import os
|
||||
from neo4j import GraphDatabase
|
||||
from threading import Lock
|
||||
from typing import Optional, Dict, Any, List
|
||||
from neo4j import GraphDatabase, Driver, Session
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
class Neo4jConnection:
|
||||
def __init__(self, uri: str, user: str, password: str):
|
||||
self.__uri = uri
|
||||
self.__user = user
|
||||
self.__password = password
|
||||
self.__driver = GraphDatabase.driver(
|
||||
self.__uri, auth=(self.__user, self.__password)
|
||||
)
|
||||
"""
|
||||
Singleton Neo4j connection manager with proper connection pooling.
|
||||
|
||||
def query(self, query: str, parameters: dict = None):
|
||||
with self.__driver.session() as session:
|
||||
result = session.run(query, parameters)
|
||||
This class implements the Singleton pattern to ensure only one Neo4j driver
|
||||
instance exists throughout the application lifecycle, providing efficient
|
||||
connection pooling and resource management.
|
||||
"""
|
||||
|
||||
_instance: Optional['Neo4jConnection'] = None
|
||||
_lock: Lock = Lock()
|
||||
_driver: Optional[Driver] = None
|
||||
|
||||
def __new__(cls, uri: str = None, user: str = None, password: str = None):
|
||||
"""
|
||||
Create or return the singleton instance.
|
||||
|
||||
Thread-safe singleton implementation using double-checked locking.
|
||||
"""
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
instance = super().__new__(cls)
|
||||
cls._instance = instance
|
||||
return cls._instance
|
||||
|
||||
def __init__(self, uri: str = None, user: str = None, password: str = None):
|
||||
"""
|
||||
Initialize the Neo4j connection (only once).
|
||||
|
||||
Args:
|
||||
uri: Neo4j connection URI
|
||||
user: Neo4j username
|
||||
password: Neo4j password
|
||||
"""
|
||||
# Only initialize once
|
||||
if self._driver is None:
|
||||
self._uri = uri or os.getenv("NEO4J_URI_BOLT")
|
||||
self._user = user or os.getenv("NEO4J_USERNAME")
|
||||
self._password = password or os.getenv("NEO4J_PASSWORD")
|
||||
|
||||
if not all([self._uri, self._user, self._password]):
|
||||
raise ValueError("Neo4j connection credentials are required")
|
||||
|
||||
self._driver = GraphDatabase.driver(
|
||||
self._uri,
|
||||
auth=(self._user, self._password),
|
||||
max_connection_pool_size=50,
|
||||
connection_acquisition_timeout=60.0
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> 'Neo4jConnection':
|
||||
"""
|
||||
Get the singleton instance.
|
||||
|
||||
Returns:
|
||||
The singleton Neo4jConnection instance
|
||||
"""
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
def get_driver(self) -> Driver:
|
||||
"""
|
||||
Get the Neo4j driver instance.
|
||||
|
||||
Returns:
|
||||
Neo4j Driver instance
|
||||
"""
|
||||
return self._driver
|
||||
|
||||
def query(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Execute a single query.
|
||||
|
||||
Args:
|
||||
query: Cypher query string
|
||||
parameters: Query parameters
|
||||
|
||||
Returns:
|
||||
List of result records as dictionaries
|
||||
"""
|
||||
with self._driver.session() as session:
|
||||
result = session.run(query, parameters or {})
|
||||
return result.data()
|
||||
|
||||
def execute_write(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Execute a write query within a write transaction.
|
||||
|
||||
URI = os.getenv("NEO4J_URI_BOLT")
|
||||
USERNAME = os.getenv("NEO4J_USERNAME")
|
||||
PASSWORD = os.getenv("NEO4J_PASSWORD")
|
||||
Args:
|
||||
query: Cypher query string
|
||||
parameters: Query parameters
|
||||
|
||||
neo4j_connection = Neo4jConnection(URI, USERNAME, PASSWORD)
|
||||
Returns:
|
||||
List of result records as dictionaries
|
||||
"""
|
||||
def _execute(tx):
|
||||
result = tx.run(query, parameters or {})
|
||||
return result.data()
|
||||
|
||||
with self._driver.session() as session:
|
||||
return session.execute_write(_execute)
|
||||
|
||||
def execute_batch(self, queries: List[tuple[str, Dict[str, Any]]]) -> None:
|
||||
"""
|
||||
Execute multiple queries in a single transaction.
|
||||
|
||||
Args:
|
||||
queries: List of (query, parameters) tuples
|
||||
"""
|
||||
def _execute_batch(tx):
|
||||
for query, params in queries:
|
||||
tx.run(query, params or {})
|
||||
|
||||
with self._driver.session() as session:
|
||||
session.execute_write(_execute_batch)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the driver connection."""
|
||||
if self._driver:
|
||||
self._driver.close()
|
||||
self._driver = None
|
||||
|
||||
def verify_connectivity(self) -> bool:
|
||||
"""
|
||||
Verify the connection to Neo4j.
|
||||
|
||||
Returns:
|
||||
True if connection is successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
with self._driver.session() as session:
|
||||
result = session.run("RETURN 1")
|
||||
return result.single()[0] == 1
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def reset_instance(cls) -> None:
|
||||
"""
|
||||
Reset the singleton instance (mainly for testing).
|
||||
|
||||
WARNING: This should only be used in test scenarios.
|
||||
"""
|
||||
with cls._lock:
|
||||
if cls._instance and cls._instance._driver:
|
||||
cls._instance._driver.close()
|
||||
cls._instance = None
|
||||
cls._driver = None
|
||||
|
||||
|
||||
# import os
|
||||
# from neo4j import GraphDatabase
|
||||
# from dotenv import load_dotenv
|
||||
# from pydantic_neo4j import PydanticNeo4j
|
||||
# load_dotenv()
|
||||
# Create the default singleton instance
|
||||
try:
|
||||
URI = os.getenv("NEO4J_URI_BOLT")
|
||||
USERNAME = os.getenv("NEO4J_USERNAME")
|
||||
PASSWORD = os.getenv("NEO4J_PASSWORD")
|
||||
|
||||
# class Neo4jConnection:
|
||||
# def __init__(self, uri: str, user: str, password: str):
|
||||
# self.__uri = uri
|
||||
# self.__user = user
|
||||
# self.__password = password
|
||||
# self.__driver = GraphDatabase.driver(self.__uri, auth=(self.__user, self.__password))
|
||||
|
||||
# def query(self, query: str, parameters: dict = None):
|
||||
# with self.__driver.session() as session:
|
||||
# result = session.run(query, parameters)
|
||||
# return result.data()
|
||||
|
||||
|
||||
# URI = os.getenv("NEO4J_URI_BOLT")
|
||||
# USERNAME = os.getenv("NEO4J_USERNAME")
|
||||
# PASSWORD = os.getenv("NEO4J_PASSWORD")
|
||||
|
||||
# neo4j_connection = PydanticNeo4j(URI, USERNAME, PASSWORD)
|
||||
# match_util = neo4j_connection.match_utilities
|
||||
# create_util = neo4j_connection.create_utilities
|
||||
# database_operations = neo4j_connection.database_operations
|
||||
if all([URI, USERNAME, PASSWORD]):
|
||||
neo4j_connection = Neo4jConnection(URI, USERNAME, PASSWORD)
|
||||
else:
|
||||
# Don't create instance if credentials are missing
|
||||
neo4j_connection = None
|
||||
except Exception as e:
|
||||
import logging
|
||||
logging.getLogger(__name__).warning(f"Failed to initialize Neo4j connection: {e}")
|
||||
neo4j_connection = None
|
||||
|
||||
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
Graph database repository for Neo4j operations.
|
||||
|
||||
This module provides a repository pattern implementation for Neo4j,
|
||||
handling node and relationship operations with batching support.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, List, Optional, Tuple
|
||||
from .graph_db import Neo4jConnection
|
||||
from .graph_serializer import GraphSerializer
|
||||
|
||||
|
||||
class GraphRepository:
|
||||
"""
|
||||
Repository for Neo4j graph database operations.
|
||||
|
||||
This class follows the Repository pattern, providing a clean abstraction
|
||||
over Neo4j operations and handling batching for improved performance.
|
||||
"""
|
||||
|
||||
def __init__(self, neo4j_connection: Optional[Neo4jConnection] = None):
|
||||
"""
|
||||
Initialize the graph repository.
|
||||
|
||||
Args:
|
||||
neo4j_connection: Optional Neo4j connection instance.
|
||||
If None, uses the singleton instance.
|
||||
"""
|
||||
self._connection = neo4j_connection or Neo4jConnection.get_instance()
|
||||
self._batch_operations: List[Tuple[str, Dict[str, Any]]] = []
|
||||
self._batch_size = 100
|
||||
|
||||
def create_node(
|
||||
self,
|
||||
node_type: str,
|
||||
key_prop: str,
|
||||
key_value: str,
|
||||
sketch_id: str,
|
||||
**properties: Any
|
||||
) -> None:
|
||||
"""
|
||||
Create or update a single node in Neo4j.
|
||||
|
||||
Args:
|
||||
node_type: Node label (e.g., "domain", "ip")
|
||||
key_prop: Property name used as unique identifier
|
||||
key_value: Value of the key property
|
||||
sketch_id: Investigation sketch ID
|
||||
**properties: Additional node properties
|
||||
"""
|
||||
if not self._connection:
|
||||
return
|
||||
|
||||
# Serialize properties
|
||||
serialized_props = GraphSerializer.serialize_properties(properties)
|
||||
|
||||
# Add required properties
|
||||
serialized_props["type"] = node_type.lower()
|
||||
serialized_props["sketch_id"] = sketch_id
|
||||
serialized_props["label"] = serialized_props.get("label", key_value)
|
||||
|
||||
# Build SET clauses
|
||||
set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()]
|
||||
params = {key_prop: key_value, **serialized_props}
|
||||
|
||||
# Build and execute query
|
||||
query = f"""
|
||||
MERGE (n:{node_type} {{{key_prop}: ${key_prop}}})
|
||||
SET {', '.join(set_clauses)}
|
||||
"""
|
||||
|
||||
self._connection.execute_write(query, params)
|
||||
|
||||
def create_relationship(
|
||||
self,
|
||||
from_type: str,
|
||||
from_key: str,
|
||||
from_value: str,
|
||||
to_type: str,
|
||||
to_key: str,
|
||||
to_value: str,
|
||||
rel_type: str,
|
||||
sketch_id: str,
|
||||
**properties: Any
|
||||
) -> None:
|
||||
"""
|
||||
Create a relationship between two nodes.
|
||||
|
||||
Args:
|
||||
from_type: Source node label
|
||||
from_key: Source node key property
|
||||
from_value: Source node key value
|
||||
to_type: Target node label
|
||||
to_key: Target node key property
|
||||
to_value: Target node key value
|
||||
rel_type: Relationship type
|
||||
sketch_id: Investigation sketch ID
|
||||
**properties: Additional relationship properties
|
||||
"""
|
||||
if not self._connection:
|
||||
return
|
||||
|
||||
# Serialize relationship properties
|
||||
serialized_props = GraphSerializer.serialize_properties(properties)
|
||||
serialized_props["sketch_id"] = sketch_id
|
||||
|
||||
# Build relationship properties string
|
||||
if serialized_props:
|
||||
props_str = ", ".join([f"{k}: ${k}" for k in serialized_props.keys()])
|
||||
rel_props = f"{{{props_str}}}"
|
||||
else:
|
||||
rel_props = "{sketch_id: $sketch_id}"
|
||||
|
||||
query = f"""
|
||||
MATCH (from:{from_type} {{{from_key}: $from_value}})
|
||||
MATCH (to:{to_type} {{{to_key}: $to_value}})
|
||||
MERGE (from)-[:{rel_type} {rel_props}]->(to)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"from_value": from_value,
|
||||
"to_value": to_value,
|
||||
**serialized_props
|
||||
}
|
||||
|
||||
self._connection.execute_write(query, params)
|
||||
|
||||
def add_to_batch(
|
||||
self,
|
||||
operation_type: str,
|
||||
**kwargs: Any
|
||||
) -> None:
|
||||
"""
|
||||
Add an operation to the batch queue.
|
||||
|
||||
Args:
|
||||
operation_type: Type of operation ("node" or "relationship")
|
||||
**kwargs: Operation parameters
|
||||
"""
|
||||
if operation_type == "node":
|
||||
query, params = self._build_node_query(**kwargs)
|
||||
elif operation_type == "relationship":
|
||||
query, params = self._build_relationship_query(**kwargs)
|
||||
else:
|
||||
raise ValueError(f"Unknown operation type: {operation_type}")
|
||||
|
||||
self._batch_operations.append((query, params))
|
||||
|
||||
# Auto-flush if batch is full
|
||||
if len(self._batch_operations) >= self._batch_size:
|
||||
self.flush_batch()
|
||||
|
||||
def _build_node_query(
|
||||
self,
|
||||
node_type: str,
|
||||
key_prop: str,
|
||||
key_value: str,
|
||||
sketch_id: str,
|
||||
**properties: Any
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""Build a node creation query."""
|
||||
serialized_props = GraphSerializer.serialize_properties(properties)
|
||||
serialized_props["type"] = node_type.lower()
|
||||
serialized_props["sketch_id"] = sketch_id
|
||||
serialized_props["label"] = serialized_props.get("label", key_value)
|
||||
|
||||
set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()]
|
||||
params = {key_prop: key_value, **serialized_props}
|
||||
|
||||
query = f"""
|
||||
MERGE (n:{node_type} {{{key_prop}: ${key_prop}}})
|
||||
SET {', '.join(set_clauses)}
|
||||
"""
|
||||
|
||||
return query, params
|
||||
|
||||
def _build_relationship_query(
|
||||
self,
|
||||
from_type: str,
|
||||
from_key: str,
|
||||
from_value: str,
|
||||
to_type: str,
|
||||
to_key: str,
|
||||
to_value: str,
|
||||
rel_type: str,
|
||||
sketch_id: str,
|
||||
**properties: Any
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""Build a relationship creation query."""
|
||||
serialized_props = GraphSerializer.serialize_properties(properties)
|
||||
serialized_props["sketch_id"] = sketch_id
|
||||
|
||||
if serialized_props:
|
||||
props_str = ", ".join([f"{k}: ${k}" for k in serialized_props.keys()])
|
||||
rel_props = f"{{{props_str}}}"
|
||||
else:
|
||||
rel_props = "{sketch_id: $sketch_id}"
|
||||
|
||||
query = f"""
|
||||
MATCH (from:{from_type} {{{from_key}: $from_value}})
|
||||
MATCH (to:{to_type} {{{to_key}: $to_value}})
|
||||
MERGE (from)-[:{rel_type} {rel_props}]->(to)
|
||||
"""
|
||||
|
||||
params = {
|
||||
"from_value": from_value,
|
||||
"to_value": to_value,
|
||||
**serialized_props
|
||||
}
|
||||
|
||||
return query, params
|
||||
|
||||
def flush_batch(self) -> None:
|
||||
"""Execute all batched operations in a single transaction."""
|
||||
if not self._batch_operations:
|
||||
return
|
||||
|
||||
if not self._connection:
|
||||
self._batch_operations.clear()
|
||||
return
|
||||
|
||||
try:
|
||||
self._connection.execute_batch(self._batch_operations)
|
||||
finally:
|
||||
self._batch_operations.clear()
|
||||
|
||||
def clear_batch(self) -> None:
|
||||
"""Clear the batch without executing."""
|
||||
self._batch_operations.clear()
|
||||
|
||||
def set_batch_size(self, size: int) -> None:
|
||||
"""
|
||||
Set the batch size for auto-flushing.
|
||||
|
||||
Args:
|
||||
size: Number of operations to batch before auto-flush
|
||||
"""
|
||||
if size < 1:
|
||||
raise ValueError("Batch size must be at least 1")
|
||||
self._batch_size = size
|
||||
|
||||
def query(self, cypher: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Execute a custom Cypher query.
|
||||
|
||||
Args:
|
||||
cypher: Cypher query string
|
||||
parameters: Query parameters
|
||||
|
||||
Returns:
|
||||
List of result records
|
||||
"""
|
||||
if not self._connection:
|
||||
return []
|
||||
|
||||
return self._connection.query(cypher, parameters)
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry."""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit - auto-flush batch."""
|
||||
if exc_type is None:
|
||||
self.flush_batch()
|
||||
else:
|
||||
self.clear_batch()
|
||||
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
Graph property serialization utilities.
|
||||
|
||||
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
|
||||
|
||||
|
||||
class GraphSerializer:
|
||||
"""
|
||||
Handles serialization of complex objects to Neo4j-compatible types.
|
||||
|
||||
This class is responsible for converting Pydantic models, nested objects,
|
||||
and other complex types into primitive types that can be stored in Neo4j.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def serialize_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Convert properties to Neo4j-compatible values.
|
||||
|
||||
Handles:
|
||||
- Pydantic models (flattened into individual properties)
|
||||
- Nested objects (flattened with prefixed keys)
|
||||
- Lists (converted to primitive types)
|
||||
- Dictionaries (flattened with prefixed keys)
|
||||
- None values (converted to empty strings)
|
||||
|
||||
Args:
|
||||
properties: Dictionary of properties to serialize
|
||||
|
||||
Returns:
|
||||
Dictionary of Neo4j-compatible properties
|
||||
"""
|
||||
serialized = {}
|
||||
|
||||
for key, value in properties.items():
|
||||
if value is None:
|
||||
serialized[key] = ""
|
||||
elif GraphSerializer._is_pydantic_model(value):
|
||||
# Flatten Pydantic models
|
||||
flattened = GraphSerializer._flatten_pydantic(key, value)
|
||||
serialized.update(flattened)
|
||||
elif isinstance(value, dict):
|
||||
# Flatten dictionaries
|
||||
flattened = GraphSerializer._flatten_dict(key, value)
|
||||
serialized.update(flattened)
|
||||
elif isinstance(value, list):
|
||||
# Handle lists
|
||||
serialized[key] = GraphSerializer._serialize_list(value)
|
||||
elif isinstance(value, (str, int, float, bool)):
|
||||
# Keep primitive types as-is
|
||||
serialized[key] = "" if value == "None" else value
|
||||
else:
|
||||
# Convert other complex types to strings
|
||||
serialized[key] = str(value) if value != "None" else ""
|
||||
|
||||
return serialized
|
||||
|
||||
@staticmethod
|
||||
def _is_pydantic_model(obj: Any) -> bool:
|
||||
"""Check if an object is a Pydantic model."""
|
||||
return (
|
||||
hasattr(obj, "__dict__")
|
||||
and not isinstance(obj, (str, int, float, bool))
|
||||
and (hasattr(obj, "model_dump") or hasattr(obj, "dict"))
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _flatten_pydantic(prefix: str, model: Any) -> Dict[str, Any]:
|
||||
"""
|
||||
Flatten a Pydantic model into individual properties.
|
||||
|
||||
Args:
|
||||
prefix: Key prefix for nested properties
|
||||
model: Pydantic model instance
|
||||
|
||||
Returns:
|
||||
Flattened dictionary
|
||||
"""
|
||||
flattened = {}
|
||||
|
||||
# Try Pydantic v2 first, then v1
|
||||
if hasattr(model, "model_dump"):
|
||||
data = model.model_dump()
|
||||
elif hasattr(model, "dict"):
|
||||
data = model.dict()
|
||||
else:
|
||||
# Fallback to __dict__
|
||||
data = model.__dict__
|
||||
|
||||
for nested_key, nested_value in data.items():
|
||||
if nested_value is not None:
|
||||
new_key = f"{prefix}_{nested_key}"
|
||||
if isinstance(nested_value, (str, int, float, bool)):
|
||||
flattened[new_key] = nested_value
|
||||
else:
|
||||
# Recursively handle nested complex types
|
||||
flattened[new_key] = str(nested_value)
|
||||
|
||||
return flattened
|
||||
|
||||
@staticmethod
|
||||
def _flatten_dict(prefix: str, data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""
|
||||
Flatten a dictionary into individual properties.
|
||||
|
||||
Args:
|
||||
prefix: Key prefix for nested properties
|
||||
data: Dictionary to flatten
|
||||
|
||||
Returns:
|
||||
Flattened dictionary
|
||||
"""
|
||||
flattened = {}
|
||||
|
||||
for dict_key, dict_value in data.items():
|
||||
if dict_value is not None:
|
||||
new_key = f"{prefix}_{dict_key}"
|
||||
if isinstance(dict_value, (str, int, float, bool)):
|
||||
flattened[new_key] = dict_value
|
||||
else:
|
||||
flattened[new_key] = str(dict_value)
|
||||
|
||||
return flattened
|
||||
|
||||
@staticmethod
|
||||
def _serialize_list(items: List[Any]) -> List[Any]:
|
||||
"""
|
||||
Serialize a list to Neo4j-compatible types.
|
||||
|
||||
Args:
|
||||
items: List of items to serialize
|
||||
|
||||
Returns:
|
||||
List of serialized items
|
||||
"""
|
||||
serialized = []
|
||||
|
||||
for item in items:
|
||||
if isinstance(item, (str, int, float, bool)):
|
||||
serialized.append(item)
|
||||
elif GraphSerializer._is_pydantic_model(item):
|
||||
# Convert complex objects to strings
|
||||
serialized.append(str(item))
|
||||
else:
|
||||
serialized.append(str(item))
|
||||
|
||||
return serialized
|
||||
|
||||
@staticmethod
|
||||
def normalize_property_value(value: Any) -> Any:
|
||||
"""
|
||||
Normalize a single property value for Neo4j.
|
||||
|
||||
Args:
|
||||
value: Value to normalize
|
||||
|
||||
Returns:
|
||||
Normalized value
|
||||
"""
|
||||
if value is None:
|
||||
return ""
|
||||
elif isinstance(value, (str, int, float, bool)):
|
||||
return "" if value == "None" else value
|
||||
else:
|
||||
return str(value)
|
||||
@@ -0,0 +1,217 @@
|
||||
"""
|
||||
Graph service for high-level graph operations.
|
||||
|
||||
This module provides a service layer for graph operations,
|
||||
integrating repository and logging functionality.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any, Optional, Protocol
|
||||
from uuid import UUID
|
||||
from .graph_repository import GraphRepository
|
||||
from .graph_db import Neo4jConnection
|
||||
|
||||
|
||||
class LoggerProtocol(Protocol):
|
||||
"""Protocol for logger implementations."""
|
||||
|
||||
@staticmethod
|
||||
def graph_append(sketch_id: str, message: Dict[str, Any]) -> None:
|
||||
"""Log a graph append message."""
|
||||
...
|
||||
|
||||
|
||||
class GraphService:
|
||||
"""
|
||||
High-level service for graph operations.
|
||||
|
||||
This service provides a clean interface for transform operations,
|
||||
handling both graph persistence and logging with proper separation of concerns.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sketch_id: str,
|
||||
neo4j_connection: Optional[Neo4jConnection] = None,
|
||||
logger: Optional[LoggerProtocol] = None,
|
||||
enable_batching: bool = True
|
||||
):
|
||||
"""
|
||||
Initialize the graph service.
|
||||
|
||||
Args:
|
||||
sketch_id: Investigation sketch ID
|
||||
neo4j_connection: Optional Neo4j connection
|
||||
logger: Optional logger instance
|
||||
enable_batching: Enable batch operations
|
||||
"""
|
||||
self._sketch_id = sketch_id
|
||||
self._repository = GraphRepository(neo4j_connection)
|
||||
self._logger = logger
|
||||
self._enable_batching = enable_batching
|
||||
|
||||
@property
|
||||
def sketch_id(self) -> str:
|
||||
"""Get the sketch ID."""
|
||||
return self._sketch_id
|
||||
|
||||
@property
|
||||
def repository(self) -> GraphRepository:
|
||||
"""Get the underlying repository."""
|
||||
return self._repository
|
||||
|
||||
def create_node(
|
||||
self,
|
||||
node_type: str,
|
||||
key_prop: str,
|
||||
key_value: str,
|
||||
**properties: Any
|
||||
) -> None:
|
||||
"""
|
||||
Create or update a node in the graph.
|
||||
|
||||
Args:
|
||||
node_type: Node label (e.g., "domain", "ip")
|
||||
key_prop: Property name used as unique identifier
|
||||
key_value: Value of the key property
|
||||
**properties: Additional node properties
|
||||
"""
|
||||
if self._enable_batching:
|
||||
self._repository.add_to_batch(
|
||||
"node",
|
||||
node_type=node_type,
|
||||
key_prop=key_prop,
|
||||
key_value=key_value,
|
||||
sketch_id=self._sketch_id,
|
||||
**properties
|
||||
)
|
||||
else:
|
||||
self._repository.create_node(
|
||||
node_type=node_type,
|
||||
key_prop=key_prop,
|
||||
key_value=key_value,
|
||||
sketch_id=self._sketch_id,
|
||||
**properties
|
||||
)
|
||||
|
||||
def create_relationship(
|
||||
self,
|
||||
from_type: str,
|
||||
from_key: str,
|
||||
from_value: str,
|
||||
to_type: str,
|
||||
to_key: str,
|
||||
to_value: str,
|
||||
rel_type: str,
|
||||
**properties: Any
|
||||
) -> None:
|
||||
"""
|
||||
Create a relationship between two nodes.
|
||||
|
||||
Args:
|
||||
from_type: Source node label
|
||||
from_key: Source node key property
|
||||
from_value: Source node key value
|
||||
to_type: Target node label
|
||||
to_key: Target node key property
|
||||
to_value: Target node key value
|
||||
rel_type: Relationship type
|
||||
**properties: Additional relationship properties
|
||||
"""
|
||||
if self._enable_batching:
|
||||
self._repository.add_to_batch(
|
||||
"relationship",
|
||||
from_type=from_type,
|
||||
from_key=from_key,
|
||||
from_value=from_value,
|
||||
to_type=to_type,
|
||||
to_key=to_key,
|
||||
to_value=to_value,
|
||||
rel_type=rel_type,
|
||||
sketch_id=self._sketch_id,
|
||||
**properties
|
||||
)
|
||||
else:
|
||||
self._repository.create_relationship(
|
||||
from_type=from_type,
|
||||
from_key=from_key,
|
||||
from_value=from_value,
|
||||
to_type=to_type,
|
||||
to_key=to_key,
|
||||
to_value=to_value,
|
||||
rel_type=rel_type,
|
||||
sketch_id=self._sketch_id,
|
||||
**properties
|
||||
)
|
||||
|
||||
def log_graph_message(self, message: str) -> None:
|
||||
"""
|
||||
Log a graph operation message.
|
||||
|
||||
Args:
|
||||
message: Message to log
|
||||
"""
|
||||
if self._logger:
|
||||
self._logger.graph_append(self._sketch_id, {"message": message})
|
||||
|
||||
def flush(self) -> None:
|
||||
"""Flush any pending batch operations."""
|
||||
if self._enable_batching:
|
||||
self._repository.flush_batch()
|
||||
|
||||
def query(self, cypher: str, parameters: Dict[str, Any] = None) -> list:
|
||||
"""
|
||||
Execute a custom Cypher query.
|
||||
|
||||
Args:
|
||||
cypher: Cypher query string
|
||||
parameters: Query parameters
|
||||
|
||||
Returns:
|
||||
List of result records
|
||||
"""
|
||||
return self._repository.query(cypher, parameters)
|
||||
|
||||
def set_batch_size(self, size: int) -> None:
|
||||
"""
|
||||
Set the batch size for operations.
|
||||
|
||||
Args:
|
||||
size: Number of operations to batch
|
||||
"""
|
||||
self._repository.set_batch_size(size)
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry."""
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit - auto-flush batch."""
|
||||
if exc_type is None:
|
||||
self.flush()
|
||||
|
||||
|
||||
def create_graph_service(
|
||||
sketch_id: str,
|
||||
neo4j_connection: Optional[Neo4jConnection] = None,
|
||||
enable_batching: bool = True
|
||||
) -> GraphService:
|
||||
"""
|
||||
Factory function to create a GraphService instance.
|
||||
|
||||
Args:
|
||||
sketch_id: Investigation sketch ID
|
||||
neo4j_connection: Optional Neo4j connection
|
||||
enable_batching: Enable batch operations
|
||||
|
||||
Returns:
|
||||
Configured GraphService instance
|
||||
"""
|
||||
# Import Logger here to avoid circular imports
|
||||
from .logger import Logger
|
||||
|
||||
return GraphService(
|
||||
sketch_id=sketch_id,
|
||||
neo4j_connection=neo4j_connection,
|
||||
logger=Logger,
|
||||
enable_batching=enable_batching
|
||||
)
|
||||
@@ -5,7 +5,9 @@ from pydantic.config import ConfigDict
|
||||
from .graph_db import Neo4jConnection
|
||||
from .logger import Logger
|
||||
from .vault import VaultProtocol
|
||||
from .graph_service import GraphService, create_graph_service
|
||||
from ..utils import resolve_type
|
||||
import warnings
|
||||
|
||||
|
||||
class InvalidTransformParams(Exception):
|
||||
@@ -112,14 +114,27 @@ class Transform(ABC):
|
||||
params_schema: Optional[List[Dict[str, Any]]] = None,
|
||||
vault: Optional[VaultProtocol] = None,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
graph_service: Optional[GraphService] = None,
|
||||
):
|
||||
self.scan_id = scan_id or "default"
|
||||
self.sketch_id = sketch_id or "system"
|
||||
self.neo4j_conn = neo4j_conn
|
||||
self.neo4j_conn = neo4j_conn # Kept for backward compatibility
|
||||
self.vault = vault
|
||||
self.params_schema = params_schema or []
|
||||
self.ParamsModel = build_params_model(self.params_schema)
|
||||
self.params: Dict[str, Any] = params or {}
|
||||
|
||||
# Initialize graph service (new architecture)
|
||||
if graph_service:
|
||||
self._graph_service = graph_service
|
||||
else:
|
||||
# Create graph service with the provided or singleton connection
|
||||
self._graph_service = create_graph_service(
|
||||
sketch_id=self.sketch_id,
|
||||
neo4j_connection=neo4j_conn,
|
||||
enable_batching=True
|
||||
)
|
||||
|
||||
# Params is filled synchronously by the constructor. This params is generally constructed of
|
||||
# vaultSecret references, not the key directly. The idea is that the real key values are resolved after calling
|
||||
# async_init(), right before the execution.
|
||||
@@ -356,6 +371,9 @@ class Transform(ABC):
|
||||
results = await self.scan(preprocessed)
|
||||
processed = self.postprocess(results, preprocessed)
|
||||
|
||||
# Flush any pending batch operations
|
||||
self._graph_service.flush()
|
||||
|
||||
if self.name() != "transform_orchestrator":
|
||||
Logger.completed(
|
||||
self.sketch_id, {"message": f"Transform {self.name()} finished."}
|
||||
@@ -374,90 +392,40 @@ class Transform(ABC):
|
||||
def create_node(
|
||||
self, node_type: str, key_prop: str, key_value: str, **properties
|
||||
) -> None:
|
||||
"""Simple helper to create a single Neo4j node."""
|
||||
if not self.neo4j_conn:
|
||||
return
|
||||
|
||||
# Serialize properties to handle nested Pydantic objects
|
||||
serialized_properties = self._serialize_properties(properties)
|
||||
|
||||
# Ensure all values are Neo4j-compatible primitive types
|
||||
final_properties = {}
|
||||
for key, value in serialized_properties.items():
|
||||
if value is None:
|
||||
final_properties[key] = ""
|
||||
elif isinstance(value, (str, int, float, bool)):
|
||||
final_properties[key] = "" if value == "None" else value
|
||||
else:
|
||||
# Convert any remaining complex types to strings
|
||||
final_properties[key] = str(value)
|
||||
|
||||
final_properties["type"] = node_type.lower()
|
||||
final_properties["sketch_id"] = self.sketch_id
|
||||
final_properties["label"] = final_properties.get("label", key_value)
|
||||
|
||||
set_clauses = [f"n.{prop} = ${prop}" for prop in final_properties.keys()]
|
||||
params = {key_prop: key_value, **final_properties}
|
||||
|
||||
query = f"""
|
||||
MERGE (n:{node_type} {{{key_prop}: ${key_prop}}})
|
||||
SET {', '.join(set_clauses)}
|
||||
"""
|
||||
self.neo4j_conn.query(query, params)
|
||||
Create a single Neo4j node.
|
||||
|
||||
This method now uses the GraphService for improved performance and
|
||||
better separation of concerns.
|
||||
|
||||
Args:
|
||||
node_type: Node label (e.g., "domain", "ip")
|
||||
key_prop: Property name used as unique identifier
|
||||
key_value: Value of the key property
|
||||
**properties: Additional node properties
|
||||
"""
|
||||
self._graph_service.create_node(
|
||||
node_type=node_type,
|
||||
key_prop=key_prop,
|
||||
key_value=key_value,
|
||||
**properties
|
||||
)
|
||||
|
||||
def _serialize_properties(self, properties: dict) -> dict:
|
||||
"""Convert properties to Neo4j-compatible values, handling nested Pydantic objects."""
|
||||
serialized = {}
|
||||
"""
|
||||
Convert properties to Neo4j-compatible values.
|
||||
|
||||
for key, value in properties.items():
|
||||
DEPRECATED: This method is kept for backward compatibility.
|
||||
New code should use GraphSerializer directly.
|
||||
|
||||
if hasattr(value, "__dict__") and not isinstance(
|
||||
value, (str, int, float, bool)
|
||||
):
|
||||
# Handle Pydantic objects and other complex objects
|
||||
if hasattr(value, "model_dump"): # Pydantic v2
|
||||
# Flatten the Pydantic object into individual properties
|
||||
flattened = value.model_dump()
|
||||
for nested_key, nested_value in flattened.items():
|
||||
if nested_value is not None:
|
||||
serialized[f"{key}_{nested_key}"] = nested_value
|
||||
elif hasattr(value, "dict"): # Pydantic v1
|
||||
# Flatten the Pydantic object into individual properties
|
||||
flattened = value.dict()
|
||||
for nested_key, nested_value in flattened.items():
|
||||
if nested_value is not None:
|
||||
serialized[f"{key}_{nested_key}"] = nested_value
|
||||
else:
|
||||
# For other objects, try to convert to dict or string
|
||||
try:
|
||||
flattened = value.__dict__
|
||||
for nested_key, nested_value in flattened.items():
|
||||
if nested_value is not None:
|
||||
serialized[f"{key}_{nested_key}"] = nested_value
|
||||
except:
|
||||
serialized[key] = str(value)
|
||||
elif isinstance(value, list):
|
||||
# Handle lists - convert all items to primitive types
|
||||
serialized_list = []
|
||||
for item in value:
|
||||
if hasattr(item, "__dict__") and not isinstance(
|
||||
item, (str, int, float, bool)
|
||||
):
|
||||
# Convert complex objects to strings
|
||||
serialized_list.append(str(item))
|
||||
else:
|
||||
serialized_list.append(item)
|
||||
serialized[key] = serialized_list
|
||||
elif isinstance(value, dict):
|
||||
# Handle dictionaries - flatten them
|
||||
for dict_key, dict_value in value.items():
|
||||
if dict_value is not None:
|
||||
serialized[f"{key}_{dict_key}"] = dict_value
|
||||
else:
|
||||
# Keep primitive types as-is
|
||||
serialized[key] = value
|
||||
Args:
|
||||
properties: Dictionary of properties to serialize
|
||||
|
||||
return serialized
|
||||
Returns:
|
||||
Dictionary of serialized properties
|
||||
"""
|
||||
from .graph_serializer import GraphSerializer
|
||||
return GraphSerializer.serialize_properties(properties)
|
||||
|
||||
def create_relationship(
|
||||
self,
|
||||
@@ -469,25 +437,46 @@ class Transform(ABC):
|
||||
to_value: str,
|
||||
rel_type: str,
|
||||
) -> None:
|
||||
"""Simple helper to create a relationship between two nodes."""
|
||||
if not self.neo4j_conn:
|
||||
return
|
||||
|
||||
query = f"""
|
||||
MATCH (from:{from_type} {{{from_key}: $from_value}})
|
||||
MATCH (to:{to_type} {{{to_key}: $to_value}})
|
||||
MERGE (from)-[:{rel_type} {{sketch_id: $sketch_id}}]->(to)
|
||||
"""
|
||||
Create a relationship between two nodes.
|
||||
|
||||
self.neo4j_conn.query(
|
||||
query,
|
||||
{
|
||||
"from_value": from_value,
|
||||
"to_value": to_value,
|
||||
"sketch_id": self.sketch_id,
|
||||
},
|
||||
This method now uses the GraphService for improved performance and
|
||||
better separation of concerns.
|
||||
|
||||
Args:
|
||||
from_type: Source node label
|
||||
from_key: Source node key property
|
||||
from_value: Source node key value
|
||||
to_type: Target node label
|
||||
to_key: Target node key property
|
||||
to_value: Target node key value
|
||||
rel_type: Relationship type
|
||||
"""
|
||||
self._graph_service.create_relationship(
|
||||
from_type=from_type,
|
||||
from_key=from_key,
|
||||
from_value=from_value,
|
||||
to_type=to_type,
|
||||
to_key=to_key,
|
||||
to_value=to_value,
|
||||
rel_type=rel_type
|
||||
)
|
||||
|
||||
def log_graph_message(self, message: str) -> None:
|
||||
"""Simple helper to log a graph message."""
|
||||
Logger.graph_append(self.sketch_id, {"message": message})
|
||||
"""
|
||||
Log a graph operation message.
|
||||
|
||||
Args:
|
||||
message: Message to log
|
||||
"""
|
||||
self._graph_service.log_graph_message(message)
|
||||
|
||||
@property
|
||||
def graph_service(self) -> GraphService:
|
||||
"""
|
||||
Get the graph service instance.
|
||||
|
||||
Returns:
|
||||
GraphService instance for advanced operations
|
||||
"""
|
||||
return self._graph_service
|
||||
|
||||
Reference in New Issue
Block a user