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("

BugHunter AI

") 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)