pentest_agent_python.py

import os
import json
import time
from datetime import datetime
from typing import List, Dict
import requests
from colorama import Fore, Style, init

# Initialize colorama for cross-platform colored output
init(autoreset=True)


class PenTestAgent:
    def __init__(self, api_key: str):
        """Initialize the penetration testing agent with OpenAI API key."""
        self.api_key = api_key
        self.target = ""
        self.scope = ""
        self.results = []
        self.phases = [
            {"id": "recon", "name": "Reconnaissance", "icon": "🔍"},
            {"id": "scanning", "name": "Scanning & Enumeration", "icon": "🎯"},
            {"id": "vulnerability", "name": "Vulnerability Analysis", "icon": "⚠️"},
            {"id": "exploitation", "name": "Exploitation Assessment", "icon": "🔓"},
            {"id": "reporting", "name": "Report Generation", "icon": "🛡️"}
        ]

    def print_banner(self):
        """Print the application banner."""
        print(f"\n{Fore.CYAN}{'='*70}")
        print(f"{Fore.CYAN}🛡️  AI PENETRATION TESTING AGENT")
        print(f"{Fore.CYAN}{'='*70}{Style.RESET_ALL}\n")
        print(f"{Fore.RED}⚠️  LEGAL WARNING:")
        print(f"{Fore.RED}Only use this tool on systems you own or have explicit written")
        print(f"{Fore.RED}authorization to test. Unauthorized penetration testing is illegal!")
        print(f"{Fore.RED}{'='*70}{Style.RESET_ALL}\n")

    def call_openai(self, prompt: str) -> str:
        """Make a call to OpenAI API."""
        url = "https://api.openai.com/v1/chat/completions"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        payload = {
            "model": "gpt-4",
            "messages": [
                {
                    "role": "system",
                    "content": "You are a professional penetration testing expert. Provide detailed, methodical analysis following industry standards (OWASP, PTES). Focus on educational and ethical testing approaches. Always emphasize legal authorization requirements."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            "temperature": 0.7,
            "max_tokens": 2000
        }

        try:
            response = requests.post(url, headers=headers, json=payload)
            response.raise_for_status()
            data = response.json()
            return data["choices"][0]["message"]["content"]
        except requests.exceptions.RequestException as e:
            return f"Error calling OpenAI API: {str(e)}"

    def generate_prompt(self, phase_id: str) -> str:
        """Generate appropriate prompt for each phase."""
        prompts = {
            "recon": f"""Conduct a reconnaissance analysis for target: {self.target}
Scope: {self.scope}

Provide a detailed reconnaissance plan including:
1. Information gathering methodology (OSINT techniques)
2. Passive reconnaissance approaches
3. DNS enumeration strategies
4. Subdomain discovery methods
5. Technology stack identification
6. Social engineering considerations

Format as a structured analysis with specific techniques and tools.""",

            "scanning": f"""Based on target {self.target}, create a comprehensive scanning and enumeration plan:

1. Network mapping approach
2. Port scanning methodology (TCP/UDP)
3. Service version detection
4. OS fingerprinting techniques
5. Network architecture analysis
6. Common ports and services to investigate

Provide specific scanning strategies and what to look for in results.""",

            "vulnerability": f"""Perform vulnerability analysis for {self.target}:

1. Identify common vulnerability categories (OWASP Top 10, CVE databases)
2. Web application vulnerabilities (SQLi, XSS, CSRF, etc.)
3. Network service vulnerabilities
4. Configuration weaknesses
5. Authentication/authorization flaws
6. Encryption and data protection issues

For each category, explain detection methods and potential impact.""",

            "exploitation": f"""Create an exploitation assessment plan for {self.target}:

1. Prioritize vulnerabilities by severity and exploitability
2. Proof-of-concept exploitation approaches (ethical and authorized)
3. Privilege escalation vectors
4. Lateral movement possibilities
5. Data exfiltration scenarios
6. Persistence mechanisms

Emphasize this is for AUTHORIZED testing only. Explain each exploit category conceptually.""",

            "reporting": f"""Generate a comprehensive penetration testing report for {self.target}:

Include:
1. Executive Summary
2. Methodology Overview
3. Findings Summary (Critical, High, Medium, Low)
4. Detailed Vulnerability Descriptions
5. Proof of Concept (where applicable)
6. Remediation Recommendations
7. Risk Assessment
8. Conclusion and Next Steps

Format as a professional security assessment report."""
        }
        
        return prompts.get(phase_id, "")

    def run_phase(self, phase: Dict) -> Dict:
        """Execute a single penetration testing phase."""
        print(f"\n{Fore.YELLOW}{phase['icon']} Running Phase: {phase['name']}")
        print(f"{Fore.YELLOW}{'-'*70}{Style.RESET_ALL}")
        
        prompt = self.generate_prompt(phase["id"])
        timestamp = datetime.now().strftime("%H:%M:%S")
        
        try:
            result = self.call_openai(prompt)
            
            phase_result = {
                "phase": phase["name"],
                "content": result,
                "timestamp": timestamp,
                "error": False
            }
            
            print(f"{Fore.GREEN}✓ Phase completed successfully at {timestamp}{Style.RESET_ALL}")
            
        except Exception as e:
            phase_result = {
                "phase": phase["name"],
                "content": f"Error: {str(e)}",
                "timestamp": timestamp,
                "error": True
            }
            print(f"{Fore.RED}✗ Phase failed: {str(e)}{Style.RESET_ALL}")
        
        self.results.append(phase_result)
        return phase_result

    def display_result(self, result: Dict):
        """Display a single phase result."""
        color = Fore.RED if result["error"] else Fore.CYAN
        print(f"\n{color}{'='*70}")
        print(f"{result['phase']} - {result['timestamp']}")
        print(f"{'='*70}{Style.RESET_ALL}")
        print(f"\n{result['content']}\n")

    def start_pentest(self):
        """Execute the full penetration testing process."""
        self.print_banner()
        
        # Get user input
        self.target = input(f"{Fore.CYAN}Enter target system/URL: {Style.RESET_ALL}").strip()
        if not self.target:
            print(f"{Fore.RED}Error: Target cannot be empty{Style.RESET_ALL}")
            return
        
        self.scope = input(f"{Fore.CYAN}Enter testing scope: {Style.RESET_ALL}").strip()
        if not self.scope:
            print(f"{Fore.RED}Error: Scope cannot be empty{Style.RESET_ALL}")
            return
        
        print(f"\n{Fore.GREEN}Starting penetration test for: {self.target}{Style.RESET_ALL}")
        print(f"{Fore.GREEN}Scope: {self.scope}{Style.RESET_ALL}")
        
        # Run all phases
        for phase in self.phases:
            result = self.run_phase(phase)
            time.sleep(1)  # Brief pause between phases
        
        # Generate final report
        print(f"\n{Fore.MAGENTA}{'='*70}")
        print(f"Generating Executive Summary Report...")
        print(f"{'='*70}{Style.RESET_ALL}")
        
        final_prompt = f"""Create a final executive summary report combining all phases of penetration testing for {self.target}. Include:
- Overall security posture
- Critical findings count
- Key recommendations
- Risk level assessment
Keep it concise and executive-friendly."""
        
        final_report = self.call_openai(final_prompt)
        
        print(f"\n{Fore.GREEN}{'='*70}")
        print(f"EXECUTIVE SUMMARY REPORT")
        print(f"{'='*70}{Style.RESET_ALL}")
        print(f"\n{final_report}\n")
        
        # Save results to file
        self.save_results(final_report)

    def save_results(self, final_report: str):
        """Save all results to a JSON file."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"pentest_report_{self.target.replace('.', '_')}_{timestamp}.json"
        
        report_data = {
            "target": self.target,
            "scope": self.scope,
            "timestamp": datetime.now().isoformat(),
            "phases": self.results,
            "executive_summary": final_report
        }
        
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(report_data, f, indent=2, ensure_ascii=False)
            print(f"{Fore.GREEN}✓ Full report saved to: {filename}{Style.RESET_ALL}\n")
        except Exception as e:
            print(f"{Fore.RED}✗ Error saving report: {str(e)}{Style.RESET_ALL}\n")

    def interactive_mode(self):
        """Run in interactive mode with menu."""
        self.print_banner()
        
        while True:
            print(f"\n{Fore.CYAN}MENU:{Style.RESET_ALL}")
            print("1. Start New Penetration Test")
            print("2. View Previous Results")
            print("3. Exit")
            
            choice = input(f"\n{Fore.CYAN}Select option: {Style.RESET_ALL}").strip()
            
            if choice == "1":
                self.results = []  # Clear previous results
                self.start_pentest()
            elif choice == "2":
                self.view_results()
            elif choice == "3":
                print(f"\n{Fore.YELLOW}Exiting... Stay safe and test responsibly!{Style.RESET_ALL}\n")
                break
            else:
                print(f"{Fore.RED}Invalid option. Please try again.{Style.RESET_ALL}")

    def view_results(self):
        """Display stored results."""
        if not self.results:
            print(f"\n{Fore.YELLOW}No results available. Run a penetration test first.{Style.RESET_ALL}")
            return
        
        for result in self.results:
            self.display_result(result)


def main():
    """Main entry point."""
    print(f"{Fore.CYAN}AI Penetration Testing Agent - Python Version{Style.RESET_ALL}")
    
    # Get API key from environment or user input
    api_key = os.getenv("OPENAI_API_KEY")
    
    if not api_key:
        api_key = input(f"\n{Fore.CYAN}Enter your OpenAI API key: {Style.RESET_ALL}").strip()
    
    if not api_key:
        print(f"{Fore.RED}Error: API key is required{Style.RESET_ALL}")
        return
    
    # Create agent and start
    agent = PenTestAgent(api_key)
    agent.interactive_mode()


if __name__ == "__main__":
    main()