main.py

# main.py
from autogen import AssistantAgent, UserProxyAgent
from tenacity import retry, stop_after_attempt, wait_exponential
import subprocess
import shlex
import xml.etree.ElementTree as ET
import requests
import os
import getpass
import argparse
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# ======================
# Configuration
# ======================
api_key = os.getenv("DEEPSEEK_API_KEY")
if not api_key:
    raise ValueError("DEEPSEEK_API_KEY not found in .env file")

config_list = [
    {
        "model": "deepseek-chat",
        "api_key": api_key,
        "base_url": "https://api.deepseek.com/v1"
    }
]

llm_config = {
    "config_list": config_list,
    "temperature": 0.3,
    "max_tokens": 2000
}

def parse_arguments():
    parser = argparse.ArgumentParser(
        description='PentestAgent - Automated Security Assessment Tool'
    )
    parser.add_argument(
        '-t', '--target',
        help='Target IP address or hostname',
        required=True
    )
    parser.add_argument(
        '-o', '--output',
        help='Output file for the report (default: report.txt)',
        default='report.txt'
    )
    return parser.parse_args()

# ======================
# Tool Definitions
# ======================
def nmap_scan(target: str, options: str) -> str:
    """Perform Nmap network scan using command line and return parsed results."""
    try:
        # Split options into list, handling quoted arguments
        options_list = shlex.split(options)
        # Build command with XML output
        command = ["nmap"] + options_list + ["-oX", "-", target]
        
        # Execute command
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            check=True
        )
        
        # Parse XML output
        root = ET.fromstring(result.stdout)
        results = []
        
        for host in root.findall('host'):
            # Get host address
            address_elem = host.find('address[@addrtype="ipv4"]')
            address = address_elem.get('addr') if address_elem is not None else target
            results.append(f"Scan Results for {address}")
            
            # Process ports
            for port in host.findall('.//port'):
                portid = port.get('portid')
                protocol = port.get('protocol')
                state = port.find('state').get('state')
                service_elem = port.find('service')
                service_name = service_elem.get('name') if service_elem else 'unknown'
                product = service_elem.get('product', 'N/A') if service_elem else 'N/A'
                version = service_elem.get('version', 'N/A') if service_elem else 'N/A'
                
                version_info = product
                if version != 'N/A':
                    version_info += f" {version}"
                
                results.append(
                    f"Port {portid}/{protocol}: {state} | Service: {service_name} | Version: {version_info}"
                )
        
        return "\n".join(results)
    
    except subprocess.CalledProcessError as e:
        return f"Scan failed: {e.stderr}"
    except FileNotFoundError:
        return "Error: nmap is not installed or not found in PATH."
    except ET.ParseError as e:
        return f"Error parsing XML output: {str(e)}"
    except Exception as e:
        return f"Unexpected error: {str(e)}"

def file_write(filepath: str, content: str) -> str:
    """Write content to specified file."""
    try:
        with open(filepath, "w") as f:
            f.write(content)
        return f"Successfully wrote to {filepath}"
    except Exception as e:
        return f"Write error: {str(e)}"

def search_cves(scan_results: str) -> str:
    """Search for potential CVEs based on scan results."""
    cve_list = ["Potential Vulnerabilities:"]
    lines = scan_results.split('\n')
    
    for line in lines:
        if "Service:" in line and "Version:" in line:
            try:
                parts = [p.strip() for p in line.split('|')]
                service = parts[1].split(': ')[1]
                version = parts[2].split(': ')[1]
                
                if version == 'N/A':
                    continue
                
                # Search CVE database
                response = requests.get(
                    f"https://cve.circl.lu/api/search/{service}/{version}",
                    timeout=10
                )
                
                if response.status_code == 200:
                    data = response.json()
                    if data:
                        for cve in data:
                            cve_id = cve.get('id', 'Unknown CVE')
                            summary = cve.get('summary', 'No description available')
                            cve_list.append(f"{cve_id}: {summary}")
            except Exception as e:
                cve_list.append(f"Error checking {service} {version}: {str(e)}")
    
    return "\n".join(cve_list) if len(cve_list) > 1 else "No potential CVEs found"

# ======================
# Agent Initialization
# ======================
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def init_agents():
    try:
        # Initialize agents
        assistant = AssistantAgent(
            name="CyberSec-Assistant",
            llm_config=llm_config,
            system_message="""
            You are a cybersecurity AI assistant. Follow this workflow:
            1. Generate appropriate nmap command with options including service/version detection (e.g., -sV)
            2. Perform network scan with nmap_scan
            3. Check for vulnerabilities with search_cves
            4. Save results with file_write
            5. Summarize findings
            
            When generating nmap commands:
            - Always include version detection (-sV)
            - Include timing template (-T4) for faster scans
            - Consider adding -sC for default scripts when appropriate
            - Use -p- to scan all ports if needed
            """
        )

        user_proxy = UserProxyAgent(
            name="User-Proxy",
            human_input_mode="NEVER",
            max_consecutive_auto_reply=5,
            code_execution_config={"work_dir": "ops"},
        )

        # Register tools
        for func in [nmap_scan, file_write, search_cves]:
            assistant.register_for_llm(name=func.__name__, description=func.__doc__)(func)
            user_proxy.register_function(function_map={func.__name__: func})

        return assistant, user_proxy
    
    except Exception as e:
        print(f"Initialization error: {str(e)}")
        raise

# ======================
# Main Execution
# ======================
if __name__ == "__main__":
    try:
        args = parse_arguments()
        assistant, user_proxy = init_agents()
        
        user_proxy.initiate_chat(
            assistant,
            message=f"""
            Perform security assessment on {args.target}:
            1. Perform comprehensive network scan with version detection
            2. Identify potential vulnerabilities
            3. Save results to {args.output}
            4. Provide executive summary
            
            Important: Generate nmap options that include service version detection.
            """
        )
        
    except Exception as e:
        print(f"Critical failure: {str(e)}")