import gradio as gr
from anthropic import Anthropic
import os
import sys

client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY", "sk-test"))

def generate_payload(target, technique, temperature):
    """Generate pentesting payload"""
    try:
        response = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=1000,
            temperature=temperature,
            system="Tu es un expert en pentesting. Génère des payloads variés et techniques d'énumération créatives.",
            messages=[{
                "role": "user",
                "content": f"Generate {technique} payload for: {target}"
            }]
        )
        return response.content[0].text
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio interface
with gr.Blocks(title="Pentesting Copilot - Gradio") as demo:
    gr.Markdown("# 🔐 Pentesting Copilot (Gradio)")

    with gr.Row():
        target = gr.Textbox(label="Target", placeholder="Enter target URL/IP")
        technique = gr.Dropdown(
            ["SQL injection", "XSS", "RCE", "SSRF", "Path Traversal", "XXE"],
            label="Technique"
        )

    temperature = gr.Slider(0, 1, step=0.1, value=0.7, label="Temperature")

    submit_btn = gr.Button("Generate Payload")
    output = gr.Textbox(label="Output", lines=10)

    submit_btn.click(
        fn=generate_payload,
        inputs=[target, technique, temperature],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(share=True, show_error=True)