large entity real fix
This commit is contained in:
overcuriousity
parent
b2c5d2331c
commit
b2629de055
175
core/scanner.py
175
core/scanner.py
@ -824,18 +824,34 @@ class Scanner:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
def _create_large_entity_from_result(self, source_node: str, provider_name: str,
|
def _create_large_entity_from_result(self, source_node: str, provider_name: str,
|
||||||
provider_result: ProviderResult, depth: int) -> Tuple[str, Set[str]]:
|
provider_result: ProviderResult, depth: int) -> Tuple[str, Set[str]]:
|
||||||
"""
|
"""
|
||||||
Creates a large entity node, tags all member nodes, and returns its ID and members.
|
Creates a large entity node, tags all member nodes, and stores original relationships.
|
||||||
|
FIXED: Now stores original relationships for later restoration during extraction.
|
||||||
"""
|
"""
|
||||||
members = {rel.target_node for rel in provider_result.relationships
|
members = {rel.target_node for rel in provider_result.relationships
|
||||||
if _is_valid_domain(rel.target_node) or _is_valid_ip(rel.target_node)}
|
if _is_valid_domain(rel.target_node) or _is_valid_ip(rel.target_node)}
|
||||||
|
|
||||||
if not members:
|
if not members:
|
||||||
return "", set()
|
return "", set()
|
||||||
|
|
||||||
large_entity_id = f"le_{provider_name}_{source_node}"
|
large_entity_id = f"le_{provider_name}_{source_node}"
|
||||||
|
|
||||||
|
# FIXED: Store original relationships for each member
|
||||||
|
member_relationships = {}
|
||||||
|
for rel in provider_result.relationships:
|
||||||
|
if rel.target_node in members:
|
||||||
|
if rel.target_node not in member_relationships:
|
||||||
|
member_relationships[rel.target_node] = []
|
||||||
|
member_relationships[rel.target_node].append({
|
||||||
|
'source_node': rel.source_node,
|
||||||
|
'target_node': rel.target_node,
|
||||||
|
'relationship_type': rel.relationship_type,
|
||||||
|
'confidence': rel.confidence,
|
||||||
|
'provider': rel.provider,
|
||||||
|
'raw_data': rel.raw_data
|
||||||
|
})
|
||||||
|
|
||||||
self.graph.add_node(
|
self.graph.add_node(
|
||||||
node_id=large_entity_id,
|
node_id=large_entity_id,
|
||||||
node_type=NodeType.LARGE_ENTITY,
|
node_type=NodeType.LARGE_ENTITY,
|
||||||
@ -843,7 +859,8 @@ class Scanner:
|
|||||||
{"name": "count", "value": len(members), "type": "statistic"},
|
{"name": "count", "value": len(members), "type": "statistic"},
|
||||||
{"name": "source_provider", "value": provider_name, "type": "metadata"},
|
{"name": "source_provider", "value": provider_name, "type": "metadata"},
|
||||||
{"name": "discovery_depth", "value": depth, "type": "metadata"},
|
{"name": "discovery_depth", "value": depth, "type": "metadata"},
|
||||||
{"name": "nodes", "value": list(members), "type": "metadata"}
|
{"name": "nodes", "value": list(members), "type": "metadata"},
|
||||||
|
{"name": "original_relationships", "value": member_relationships, "type": "metadata"} # FIXED: Store original relationships
|
||||||
],
|
],
|
||||||
description=f"A collection of {len(members)} nodes discovered from {source_node} via {provider_name}."
|
description=f"A collection of {len(members)} nodes discovered from {source_node} via {provider_name}."
|
||||||
)
|
)
|
||||||
@ -860,7 +877,8 @@ class Scanner:
|
|||||||
|
|
||||||
def extract_node_from_large_entity(self, large_entity_id: str, node_id: str) -> bool:
|
def extract_node_from_large_entity(self, large_entity_id: str, node_id: str) -> bool:
|
||||||
"""
|
"""
|
||||||
FIXED: Extract a node from a large entity with proper backend updates and edge re-routing.
|
Removes a node from a large entity and restores its original relationships.
|
||||||
|
FIXED: Now restores original relationships to make the node reachable.
|
||||||
"""
|
"""
|
||||||
if not self.graph.graph.has_node(node_id):
|
if not self.graph.graph.has_node(node_id):
|
||||||
return False
|
return False
|
||||||
@ -871,51 +889,45 @@ class Scanner:
|
|||||||
if metadata.get('large_entity_id') != large_entity_id:
|
if metadata.get('large_entity_id') != large_entity_id:
|
||||||
return False
|
return False
|
||||||
|
|
||||||
# FIXED: Update the large entity's attributes to remove the extracted node
|
# Remove the large entity tag
|
||||||
if self.graph.graph.has_node(large_entity_id):
|
del metadata['large_entity_id']
|
||||||
le_node_data = self.graph.graph.nodes[large_entity_id]
|
self.graph.add_node(node_id, NodeType(node_data['type']), metadata=metadata)
|
||||||
le_attributes = le_node_data.get('attributes', [])
|
|
||||||
|
|
||||||
# Update the 'nodes' attribute to remove extracted node
|
# FIXED: Restore original relationships if they exist
|
||||||
nodes_attr = next((attr for attr in le_attributes if attr['name'] == 'nodes'), None)
|
if self.graph.graph.has_node(large_entity_id):
|
||||||
if nodes_attr and isinstance(nodes_attr['value'], list):
|
le_attrs = self.graph.graph.nodes[large_entity_id].get('attributes', [])
|
||||||
if node_id in nodes_attr['value']:
|
original_relationships_attr = next((a for a in le_attrs if a['name'] == 'original_relationships'), None)
|
||||||
|
|
||||||
|
if original_relationships_attr and node_id in original_relationships_attr['value']:
|
||||||
|
# Restore all original relationships for this node
|
||||||
|
for rel_data in original_relationships_attr['value'][node_id]:
|
||||||
|
self.graph.add_edge(
|
||||||
|
source_id=rel_data['source_node'],
|
||||||
|
target_id=rel_data['target_node'],
|
||||||
|
relationship_type=rel_data['relationship_type'],
|
||||||
|
confidence_score=rel_data['confidence'],
|
||||||
|
source_provider=rel_data['provider'],
|
||||||
|
raw_data=rel_data['raw_data']
|
||||||
|
)
|
||||||
|
|
||||||
|
# Ensure both nodes exist in the graph
|
||||||
|
source_type = NodeType.IP if _is_valid_ip(rel_data['source_node']) else NodeType.DOMAIN
|
||||||
|
target_type = NodeType.IP if _is_valid_ip(rel_data['target_node']) else NodeType.DOMAIN
|
||||||
|
self.graph.add_node(rel_data['source_node'], source_type)
|
||||||
|
self.graph.add_node(rel_data['target_node'], target_type)
|
||||||
|
|
||||||
|
# Update the large entity to remove this node from its list
|
||||||
|
nodes_attr = next((a for a in le_attrs if a['name'] == 'nodes'), None)
|
||||||
|
if nodes_attr and node_id in nodes_attr['value']:
|
||||||
nodes_attr['value'].remove(node_id)
|
nodes_attr['value'].remove(node_id)
|
||||||
|
|
||||||
# Update the 'count' attribute
|
count_attr = next((a for a in le_attrs if a['name'] == 'count'), None)
|
||||||
count_attr = next((attr for attr in le_attributes if attr['name'] == 'count'), None)
|
if count_attr:
|
||||||
if count_attr and isinstance(count_attr['value'], (int, float)):
|
count_attr['value'] = max(0, count_attr['value'] - 1)
|
||||||
count_attr['value'] = max(0, count_attr['value'] - 1)
|
|
||||||
|
|
||||||
# Update the large entity node
|
# Remove from original relationships tracking
|
||||||
self.graph.add_node(
|
if node_id in original_relationships_attr['value']:
|
||||||
large_entity_id,
|
del original_relationships_attr['value'][node_id]
|
||||||
NodeType.LARGE_ENTITY,
|
|
||||||
attributes=le_attributes,
|
|
||||||
description=le_node_data.get('description', ''),
|
|
||||||
metadata=le_node_data.get('metadata', {})
|
|
||||||
)
|
|
||||||
|
|
||||||
# Remove the large entity tag from extracted node
|
|
||||||
updated_metadata = metadata.copy()
|
|
||||||
del updated_metadata['large_entity_id']
|
|
||||||
|
|
||||||
# Add extraction history for forensic integrity
|
|
||||||
extraction_record = {
|
|
||||||
'extracted_at': datetime.now(timezone.utc).isoformat(),
|
|
||||||
'extracted_from': large_entity_id,
|
|
||||||
'extraction_method': 'manual'
|
|
||||||
}
|
|
||||||
|
|
||||||
if 'extraction_history' not in updated_metadata:
|
|
||||||
updated_metadata['extraction_history'] = []
|
|
||||||
updated_metadata['extraction_history'].append(extraction_record)
|
|
||||||
|
|
||||||
# Update the extracted node
|
|
||||||
self.graph.add_node(node_id, NodeType(node_data['type']), metadata=updated_metadata)
|
|
||||||
|
|
||||||
# FIXED: Re-route edges that were pointing to the large entity
|
|
||||||
self._reroute_large_entity_edges(large_entity_id, node_id)
|
|
||||||
|
|
||||||
# Re-enqueue the node for full processing
|
# Re-enqueue the node for full processing
|
||||||
is_ip = _is_valid_ip(node_id)
|
is_ip = _is_valid_ip(node_id)
|
||||||
@ -923,7 +935,6 @@ class Scanner:
|
|||||||
for provider in eligible_providers:
|
for provider in eligible_providers:
|
||||||
provider_name = provider.get_name()
|
provider_name = provider.get_name()
|
||||||
priority = self._get_priority(provider_name)
|
priority = self._get_priority(provider_name)
|
||||||
|
|
||||||
# Use current depth of the large entity if available, else 0
|
# Use current depth of the large entity if available, else 0
|
||||||
depth = 0
|
depth = 0
|
||||||
if self.graph.graph.has_node(large_entity_id):
|
if self.graph.graph.has_node(large_entity_id):
|
||||||
@ -935,78 +946,8 @@ class Scanner:
|
|||||||
self.task_queue.put((time.time(), priority, (provider_name, node_id, depth)))
|
self.task_queue.put((time.time(), priority, (provider_name, node_id, depth)))
|
||||||
self.total_tasks_ever_enqueued += 1
|
self.total_tasks_ever_enqueued += 1
|
||||||
|
|
||||||
# Force session state update for immediate frontend sync
|
|
||||||
self._update_session_state()
|
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
def _reroute_large_entity_edges(self, large_entity_id: str, extracted_node_id: str) -> None:
|
|
||||||
"""
|
|
||||||
FIXED: Re-route edges from large entity to extracted node where appropriate.
|
|
||||||
"""
|
|
||||||
if not self.graph.graph.has_node(large_entity_id) or not self.graph.graph.has_node(extracted_node_id):
|
|
||||||
return
|
|
||||||
|
|
||||||
edges_to_reroute = []
|
|
||||||
|
|
||||||
# Find edges pointing TO the large entity that should point to the extracted node
|
|
||||||
for source, target, edge_data in self.graph.graph.in_edges(large_entity_id, data=True):
|
|
||||||
# Check if this edge was originally meant for the extracted node
|
|
||||||
raw_data = edge_data.get('raw_data', {})
|
|
||||||
|
|
||||||
# If the raw data suggests this edge was for the extracted node, re-route it
|
|
||||||
if (raw_data.get('original_target') == extracted_node_id or
|
|
||||||
self._should_reroute_edge(edge_data, extracted_node_id)):
|
|
||||||
edges_to_reroute.append(('in', source, target, edge_data))
|
|
||||||
|
|
||||||
# Find edges pointing FROM the large entity that should point from the extracted node
|
|
||||||
for source, target, edge_data in self.graph.graph.out_edges(large_entity_id, data=True):
|
|
||||||
raw_data = edge_data.get('raw_data', {})
|
|
||||||
|
|
||||||
if (raw_data.get('original_source') == extracted_node_id or
|
|
||||||
self._should_reroute_edge(edge_data, extracted_node_id)):
|
|
||||||
edges_to_reroute.append(('out', source, target, edge_data))
|
|
||||||
|
|
||||||
# Re-route the edges
|
|
||||||
for direction, source, target, edge_data in edges_to_reroute:
|
|
||||||
# Remove old edge
|
|
||||||
self.graph.graph.remove_edge(source, target)
|
|
||||||
|
|
||||||
# Add new edge with extracted node
|
|
||||||
if direction == 'in':
|
|
||||||
new_target = extracted_node_id
|
|
||||||
new_source = source
|
|
||||||
else: # direction == 'out'
|
|
||||||
new_source = extracted_node_id
|
|
||||||
new_target = target
|
|
||||||
|
|
||||||
# Add the re-routed edge
|
|
||||||
self.graph.add_edge(
|
|
||||||
source_id=new_source,
|
|
||||||
target_id=new_target,
|
|
||||||
relationship_type=edge_data.get('relationship_type', 'unknown'),
|
|
||||||
confidence_score=edge_data.get('confidence_score', 0.5),
|
|
||||||
source_provider=edge_data.get('source_provider', 'rerouted'),
|
|
||||||
raw_data=dict(edge_data.get('raw_data', {}), **{'rerouted_from_large_entity': large_entity_id})
|
|
||||||
)
|
|
||||||
|
|
||||||
def _should_reroute_edge(self, edge_data: dict, extracted_node_id: str) -> bool:
|
|
||||||
"""
|
|
||||||
Determine if an edge should be re-routed to an extracted node.
|
|
||||||
This is a heuristic-based approach since we don't store original targets.
|
|
||||||
"""
|
|
||||||
relationship_type = edge_data.get('relationship_type', '')
|
|
||||||
|
|
||||||
# For now, re-route DNS and certificate-based relationships
|
|
||||||
# These are likely to be node-specific rather than entity-wide
|
|
||||||
reroutable_types = [
|
|
||||||
'dns_a_record', 'dns_aaaa_record', 'dns_cname_record',
|
|
||||||
'dns_mx_record', 'dns_ptr_record',
|
|
||||||
'crtsh_san_certificate', 'crtsh_cert_issuer'
|
|
||||||
]
|
|
||||||
|
|
||||||
return any(rtype in relationship_type for rtype in reroutable_types)
|
|
||||||
|
|
||||||
def _process_provider_result_unified(self, target: str, provider: BaseProvider,
|
def _process_provider_result_unified(self, target: str, provider: BaseProvider,
|
||||||
provider_result: ProviderResult, current_depth: int) -> Tuple[Set[str], bool]:
|
provider_result: ProviderResult, current_depth: int) -> Tuple[Set[str], bool]:
|
||||||
"""
|
"""
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user