APIがない企業SSOにPlaywrightでサーバーサイド認証プロキシを作った話

import { chromium, type Browser, type BrowserContext } from "playwright";
import { TOTP } from "totp-generator";

export class PlaywrightAuthService {
  private static credentials: LoginCredentials | null = null;
  private browser: Browser | null = null;
  private context: BrowserContext | null = null;

  async initialize() {
    if (!this.browser) {
      this.browser = await chromium.launch({
        headless: false,
        args: ["--no-sandbox", "--disable-setuid-sandbox"]
      });
      this.context = await this.browser.newContext({ locale: "en-US" });
    }
  }

async login(credentials: LoginCredentials): Promise<AuthResult> {
  await this.initialize();
  const page = await this.context!.newPage();

  try {
    await page.goto(GPT_SERVICE_URL);

    // SSOログインボタンをクリック
    await page.getByRole("button", { name: "Login with SSO" }).click();

    // ユーザー名を入力
    await page.getByLabel("Username").fill(credentials.username);
    await page.getByRole("button", { name: "Next" }).click();

    // TOTP MFAの処理
    if (credentials.totpSecret) {
      await page.getByRole("link", { name: "Select to enter a code" }).click();

      const otpInput = await page.getByRole("textbox", {
        name: "Enter verification code"
      });
      const mfaOtp = (await TOTP.generate(credentials.totpSecret)).otp;

      await otpInput.fill(mfaOtp);
      await page.getByRole("button", { name: "Verify" }).click();
    }

    // ログイン完了を待機
    await page.waitForURL(
      (url) => url.toString().includes("/chats") || url.toString().includes("/dashboard"),
      { timeout: 60000 }
    );

    // ブラウザストレージからトークンを抽出
    const storage = await page.context().storageState();
    const tokenStorage = storage.origins
      .find((origin) => origin.origin === GPT_SERVICE_URL)
      ?.localStorage.find((storage) => storage.name === "idp-token-storage");

    if (!tokenStorage?.value) {
      throw new AuthError("Failed to extract token from browser storage", "LOGIN_FAILED");
    }

    const tokenData = JSON.parse(tokenStorage.value);
    const jwt = tokenData.idToken;

    tokenStore.setToken("default", jwt.idToken);

    return {
      accessToken: jwt.idToken,
      userInfo: jwt.claims || {}
    };

  static async authenticateAI(): Promise<JWTTokens> {
    if (!this.credentials) {
      throw new AuthError("No credentials available. Please login first.", "NO_CREDENTIALS");
    }

    const authService = new PlaywrightAuthService();
    const result = await authService.login(this.credentials);

    // JWTをパースして有効期限を取得
    const tokenParts = result.accessToken.split(".");
    const payload = JSON.parse(atob(tokenParts[1]));

    return {
      idToken: result.accessToken,
      expiresAt: payload.exp || Math.floor(Date.now() / 1000) + 3600
    };
  }

  async cleanup() {
    try {
      await this.context?.close();
      await this.browser?.close();
    } catch (error) {
      console.error("Error during cleanup:", error);
    } finally {
      this.context = null;
      this.browser = null;
    }
  }