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

// 토큰 누출 검사 결과를 담는 구조
export interface TokenLeakResult {
  found: boolean;        // 토큰이 발견되었는지 여부 (true/false)
  location: 'url' | 'body' | 'header';  // 토큰이 발견된 위치 (url, body, header 중 하나)
  title: string;         // 경고 제목
  description: string;   // 상세 설명
  value?: string;        // 실제 발견된 값 (선택적)
}

// 액세스 토큰 누출 검사 클래스
export class AccessTokenLeakController {
  async testReq(sdk: SDK<DefineAPI<{}>>, request: Request): Promise<void> {
    const result = await this._scanRequest(request);
    if (result) {
      await sdk.findings.create({
        title: result.title,
        description: result.description,
        request,
        reporter: "AccessTokenLeak",
      });
    }
  }

  async testResp(sdk: SDK<DefineAPI<{}>>, response: Response, request: Request): Promise<void> {
    const result = await this._scanResponse(response);
    if (result) {
      await sdk.findings.create({
        title: result.title,
        description: result.description,
        request,
        reporter: "AccessTokenLeak",
      });
    }
  }

  /**
   * @param request - 검사할 HTTP 요청 객체
   * @returns 토큰이 발견되면 결과 객체, 없으면 null
   */
  async _scanRequest(request: Request): Promise<TokenLeakResult | null> {
    
    // === 1. URL에서 토큰 검사 ===
    const url = request.getUrl();
    
    const extractedTokenFromUrl = this.extractTokenFromText(url);
    
    if (extractedTokenFromUrl) {
      return {
        found: true,
        location: 'url',
        title: "Token Leak in URL",
        description: `요청 URL에 토큰이 포함되어 있습니다. (토큰: ${extractedTokenFromUrl.substring(0, 20)}...)`,
        value: url
      };
    }

    // === 2. 요청 본문(Body)에서 토큰 검사 ===
    const body = request.getBody();
    
    if (body) {
      const bodyText = await body.toText();
      
      const extractedTokenFromBody = this.extractTokenFromText(bodyText);
      
      if (extractedTokenFromBody) {
        return {
          found: true,
          location: 'body',
          title: "Token Leak in Request Body",
          description: `요청 Body에 토큰이이 포함되어 있습니다. (토큰: ${extractedTokenFromBody.substring(0, 20)}...)`,
          value: bodyText
        };
      }
    }

    return null;
  }

  /**
   * HTTP 응답에서 액세스 토큰 누출 검사
   * @param response - 검사할 HTTP 응답 객체
   * @returns 토큰이 발견되면 결과 객체, 없으면 null
   */
  async _scanResponse(response: Response): Promise<TokenLeakResult | null> {
    
    // === 1. Location 헤더에서 토큰 검사 ===
    const locationHeader = response.getHeader("Location");
    
    const locationHeaderStr = Array.isArray(locationHeader) ? locationHeader.join(', ') : locationHeader;
    
    if (locationHeaderStr) {
      const extractedTokenFromHeader = this.extractTokenFromText(locationHeaderStr);
      
      if (extractedTokenFromHeader) {
        return {
          found: true,
          location: 'header',
          title: "Token Leak in Redirect URL",
          description: `Location 헤더에 토큰이 노출되었습니다: ${locationHeaderStr} (토큰: ${extractedTokenFromHeader.substring(0, 20)}...)`,
          value: locationHeaderStr
        };
      }
    }

    // === 2. 응답 본문에서 토큰 검사 ===
    const bodyBytes = response.getBody();
    
    if (bodyBytes) {
      const bodyText = await bodyBytes.toText();
      
      const extractedTokenFromBody = this.extractTokenFromText(bodyText);
    
      if (extractedTokenFromBody) {
        return {
          found: true,
          location: 'body',
          title: "Token Leak in Response Body",
          description: `HTTP 응답 본문에 토큰이 노출되었습니다. (토큰: ${extractedTokenFromBody.substring(0, 20)}...)`,
          value: bodyText
        };
      }
    }

    return null;
  }

  /**
   * 텍스트에서 실제 토큰 값을 추출
   * @param text - 검사할 텍스트
   * @returns 토큰 값이 있으면 해당 값, 없으면 null
   */
  private extractTokenFromText(text: string): string | null {
     // 토큰 관련 키워드 리스트
    const tokenKeys = [
      'access_token',
      'accesstoken',
      'Access-Token',
      'Refresh_Token',
      'Refresh-Token',
      'RefreshToken',
      'Secret_Token',
      'Secret-Token',
      'SecretToken',
      'SSO_Auth',
      'SSO-Auth',
      'SSOAuth',
      'auth_token',
      'session_token'
    ];

    const tokenTypeKeys = [
      'token_type',
      'tokenType'
    ];

     // 정규표현식 토큰 타입 유무 패턴 리스트 생성
    const tokenTypeRegexes: RegExp[] = [];
    for (const key of tokenTypeKeys) {
      // JSON 형식: "token_type": "Bearer"
      tokenTypeRegexes.push(new RegExp(`"${key}"\\s*:\\s*"bearer"`, 'i'));
      // 일반 key=value 형식: token_type=Bearer
      tokenTypeRegexes.push(new RegExp(`${key}[=:]\\s*bearer`, 'i'));
      // 공백 있는 형식: token_type : Bearer
      tokenTypeRegexes.push(new RegExp(`${key}\\s*:\\s*bearer`, 'i'));
    }

    // token_type=bearer 형태 중 하나라도 포함되는지 확인
    const hasTokenTypeBearer = tokenTypeRegexes.some(rx => rx.test(text));

    // 정규표현식 토큰 유무 패턴 리스트 생성
    const tokenPatterns: RegExp[] = [];

    for (const key of tokenKeys) {
      // 1. key=token 또는 key: token
      tokenPatterns.push(new RegExp(`${key}[=:]\\s*([a-zA-Z0-9\\-._~+/]+=*)`, 'i'));

      // 2. JSON 형태의 "key": "token"
      tokenPatterns.push(new RegExp(`"${key}"\\s*:\\s*"([^"]+)"`, 'i'));
    }

    // 3. Authorization: Bearer <token> 형태
    tokenPatterns.push(/bearer\s+([a-zA-Z0-9\-._~+/]+=*)/i);

    // 모든 패턴에 대해 검사
    for (const pattern of tokenPatterns) {
      const match = pattern.exec(text);
      if (match && match[1]) {
        if(hasTokenTypeBearer){
          return match[1];
        }
      }
    }

    return null;
  }
}