from praisonaiagents import Agent, Task, AgentTeam
import time
from typing import Dict, List
import asyncio

def scan_code_patterns():
    """Simulates code pattern scanning"""
    patterns = [
        {
            "type": "buffer_overflow",
            "severity": "critical",
            "location": "memory_management",
            "frequency": 3
        },
        {
            "type": "sql_injection",
            "severity": "high",
            "location": "database_interface",
            "frequency": 2
        },
        {
            "type": "race_condition",
            "severity": "medium",
            "location": "thread_handling",
            "frequency": 1
        }
    ]
    return patterns[int(time.time()) % 3]

def analyze_attack_vectors(pattern: Dict):
    """Analyzes potential attack vectors"""
    vectors = []
    severity_scores = {"critical": 9.5, "high": 8.0, "medium": 6.5}
    base_score = severity_scores[pattern["type"]]
    
    vectors.append({
        "type": pattern["type"],
        "attack_path": f"exploit_{pattern['location']}",
        "complexity": "high" if pattern["frequency"] < 2 else "medium",
        "cvss_score": base_score + (time.time() % 1)
    })
    return vectors

def simulate_exploitation(vectors: List[Dict]):
    """Simulates exploitation attempts"""
    results = []
    for vector in vectors:
        results.append({
            "vector": vector["type"],
            "success_rate": 0.7 + (time.time() % 3) / 10,
            "detection_time": 100 + (int(time.time()) % 50),
            "mitigation_possible": vector["complexity"] != "high"
        })
    return results

def generate_signatures(results: List[Dict]):
    """Generates vulnerability signatures"""
    signatures = []
    for result in results:
        signatures.append({
            "pattern": f"signature_{result['vector']}",
            "reliability": result['success_rate'],
            "false_positive_rate": 0.1 + (time.time() % 2) / 10,
            "detection_speed": result['detection_time']
        })
    return signatures

def validate_findings(signatures: List[Dict]):
    """Validates vulnerability findings"""
    validations = []
    for sig in signatures:
        validations.append({
            "signature": sig["pattern"],
            "confidence": sig["reliability"] * (1 - sig["false_positive_rate"]),
            "reproducibility": 0.8 + (time.time() % 2) / 10,
            "severity": "critical" if sig["reliability"] > 0.8 else "high"
        })
    return validations

# Create specialized agents
pattern_scanner = Agent(
    name="Pattern Scanner",
    role="Code Analysis",
    goal="Scan for suspicious patterns",
    instructions="Analyze code for potential vulnerabilities",
    tools=[scan_code_patterns]
)

vector_analyzer = Agent(
    name="Vector Analyzer",
    role="Attack Vector Analysis",
    goal="Analyze attack vectors",
    instructions="Identify potential attack paths",
    tools=[analyze_attack_vectors]
)

exploit_simulator = Agent(
    name="Exploit Simulator",
    role="Exploitation Testing",
    goal="Simulate exploitation attempts",
    instructions="Test potential vulnerabilities",
    tools=[simulate_exploitation]
)

signature_generator = Agent(
    name="Signature Generator",
    role="Signature Creation",
    goal="Generate vulnerability signatures",
    instructions="Create detection signatures",
    tools=[generate_signatures]
)

finding_validator = Agent(
    name="Finding Validator",
    role="Validation",
    goal="Validate findings",
    instructions="Verify vulnerability findings",
    tools=[validate_findings]
)

# Create workflow tasks
scanning_task = Task(
    name="scan_patterns",
    description="Scan code patterns",
    expected_output="Suspicious patterns",
    agent=pattern_scanner,
    is_start=True,
    next_tasks=["analyze_vectors"]
)

vector_task = Task(
    name="analyze_vectors",
    description="Analyze attack vectors",
    expected_output="Attack vectors",
    agent=vector_analyzer,
    next_tasks=["simulate_exploits"]
)

simulation_task = Task(
    name="simulate_exploits",
    description="Simulate exploits",
    expected_output="Exploitation results",
    agent=exploit_simulator,
    next_tasks=["generate_signatures"]
)

signature_task = Task(
    name="generate_signatures",
    description="Generate signatures",
    expected_output="Vulnerability signatures",
    agent=signature_generator,
    next_tasks=["validate_findings"]
)

validation_task = Task(
    name="validate_findings",
    description="Validate findings",
    expected_output="Validation results",
    agent=finding_validator,
    task_type="decision",
    condition={
        "critical": ["scan_patterns"],  # Continue scanning if critical
        "high": ["analyze_vectors"],  # Reanalyze if high
        "medium": ""  # End workflow if medium
    }
)

# Create workflow
workflow = AgentTeam(
    agents=[pattern_scanner, vector_analyzer, exploit_simulator,
            signature_generator, finding_validator],
    tasks=[scanning_task, vector_task, simulation_task,
           signature_task, validation_task],
    process="workflow", output="verbose"
)

async def main():
    print("\nStarting Vulnerability Detection Workflow...")
    print("=" * 50)
    
    # Run workflow
    results = await workflow.astart()
    
    # Print results
    print("\nVulnerability Detection Results:")
    print("=" * 50)
    task_results = (
        results.get("task_results", {})
        if isinstance(results, dict)
        else {}
    )
    if not task_results and isinstance(results, str):
        print(f"\nWorkflow output:\n{results}")
    for task_id, result in task_results.items():
        if result:
            print(f"\nTask: {task_id}")
            print(f"Result: {result.raw}")
            print("-" * 50)

if __name__ == "__main__":
    asyncio.run(main())