interactive_vulnerability_chatbot.py

#!/usr/bin/env python3
"""
Interactive Vulnerability Chatbot
Answers questions about vulnerability findings and provides remediation advice
"""

import json
import re
import sys
from pathlib import Path
from typing import Dict, List, Any, Optional
from datetime import datetime

class VulnerabilityChatbot:
    """Interactive chatbot for vulnerability analysis results"""
    
    def __init__(self, knowledge_base_file: str):
        self.knowledge_base = self._load_knowledge_base(knowledge_base_file)
        self.conversation_history = []
        
    def _load_knowledge_base(self, kb_file: str) -> Dict:
        """Load knowledge base from JSON file"""
        try:
            with open(kb_file, 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            print(f"Knowledge base file not found: {kb_file}")
            return self._create_default_kb()
        except json.JSONDecodeError:
            print(f"Invalid JSON in knowledge base file: {kb_file}")
            return self._create_default_kb()
    
    def _create_default_kb(self) -> Dict:
        """Create default knowledge base if file not found"""
        return {
            'target_info': {
                'website': 'example.com',
                'analysis_date': datetime.now().isoformat(),
                'total_vulnerabilities': 0
            },
            'vulnerabilities': {},
            'attack_paths': [],
            'recommendations': ['No analysis data available'],
            'severity_summary': {'Critical': 0, 'High': 0, 'Medium': 0, 'Low': 0}
        }
    
    def process_query(self, query: str) -> str:
        """Process user query and return response"""
        query_lower = query.lower().strip()
        
        # Store query in conversation history
        self.conversation_history.append({
            'timestamp': datetime.now().isoformat(),
            'query': query,
            'type': self._classify_query(query_lower)
        })
        
        # Route to appropriate handler
        if self._is_greeting(query_lower):
            return self._handle_greeting()
        elif self._is_summary_query(query_lower):
            return self._handle_summary_query()
        elif self._is_critical_query(query_lower):
            return self._handle_critical_vulnerabilities()
        elif self._is_cve_query(query_lower):
            return self._handle_cve_query(query_lower)
        elif self._is_attack_path_query(query_lower):
            return self._handle_attack_path_query()
        elif self._is_remediation_query(query_lower):
            return self._handle_remediation_query(query_lower)
        elif self._is_priority_query(query_lower):
            return self._handle_priority_query()
        elif self._is_service_query(query_lower):
            return self._handle_service_query(query_lower)
        elif self._is_help_query(query_lower):
            return self._handle_help_query()
        else:
            return self._handle_general_query(query_lower)
    
    def _classify_query(self, query: str) -> str:
        """Classify the type of query"""
        if any(word in query for word in ['hello', 'hi', 'hey']):
            return 'greeting'
        elif any(word in query for word in ['critical', 'severe', 'urgent']):
            return 'critical'
        elif 'cve-' in query or 'vulnerability' in query:
            return 'vulnerability'
        elif any(word in query for word in ['fix', 'remediate', 'patch', 'solve']):
            return 'remediation'
        elif any(word in query for word in ['attack', 'path', 'exploit']):
            return 'attack_path'
        elif any(word in query for word in ['priority', 'prioritize', 'order']):
            return 'priority'
        elif any(word in query for word in ['summary', 'overview', 'total']):
            return 'summary'
        else:
            return 'general'
    
    def _is_greeting(self, query: str) -> bool:
        return any(word in query for word in ['hello', 'hi', 'hey', 'start'])
    
    def _is_summary_query(self, query: str) -> bool:
        return any(word in query for word in ['summary', 'overview', 'total', 'how many', 'count'])
    
    def _is_critical_query(self, query: str) -> bool:
        return any(word in query for word in ['critical', 'severe', 'urgent', 'worst'])
    
    def _is_cve_query(self, query: str) -> bool:
        return 'cve-' in query or ('vulnerability' in query and any(cve in query for cve in ['log4j', 'apache', 'ssh']))
    
    def _is_attack_path_query(self, query: str) -> bool:
        return any(word in query for word in ['attack', 'path', 'exploit', 'compromise', 'chain'])
    
    def _is_remediation_query(self, query: str) -> bool:
        return any(word in query for word in ['fix', 'remediate', 'patch', 'solve', 'update', 'secure'])
    
    def _is_priority_query(self, query: str) -> bool:
        return any(word in query for word in ['priority', 'prioritize', 'order', 'first', 'important'])
    
    def _is_service_query(self, query: str) -> bool:
        return any(service in query for service in ['http', 'https', 'ssh', 'web', 'server'])
    
    def _is_help_query(self, query: str) -> bool:
        return any(word in query for word in ['help', 'commands', 'what can', 'options'])
    
    def _handle_greeting(self) -> str:
        website = self.knowledge_base['target_info']['website']
        total_vulns = self.knowledge_base['target_info']['total_vulnerabilities']
        
        return f"""Hello! I'm your SecureChain vulnerability analysis assistant.

I have analyzed {website} and found {total_vulns} vulnerabilities.

You can ask me about:
• Critical vulnerabilities and their fixes
• Specific CVEs (e.g., "Tell me about CVE-2021-44228")
• Attack paths and exploitation scenarios
• Remediation priorities and recommendations
• Service-specific vulnerabilities

What would you like to know about the security analysis?"""
    
    def _handle_summary_query(self) -> str:
        target_info = self.knowledge_base['target_info']
        severity_summary = self.knowledge_base['severity_summary']
        
        response = f"""📊 VULNERABILITY ANALYSIS SUMMARY

🎯 Target: {target_info['website']}
📅 Analysis Date: {target_info['analysis_date'][:10]}
🔍 Total Vulnerabilities: {target_info['total_vulnerabilities']}

📈 Severity Breakdown:"""
        
        for severity in ['Critical', 'High', 'Medium', 'Low']:
            count = severity_summary.get(severity, 0)
            if count > 0:
                emoji = {'Critical': '🔴', 'High': '🟠', 'Medium': '🟡', 'Low': '🟢'}[severity]
                response += f"\n{emoji} {severity}: {count} vulnerabilities"
        
        # Add top recommendations
        recommendations = self.knowledge_base.get('recommendations', [])
        if recommendations:
            response += f"\n\n💡 Top Recommendations:\n"
            for i, rec in enumerate(recommendations[:3], 1):
                response += f"{i}. {rec}\n"
        
        return response
    
    def _handle_critical_vulnerabilities(self) -> str:
        critical_vulns = {k: v for k, v in self.knowledge_base['vulnerabilities'].items() 
                         if v.get('severity') == 'Critical'}
        
        if not critical_vulns:
            return "✅ Good news! No critical vulnerabilities were found in the analysis."
        
        response = f"🚨 CRITICAL VULNERABILITIES FOUND ({len(critical_vulns)})\n\n"
        
        for i, (vuln_id, vuln_data) in enumerate(critical_vulns.items(), 1):
            cvss = vuln_data.get('cvss', 'N/A')
            service = vuln_data.get('service', 'Unknown')
            description = vuln_data.get('description', 'No description available')
            remediation = vuln_data.get('remediation', 'Contact security team')
            
            response += f"""🔴 {i}. {vuln_id}
   CVSS Score: {cvss}
   Service: {service}
   Description: {description}
   
   🛠️ IMMEDIATE ACTION REQUIRED:
   {remediation}
   
"""
        
        response += "⚠️ These critical vulnerabilities require immediate attention!"
        return response
    
    def _handle_cve_query(self, query: str) -> str:
        # Extract CVE from query
        cve_match = re.search(r'cve-\d{4}-\d{4,}', query)
        if cve_match:
            cve = cve_match.group().upper()
            if cve in self.knowledge_base['vulnerabilities']:
                vuln = self.knowledge_base['vulnerabilities'][cve]
                return f"""🔍 {cve} DETAILS

🎯 Severity: {vuln.get('severity', 'Unknown')} (CVSS: {vuln.get('cvss', 'N/A')})
🖥️ Affected Service: {vuln.get('service', 'Unknown')}
📝 Description: {vuln.get('description', 'No description available')}

🛠️ REMEDIATION:
{vuln.get('remediation', 'Contact security team for specific remediation steps')}

💡 This vulnerability should be addressed based on its {vuln.get('severity', 'unknown')} severity level."""
            else:
                return f"❌ {cve} was not found in the current analysis results."
        
        # Handle specific vulnerability mentions
        if 'log4j' in query or '44228' in query:
            return self._get_log4j_info()
        elif 'apache' in query and ('path' in query or '41773' in query):
            return self._get_apache_path_traversal_info()
        elif 'ssh' in query and ('enum' in query or '15473' in query):
            return self._get_ssh_enum_info()
        else:
            return "Please specify a CVE number (e.g., CVE-2021-44228) or vulnerability name for detailed information."
    
    def _get_log4j_info(self) -> str:
        return """🚨 LOG4J VULNERABILITY (CVE-2021-44228) - "Log4Shell"

🔴 Severity: CRITICAL (CVSS 9.8)
🎯 Impact: Remote Code Execution
📅 Discovered: December 2021

📝 DESCRIPTION:
Apache Log4j2 versions 2.0-beta9 through 2.15.0 are vulnerable to remote code execution via JNDI lookup. Attackers can execute arbitrary code by sending crafted requests containing malicious JNDI lookup strings.

🛠️ IMMEDIATE REMEDIATION:
1. Update Log4j to version 2.17.0 or later
2. If immediate update not possible:
   - Remove JndiLookup class: zip -q -d log4j-core-*.jar org/apache/logging/log4j/core/lookup/JndiLookup.class
   - Set system property: -Dlog4j2.formatMsgNoLookups=true
3. Implement WAF rules to block JNDI lookup patterns
4. Monitor for exploitation attempts

⚠️ This vulnerability is actively exploited in the wild. Patch immediately!"""
    
    def _get_apache_path_traversal_info(self) -> str:
        return """🟠 APACHE PATH TRAVERSAL (CVE-2021-41773)

🟠 Severity: HIGH (CVSS 7.5)
🎯 Impact: Path Traversal / Information Disclosure
📅 Discovered: October 2021

📝 DESCRIPTION:
Apache HTTP Server 2.4.49 is vulnerable to path traversal attacks. Attackers can access files outside the document root by using encoded path sequences.

🛠️ REMEDIATION:
1. Update Apache HTTP Server to version 2.4.51 or later
2. Review and restrict file system permissions
3. Implement proper input validation
4. Monitor access logs for traversal attempts

💡 While not as critical as Log4j, this should be patched promptly."""
    
    def _get_ssh_enum_info(self) -> str:
        return """🟡 SSH USER ENUMERATION (CVE-2018-15473)

🟡 Severity: MEDIUM (CVSS 5.3)
🎯 Impact: Information Disclosure
📅 Discovered: August 2018

📝 DESCRIPTION:
OpenSSH versions before 7.7 are vulnerable to user enumeration. Attackers can determine valid usernames by observing timing differences in authentication responses.

🛠️ REMEDIATION:
1. Update OpenSSH to version 7.7 or later
2. Implement fail2ban or similar brute-force protection
3. Use key-based authentication instead of passwords
4. Restrict SSH access to specific IP ranges
5. Monitor SSH logs for enumeration attempts

💡 While medium severity, this can aid attackers in reconnaissance."""
    
    def _handle_attack_path_query(self) -> str:
        attack_paths = self.knowledge_base.get('attack_paths', [])
        
        if not attack_paths:
            # Generate generic attack path information
            critical_vulns = [k for k, v in self.knowledge_base['vulnerabilities'].items() 
                            if v.get('severity') == 'Critical']
            
            if critical_vulns:
                response = """🕸️ POTENTIAL ATTACK PATHS

Based on the vulnerabilities found, here are likely attack scenarios:

🔴 HIGH-RISK ATTACK PATH:
1. Initial Access: Exploit web application vulnerabilities (HTTP/HTTPS services)
2. Code Execution: Leverage critical vulnerabilities like Log4j for remote code execution
3. Persistence: Establish foothold on the compromised system
4. Lateral Movement: Use system access to explore internal network
5. Data Exfiltration: Access sensitive data or systems

🛡️ DEFENSIVE MEASURES:
• Patch critical vulnerabilities immediately
• Implement network segmentation
• Deploy endpoint detection and response (EDR)
• Monitor for suspicious network activity
• Implement zero-trust architecture"""
                return response
            else:
                return "✅ No high-risk attack paths identified. The current vulnerability profile shows lower risk for direct exploitation."
        
        # If we have specific attack paths from analysis
        response = f"🕸️ IDENTIFIED ATTACK PATHS ({len(attack_paths)})\n\n"
        for i, path in enumerate(attack_paths[:3], 1):  # Show top 3
            response += f"{i}. {path.get('description', 'Attack path details not available')}\n"
        
        return response
    
    def _handle_remediation_query(self, query: str) -> str:
        # Check if asking about specific vulnerability
        if 'log4j' in query:
            return self._get_log4j_info()
        elif 'apache' in query:
            return self._get_apache_path_traversal_info()
        elif 'ssh' in query:
            return self._get_ssh_enum_info()
        
        # General remediation advice
        recommendations = self.knowledge_base.get('recommendations', [])
        severity_summary = self.knowledge_base['severity_summary']
        
        response = """🛠️ REMEDIATION ROADMAP

📋 IMMEDIATE ACTIONS (0-7 days):"""
        
        if severity_summary.get('Critical', 0) > 0:
            response += f"\n🚨 Patch {severity_summary['Critical']} critical vulnerabilities"
        if severity_summary.get('High', 0) > 0:
            response += f"\n🟠 Address {severity_summary['High']} high-severity issues"
        
        response += "\n\n📋 SHORT-TERM ACTIONS (1-4 weeks):"
        if severity_summary.get('Medium', 0) > 0:
            response += f"\n🟡 Remediate {severity_summary['Medium']} medium-severity vulnerabilities"
        response += "\n🔍 Implement vulnerability scanning automation"
        response += "\n🛡️ Deploy additional security controls"
        
        response += "\n\n📋 LONG-TERM IMPROVEMENTS (1-3 months):"
        response += "\n📊 Establish regular security assessments"
        response += "\n🎓 Conduct security awareness training"
        response += "\n🏗️ Implement security-by-design practices"
        
        if recommendations:
            response += "\n\n💡 SPECIFIC RECOMMENDATIONS:\n"
            for i, rec in enumerate(recommendations, 1):
                response += f"{i}. {rec}\n"
        
        return response
    
    def _handle_priority_query(self) -> str:
        severity_summary = self.knowledge_base['severity_summary']
        
        response = """🎯 VULNERABILITY PRIORITIZATION

📊 PRIORITY MATRIX (Address in this order):

1️⃣ CRITICAL PRIORITY:"""
        
        if severity_summary.get('Critical', 0) > 0:
            response += f"\n   🔴 {severity_summary['Critical']} Critical vulnerabilities"
            response += "\n   ⏰ Timeline: Immediate (0-24 hours)"
            response += "\n   💼 Business Impact: Severe"
        else:
            response += "\n   ✅ No critical vulnerabilities found"
        
        response += "\n\n2️⃣ HIGH PRIORITY:"
        if severity_summary.get('High', 0) > 0:
            response += f"\n   🟠 {severity_summary['High']} High-severity vulnerabilities"
            response += "\n   ⏰ Timeline: 1-7 days"
            response += "\n   💼 Business Impact: Significant"
        else:
            response += "\n   ✅ No high-severity vulnerabilities found"
        
        response += "\n\n3️⃣ MEDIUM PRIORITY:"
        if severity_summary.get('Medium', 0) > 0:
            response += f"\n   🟡 {severity_summary['Medium']} Medium-severity vulnerabilities"
            response += "\n   ⏰ Timeline: 1-4 weeks"
            response += "\n   💼 Business Impact: Moderate"
        else:
            response += "\n   ✅ No medium-severity vulnerabilities found"
        
        response += "\n\n💡 PRIORITIZATION FACTORS:"
        response += "\n• CVSS Score (higher = more urgent)"
        response += "\n• Exploitability (public exploits available)"
        response += "\n• Asset criticality (business importance)"
        response += "\n• Network exposure (internet-facing services)"
        
        return response
    
    def _handle_service_query(self, query: str) -> str:
        service_vulns = {}
        
        # Group vulnerabilities by service
        for vuln_id, vuln_data in self.knowledge_base['vulnerabilities'].items():
            service = vuln_data.get('service', 'unknown')
            if service not in service_vulns:
                service_vulns[service] = []
            service_vulns[service].append((vuln_id, vuln_data))
        
        # Check for specific service in query
        target_service = None
        if 'http' in query or 'web' in query:
            target_service = 'http'
        elif 'ssh' in query:
            target_service = 'ssh'
        elif 'server' in query:
            target_service = None  # Show all services
        
        if target_service and target_service in service_vulns:
            vulns = service_vulns[target_service]
            response = f"🖥️ {target_service.upper()} SERVICE VULNERABILITIES ({len(vulns)})\n\n"
            
            for vuln_id, vuln_data in vulns:
                severity = vuln_data.get('severity', 'Unknown')
                cvss = vuln_data.get('cvss', 'N/A')
                response += f"• {vuln_id} - {severity} (CVSS: {cvss})\n"
            
            return response
        else:
            # Show all services
            response = "🖥️ VULNERABILITIES BY SERVICE\n\n"
            
            for service, vulns in service_vulns.items():
                if vulns:
                    response += f"📡 {service.upper()}: {len(vulns)} vulnerabilities\n"
                    for vuln_id, vuln_data in vulns[:2]:  # Show first 2
                        severity = vuln_data.get('severity', 'Unknown')
                        response += f"   • {vuln_id} ({severity})\n"
                    if len(vulns) > 2:
                        response += f"   • ... and {len(vulns) - 2} more\n"
                    response += "\n"
            
            return response
    
    def _handle_help_query(self) -> str:
        return """🤖 SECURECHAIN VULNERABILITY CHATBOT HELP

I can help you understand and address the security vulnerabilities found in your analysis.

📋 AVAILABLE COMMANDS:

🔍 VULNERABILITY QUERIES:
• "What vulnerabilities were found?"
• "Show me critical vulnerabilities"
• "Tell me about CVE-2021-44228"
• "What are the Log4j issues?"

🛠️ REMEDIATION HELP:
• "How do I fix the critical vulnerabilities?"
• "What should I patch first?"
• "How do I secure the web server?"
• "Give me remediation steps"

🎯 PRIORITIZATION:
• "What should I prioritize?"
• "Which vulnerabilities are most urgent?"
• "How should I order the fixes?"

🕸️ ATTACK ANALYSIS:
• "What attack paths exist?"
• "How could an attacker exploit this?"
• "Show me the attack scenarios"

📊 SUMMARY & OVERVIEW:
• "Give me a summary"
• "How many vulnerabilities were found?"
• "What's the overall risk level?"

💡 EXAMPLES:
• "What are the critical vulnerabilities and how do I fix them?"
• "Tell me about the Log4j vulnerability"
• "Which vulnerabilities should I prioritize?"
• "How could an attacker exploit these issues?"

Just ask me anything about the security analysis results!"""
    
    def _handle_general_query(self, query: str) -> str:
        # Try to provide a helpful response based on keywords
        if 'risk' in query:
            critical_count = self.knowledge_base['severity_summary'].get('Critical', 0)
            high_count = self.knowledge_base['severity_summary'].get('High', 0)
            
            if critical_count > 0:
                risk_level = "HIGH RISK"
            elif high_count > 0:
                risk_level = "MEDIUM-HIGH RISK"
            else:
                risk_level = "LOW-MEDIUM RISK"
            
            return f"""📊 RISK ASSESSMENT

🎯 Overall Risk Level: {risk_level}

This assessment is based on:
• {critical_count} Critical vulnerabilities
• {high_count} High-severity vulnerabilities
• Total of {self.knowledge_base['target_info']['total_vulnerabilities']} vulnerabilities found

Ask me about specific vulnerabilities or remediation steps for more details."""
        
        elif 'secure' in query or 'protect' in query:
            return """🛡️ SECURITY RECOMMENDATIONS

To improve your security posture:

1. 🚨 Address critical vulnerabilities immediately
2. 🔄 Implement regular patch management
3. 🔍 Set up continuous vulnerability scanning
4. 🛡️ Deploy web application firewall (WAF)
5. 📊 Monitor security logs and alerts
6. 🎓 Train staff on security best practices

Ask me about specific vulnerabilities or "How do I fix [vulnerability]?" for detailed guidance."""
        
        else:
            return f"""I'm here to help with your vulnerability analysis results for {self.knowledge_base['target_info']['website']}.

You can ask me about:
• Specific vulnerabilities (e.g., "Tell me about CVE-2021-44228")
• Remediation steps (e.g., "How do I fix the critical issues?")
• Risk prioritization (e.g., "What should I patch first?")
• Attack scenarios (e.g., "How could this be exploited?")

Type 'help' for more detailed command options, or just ask me a question!"""
    
    def start_interactive_session(self):
        """Start interactive chatbot session"""
        print("="*80)
        print("🤖 SECURECHAIN VULNERABILITY ANALYSIS CHATBOT")
        print("="*80)
        print(self._handle_greeting())
        print("\nType 'quit' or 'exit' to end the session.")
        print("="*80)
        
        while True:
            try:
                user_input = input("\n🔍 Your question: ").strip()
                
                if user_input.lower() in ['quit', 'exit', 'bye']:
                    print("\n👋 Thank you for using SecureChain! Stay secure!")
                    break
                
                if not user_input:
                    print("Please enter a question or type 'help' for available commands.")
                    continue
                
                response = self.process_query(user_input)
                print(f"\n🤖 {response}")
                
            except KeyboardInterrupt:
                print("\n\n👋 Session ended. Stay secure!")
                break
            except Exception as e:
                print(f"\n❌ Error processing query: {str(e)}")
                print("Please try rephrasing your question.")

def main():
    """Main function"""
    if len(sys.argv) > 1:
        kb_file = sys.argv[1]
    else:
        # Look for the most recent knowledge base file
        kb_files = list(Path(".").glob("*_chatbot_kb.json"))
        if kb_files:
            kb_file = str(sorted(kb_files)[-1])  # Most recent
            print(f"Using knowledge base: {kb_file}")
        else:
            print("No knowledge base file found. Please run complete_website_analysis.py first.")
            sys.exit(1)
    
    try:
        chatbot = VulnerabilityChatbot(kb_file)
        chatbot.start_interactive_session()
    except Exception as e:
        print(f"Error starting chatbot: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()