Initial commit: MCP Red-Team Security Server setup and results
هذا الالتزام موجود في:
ZiadMahmoud2003
492
src/client.py
Normal file
492
src/client.py
Normal file
@@ -0,0 +1,492 @@
|
||||
"""
|
||||
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"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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)
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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
|
||||
المرجع في مشكلة جديدة
حظر مستخدم