502 أسطر
17 KiB
Python
502 أسطر
17 KiB
Python
"""
|
|
client.py — Browser Automation & Payload Execution Layer
|
|
=========================================================
|
|
Uses Playwright (sync API) to drive headless Chromium against https://os.solidpoint.ai.
|
|
Handles login, chat-message injection, auto-continue monitoring, and response extraction.
|
|
|
|
⚠️ NO print() calls anywhere — all output goes to logging (stderr / file).
|
|
"""
|
|
|
|
import csv
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Logging — stderr + file, NEVER stdout
|
|
# ---------------------------------------------------------------------------
|
|
LOG_FILE = Path(__file__).parent.parent / "logs" / "mcp_manager.log"
|
|
|
|
_fmt = logging.Formatter(
|
|
"[%(asctime)s] %(levelname)-7s %(name)s — %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
|
|
_sh = logging.StreamHandler(sys.stderr)
|
|
_sh.setFormatter(_fmt)
|
|
|
|
_fh = logging.FileHandler(LOG_FILE, encoding="utf-8")
|
|
_fh.setFormatter(_fmt)
|
|
|
|
log = logging.getLogger("redteam.client")
|
|
log.setLevel(logging.DEBUG)
|
|
log.addHandler(_sh)
|
|
log.addHandler(_fh)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Paths
|
|
# ---------------------------------------------------------------------------
|
|
BASE_DIR = Path(__file__).parent.parent
|
|
ENV_FILE = BASE_DIR / "config" / ".env"
|
|
ACCOUNTS_CSV = BASE_DIR / "data" / "accounts.csv"
|
|
TESTCASES_JSON = BASE_DIR / "data" / "redteam_testcases.json"
|
|
EVIDENCE_DIR = BASE_DIR / "evidence_screenshots"
|
|
EVIDENCE_DIR.mkdir(exist_ok=True)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Configuration loader
|
|
# ---------------------------------------------------------------------------
|
|
def load_env() -> dict:
|
|
"""Parse .enc file into a dict (simple KEY=VALUE format)."""
|
|
env = {}
|
|
if not ENV_FILE.exists():
|
|
log.warning(".env file not found at %s", ENV_FILE)
|
|
return env
|
|
for line in ENV_FILE.read_text(encoding="utf-8").splitlines():
|
|
line = line.strip()
|
|
if not line or line.startswith("#"):
|
|
continue
|
|
if "=" in line:
|
|
k, v = line.split("=", 1)
|
|
env[k.strip()] = v.strip()
|
|
return env
|
|
|
|
|
|
def load_accounts() -> list[dict]:
|
|
"""Load test accounts from accounts.csv, filtering only SUCCESS rows."""
|
|
accounts = []
|
|
if not ACCOUNTS_CSV.exists():
|
|
log.warning("accounts.csv not found at %s", ACCOUNTS_CSV)
|
|
return accounts
|
|
with open(ACCOUNTS_CSV, newline="", encoding="utf-8") as f:
|
|
reader = csv.DictReader(f)
|
|
for row in reader:
|
|
if row.get("Status", "").strip().upper() == "SUCCESS":
|
|
accounts.append({
|
|
"email": row["Email"].strip().strip('"'),
|
|
"password": row["Password"].strip().strip('"'),
|
|
})
|
|
log.info("Loaded %d valid accounts from CSV", len(accounts))
|
|
return accounts
|
|
|
|
|
|
def load_testcases() -> list[dict]:
|
|
"""Load red-team test cases from JSON."""
|
|
if not TESTCASES_JSON.exists():
|
|
log.error("redteam_testcases.json not found at %s", TESTCASES_JSON)
|
|
return []
|
|
with open(TESTCASES_JSON, encoding="utf-8") as f:
|
|
cases = json.load(f)
|
|
log.info("Loaded %d test cases", len(cases))
|
|
return cases
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Browser helpers
|
|
# ---------------------------------------------------------------------------
|
|
# Selectors observed from the os.solidpoint.ai UI (Preact SPA)
|
|
SEL_EMAIL = '#login-email, input[type="email"]'
|
|
SEL_PASS = '#login-pass, input[type="password"]'
|
|
SEL_LOGIN_BTN = ".login-btn"
|
|
SEL_SIGNUP_TOGGLE = ".login-toggle button"
|
|
SEL_LOGIN_ERROR = ".login-error"
|
|
SEL_TEXTAREA = ".input-row textarea"
|
|
SEL_SEND_BTN = ".send-btn"
|
|
SEL_MSG_BODY = ".m-body"
|
|
SEL_TYPING = ".typing"
|
|
SEL_WELCOME = ".welcome"
|
|
SEL_SIDEBAR = ".sidebar"
|
|
SEL_CONTINUE_BTN = 'button:has-text("Continue"), button:has-text("continue")'
|
|
SEL_UPGRADE_CLOSE = ".upgrade-cancel, .upgrade-overlay"
|
|
|
|
|
|
def _wait_for_app(page, timeout: int = 15_000):
|
|
"""Wait until the SPA renders (login or chat view)."""
|
|
page.wait_for_load_state("networkidle", timeout=timeout)
|
|
# Either login-box or sidebar/welcome must appear
|
|
page.wait_for_selector(
|
|
f"{SEL_EMAIL}, {SEL_SIDEBAR}, {SEL_WELCOME}",
|
|
timeout=timeout,
|
|
)
|
|
|
|
|
|
def login_with_credentials(page, email: str, password: str, is_signup: bool = False) -> bool:
|
|
"""
|
|
Log in (or sign up) via the UI.
|
|
Returns True on success, False on failure.
|
|
"""
|
|
log.info("Attempting %s for %s", "signup" if is_signup else "signin", email)
|
|
|
|
page.goto("https://os.solidpoint.ai", wait_until="networkidle", timeout=30_000)
|
|
_wait_for_app(page)
|
|
|
|
# Switch to signup mode if needed
|
|
if is_signup:
|
|
toggle = page.locator(SEL_SIGNUP_TOGGLE)
|
|
if toggle.is_visible():
|
|
txt = (toggle.text_content() or "").lower()
|
|
if "sign up" in txt:
|
|
toggle.click()
|
|
page.wait_for_timeout(500)
|
|
|
|
# Fill credentials
|
|
page.locator(SEL_EMAIL).first.fill(email)
|
|
page.locator(SEL_PASS).first.fill(password)
|
|
page.locator(SEL_LOGIN_BTN).first.click()
|
|
|
|
# Wait for transition
|
|
page.wait_for_timeout(4_000)
|
|
|
|
# Check for errors
|
|
err_el = page.locator(SEL_LOGIN_ERROR)
|
|
if err_el.is_visible():
|
|
err_text = err_el.text_content() or ""
|
|
log.warning("Login error for %s: %s", email, err_text.strip())
|
|
return False
|
|
|
|
# Success indicators: sidebar, welcome screen, or chat area visible
|
|
success = (
|
|
page.locator(SEL_SIDEBAR).is_visible()
|
|
or page.locator(SEL_WELCOME).is_visible()
|
|
or page.locator(SEL_TEXTAREA).is_visible()
|
|
)
|
|
if success:
|
|
log.info("Login succeeded for %s", email)
|
|
else:
|
|
log.warning("Login outcome unclear for %s (no sidebar/welcome)", email)
|
|
|
|
return success
|
|
|
|
|
|
def login_with_token(page, token: str) -> bool:
|
|
"""
|
|
Inject PRO_CHAT_TOKEN into localStorage and reload to authenticate.
|
|
"""
|
|
log.info("Injecting token via localStorage…")
|
|
page.goto("https://os.solidpoint.ai", wait_until="networkidle", timeout=30_000)
|
|
page.evaluate(f"localStorage.setItem('chat_token', '{token}')")
|
|
page.reload(wait_until="networkidle", timeout=30_000)
|
|
_wait_for_app(page)
|
|
|
|
success = (
|
|
page.locator(SEL_SIDEBAR).is_visible()
|
|
or page.locator(SEL_WELCOME).is_visible()
|
|
or page.locator(SEL_TEXTAREA).is_visible()
|
|
)
|
|
if success:
|
|
log.info("Token-based auth succeeded")
|
|
else:
|
|
log.warning("Token-based auth outcome unclear")
|
|
return success
|
|
|
|
|
|
def _dismiss_upgrade_modal(page):
|
|
"""Dismiss the upgrade modal if it appears."""
|
|
cancel = page.locator(".upgrade-cancel")
|
|
if cancel.is_visible():
|
|
cancel.click()
|
|
page.wait_for_timeout(300)
|
|
|
|
|
|
def _create_new_chat(page):
|
|
"""Click the new-chat button in sidebar."""
|
|
new_btn = page.locator(".sb-new").first
|
|
if new_btn.is_visible():
|
|
new_btn.click()
|
|
page.wait_for_timeout(1_000)
|
|
log.debug("Created new chat session")
|
|
|
|
|
|
def _upload_file_for_testcase(page, file_spec: dict) -> Optional[str]:
|
|
"""
|
|
For indirect-injection test cases that include a file payload,
|
|
create a temp file and attach it via the file input.
|
|
Returns the temp file path or None.
|
|
"""
|
|
if not file_spec:
|
|
return None
|
|
|
|
fname = file_spec.get("name", "payload.txt")
|
|
content = file_spec.get("content", "")
|
|
|
|
tmp_dir = Path(tempfile.mkdtemp(prefix="redteam_"))
|
|
tmp_path = tmp_dir / fname
|
|
tmp_path.write_text(content, encoding="utf-8")
|
|
log.debug("Created temp payload file: %s", tmp_path)
|
|
|
|
# Try to locate file input (hidden) — the attach button triggers it
|
|
file_input = page.locator('input[type="file"]')
|
|
if file_input.count() > 0:
|
|
file_input.first.set_input_files(str(tmp_path))
|
|
page.wait_for_timeout(1_000)
|
|
log.debug("Attached file via input[type=file]")
|
|
else:
|
|
# Fallback: inject file content as a paste-style attachment in the message
|
|
log.debug("No file input found; prepending file content to message")
|
|
|
|
return str(tmp_path)
|
|
|
|
|
|
def _wait_for_response_complete(page, timeout_sec: int = 120) -> str:
|
|
"""
|
|
Monitor the chat for response completion.
|
|
Handles auto-continue clicks and waits for typing indicator to disappear.
|
|
Returns the extracted final assistant response text.
|
|
"""
|
|
log.debug("Waiting for response (timeout=%ds)…", timeout_sec)
|
|
start = time.time()
|
|
last_content = ""
|
|
stable_count = 0
|
|
|
|
while time.time() - start < timeout_sec:
|
|
# Dismiss upgrade modals
|
|
_dismiss_upgrade_modal(page)
|
|
|
|
# Auto-continue: click "Continue" if it appears
|
|
continue_btn = page.locator(SEL_CONTINUE_BTN)
|
|
if continue_btn.first.is_visible():
|
|
log.debug("Clicking 'Continue' button…")
|
|
continue_btn.first.click()
|
|
page.wait_for_timeout(1_000)
|
|
stable_count = 0
|
|
continue
|
|
|
|
# Check if typing indicator is present
|
|
typing_visible = page.locator(SEL_TYPING).is_visible()
|
|
|
|
# Check if send button is in "busy" (stop) state
|
|
send_busy = page.locator(".send-btn.busy").is_visible()
|
|
|
|
if typing_visible or send_busy:
|
|
stable_count = 0
|
|
page.wait_for_timeout(500)
|
|
continue
|
|
|
|
# Extract latest assistant message
|
|
msgs = page.locator(f".m-ast {SEL_MSG_BODY}")
|
|
count = msgs.count()
|
|
if count > 0:
|
|
current = msgs.nth(count - 1).inner_text()
|
|
if current == last_content and current.strip():
|
|
stable_count += 1
|
|
else:
|
|
last_content = current
|
|
stable_count = 0
|
|
|
|
# Stable for ~2 seconds → response complete
|
|
if stable_count >= 4:
|
|
log.debug("Response stable — extracting")
|
|
return current.strip()
|
|
|
|
page.wait_for_timeout(500)
|
|
|
|
log.warning("Response extraction timed out after %ds", timeout_sec)
|
|
# Return whatever we have
|
|
msgs = page.locator(f".m-ast {SEL_MSG_BODY}")
|
|
if msgs.count() > 0:
|
|
return msgs.nth(msgs.count() - 1).inner_text().strip()
|
|
return "[TIMEOUT — no response extracted]"
|
|
|
|
|
|
def execute_payload(page, testcase: dict, timeout_sec: int = 120) -> dict:
|
|
"""
|
|
Send a single red-team payload into the active chat and extract the response.
|
|
|
|
Args:
|
|
page: Playwright page with an authenticated session.
|
|
testcase: Dict with keys: id, cat, title, message, file.
|
|
timeout_sec: Max seconds to wait for response.
|
|
|
|
Returns:
|
|
Dict with: id, cat, title, payload, response, elapsed_sec, has_file.
|
|
"""
|
|
tc_id = testcase["id"]
|
|
message = testcase["message"]
|
|
file_spec = testcase.get("file")
|
|
|
|
log.info("▶ Executing %s [%s]: %s", tc_id, testcase["cat"], testcase["title"])
|
|
|
|
# Create a fresh chat for isolation
|
|
_create_new_chat(page)
|
|
_dismiss_upgrade_modal(page)
|
|
|
|
# Handle file attachment for indirect injection tests
|
|
tmp_file = None
|
|
if file_spec:
|
|
tmp_file = _upload_file_for_testcase(page, file_spec)
|
|
# If no file input was available, prepend file content to message
|
|
file_input = page.locator('input[type="file"]')
|
|
if file_input.count() == 0:
|
|
message = (
|
|
f"[File: {file_spec['name']}]\n"
|
|
f"```\n{file_spec['content']}\n```\n\n"
|
|
f"{message}"
|
|
)
|
|
|
|
# Type the message
|
|
textarea = page.locator(SEL_TEXTAREA).first
|
|
textarea.click()
|
|
textarea.fill(message)
|
|
page.wait_for_timeout(300)
|
|
|
|
# Send
|
|
t0 = time.time()
|
|
send_btn = page.locator(SEL_SEND_BTN).first
|
|
send_btn.click()
|
|
page.wait_for_timeout(1_500)
|
|
|
|
# Wait for and extract response
|
|
response_text = _wait_for_response_complete(page, timeout_sec)
|
|
elapsed = round(time.time() - t0, 2)
|
|
|
|
page.wait_for_timeout(1000)
|
|
screenshot_name = f"{tc_id}_{testcase['cat']}_evidence.png"
|
|
screenshot_path = EVIDENCE_DIR / screenshot_name
|
|
page.screenshot(path=str(screenshot_path), full_page=True)
|
|
log.info("Captured evidence screenshot: %s", screenshot_path)
|
|
|
|
log.info("◀ %s completed in %.1fs (%d chars)", tc_id, elapsed, len(response_text))
|
|
|
|
return {
|
|
"id": tc_id,
|
|
"cat": testcase["cat"],
|
|
"title": testcase["title"],
|
|
"payload": testcase["message"],
|
|
"response": response_text,
|
|
"elapsed_sec": elapsed,
|
|
"has_file": file_spec is not None,
|
|
"screenshot_path": str(screenshot_path),
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# High-level orchestration (called by server.py)
|
|
# ---------------------------------------------------------------------------
|
|
def run_full_suite(
|
|
accounts: list[dict],
|
|
testcases: list[dict],
|
|
env: dict,
|
|
headless: bool = True,
|
|
timeout_per_test: int = 120,
|
|
) -> list[dict]:
|
|
"""
|
|
Execute all test cases across available accounts.
|
|
Uses Pro account (token injection) if available, falls back to CSV accounts.
|
|
|
|
Returns list of result dicts.
|
|
"""
|
|
from playwright.sync_api import sync_playwright
|
|
|
|
results = []
|
|
target_url = env.get("TARGET_URL", "https://os.solidpoint.ai")
|
|
pro_token = env.get("PRO_CHAT_TOKEN", "")
|
|
pro_email = env.get("PRO_EMAIL", "")
|
|
pro_password = env.get("PRO_PASSWORD", "")
|
|
|
|
log.info("=" * 60)
|
|
log.info("STARTING RED-TEAM SECURITY SUITE")
|
|
log.info("Target: %s | Tests: %d | Accounts: %d", target_url, len(testcases), len(accounts))
|
|
log.info("=" * 60)
|
|
|
|
with sync_playwright() as pw:
|
|
browser = pw.chromium.launch(headless=headless)
|
|
|
|
# Determine which accounts to cycle through
|
|
auth_methods = []
|
|
|
|
# Priority 1: Pro token
|
|
if pro_token:
|
|
auth_methods.append({"type": "token", "token": pro_token, "label": "PRO (token)"})
|
|
|
|
# Priority 2: Pro credentials
|
|
if pro_email and pro_password:
|
|
auth_methods.append({"type": "creds", "email": pro_email, "password": pro_password, "label": f"PRO ({pro_email})"})
|
|
|
|
# Priority 3: CSV accounts
|
|
for acc in accounts:
|
|
auth_methods.append({"type": "creds", "email": acc["email"], "password": acc["password"], "label": acc["email"]})
|
|
|
|
if not auth_methods:
|
|
log.error("No authentication methods available!")
|
|
return results
|
|
|
|
# Round-robin through accounts for test distribution
|
|
account_idx = 0
|
|
|
|
for i, tc in enumerate(testcases):
|
|
auth = auth_methods[account_idx % len(auth_methods)]
|
|
log.info(
|
|
"--- Test %d/%d [%s] via %s ---",
|
|
i + 1, len(testcases), tc["id"], auth["label"],
|
|
)
|
|
|
|
context = browser.new_context(
|
|
viewport={"width": 1280, "height": 800},
|
|
user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
|
|
)
|
|
page = context.new_page()
|
|
|
|
try:
|
|
# Authenticate
|
|
logged_in = False
|
|
if auth["type"] == "token":
|
|
logged_in = login_with_token(page, auth["token"])
|
|
else:
|
|
logged_in = login_with_credentials(page, auth["email"], auth["password"])
|
|
|
|
if not logged_in:
|
|
log.warning("Auth failed for %s — skipping test %s", auth["label"], tc["id"])
|
|
results.append({
|
|
"id": tc["id"],
|
|
"cat": tc["cat"],
|
|
"title": tc["title"],
|
|
"payload": tc["message"],
|
|
"response": "[AUTH_FAILED]",
|
|
"elapsed_sec": 0,
|
|
"has_file": tc.get("file") is not None,
|
|
"account": auth["label"],
|
|
"error": "Authentication failed",
|
|
})
|
|
# Try next account for subsequent tests
|
|
account_idx += 1
|
|
continue
|
|
|
|
# Execute payload
|
|
result = execute_payload(page, tc, timeout_per_test)
|
|
result["account"] = auth["label"]
|
|
results.append(result)
|
|
|
|
except Exception as exc:
|
|
log.exception("Exception during test %s: %s", tc["id"], exc)
|
|
results.append({
|
|
"id": tc["id"],
|
|
"cat": tc["cat"],
|
|
"title": tc["title"],
|
|
"payload": tc["message"],
|
|
"response": f"[ERROR: {exc}]",
|
|
"elapsed_sec": 0,
|
|
"has_file": tc.get("file") is not None,
|
|
"account": auth["label"],
|
|
"error": str(exc),
|
|
})
|
|
finally:
|
|
context.close()
|
|
|
|
# Rotate account every N tests to distribute load
|
|
if (i + 1) % 5 == 0:
|
|
account_idx += 1
|
|
|
|
browser.close()
|
|
|
|
log.info("=" * 60)
|
|
log.info("SUITE COMPLETE — %d results collected", len(results))
|
|
log.info("=" * 60)
|
|
|
|
return results
|