import type { Request, Response } from "caido:utils";
import type { SDK, DefineAPI } from "caido:plugin";
import { HttpUtils } from "../utils/http";

const httpUtils = new HttpUtils();

export class CsrfCheck {
  private nonceParam = [
    "state",
    "nonce",
    "as",
    "frame_id",
    "csrf_token",
    "csrf",
  ];

  private isTargetUri(uri: string): boolean {
    if (
      httpUtils.getQueryParamFromURI(uri, "client_id") !== null &&
      (httpUtils.getQueryParamFromURI(uri, "response_type") !== null ||
        httpUtils.getQueryParamFromURI(uri, "grant_type") !== null ||
        httpUtils.getQueryParamFromURI(uri, "redirect_uri") !== null ||
        httpUtils.getQueryParamFromURI(uri, "scope") !== null ||
        httpUtils.getQueryParamFromURI(uri, "state") !== null ||
        httpUtils.getQueryParamFromURI(uri, "nonce") !== null)
    ) {
      return true;
    }

    return false;
  }

  private isOauthUri(request: Request): boolean {
    const uri = request.getUrl() || "";

    // Check if the request is an OAuth authorization request
    if (this.isTargetUri(uri)) {
      return true;
    }

    return false;
  }

  private isOauthRedirectResponse(response: Response): boolean {
    const status = response.getCode();
    const uri =
      httpUtils.getHeaderValue(response.getHeaders(), "location") || "";

    if (status >= 300 && status < 400 && this.isTargetUri(uri)) {
      return true;
    }
    return false;
  }

  private isNonceInQuery(request: Request): boolean {
    const query = request.getQuery() || "";

    for (const param of this.nonceParam) {
      if (httpUtils.getQueryParam(query, param) !== null) {
        return true; // Nonce parameter is present in the query
      }
    }

    return false; // No nonce parameter found in the query
  }

  private getNonceParamName(url: string): string | null {
    for (const param of this.nonceParam) {
      if (httpUtils.getQueryParamFromURI(url, param) !== null) {
        return param; // Return the first matching nonce parameter
      }
    }

    return null; // No nonce parameter found
  }

  private checkNonceAtResponseLocationHeader(
    request: Request,
    response: Response
  ): string[] | 0 {
    const nonceParamName = this.getNonceParamName(request.getUrl() || "");

    if (
      !this.isOauthUri(request) ||
      !this.isNonceInQuery(request) ||
      !this.isOauthRedirectResponse(response) ||
      !nonceParamName
    ) {
      return 0; // Not a target, no CSRF risk
    }

    // 요청에서 보낸 Nonce 추출
    const query = request.getQuery() || "";
    const originalNonce = httpUtils.getQueryParam(query, nonceParamName);

    // 리다이렉트 URL에서 쿼리 부분만 추출
    const locationHeader =
      httpUtils.getHeaderValue(response.getHeaders(), "location") || "";

    const responseNonce = httpUtils.getQueryParamFromURI(
      locationHeader || "",
      nonceParamName
    );

    // Nonce가 없거나, 요청값과 다르면 CSRF 위험
    if (!responseNonce) {
      // missing state
      return ["Nonce parameter is missing in the response location header"];
    }
    if (originalNonce !== responseNonce) {
      // mismatch
      return ["Nonce parameter mismatch between request and response"];
    }

    return 0; // no CSRF risk detected
  }

  private async checkNonceReuse(
    sdk: SDK<DefineAPI<{}>, {}>,
    request: Request,
    originResponse: Response
  ): Promise<string[] | 0> {
    // uri에 oauth 관련 파라미터가 없지만, 응답이 oauth 리다이렉트 응답인지 확인
    // 즉, 처음으로 Nonce를 발급한 요청인지 확인
    if (
      this.isOauthUri(request) ||
      !this.isOauthRedirectResponse(originResponse)
    ) {
      return 0; // Not a target, no CSRF risk
    }

    // 기존 응답의 location 헤더의 url에서 Nonce 파라미터 이름, nonce 파라미터 값, 쿼리 추출
    const originResponseLocationHeader =
      httpUtils.getHeaderValue(originResponse.getHeaders(), "location") || "";
    const nonceParamName =
      this.getNonceParamName(originResponseLocationHeader || "") || "state";
    const originLocationQuery =
      httpUtils.getQueryFromURI(originResponseLocationHeader || "") || "";
    const originLocationNonce = httpUtils.getQueryParam(
      originLocationQuery,
      nonceParamName
    );

    // 쿠키가 없는 헤더로 새로운 nonce를 발급받기 위해 요청
    const noCookieHeaders = httpUtils.removeHeaders(request.getHeaders(), [
      "cookie",
    ]);
    const noCookieResponse = await httpUtils.resend(sdk, request, {
      headers: noCookieHeaders,
    });
    if (!noCookieResponse || noCookieResponse?.getCode() >= 400) {
      return 0;
    }

    // 쿠키가 없는 응답의 location 헤더 추출 및 Nonce 추출
    const noCookieLocationHeader = httpUtils.getHeaderValue(
      noCookieResponse?.getHeaders() || {},
      "location"
    );
    const newNonce =
      httpUtils.getQueryParamFromURI(
        noCookieLocationHeader || "",
        nonceParamName
      ) || "";

    if (originLocationNonce === newNonce) {
      return [
        "State parameter reused in the response location header, indicating a potential CSRF risk",
      ];
    }

    // 기존 쿠키와 함께 새로운 Nonce로 요청
    const newQuery = httpUtils.setQueryParam(
      originLocationQuery,
      nonceParamName,
      newNonce
    );

    // 기존 location 헤더의 uri 요청과 location 헤더에서 nonce값만 새로 발급한 값으로 바꾸어 요청한 결과를 비교
    const res1 = await httpUtils.customFetch(
      sdk,
      originResponseLocationHeader,
      "GET",
      originLocationQuery,
      request.getHeaders()
    );

    const res2 = await httpUtils.customFetch(
      sdk,
      originResponseLocationHeader,
      "GET",
      newQuery,
      request.getHeaders()
    );

    if (
      !res1 ||
      !res2 ||
      res1.getCode() >= 400 ||
      res2.getCode() >= 400 ||
      res1.getCode() !== res2.getCode()
    ) {
      return 0;
    }

    if (
      res1.getCode() === res2.getCode() &&
      300 <= res1.getCode() &&
      res1.getCode() < 400
    ) {
      const res1LocationHeader =
        httpUtils.getHeaderValue(res1.getHeaders(), "location") || "";
      const res2LocationHeader =
        httpUtils.getHeaderValue(res2.getHeaders(), "location") || "";
      const res1ReirectPath = httpUtils.getPathFromURI(res1LocationHeader);
      const res2ReirectPath = httpUtils.getPathFromURI(res2LocationHeader);

      if (res1ReirectPath === res2ReirectPath) {
        return [
          "When nonce parameter reused in the response location header, it might not be verified. Indicating a potential CSRF risk",
        ];
      }
    }

    return 0; // no CSRF risk detected
  }

  async checker(
    sdk: SDK<DefineAPI<{}>, {}>,
    request: Request,
    response: Response
  ): Promise<void> {
    let result = ``;

    // 쿼리에 state 파라미터가 없으면 CSRF 위험
    try {
      if (this.isOauthUri(request) && !this.isNonceInQuery(request)) {
        result += "CSRF risk: missing state parameter"; // CSRF risk: missing state parameter
      }
    } catch (error) {
      sdk.console.error(`Error checking state in query: ${error}`);
    }

    // location 헤더에 state 파라미터가 없거나, 요청에서 보낸 state와 다르면 CSRF 위험
    try {
      const stateAtResponseLocationHeaderCheck =
        this.checkNonceAtResponseLocationHeader(request, response);
      if (stateAtResponseLocationHeaderCheck !== 0) {
        result += `, ${stateAtResponseLocationHeaderCheck.join(", ")}`;
      }
    } catch (error) {
      sdk.console.error(
        `Error checking state in response location header: ${error}`
      );
    }

    // 처음으로 state를 발급한 요청에서 state 파라미터를 바꿔서 보내기
    const reusedStateCheck = await this.checkNonceReuse(sdk, request, response);
    if (reusedStateCheck !== 0) {
      result += `, ${reusedStateCheck.join(", ")}`;
    }

    try {
      result.replace(/^\s*,\s*|\s*$/, ""); // Remove leading/trailing commas
      if (result) {
        await sdk.findings.create({
          title: "csrf vuln",
          description: `SSO-related parameters detected in response:\n\n${request.getMethod()} ${request.getUrl()} : ${result}`,
          request,
          reporter: "csrf reporter",
        });
      }
    } catch (error) {
      sdk.console.error(`Error creating finding: ${error}`);
    }
  }
}