From f775c61731aae908c7253097cc3bb1361002f1d2 Mon Sep 17 00:00:00 2001 From: overcuriousity Date: Wed, 17 Sep 2025 11:08:50 +0200 Subject: [PATCH] iterating on fixes --- app.py | 86 ++++++++++++++++++++++++++++++++------- core/graph_manager.py | 59 +++++++++++++++++++-------- providers/dns_provider.py | 56 ++++++++++++++----------- static/js/graph.js | 19 ++++++++- static/js/main.js | 78 +++++++++++++++++++++++++---------- 5 files changed, 220 insertions(+), 78 deletions(-) diff --git a/app.py b/app.py index b2233ff..1955e36 100644 --- a/app.py +++ b/app.py @@ -35,6 +35,29 @@ def get_user_scanner(): existing_scanner = session_manager.get_session(current_flask_session_id) if existing_scanner: return current_flask_session_id, existing_scanner + else: + print(f"Session {current_flask_session_id} not found in Redis, checking for active sessions...") + + # This prevents creating duplicate sessions when Flask session is lost but Redis session exists + stats = session_manager.get_statistics() + if stats['running_scans'] > 0: + # Get all session keys and find running ones + try: + import redis + redis_client = redis.StrictRedis(db=0, decode_responses=False) + session_keys = redis_client.keys("dnsrecon:session:*") + + for session_key in session_keys: + session_id = session_key.decode('utf-8').split(':')[-1] + scanner = session_manager.get_session(session_id) + if scanner and scanner.status in ['running', 'completed']: + print(f"Reusing active session: {session_id}") + # Update Flask session to match + session['dnsrecon_session_id'] = session_id + session.permanent = True + return session_id, scanner + except Exception as e: + print(f"Error finding active session: {e}") # Create new session if none exists print("Creating new session as none was found...") @@ -44,10 +67,11 @@ def get_user_scanner(): if not new_scanner: raise Exception("Failed to create new scanner session") - # Store in Flask session + # Store in Flask session with explicit settings session['dnsrecon_session_id'] = new_session_id session.permanent = True + print(f"Created new session: {new_session_id}") return new_session_id, new_scanner @app.route('/') @@ -240,29 +264,56 @@ def get_scan_status(): @app.route('/api/graph', methods=['GET']) def get_graph_data(): - """Get current graph data with error handling.""" + """Get current graph data with error handling and proper empty graph structure.""" try: # Get user-specific scanner user_session_id, scanner = get_user_scanner() if not scanner: - # Return empty graph if no scanner + # FIXED: Return proper empty graph structure instead of None + empty_graph = { + 'nodes': [], + 'edges': [], + 'statistics': { + 'node_count': 0, + 'edge_count': 0, + 'creation_time': datetime.now(timezone.utc).isoformat(), + 'last_modified': datetime.now(timezone.utc).isoformat() + } + } return jsonify({ 'success': True, - 'graph': { - 'nodes': [], - 'edges': [], - 'statistics': { - 'node_count': 0, - 'edge_count': 0, - 'creation_time': datetime.now(timezone.utc).isoformat(), - 'last_modified': datetime.now(timezone.utc).isoformat() - } - }, + 'graph': empty_graph, 'user_session_id': user_session_id }) graph_data = scanner.get_graph_data() + + if not graph_data: + graph_data = { + 'nodes': [], + 'edges': [], + 'statistics': { + 'node_count': 0, + 'edge_count': 0, + 'creation_time': datetime.now(timezone.utc).isoformat(), + 'last_modified': datetime.now(timezone.utc).isoformat() + } + } + + # FIXED: Ensure required fields exist + if 'nodes' not in graph_data: + graph_data['nodes'] = [] + if 'edges' not in graph_data: + graph_data['edges'] = [] + if 'statistics' not in graph_data: + graph_data['statistics'] = { + 'node_count': len(graph_data['nodes']), + 'edge_count': len(graph_data['edges']), + 'creation_time': datetime.now(timezone.utc).isoformat(), + 'last_modified': datetime.now(timezone.utc).isoformat() + } + return jsonify({ 'success': True, 'graph': graph_data, @@ -272,13 +323,20 @@ def get_graph_data(): except Exception as e: print(f"ERROR: Exception in get_graph_data endpoint: {e}") traceback.print_exc() + + # FIXED: Return proper error structure with empty graph fallback return jsonify({ 'success': False, 'error': f'Internal server error: {str(e)}', 'fallback_graph': { 'nodes': [], 'edges': [], - 'statistics': {'node_count': 0, 'edge_count': 0} + 'statistics': { + 'node_count': 0, + 'edge_count': 0, + 'creation_time': datetime.now(timezone.utc).isoformat(), + 'last_modified': datetime.now(timezone.utc).isoformat() + } } }), 500 diff --git a/core/graph_manager.py b/core/graph_manager.py index cd84cca..f0629cb 100644 --- a/core/graph_manager.py +++ b/core/graph_manager.py @@ -477,8 +477,13 @@ class GraphManager: } def _get_confidence_distribution(self) -> Dict[str, int]: - """Get distribution of edge confidence scores.""" + """Get distribution of edge confidence scores with empty graph handling.""" distribution = {'high': 0, 'medium': 0, 'low': 0} + + # FIXED: Handle empty graph case + if self.get_edge_count() == 0: + return distribution + for _, _, data in self.graph.edges(data=True): confidence = data.get('confidence_score', 0) if confidence >= 0.8: @@ -490,22 +495,42 @@ class GraphManager: return distribution def get_statistics(self) -> Dict[str, Any]: - """Get comprehensive statistics about the graph.""" - stats = {'basic_metrics': {'total_nodes': self.get_node_count(), - 'total_edges': self.get_edge_count(), - 'creation_time': self.creation_time, - 'last_modified': self.last_modified}, - 'node_type_distribution': {}, 'relationship_type_distribution': {}, - 'confidence_distribution': self._get_confidence_distribution(), - 'provider_distribution': {}} - # Calculate distributions - for node_type in NodeType: - stats['node_type_distribution'][node_type.value] = self.get_nodes_by_type(node_type).__len__() - for _, _, data in self.graph.edges(data=True): - rel_type = data.get('relationship_type', 'unknown') - stats['relationship_type_distribution'][rel_type] = stats['relationship_type_distribution'].get(rel_type, 0) + 1 - provider = data.get('source_provider', 'unknown') - stats['provider_distribution'][provider] = stats['provider_distribution'].get(provider, 0) + 1 + """Get comprehensive statistics about the graph with proper empty graph handling.""" + + # FIXED: Handle empty graph case properly + node_count = self.get_node_count() + edge_count = self.get_edge_count() + + stats = { + 'basic_metrics': { + 'total_nodes': node_count, + 'total_edges': edge_count, + 'creation_time': self.creation_time, + 'last_modified': self.last_modified + }, + 'node_type_distribution': {}, + 'relationship_type_distribution': {}, + 'confidence_distribution': self._get_confidence_distribution(), + 'provider_distribution': {} + } + + # FIXED: Only calculate distributions if we have data + if node_count > 0: + # Calculate node type distributions + for node_type in NodeType: + count = len(self.get_nodes_by_type(node_type)) + if count > 0: # Only include types that exist + stats['node_type_distribution'][node_type.value] = count + + if edge_count > 0: + # Calculate edge distributions + for _, _, data in self.graph.edges(data=True): + rel_type = data.get('relationship_type', 'unknown') + stats['relationship_type_distribution'][rel_type] = stats['relationship_type_distribution'].get(rel_type, 0) + 1 + + provider = data.get('source_provider', 'unknown') + stats['provider_distribution'][provider] = stats['provider_distribution'].get(provider, 0) + 1 + return stats def clear(self) -> None: diff --git a/providers/dns_provider.py b/providers/dns_provider.py index 7abaf58..6a855d9 100644 --- a/providers/dns_provider.py +++ b/providers/dns_provider.py @@ -50,6 +50,7 @@ class DNSProvider(BaseProvider): def query_domain(self, domain: str) -> ProviderResult: """ Query DNS records for the domain to discover relationships and attributes. + FIXED: Now creates separate attributes for each DNS record type. Args: domain: Domain to investigate @@ -62,7 +63,7 @@ class DNSProvider(BaseProvider): result = ProviderResult() - # Query all record types + # Query all record types - each gets its own attribute for record_type in ['A', 'AAAA', 'CNAME', 'MX', 'NS', 'SOA', 'TXT', 'SRV', 'CAA']: try: self._query_record(domain, record_type, result) @@ -97,6 +98,7 @@ class DNSProvider(BaseProvider): response = self.resolver.resolve(reverse_name, 'PTR') self.successful_requests += 1 + ptr_records = [] for ptr_record in response: hostname = str(ptr_record).rstrip('.') @@ -116,16 +118,8 @@ class DNSProvider(BaseProvider): } ) - # Add PTR record as attribute to the IP - result.add_attribute( - target_node=ip, - name='ptr_record', - value=hostname, - attr_type='dns_record', - provider=self.name, - confidence=0.8, - metadata={'ttl': response.ttl} - ) + # Add to PTR records list + ptr_records.append(f"PTR: {hostname}") # Log the relationship discovery self.log_relationship_discovery( @@ -142,6 +136,18 @@ class DNSProvider(BaseProvider): discovery_method="reverse_dns_lookup" ) + # Add PTR records as separate attribute + if ptr_records: + result.add_attribute( + target_node=ip, + name='ptr_records', # Specific name for PTR records + value=ptr_records, + attr_type='dns_record', + provider=self.name, + confidence=0.8, + metadata={'ttl': response.ttl} + ) + except resolver.NXDOMAIN: self.failed_requests += 1 self.logger.logger.debug(f"Reverse DNS lookup failed for {ip}: NXDOMAIN") @@ -155,7 +161,7 @@ class DNSProvider(BaseProvider): def _query_record(self, domain: str, record_type: str, result: ProviderResult) -> None: """ - UPDATED: Query DNS records with minimal formatting - keep raw values. + FIXED: Query DNS records with unique attribute names for each record type. """ try: self.total_requests += 1 @@ -175,16 +181,16 @@ class DNSProvider(BaseProvider): elif record_type == 'SOA': target = str(record.mname).rstrip('.') elif record_type in ['TXT']: - # UPDATED: Keep raw TXT record value + # Keep raw TXT record value txt_value = str(record).strip('"') - dns_records.append(f"TXT: {txt_value}") + dns_records.append(txt_value) # Just the value for TXT continue elif record_type == 'SRV': target = str(record.target).rstrip('.') elif record_type == 'CAA': - # UPDATED: Keep raw CAA record format + # Keep raw CAA record format caa_value = f"{record.flags} {record.tag.decode('utf-8')} \"{record.value.decode('utf-8')}\"" - dns_records.append(f"CAA: {caa_value}") + dns_records.append(caa_value) # Just the value for CAA continue else: target = str(record) @@ -196,7 +202,7 @@ class DNSProvider(BaseProvider): 'value': target, 'ttl': response.ttl } - relationship_type = f"{record_type.lower()}_record" # Raw relationship type + relationship_type = f"{record_type.lower()}_record" confidence = 0.8 # Add relationship @@ -209,8 +215,8 @@ class DNSProvider(BaseProvider): raw_data=raw_data ) - # UPDATED: Keep raw DNS record format - dns_records.append(f"{record_type}: {target}") + # Add target to records list + dns_records.append(target) # Log relationship discovery self.log_relationship_discovery( @@ -222,20 +228,22 @@ class DNSProvider(BaseProvider): discovery_method=f"dns_{record_type.lower()}_record" ) - # Add DNS records as a consolidated attribute (raw format) + # FIXED: Create attribute with specific name for each record type if dns_records: + # Use record type specific attribute name (e.g., 'a_records', 'mx_records', etc.) + attribute_name = f"{record_type.lower()}_records" + result.add_attribute( target_node=domain, - name='dns_records', + name=attribute_name, # UNIQUE name for each record type! value=dns_records, attr_type='dns_record_list', provider=self.name, confidence=0.8, - metadata={'record_types': [record_type]} + metadata={'record_type': record_type, 'ttl': response.ttl} ) except Exception as e: self.failed_requests += 1 self.logger.logger.debug(f"{record_type} record query failed for {domain}: {e}") - raise e - + raise e \ No newline at end of file diff --git a/static/js/graph.js b/static/js/graph.js index 9391150..9273ddf 100644 --- a/static/js/graph.js +++ b/static/js/graph.js @@ -377,6 +377,21 @@ class GraphManager { this.initialize(); } + // Check if we have actual data to display + const hasData = graphData.nodes.length > 0 || graphData.edges.length > 0; + + // Handle placeholder visibility + const placeholder = this.container.querySelector('.graph-placeholder'); + if (placeholder) { + if (hasData) { + placeholder.style.display = 'none'; + } else { + placeholder.style.display = 'flex'; + // Early return if no data to process + return; + } + } + this.largeEntityMembers.clear(); const largeEntityMap = new Map(); @@ -398,11 +413,11 @@ class GraphManager { console.log(`Filtered ${graphData.nodes.length - filteredNodes.length} large entity member nodes from visualization`); - // FIXED: Process nodes with proper certificate coloring + // Process nodes with proper certificate coloring const processedNodes = filteredNodes.map(node => { const processed = this.processNode(node); - // FIXED: Apply certificate-based coloring here in frontend + // Apply certificate-based coloring here in frontend if (node.type === 'domain' && Array.isArray(node.attributes)) { const certInfo = this.analyzeCertificateInfo(node.attributes); diff --git a/static/js/main.js b/static/js/main.js index 72f5fdf..ea96e9e 100644 --- a/static/js/main.js +++ b/static/js/main.js @@ -35,6 +35,9 @@ class DNSReconApp { this.loadProviders(); this.initializeEnhancedModals(); + // FIXED: Force initial graph update to handle empty sessions properly + this.updateGraph(); + console.log('DNSRecon application initialized successfully'); } catch (error) { console.error('Failed to initialize DNSRecon application:', error); @@ -42,7 +45,7 @@ class DNSReconApp { } }); } - + /** * Initialize DOM element references */ @@ -484,9 +487,8 @@ class DNSReconApp { console.log('- Nodes:', graphData.nodes ? graphData.nodes.length : 0); console.log('- Edges:', graphData.edges ? graphData.edges.length : 0); - // Only update if data has changed - if (this.hasGraphChanged(graphData)) { - console.log('*** GRAPH DATA CHANGED - UPDATING VISUALIZATION ***'); + // FIXED: Always update graph, even if empty - let GraphManager handle placeholder + if (this.graphManager) { this.graphManager.updateGraph(graphData); this.lastGraphUpdate = Date.now(); @@ -495,18 +497,30 @@ class DNSReconApp { if (this.elements.relationshipsDisplay) { this.elements.relationshipsDisplay.textContent = edgeCount; } - } else { - console.log('Graph data unchanged, skipping update'); } } else { console.error('Graph update failed:', response); + // FIXED: Show placeholder when graph update fails + if (this.graphManager && this.graphManager.container) { + const placeholder = this.graphManager.container.querySelector('.graph-placeholder'); + if (placeholder) { + placeholder.style.display = 'flex'; + } + } } } catch (error) { console.error('Failed to update graph:', error); - // Don't show error for graph updates to avoid spam + // FIXED: Show placeholder on error + if (this.graphManager && this.graphManager.container) { + const placeholder = this.graphManager.container.querySelector('.graph-placeholder'); + if (placeholder) { + placeholder.style.display = 'flex'; + } + } } } + /** * Update status display elements @@ -925,18 +939,32 @@ class DNSReconApp { return 'Empty Array'; } - // Special handling for DNS records and similar arrays - if (name === 'dns_records' || name.includes('record') || name.includes('hostname')) { - // Show DNS records as a readable list + // ENHANCED: Special handling for specific DNS record types + if (name.endsWith('_records') || name.includes('record')) { + const recordType = name.replace('_records', '').toUpperCase(); + + // Format nicely for DNS records if (value.length <= 5) { - return this.escapeHtml(value.join(', ')); + const formattedRecords = value.map(record => { + // Add record type prefix if not already present + if (recordType !== 'DNS' && !record.includes(':')) { + return `${recordType}: ${record}`; + } + return record; + }); + return this.escapeHtml(formattedRecords.join('\n')); } else { - const preview = value.slice(0, 3).join(', '); - return this.escapeHtml(`${preview} ... (+${value.length - 3} more)`); + const preview = value.slice(0, 3).map(record => { + if (recordType !== 'DNS' && !record.includes(':')) { + return `${recordType}: ${record}`; + } + return record; + }).join('\n'); + return this.escapeHtml(`${preview}\n... (+${value.length - 3} more ${recordType} records)`); } } - // For other arrays + // For other arrays (existing logic) if (value.length <= 3) { return this.escapeHtml(value.join(', ')); } else { @@ -953,6 +981,10 @@ class DNSReconApp { } groupAttributesByProviderAndType(attributes, nodeType) { + if (!Array.isArray(attributes) || attributes.length === 0) { + return {}; + } + const groups = { 'DNS Records': { icon: '📋', priority: 'high', attributes: [] }, 'Certificate Information': { icon: '🔒', priority: 'high', attributes: [] }, @@ -968,11 +1000,11 @@ class DNSReconApp { let assigned = false; - // DNS-related attributes (better detection) + // ENHANCED: Better DNS record detection for specific record types if (provider === 'dns' || - name === 'dns_records' || + name.endsWith('_records') || // Catches a_records, mx_records, txt_records, etc. name.includes('record') || - ['ptr', 'mx', 'cname', 'ns', 'txt', 'soa'].some(keyword => name.includes(keyword))) { + ['ptr', 'mx', 'cname', 'ns', 'txt', 'soa', 'srv', 'caa', 'a_records', 'aaaa_records'].some(keyword => name.includes(keyword))) { groups['DNS Records'].attributes.push(attr); assigned = true; } @@ -1013,7 +1045,6 @@ class DNSReconApp { formatEdgeLabel(relationshipType, confidence) { if (!relationshipType) return ''; - // UPDATED: No formatting of relationship type - use raw values const confidenceText = confidence >= 0.8 ? '●' : confidence >= 0.6 ? '◐' : '○'; return `${relationshipType} ${confidenceText}`; } @@ -1695,13 +1726,18 @@ class DNSReconApp { const newNodeCount = graphData.nodes ? graphData.nodes.length : 0; const newEdgeCount = graphData.edges ? graphData.edges.length : 0; + // FIXED: Always update if we currently have no data (ensures placeholder is handled correctly) + if (currentStats.nodeCount === 0 && currentStats.edgeCount === 0) { + return true; + } + // Check if counts changed const countsChanged = currentStats.nodeCount !== newNodeCount || currentStats.edgeCount !== newEdgeCount; // Also check if we have new timestamp data const hasNewTimestamp = graphData.statistics && - graphData.statistics.last_modified && - graphData.statistics.last_modified !== this.lastGraphTimestamp; + graphData.statistics.last_modified && + graphData.statistics.last_modified !== this.lastGraphTimestamp; if (hasNewTimestamp) { this.lastGraphTimestamp = graphData.statistics.last_modified; @@ -1713,7 +1749,7 @@ class DNSReconApp { return changed; } - + /** * Make API call to server * @param {string} endpoint - API endpoint