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 isOauthUri(request: Request): boolean {
    const query = request.getQuery() || "";

    // Check if the request is an OAuth authorization request
    if (
      query.includes("client_id=") &&
      (query.includes("response_type=") ||
        query.includes("grant_type=") ||
        query.includes("redirect_uri=") ||
        query.includes("scope=") ||
        query.includes("state="))
    ) {
      return true;
    }

    return false;
  }

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

    if (
      status >= 300 &&
      status < 400 &&
      locationHeader &&
      (locationHeader.includes("client_id=") ||
        locationHeader.includes("response_type=") ||
        locationHeader.includes("grant_type=") ||
        locationHeader.includes("redirect_uri=") ||
        locationHeader.includes("scope=") ||
        locationHeader.includes("state=") ||
        locationHeader.includes("code=")) // code is also common in OAuth redirects
    ) {
      return true;
    }
    return false;
  }

  private isStateInQuery(request: Request): boolean {
    const query = request.getQuery();
    const stateValue = httpUtils.getQueryParam(query || "", "state");
    if (!stateValue) {
      return false;
    }
    return true;
  }

  private checkStateAtResponseLocationHeader(
    request: Request,
    response: Response
  ): string[] | 0 {
    if (
      !(
        this.isOauthUri(request) &&
        this.isStateInQuery(request) &&
        this.isOauthRedirectResponse(response)
      )
    ) {
      return 0; // Not a target, no CSRF risk
    }

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

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

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

    return 0; // no CSRF risk detected
  }

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

  //   const originResponseLocationHeader = httpUtils.getHeaderValue(
  //     originResponse.getHeaders(),
  //     "location"
  //   );
  //   const originState = httpUtils.getQueryParamFromURI(
  //     originResponseLocationHeader || "",
  //     "state"
  //   );

  //   const requestHeaders = request.getHeaders();
  //   const noCookieHeaders = httpUtils.removeHeaders(requestHeaders, ["cookie"]);
  //   const newResponse = await httpUtils.resend(request, {
  //     headers: noCookieHeaders,
  //   });
  //   const newLocationHeader = httpUtils.getHeaderValue(
  //     newResponse.getHeaders(),
  //     "location"
  //   );
  //   const newState = httpUtils.getQueryParamFromURI(
  //     newLocationHeader || "",
  //     "state"
  //   );

  //   if (originState === newState) {
  //     return [
  //       "State parameter reused in the response location header, 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 위험
    if (this.isOauthUri(request) && !this.isStateInQuery(request)) {
      result += "CSRF risk: missing state parameter"; // CSRF risk: missing state parameter
    }

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

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

    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",
      });
    }
  }
}