app.py

import os
import asyncio
import gradio as gr
from playwright.async_api import async_playwright
from google import genai
from google.genai import types
from groq import Groq
from datetime import datetime
from bs4 import BeautifulSoup
import traceback
from fpdf import FPDF
import re

# --- CONSTANTS & CONFIG ---
SKILL_FILE_PATH = "ADVANCED_QA_AGENT_SKILL.md"
REPORTS_DIR = "reports"
os.makedirs(REPORTS_DIR, exist_ok=True)

# --- RATE LIMITING ---
IP_TRACKER = {}
MAX_REQUESTS_PER_HOUR = 5

def check_rate_limit(request: gr.Request):
    if request is None: return True
    ip = getattr(request.client, "host", "unknown")
    now = datetime.now()
    if ip not in IP_TRACKER:
        IP_TRACKER[ip] = []
    
    # Filter out requests older than 1 hour
    IP_TRACKER[ip] = [t for t in IP_TRACKER[ip] if (now - t).total_seconds() < 3600]
    
    if len(IP_TRACKER[ip]) >= MAX_REQUESTS_PER_HOUR:
        return False
    IP_TRACKER[ip].append(now)
    return True

# THEME FIX: Aggressively forcing white text for readability
CSS = """
body { background-color: #0b0f19; color: #e2e8f0; font-family: 'Inter', sans-serif; }
.gradio-container { background: rgba(15, 23, 42, 0.7) !important; backdrop-filter: blur(20px); border: 1px solid rgba(255, 255, 255, 0.1); border-radius: 20px; }
.glass-panel { background: rgba(30, 41, 59, 0.7) !important; padding: 30px !important; border-radius: 12px; border: 1px solid rgba(255, 255, 255, 0.1); }
/* AGGRESSIVE TEXT OVERRIDE */
.glass-panel *, .glass-panel p, .glass-panel li, .glass-panel span, .glass-panel div { 
    color: #ffffff !important; 
    font-size: 1.05rem !important;
}
.glass-panel h1, .glass-panel h2, .glass-panel h3 { 
    color: #60a5fa !important; 
    font-weight: 700 !important;
    margin-top: 1.5rem !important;
}
.primary-btn { background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%) !important; border: none !important; color: white !important; font-weight: 600 !important; }
.status-log { font-family: 'JetBrains Mono', monospace; font-size: 0.85rem; color: #94a3b8; background: #020617; padding: 15px; border: 1px solid #1e293b; max-height: 400px; overflow-y: auto; }
h1 { background: linear-gradient(to right, #60a5fa, #a78bfa); -webkit-background-clip: text; -webkit-text-fill-color: transparent; font-weight: 800; font-size: 2.5rem; text-align: center; }
"""

class BugHunterAI:
    def __init__(self):
        self.gemini_key = os.environ.get("GEMINI_API_KEY")
        self.groq_key = os.environ.get("GROQ_API_KEY")
        self.skill_content = self._load_skill()
        self.logs = []

    def _load_skill(self):
        try:
            if os.path.exists(SKILL_FILE_PATH):
                with open(SKILL_FILE_PATH, "r", encoding="utf-8") as f:
                    return f.read()
            return "Skill file not found."
        except: return "Error loading skill."

    def log(self, message):
        timestamp = datetime.now().strftime("%H:%M:%S")
        log_entry = f"[{timestamp}] {message}"
        self.logs.append(log_entry)
        return "\n".join(self.logs)

    def generate_pdf(self, text, filename):
        try:
            pdf = FPDF()
            pdf.add_page()
            pdf.set_font("Helvetica", size=12)
            # Remove markdown syntax for cleaner PDF
            clean_text = re.sub(r'[*#_`]', '', text)
            pdf.multi_cell(0, 10, clean_text)
            path = os.path.join(REPORTS_DIR, filename)
            pdf.output(path)
            return path
        except Exception as e:
            print(f"PDF creation error: {e}")
            return None

    async def get_ai_response(self, prompt, system_instruction):
        try:
            if self.gemini_key:
                client = genai.Client(api_key=self.gemini_key)
                for model_name in ['gemini-3.1-pro-preview', 'gemini-2.5-flash', 'gemini-1.5-flash-latest']:
                    try:
                        response = await asyncio.to_thread(
                            client.models.generate_content,
                            model=model_name,
                            contents=prompt,
                            config=types.GenerateContentConfig(
                                system_instruction=system_instruction
                            )
                        )
                        return response.text
                    except: continue
            
            if self.groq_key:
                client = Groq(api_key=self.groq_key)
                for model_name in ['llama-4-scout', 'gemma-4-31b-it', 'llama-3.3-70b-versatile']:
                    try:
                        res = client.chat.completions.create(
                            messages=[{"role":"system","content":system_instruction},{"role":"user","content":prompt}],
                            model=model_name
                        )
                        return res.choices[0].message.content
                    except: continue
            return "No valid API config."
        except Exception as e: return f"Error: {str(e)}"

    async def run_audit(self, url):
        self.logs = []
        try:
            if not url.startswith("http"): url = "https://" + url
            yield self.log(f"🚀 AUDIT START: {url}"), "", None
            
            async with async_playwright() as p:
                yield self.log("🌐 Launching Browser..."), "", None
                browser = await p.chromium.launch(headless=True)
                page = await browser.new_page()
                await page.goto(url, wait_until="load", timeout=60000)
                
                yield self.log("📸 Scanning..."), "", None
                content = await page.content()
                soup = BeautifulSoup(content, 'html.parser')
                for el in soup(["script", "style", "nav", "footer"]): el.decompose()
                
                yield self.log("🧠 Consulting Skills..."), "", None
                prompt = f"Target: {url}\n\nCONTENT:\n{soup.get_text()[:4000]}\n\nReport findings."
                report = await self.get_ai_response(prompt, self.skill_content)
                
                # Save PDF
                pdf_path = self.generate_pdf(report, f"audit_{datetime.now().strftime('%H%M%S')}.pdf")
                
                yield self.log("✅ Audit Complete."), report, pdf_path
                await browser.close()
        except Exception as e:
            yield self.log(f"❌ ERROR: {str(e)}"), f"Audit Failed: {str(e)}", None

def create_ui():
    with gr.Blocks(css=CSS) as demo:
        gr.Markdown("<h1>BugHunter AI</h1>")
        with gr.Row():
            with gr.Column(scale=2):
                url_input = gr.Textbox(label="Website URL")
                with gr.Row():
                    captcha_input = gr.Textbox(label="Anti-Bot Check: What is 7 + 4?")
                btn = gr.Button("🚀 Run Audit", variant="primary", elem_classes=["primary-btn"])
                out_file = gr.File(label="📄 Download PDF Report")
            with gr.Column(scale=1):
                logs = gr.Markdown("### 🕒 Logs", elem_classes=["status-log"])
        report = gr.Markdown("### 🔍 Report", elem_classes=["glass-panel"])

        async def start(u, captcha, request: gr.Request):
            if captcha.strip() != "11":
                yield "❌ CAPTCHA Failed. Are you a bot?", "Audit Failed: Incorrect Captcha.", None
                return
                
            if not check_rate_limit(request):
                yield "❌ Rate Limit Exceeded.", "Audit Failed: You have exceeded the maximum number of requests per hour (5).", None
                return
                
            hunter = BugHunterAI()
            async for l, r, f in hunter.run_audit(u):
                yield l, r, f

        btn.click(start, inputs=[url_input, captcha_input], outputs=[logs, report, out_file])
    return demo

if __name__ == "__main__":
    demo = create_ui()
    demo.queue(default_concurrency_limit=2).launch(server_name="0.0.0.0", server_port=7860)