attempt to fix some logic
This commit is contained in:
overcuriousity
@@ -41,7 +41,7 @@ class GraphManager:
|
||||
self.correlation_index = {}
|
||||
# Compile regex for date filtering for efficiency
|
||||
self.date_pattern = re.compile(r'^\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2}')
|
||||
self.EXCLUDED_KEYS = ['confidence', 'provider', 'timestamp', 'type','crtsh_cert_validity_period_days']
|
||||
self.EXCLUDED_KEYS = ['confidence', 'provider', 'timestamp', 'type','crtsh_cert_validity_period_days','crtsh_cert_source']
|
||||
|
||||
def __getstate__(self):
|
||||
"""Prepare GraphManager for pickling, excluding compiled regex."""
|
||||
@@ -112,7 +112,7 @@ class GraphManager:
|
||||
|
||||
def _create_enhanced_correlation_node_and_edges(self, value, correlation_data):
|
||||
"""
|
||||
UPDATED: Create correlation node and edges with detailed provider tracking.
|
||||
UPDATED: Create correlation node and edges with raw provider data (no formatting).
|
||||
"""
|
||||
correlation_node_id = f"corr_{hash(str(value)) & 0x7FFFFFFF}"
|
||||
nodes = correlation_data['nodes']
|
||||
@@ -120,13 +120,14 @@ class GraphManager:
|
||||
|
||||
# Create or update correlation node
|
||||
if not self.graph.has_node(correlation_node_id):
|
||||
# Determine the most common provider/attribute combination
|
||||
# Use raw provider/attribute data - no formatting
|
||||
provider_counts = {}
|
||||
for source in sources:
|
||||
# Keep original provider and attribute names
|
||||
key = f"{source['provider']}_{source['attribute']}"
|
||||
provider_counts[key] = provider_counts.get(key, 0) + 1
|
||||
|
||||
# Use the most common provider/attribute as the primary label
|
||||
# Use the most common provider/attribute as the primary label (raw)
|
||||
primary_source = max(provider_counts.items(), key=lambda x: x[1])[0] if provider_counts else "unknown_correlation"
|
||||
|
||||
metadata = {
|
||||
@@ -303,18 +304,18 @@ class GraphManager:
|
||||
return is_new_node
|
||||
|
||||
def add_edge(self, source_id: str, target_id: str, relationship_type: str,
|
||||
confidence_score: float = 0.5, source_provider: str = "unknown",
|
||||
raw_data: Optional[Dict[str, Any]] = None) -> bool:
|
||||
"""Add or update an edge between two nodes, ensuring nodes exist."""
|
||||
confidence_score: float = 0.5, source_provider: str = "unknown",
|
||||
raw_data: Optional[Dict[str, Any]] = None) -> bool:
|
||||
"""
|
||||
UPDATED: Add or update an edge between two nodes with raw relationship labels.
|
||||
"""
|
||||
if not self.graph.has_node(source_id) or not self.graph.has_node(target_id):
|
||||
return False
|
||||
|
||||
new_confidence = confidence_score
|
||||
|
||||
if relationship_type.startswith("c_"):
|
||||
edge_label = relationship_type
|
||||
else:
|
||||
edge_label = f"{source_provider}_{relationship_type}"
|
||||
# UPDATED: Use raw relationship type - no formatting
|
||||
edge_label = relationship_type
|
||||
|
||||
if self.graph.has_edge(source_id, target_id):
|
||||
# If edge exists, update confidence if the new score is higher.
|
||||
@@ -324,7 +325,7 @@ class GraphManager:
|
||||
self.graph.edges[source_id, target_id]['updated_by'] = source_provider
|
||||
return False
|
||||
|
||||
# Add a new edge with all attributes.
|
||||
# Add a new edge with raw attributes
|
||||
self.graph.add_edge(source_id, target_id,
|
||||
relationship_type=edge_label,
|
||||
confidence_score=new_confidence,
|
||||
@@ -333,7 +334,7 @@ class GraphManager:
|
||||
raw_data=raw_data or {})
|
||||
self.last_modified = datetime.now(timezone.utc).isoformat()
|
||||
return True
|
||||
|
||||
|
||||
def extract_node_from_large_entity(self, large_entity_id: str, node_id_to_extract: str) -> bool:
|
||||
"""
|
||||
Removes a node from a large entity's internal lists and updates its count.
|
||||
@@ -417,73 +418,45 @@ class GraphManager:
|
||||
def get_graph_data(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Export graph data formatted for frontend visualization.
|
||||
UPDATED: Fixed certificate validity styling logic for unified data model.
|
||||
SIMPLIFIED: No certificate styling - frontend handles all visual styling.
|
||||
"""
|
||||
nodes = []
|
||||
for node_id, attrs in self.graph.nodes(data=True):
|
||||
node_data = {'id': node_id, 'label': node_id, 'type': attrs.get('type', 'unknown'),
|
||||
'attributes': attrs.get('attributes', []), # Ensure attributes is a list
|
||||
'description': attrs.get('description', ''),
|
||||
'metadata': attrs.get('metadata', {}),
|
||||
'added_timestamp': attrs.get('added_timestamp')}
|
||||
|
||||
# UPDATED: Fixed certificate validity styling logic
|
||||
node_type = node_data['type']
|
||||
attributes_list = node_data['attributes']
|
||||
|
||||
if node_type == 'domain' and isinstance(attributes_list, list):
|
||||
# Check for certificate-related attributes
|
||||
has_certificates = False
|
||||
has_valid_certificates = False
|
||||
has_expired_certificates = False
|
||||
|
||||
for attr in attributes_list:
|
||||
attr_name = attr.get('name', '').lower()
|
||||
attr_provider = attr.get('provider', '').lower()
|
||||
attr_value = attr.get('value')
|
||||
|
||||
# Look for certificate attributes from crt.sh provider
|
||||
if attr_provider == 'crtsh' or 'cert' in attr_name:
|
||||
has_certificates = True
|
||||
|
||||
# Check certificate validity
|
||||
if attr_name == 'cert_is_currently_valid':
|
||||
if attr_value is True:
|
||||
has_valid_certificates = True
|
||||
elif attr_value is False:
|
||||
has_expired_certificates = True
|
||||
|
||||
# Also check for certificate expiry indicators
|
||||
elif 'expires_soon' in attr_name and attr_value is True:
|
||||
has_expired_certificates = True
|
||||
elif 'expired' in attr_name and attr_value is True:
|
||||
has_expired_certificates = True
|
||||
|
||||
# Apply styling based on certificate status
|
||||
if has_expired_certificates and not has_valid_certificates:
|
||||
# Red for expired/invalid certificates
|
||||
node_data['color'] = {'background': '#ff6b6b', 'border': '#cc5555'}
|
||||
elif not has_certificates:
|
||||
# Grey for domains with no certificates
|
||||
node_data['color'] = {'background': '#c7c7c7', 'border': '#999999'}
|
||||
# Default green styling is handled by the frontend for domains with valid certificates
|
||||
node_data = {
|
||||
'id': node_id,
|
||||
'label': node_id,
|
||||
'type': attrs.get('type', 'unknown'),
|
||||
'attributes': attrs.get('attributes', []), # Raw attributes list
|
||||
'description': attrs.get('description', ''),
|
||||
'metadata': attrs.get('metadata', {}),
|
||||
'added_timestamp': attrs.get('added_timestamp')
|
||||
}
|
||||
|
||||
# Add incoming and outgoing edges to node data
|
||||
if self.graph.has_node(node_id):
|
||||
node_data['incoming_edges'] = [{'from': u, 'data': d} for u, _, d in self.graph.in_edges(node_id, data=True)]
|
||||
node_data['outgoing_edges'] = [{'to': v, 'data': d} for _, v, d in self.graph.out_edges(node_id, data=True)]
|
||||
node_data['incoming_edges'] = [
|
||||
{'from': u, 'data': d} for u, _, d in self.graph.in_edges(node_id, data=True)
|
||||
]
|
||||
node_data['outgoing_edges'] = [
|
||||
{'to': v, 'data': d} for _, v, d in self.graph.out_edges(node_id, data=True)
|
||||
]
|
||||
|
||||
nodes.append(node_data)
|
||||
|
||||
edges = []
|
||||
for source, target, attrs in self.graph.edges(data=True):
|
||||
edges.append({'from': source, 'to': target,
|
||||
'label': attrs.get('relationship_type', ''),
|
||||
'confidence_score': attrs.get('confidence_score', 0),
|
||||
'source_provider': attrs.get('source_provider', ''),
|
||||
'discovery_timestamp': attrs.get('discovery_timestamp')})
|
||||
edges.append({
|
||||
'from': source,
|
||||
'to': target,
|
||||
'label': attrs.get('relationship_type', ''),
|
||||
'confidence_score': attrs.get('confidence_score', 0),
|
||||
'source_provider': attrs.get('source_provider', ''),
|
||||
'discovery_timestamp': attrs.get('discovery_timestamp')
|
||||
})
|
||||
|
||||
return {
|
||||
'nodes': nodes, 'edges': edges,
|
||||
'nodes': nodes,
|
||||
'edges': edges,
|
||||
'statistics': self.get_statistics()['basic_metrics']
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user