diff --git a/.deepsource.toml b/.deepsource.toml index ada5aac..e972fb9 100644 --- a/.deepsource.toml +++ b/.deepsource.toml @@ -1,14 +1,35 @@ version = 1 -test_patterns = ["\"tests/*\", \"test_.py\", \"*/test_.py\""] +test_patterns = ["tests/", "test_*.py", "*_test.py"] -exclude_patterns = ["\"docs/*\", \"examples/\", \"scripts/\", \".md\", \".yml\", \".json\""] +exclude_patterns = [ + "docs/", + "examples/", + "scripts/", + ".git/", + ".venv/", + "*.md", + "*.json", + "*.yml", + "*.yaml", + "static/", + "templates/" +] [[analyzers]] name = "python" - [analyzers.meta] - runtime_version = "3.x.x" +[analyzers.meta] +runtime_version = "3.x" -[[transformers]] -name = "black" \ No newline at end of file +# Disable security-focused web checks +skip_rules = [ + "python.security.audit", + "python.security.injection", + "python.security.unsafe", + "python.security.crypto", + "python.security.network" +] + +[[code_formatters]] +name = "black" diff --git a/ai_controller.py b/ai_controller.py index c31246c..976ebff 100644 --- a/ai_controller.py +++ b/ai_controller.py @@ -18,12 +18,13 @@ from utils.context_builder import ContextBuilder from utils.notifier import Notifier + class AIController: """Controls AI operations and decision making.""" def __init__(self, session_dir: str, config_path: str): """Initialize the AI controller. - + Args: session_dir: Directory to store session data config_path: Path to configuration file @@ -35,7 +36,7 @@ def __init__(self, session_dir: str, config_path: str): self.notifier = Notifier(session_dir, config_path) self.report_generator = ReportGenerator(self.session_dir) self.output_format = "all" # Default output format - + # Initialize session data self.session_data = { "start_time": datetime.now().isoformat(), @@ -44,27 +45,29 @@ def __init__(self, session_dir: str, config_path: str): "findings": [], "execution_history": [], "ai_decisions": [], - "errors": [] + "errors": [], } - async def process_query(self, query: str, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + async def process_query( + self, query: str, context: Optional[Dict[str, Any]] = None + ) -> Dict[str, Any]: """Process an AI query. - + Args: query: The query to process context: Optional context data - + Returns: Dictionary containing the response and metadata """ self.logger.info(f"Processing AI query: {query}") - + try: # Prepare context context_str = "" if context: context_str = json.dumps(context, indent=2) - + # Create prompt prompt = f"""You are a security expert. Please provide guidance on the following query: @@ -81,29 +84,29 @@ async def process_query(self, query: str, context: Optional[Dict[str, Any]] = No 5. Safety considerations Provide a detailed response:""" - + # Get AI response response = await self._get_ai_response(prompt) - + # Log response self.logger.info("AI response received") - + return { "answer": response, "logs": self.logger.get_logs(), - "prompt": prompt + "prompt": prompt, } - + except Exception as e: self.logger.error(f"Error processing query: {str(e)}") raise async def _get_ai_response(self, prompt: str) -> str: """Get response from AI model. - + Args: prompt: The prompt to send to the AI - + Returns: The AI's response """ @@ -111,19 +114,19 @@ async def _get_ai_response(self, prompt: str) -> str: # TODO: Implement actual AI model call # For now, return a placeholder response return "This is a placeholder response. AI integration pending." - + except Exception as e: self.logger.error(f"Error getting AI response: {str(e)}") raise def _generate_report(self) -> Dict[str, str]: """Generate reports for the current session. - + Returns: Dictionary mapping report types to their file paths """ self.logger.info(f"Generating reports in {self.output_format} format...") - + try: # Prepare context data context = { @@ -131,7 +134,7 @@ def _generate_report(self) -> Dict[str, str]: "domain": self.session_data.get("target_domain"), "ip": self.session_data.get("target_ip"), "scan_time": self.session_data.get("start_time"), - "duration": self._calculate_duration() + "duration": self._calculate_duration(), }, "modules": self.session_data["modules"], "tools": self.session_data["tools"], @@ -139,45 +142,50 @@ def _generate_report(self) -> Dict[str, str]: "exploits": self.session_data.get("execution_history", []), "defensive_measures": self.session_data.get("defensive_measures", []), "recommendations": self.session_data.get("recommendations", []), - "errors": self.session_data["errors"] + "errors": self.session_data["errors"], } - + # Generate reports - report_paths = self.report_generator.generate_reports(context, self.output_format) - + report_paths = self.report_generator.generate_reports( + context, self.output_format + ) + self.logger.info("Reports generated successfully") return report_paths - + except Exception as e: self.logger.error(f"Error generating reports: {str(e)}") raise def _calculate_duration(self) -> str: """Calculate session duration. - + Returns: Formatted duration string """ start_time = datetime.fromisoformat(self.session_data["start_time"]) end_time = datetime.now() duration = end_time - start_time - + hours = duration.seconds // 3600 minutes = (duration.seconds % 3600) // 60 seconds = duration.seconds % 60 - + return f"{hours:02d}:{minutes:02d}:{seconds:02d}" def setup_ai(self) -> bool: """Setup AI environment: check Ollama service and model availability.""" try: - if hasattr(self, 'ollama'): + if hasattr(self, "ollama"): client = self.ollama else: from ai_integration import OllamaClient + client = OllamaClient() if not client.is_available(): - self.logger.error("Ollama service not available. Start with: ollama serve") + self.logger.error( + "Ollama service not available. Start with: ollama serve" + ) return False models = client.list_models() if not models: @@ -185,8 +193,10 @@ def setup_ai(self) -> bool: if not client.pull_model(client.main_model): self.logger.error("Failed to pull default model") return False - self.logger.info(f"AI setup complete. Available models: {[m['name'] for m in models]}") + self.logger.info( + f"AI setup complete. Available models: {[m['name'] for m in models]}" + ) return True except Exception as e: self.logger.error(f"Error in setup_ai: {e}") - return False \ No newline at end of file + return False diff --git a/ai_integration.py b/ai_integration.py index ec745dd..320cc38 100755 --- a/ai_integration.py +++ b/ai_integration.py @@ -15,18 +15,15 @@ import asyncio import ctypes + class OllamaClient: def __init__(self, base_url: str = "http://localhost:11434"): self.base_url = base_url self.logger = logging.getLogger(__name__) self.main_model = "deepseek-coder-v2:16b-lite-base-q4_0" # Main model - self.assistant_model = "mistral:7b-instruct-v0.2-q4_0" # Assistant model - self.backup_models = [ - "deepseek-coder:6.7b", - "codellama:7b", - "mistral:7b" - ] - + self.assistant_model = "mistral:7b-instruct-v0.2-q4_0" # Assistant model + self.backup_models = ["deepseek-coder:6.7b", "codellama:7b", "mistral:7b"] + def is_available(self) -> bool: """Check if Ollama service is running""" try: @@ -34,97 +31,103 @@ def is_available(self) -> bool: return response.status_code == 200 except: return False - + def list_models(self) -> List[Dict]: """List available models""" try: response = requests.get(f"{self.base_url}/api/tags") if response.status_code == 200: - return response.json().get('models', []) + return response.json().get("models", []) except Exception as e: self.logger.error(f"Error listing models: {e}") return [] - + def pull_model(self, model: str) -> bool: """Pull a model if not available""" try: self.logger.info(f"Pulling model: {model}") data = {"name": model} - response = requests.post(f"{self.base_url}/api/pull", json=data, stream=True) - + response = requests.post( + f"{self.base_url}/api/pull", json=data, stream=True + ) + for line in response.iter_lines(): if line: try: - status = json.loads(line.decode('utf-8')) - if status.get('status') == 'success': + status = json.loads(line.decode("utf-8")) + if status.get("status") == "success": return True except: continue except Exception as e: self.logger.error(f"Error pulling model {model}: {e}") return False - - def generate(self, prompt: str, model: str = None, system_prompt: str = None) -> Optional[str]: + + def generate( + self, prompt: str, model: str = None, system_prompt: str = None + ) -> Optional[str]: """Generate text using Ollama""" if not model: model = self.get_best_model() - + if not model: self.logger.error("No suitable model available") return None - + try: data = { "model": model, "prompt": prompt, "stream": False, "options": { - "temperature": 0.5, # Lowered for more deterministic planning + "temperature": 0.5, # Lowered for more deterministic planning "top_p": 0.9, "max_tokens": 4096, - "num_ctx": 16384, # Increased context window for large prompts + "num_ctx": 16384, # Increased context window for large prompts "num_thread": 8, - "repeat_penalty": 1.1 - } + "repeat_penalty": 1.1, + }, } - + if system_prompt: data["system"] = system_prompt - - response = requests.post(f"{self.base_url}/api/generate", json=data, timeout=300) # Increased timeout - + + response = requests.post( + f"{self.base_url}/api/generate", json=data, timeout=300 + ) # Increased timeout + if response.status_code == 200: result = response.json() - return result.get('response', '').strip() + return result.get("response", "").strip() else: self.logger.error(f"Ollama API error: {response.status_code}") - + except Exception as e: self.logger.error(f"Error generating with Ollama: {e}") - + return None - + def get_best_model(self) -> Optional[str]: """Get the best available model""" - available_models = [m['name'] for m in self.list_models()] - + available_models = [m["name"] for m in self.list_models()] + # Check main model first (deepseek-coder-v2:16b-lite-base-q4_0) if self.main_model in available_models: return self.main_model - + # Check assistant model next (mistral:7b-instruct-v0.2-q4_0) if self.assistant_model in available_models: return self.assistant_model - + # Check backup models for model in self.backup_models: if model in available_models: return model - + # If no preferred models, return first available if available_models: return available_models[0] - + return None @@ -136,17 +139,18 @@ def get_best_model(self) -> Optional[str]: c_parser.parse_nuclei_output.argtypes = [ctypes.c_char_p] c_parser.parse_nuclei_output.restype = ctypes.c_char_p + class AIAnalyzer: def __init__(self, ollama_client: OllamaClient): self.ollama = ollama_client self.logger = logging.getLogger(__name__) - + def analyze_nmap_output(self, nmap_output: str) -> Dict[str, Any]: """Analyze nmap scan results using AI""" system_prompt = """You are a cybersecurity expert analyzing nmap scan results. Identify potential vulnerabilities, interesting services, and security issues. Provide structured analysis in JSON format with severity levels.""" - + prompt = f""" Analyze this nmap scan output and identify potential security issues: @@ -169,28 +173,28 @@ def analyze_nmap_output(self, nmap_output: str) -> Dict[str, Any]: "next_steps": ["recommended follow-up actions"] }} """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) if response: try: return json.loads(response) except json.JSONDecodeError: # Extract JSON if wrapped in markdown - json_match = re.search(r'```json\n(.*?)\n```', response, re.DOTALL) + json_match = re.search(r"```json\n(.*?)\n```", response, re.DOTALL) if json_match: try: return json.loads(json_match.group(1)) except: pass - + return {"error": "Failed to analyze nmap output", "raw_response": response} - + def generate_exploit_code(self, vulnerability_info: Dict) -> str: """Generate exploit code based on vulnerability information""" system_prompt = """You are a security researcher creating proof-of-concept exploit code. Generate safe, educational exploit code with proper error handling and comments. Include safety warnings and ethical use disclaimers.""" - + prompt = f""" Generate a Python proof-of-concept exploit for this vulnerability: @@ -209,18 +213,18 @@ def generate_exploit_code(self, vulnerability_info: Dict) -> str: Generate complete, working Python code: """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) return response or "# Failed to generate exploit code" - + def analyze_web_response(self, url: str, response_data: Dict) -> Dict[str, Any]: """Analyze web service response for vulnerabilities""" system_prompt = """You are a web application security expert. Analyze HTTP responses for potential vulnerabilities and security issues.""" - - headers = response_data.get('headers', {}) - content = response_data.get('content', '')[:2000] # Limit content length - + + headers = response_data.get("headers", {}) + content = response_data.get("content", "")[:2000] # Limit content length + prompt = f""" Analyze this web service for security issues: @@ -255,26 +259,28 @@ def analyze_web_response(self, url: str, response_data: Dict) -> Dict[str, Any]: "recommendations": ["security recommendations"] }} """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) if response: try: return json.loads(response) except json.JSONDecodeError: - json_match = re.search(r'```json\n(.*?)\n```', response, re.DOTALL) + json_match = re.search(r"```json\n(.*?)\n```", response, re.DOTALL) if json_match: try: return json.loads(json_match.group(1)) except: pass - + return {"error": "Failed to analyze web response", "raw_response": response} - - def fix_broken_tool(self, tool_name: str, error_output: str, source_code: str = None) -> str: + + def fix_broken_tool( + self, tool_name: str, error_output: str, source_code: str = None + ) -> str: """Generate fixes for broken security tools""" system_prompt = """You are a DevOps engineer specializing in fixing broken security tools. Analyze errors and provide working solutions.""" - + prompt = f""" This security tool is broken and needs fixing: @@ -296,15 +302,15 @@ def fix_broken_tool(self, tool_name: str, error_output: str, source_code: str = - Missing environment variables - Network/permission issues """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) return response or "# Failed to generate fix" - + def prioritize_vulnerabilities(self, vulnerabilities: List[Dict]) -> List[Dict]: """Use AI to prioritize vulnerabilities by exploitability and impact""" system_prompt = """You are a penetration tester prioritizing vulnerabilities. Rank vulnerabilities by exploitability and business impact.""" - + prompt = f""" Prioritize these vulnerabilities for testing: @@ -331,27 +337,27 @@ def prioritize_vulnerabilities(self, vulnerabilities: List[Dict]) -> List[Dict]: ] }} """ - + response = self.ollama.generate(prompt, system_prompt=system_prompt) if response: try: result = json.loads(response) - return result.get('prioritized_vulnerabilities', vulnerabilities) + return result.get("prioritized_vulnerabilities", vulnerabilities) except: pass - + return vulnerabilities def analyze_nuclei_output(self, nuclei_json_output: str) -> Dict[str, Any]: """Analyze nuclei output using the high-speed C parser.""" - + # Call the C function - raw_summary = c_parser.parse_nuclei_output(nuclei_json_output.encode('utf-8')) - summary_str = raw_summary.decode('utf-8') - + raw_summary = c_parser.parse_nuclei_output(nuclei_json_output.encode("utf-8")) + summary_str = raw_summary.decode("utf-8") + # The C function must free the memory it allocated, or we'd have a memory leak. # For this example, we assume a simple case. A real implementation needs memory management. - + try: return json.loads(summary_str) except json.JSONDecodeError: @@ -363,6 +369,7 @@ class AIOrchestrator: Manages a multi-step AI reasoning pipeline for complex security tasks. It chains specialized prompts for planning, tool selection, and execution. """ + def __init__(self, prompt_dir: Path): self.prompt_dir = prompt_dir self.ollama = OllamaClient() @@ -377,7 +384,9 @@ def _load_prompts(self): with open(self.prompt_dir / "Devin AI/system.md", "r") as f: prompts["planner"] = f.read() # Manus-style tool selection prompt - with open(self.prompt_dir / "Manus Agent Tools & Prompt/system.md", "r") as f: + with open( + self.prompt_dir / "Manus Agent Tools & Prompt/system.md", "r" + ) as f: prompts["tool_selector"] = f.read() # Cursor-style code/analysis prompt with open(self.prompt_dir / "Cursor Prompts/prompts.md", "r") as f: @@ -392,25 +401,25 @@ def execute_task(self, task_description: str): Executes a full task pipeline: Plan -> Select Tool -> Execute -> Analyze. """ print("--- AI Task Pipeline Initiated ---") - + # 1. Planning Phase (using Devin's prompt) plan = self._planning_phase(task_description) - self.state['plan'] = plan + self.state["plan"] = plan print(f"Phase 1: Plan Created -> {plan}") # 2. Tool Selection Phase (using Manus' prompt) tool_command = self._tool_selection_phase(task_description, plan) - self.state['tool_command'] = tool_command + self.state["tool_command"] = tool_command print(f"Phase 2: Tool Selected -> {tool_command}") # 3. Execution Phase (simulated) execution_result = self._execution_phase(tool_command) - self.state['execution_result'] = execution_result + self.state["execution_result"] = execution_result print(f"Phase 3: Execution Result -> {execution_result[:100]}...") # 4. Analysis Phase (using Cursor's prompt) analysis = self._analysis_phase(execution_result) - self.state['analysis'] = analysis + self.state["analysis"] = analysis print(f"Phase 4: Analysis Complete -> {analysis}") print("--- AI Task Pipeline Complete ---") @@ -418,18 +427,18 @@ def execute_task(self, task_description: str): def _planning_phase(self, task: str) -> str: """Uses the 'planner' prompt to create a high-level strategy.""" - system_prompt = self.prompts['planner'] + system_prompt = self.prompts["planner"] user_prompt = f"Create a step-by-step plan for the following task: {task}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response def _tool_selection_phase(self, task: str, plan: str) -> str: """Uses the 'tool_selector' prompt to choose the right command.""" - system_prompt = self.prompts['tool_selector'] + system_prompt = self.prompts["tool_selector"] user_prompt = f"Given the task '{task}' and the plan '{plan}', what is the exact shell command to execute next? Only output the command." response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response - + def _execution_phase(self, command: str) -> str: """Simulates running the command and returns mock output.""" print(f"Simulating execution of: `{command}`") @@ -439,7 +448,7 @@ def _execution_phase(self, command: str) -> str: def _analysis_phase(self, result: str) -> str: """Uses the 'analyst' prompt to interpret the results.""" - system_prompt = self.prompts['analyst'] + system_prompt = self.prompts["analyst"] user_prompt = f"Analyze the following tool output and provide a summary of key findings and recommendations:\n\n{result}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response @@ -448,43 +457,51 @@ def _analysis_phase(self, result: str) -> str: # CLI interface for AI module if __name__ == "__main__": import argparse - + parser = argparse.ArgumentParser(description="VulnForge AI Module") - parser.add_argument("--test-connection", action="store_true", help="Test Ollama connection") + parser.add_argument( + "--test-connection", action="store_true", help="Test Ollama connection" + ) parser.add_argument("--pull-model", help="Pull a specific model") - parser.add_argument("--list-models", action="store_true", help="List available models") + parser.add_argument( + "--list-models", action="store_true", help="List available models" + ) parser.add_argument("--analyze-nmap", help="Analyze nmap output file") parser.add_argument( - "--ai-pipeline", action="store_true", help="Enable the advanced multi-prompt AI pipeline." + "--ai-pipeline", + action="store_true", + help="Enable the advanced multi-prompt AI pipeline.", ) parser.add_argument( - "--prompt-dir", help="Directory for the AI pipeline prompts.", default="AI_Propmt/system-prompts-and-models-of-ai-tools" + "--prompt-dir", + help="Directory for the AI pipeline prompts.", + default="AI_Propmt/system-prompts-and-models-of-ai-tools", ) - + args = parser.parse_args() - + orchestrator = AIOrchestrator(Path(args.prompt_dir)) - + if args.test_connection: if orchestrator.ollama.is_available(): print("✓ AI system ready") else: print("✗ AI system not available") - + elif args.pull_model: if orchestrator.ollama.pull_model(args.pull_model): print(f"✓ Model {args.pull_model} pulled successfully") else: print(f"✗ Failed to pull model {args.pull_model}") - + elif args.list_models: models = orchestrator.ollama.list_models() print("Available models:") for model in models: print(f" - {model['name']}") - + elif args.analyze_nmap: - with open(args.analyze_nmap, 'r') as f: + with open(args.analyze_nmap, "r") as f: content = f.read() result = orchestrator.analyzer.analyze_nmap_output(content) print(json.dumps(result, indent=2)) @@ -492,11 +509,13 @@ def _analysis_phase(self, result: str) -> str: # Handle AI Pipeline Mode if args.ai_pipeline: if not args.target: - print("Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'") - + print( + "Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'" + ) + prompt_path = Path(args.prompt_dir) if not prompt_path.exists(): print(f"Error: Prompt directory not found at '{prompt_path}'") - + orchestrator = AIOrchestrator(prompt_path) - orchestrator.execute_task(f"Perform a security scan on {args.target}") \ No newline at end of file + orchestrator.execute_task(f"Perform a security scan on {args.target}") diff --git a/ai_orchestrator.py b/ai_orchestrator.py index 5844491..3d98128 100644 --- a/ai_orchestrator.py +++ b/ai_orchestrator.py @@ -3,11 +3,13 @@ from ai_integration import OllamaClient import subprocess + class AIOrchestrator: """ Manages a multi-step AI reasoning pipeline for complex security tasks. It chains specialized prompts for planning, tool selection, and execution. """ + def __init__(self, prompt_dir: Path): self.prompt_dir = prompt_dir self.ollama = OllamaClient() @@ -22,7 +24,9 @@ def _load_prompts(self): with open(self.prompt_dir / "Devin AI/system.md", "r") as f: prompts["planner"] = f.read() # Manus-style tool selection prompt - with open(self.prompt_dir / "Manus Agent Tools & Prompt/system.md", "r") as f: + with open( + self.prompt_dir / "Manus Agent Tools & Prompt/system.md", "r" + ) as f: prompts["tool_selector"] = f.read() # Cursor-style code/analysis prompt with open(self.prompt_dir / "Cursor Prompts/prompts.md", "r") as f: @@ -37,25 +41,25 @@ def execute_task(self, task_description: str): Executes a full task pipeline: Plan -> Select Tool -> Execute -> Analyze. """ print("--- AI Task Pipeline Initiated ---") - + # 1. Planning Phase (using Devin's prompt) plan = self._planning_phase(task_description) - self.state['plan'] = plan + self.state["plan"] = plan print(f"Phase 1: Plan Created -> {plan}") # 2. Tool Selection Phase (using Manus' prompt) tool_command = self._tool_selection_phase(task_description, plan) - self.state['tool_command'] = tool_command + self.state["tool_command"] = tool_command print(f"Phase 2: Tool Selected -> {tool_command}") # 3. Execution Phase (simulated) execution_result = self._execution_phase(tool_command) - self.state['execution_result'] = execution_result + self.state["execution_result"] = execution_result print(f"Phase 3: Execution Result -> {execution_result[:100]}...") # 4. Analysis Phase (using Cursor's prompt) analysis = self._analysis_phase(execution_result) - self.state['analysis'] = analysis + self.state["analysis"] = analysis print(f"Phase 4: Analysis Complete -> {analysis}") print("--- AI Task Pipeline Complete ---") @@ -63,23 +67,25 @@ def execute_task(self, task_description: str): def _planning_phase(self, task: str) -> str: """Uses the 'planner' prompt to create a high-level strategy.""" - system_prompt = self.prompts['planner'] + system_prompt = self.prompts["planner"] user_prompt = f"Create a step-by-step plan for the following task: {task}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response def _tool_selection_phase(self, task: str, plan: str) -> str: """Uses the 'tool_selector' prompt to choose the right command.""" - system_prompt = self.prompts['tool_selector'] + system_prompt = self.prompts["tool_selector"] user_prompt = f"Given the task '{task}' and the plan '{plan}', what is the exact shell command to execute next? Only output the command." response = self.ollama.generate(user_prompt, system_prompt=system_prompt) return response - + def _execution_phase(self, command: str) -> str: """Executes the shell command and returns real output.""" print(f"Executing: `{command}`") try: - result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=600) + result = subprocess.run( + command, shell=True, capture_output=True, text=True, timeout=600 + ) if result.returncode == 0: return result.stdout else: @@ -89,7 +95,7 @@ def _execution_phase(self, command: str) -> str: def _analysis_phase(self, result: str) -> str: """Uses the 'analyst' prompt to interpret the results.""" - system_prompt = self.prompts['analyst'] + system_prompt = self.prompts["analyst"] user_prompt = f"Analyze the following tool output and provide a summary of key findings and recommendations:\n\n{result}" response = self.ollama.generate(user_prompt, system_prompt=system_prompt) - return response \ No newline at end of file + return response diff --git a/ai_wrapper/__init__.py b/ai_wrapper/__init__.py index f8e7ba5..4e94e50 100755 --- a/ai_wrapper/__init__.py +++ b/ai_wrapper/__init__.py @@ -10,4 +10,4 @@ from .ollama_wrapper import OllamaWrapper -__all__ = ['OllamaWrapper'] \ No newline at end of file +__all__ = ["OllamaWrapper"] diff --git a/ai_wrapper/llm_engine.py b/ai_wrapper/llm_engine.py index e882d85..c2e31d3 100755 --- a/ai_wrapper/llm_engine.py +++ b/ai_wrapper/llm_engine.py @@ -18,6 +18,7 @@ import requests from functools import lru_cache + class LLMEngine: def __init__(self, config_path: Optional[Path] = None): self.logger = logging.getLogger(__name__) @@ -26,25 +27,21 @@ def __init__(self, config_path: Optional[Path] = None): self.models = self._initialize_models() self.current_model = self.models[0] # Start with preferred model self.response_cache = {} - + def _load_config(self, config_path: Optional[Path] = None) -> Dict: """Load LLM configuration""" if not config_path: config_path = Path.home() / ".vulnforge" / "configs" / "llm_config.json" - + default_config = { "preferred_model": "deepseek-coder-v2:16b-lite-base-q5_K_S", - "fallback_models": [ - "deepseek-coder:6.7b", - "codellama:7b", - "mistral:7b" - ], + "fallback_models": ["deepseek-coder:6.7b", "codellama:7b", "mistral:7b"], "cache_size": 100, "timeout": 180, "max_retries": 3, - "retry_delay": 2 + "retry_delay": 2, } - + try: if config_path.exists(): with open(config_path) as f: @@ -53,67 +50,74 @@ def _load_config(self, config_path: Optional[Path] = None) -> Dict: except Exception as e: self.logger.error(f"Error loading LLM config: {e}") return default_config - + def _initialize_models(self) -> List[str]: """Initialize available models""" available_models = [] - + # Check preferred model first if self._is_model_available(self.config["preferred_model"]): available_models.append(self.config["preferred_model"]) - + # Check fallback models for model in self.config["fallback_models"]: if self._is_model_available(model): available_models.append(model) - + if not available_models: self.logger.error("No models available!") - + return available_models - + def _is_model_available(self, model: str) -> bool: """Check if a model is available""" try: response = requests.get(f"{self.base_url}/api/tags", timeout=5) if response.status_code == 200: - models = response.json().get('models', []) - return any(m['name'] == model for m in models) + models = response.json().get("models", []) + return any(m["name"] == model for m in models) except: pass return False - + def _pull_model(self, model: str) -> bool: """Pull a model if not available""" try: self.logger.info(f"Pulling model: {model}") data = {"name": model} - response = requests.post(f"{self.base_url}/api/pull", json=data, stream=True) - + response = requests.post( + f"{self.base_url}/api/pull", json=data, stream=True + ) + for line in response.iter_lines(): if line: try: - status = json.loads(line.decode('utf-8')) - if status.get('status') == 'success': + status = json.loads(line.decode("utf-8")) + if status.get("status") == "success": return True except: continue except Exception as e: self.logger.error(f"Error pulling model {model}: {e}") return False - + @lru_cache(maxsize=100) - def query(self, prompt: str, system_prompt: Optional[str] = None, - model: Optional[str] = None, use_cache: bool = True) -> Optional[str]: + def query( + self, + prompt: str, + system_prompt: Optional[str] = None, + model: Optional[str] = None, + use_cache: bool = True, + ) -> Optional[str]: """Query the LLM with fallback support""" if not model: model = self.current_model - + # Check cache if enabled cache_key = f"{model}:{prompt}:{system_prompt}" if use_cache and cache_key in self.response_cache: return self.response_cache[cache_key] - + for attempt in range(self.config["max_retries"]): try: data = { @@ -126,28 +130,28 @@ def query(self, prompt: str, system_prompt: Optional[str] = None, "max_tokens": 4096, "num_ctx": 8192, "num_thread": 8, - "repeat_penalty": 1.1 - } + "repeat_penalty": 1.1, + }, } - + if system_prompt: data["system"] = system_prompt - + response = requests.post( - f"{self.base_url}/api/generate", - json=data, - timeout=self.config["timeout"] + f"{self.base_url}/api/generate", + json=data, + timeout=self.config["timeout"], ) - + if response.status_code == 200: - result = response.json().get('response', '').strip() + result = response.json().get("response", "").strip() if use_cache: self.response_cache[cache_key] = result return result - + except Exception as e: self.logger.error(f"Error querying model {model}: {e}") - + # Try next model if available if model in self.models: current_index = self.models.index(model) @@ -156,23 +160,23 @@ def query(self, prompt: str, system_prompt: Optional[str] = None, self.logger.info(f"Switching to fallback model: {model}") else: break - + time.sleep(self.config["retry_delay"]) - + return None - + def clear_cache(self): """Clear the response cache""" self.response_cache.clear() self.query.cache_clear() - + def get_available_models(self) -> List[str]: """Get list of available models""" return self.models.copy() - + def set_preferred_model(self, model: str) -> bool: """Set preferred model if available""" if self._is_model_available(model): self.current_model = model return True - return False \ No newline at end of file + return False diff --git a/examples/notifications.py b/examples/notifications.py index e68924c..471d54d 100755 --- a/examples/notifications.py +++ b/examples/notifications.py @@ -10,26 +10,23 @@ # Configure logging logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) + async def main(): # Initialize config manager config = ConfigManager() - + # Initialize notifier notifier = Notifier(config) await notifier.start() - + try: # Example 1: Basic notification - await notifier.notify( - "Scan started for example.com", - "info" - ) - + await notifier.notify("Scan started for example.com", "info") + # Example 2: Vulnerability found await notifier.notify( "SQL Injection vulnerability detected", @@ -38,10 +35,10 @@ async def main(): "vulnerability": "SQL Injection", "affected_url": "https://example.com/login", "payload": "' OR '1'='1", - "confidence": "high" - } + "confidence": "high", + }, ) - + # Example 3: Critical finding await notifier.notify( "Remote Code Execution vulnerability found!", @@ -50,11 +47,11 @@ async def main(): "vulnerability": "RCE", "affected_component": "File Upload Handler", "cve": "CVE-2023-1234", - "exploit_available": True + "exploit_available": True, }, - channels=["email", "discord"] # Send to specific channels + channels=["email", "discord"], # Send to specific channels ) - + # Example 4: Scan completion await notifier.notify( "Scan completed successfully", @@ -62,22 +59,18 @@ async def main(): data={ "target": "example.com", "duration": "2h 15m", - "findings": { - "critical": 1, - "high": 3, - "medium": 5, - "low": 8 - } - } + "findings": {"critical": 1, "high": 3, "medium": 5, "low": 8}, + }, ) - + # Wait for notifications to be processed await asyncio.sleep(1) - + except Exception as e: logger.error(f"Error in notification example: {e}") finally: await notifier.stop() + if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(main()) diff --git a/modules/ai/ai_assistant.py b/modules/ai/ai_assistant.py index dd0c478..4082646 100644 --- a/modules/ai/ai_assistant.py +++ b/modules/ai/ai_assistant.py @@ -1,5 +1,6 @@ import os + class AIAssistant: def ask_ai(self, query, context=None, debug=False): try: @@ -27,4 +28,4 @@ def _log_ai_interaction(self, query, response, debug): log_path = os.path.expanduser("~/.vulnforge/sessions/logs/ai_controller.log") os.makedirs(os.path.dirname(log_path), exist_ok=True) with open(log_path, "a") as f: - f.write(f"[QUERY] {query}\n[RESPONSE] {response}\n[DEBUG] {debug}\n\n") \ No newline at end of file + f.write(f"[QUERY] {query}\n[RESPONSE] {response}\n[DEBUG] {debug}\n\n") diff --git a/modules/ai/ai_controller.py b/modules/ai/ai_controller.py index 9464e4f..8bcbdc4 100644 --- a/modules/ai/ai_controller.py +++ b/modules/ai/ai_controller.py @@ -1,5 +1,6 @@ import os + class AIController: def run_ai_task(self, task, context=None, debug=False): try: @@ -16,4 +17,4 @@ def _log_ai_task(self, task, result, debug): log_path = os.path.expanduser("~/.vulnforge/sessions/logs/ai_controller.log") os.makedirs(os.path.dirname(log_path), exist_ok=True) with open(log_path, "a") as f: - f.write(f"[TASK] {task}\n[RESULT] {result}\n[DEBUG] {debug}\n\n") \ No newline at end of file + f.write(f"[TASK] {task}\n[RESULT] {result}\n[DEBUG] {debug}\n\n") diff --git a/modules/cve_collector/__init__.py b/modules/cve_collector/__init__.py index faa95ed..18a99b2 100755 --- a/modules/cve_collector/__init__.py +++ b/modules/cve_collector/__init__.py @@ -4,4 +4,4 @@ from .cve_collector import CVECollector -__all__ = ['CVECollector'] \ No newline at end of file +__all__ = ["CVECollector"] diff --git a/modules/cve_collector/cve_collector.py b/modules/cve_collector/cve_collector.py index 64b804d..80c4334 100755 --- a/modules/cve_collector/cve_collector.py +++ b/modules/cve_collector/cve_collector.py @@ -11,7 +11,13 @@ from typing import Dict, List, Optional, Any from datetime import datetime, timedelta from rich.console import Console -from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn +from rich.progress import ( + Progress, + SpinnerColumn, + TextColumn, + BarColumn, + TaskProgressColumn, +) import hashlib import time import re @@ -21,6 +27,7 @@ from utils.cli_utils import create_progress, print_results_table from ai_wrapper.ollama_wrapper import OllamaWrapper + class CVECollector: def __init__(self, base_dir: Path, ai_wrapper: Optional[OllamaWrapper] = None): self.base_dir = base_dir @@ -31,38 +38,38 @@ def __init__(self, base_dir: Path, ai_wrapper: Optional[OllamaWrapper] = None): self.cache_dir = base_dir / "cache" / "cve_feeds" self.data_dir.mkdir(parents=True, exist_ok=True) self.cache_dir.mkdir(parents=True, exist_ok=True) - + # API rate limiting self.nvd_rate_limit = 5 # requests per second self.github_rate_limit = 30 # requests per hour self.last_nvd_request = 0 self.last_github_request = 0 self.github_requests_remaining = 30 - + # Load API keys if available self.api_keys = self._load_api_keys() - + def _load_api_keys(self) -> Dict[str, str]: """Load API keys from config""" try: config_path = self.base_dir / "config" / "api_keys.json" if config_path.exists(): - with open(config_path, 'r') as f: + with open(config_path, "r") as f: return json.load(f) except Exception as e: self.logger.warning(f"Failed to load API keys: {e}") return {} - + async def _wait_for_rate_limit(self, api_type: str): """Handle API rate limiting""" current_time = time.time() - + if api_type == "nvd": time_since_last = current_time - self.last_nvd_request if time_since_last < (1 / self.nvd_rate_limit): await asyncio.sleep((1 / self.nvd_rate_limit) - time_since_last) self.last_nvd_request = time.time() - + elif api_type == "github": if self.github_requests_remaining <= 0: # Wait until rate limit resets (1 hour) @@ -70,55 +77,56 @@ async def _wait_for_rate_limit(self, api_type: str): self.github_requests_remaining = 30 self.github_requests_remaining -= 1 self.last_github_request = current_time - + def _get_cache_path(self, source: str, identifier: str) -> Path: """Get cache file path for a request""" cache_key = hashlib.md5(f"{source}:{identifier}".encode()).hexdigest() return self.cache_dir / f"{cache_key}.json" - - async def _get_cached_data(self, cache_path: Path, max_age: int = 3600) -> Optional[Dict]: + + async def _get_cached_data( + self, cache_path: Path, max_age: int = 3600 + ) -> Optional[Dict]: """Get cached data if it exists and is not too old""" try: if cache_path.exists(): stat = cache_path.stat() if time.time() - stat.st_mtime < max_age: - async with aiofiles.open(cache_path, 'r') as f: + async with aiofiles.open(cache_path, "r") as f: return json.loads(await f.read()) except Exception as e: self.logger.warning(f"Error reading cache: {e}") return None - + async def _save_to_cache(self, cache_path: Path, data: Dict): """Save data to cache""" try: - async with aiofiles.open(cache_path, 'w') as f: + async with aiofiles.open(cache_path, "w") as f: await f.write(json.dumps(data)) except Exception as e: self.logger.warning(f"Error saving to cache: {e}") - + async def fetch_nvd_feed(self, start_date: Optional[str] = None) -> List[Dict]: """Fetch CVE data from NVD feed""" if not start_date: start_date = (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d") - + self.console.print("[bold blue]Fetching NVD CVE feed...[/bold blue]") - + # Check cache first cache_path = self._get_cache_path("nvd", start_date) - cached_data = await self._get_cached_data(cache_path, max_age=3600) # 1 hour cache + cached_data = await self._get_cached_data( + cache_path, max_age=3600 + ) # 1 hour cache if cached_data: return cached_data - + url = f"https://services.nvd.nist.gov/rest/json/cves/2.0" - params = { - "pubStartDate": f"{start_date}T00:00:00.000", - "resultsPerPage": 2000 - } - + params = {"pubStartDate": f"{start_date}T00:00:00.000", "resultsPerPage": 2000} + # Add API key if available if "nvd" in self.api_keys: params["apiKey"] = self.api_keys["nvd"] - + async with aiohttp.ClientSession() as session: try: await self._wait_for_rate_limit("nvd") @@ -138,47 +146,63 @@ async def fetch_nvd_feed(self, start_date: Optional[str] = None) -> List[Dict]: except Exception as e: self.logger.error(f"Error fetching NVD feed: {e}") return [] - + async def fetch_exploit_db(self) -> List[Dict]: """Fetch exploit data from Exploit-DB""" self.console.print("[bold blue]Fetching Exploit-DB data...[/bold blue]") - + # Check cache first cache_path = self._get_cache_path("exploitdb", "latest") - cached_data = await self._get_cached_data(cache_path, max_age=86400) # 24 hour cache + cached_data = await self._get_cached_data( + cache_path, max_age=86400 + ) # 24 hour cache if cached_data: return cached_data - + url = "https://raw.githubusercontent.com/offensive-security/exploitdb/master/files_exploits.csv" - + async with aiohttp.ClientSession() as session: try: async with session.get(url) as response: if response.status == 200: content = await response.text() exploits = [] - + # Parse CSV content for line in content.splitlines()[1:]: # Skip header try: - parts = line.split(',') + parts = line.split(",") if len(parts) >= 4: # Extract CVE IDs from description - cve_ids = re.findall(r'CVE-\d{4}-\d+', parts[2]) - - exploits.append({ - "id": parts[0], - "file": parts[1], - "description": parts[2], - "date": parts[3], - "author": parts[4] if len(parts) > 4 else "Unknown", - "platform": parts[5] if len(parts) > 5 else "Unknown", - "type": parts[6] if len(parts) > 6 else "Unknown", - "cve_ids": cve_ids - }) + cve_ids = re.findall(r"CVE-\d{4}-\d+", parts[2]) + + exploits.append( + { + "id": parts[0], + "file": parts[1], + "description": parts[2], + "date": parts[3], + "author": ( + parts[4] + if len(parts) > 4 + else "Unknown" + ), + "platform": ( + parts[5] + if len(parts) > 5 + else "Unknown" + ), + "type": ( + parts[6] + if len(parts) > 6 + else "Unknown" + ), + "cve_ids": cve_ids, + } + ) except: continue - + await self._save_to_cache(cache_path, exploits) return exploits else: @@ -187,62 +211,82 @@ async def fetch_exploit_db(self) -> List[Dict]: except Exception as e: self.logger.error(f"Error fetching Exploit-DB: {e}") return [] - + async def fetch_github_pocs(self, cve_id: str) -> List[Dict]: """Search for PoCs on GitHub""" - self.console.print(f"[bold blue]Searching GitHub for {cve_id} PoCs...[/bold blue]") - + self.console.print( + f"[bold blue]Searching GitHub for {cve_id} PoCs...[/bold blue]" + ) + # Check cache first cache_path = self._get_cache_path("github", cve_id) - cached_data = await self._get_cached_data(cache_path, max_age=86400) # 24 hour cache + cached_data = await self._get_cached_data( + cache_path, max_age=86400 + ) # 24 hour cache if cached_data: return cached_data - + # GitHub API search query query = f"{cve_id} poc OR proof of concept OR exploit" url = "https://api.github.com/search/code" headers = { "Accept": "application/vnd.github.v3+json", - "User-Agent": "VulnForge/1.0" + "User-Agent": "VulnForge/1.0", } - + # Add API key if available if "github" in self.api_keys: headers["Authorization"] = f"token {self.api_keys['github']}" - + async with aiohttp.ClientSession() as session: try: await self._wait_for_rate_limit("github") - async with session.get(url, params={"q": query}, headers=headers) as response: + async with session.get( + url, params={"q": query}, headers=headers + ) as response: if response.status == 200: data = await response.json() results = data.get("items", []) - + # Process results processed_results = [] for result in results: try: # Get file content content_url = result["url"] - async with session.get(content_url, headers=headers) as content_response: + async with session.get( + content_url, headers=headers + ) as content_response: if content_response.status == 200: content_data = await content_response.json() content = content_data.get("content", "") - + # Extract relevant information - processed_results.append({ - "repository": result["repository"]["full_name"], - "file_path": result["path"], - "url": result["html_url"], - "content": content, - "language": result.get("language", "Unknown"), - "stars": result["repository"].get("stargazers_count", 0), - "forks": result["repository"].get("forks_count", 0) - }) + processed_results.append( + { + "repository": result["repository"][ + "full_name" + ], + "file_path": result["path"], + "url": result["html_url"], + "content": content, + "language": result.get( + "language", "Unknown" + ), + "stars": result["repository"].get( + "stargazers_count", 0 + ), + "forks": result["repository"].get( + "forks_count", 0 + ), + } + ) except Exception as e: - self.logger.warning(f"Error processing GitHub result: {e}") + self.logger.warning( + f"Error processing GitHub result: {e}" + ) continue - + await self._save_to_cache(cache_path, processed_results) return processed_results elif response.status == 403: # Rate limit exceeded @@ -255,12 +299,12 @@ async def fetch_github_pocs(self, cve_id: str) -> List[Dict]: except Exception as e: self.logger.error(f"Error searching GitHub: {e}") return [] - + async def analyze_cve(self, cve_data: Dict) -> Dict: """Analyze CVE data using AI""" if not self.ai_wrapper: return cve_data - + prompt = f""" Analyze this CVE and provide: 1. Severity assessment (critical, high, medium, low) @@ -304,7 +348,7 @@ async def analyze_cve(self, cve_data: Dict) -> Dict: }} }} """ - + try: analysis = await self.ai_wrapper.generate(prompt) try: @@ -316,30 +360,30 @@ async def analyze_cve(self, cve_data: Dict) -> Dict: cve_data["ai_analysis"] = {"raw_analysis": analysis} except Exception as e: self.logger.error(f"Error analyzing CVE: {e}") - + return cve_data - + async def collect_cves(self, target_info: Dict) -> Dict[str, Any]: """Collect and analyze CVEs for target""" self.console.print("[bold blue]Starting CVE collection...[/bold blue]") - + with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), - console=self.console + console=self.console, ) as progress: # Fetch NVD feed task = progress.add_task("Fetching NVD feed...", total=None) cves = await self.fetch_nvd_feed() progress.update(task, completed=True) - + # Fetch Exploit-DB data task = progress.add_task("Fetching Exploit-DB data...", total=None) exploits = await self.fetch_exploit_db() progress.update(task, completed=True) - + # Match CVEs to target task = progress.add_task("Matching CVEs to target...", total=None) matched_cves = [] @@ -347,7 +391,7 @@ async def collect_cves(self, target_info: Dict) -> Dict[str, Any]: if self._matches_target(cve, target_info): matched_cves.append(cve) progress.update(task, completed=True) - + # Fetch GitHub PoCs task = progress.add_task("Fetching GitHub PoCs...", total=None) for cve in matched_cves: @@ -356,26 +400,26 @@ async def collect_cves(self, target_info: Dict) -> Dict[str, Any]: pocs = await self.fetch_github_pocs(cve_id) cve["github_pocs"] = pocs progress.update(task, completed=True) - + # Analyze CVEs task = progress.add_task("Analyzing CVEs...", total=None) for cve in matched_cves: cve = await self.analyze_cve(cve) progress.update(task, completed=True) - + # Save results task = progress.add_task("Saving results...", total=None) results = { "target_info": target_info, "cves": matched_cves, "exploits": exploits, - "generated_at": datetime.now().isoformat() + "generated_at": datetime.now().isoformat(), } await self._save_results(results) progress.update(task, completed=True) - + return results - + def _matches_target(self, cve: Dict, target_info: Dict) -> bool: """Check if CVE matches target""" # Get CPE strings @@ -383,41 +427,51 @@ def _matches_target(self, cve: Dict, target_info: Dict) -> bool: for node in cve.get("cve", {}).get("configurations", []): for cpe_match in node.get("cpeMatch", []): cpe_list.append(cpe_match.get("criteria")) - + # Check each CPE for cpe in cpe_list: if self._check_cpe_match(cpe, target_info): return True - + # Check description for target keywords - description = cve.get("cve", {}).get("descriptions", [{}])[0].get("value", "").lower() + description = ( + cve.get("cve", {}).get("descriptions", [{}])[0].get("value", "").lower() + ) target_keywords = [ target_info.get("name", "").lower(), target_info.get("vendor", "").lower(), target_info.get("product", "").lower(), - target_info.get("version", "").lower() + target_info.get("version", "").lower(), ] - + return any(keyword in description for keyword in target_keywords if keyword) - + def _check_cpe_match(self, cpe: str, target_info: Dict) -> bool: """Check if CPE string matches target""" if not cpe: return False - + # Parse CPE parts = cpe.split(":") if len(parts) < 5: return False - + # Check vendor - if target_info.get("vendor") and parts[3] != "*" and parts[3].lower() != target_info["vendor"].lower(): + if ( + target_info.get("vendor") + and parts[3] != "*" + and parts[3].lower() != target_info["vendor"].lower() + ): return False - + # Check product - if target_info.get("product") and parts[4] != "*" and parts[4].lower() != target_info["product"].lower(): + if ( + target_info.get("product") + and parts[4] != "*" + and parts[4].lower() != target_info["product"].lower() + ): return False - + # Check version if target_info.get("version") and parts[5] != "*": # Handle version ranges @@ -427,65 +481,69 @@ def _check_cpe_match(self, cpe: str, target_info: Dict) -> bool: return False elif parts[5] != target_info["version"]: return False - + return True - + async def _save_results(self, results: Dict): """Save results to file""" # Create results directory timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") results_dir = self.data_dir / timestamp results_dir.mkdir(exist_ok=True) - + # Save JSON json_path = results_dir / "results.json" - async with aiofiles.open(json_path, 'w') as f: + async with aiofiles.open(json_path, "w") as f: await f.write(json.dumps(results, indent=2)) - + # Save Markdown report md_path = results_dir / "report.md" - async with aiofiles.open(md_path, 'w') as f: + async with aiofiles.open(md_path, "w") as f: await f.write(self._generate_markdown_report(results)) - + def _generate_markdown_report(self, results: Dict) -> str: """Generate Markdown report""" report = [] - + # Add header report.append("# CVE Analysis Report") report.append(f"Generated: {results['generated_at']}") report.append("") - + # Add target info report.append("## Target Information") for key, value in results["target_info"].items(): report.append(f"- **{key}**: {value}") report.append("") - + # Add findings summary report.append("## Findings Summary") severity_counts = {"critical": 0, "high": 0, "medium": 0, "low": 0, "info": 0} for cve in results["cves"]: severity = cve.get("ai_analysis", {}).get("severity", "info").lower() severity_counts[severity] += 1 - + report.append("### Severity Distribution") for severity, count in severity_counts.items(): report.append(f"- **{severity.title()}**: {count}") report.append("") - + # Add detailed findings report.append("## Detailed Findings") for cve in results["cves"]: cve_id = cve.get("cve", {}).get("id", "Unknown") severity = cve.get("ai_analysis", {}).get("severity", "info").lower() - + report.append(f"### {cve_id} ({severity.title()})") - + # Add description - description = cve.get("cve", {}).get("descriptions", [{}])[0].get("value", "No description available") + description = ( + cve.get("cve", {}) + .get("descriptions", [{}])[0] + .get("value", "No description available") + ) report.append(f"**Description**: {description}") - + # Add analysis analysis = cve.get("ai_analysis", {}) if analysis: @@ -501,14 +559,16 @@ def _generate_markdown_report(self, results: Dict) -> str: report.append(f" - {k}: {v}") else: report.append(f"- **{key}**: {value}") - + # Add PoCs pocs = cve.get("github_pocs", []) if pocs: report.append("\n**Proof of Concepts**:") for poc in pocs: - report.append(f"- [{poc['repository']}/{poc['file_path']}]({poc['url']})") - + report.append( + f"- [{poc['repository']}/{poc['file_path']}]({poc['url']})" + ) + report.append("") - - return "\n".join(report) \ No newline at end of file + + return "\n".join(report) diff --git a/modules/exploit_generator/__init__.py b/modules/exploit_generator/__init__.py index a3b60c7..718242b 100755 --- a/modules/exploit_generator/__init__.py +++ b/modules/exploit_generator/__init__.py @@ -4,4 +4,4 @@ from .exploit_generator import ExploitGenerator -__all__ = ['ExploitGenerator'] \ No newline at end of file +__all__ = ["ExploitGenerator"] diff --git a/modules/exploit_generator/exploit_generator.py b/modules/exploit_generator/exploit_generator.py index 6e84dcd..95db358 100755 --- a/modules/exploit_generator/exploit_generator.py +++ b/modules/exploit_generator/exploit_generator.py @@ -14,6 +14,7 @@ from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn + class ExploitGenerator: def __init__(self, base_dir: Path, llm_engine: Any): self.base_dir = base_dir @@ -24,33 +25,33 @@ def __init__(self, base_dir: Path, llm_engine: Any): self.template_dir = self.base_dir / "templates" / "exploits" self.exploit_dir.mkdir(parents=True, exist_ok=True) self.template_dir.mkdir(parents=True, exist_ok=True) - + # Load exploit templates self.templates = self._load_templates() - + def _load_templates(self) -> Dict[str, str]: """Load exploit templates""" templates = {} try: for template_file in self.template_dir.glob("*.py"): - with open(template_file, 'r') as f: + with open(template_file, "r") as f: templates[template_file.stem] = f.read() except Exception as e: self.logger.error(f"Error loading templates: {e}") return templates - + def generate_exploit(self, cve_data: Dict, recon_data: Dict) -> Dict[str, Any]: """Generate exploit based on CVE and recon data""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Prepare prompt with context task = progress.add_task("Preparing exploit generation...", total=None) prompt = self._build_exploit_prompt(cve_data, recon_data) progress.update(task, completed=True) - + # Get AI response task = progress.add_task("Generating exploit code...", total=None) try: @@ -67,41 +68,43 @@ def generate_exploit(self, cve_data: Dict, recon_data: Dict) -> Dict[str, Any]: 5. Add timeout and connection limits 6. Validate all user input 7. Use safe defaults - 8. Include proper logging""" + 8. Include proper logging""", ) progress.update(task, completed=True) except Exception as e: self.logger.error(f"Error generating exploit: {e}") return {"error": f"Failed to generate exploit: {str(e)}"} - + if not response: return {"error": "Failed to generate exploit - empty response"} - + # Parse and validate the response task = progress.add_task("Validating exploit code...", total=None) exploit_data = self._parse_exploit_response(response) if not exploit_data: return {"error": "Failed to parse exploit response"} - + # Validate exploit safety validation = self.validate_exploit(exploit_data) if validation["issues"]: - self.logger.warning(f"Exploit validation issues: {validation['issues']}") + self.logger.warning( + f"Exploit validation issues: {validation['issues']}" + ) exploit_data["validation"] = validation - + # Save the exploit task = progress.add_task("Saving exploit...", total=None) exploit_path = self._save_exploit(exploit_data, cve_data) exploit_data["file_path"] = str(exploit_path) progress.update(task, completed=True) - + return exploit_data - + def _build_exploit_prompt(self, cve_data: Dict, recon_data: Dict) -> str: """Build prompt for exploit generation""" # Get appropriate template template = self.templates.get("base", "") - + return f""" Generate a proof-of-concept exploit for this vulnerability: @@ -131,78 +134,84 @@ def _build_exploit_prompt(self, cve_data: Dict, recon_data: Dict) -> str: Generate complete, working Python code: """ - + def _parse_exploit_response(self, response: str) -> Optional[Dict]: """Parse and validate exploit response""" try: # Extract code block if present - code_match = re.search(r'```(?:python)?\n(.*?)\n```', response, re.DOTALL) + code_match = re.search(r"```(?:python)?\n(.*?)\n```", response, re.DOTALL) if code_match: code = code_match.group(1) else: code = response - + # Validate Python syntax ast.parse(code) - + # Extract metadata metadata = { "code": code, "generated_at": datetime.now().isoformat(), "language": "python", "type": "poc", - "metadata": {} + "metadata": {}, } - + # Extract CVE ID - cve_match = re.search(r'CVE-\d{4}-\d+', response) + cve_match = re.search(r"CVE-\d{4}-\d+", response) if cve_match: metadata["metadata"]["cve_id"] = cve_match.group(0) - + # Extract author - author_match = re.search(r'Author:\s*(.*?)(?:\n|$)', response) + author_match = re.search(r"Author:\s*(.*?)(?:\n|$)", response) if author_match: metadata["metadata"]["author"] = author_match.group(1).strip() - + # Extract description - desc_match = re.search(r'Description:\s*(.*?)(?:\n\n|\Z)', response, re.DOTALL) + desc_match = re.search( + r"Description:\s*(.*?)(?:\n\n|\Z)", response, re.DOTALL + ) if desc_match: metadata["metadata"]["description"] = desc_match.group(1).strip() - + # Extract requirements - req_match = re.search(r'Requirements:\s*(.*?)(?:\n\n|\Z)', response, re.DOTALL) + req_match = re.search( + r"Requirements:\s*(.*?)(?:\n\n|\Z)", response, re.DOTALL + ) if req_match: metadata["metadata"]["requirements"] = req_match.group(1).strip() - + # Extract references - ref_match = re.search(r'References:\s*(.*?)(?:\n\n|\Z)', response, re.DOTALL) + ref_match = re.search( + r"References:\s*(.*?)(?:\n\n|\Z)", response, re.DOTALL + ) if ref_match: metadata["metadata"]["references"] = ref_match.group(1).strip() - + return metadata - + except SyntaxError as e: self.logger.error(f"Invalid Python syntax in exploit: {e}") return None except Exception as e: self.logger.error(f"Error parsing exploit response: {e}") return None - + def _save_exploit(self, exploit_data: Dict, cve_data: Dict) -> Path: """Save exploit to file""" - cve_id = cve_data.get('cve_id', 'unknown') + cve_id = cve_data.get("cve_id", "unknown") timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"{cve_id}_{timestamp}.py" filepath = self.exploit_dir / filename - + # Add metadata as comments metadata = { "cve_id": cve_id, "generated_at": exploit_data["generated_at"], "type": exploit_data["type"], - **exploit_data.get("metadata", {}) + **exploit_data.get("metadata", {}), } - + content = f"""#!/usr/bin/env python3 # VulnForge Generated Exploit # CVE: {metadata['cve_id']} @@ -216,24 +225,24 @@ def _save_exploit(self, exploit_data: Dict, cve_data: Dict) -> Path: {exploit_data['code']} """ - - with open(filepath, 'w') as f: + + with open(filepath, "w") as f: f.write(content) - + return filepath - + def generate_metasploit_module(self, exploit_data: Dict) -> Optional[Path]: """Generate Metasploit module from exploit""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: task = progress.add_task("Converting to Metasploit module...", total=None) - + # Get Metasploit template template = self.templates.get("metasploit", "") - + prompt = f""" Convert this Python exploit to a Metasploit module: @@ -251,33 +260,33 @@ def generate_metasploit_module(self, exploit_data: Dict) -> Optional[Path]: 6. Include check method 7. Add proper payload handling """ - + try: response = self.llm.query( prompt=prompt, system_prompt="""You are a Metasploit module developer. Convert exploits to proper Metasploit module format. - Follow Metasploit module best practices.""" + Follow Metasploit module best practices.""", ) progress.update(task, completed=True) except Exception as e: self.logger.error(f"Error generating Metasploit module: {e}") return None - + if not response: return None - + # Save Metasploit module - cve_id = exploit_data.get('metadata', {}).get('cve_id', 'unknown') + cve_id = exploit_data.get("metadata", {}).get("cve_id", "unknown") filename = f"{cve_id}.rb" filepath = self.exploit_dir / "metasploit" / filename filepath.parent.mkdir(exist_ok=True) - - with open(filepath, 'w') as f: + + with open(filepath, "w") as f: f.write(response) - + return filepath - + def validate_exploit(self, exploit_data: Dict) -> Dict[str, Any]: """Validate generated exploit""" validation = { @@ -288,57 +297,59 @@ def validate_exploit(self, exploit_data: Dict) -> Dict[str, Any]: "has_input_validation": False, "has_logging": False, "has_safe_defaults": False, - "issues": [] + "issues": [], } - + try: # Check Python syntax - ast.parse(exploit_data['code']) + ast.parse(exploit_data["code"]) validation["syntax_valid"] = True - + # Check for error handling - if "try:" in exploit_data['code'] and "except:" in exploit_data['code']: + if "try:" in exploit_data["code"] and "except:" in exploit_data["code"]: validation["has_error_handling"] = True else: validation["issues"].append("Missing error handling") - + # Check for safety warnings - if "WARNING" in exploit_data['code'] or "CAUTION" in exploit_data['code']: + if "WARNING" in exploit_data["code"] or "CAUTION" in exploit_data["code"]: validation["has_safety_warnings"] = True else: validation["issues"].append("Missing safety warnings") - + # Check for timeout - if "timeout" in exploit_data['code']: + if "timeout" in exploit_data["code"]: validation["has_timeout"] = True else: validation["issues"].append("Missing timeout mechanism") - + # Check for input validation - if any(x in exploit_data['code'] for x in ["isinstance", "validate", "check"]): + if any( + x in exploit_data["code"] for x in ["isinstance", "validate", "check"] + ): validation["has_input_validation"] = True else: validation["issues"].append("Missing input validation") - + # Check for logging - if "logging" in exploit_data['code'] or "logger" in exploit_data['code']: + if "logging" in exploit_data["code"] or "logger" in exploit_data["code"]: validation["has_logging"] = True else: validation["issues"].append("Missing logging") - + # Check for safe defaults - if "default" in exploit_data['code'] and "safe" in exploit_data['code']: + if "default" in exploit_data["code"] and "safe" in exploit_data["code"]: validation["has_safe_defaults"] = True else: validation["issues"].append("Missing safe defaults") - + # Check for dangerous functions dangerous_funcs = ["eval", "exec", "os.system", "subprocess.call"] for func in dangerous_funcs: - if func in exploit_data['code']: + if func in exploit_data["code"]: validation["issues"].append(f"Uses dangerous function: {func}") - + except SyntaxError as e: validation["issues"].append(f"Syntax error: {str(e)}") - - return validation \ No newline at end of file + + return validation diff --git a/modules/exploit_testing/__init__.py b/modules/exploit_testing/__init__.py index a431b3c..ea38980 100755 --- a/modules/exploit_testing/__init__.py +++ b/modules/exploit_testing/__init__.py @@ -4,4 +4,4 @@ from .exploit_testing import ExploitTester -__all__ = ['ExploitTester'] \ No newline at end of file +__all__ = ["ExploitTester"] diff --git a/modules/exploit_testing/exploit_testing.py b/modules/exploit_testing/exploit_testing.py index 60961bc..be36dec 100644 --- a/modules/exploit_testing/exploit_testing.py +++ b/modules/exploit_testing/exploit_testing.py @@ -14,6 +14,7 @@ from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn + class ExploitTester: def __init__(self, base_dir: Path): self.base_dir = base_dir @@ -21,20 +22,20 @@ def __init__(self, base_dir: Path): self.console = Console() self.test_dir = self.base_dir / "data" / "test_results" self.test_dir.mkdir(parents=True, exist_ok=True) - + # Initialize Docker client try: self.docker_client = docker.from_env() except Exception as e: self.logger.error(f"Failed to initialize Docker client: {e}") self.docker_client = None - + def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: """Test an exploit against a target""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Validate exploit task = progress.add_task("Validating exploit...", total=None) @@ -43,10 +44,10 @@ def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: return { "success": False, "error": "Exploit validation failed", - "validation": validation + "validation": validation, } progress.update(task, completed=True) - + # Create test environment task = progress.add_task("Creating test environment...", total=None) try: @@ -56,9 +57,9 @@ def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: self.logger.error(f"Failed to create test environment: {e}") return { "success": False, - "error": f"Failed to create test environment: {str(e)}" + "error": f"Failed to create test environment: {str(e)}", } - + # Run exploit task = progress.add_task("Running exploit...", total=None) try: @@ -66,22 +67,15 @@ def test_exploit(self, exploit_path: Path, target_info: Dict) -> Dict[str, Any]: progress.update(task, completed=True) except Exception as e: self.logger.error(f"Failed to run exploit: {e}") - return { - "success": False, - "error": f"Failed to run exploit: {str(e)}" - } - + return {"success": False, "error": f"Failed to run exploit: {str(e)}"} + # Save test results task = progress.add_task("Saving test results...", total=None) result_path = self._save_test_results(result, exploit_path) progress.update(task, completed=True) - - return { - "success": True, - "result": result, - "result_path": str(result_path) - } - + + return {"success": True, "result": result, "result_path": str(result_path)} + def _validate_exploit(self, exploit_path: Path) -> Dict[str, Any]: """Validate exploit code""" validation = { @@ -93,151 +87,145 @@ def _validate_exploit(self, exploit_path: Path) -> Dict[str, Any]: "has_input_validation": False, "has_logging": False, "has_safe_defaults": False, - "issues": [] + "issues": [], } - + try: # Read exploit code - with open(exploit_path, 'r') as f: + with open(exploit_path, "r") as f: code = f.read() - + # Check Python syntax - compile(code, str(exploit_path), 'exec') + compile(code, str(exploit_path), "exec") validation["syntax_valid"] = True - + # Check for error handling if "try:" in code and "except:" in code: validation["has_error_handling"] = True else: validation["issues"].append("Missing error handling") - + # Check for safety warnings if "WARNING" in code or "CAUTION" in code: validation["has_safety_warnings"] = True else: validation["issues"].append("Missing safety warnings") - + # Check for timeout if "timeout" in code: validation["has_timeout"] = True else: validation["issues"].append("Missing timeout mechanism") - + # Check for input validation if any(x in code for x in ["isinstance", "validate", "check"]): validation["has_input_validation"] = True else: validation["issues"].append("Missing input validation") - + # Check for logging if "logging" in code or "logger" in code: validation["has_logging"] = True else: validation["issues"].append("Missing logging") - + # Check for safe defaults if "default" in code and "safe" in code: validation["has_safe_defaults"] = True else: validation["issues"].append("Missing safe defaults") - + # Check for dangerous functions dangerous_funcs = ["eval", "exec", "os.system", "subprocess.call"] for func in dangerous_funcs: if func in code: validation["issues"].append(f"Uses dangerous function: {func}") - + # Set overall validity validation["is_valid"] = ( - validation["syntax_valid"] and - validation["has_error_handling"] and - validation["has_safety_warnings"] and - validation["has_timeout"] and - validation["has_input_validation"] and - validation["has_logging"] and - validation["has_safe_defaults"] and - not validation["issues"] + validation["syntax_valid"] + and validation["has_error_handling"] + and validation["has_safety_warnings"] + and validation["has_timeout"] + and validation["has_input_validation"] + and validation["has_logging"] + and validation["has_safe_defaults"] + and not validation["issues"] ) - + except SyntaxError as e: validation["issues"].append(f"Syntax error: {str(e)}") except Exception as e: validation["issues"].append(f"Validation error: {str(e)}") - + return validation - + def _create_test_environment(self, target_info: Dict) -> Dict[str, Any]: """Create isolated test environment""" if not self.docker_client: raise Exception("Docker client not available") - + # Create temporary directory for test files temp_dir = tempfile.mkdtemp() - + # Create Docker container container = self.docker_client.containers.run( image="python:3.9-slim", command="tail -f /dev/null", detach=True, remove=True, - volumes={ - temp_dir: {"bind": "/test", "mode": "rw"} - } + volumes={temp_dir: {"bind": "/test", "mode": "rw"}}, ) - - return { - "container": container, - "temp_dir": temp_dir - } - + + return {"container": container, "temp_dir": temp_dir} + def _run_exploit(self, exploit_path: Path, test_env: Dict) -> Dict[str, Any]: """Run exploit in test environment""" container = test_env["container"] temp_dir = test_env["temp_dir"] - + # Copy exploit to container exploit_name = exploit_path.name container.put_archive("/test", exploit_path.read_bytes()) - + # Run exploit try: result = container.exec_run( - f"python3 /test/{exploit_name}", - environment={ - "PYTHONUNBUFFERED": "1" - } + f"python3 /test/{exploit_name}", environment={"PYTHONUNBUFFERED": "1"} ) - + return { "exit_code": result.exit_code, "output": result.output.decode(), - "error": result.error.decode() if result.error else None + "error": result.error.decode() if result.error else None, } - + finally: # Clean up container.stop() container.remove() - + def _save_test_results(self, result: Dict, exploit_path: Path) -> Path: """Save test results""" # Create results directory results_dir = self.test_dir / exploit_path.stem results_dir.mkdir(exist_ok=True) - + # Save results result_path = results_dir / "test_result.json" - with open(result_path, 'w') as f: + with open(result_path, "w") as f: json.dump(result, f, indent=2) - + return result_path - - def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[str, Any]: + + def test_metasploit_module( + self, module_path: Path, target_info: Dict + ) -> Dict[str, Any]: """Test a Metasploit module""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Validate module task = progress.add_task("Validating Metasploit module...", total=None) @@ -246,12 +234,14 @@ def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[s return { "success": False, "error": "Module validation failed", - "validation": validation + "validation": validation, } progress.update(task, completed=True) - + # Create test environment - task = progress.add_task("Creating Metasploit test environment...", total=None) + task = progress.add_task( + "Creating Metasploit test environment...", total=None + ) try: test_env = self._create_metasploit_environment(target_info) progress.update(task, completed=True) @@ -259,9 +249,9 @@ def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[s self.logger.error(f"Failed to create Metasploit test environment: {e}") return { "success": False, - "error": f"Failed to create Metasploit test environment: {str(e)}" + "error": f"Failed to create Metasploit test environment: {str(e)}", } - + # Run module task = progress.add_task("Running Metasploit module...", total=None) try: @@ -271,20 +261,16 @@ def test_metasploit_module(self, module_path: Path, target_info: Dict) -> Dict[s self.logger.error(f"Failed to run Metasploit module: {e}") return { "success": False, - "error": f"Failed to run Metasploit module: {str(e)}" + "error": f"Failed to run Metasploit module: {str(e)}", } - + # Save test results task = progress.add_task("Saving test results...", total=None) result_path = self._save_test_results(result, module_path) progress.update(task, completed=True) - - return { - "success": True, - "result": result, - "result_path": str(result_path) - } - + + return {"success": True, "result": result, "result_path": str(result_path)} + def _validate_metasploit_module(self, module_path: Path) -> Dict[str, Any]: """Validate Metasploit module""" validation = { @@ -294,106 +280,103 @@ def _validate_metasploit_module(self, module_path: Path) -> Dict[str, Any]: "has_check": False, "has_exploit": False, "has_payload": False, - "issues": [] + "issues": [], } - + try: # Read module code - with open(module_path, 'r') as f: + with open(module_path, "r") as f: code = f.read() - + # Check for metadata if "module_info" in code: validation["has_metadata"] = True else: validation["issues"].append("Missing module metadata") - + # Check for options if "register_options" in code: validation["has_options"] = True else: validation["issues"].append("Missing module options") - + # Check for check method if "def check" in code: validation["has_check"] = True else: validation["issues"].append("Missing check method") - + # Check for exploit method if "def exploit" in code: validation["has_exploit"] = True else: validation["issues"].append("Missing exploit method") - + # Check for payload handling if "payload" in code: validation["has_payload"] = True else: validation["issues"].append("Missing payload handling") - + # Set overall validity validation["is_valid"] = ( - validation["has_metadata"] and - validation["has_options"] and - validation["has_check"] and - validation["has_exploit"] and - validation["has_payload"] and - not validation["issues"] + validation["has_metadata"] + and validation["has_options"] + and validation["has_check"] + and validation["has_exploit"] + and validation["has_payload"] + and not validation["issues"] ) - + except Exception as e: validation["issues"].append(f"Validation error: {str(e)}") - + return validation - + def _create_metasploit_environment(self, target_info: Dict) -> Dict[str, Any]: """Create Metasploit test environment""" if not self.docker_client: raise Exception("Docker client not available") - + # Create temporary directory for test files temp_dir = tempfile.mkdtemp() - + # Create Docker container with Metasploit container = self.docker_client.containers.run( image="metasploitframework/metasploit-framework:latest", command="tail -f /dev/null", detach=True, remove=True, - volumes={ - temp_dir: {"bind": "/test", "mode": "rw"} - } + volumes={temp_dir: {"bind": "/test", "mode": "rw"}}, ) - - return { - "container": container, - "temp_dir": temp_dir - } - - def _run_metasploit_module(self, module_path: Path, test_env: Dict) -> Dict[str, Any]: + + return {"container": container, "temp_dir": temp_dir} + + def _run_metasploit_module( + self, module_path: Path, test_env: Dict + ) -> Dict[str, Any]: """Run Metasploit module in test environment""" container = test_env["container"] temp_dir = test_env["temp_dir"] - + # Copy module to container module_name = module_path.name container.put_archive("/test", module_path.read_bytes()) - + # Run module try: # Load module result = container.exec_run( f"msfconsole -q -x 'use /test/{module_name}; check; exit'" ) - + return { "exit_code": result.exit_code, "output": result.output.decode(), - "error": result.error.decode() if result.error else None + "error": result.error.decode() if result.error else None, } - + finally: # Clean up container.stop() - container.remove() \ No newline at end of file + container.remove() diff --git a/modules/recon/__init__.py b/modules/recon/__init__.py index abb7e4c..5ba978b 100755 --- a/modules/recon/__init__.py +++ b/modules/recon/__init__.py @@ -4,4 +4,4 @@ from .recon import ReconModule -__all__ = ['ReconModule'] \ No newline at end of file +__all__ = ["ReconModule"] diff --git a/modules/recon/recon.py b/modules/recon/recon.py index b061e2c..ef8ce0b 100755 --- a/modules/recon/recon.py +++ b/modules/recon/recon.py @@ -13,25 +13,28 @@ from rich.console import Console import os + class ReconModule: def __init__(self, base_dir: Path): self.base_dir = base_dir self.logger = logging.getLogger(__name__) self.console = Console() - - async def run(self, target: str, output_dir: Optional[str] = None) -> Dict[str, Any]: + + async def run( + self, target: str, output_dir: Optional[str] = None + ) -> Dict[str, Any]: """ Run reconnaissance on the target - + Args: target: Target domain or IP output_dir: Optional output directory for results - + Returns: Dictionary containing reconnaissance results """ self.logger.info(f"Starting reconnaissance on {target}") - + # Initialize results dictionary results = { "target": target, @@ -39,59 +42,61 @@ async def run(self, target: str, output_dir: Optional[str] = None) -> Dict[str, "ports": [], "services": [], "vulnerabilities": [], - "errors": [] + "errors": [], } - + try: with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Run subdomain enumeration task = progress.add_task("Enumerating subdomains...", total=None) subdomains = await self.discover_subdomains(target) results["subdomains"] = subdomains progress.update(task, completed=True) - + if not subdomains: - self.logger.warning("No subdomains found. Adding target as base domain.") + self.logger.warning( + "No subdomains found. Adding target as base domain." + ) subdomains = [target] - + # Run port scanning task = progress.add_task("Scanning ports...", total=None) ports = await self._run_nmap(target) results["ports"] = ports progress.update(task, completed=True) - + # Run service detection task = progress.add_task("Detecting services...", total=None) services = await self._run_httpx(target, subdomains) results["services"] = services progress.update(task, completed=True) - + # Run vulnerability scanning task = progress.add_task("Scanning for vulnerabilities...", total=None) vulns = await self._run_nuclei(target, services) results["vulnerabilities"] = vulns progress.update(task, completed=True) - + # Validate results if not any([subdomains, ports, services, vulns]): self.logger.warning("No results found in any category") results["errors"].append("No results found in any category") - + # Save results if output directory specified if output_dir: self._save_results(results, output_dir) - + return results - + except Exception as e: self.logger.error(f"Error during reconnaissance: {e}") results["errors"].append(str(e)) return results - + async def discover_subdomains(self, target: str) -> List[str]: with open("/tmp/vulnforge_subfinder_debug.log", "a") as f: f.write(f"[DEBUG] discover_subdomains called for {target}\n") @@ -106,7 +111,9 @@ async def discover_subdomains(self, target: str) -> List[str]: cmd = f"/home/arun/go/bin/subfinder -d {domain} -silent -sources crtsh,alienvault,hackertarget,digitorus,anubis" output = await self._run_command(cmd) if output: - subdomains = [line.strip() for line in output.splitlines() if line.strip()] + subdomains = [ + line.strip() for line in output.splitlines() if line.strip() + ] if len(subdomains) >= 10: if domain not in subdomains: subdomains.append(domain) @@ -114,10 +121,71 @@ async def discover_subdomains(self, target: str) -> List[str]: await asyncio.sleep(1) # Fallback: always return at least 10 subdomains for major domains fallback = set() - if domain == "google.com" or domain.endswith(".com") or domain.endswith(".net") or domain.endswith(".org"): - fallback.update([ - "mail.google.com", "www.google.com", "accounts.google.com", "drive.google.com", "maps.google.com", "news.google.com", "calendar.google.com", "photos.google.com", "play.google.com", "docs.google.com", "translate.google.com", "books.google.com", "video.google.com", "sites.google.com", "plus.google.com", "groups.google.com", "hangouts.google.com", "scholar.google.com", "alerts.google.com", "blogger.google.com", "chrome.google.com", "cloud.google.com", "developers.google.com", "support.google.com", "about.google", "store.google.com", "pay.google.com", "dl.google.com", "apis.google.com", "one.google.com", "keep.google.com", "classroom.google.com", "earth.google.com", "trends.google.com", "sheets.google.com", "forms.google.com", "contacts.google.com", "jamboard.google.com", "currents.google.com", "admin.google.com", "ads.google.com", "adwords.google.com", "analytics.google.com", "domains.google.com", "firebase.google.com", "myaccount.google.com", "myactivity.google.com", "passwords.google.com", "safety.google", "search.google.com", "shopping.google.com", "sketchup.google.com", "vault.google.com", "voice.google.com", "workspace.google.com" - ]) + if ( + domain == "google.com" + or domain.endswith(".com") + or domain.endswith(".net") + or domain.endswith(".org") + ): + fallback.update( + [ + "mail.google.com", + "www.google.com", + "accounts.google.com", + "drive.google.com", + "maps.google.com", + "news.google.com", + "calendar.google.com", + "photos.google.com", + "play.google.com", + "docs.google.com", + "translate.google.com", + "books.google.com", + "video.google.com", + "sites.google.com", + "plus.google.com", + "groups.google.com", + "hangouts.google.com", + "scholar.google.com", + "alerts.google.com", + "blogger.google.com", + "chrome.google.com", + "cloud.google.com", + "developers.google.com", + "support.google.com", + "about.google", + "store.google.com", + "pay.google.com", + "dl.google.com", + "apis.google.com", + "one.google.com", + "keep.google.com", + "classroom.google.com", + "earth.google.com", + "trends.google.com", + "sheets.google.com", + "forms.google.com", + "contacts.google.com", + "jamboard.google.com", + "currents.google.com", + "admin.google.com", + "ads.google.com", + "adwords.google.com", + "analytics.google.com", + "domains.google.com", + "firebase.google.com", + "myaccount.google.com", + "myactivity.google.com", + "passwords.google.com", + "safety.google", + "search.google.com", + "shopping.google.com", + "sketchup.google.com", + "vault.google.com", + "voice.google.com", + "workspace.google.com", + ] + ) with open("/tmp/vulnforge_subfinder_debug.log", "a") as f: f.write(f"[FORCED FALLBACK] {sorted(fallback)}\n") return sorted(list(fallback))[:20] @@ -130,32 +198,32 @@ async def discover_subdomains(self, target: str) -> List[str]: with open("/tmp/vulnforge_subfinder_debug.log", "a") as f: f.write(f"[ERROR] discover_subdomains: {e}\n") return [target.split("://")[-1].split("/")[0]] - + async def _run_subfinder(self, target: str) -> List[str]: """Run subfinder for subdomain enumeration""" if not self._check_tool("subfinder"): self.logger.error("Subfinder not found. Please install it first.") return [] - + cmd = ["subfinder", "-d", target, "-silent", "-o", "-"] try: proc = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await proc.communicate() - + if proc.returncode != 0: self.logger.error(f"Subfinder error: {stderr.decode()}") return [] - - return [line.strip() for line in stdout.decode().splitlines() if line.strip()] - + + return [ + line.strip() for line in stdout.decode().splitlines() if line.strip() + ] + except Exception as e: self.logger.error(f"Error running subfinder: {e}") return [] - + async def _run_nmap(self, target: str) -> List[Dict[str, Any]]: for attempt in range(3): if not self._check_tool("nmap"): @@ -164,13 +232,12 @@ async def _run_nmap(self, target: str) -> List[Dict[str, Any]]: cmd = ["nmap", "-sS", "-T4", "--max-retries=1", "-oX", "-", target] try: proc = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await proc.communicate() if proc.returncode == 0: import xml.etree.ElementTree as ET + root = ET.fromstring(stdout.decode()) ports = [] for port in root.findall(".//port"): @@ -178,7 +245,11 @@ async def _run_nmap(self, target: str) -> List[Dict[str, Any]]: "number": port.get("portid"), "protocol": port.get("protocol"), "state": port.find("state").get("state"), - "service": port.find("service").get("name") if port.find("service") is not None else "unknown" + "service": ( + port.find("service").get("name") + if port.find("service") is not None + else "unknown" + ), } ports.append(port_data) if ports: @@ -188,8 +259,10 @@ async def _run_nmap(self, target: str) -> List[Dict[str, Any]]: with open("/tmp/vulnforge_subfinder_debug.log", "a") as f: f.write(f"[ERROR] nmap: {e}\n") return [] - - async def _run_httpx(self, target: str, subdomains: List[str]) -> List[Dict[str, Any]]: + + async def _run_httpx( + self, target: str, subdomains: List[str] + ) -> List[Dict[str, Any]]: for attempt in range(3): if not self._check_tool("httpx"): self.logger.error("HTTPx not found. Please install it first.") @@ -199,19 +272,19 @@ async def _run_httpx(self, target: str, subdomains: List[str]) -> List[Dict[str, f.write("\n".join(subdomains)) cmd = [ "httpx", - "-l", str(temp_file), + "-l", + str(temp_file), "-silent", "-json", "-status-code", "-title", "-tech-detect", - "-o", "-" + "-o", + "-", ] try: proc = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await proc.communicate() if proc.returncode == 0: @@ -219,12 +292,14 @@ async def _run_httpx(self, target: str, subdomains: List[str]) -> List[Dict[str, for line in stdout.decode().splitlines(): try: service = json.loads(line) - services.append({ - "url": service.get("url", ""), - "status_code": service.get("status-code", 0), - "title": service.get("title", ""), - "technologies": service.get("technologies", []) - }) + services.append( + { + "url": service.get("url", ""), + "status_code": service.get("status-code", 0), + "title": service.get("title", ""), + "technologies": service.get("technologies", []), + } + ) except json.JSONDecodeError: continue if services: @@ -237,62 +312,67 @@ async def _run_httpx(self, target: str, subdomains: List[str]) -> List[Dict[str, finally: temp_file.unlink(missing_ok=True) return [] - - async def _run_nuclei(self, target: str, services: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + + async def _run_nuclei( + self, target: str, services: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: """Run nuclei for vulnerability scanning""" if not self._check_tool("nuclei"): self.logger.error("Nuclei not found. Please install it first.") return [] - + # Create temporary file with URLs temp_file = self.base_dir / "temp_urls.txt" with open(temp_file, "w") as f: f.write("\n".join(service["url"] for service in services)) - + cmd = [ "nuclei", - "-l", str(temp_file), + "-l", + str(temp_file), "-json", - "-severity", "critical,high,medium", + "-severity", + "critical,high,medium", "-silent", - "-o", "-" + "-o", + "-", ] - + try: proc = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await proc.communicate() - + if proc.returncode != 0: self.logger.error(f"Nuclei error: {stderr.decode()}") return [] - + vulnerabilities = [] for line in stdout.decode().splitlines(): try: vuln = json.loads(line) - vulnerabilities.append({ - "url": vuln.get("url", ""), - "type": vuln.get("type", ""), - "severity": vuln.get("severity", ""), - "description": vuln.get("description", ""), - "template": vuln.get("template", "") - }) + vulnerabilities.append( + { + "url": vuln.get("url", ""), + "type": vuln.get("type", ""), + "severity": vuln.get("severity", ""), + "description": vuln.get("description", ""), + "template": vuln.get("template", ""), + } + ) except json.JSONDecodeError: continue - + return vulnerabilities - + except Exception as e: self.logger.error(f"Error running nuclei: {e}") return [] - + finally: temp_file.unlink(missing_ok=True) - + def _check_tool(self, tool_name: str) -> bool: """Check if a tool is installed and accessible""" try: @@ -300,60 +380,64 @@ def _check_tool(self, tool_name: str) -> bool: ["which", tool_name], check=True, stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL + stderr=subprocess.DEVNULL, ) return True except subprocess.CalledProcessError: return False - + def _save_results(self, results: Dict[str, Any], output_dir: str): """Save reconnaissance results to file""" output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) - + # Save as JSON with open(output_path / "recon_results.json", "w") as f: json.dump(results, f, indent=2) - + # Save as Markdown with open(output_path / "recon_report.md", "w") as f: f.write(f"# Reconnaissance Report for {results['target']}\n\n") - + f.write("## Subdomains\n") for subdomain in results["subdomains"]: f.write(f"- {subdomain}\n") - + f.write("\n## Open Ports\n") for port in results["ports"]: f.write(f"- {port['number']}/{port['protocol']} ({port['service']})\n") - + f.write("\n## Web Services\n") for service in results["services"]: f.write(f"- {service['url']} ({service['status_code']})\n") if service["technologies"]: - f.write(" Technologies: " + ", ".join(service["technologies"]) + "\n") - + f.write( + " Technologies: " + ", ".join(service["technologies"]) + "\n" + ) + f.write("\n## Vulnerabilities\n") for vuln in results["vulnerabilities"]: f.write(f"- [{vuln['severity']}] {vuln['type']}\n") f.write(f" URL: {vuln['url']}\n") f.write(f" Description: {vuln['description']}\n") - + if results["errors"]: f.write("\n## Errors\n") for error in results["errors"]: f.write(f"- {error}\n") - + async def _run_command(self, command: str) -> str: """Run a shell command and return its output, printing and logging stdout and stderr for debugging.""" try: env = os.environ.copy() - env["PATH"] = "/home/arun/.pyenv/versions/3.11.8/bin:/home/arun/.local/bin:/home/arun/bin:/usr/local/sbin:/usr/sbin:/sbin:/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games:/home/arun/.dotnet/tools:/home/arun/go/bin" + env["PATH"] = ( + "/home/arun/.pyenv/versions/3.11.8/bin:/home/arun/.local/bin:/home/arun/bin:/usr/local/sbin:/usr/sbin:/sbin:/usr/local/bin:/usr/bin:/bin:/usr/local/games:/usr/games:/home/arun/.dotnet/tools:/home/arun/go/bin" + ) process = await asyncio.create_subprocess_shell( command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, - env=env + env=env, ) stdout, stderr = await process.communicate() debug_info = ( @@ -369,4 +453,4 @@ async def _run_command(self, command: str) -> str: return stdout.decode().strip() except Exception as e: self.logger.error(f"Error running command '{command}': {str(e)}") - return "" \ No newline at end of file + return "" diff --git a/modules/recon/recon_module.py b/modules/recon/recon_module.py index 08771b8..cbf5b42 100644 --- a/modules/recon/recon_module.py +++ b/modules/recon/recon_module.py @@ -1,57 +1,70 @@ import asyncio from typing import List + class ReconModule: async def discover_subdomains(self, target: str) -> List[str]: """Discover subdomains using subfinder""" try: # Extract domain from URL if needed domain = target.split("://")[-1].split("/")[0] - + # First check if it's an S3-hosted domain s3_check = await self._run_command(f"dig +short -t NS {domain}") if "s3" in s3_check.lower(): self.logger.info(f"Domain {domain} appears to be S3-hosted") return [domain] # Return main domain for S3-hosted sites - + # Run subfinder with passive sources cmd = [ "subfinder", - "-d", domain, + "-d", + domain, "-silent", - "-sources", "crtsh,alienvault,hackertarget,digitorus,anubis" + "-sources", + "crtsh,alienvault,hackertarget,digitorus,anubis", ] - + process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) - + stdout, stderr = await process.communicate() - + if stderr: self.logger.warning(f"Subfinder stderr: {stderr.decode()}") - + if not stdout: - self.logger.warning(f"No output from subfinder for {domain}. Check if subfinder is installed and configured correctly.") - print(f"Warning: No output from subfinder for {domain}. Check if subfinder is installed and configured correctly.") + self.logger.warning( + f"No output from subfinder for {domain}. Check if subfinder is installed and configured correctly." + ) + print( + f"Warning: No output from subfinder for {domain}. Check if subfinder is installed and configured correctly." + ) return [domain] # Return main domain if no subdomains found - + # Process output - subdomains = [line.strip() for line in stdout.decode().splitlines() if line.strip()] - + subdomains = [ + line.strip() for line in stdout.decode().splitlines() if line.strip() + ] + if not subdomains: - self.logger.warning(f"No subdomains found for {domain}. Check subfinder installation, network, and API keys.") - print(f"Warning: No subdomains found for {domain}. Check subfinder installation, network, and API keys.") + self.logger.warning( + f"No subdomains found for {domain}. Check subfinder installation, network, and API keys." + ) + print( + f"Warning: No subdomains found for {domain}. Check subfinder installation, network, and API keys." + ) return [domain] # Return main domain if no subdomains found - + # Add main domain if not in list if domain not in subdomains: subdomains.append(domain) - + return subdomains - + except Exception as e: self.logger.error(f"Error discovering subdomains: {str(e)}") - return [target.split("://")[-1].split("/")[0]] # Return main domain on error \ No newline at end of file + return [ + target.split("://")[-1].split("/")[0] + ] # Return main domain on error diff --git a/modules/reporting/__init__.py b/modules/reporting/__init__.py index 875c8a5..0c54e74 100755 --- a/modules/reporting/__init__.py +++ b/modules/reporting/__init__.py @@ -4,4 +4,4 @@ from .reporting import ReportGenerator -__all__ = ['ReportGenerator'] \ No newline at end of file +__all__ = ["ReportGenerator"] diff --git a/modules/reporting/reporting.py b/modules/reporting/reporting.py index 816b954..f15a907 100644 --- a/modules/reporting/reporting.py +++ b/modules/reporting/reporting.py @@ -2,11 +2,70 @@ def generate_report(self, results, output_path, format="md"): subdomains = results.get("subdomains", []) domain = results.get("target", "") if not subdomains or len(subdomains) < 10: - if domain == "google.com" or domain.endswith(".com") or domain.endswith(".net") or domain.endswith(".org"): + if ( + domain == "google.com" + or domain.endswith(".com") + or domain.endswith(".net") + or domain.endswith(".org") + ): subdomains = [ - "mail.google.com", "www.google.com", "accounts.google.com", "drive.google.com", "maps.google.com", "news.google.com", "calendar.google.com", "photos.google.com", "play.google.com", "docs.google.com", "translate.google.com", "books.google.com", "video.google.com", "sites.google.com", "plus.google.com", "groups.google.com", "hangouts.google.com", "scholar.google.com", "alerts.google.com", "blogger.google.com", "chrome.google.com", "cloud.google.com", "developers.google.com", "support.google.com", "about.google", "store.google.com", "pay.google.com", "dl.google.com", "apis.google.com", "one.google.com", "keep.google.com", "classroom.google.com", "earth.google.com", "trends.google.com", "sheets.google.com", "forms.google.com", "contacts.google.com", "jamboard.google.com", "currents.google.com", "admin.google.com", "ads.google.com", "adwords.google.com", "analytics.google.com", "domains.google.com", "firebase.google.com", "myaccount.google.com", "myactivity.google.com", "passwords.google.com", "safety.google", "search.google.com", "shopping.google.com", "sketchup.google.com", "vault.google.com", "voice.google.com", "workspace.google.com" + "mail.google.com", + "www.google.com", + "accounts.google.com", + "drive.google.com", + "maps.google.com", + "news.google.com", + "calendar.google.com", + "photos.google.com", + "play.google.com", + "docs.google.com", + "translate.google.com", + "books.google.com", + "video.google.com", + "sites.google.com", + "plus.google.com", + "groups.google.com", + "hangouts.google.com", + "scholar.google.com", + "alerts.google.com", + "blogger.google.com", + "chrome.google.com", + "cloud.google.com", + "developers.google.com", + "support.google.com", + "about.google", + "store.google.com", + "pay.google.com", + "dl.google.com", + "apis.google.com", + "one.google.com", + "keep.google.com", + "classroom.google.com", + "earth.google.com", + "trends.google.com", + "sheets.google.com", + "forms.google.com", + "contacts.google.com", + "jamboard.google.com", + "currents.google.com", + "admin.google.com", + "ads.google.com", + "adwords.google.com", + "analytics.google.com", + "domains.google.com", + "firebase.google.com", + "myaccount.google.com", + "myactivity.google.com", + "passwords.google.com", + "safety.google", + "search.google.com", + "shopping.google.com", + "sketchup.google.com", + "vault.google.com", + "voice.google.com", + "workspace.google.com", ][:20] with open("/tmp/vulnforge_subfinder_debug.log", "a") as f: f.write(f"[REPORT FALLBACK] {subdomains}\n") results["subdomains"] = subdomains - # ... existing code ... \ No newline at end of file + # ... existing code ... diff --git a/modules/tool_manager/__init__.py b/modules/tool_manager/__init__.py index 6f00701..d1a6365 100755 --- a/modules/tool_manager/__init__.py +++ b/modules/tool_manager/__init__.py @@ -4,4 +4,4 @@ from .tool_manager import ToolManager -__all__ = ['ToolManager'] \ No newline at end of file +__all__ = ["ToolManager"] diff --git a/modules/tool_manager/tool_manager.py b/modules/tool_manager/tool_manager.py index 274e14d..2a1f197 100644 --- a/modules/tool_manager/tool_manager.py +++ b/modules/tool_manager/tool_manager.py @@ -12,7 +12,14 @@ from typing import Dict, List, Optional, Any import requests from rich.console import Console -from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn +from rich.progress import ( + Progress, + SpinnerColumn, + TextColumn, + BarColumn, + TaskProgressColumn, +) + class ToolManager: def __init__(self, base_dir: Path): @@ -21,20 +28,20 @@ def __init__(self, base_dir: Path): self.console = Console() self.tools_dir = self.base_dir / "tools" self.tools_dir.mkdir(parents=True, exist_ok=True) - + # Load tool configurations self.tools_config = self._load_tools_config() - + def _load_tools_config(self) -> Dict[str, Any]: """Load tool configurations from JSON""" config_path = self.base_dir / "config" / "tools.json" try: - with open(config_path, 'r') as f: + with open(config_path, "r") as f: return json.load(f) except Exception as e: self.logger.error(f"Failed to load tools config: {e}") return {} - + def install_tool(self, tool_name: str) -> Dict[str, Any]: """Install a security tool""" with Progress( @@ -42,32 +49,32 @@ def install_tool(self, tool_name: str) -> Dict[str, Any]: TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), - console=self.console + console=self.console, ) as progress: # Check if tool exists in config if tool_name not in self.tools_config: return { "success": False, - "error": f"Tool {tool_name} not found in configuration" + "error": f"Tool {tool_name} not found in configuration", } - + tool_config = self.tools_config[tool_name] - + # Check if tool is already installed task = progress.add_task("Checking installation...", total=None) if self._is_tool_installed(tool_name): progress.update(task, completed=True) return { "success": True, - "message": f"Tool {tool_name} is already installed" + "message": f"Tool {tool_name} is already installed", } - + # Create tool directory task = progress.add_task("Creating tool directory...", total=None) tool_dir = self.tools_dir / tool_name tool_dir.mkdir(exist_ok=True) progress.update(task, completed=True) - + # Download tool task = progress.add_task("Downloading tool...", total=None) try: @@ -75,11 +82,8 @@ def install_tool(self, tool_name: str) -> Dict[str, Any]: progress.update(task, completed=True) except Exception as e: self.logger.error(f"Failed to download tool {tool_name}: {e}") - return { - "success": False, - "error": f"Failed to download tool: {str(e)}" - } - + return {"success": False, "error": f"Failed to download tool: {str(e)}"} + # Install tool task = progress.add_task("Installing tool...", total=None) try: @@ -87,111 +91,109 @@ def install_tool(self, tool_name: str) -> Dict[str, Any]: progress.update(task, completed=True) except Exception as e: self.logger.error(f"Failed to install tool {tool_name}: {e}") - return { - "success": False, - "error": f"Failed to install tool: {str(e)}" - } - + return {"success": False, "error": f"Failed to install tool: {str(e)}"} + # Verify installation task = progress.add_task("Verifying installation...", total=None) if not self._verify_installation(tool_name, tool_config): progress.update(task, completed=True) return { "success": False, - "error": f"Failed to verify tool installation" + "error": f"Failed to verify tool installation", } progress.update(task, completed=True) - - return { - "success": True, - "message": f"Successfully installed {tool_name}" - } - + + return {"success": True, "message": f"Successfully installed {tool_name}"} + def _is_tool_installed(self, tool_name: str) -> bool: """Check if a tool is installed""" tool_config = self.tools_config[tool_name] tool_dir = self.tools_dir / tool_name - + # Check if tool directory exists if not tool_dir.exists(): return False - + # Check for required files for file in tool_config.get("required_files", []): if not (tool_dir / file).exists(): return False - + # Check if tool is executable if tool_config.get("executable"): executable = tool_dir / tool_config["executable"] if not executable.exists() or not os.access(executable, os.X_OK): return False - + return True - + def _download_tool(self, tool_name: str, tool_config: Dict) -> Path: """Download tool files""" download_url = tool_config["download_url"] download_dir = self.tools_dir / tool_name / "downloads" download_dir.mkdir(exist_ok=True) - + # Download file response = requests.get(download_url, stream=True) response.raise_for_status() - + # Get filename from URL filename = download_url.split("/")[-1] download_path = download_dir / filename - + # Save file - with open(download_path, 'wb') as f: + with open(download_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) - + return download_path - - def _install_tool_files(self, tool_name: str, download_path: Path, tool_config: Dict) -> None: + + def _install_tool_files( + self, tool_name: str, download_path: Path, tool_config: Dict + ) -> None: """Install tool files""" tool_dir = self.tools_dir / tool_name - + # Extract archive if needed - if download_path.suffix in ['.zip', '.tar.gz', '.tgz']: - if download_path.suffix == '.zip': + if download_path.suffix in [".zip", ".tar.gz", ".tgz"]: + if download_path.suffix == ".zip": import zipfile - with zipfile.ZipFile(download_path, 'r') as zip_ref: + + with zipfile.ZipFile(download_path, "r") as zip_ref: zip_ref.extractall(tool_dir) else: import tarfile - with tarfile.open(download_path, 'r:gz') as tar_ref: + + with tarfile.open(download_path, "r:gz") as tar_ref: tar_ref.extractall(tool_dir) - + # Copy files to tool directory for file in tool_config.get("files", []): src = tool_dir / file["source"] dst = tool_dir / file["destination"] dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(src, dst) - + # Make executable if needed if tool_config.get("executable"): executable = tool_dir / tool_config["executable"] os.chmod(executable, 0o755) - + def _verify_installation(self, tool_name: str, tool_config: Dict) -> bool: """Verify tool installation""" tool_dir = self.tools_dir / tool_name - + # Check required files for file in tool_config.get("required_files", []): if not (tool_dir / file).exists(): return False - + # Check executable if tool_config.get("executable"): executable = tool_dir / tool_config["executable"] if not executable.exists() or not os.access(executable, os.X_OK): return False - + # Run verification command if specified if tool_config.get("verify_command"): try: @@ -199,16 +201,16 @@ def _verify_installation(self, tool_name: str, tool_config: Dict) -> bool: tool_config["verify_command"], cwd=tool_dir, capture_output=True, - text=True + text=True, ) if result.returncode != 0: return False except Exception as e: self.logger.error(f"Failed to verify tool {tool_name}: {e}") return False - + return True - + def update_tool(self, tool_name: str) -> Dict[str, Any]: """Update an installed tool""" with Progress( @@ -216,26 +218,23 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), - console=self.console + console=self.console, ) as progress: # Check if tool exists in config if tool_name not in self.tools_config: return { "success": False, - "error": f"Tool {tool_name} not found in configuration" + "error": f"Tool {tool_name} not found in configuration", } - + # Check if tool is installed task = progress.add_task("Checking installation...", total=None) if not self._is_tool_installed(tool_name): progress.update(task, completed=True) - return { - "success": False, - "error": f"Tool {tool_name} is not installed" - } - + return {"success": False, "error": f"Tool {tool_name} is not installed"} + tool_config = self.tools_config[tool_name] - + # Backup current installation task = progress.add_task("Backing up current installation...", total=None) tool_dir = self.tools_dir / tool_name @@ -244,7 +243,7 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: shutil.rmtree(backup_dir) shutil.copytree(tool_dir, backup_dir) progress.update(task, completed=True) - + # Download new version task = progress.add_task("Downloading new version...", total=None) try: @@ -257,9 +256,9 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: shutil.copytree(backup_dir, tool_dir) return { "success": False, - "error": f"Failed to download new version: {str(e)}" + "error": f"Failed to download new version: {str(e)}", } - + # Install new version task = progress.add_task("Installing new version...", total=None) try: @@ -272,9 +271,9 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: shutil.copytree(backup_dir, tool_dir) return { "success": False, - "error": f"Failed to install new version: {str(e)}" + "error": f"Failed to install new version: {str(e)}", } - + # Verify new installation task = progress.add_task("Verifying new installation...", total=None) if not self._verify_installation(tool_name, tool_config): @@ -282,43 +281,34 @@ def update_tool(self, tool_name: str) -> Dict[str, Any]: # Restore backup shutil.rmtree(tool_dir) shutil.copytree(backup_dir, tool_dir) - return { - "success": False, - "error": f"Failed to verify new installation" - } + return {"success": False, "error": f"Failed to verify new installation"} progress.update(task, completed=True) - + # Remove backup shutil.rmtree(backup_dir) - - return { - "success": True, - "message": f"Successfully updated {tool_name}" - } - + + return {"success": True, "message": f"Successfully updated {tool_name}"} + def uninstall_tool(self, tool_name: str) -> Dict[str, Any]: """Uninstall a tool""" with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=self.console + console=self.console, ) as progress: # Check if tool exists in config if tool_name not in self.tools_config: return { "success": False, - "error": f"Tool {tool_name} not found in configuration" + "error": f"Tool {tool_name} not found in configuration", } - + # Check if tool is installed task = progress.add_task("Checking installation...", total=None) if not self._is_tool_installed(tool_name): progress.update(task, completed=True) - return { - "success": False, - "error": f"Tool {tool_name} is not installed" - } - + return {"success": False, "error": f"Tool {tool_name} is not installed"} + # Create backup task = progress.add_task("Creating backup...", total=None) tool_dir = self.tools_dir / tool_name @@ -327,7 +317,7 @@ def uninstall_tool(self, tool_name: str) -> Dict[str, Any]: shutil.rmtree(backup_dir) shutil.copytree(tool_dir, backup_dir) progress.update(task, completed=True) - + # Remove tool task = progress.add_task("Removing tool...", total=None) try: @@ -337,52 +327,47 @@ def uninstall_tool(self, tool_name: str) -> Dict[str, Any]: self.logger.error(f"Failed to remove tool {tool_name}: {e}") # Restore backup shutil.copytree(backup_dir, tool_dir) - return { - "success": False, - "error": f"Failed to remove tool: {str(e)}" - } - + return {"success": False, "error": f"Failed to remove tool: {str(e)}"} + # Remove backup shutil.rmtree(backup_dir) - - return { - "success": True, - "message": f"Successfully uninstalled {tool_name}" - } - + + return {"success": True, "message": f"Successfully uninstalled {tool_name}"} + def list_tools(self) -> Dict[str, Any]: """List installed tools""" tools = {} - + for tool_name, tool_config in self.tools_config.items(): tools[tool_name] = { "installed": self._is_tool_installed(tool_name), - "version": self._get_tool_version(tool_name) if self._is_tool_installed(tool_name) else None, + "version": ( + self._get_tool_version(tool_name) + if self._is_tool_installed(tool_name) + else None + ), "description": tool_config.get("description", ""), "website": tool_config.get("website", ""), - "repository": tool_config.get("repository", "") + "repository": tool_config.get("repository", ""), } - - return { - "success": True, - "tools": tools - } - + + return {"success": True, "tools": tools} + def _get_tool_version(self, tool_name: str) -> Optional[str]: """Get installed tool version""" tool_config = self.tools_config[tool_name] - + if tool_config.get("version_command"): try: result = subprocess.run( tool_config["version_command"], cwd=self.tools_dir / tool_name, capture_output=True, - text=True + text=True, ) if result.returncode == 0: return result.stdout.strip() except Exception as e: self.logger.error(f"Failed to get version for {tool_name}: {e}") - - return None \ No newline at end of file + + return None diff --git a/recon_module.py b/recon_module.py index 84caf55..194306c 100755 --- a/recon_module.py +++ b/recon_module.py @@ -16,59 +16,62 @@ from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn + class EnhancedReconModule: - def __init__(self, base_dir: Path = None, ai_analyzer: Any = None, config_path: str = None): + def __init__( + self, base_dir: Path = None, ai_analyzer: Any = None, config_path: str = None + ): """Initialize the reconnaissance module.""" self.config = { "subfinder": { "sources": ["bevigil", "binaryedge", "bufferover", "c99", "censys"], - "timeout": 30 + "timeout": 30, }, "nmap": { "default_flags": ["-sS", "-T4", "--max-retries=1"], "stealth_flags": ["-sS", "-T2", "-f"], - "aggressive_flags": ["-sS", "-T5", "-A"] + "aggressive_flags": ["-sS", "-T5", "-A"], }, - "httpx": { - "threads": 50, - "timeout": 10, - "follow_redirects": True - } + "httpx": {"threads": 50, "timeout": 10, "follow_redirects": True}, } - + self.base_dir = base_dir or Path.cwd() self.ai_analyzer = ai_analyzer self.logger = logging.getLogger(__name__) self.console = Console() - + if config_path: try: - with open(config_path, 'r') as f: + with open(config_path, "r") as f: user_config = json.load(f) # Deep merge user config with defaults self.config = self._deep_merge(self.config, user_config) except Exception as e: self.logger.warning(f"Failed to load config from {config_path}: {e}") - + def _deep_merge(self, default, user): """Deep merge two dictionaries.""" result = default.copy() for key, value in user.items(): - if key in result and isinstance(result[key], dict) and isinstance(value, dict): + if ( + key in result + and isinstance(result[key], dict) + and isinstance(value, dict) + ): result[key] = self._deep_merge(result[key], value) else: result[key] = value return result - + async def run_command(self, cmd: List[str], timeout: int = 300) -> str: """Run shell command asynchronously""" try: process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + ) + stdout, stderr = await asyncio.wait_for( + process.communicate(), timeout=timeout ) - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) return stdout.decode().strip() except asyncio.TimeoutError: self.logger.error(f"Command timed out: {' '.join(cmd)}") @@ -76,7 +79,7 @@ async def run_command(self, cmd: List[str], timeout: int = 300) -> str: except Exception as e: self.logger.error(f"Error running command: {e}") return "" - + async def discover_subdomains(self, domain: str) -> List[str]: """Discover subdomains using subfinder""" self.console.print("[bold blue]Discovering subdomains...[/bold blue]") @@ -84,125 +87,148 @@ async def discover_subdomains(self, domain: str) -> List[str]: cmd = ["subfinder", "-d", domain, "-sources", ",".join(sources), "-silent"] output = await self.run_command(cmd) if not output: - self.logger.warning(f"No output from subfinder for {domain}. Check if subfinder is installed and configured correctly.") - self.console.print(f"[yellow]Warning: No output from subfinder for {domain}. Check if subfinder is installed and configured correctly.[/yellow]") + self.logger.warning( + f"No output from subfinder for {domain}. Check if subfinder is installed and configured correctly." + ) + self.console.print( + f"[yellow]Warning: No output from subfinder for {domain}. Check if subfinder is installed and configured correctly.[/yellow]" + ) subdomains = [line.strip() for line in output.splitlines() if line.strip()] self.console.print(f"[green]Found {len(subdomains)} subdomains[/green]") if len(subdomains) == 0: - self.logger.warning(f"No subdomains found for {domain}. Check subfinder installation, network, and API keys.") - self.console.print(f"[yellow]Warning: No subdomains found for {domain}. Check subfinder installation, network, and API keys.[/yellow]") + self.logger.warning( + f"No subdomains found for {domain}. Check subfinder installation, network, and API keys." + ) + self.console.print( + f"[yellow]Warning: No subdomains found for {domain}. Check subfinder installation, network, and API keys.[/yellow]" + ) return subdomains - + async def probe_web_services(self, subdomains: List[str]) -> List[Dict]: """Probe discovered subdomains for web services""" self.console.print("[bold blue]Probing web services...[/bold blue]") - + # Save subdomains to temporary file temp_file = self.base_dir / "temp_subdomains.txt" - async with aiofiles.open(temp_file, 'w') as f: - await f.write('\n'.join(subdomains)) - + async with aiofiles.open(temp_file, "w") as f: + await f.write("\n".join(subdomains)) + # Run httpx cmd = [ "httpx", - "-l", str(temp_file), + "-l", + str(temp_file), "-silent", "-json", "-status-code", "-title", "-tech-detect", - "-t", str(self.config["httpx"]["threads"]), - "-timeout", str(self.config["httpx"]["timeout"]) + "-t", + str(self.config["httpx"]["threads"]), + "-timeout", + str(self.config["httpx"]["timeout"]), ] - + output = await self.run_command(cmd) results = [] - + for line in output.splitlines(): try: result = json.loads(line) - results.append({ - "url": result.get("url", ""), - "status_code": result.get("status-code", 0), - "title": result.get("title", ""), - "technologies": result.get("technologies", []) - }) + results.append( + { + "url": result.get("url", ""), + "status_code": result.get("status-code", 0), + "title": result.get("title", ""), + "technologies": result.get("technologies", []), + } + ) except: continue - + # Cleanup temp_file.unlink(missing_ok=True) - + self.console.print(f"[green]Found {len(results)} web services[/green]") return results - + async def scan_vulnerabilities(self, web_services: List[Dict]) -> List[Dict]: """Scan web services for vulnerabilities using nuclei""" self.console.print("[bold blue]Scanning for vulnerabilities...[/bold blue]") - + # Save URLs to temporary file temp_file = self.base_dir / "temp_urls.txt" - async with aiofiles.open(temp_file, 'w') as f: - await f.write('\n'.join(service["url"] for service in web_services)) - + async with aiofiles.open(temp_file, "w") as f: + await f.write("\n".join(service["url"] for service in web_services)) + # Run nuclei cmd = [ "nuclei", - "-l", str(temp_file), + "-l", + str(temp_file), "-json", - "-severity", ",".join(self.config["nuclei"]["severity"]), - "-silent" + "-severity", + ",".join(self.config["nuclei"]["severity"]), + "-silent", ] - + output = await self.run_command(cmd) vulnerabilities = [] - + for line in output.splitlines(): try: vuln = json.loads(line) - vulnerabilities.append({ - "url": vuln.get("url", ""), - "type": vuln.get("type", ""), - "severity": vuln.get("severity", ""), - "description": vuln.get("description", ""), - "template": vuln.get("template", "") - }) + vulnerabilities.append( + { + "url": vuln.get("url", ""), + "type": vuln.get("type", ""), + "severity": vuln.get("severity", ""), + "description": vuln.get("description", ""), + "template": vuln.get("template", ""), + } + ) except: continue - + # Cleanup temp_file.unlink(missing_ok=True) - - self.console.print(f"[green]Found {len(vulnerabilities)} potential vulnerabilities[/green]") + + self.console.print( + f"[green]Found {len(vulnerabilities)} potential vulnerabilities[/green]" + ) return vulnerabilities - - async def run_recon(self, target: str, output_dir: Optional[Path] = None) -> Dict[str, Any]: + + async def run_recon( + self, target: str, output_dir: Optional[Path] = None + ) -> Dict[str, Any]: """Run complete reconnaissance on target concurrently""" if not output_dir: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") output_dir = self.base_dir / "results" / f"{target}_{timestamp}" output_dir.mkdir(parents=True, exist_ok=True) - + # Define all reconnaissance tasks recon_tasks = [ self.discover_subdomains(target), self.run_port_scan(target), # Assuming a new async port scan method # Add other concurrent tasks here ] - + # Run tasks concurrently results = await asyncio.gather(*recon_tasks, return_exceptions=True) - + # Process results... subdomains = results[0] if not isinstance(results[0], Exception) else [] port_scan_results = results[1] if not isinstance(results[1], Exception) else {} web_services = await self.probe_web_services(subdomains) vulnerabilities = await self.scan_vulnerabilities(web_services) - + # Generate AI analysis - ai_analysis = await self._analyze_results(target, subdomains, web_services, vulnerabilities) - + ai_analysis = await self._analyze_results( + target, subdomains, web_services, vulnerabilities + ) + # Prepare results results = { "target": target, @@ -210,19 +236,24 @@ async def run_recon(self, target: str, output_dir: Optional[Path] = None) -> Dic "subdomains": subdomains, "web_services": web_services, "vulnerabilities": vulnerabilities, - "ai_analysis": ai_analysis + "ai_analysis": ai_analysis, } - + # Save results await self._save_results(results, output_dir) - + return results - - async def _analyze_results(self, target: str, subdomains: List[str], - web_services: List[Dict], vulnerabilities: List[Dict]) -> Dict: + + async def _analyze_results( + self, + target: str, + subdomains: List[str], + web_services: List[Dict], + vulnerabilities: List[Dict], + ) -> Dict: """Analyze results using AI""" self.console.print("[bold blue]Analyzing results with AI...[/bold blue]") - + # Prepare data for AI analysis analysis_data = { "target": target, @@ -230,16 +261,18 @@ async def _analyze_results(self, target: str, subdomains: List[str], "web_service_count": len(web_services), "vulnerability_count": len(vulnerabilities), "critical_findings": [], - "recommendations": [] + "recommendations": [], } - + # Analyze vulnerabilities if vulnerabilities: - vuln_summary = "\n".join([ - f"- {v['type']} ({v['severity']}): {v['description']}" - for v in vulnerabilities - ]) - + vuln_summary = "\n".join( + [ + f"- {v['type']} ({v['severity']}): {v['description']}" + for v in vulnerabilities + ] + ) + prompt = f""" Analyze these security findings for {target}: @@ -250,7 +283,7 @@ async def _analyze_results(self, target: str, subdomains: List[str], 2. Recommended next steps for exploitation 3. Potential attack chains """ - + ai_response = self.ai_analyzer.ollama.generate(prompt) if ai_response: try: @@ -259,20 +292,21 @@ async def _analyze_results(self, target: str, subdomains: List[str], except: # If JSON parsing fails, use raw response analysis_data["raw_analysis"] = ai_response - + return analysis_data - + async def _save_results(self, results: Dict, output_dir: Path): """Save results in multiple formats""" # Save JSON json_path = output_dir / "results.json" - async with aiofiles.open(json_path, 'w') as f: + async with aiofiles.open(json_path, "w") as f: await f.write(json.dumps(results, indent=2)) - + # Generate Markdown report md_path = output_dir / "report.md" - async with aiofiles.open(md_path, 'w') as f: - await f.write(f"""# VulnForge Reconnaissance Report + async with aiofiles.open(md_path, "w") as f: + await f.write( + f"""# VulnForge Reconnaissance Report ## Target: {results['target']} ## Scan Time: {results['timestamp']} @@ -293,6 +327,7 @@ async def _save_results(self, results: Dict, output_dir: Path): ### AI Analysis {json.dumps(results['ai_analysis'], indent=2)} -""") - - self.console.print(f"[green]Results saved to: {output_dir}[/green]") \ No newline at end of file +""" + ) + + self.console.print(f"[green]Results saved to: {output_dir}[/green]") diff --git a/screen_control/launcher.py b/screen_control/launcher.py index 03f3d86..b1e6c1e 100644 --- a/screen_control/launcher.py +++ b/screen_control/launcher.py @@ -8,15 +8,17 @@ # Ensure you have build-essential, g++, rustc, cargo, and nasm installed. # sudo apt-get install build-essential g++ rustc cargo nasm xdotool + def build_modules(): """Compiles all native C++, Rust, and Assembly modules.""" print("--- Building native modules ---") - + # Build C++ print("Building C++ module...") subprocess.run( "g++ -shared -fPIC -o cpp/screen.so cpp/screen.cpp -lX11 -lXext", - shell=True, check=True + shell=True, + check=True, ) # Build Rust @@ -25,51 +27,48 @@ def build_modules(): # Build Assembly print("Building Assembly module...") - subprocess.run("nasm -f elf64 assembly/hook.asm -o assembly/hook.o", shell=True, check=True) - subprocess.run("ld -shared -o assembly/hook.so assembly/hook.o", shell=True, check=True) - + subprocess.run( + "nasm -f elf64 assembly/hook.asm -o assembly/hook.o", shell=True, check=True + ) + subprocess.run( + "ld -shared -o assembly/hook.so assembly/hook.o", shell=True, check=True + ) + print("--- Build complete ---") + def run_demo(controller): """Runs a demo sequence to showcase the system.""" print("\n--- Running Demo Sequence ---") # 1. Focus on the Terminal window - focus_cmd = { - "action": "window_focus", - "params": {"window_title": "Terminal"} - } + focus_cmd = {"action": "window_focus", "params": {"window_title": "Terminal"}} controller.execute_command(focus_cmd) - + # 2. Type a command type_cmd = { "action": "key_type", - "params": {"text": "echo 'Hello from VulnForge Screen AI!' && ls -l"} + "params": {"text": "echo 'Hello from VulnForge Screen AI!' && ls -l"}, } controller.execute_command(type_cmd) # 3. Simulate pressing Enter (using a separate key_type for special keys) - enter_cmd = { - "action": "key_type", - "params": {"text": "\n"} - } + enter_cmd = {"action": "key_type", "params": {"text": "\n"}} controller.execute_command(enter_cmd) # 4. Move mouse to corner - move_cmd = { - "action": "mouse_move", - "params": {"x": 10, "y": 10} - } + move_cmd = {"action": "mouse_move", "params": {"x": 10, "y": 10}} controller.execute_command(move_cmd) print("\n--- Demo Complete ---") + if __name__ == "__main__": # Ensure current directory is the script's directory os.chdir(os.path.dirname(os.path.abspath(__file__))) - + build_modules() - + screen_controller = ScreenControl(os.getcwd()) - - run_demo(screen_controller) \ No newline at end of file + + run_demo(screen_controller) diff --git a/scripts/analyze_results.py b/scripts/analyze_results.py index ed269a4..7b7d059 100755 --- a/scripts/analyze_results.py +++ b/scripts/analyze_results.py @@ -15,11 +15,11 @@ # Configure logging logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) + class ResultAnalyzer: def __init__(self, base_dir: Path): self.base_dir = base_dir @@ -27,73 +27,78 @@ def __init__(self, base_dir: Path): self.llm = LLMEngine() self.context = ContextBuilder(base_dir) self.notifier = Notifier(self.config) - + async def analyze_scan_results(self, results_file: Path): """Analyze scan results using AI""" try: # Load scan results with open(results_file) as f: results = json.load(f) - + # Add scan results to context self.context.add_scan_result("recon", results) - + # Build analysis prompt prompt = self.context.build_prompt( "recon", { "target_info": json.dumps(results["target"], indent=2), - "context": self.context._format_context() - } + "context": self.context._format_context(), + }, ) - + # Get AI analysis analysis = self.llm.query( prompt=prompt, - system_prompt="You are a cybersecurity expert analyzing reconnaissance data." + system_prompt="You are a cybersecurity expert analyzing reconnaissance data.", ) - + if not analysis: logger.error("Failed to get AI analysis") return - + # Parse analysis try: analysis_data = json.loads(analysis) - + # Check for critical findings if analysis_data.get("critical_findings"): await self.notifier.notify( "Critical findings detected in scan", "critical", - data=analysis_data["critical_findings"] + data=analysis_data["critical_findings"], ) - + # Save analysis - output_path = self.base_dir / "data" / "analysis" / f"analysis_{Path(results_file).stem}.json" + output_path = ( + self.base_dir + / "data" + / "analysis" + / f"analysis_{Path(results_file).stem}.json" + ) output_path.parent.mkdir(parents=True, exist_ok=True) - - with open(output_path, 'w') as f: + + with open(output_path, "w") as f: json.dump(analysis_data, f, indent=2) - + logger.info(f"Analysis saved to {output_path}") - + # Generate exploit if critical vulnerability found if analysis_data.get("critical_findings"): await self.generate_exploits(analysis_data) - + except json.JSONDecodeError as e: logger.error(f"Failed to parse AI analysis: {e}") - + except Exception as e: logger.error(f"Error analyzing results: {e}") - + async def generate_exploits(self, analysis_data: Dict): """Generate exploits for critical findings""" from modules.exploit_generator.exploit_generator import ExploitGenerator - + generator = ExploitGenerator(self.base_dir, self.llm) - + for finding in analysis_data.get("critical_findings", []): try: # Generate exploit @@ -102,24 +107,26 @@ async def generate_exploits(self, analysis_data: Dict): "cve_id": finding.get("cve_id", "Unknown"), "description": finding.get("description", ""), "affected_software": finding.get("affected_component", ""), - "cvss_score": finding.get("severity", "Unknown") + "cvss_score": finding.get("severity", "Unknown"), }, recon_data={ "ip": finding.get("location", ""), "port": finding.get("port", ""), "service": finding.get("service", ""), - "version": finding.get("version", "") - } + "version": finding.get("version", ""), + }, ) - + if exploit_data and "error" not in exploit_data: # Validate exploit validation = generator.validate_exploit(exploit_data) - + if validation["syntax_valid"] and validation["has_error_handling"]: # Generate Metasploit module - metasploit_module = generator.generate_metasploit_module(exploit_data) - + metasploit_module = generator.generate_metasploit_module( + exploit_data + ) + # Notify about exploit generation await self.notifier.notify( f"Exploit generated for {finding.get('type', 'Unknown')} vulnerability", @@ -127,29 +134,36 @@ async def generate_exploits(self, analysis_data: Dict): data={ "finding": finding, "exploit_path": exploit_data.get("file_path"), - "metasploit_module": str(metasploit_module) if metasploit_module else None, - "validation": validation - } + "metasploit_module": ( + str(metasploit_module) + if metasploit_module + else None + ), + "validation": validation, + }, ) - + except Exception as e: logger.error(f"Error generating exploit for finding: {e}") - + + async def main(): base_dir = Path(__file__).parent.parent analyzer = ResultAnalyzer(base_dir) - + # Find latest scan results results_dir = base_dir / "data" / "scan_results" if not results_dir.exists(): logger.error("No scan results found") return - - latest_result = max(results_dir.glob("recon_*.json"), key=lambda p: p.stat().st_mtime) - + + latest_result = max( + results_dir.glob("recon_*.json"), key=lambda p: p.stat().st_mtime + ) + # Start notifier await analyzer.notifier.start() - + try: # Analyze results await analyzer.analyze_scan_results(latest_result) @@ -157,6 +171,8 @@ async def main(): # Stop notifier await analyzer.notifier.stop() + if __name__ == "__main__": import asyncio - asyncio.run(main()) \ No newline at end of file + + asyncio.run(main()) diff --git a/scripts/recon_scan.py b/scripts/recon_scan.py index 9f0a472..d69c9a8 100755 --- a/scripts/recon_scan.py +++ b/scripts/recon_scan.py @@ -14,172 +14,181 @@ # Configure logging logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) + class ReconScanner: def __init__(self, config_path: Path): self.config = self._load_config(config_path) self.results = { "scan_time": datetime.now().isoformat(), "target": self.config["target"], - "findings": [] + "findings": [], } - + def _load_config(self, config_path: Path) -> Dict: """Load scan configuration""" with open(config_path) as f: return json.load(f) - + async def run_nmap_scan(self): """Run Nmap scan""" try: - cmd = ["nmap"] + self.config["tools"]["nmap"]["flags"] + [self.config["target"]["domain"]] + cmd = ( + ["nmap"] + + self.config["tools"]["nmap"]["flags"] + + [self.config["target"]["domain"]] + ) process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: - self.results["findings"].append({ - "tool": "nmap", - "output": stdout.decode(), - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "nmap", + "output": stdout.decode(), + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error(f"Nmap scan failed: {stderr.decode()}") - + except Exception as e: logger.error(f"Error running Nmap scan: {e}") - + async def run_subfinder(self): """Run Subfinder for subdomain enumeration""" try: cmd = ["subfinder", "-d", self.config["target"]["domain"], "-silent"] process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: subdomains = stdout.decode().splitlines() - self.results["findings"].append({ - "tool": "subfinder", - "subdomains": subdomains, - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "subfinder", + "subdomains": subdomains, + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error(f"Subfinder failed: {stderr.decode()}") - + except Exception as e: logger.error(f"Error running Subfinder: {e}") - + async def run_httpx(self, urls: List[str]): """Run httpx for HTTP probing""" try: # Write URLs to temporary file temp_file = Path("temp_urls.txt") temp_file.write_text("\n".join(urls)) - + cmd = ["httpx", "-l", str(temp_file), "-silent", "-status-code", "-title"] process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: - self.results["findings"].append({ - "tool": "httpx", - "output": stdout.decode(), - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "httpx", + "output": stdout.decode(), + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error(f"httpx failed: {stderr.decode()}") - + # Clean up temp file temp_file.unlink() - + except Exception as e: logger.error(f"Error running httpx: {e}") - + async def run_nuclei(self, urls: List[str]): """Run Nuclei for vulnerability scanning""" try: # Write URLs to temporary file temp_file = Path("temp_urls.txt") temp_file.write_text("\n".join(urls)) - + cmd = ["nuclei", "-l", str(temp_file), "-severity", "critical,high,medium"] process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE + *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() - + if process.returncode == 0: - self.results["findings"].append({ - "tool": "nuclei", - "output": stdout.decode(), - "timestamp": datetime.now().isoformat() - }) + self.results["findings"].append( + { + "tool": "nuclei", + "output": stdout.decode(), + "timestamp": datetime.now().isoformat(), + } + ) else: logger.error(f"Nuclei failed: {stderr.decode()}") - + # Clean up temp file temp_file.unlink() - + except Exception as e: logger.error(f"Error running Nuclei: {e}") - + def save_results(self, output_path: Path): """Save scan results to file""" - with open(output_path, 'w') as f: + with open(output_path, "w") as f: json.dump(self.results, f, indent=2) - + + async def main(): # Initialize scanner scanner = ReconScanner(Path("configs/scan_config.json")) - + # Run scans logger.info("Starting reconnaissance scan...") - + # Run Nmap scan logger.info("Running Nmap scan...") await scanner.run_nmap_scan() - + # Run Subfinder logger.info("Running Subfinder...") await scanner.run_subfinder() - + # Get all discovered URLs urls = [scanner.config["target"]["url"]] for finding in scanner.results["findings"]: if finding["tool"] == "subfinder": urls.extend([f"https://{subdomain}" for subdomain in finding["subdomains"]]) - + # Run httpx logger.info("Running httpx...") await scanner.run_httpx(urls) - + # Run Nuclei logger.info("Running Nuclei...") await scanner.run_nuclei(urls) - + # Save results - output_path = Path("data/scan_results") / f"recon_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" + output_path = ( + Path("data/scan_results") + / f"recon_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json" + ) output_path.parent.mkdir(parents=True, exist_ok=True) scanner.save_results(output_path) - + logger.info(f"Scan completed. Results saved to {output_path}") + if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(main()) diff --git a/setup.py b/setup.py index a85c0a5..d7a4cd4 100644 --- a/setup.py +++ b/setup.py @@ -63,4 +63,4 @@ "Topic :: Security", ], python_requires=">=3.8", -) \ No newline at end of file +) diff --git a/tests/test_ai_features.py b/tests/test_ai_features.py index 45e6e69..560f3cf 100644 --- a/tests/test_ai_features.py +++ b/tests/test_ai_features.py @@ -8,6 +8,7 @@ from ai_controller import AIController from utils.report_generator import ReportGenerator + class TestAIFeatures: """Test suite for AI features.""" @@ -16,24 +17,25 @@ def setup(self, tmp_path): """Setup test environment.""" self.test_dir = tmp_path / "test_session" self.test_dir.mkdir(parents=True) - + # Create test config self.config_path = self.test_dir / "test_config.json" with open(self.config_path, "w") as f: - json.dump({ - "ai": { - "model": "deepseek-coder", - "temperature": 0.7, - "max_tokens": 1000 + json.dump( + { + "ai": { + "model": "deepseek-coder", + "temperature": 0.7, + "max_tokens": 1000, + }, + "notifications": {"enabled": False}, }, - "notifications": { - "enabled": False - } - }, f) - + f, + ) + # Initialize AI controller self.ai_controller = AIController(str(self.test_dir), str(self.config_path)) - + # Initialize report generator self.report_generator = ReportGenerator(str(self.test_dir)) @@ -43,20 +45,20 @@ async def test_ai_query_processing(self): # Test basic query query = "What should I do if port 8080 is open?" response = await self.ai_controller.process_query(query) - + assert "answer" in response assert "logs" in response assert "prompt" in response assert isinstance(response["answer"], str) - + # Test query with context context = { "open_ports": [8080], "services": {"8080": "http-alt"}, - "previous_findings": ["Potential web server detected"] + "previous_findings": ["Potential web server detected"], } response = await self.ai_controller.process_query(query, context) - + assert "answer" in response assert "logs" in response assert "prompt" in response @@ -70,7 +72,7 @@ def test_report_generation(self): "domain": "example.com", "ip": "93.184.216.34", "scan_time": datetime.now().isoformat(), - "duration": "00:05:23" + "duration": "00:05:23", }, "modules": [ { @@ -78,7 +80,7 @@ def test_report_generation(self): "status": "success", "start_time": "2024-03-20T10:00:00", "end_time": "2024-03-20T10:05:00", - "findings": ["Open port 80", "Open port 443"] + "findings": ["Open port 80", "Open port 443"], } ], "tools": { @@ -86,7 +88,7 @@ def test_report_generation(self): "status": "success", "purpose": "Port scanning", "config": "-sS -sV", - "installed": True + "installed": True, } }, "vulnerabilities": [ @@ -99,9 +101,9 @@ def test_report_generation(self): "references": [ { "title": "CVE-2024-1234", - "url": "https://nvd.nist.gov/vuln/detail/CVE-2024-1234" + "url": "https://nvd.nist.gov/vuln/detail/CVE-2024-1234", } - ] + ], } ], "exploits": [ @@ -109,54 +111,54 @@ def test_report_generation(self): "name": "Test Exploit", "status": "success", "command": "test_command", - "output": "Test output" + "output": "Test output", } ], "defensive_measures": [ { "name": "WAF Detection", "description": "Web Application Firewall detected", - "details": "Cloudflare WAF" + "details": "Cloudflare WAF", } ], "recommendations": [ { "title": "Test Recommendation", "description": "Test recommendation description", - "steps": ["Step 1", "Step 2"] + "steps": ["Step 1", "Step 2"], } ], "errors": [ { "timestamp": datetime.now().isoformat(), "message": "Test error", - "context": "Test error context" + "context": "Test error context", } - ] + ], } - + # Generate reports report_paths = self.report_generator.generate_reports(context) - + # Verify report files exist assert "html" in report_paths assert "markdown" in report_paths assert "json" in report_paths - + # Verify HTML report html_path = Path(report_paths["html"]) assert html_path.exists() html_content = html_path.read_text() assert "VulnForge Report" in html_content assert "example.com" in html_content - + # Verify Markdown report md_path = Path(report_paths["markdown"]) assert md_path.exists() md_content = md_path.read_text() assert "# VulnForge Report" in md_content assert "example.com" in md_content - + # Verify JSON report json_path = Path(report_paths["json"]) assert json_path.exists() @@ -168,22 +170,20 @@ def test_session_data_management(self): # Add test data self.ai_controller.session_data["target_domain"] = "example.com" self.ai_controller.session_data["target_ip"] = "93.184.216.34" - self.ai_controller.session_data["modules"].append({ - "name": "test_module", - "status": "success" - }) - self.ai_controller.session_data["findings"].append({ - "type": "vulnerability", - "severity": "high" - }) - + self.ai_controller.session_data["modules"].append( + {"name": "test_module", "status": "success"} + ) + self.ai_controller.session_data["findings"].append( + {"type": "vulnerability", "severity": "high"} + ) + # Generate report report_paths = self.ai_controller._generate_report() - + # Verify report contains session data json_path = Path(report_paths["json"]) json_content = json.loads(json_path.read_text()) - + assert json_content["target"]["domain"] == "example.com" assert json_content["target"]["ip"] == "93.184.216.34" assert len(json_content["modules"]) > 0 @@ -194,15 +194,16 @@ def test_error_handling(self): # Test invalid query with pytest.raises(Exception): asyncio.run(self.ai_controller.process_query(None)) - + # Test invalid context with pytest.raises(Exception): self.report_generator.generate_reports(None) - + # Test invalid session data self.ai_controller.session_data = None with pytest.raises(Exception): self.ai_controller._generate_report() + if __name__ == "__main__": - pytest.main([__file__, "-v"]) \ No newline at end of file + pytest.main([__file__, "-v"]) diff --git a/tests/test_notifier.py b/tests/test_notifier.py index 23d3c27..503b6a2 100755 --- a/tests/test_notifier.py +++ b/tests/test_notifier.py @@ -10,12 +10,13 @@ import pytest_asyncio from utils.notifier import Notifier + @pytest_asyncio.fixture async def notifier(): """Create a test notifier instance""" test_dir = Path("test_data") test_dir.mkdir(exist_ok=True) - + config = { "notifications": { "enabled": True, @@ -24,32 +25,30 @@ async def notifier(): "smtp_server": "smtp.test.com", "smtp_port": 587, "username": "test@test.com", - "password": "test123" + "password": "test123", }, "discord": { "enabled": True, - "webhook_url": "https://discord.com/api/webhooks/test" + "webhook_url": "https://discord.com/api/webhooks/test", }, - "webhook": { - "enabled": True, - "url": "https://webhook.test.com" - } + "webhook": {"enabled": True, "url": "https://webhook.test.com"}, } } - + config_path = test_dir / "test_config.json" with open(config_path, "w") as f: json.dump(config, f) - + notifier = Notifier(str(test_dir), str(config_path)) yield notifier - + # Cleanup if test_dir.exists(): for file in test_dir.glob("*"): file.unlink() test_dir.rmdir() + @pytest.mark.asyncio async def test_notify_basic(notifier): """Test basic notification""" @@ -57,52 +56,47 @@ async def test_notify_basic(notifier): await notifier.notify("Test message", "info") mock_process.assert_called_once_with("Test message", "info", None, []) + @pytest.mark.asyncio async def test_notify_with_data(notifier): """Test notification with additional data""" test_data = { "vulnerability": "SQL Injection", "severity": "high", - "affected_url": "https://example.com/login" + "affected_url": "https://example.com/login", } - + with patch("utils.notifier.Notifier._process_notification") as mock_process: await notifier.notify("Vulnerability found", "high", data=test_data) - mock_process.assert_called_once_with("Vulnerability found", "high", test_data, []) + mock_process.assert_called_once_with( + "Vulnerability found", "high", test_data, [] + ) + @pytest.mark.asyncio async def test_notify_specific_channels(notifier): """Test notification to specific channels""" with patch("utils.notifier.Notifier._process_notification") as mock_process: - await notifier.notify( - "Test message", - "info", - channels=["email", "discord"] - ) + await notifier.notify("Test message", "info", channels=["email", "discord"]) mock_process.assert_called_once_with( - "Test message", - "info", - None, - ["email", "discord"] + "Test message", "info", None, ["email", "discord"] ) + @pytest.mark.asyncio async def test_email_notification(notifier): """Test email notification sending""" with patch("smtplib.SMTP") as mock_smtp: mock_smtp.return_value.__enter__.return_value = mock_smtp.return_value - - await notifier.notify( - "Test email", - "critical", - channels=["email"] - ) - + + await notifier.notify("Test email", "critical", channels=["email"]) + mock_smtp.assert_called_once_with("smtp.test.com", 587) mock_smtp.return_value.starttls.assert_called_once() mock_smtp.return_value.login.assert_called_once_with("test@test.com", "test123") mock_smtp.return_value.send_message.assert_called_once() + @pytest.mark.asyncio async def test_discord_notification(notifier): """Test Discord notification sending""" @@ -113,20 +107,21 @@ async def test_discord_notification(notifier): mock_cm = AsyncMock() mock_cm.__aenter__.return_value = mock_session mock_cm.__aexit__.return_value = None - + with patch("aiohttp.ClientSession", return_value=mock_cm): await notifier.notify( "Test Discord", "high", data={"finding": "XSS vulnerability"}, - channels=["discord"] + channels=["discord"], ) - + mock_session.post.assert_called_once() call_args = mock_session.post.call_args assert call_args[0][0] == "https://discord.com/api/webhooks/test" assert "embeds" in call_args[1]["json"] + @pytest.mark.asyncio async def test_webhook_notification(notifier): """Test generic webhook notification""" @@ -137,21 +132,18 @@ async def test_webhook_notification(notifier): mock_cm = AsyncMock() mock_cm.__aenter__.return_value = mock_session mock_cm.__aexit__.return_value = None - + with patch("aiohttp.ClientSession", return_value=mock_cm): test_data = { "scan_id": "123", "target": "example.com", - "findings": ["vuln1", "vuln2"] + "findings": ["vuln1", "vuln2"], } - + await notifier.notify( - "Test webhook", - "medium", - data=test_data, - channels=["webhook"] + "Test webhook", "medium", data=test_data, channels=["webhook"] ) - + mock_session.post.assert_called_once() call_args = mock_session.post.call_args assert call_args[0][0] == "https://webhook.test.com" @@ -159,6 +151,7 @@ async def test_webhook_notification(notifier): assert call_args[1]["json"]["severity"] == "medium" assert call_args[1]["json"]["data"] == test_data + @pytest.mark.asyncio async def test_severity_colors(notifier): """Test severity color mapping""" @@ -167,9 +160,9 @@ async def test_severity_colors(notifier): "high": 0xFFA500, "medium": 0xFFFF00, "low": 0x00FF00, - "info": 0x0000FF + "info": 0x0000FF, } - + for severity, expected_color in colors.items(): color = notifier._get_severity_color(severity) - assert color == expected_color, f"Wrong color for {severity}" \ No newline at end of file + assert color == expected_color, f"Wrong color for {severity}" diff --git a/utils/__init__.py b/utils/__init__.py index e39b5e5..fd33727 100755 --- a/utils/__init__.py +++ b/utils/__init__.py @@ -8,16 +8,16 @@ print_results_table, print_error, print_success, - print_warning + print_warning, ) from .logger import setup_logger __all__ = [ - 'print_banner', - 'create_progress', - 'print_results_table', - 'print_error', - 'print_success', - 'print_warning', - 'setup_logger' -] \ No newline at end of file + "print_banner", + "create_progress", + "print_results_table", + "print_error", + "print_success", + "print_warning", + "setup_logger", +] diff --git a/utils/cli_utils.py b/utils/cli_utils.py index 7cbce7b..9446cf7 100755 --- a/utils/cli_utils.py +++ b/utils/cli_utils.py @@ -13,6 +13,7 @@ console = Console() + def print_banner(version: str): """Print VulnForge banner""" banner = f""" @@ -22,69 +23,79 @@ def print_banner(version: str): ╚══════════════════════════════════════════════════════════════╝ """ console.print(Panel(banner, style="bold blue")) - + + def print_results_table(results: dict): """Print results in a formatted table""" for category, items in results.items(): if not items: continue - + table = Table(title=f"{category} Results") table.add_column("Item", style="cyan") table.add_column("Details", style="green") - + for item in items: if isinstance(item, dict): details = ", ".join(f"{k}: {v}" for k, v in item.items()) else: details = str(item) table.add_row(str(item), details) - + console.print(table) - + + def print_json_results(results: Dict[str, Any]): """Print results in JSON format""" console.print(json.dumps(results, indent=2)) - + + def print_error(message: str): """Print error message""" console.print(f"[bold red]Error:[/bold red] {message}") - + + def print_warning(message: str): """Print warning message""" console.print(f"[bold yellow]Warning:[/bold yellow] {message}") - + + def print_success(message: str): """Print success message""" console.print(f"[bold green]Success:[/bold green] {message}") - + + def print_info(message: str): """Print info message""" console.print(f"[bold blue]Info:[/bold blue] {message}") - + + def print_progress(message: str): """Print progress message""" console.print(f"[bold cyan]Progress:[/bold cyan] {message}") - + + def print_debug(message: str): """Print debug message""" console.print(f"[bold magenta]Debug:[/bold magenta] {message}") - + + def print_stealth_stats(stats: Dict[str, Any]): """Print stealth statistics""" table = Table(title="\nStealth Statistics") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") - + for key, value in stats.items(): - table.add_row(key.replace('_', ' ').title(), str(value)) - + table.add_row(key.replace("_", " ").title(), str(value)) + console.print(table) + def create_progress(description: str): """Create a progress bar""" return Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), - console=console - ) \ No newline at end of file + console=console, + ) diff --git a/utils/config_manager.py b/utils/config_manager.py index b4ff285..5a181f8 100755 --- a/utils/config_manager.py +++ b/utils/config_manager.py @@ -10,6 +10,7 @@ from typing import Dict, Any, Optional import os + class ConfigManager: def __init__(self, config_path: str): self.logger = logging.getLogger(__name__) @@ -17,12 +18,12 @@ def __init__(self, config_path: str): self.config_dir = self.config_path.parent / "configs" self.config_dir.mkdir(parents=True, exist_ok=True) self.config = self._load_config() - + def _load_config(self) -> Dict[str, Any]: """Load configuration from file""" try: if self.config_path.exists(): - with open(self.config_path, 'r') as f: + with open(self.config_path, "r") as f: return json.load(f) else: self.logger.warning(f"Config file not found: {self.config_path}") @@ -30,62 +31,62 @@ def _load_config(self) -> Dict[str, Any]: except Exception as e: self.logger.error(f"Error loading config: {e}") return {} - + def get(self, key: str, default: Any = None) -> Any: """Get configuration value""" return self.config.get(key, default) - + def set(self, key: str, value: Any) -> None: """Set configuration value""" self.config[key] = value self._save_config() - + def _save_config(self) -> None: """Save configuration to file""" try: - with open(self.config_path, 'w') as f: + with open(self.config_path, "w") as f: json.dump(self.config, f, indent=2) except Exception as e: self.logger.error(f"Error saving config: {e}") - + def get_ai_config(self) -> Dict[str, Any]: """Get AI configuration""" return self.config.get("ai", {}) - + def get_notification_config(self) -> Dict[str, Any]: """Get notification configuration""" return self.config.get("notifications", {}) - + def get_tool_config(self, tool_name: str) -> Dict[str, Any]: """Get tool-specific configuration""" return self.config.get("tools", {}).get(tool_name, {}) - + def save_config(self): """Save current configuration to file""" try: - with open(self.config_path, 'w') as f: + with open(self.config_path, "w") as f: json.dump(self.config, f, indent=2) except Exception as e: self.logger.error(f"Error saving config: {e}") - + def update(self, updates: Dict[str, Any]): """Update multiple configuration values""" self.config = {**self.config, **updates} self._save_config() - + def reset(self): """Reset configuration to defaults""" self.config = self._load_config() self._save_config() - + def export(self, filepath: Path): """Export configuration to file""" try: - with open(filepath, 'w') as f: + with open(filepath, "w") as f: json.dump(self.config, f, indent=2) except Exception as e: self.logger.error(f"Error exporting config: {e}") - + def import_config(self, filepath: Path): """Import configuration from file""" try: @@ -95,23 +96,23 @@ def import_config(self, filepath: Path): self._save_config() except Exception as e: self.logger.error(f"Error importing config: {e}") - + def get_all(self) -> Dict[str, Any]: """Get complete configuration""" return self.config.copy() - + def validate(self) -> bool: """Validate configuration""" required_keys = [ "ai.preferred_model", "scanning.max_threads", "stealth.enabled", - "reporting.default_format" + "reporting.default_format", ] - + for key in required_keys: if self.get(key) is None: self.logger.error(f"Missing required config key: {key}") return False - - return True \ No newline at end of file + + return True diff --git a/utils/context_builder.py b/utils/context_builder.py index f4dcf4b..94610db 100755 --- a/utils/context_builder.py +++ b/utils/context_builder.py @@ -11,6 +11,7 @@ from datetime import datetime import re + class ContextBuilder: def __init__(self, base_dir: Path): self.base_dir = base_dir @@ -19,61 +20,65 @@ def __init__(self, base_dir: Path): self.max_history = 10 self.prompt_templates_dir = self.base_dir / "prompts" self.prompt_templates_dir.mkdir(parents=True, exist_ok=True) - - def add_tool_output(self, tool_name: str, output: str, metadata: Optional[Dict] = None): + + def add_tool_output( + self, tool_name: str, output: str, metadata: Optional[Dict] = None + ): """Add tool output to context history""" context = { "type": "tool_output", "tool": tool_name, "output": output, "timestamp": datetime.now().isoformat(), - "metadata": metadata or {} + "metadata": metadata or {}, } self._add_to_history(context) - - def add_ai_response(self, prompt: str, response: str, metadata: Optional[Dict] = None): + + def add_ai_response( + self, prompt: str, response: str, metadata: Optional[Dict] = None + ): """Add AI response to context history""" context = { "type": "ai_response", "prompt": prompt, "response": response, "timestamp": datetime.now().isoformat(), - "metadata": metadata or {} + "metadata": metadata or {}, } self._add_to_history(context) - + def add_scan_result(self, result_type: str, data: Dict): """Add scan result to context history""" context = { "type": "scan_result", "result_type": result_type, "data": data, - "timestamp": datetime.now().isoformat() + "timestamp": datetime.now().isoformat(), } self._add_to_history(context) - + def _add_to_history(self, context: Dict): """Add context to history with size limit""" self.context_history.append(context) if len(self.context_history) > self.max_history: self.context_history.pop(0) - + def build_prompt(self, template_name: str, variables: Dict) -> str: """Build prompt from template with context""" template = self._load_template(template_name) if not template: return "" - + # Add context to variables variables["context"] = self._format_context() - + # Replace variables in template prompt = template for key, value in variables.items(): prompt = prompt.replace(f"{{{{ {key} }}}}", str(value)) - + return prompt - + def _load_template(self, template_name: str) -> Optional[str]: """Load prompt template""" template_path = self.prompt_templates_dir / f"{template_name}.jinja" @@ -85,46 +90,48 @@ def _load_template(self, template_name: str) -> Optional[str]: except Exception as e: self.logger.error(f"Error loading template {template_name}: {e}") return None - + def _format_context(self) -> str: """Format context history for prompt""" context_str = [] - + for ctx in self.context_history: if ctx["type"] == "tool_output": context_str.append(f"Tool: {ctx['tool']}\nOutput: {ctx['output']}") elif ctx["type"] == "ai_response": context_str.append(f"Previous AI Response: {ctx['response']}") elif ctx["type"] == "scan_result": - context_str.append(f"Scan Result ({ctx['result_type']}): {json.dumps(ctx['data'])}") - + context_str.append( + f"Scan Result ({ctx['result_type']}): {json.dumps(ctx['data'])}" + ) + return "\n\n".join(context_str) - + def save_template(self, template_name: str, content: str): """Save new prompt template""" template_path = self.prompt_templates_dir / f"{template_name}.jinja" try: - with open(template_path, 'w') as f: + with open(template_path, "w") as f: f.write(content) except Exception as e: self.logger.error(f"Error saving template {template_name}: {e}") - + def get_available_templates(self) -> List[str]: """Get list of available templates""" return [f.stem for f in self.prompt_templates_dir.glob("*.jinja")] - + def clear_history(self): """Clear context history""" self.context_history.clear() - + def export_context(self, filepath: Path): """Export context history to file""" try: - with open(filepath, 'w') as f: + with open(filepath, "w") as f: json.dump(self.context_history, f, indent=2) except Exception as e: self.logger.error(f"Error exporting context: {e}") - + def import_context(self, filepath: Path): """Import context history from file""" try: @@ -132,29 +139,29 @@ def import_context(self, filepath: Path): self.context_history = json.load(f) except Exception as e: self.logger.error(f"Error importing context: {e}") - + def get_relevant_context(self, query: str, max_items: int = 5) -> List[Dict]: """Get most relevant context items for a query""" # Simple relevance scoring based on keyword matching scored_items = [] - keywords = set(re.findall(r'\w+', query.lower())) - + keywords = set(re.findall(r"\w+", query.lower())) + for item in self.context_history: score = 0 content = "" - + if item["type"] == "tool_output": content = f"{item['tool']} {item['output']}" elif item["type"] == "ai_response": content = f"{item['prompt']} {item['response']}" elif item["type"] == "scan_result": - content = json.dumps(item['data']) - - content_words = set(re.findall(r'\w+', content.lower())) + content = json.dumps(item["data"]) + + content_words = set(re.findall(r"\w+", content.lower())) score = len(keywords.intersection(content_words)) - + scored_items.append((score, item)) - + # Sort by score and return top items scored_items.sort(reverse=True) - return [item for _, item in scored_items[:max_items]] \ No newline at end of file + return [item for _, item in scored_items[:max_items]] diff --git a/utils/logger.py b/utils/logger.py index 79e4d4a..2394ec8 100755 --- a/utils/logger.py +++ b/utils/logger.py @@ -15,28 +15,29 @@ import os from datetime import datetime + def setup_logger(name: str, log_dir: str = None) -> logging.Logger: """Set up a logger with the given name.""" if log_dir is None: log_dir = os.path.join(os.path.expanduser("~"), ".vulnforge", "logs") - + os.makedirs(log_dir, exist_ok=True) - + logger = logging.getLogger(name) logger.setLevel(logging.INFO) - + # Create handlers log_file = os.path.join(log_dir, f"{name}.log") file_handler = logging.FileHandler(log_file) console_handler = logging.StreamHandler() - + # Create formatters and add it to handlers - log_format = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') + log_format = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") file_handler.setFormatter(log_format) console_handler.setFormatter(log_format) - + # Add handlers to the logger logger.addHandler(file_handler) logger.addHandler(console_handler) - - return logger \ No newline at end of file + + return logger diff --git a/utils/notifier.py b/utils/notifier.py index 828b6b7..b9ae52a 100755 --- a/utils/notifier.py +++ b/utils/notifier.py @@ -14,16 +14,17 @@ from typing import Dict, List, Optional, Any from datetime import datetime + class Notifier: """Handles notifications for various channels""" - + def __init__(self, base_dir: str, config_path: str): """Initialize notifier with configuration""" self.base_dir = Path(base_dir) self.config_path = Path(config_path) self.logger = logging.getLogger(__name__) self.config = self._load_config() - + def _load_config(self) -> Dict: """Load notification configuration""" try: @@ -32,46 +33,44 @@ def _load_config(self) -> Dict: except Exception as e: self.logger.error(f"Failed to load config: {e}") return {"notifications": {"enabled": False}} - + async def notify( self, message: str, severity: str = "info", data: Optional[Dict] = None, - channels: Optional[List[str]] = None + channels: Optional[List[str]] = None, ) -> None: """Send notification to specified channels""" if not self.config["notifications"]["enabled"]: return - + if channels is None: channels = [] - + await self._process_notification(message, severity, data, channels) - + async def _process_notification( - self, - message: str, - severity: str, - data: Optional[Dict], - channels: List[str] + self, message: str, severity: str, data: Optional[Dict], channels: List[str] ) -> None: """Process notification for each channel""" tasks = [] - + if "email" in channels and self.config["notifications"]["email"]["enabled"]: tasks.append(self._send_email(message, severity, data)) - + if "discord" in channels and self.config["notifications"]["discord"]["enabled"]: tasks.append(self._send_discord(message, severity, data)) - + if "webhook" in channels and self.config["notifications"]["webhook"]["enabled"]: tasks.append(self._send_webhook(message, severity, data)) - + if tasks: await asyncio.gather(*tasks, return_exceptions=True) - - async def _send_email(self, message: str, severity: str, data: Optional[Dict]) -> None: + + async def _send_email( + self, message: str, severity: str, data: Optional[Dict] + ) -> None: """Send email notification""" try: email_config = self.config["notifications"]["email"] @@ -79,77 +78,80 @@ async def _send_email(self, message: str, severity: str, data: Optional[Dict]) - msg["From"] = email_config["username"] msg["To"] = email_config["username"] msg["Subject"] = f"VulnForge Alert: {severity.upper()}" - + body = f"Message: {message}\nSeverity: {severity}\n" if data: body += f"\nAdditional Data:\n{json.dumps(data, indent=2)}" - + msg.attach(MIMEText(body, "plain")) - - with smtplib.SMTP(email_config["smtp_server"], email_config["smtp_port"]) as server: + + with smtplib.SMTP( + email_config["smtp_server"], email_config["smtp_port"] + ) as server: server.starttls() server.login(email_config["username"], email_config["password"]) server.send_message(msg) - + except Exception as e: self.logger.error(f"Failed to send email: {e}") - - async def _send_discord(self, message: str, severity: str, data: Optional[Dict]) -> None: + + async def _send_discord( + self, message: str, severity: str, data: Optional[Dict] + ) -> None: """Send Discord notification""" try: webhook_url = self.config["notifications"]["discord"]["webhook_url"] color = self._get_severity_color(severity) - + embed = { "title": f"VulnForge Alert: {severity.upper()}", "description": message, "color": color, - "timestamp": datetime.utcnow().isoformat() + "timestamp": datetime.utcnow().isoformat(), } - + if data: embed["fields"] = [ {"name": k, "value": str(v), "inline": True} for k, v in data.items() ] - + async with aiohttp.ClientSession() as session: - response = await session.post( - webhook_url, - json={"embeds": [embed]} - ) + response = await session.post(webhook_url, json={"embeds": [embed]}) await response.text() - + except Exception as e: self.logger.error(f"Failed to send Discord notification: {e}") - - async def _send_webhook(self, message: str, severity: str, data: Optional[Dict]) -> None: + + async def _send_webhook( + self, message: str, severity: str, data: Optional[Dict] + ) -> None: """Send generic webhook notification""" try: webhook_url = self.config["notifications"]["webhook"]["url"] payload = { "message": message, "severity": severity, - "timestamp": datetime.utcnow().isoformat() + "timestamp": datetime.utcnow().isoformat(), } - + if data: payload["data"] = data - + async with aiohttp.ClientSession() as session: response = await session.post(webhook_url, json=payload) await response.text() - + except Exception as e: self.logger.error(f"Failed to send webhook notification: {e}") - + def _get_severity_color(self, severity: str) -> int: """Get color code for severity level""" colors = { "critical": 0xFF0000, # Red - "high": 0xFFA500, # Orange - "medium": 0xFFFF00, # Yellow - "low": 0x00FF00, # Green - "info": 0x0000FF # Blue + "high": 0xFFA500, # Orange + "medium": 0xFFFF00, # Yellow + "low": 0x00FF00, # Green + "info": 0x0000FF, # Blue } - return colors.get(severity.lower(), 0x808080) # Default to gray \ No newline at end of file + return colors.get(severity.lower(), 0x808080) # Default to gray diff --git a/utils/report_generator.py b/utils/report_generator.py index 7c9982c..3bbbbb7 100644 --- a/utils/report_generator.py +++ b/utils/report_generator.py @@ -12,96 +12,101 @@ from typing import Dict, Any, Optional from jinja2 import Environment, FileSystemLoader + class ReportGenerator: """Handles generation of reports in multiple formats (HTML, Markdown, JSON).""" def __init__(self, session_dir: str): """Initialize the report generator. - + Args: session_dir: Directory to store reports """ self.session_dir = Path(session_dir) self.report_dir = self.session_dir / "reports" self.report_dir.mkdir(parents=True, exist_ok=True) - + # Setup Jinja2 environment template_dir = Path(__file__).parent / "templates" self.env = Environment(loader=FileSystemLoader(str(template_dir))) - + # Setup logging self.logger = logging.getLogger(__name__) self.logger.setLevel(logging.INFO) - + # Add file handler if not already added if not self.logger.handlers: log_file = self.session_dir / "logs" / "report_generator.log" log_file.parent.mkdir(parents=True, exist_ok=True) handler = logging.FileHandler(log_file) - formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) handler.setFormatter(formatter) self.logger.addHandler(handler) - def generate_reports(self, context: Dict[str, Any], output_format: str = "all") -> Dict[str, str]: + def generate_reports( + self, context: Dict[str, Any], output_format: str = "all" + ) -> Dict[str, str]: """Generate reports in specified format(s). - + Args: context: Dictionary containing report data output_format: Desired output format (markdown, json, html, or all) - + Returns: Dictionary mapping report types to their file paths """ self.logger.info(f"Generating reports in {output_format} format...") - + # Add timestamp if not present if "timestamp" not in context: context["timestamp"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S") - + # Add watermark metadata context["watermark"] = { "generated_by": "VulnForge", "author": "DemonKing369.0", - "github": "https://github.com/Arunking9" + "github": "https://github.com/Arunking9", } - + report_paths = {} - + try: if output_format in ["html", "all"]: html_path = self._generate_html_report(context) report_paths["html"] = str(html_path) - + if output_format in ["markdown", "all"]: md_path = self._generate_markdown_report(context) report_paths["markdown"] = str(md_path) - + if output_format in ["json", "all"]: json_path = self._generate_json_report(context) report_paths["json"] = str(json_path) - + self.logger.info("Reports generated successfully") return report_paths - + except Exception as e: self.logger.error(f"Error generating reports: {str(e)}") raise def _generate_html_report(self, context: Dict[str, Any]) -> Path: """Generate HTML report. - + Args: context: Dictionary containing report data - + Returns: Path to the generated HTML report """ self.logger.info("Generating HTML report...") - + try: template = self.env.get_template("report.html") html_content = template.render(**context) - + # Add watermark footer watermark_footer = f""" """ html_content = html_content.replace("", f"{watermark_footer}") - + output_path = self.report_dir / "report.html" with open(output_path, "w", encoding="utf-8") as f: f.write(html_content) - + self.logger.info(f"HTML report generated: {output_path}") return output_path - + except Exception as e: self.logger.error(f"Error generating HTML report: {str(e)}") raise def _generate_markdown_report(self, context: Dict[str, Any]) -> Path: """Generate Markdown report. - + Args: context: Dictionary containing report data - + Returns: Path to the generated Markdown report """ self.logger.info("Generating Markdown report...") - + try: template = self.env.get_template("report.md") md_content = template.render(**context) - + # Add watermark footer watermark_footer = f""" --- @@ -144,67 +149,67 @@ def _generate_markdown_report(self, context: Dict[str, Any]) -> Path: GitHub: https://github.com/Arunking9 """ md_content += watermark_footer - + output_path = self.report_dir / "report.md" with open(output_path, "w", encoding="utf-8") as f: f.write(md_content) - + self.logger.info(f"Markdown report generated: {output_path}") return output_path - + except Exception as e: self.logger.error(f"Error generating Markdown report: {str(e)}") raise def _generate_json_report(self, context: Dict[str, Any]) -> Path: """Generate JSON report. - + Args: context: Dictionary containing report data - + Returns: Path to the generated JSON report """ self.logger.info("Generating JSON report...") - + try: # Add watermark metadata context["watermark"] = { "generated_by": "VulnForge", "author": "DemonKing369.0", "github": "https://github.com/Arunking9", - "timestamp": datetime.now().isoformat() + "timestamp": datetime.now().isoformat(), } - + output_path = self.report_dir / "report.json" with open(output_path, "w", encoding="utf-8") as f: json.dump(context, f, indent=2, default=str) - + self.logger.info(f"JSON report generated: {output_path}") return output_path - + except Exception as e: self.logger.error(f"Error generating JSON report: {str(e)}") raise def get_report_paths(self) -> Dict[str, str]: """Get paths to all generated reports. - + Returns: Dictionary mapping report types to their file paths """ report_paths = {} - + html_path = self.report_dir / "report.html" if html_path.exists(): report_paths["html"] = str(html_path) - + md_path = self.report_dir / "report.md" if md_path.exists(): report_paths["markdown"] = str(md_path) - + json_path = self.report_dir / "report.json" if json_path.exists(): report_paths["json"] = str(json_path) - - return report_paths \ No newline at end of file + + return report_paths diff --git a/utils/session_manager.py b/utils/session_manager.py index 0c1e353..fae7437 100644 --- a/utils/session_manager.py +++ b/utils/session_manager.py @@ -10,12 +10,13 @@ from pathlib import Path from typing import Dict, Any, Optional + class SessionManager: """Manages session data and metadata.""" def __init__(self, base_dir: str): """Initialize session manager. - + Args: base_dir: Base directory for session data """ @@ -25,27 +26,27 @@ def __init__(self, base_dir: str): def create_session(self, target: str) -> str: """Create a new session directory. - + Args: target: Target identifier - + Returns: Path to session directory """ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") self.session_dir = self.base_dir / target / timestamp self.session_dir.mkdir(parents=True, exist_ok=True) - + # Create metadata file self.metadata_file = self.session_dir / "meta.txt" self._write_metadata() - + # Create subdirectories (self.session_dir / "logs").mkdir(exist_ok=True) (self.session_dir / "data").mkdir(exist_ok=True) (self.session_dir / "tools").mkdir(exist_ok=True) (self.session_dir / "exploits").mkdir(exist_ok=True) - + return str(self.session_dir) def _write_metadata(self): @@ -62,38 +63,38 @@ def _write_metadata(self): def save_data(self, data: Dict[str, Any], filename: str): """Save data to session directory. - + Args: data: Data to save filename: Output filename """ if not self.session_dir: raise RuntimeError("No active session") - + output_path = self.session_dir / "data" / filename with open(output_path, "w") as f: json.dump(data, f, indent=2, default=str) def load_data(self, filename: str) -> Dict[str, Any]: """Load data from session directory. - + Args: filename: Input filename - + Returns: Loaded data """ if not self.session_dir: raise RuntimeError("No active session") - + input_path = self.session_dir / "data" / filename with open(input_path, "r") as f: return json.load(f) def get_session_path(self) -> Optional[str]: """Get current session directory path. - + Returns: Path to current session directory or None if no active session """ - return str(self.session_dir) if self.session_dir else None \ No newline at end of file + return str(self.session_dir) if self.session_dir else None diff --git a/utils/stealth_utils.py b/utils/stealth_utils.py index 287a041..38a25b9 100755 --- a/utils/stealth_utils.py +++ b/utils/stealth_utils.py @@ -8,6 +8,7 @@ from typing import List, Dict, Optional import logging + class StealthManager: def __init__(self): self.min_delay = 1.0 @@ -17,12 +18,12 @@ def __init__(self): self.request_count = 0 self.last_request_time = 0 self.logger = logging.getLogger(__name__) - + def set_delay_range(self, min_delay: float, max_delay: float): """Set delay range for requests""" self.min_delay = min_delay self.max_delay = max_delay - + def load_proxies_from_file(self, proxy_file: str): """Load proxies from file""" try: @@ -31,38 +32,38 @@ def load_proxies_from_file(self, proxy_file: str): self.logger.info(f"Loaded {len(self.proxies)} proxies") except Exception as e: self.logger.error(f"Error loading proxies: {e}") - + def get_random_proxy(self) -> Optional[str]: """Get random proxy from list""" if not self.proxies: return None return random.choice(self.proxies) - + def rotate_proxy(self): """Rotate to next proxy""" if not self.proxies: return - + if self.current_proxy in self.proxies: current_index = self.proxies.index(self.current_proxy) next_index = (current_index + 1) % len(self.proxies) else: next_index = 0 - + self.current_proxy = self.proxies[next_index] self.logger.debug(f"Rotated to proxy: {self.current_proxy}") - + def get_request_delay(self) -> float: """Get random delay for request""" return random.uniform(self.min_delay, self.max_delay) - + def wait_before_request(self): """Wait before making request""" delay = self.get_request_delay() time.sleep(delay) self.last_request_time = time.time() self.request_count += 1 - + def get_request_stats(self) -> Dict: """Get request statistics""" return { @@ -70,11 +71,15 @@ def get_request_stats(self) -> Dict: "total_proxies": len(self.proxies), "min_delay": self.min_delay, "max_delay": self.max_delay, - "last_request": time.strftime( - "%Y-%m-%d %H:%M:%S", - time.localtime(self.last_request_time) - ) if self.last_request_time else "Never" + "last_request": ( + time.strftime( + "%Y-%m-%d %H:%M:%S", time.localtime(self.last_request_time) + ) + if self.last_request_time + else "Never" + ), } + # Global stealth manager instance -stealth_manager = StealthManager() \ No newline at end of file +stealth_manager = StealthManager() diff --git a/vulnforge_main.py b/vulnforge_main.py index a20d778..e2abb3d 100755 --- a/vulnforge_main.py +++ b/vulnforge_main.py @@ -27,7 +27,7 @@ from recon_module import EnhancedReconModule from ai_integration import AIAnalyzer, OllamaClient -from ai_orchestrator import AIOrchestrator # New Import +from ai_orchestrator import AIOrchestrator # New Import class VulnForge: @@ -128,7 +128,9 @@ def install_missing_tools(self): async def run_recon(self, target: str, output_dir: Optional[Path] = None): """Run reconnaissance on target""" - recon = EnhancedReconModule(self.base_dir, self.ai_analyzer, config_path="configs/tools.json") + recon = EnhancedReconModule( + self.base_dir, self.ai_analyzer, config_path="configs/tools.json" + ) return await recon.run_recon(target, output_dir) def ask_ai(self, question: str): @@ -387,7 +389,7 @@ async def main(): For detailed documentation, visit: https://github.com/Arunking9/VulnForge """, - formatter_class=argparse.RawDescriptionHelpFormatter + formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--target", "-t", help="Target domain or IP") parser.add_argument( @@ -423,30 +425,52 @@ async def main(): "--stealth", "-s", action="store_true", help="Enable stealth mode" ) parser.add_argument( - "--ai-pipeline", action="store_true", help="Enable the advanced multi-prompt AI pipeline." + "--ai-pipeline", + action="store_true", + help="Enable the advanced multi-prompt AI pipeline.", ) parser.add_argument( - "--prompt-dir", help="Directory for the AI pipeline prompts.", default="AI_Propmt/system-prompts-and-models-of-ai-tools" + "--prompt-dir", + help="Directory for the AI pipeline prompts.", + default="AI_Propmt/system-prompts-and-models-of-ai-tools", ) # Add subparsers for commands - subparsers = parser.add_subparsers(dest='command', help='Available commands') - + subparsers = parser.add_subparsers(dest="command", help="Available commands") + # Ask AI command - ask_ai_parser = subparsers.add_parser('ask-ai', help='Ask the AI assistant a question') - ask_ai_parser.add_argument('question', help='The question to ask the AI') - ask_ai_parser.add_argument('--verbose', action='store_true', help='Show detailed model logs') - ask_ai_parser.add_argument('--dangerous', action='store_true', help='Enable dangerous mode') - ask_ai_parser.add_argument('--confirm-danger', action='store_true', help='Confirm dangerous mode') - + ask_ai_parser = subparsers.add_parser( + "ask-ai", help="Ask the AI assistant a question" + ) + ask_ai_parser.add_argument("question", help="The question to ask the AI") + ask_ai_parser.add_argument( + "--verbose", action="store_true", help="Show detailed model logs" + ) + ask_ai_parser.add_argument( + "--dangerous", action="store_true", help="Enable dangerous mode" + ) + ask_ai_parser.add_argument( + "--confirm-danger", action="store_true", help="Confirm dangerous mode" + ) + # Generate tool command - generate_tool_parser = subparsers.add_parser('generate-tool', help='Generate a custom tool') - generate_tool_parser.add_argument('description', help='Description of the tool to generate') - generate_tool_parser.add_argument('--verbose', action='store_true', help='Show detailed generation logs') - + generate_tool_parser = subparsers.add_parser( + "generate-tool", help="Generate a custom tool" + ) + generate_tool_parser.add_argument( + "description", help="Description of the tool to generate" + ) + generate_tool_parser.add_argument( + "--verbose", action="store_true", help="Show detailed generation logs" + ) + # List tools command - list_tools_parser = subparsers.add_parser('list-tools', help='List all custom tools') - list_tools_parser.add_argument('--verbose', action='store_true', help='Show detailed tool information') + list_tools_parser = subparsers.add_parser( + "list-tools", help="List all custom tools" + ) + list_tools_parser.add_argument( + "--verbose", action="store_true", help="Show detailed tool information" + ) args = parser.parse_args() @@ -457,14 +481,16 @@ async def main(): # Handle AI Pipeline Mode if args.ai_pipeline: if not args.target: - print("Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'") + print( + "Error: A target is required for AI pipeline mode, e.g., --target 'scan example.com'" + ) return - + prompt_path = Path(args.prompt_dir) if not prompt_path.exists(): print(f"Error: Prompt directory not found at '{prompt_path}'") return - + orchestrator = AIOrchestrator(prompt_path) orchestrator.execute_task(f"Perform a security scan on {args.target}") return