import type { SDK } from "caido:plugin";
import { Body, RequestSpec, type Request } from "caido:utils";

export class PKCECheck {
  async test(sdk: SDK, req: Request): Promise<boolean> {
    const method = req.getMethod();
    if (method !== "GET") {
      sdk.console.log("[PKCEDowngradeCheck] Not a GET request. Skipping.");
      return false;
    }

    const query = req.getQuery();
    const searchParams = new URLSearchParams(query);
    const requiredKeys = ["client_id", "response_type", "code_challenge", "code_challenge_method"];

    if (!requiredKeys.every((key) => searchParams.has(key))) {
      sdk.console.log("[PKCEDowngradeCheck] Required PKCE parameters missing. Skipping.");
      return false;
    }

    const url = req.getUrl();
    const isOpenID = searchParams.get("scope")?.includes("openid") || url.includes("id_token");
    const methodVal = searchParams.get("code_challenge_method");
    const challengeVal = searchParams.get("code_challenge");

    if (!methodVal || !challengeVal) {
      sdk.console.log("[PKCEDowngradeCheck] code_challenge or method missing. Skipping.");
      await sdk.findings.create({
        title: isOpenID
          ? "[WARN] OpenID Flow PKCE Parameters Missing"
          : "[WARN] OAuth2 Flow PKCE Parameters Missing",
        description: `PKCE parameters are missing or incomplete for ${url}. This may indicate a misconfiguration.`,
        request: req,
        reporter: "PKCE Checker",
      });
      return false;
    }

    if (methodVal === "plain") {
      sdk.console.log("[PKCEDowngradeCheck] code_challenge_method is 'plain'. Skipping.");
      await sdk.findings.create({
        title: isOpenID
          ? "[WARN] OpenID Flow PKCE Method is 'plain'"
          : "[WARN] OAuth2 Flow PKCE Method is 'plain'",
        description: `PKCE method is set to 'plain' for ${url}. This may indicate a downgrade vulnerability.`,
        request: req,
        reporter: "PKCE Checker",
      });
      return false;
    }

    // Remove PKCE parameters to simulate a downgraded request
    searchParams.delete("code_challenge");
    searchParams.delete("code_challenge_method");
    const downgradedQuery = searchParams.toString();
    const scheme = req.getUrl().startsWith("https") ? "https" : "http";
    const downgradedUrl = `${scheme}://${req.getHost()}:${req.getPort()}${req.getPath()}?${downgradedQuery}`;

    sdk.console.log(`${req.getHost()} Original URL: ` + url);
    sdk.console.log(`${req.getHost()} Downgraded URL: ` + downgradedUrl);

    try {
      // Use Caido Replay SDK to replay the original request
      const spec = new RequestSpec(downgradedUrl);
      spec.setBody(req.getBody() as Body);
      for (const [key, value] of Object.entries(req.getHeaders())) {
        if (Array.isArray(value)) {
          spec.setHeader(key, value.join(', ')); // or another suitable delimiter
        } else {
          spec.setHeader(key, value);
        }
      }
      spec.setHost(req.getHost());
      spec.setMethod(req.getMethod());
      spec.setPath(req.getPath());
      spec.setQuery(downgradedQuery);
      spec.setTls(req.getTls());
      spec.setPort(req.getPort());

      let sendDowngradedRequest = await sdk.requests.send(spec);

      if (sendDowngradedRequest.response) {
        let domain = spec.getHost();
        let port = spec.getPort();
        let path = spec.getPath();
        let query = spec.getQuery();
        let id = sendDowngradedRequest.response.getId();
        let code = sendDowngradedRequest.response.getCode();
        sdk.console.log(`REQ ${id}: ${domain}:${port}${path}${query} received a status code of ${code}`);
      }

      if (sendDowngradedRequest.response?.getCode() === 302) {
        await sdk.findings.create({
          title: isOpenID
            ? "[CRITICAL] OpenID Flow PKCE Downgrade Vulnerability"
            : "[CRITICAL] OAuth2 Flow PKCE Downgrade Vulnerability",
          description: `The request to ${url} is vulnerable to a PKCE downgrade attack. This may indicate a configuration error.`,
          request: req,
          reporter: "PKCE Checker",
        });
      }

/*
      sdk.console.log(`${req.getHost()} Original Status: ` + resOriginal.status);
      sdk.console.log(`${req.getHost()} Downgraded Status: ` + resDowngraded.status);

      sdk.console.log(`${req.getHost()} Original Headers: ` + JSON.stringify(resOriginal.headers));
      sdk.console.log(`${req.getHost()} Downgraded Headers: ` + JSON.stringify(resDowngraded.headers));

      // Caido Dev Docs 기준으로, 리다이렉트된 URL은 Response 객체의 url 속성에 저장되어 있음
      const locationOriginal = resOriginal.url ?? "";
      const locationDowngraded = resDowngraded.url ?? "";

      sdk.console.log(`${req.getHost()} Original Location: ` + locationOriginal);
      sdk.console.log(`${req.getHost()} Downgraded Location: ` + locationDowngraded);

      const statusEqual = resOriginal.status === resDowngraded.status;
      const codeInRedirects = locationOriginal.includes("code=") && locationDowngraded.includes("code=");

      if (statusEqual && codeInRedirects) {
        const title = isOpenID
          ? "[CRITICAL] OpenID Flow PKCE Downgraded to Plaintext"
          : "[CRITICAL] OAuth2 Flow PKCE Downgraded to Plaintext";
        const reference = isOpenID
          ? "https://openid.net/specs/openid-igov-oauth2-1_0-02.html#rfc.section.3.1.7"
          : "https://datatracker.ietf.org/doc/html/rfc7636";

        await sdk.findings.create({
          title,
          description: `PKCE downgrade detected for ${url}.\n\nDowngraded URL: ${downgradedUrl}\n\nRedirect contained code=.\n\nReference: ${reference}`,
          request: req,
          reporter: "",
        });

        return true;
      }*/
    } catch (err) {
      sdk.console.error(`PKCE downgrade check failed for ${url}: ${String(err)}`);
    }

    sdk.console.log("[PKCEDowngradeCheck] No PKCE downgrade detected.");
    return false;
  }
}
