makelore/electron/api/routes/auth.ts

import type { IncomingMessage, ServerResponse } from 'http';
import { createCipheriv } from 'node:crypto';
import { shell } from 'electron';
import type { HostApiContext } from '../context';
import { parseJsonBody, sendJson } from '../route-utils';
import { proxyAwareFetch } from '../../utils/proxy-fetch';
import { NIANCODE_AUTH_CONFIG } from '../auth-config';
import { WORKS_SQUARE_CONFIG } from '../works-config';
import {
  clearWorksSquareSession,
  storeWorksSquareSession,
  storeWorksSquareSessionFromTokenPayload,
} from '../../services/works-square-session';
import { clearWorksSquareAIGatewayCredential } from '../../services/works-square-ai-gateway';
import { getProviderService } from '../../services/providers/provider-service';
import { NIANCODE_USER_MODEL_ACCOUNT_ID } from '../../../shared/user-model-config';
import { logger } from '../../utils/logger';
import type { WorksSquareTokenPayload } from '../../services/works-square-session';

type AuthClientInput = {
  authBase?: unknown;
  clientId?: unknown;
  clientSecret?: unknown;
};

type PasswordLoginInput = AuthClientInput & {
  username?: unknown;
  password?: unknown;
  code?: unknown;
  randomStr?: unknown;
  scope?: unknown;
};

type RefreshInput = AuthClientInput & {
  refreshToken?: unknown;
};

type LogoutInput = {
  authBase?: unknown;
  accessToken?: unknown;
};

type SessionSyncInput = {
  accessToken?: unknown;
  refreshToken?: unknown;
  tokenType?: unknown;
  expiresAt?: unknown;
};

type DesktopAuthStartPayload = {
  request_id?: unknown;
  device_secret?: unknown;
  authorize_url?: unknown;
  poll_interval_seconds?: unknown;
};

type DesktopAuthTokenPayload = {
  status?: unknown;
  token?: unknown;
};

const DESKTOP_AUTH_TIMEOUT_MS = 10 * 60 * 1000;

function readRequiredString(value: unknown, field: string): string {
  if (typeof value !== 'string' || !value.trim()) {
    throw new Error(`Missing ${field}`);
  }
  return value.trim();
}

function readOptionalString(value: unknown, fallback: string): string {
  return typeof value === 'string' && value.trim() ? value.trim() : fallback;
}

function readOptionalTrimmedString(value: unknown): string | null {
  return typeof value === 'string' && value.trim() ? value.trim() : null;
}

function readOptionalNumber(value: unknown): number | null {
  if (typeof value === 'number' && Number.isFinite(value)) return value;
  if (typeof value === 'string') {
    const parsed = Number(value);
    return Number.isFinite(parsed) ? parsed : null;
  }
  return null;
}

function normalizeAuthBase(value: unknown = NIANCODE_AUTH_CONFIG.gatewayAuthUrl): string {
  const authBase = readOptionalString(value, NIANCODE_AUTH_CONFIG.gatewayAuthUrl).replace(/\/+$/, '');
  if (!/^https?:\/\//i.test(authBase)) {
    throw new Error('authBase must start with http:// or https://');
  }
  return authBase;
}

function normalizeWorksBase(value = WORKS_SQUARE_CONFIG.apiBaseUrl): string {
  const apiBase = value.replace(/\/+$/, '');
  if (!/^https?:\/\//i.test(apiBase)) {
    throw new Error('Works Square API base URL must start with http:// or https://');
  }
  return apiBase;
}

function createWorksUrl(pathname: string): URL {
  return new URL(`${normalizeWorksBase()}${pathname}`);
}

function createBasicAuthHeader(clientId: string, clientSecret: string): string {
  return `Basic ${Buffer.from(`${clientId}:${clientSecret}`).toString('base64')}`;
}

function createTokenRequestBody(params: Record<string, string>): URLSearchParams {
  const body = new URLSearchParams();
  for (const [key, value] of Object.entries(params)) {
    body.set(key, value);
  }
  return body;
}

function encryptPasswordForTokenEndpoint(password: string, encodeKey: string): string {
  const key = Buffer.from(encodeKey, 'utf8');
  if (![16, 24, 32].includes(key.length)) {
    throw new Error('Auth password encode key must be 16, 24, or 32 bytes');
  }
  const cipher = createCipheriv(`aes-${key.length * 8}-cfb`, key, key);
  return Buffer.concat([
    cipher.update(password, 'utf8'),
    cipher.final(),
  ]).toString('base64');
}

function shouldEncryptPasswordForClient(clientId: string): boolean {
  return clientId !== 'customPC';
}

async function readResponsePayload(response: Response): Promise<unknown> {
  const text = await response.text();
  if (!text.trim()) return null;
  try {
    return JSON.parse(text) as unknown;
  } catch {
    return text;
  }
}

function getErrorMessage(payload: unknown, fallback: string): string {
  if (payload && typeof payload === 'object') {
    const record = payload as Record<string, unknown>;
    for (const field of ['msg', 'message', 'error_description', 'error']) {
      const value = record[field];
      if (typeof value === 'string' && value.trim()) {
        return value;
      }
    }
  }
  if (typeof payload === 'string' && payload.trim()) {
    return payload;
  }
  return fallback;
}

function readRequiredPayloadString(
  payload: Record<string, unknown>,
  field: keyof DesktopAuthStartPayload,
): string {
  return readRequiredString(payload[field], field);
}

function readPollIntervalMs(value: unknown): number {
  const seconds = typeof value === 'number'
    ? value
    : (typeof value === 'string' ? Number(value) : 2);
  if (!Number.isFinite(seconds) || seconds < 0) return 2_000;
  return Math.floor(seconds * 1000);
}

function delay(ms: number): Promise<void> {
  if (ms <= 0) return Promise.resolve();
  return new Promise((resolve) => {
    setTimeout(resolve, ms);
  });
}

async function requestToken(
  authBase: string,
  clientId: string,
  clientSecret: string,
  body: URLSearchParams,
): Promise<{ ok: boolean; status: number; payload: unknown }> {
  const response = await proxyAwareFetch(`${authBase}/oauth2/token`, {
    method: 'POST',
    headers: {
      Authorization: createBasicAuthHeader(clientId, clientSecret),
      'Content-Type': 'application/x-www-form-urlencoded',
    },
    body,
  });

  return {
    ok: response.ok,
    status: response.status,
    payload: await readResponsePayload(response),
  };
}

async function pollDesktopAuthToken(
  requestId: string,
  deviceSecret: string,
  pollIntervalMs: number,
): Promise<unknown> {
  const startedAt = Date.now();
  while (Date.now() - startedAt < DESKTOP_AUTH_TIMEOUT_MS) {
    const tokenUrl = createWorksUrl('/api/auth/desktop/token');
    tokenUrl.searchParams.set('request_id', requestId);
    tokenUrl.searchParams.set('device_secret', deviceSecret);

    const response = await proxyAwareFetch(tokenUrl.toString(), { method: 'GET' });
    const payload = await readResponsePayload(response) as DesktopAuthTokenPayload;

    if (!response.ok) {
      throw new Error(getErrorMessage(payload, `Desktop authorization failed (${response.status})`));
    }

    if (payload.status === 'approved' && payload.token) {
      return payload.token;
    }

    await delay(pollIntervalMs);
  }

  throw new Error('Authorization timed out');
}

async function handleBrowserAuthorization(res: ServerResponse): Promise<void> {
  const response = await proxyAwareFetch(createWorksUrl('/api/auth/desktop/start').toString(), {
    method: 'POST',
  });
  const payload = await readResponsePayload(response);

  if (!response.ok) {
    sendJson(res, response.status >= 400 && response.status < 500 ? response.status : 502, {
      success: false,
      error: getErrorMessage(payload, `Desktop authorization start failed (${response.status})`),
    });
    return;
  }

  if (!payload || typeof payload !== 'object' || Array.isArray(payload)) {
    throw new Error('Desktop authorization start returned an invalid payload');
  }

  const record = payload as Record<string, unknown>;
  const requestId = readRequiredPayloadString(record, 'request_id');
  const deviceSecret = readRequiredPayloadString(record, 'device_secret');
  const authorizeUrl = readRequiredPayloadString(record, 'authorize_url');
  const pollIntervalMs = readPollIntervalMs(record.poll_interval_seconds);

  await shell.openExternal(authorizeUrl);
  const token = await pollDesktopAuthToken(requestId, deviceSecret, pollIntervalMs);
  if (token && typeof token === 'object' && !Array.isArray(token)) {
    storeWorksSquareSessionFromTokenPayload(token as WorksSquareTokenPayload);
  }
  sendJson(res, 200, { success: true, token });
}

async function handlePasswordLogin(req: IncomingMessage, res: ServerResponse): Promise<void> {
  const body = await parseJsonBody<PasswordLoginInput>(req);
  const authBase = normalizeAuthBase(body.authBase);
  const clientId = readOptionalString(body.clientId, NIANCODE_AUTH_CONFIG.clientId);
  const clientSecret = readOptionalString(body.clientSecret, NIANCODE_AUTH_CONFIG.clientSecret);
  const username = readRequiredString(body.username, 'username');
  const password = readRequiredString(body.password, 'password');
  const scope = readOptionalString(body.scope, NIANCODE_AUTH_CONFIG.scope);
  const code = readOptionalTrimmedString(body.code);
  const randomStr = readOptionalTrimmedString(body.randomStr);
  const tokenPassword = shouldEncryptPasswordForClient(clientId)
    ? encryptPasswordForTokenEndpoint(password, NIANCODE_AUTH_CONFIG.passwordEncodeKey)
    : password;
  const tokenRequestParams: Record<string, string> = {
    grant_type: 'password',
    username,
    password: tokenPassword,
    scope,
  };

  if (code && randomStr) {
    tokenRequestParams.code = code;
    tokenRequestParams.randomStr = randomStr;
  }

  const tokenResult = await requestToken(
    authBase,
    clientId,
    clientSecret,
    createTokenRequestBody(tokenRequestParams),
  );

  if (!tokenResult.ok) {
    sendJson(res, tokenResult.status === 401 ? 401 : 502, {
      success: false,
      error: getErrorMessage(tokenResult.payload, `Login failed (${tokenResult.status})`),
    });
    return;
  }

  if (
    tokenResult.payload
    && typeof tokenResult.payload === 'object'
    && !Array.isArray(tokenResult.payload)
  ) {
    storeWorksSquareSessionFromTokenPayload(tokenResult.payload as WorksSquareTokenPayload);
  }

  sendJson(res, 200, { success: true, token: tokenResult.payload });
}

async function handleRefresh(req: IncomingMessage, res: ServerResponse): Promise<void> {
  const body = await parseJsonBody<RefreshInput>(req);
  const authBase = normalizeAuthBase(body.authBase);
  const clientId = readOptionalString(body.clientId, NIANCODE_AUTH_CONFIG.clientId);
  const clientSecret = readOptionalString(body.clientSecret, NIANCODE_AUTH_CONFIG.clientSecret);
  const refreshToken = readRequiredString(body.refreshToken, 'refreshToken');

  const tokenResult = await requestToken(
    authBase,
    clientId,
    clientSecret,
    createTokenRequestBody({
      grant_type: 'refresh_token',
      refresh_token: refreshToken,
    }),
  );

  if (!tokenResult.ok) {
    sendJson(res, tokenResult.status === 401 ? 401 : 502, {
      success: false,
      error: getErrorMessage(tokenResult.payload, `Refresh failed (${tokenResult.status})`),
    });
    return;
  }

  if (
    tokenResult.payload
    && typeof tokenResult.payload === 'object'
    && !Array.isArray(tokenResult.payload)
  ) {
    storeWorksSquareSessionFromTokenPayload(tokenResult.payload as WorksSquareTokenPayload, refreshToken);
  }

  sendJson(res, 200, { success: true, token: tokenResult.payload });
}

async function handleSessionSync(req: IncomingMessage, res: ServerResponse): Promise<void> {
  const body = await parseJsonBody<SessionSyncInput>(req);
  storeWorksSquareSession({
    accessToken: readRequiredString(body.accessToken, 'accessToken'),
    refreshToken: readOptionalTrimmedString(body.refreshToken),
    tokenType: readOptionalTrimmedString(body.tokenType),
    expiresAt: readOptionalNumber(body.expiresAt),
  });

  sendJson(res, 200, { success: true });
}

async function clearManagedWorksSquareRuntime(ctx: HostApiContext): Promise<void> {
  const errors: string[] = [];

  try {
    await ctx.imageWorkspace?.closeEventSessions?.();
  } catch (error) {
    errors.push(error instanceof Error ? error.message : String(error));
  }

  try {
    await ctx.opencodeManager.stop();
  } catch (error) {
    errors.push(error instanceof Error ? error.message : String(error));
  }

  clearWorksSquareAIGatewayCredential();

  try {
    await getProviderService().deleteAccountApiKey(NIANCODE_USER_MODEL_ACCOUNT_ID);
  } catch (error) {
    errors.push(error instanceof Error ? error.message : String(error));
  }

  if (errors.length > 0) {
    throw new Error(`Failed to clear managed Works Square runtime state: ${errors.join('; ')}`);
  }
}

async function handleLogout(
  req: IncomingMessage,
  res: ServerResponse,
  ctx: HostApiContext,
): Promise<void> {
  const body = await parseJsonBody<LogoutInput>(req);
  const authBase = normalizeAuthBase(body.authBase);
  const accessToken = readRequiredString(body.accessToken, 'accessToken');

  let cleanupError: Error | null = null;
  try {
    await clearManagedWorksSquareRuntime(ctx);
  } catch (error) {
    cleanupError = error instanceof Error ? error : new Error(String(error));
    logger.error('[auth] Failed to clear managed Works Square runtime state during logout', cleanupError);
  }
  clearWorksSquareSession();

  let response: Response;
  try {
    response = await proxyAwareFetch(`${authBase}/token/logout`, {
      method: 'DELETE',
      headers: {
        Authorization: `Bearer ${accessToken}`,
      },
    });
  } catch (error) {
    if (cleanupError) {
      sendJson(res, 500, {
        success: false,
        error: 'Failed to clear local AI runtime state',
      });
      return;
    }
    throw error;
  }

  if (cleanupError) {
    sendJson(res, 500, {
      success: false,
      error: 'Failed to clear local AI runtime state',
    });
    return;
  }

  if (!response.ok) {
    const payload = await readResponsePayload(response);
    sendJson(res, response.status === 401 ? 401 : 502, {
      success: false,
      error: getErrorMessage(payload, `Logout failed (${response.status})`),
    });
    return;
  }

  sendJson(res, 200, { success: true });
}

export async function handleAuthRoutes(
  req: IncomingMessage,
  res: ServerResponse,
  url: URL,
  ctx: HostApiContext,
): Promise<boolean> {
  if (!url.pathname.startsWith('/api/auth')) {
    return false;
  }

  try {
    if (url.pathname === '/api/auth/browser/start' && req.method === 'POST') {
      await handleBrowserAuthorization(res);
      return true;
    }

    if (url.pathname === '/api/auth/login' && req.method === 'POST') {
      await handlePasswordLogin(req, res);
      return true;
    }

    if (url.pathname === '/api/auth/refresh' && req.method === 'POST') {
      await handleRefresh(req, res);
      return true;
    }

    if (url.pathname === '/api/auth/session/sync' && req.method === 'POST') {
      await handleSessionSync(req, res);
      return true;
    }

    if (url.pathname === '/api/auth/logout' && req.method === 'POST') {
      await handleLogout(req, res, ctx);
      return true;
    }

    sendJson(res, 404, { success: false, error: `No route for ${req.method} ${url.pathname}` });
    return true;
  } catch (error) {
    sendJson(res, 400, {
      success: false,
      error: error instanceof Error ? error.message : String(error),
    });
    return true;
  }
}