File size: 3,458 Bytes
5a61f72 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | 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)
# Load the skill file
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:
# Use --no-sandbox for Docker/HuggingFace stability
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")
# Initial Screenshot
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...")
# Simulated Agent Thought (Compiling for speed in this demo)
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
|