| import os |
| import asyncio |
| from playwright.async_api import async_playwright |
| import google.generativeai as genai |
| from groq import Groq |
| from bs4 import BeautifulSoup |
|
|
| class BugHunterEngine: |
| def __init__(self, api_key, provider="Gemini", model_name="gemini-1.5-pro-preview-0514"): |
| self.provider = provider |
| self.api_key = api_key |
| self.model_name = model_name |
| |
| if provider == "Gemini": |
| genai.configure(api_key=api_key) |
| self.gemini_model = genai.GenerativeModel(model_name) |
| elif provider == "Groq": |
| self.groq_client = Groq(api_key=api_key) |
|
|
| |
| skill_path = os.path.join(os.path.dirname(__file__), "ADVANCED_QA_AGENT_SKILL.md") |
| with open(skill_path, "r", encoding="utf-8") as f: |
| self.skill_content = f.read() |
|
|
| def call_ai(self, prompt): |
| """Unified method to call Gemini or Groq.""" |
| if self.provider == "Gemini": |
| response = self.gemini_model.generate_content(prompt) |
| return response.text |
| else: |
| chat_completion = self.groq_client.chat.completions.create( |
| messages=[{"role": "system", "content": self.skill_content}, {"role": "user", "content": prompt}], |
| model="llama3-70b-8192", |
| ) |
| return chat_completion.choices[0].message.content |
|
|
| async def prune_dom(self, page_source): |
| soup = BeautifulSoup(page_source, "lxml") |
| for tag in soup(["script", "style", "svg", "path", "iframe"]): |
| tag.decompose() |
| interactive_tags = ["a", "button", "input", "select", "textarea", "h1", "h2", "h3", "h4", "h5", "h6", "p"] |
| pruned_html = "" |
| for tag in soup.find_all(interactive_tags): |
| attrs = {k: v for k, v in tag.attrs.items() if k in ["id", "class", "href", "name", "type", "role"]} |
| tag.attrs = attrs |
| pruned_html += str(tag) + "\n" |
| return pruned_html[:30000] |
|
|
| async def run_test(self, url, progress=None): |
| logs = [] |
| screenshots = [] |
|
|
| async with async_playwright() as p: |
| |
| browser = await p.chromium.launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox']) |
| page = await browser.new_page(viewport={"width": 1280, "height": 800}) |
| |
| if progress: progress(0.1, desc="🌍 Opening Website...") |
| await page.goto(url, wait_until="networkidle") |
| |
| |
| shot_path = "evidence_initial.png" |
| await page.screenshot(path=shot_path) |
| screenshots.append(shot_path) |
| |
| dom = await self.prune_dom(await page.content()) |
| |
| if progress: progress(0.5, desc="🧠 AI Agent Exploring...") |
| |
| |
| final_prompt = f""" |
| {self.skill_content} |
| |
| URL: {url} |
| HTML SNAPSHOT: |
| {dom} |
| |
| ACT AS AN EXPERT QA AGENT. |
| Generate a SUPERIOR SBTM v2 Exploratory Test Report in Markdown. |
| You MUST follow the reporting protocol from the skill file. |
| IDENTIFY BUGS relevant to: SFDIPOT, HICCUPPS, A11y, and Security (IDOR/JWT). |
| |
| Return ONLY the Markdown report. |
| """ |
| report_md = self.call_ai(final_prompt) |
| |
| await browser.close() |
| return report_md, screenshots |
|
|
