aurak/server/src/auth/combined-auth.guard.ts

import {
  Injectable,
  CanActivate,
  ExecutionContext,
  UnauthorizedException,
} from '@nestjs/common';
import { Reflector } from '@nestjs/core';
import { AuthGuard } from '@nestjs/passport';
import { UserService } from '../user/user.service';
import { Request } from 'express';
import { lastValueFrom, Observable } from 'rxjs';
import { IS_PUBLIC_KEY } from './public.decorator';
import { tenantStore } from '../tenant/tenant.store';
import { UserRole } from '../user/user-role.enum';
import { TenantService } from '../tenant/tenant.service';
import * as fs from 'fs';
import * as path from 'path';

/**
 * A combined authentication guard that accepts either:
 *  1. An API key via the `x-api-key` header (or `Authorization: Bearer kb_...`)
 *  2. A standard JWT Bearer token
 *
 * This replaces JwtAuthGuard on routes that should support both auth methods.
 */
@Injectable()
export class CombinedAuthGuard implements CanActivate {
  // We extend AuthGuard('jwt') functionality by composition
  private jwtGuard: ReturnType<typeof AuthGuard>;

  constructor(
    private reflector: Reflector,
    private userService: UserService,
    private tenantService: TenantService,
  ) {
    // Create a JWT guard instance
    const JwtGuardClass = AuthGuard('jwt');
    this.jwtGuard = new JwtGuardClass() as any;
  }

  async canActivate(context: ExecutionContext): Promise<boolean> {
    // Allow @Public() decorated routes
    const isPublic = this.reflector.getAllAndOverride<boolean>(IS_PUBLIC_KEY, [
      context.getHandler(),
      context.getClass(),
    ]);

    const request = context
      .switchToHttp()
      .getRequest<Request & { user?: any; tenantId?: string }>();
    const logMsg = `\n[${new Date().toISOString()}] AuthGuard: ${request.method} ${request.url} (isPublic: ${isPublic})\n`;
    fs.appendFileSync('auth_debug.log', logMsg);

    if (isPublic) {
      return true;
    }

    console.log(
      `[CombinedAuthGuard] Checking auth for route: ${request.method} ${request.url}`,
    );

    // --- Try API Key first ---
    const apiKey = this.extractApiKey(request);
    if (apiKey) {
      const user = await this.userService.findByApiKey(apiKey);
      if (user) {
        // If x-tenant-id is provided, verify membership
        const requestedTenantId = request.headers['x-tenant-id'] as string;
        let activeTenantId = user.tenantId;

        if (requestedTenantId) {
          const memberships = await this.userService.getUserTenants(user.id);
          const hasAccess = memberships.some(
            (m) => m.tenantId === requestedTenantId,
          );

          if (hasAccess || user.isAdmin) {
            activeTenantId = requestedTenantId;
          } else {
            throw new UnauthorizedException(
              'User does not belong to the requested tenant',
            );
          }
        }

        const role = await this.tenantService.getUserRole(
          user.id,
          activeTenantId,
        );

        request.user = {
          id: user.id,
          username: user.username,
          role,
          tenantId: activeTenantId,
        };
        request.tenantId = activeTenantId;

        // Update tenant context store
        const store = tenantStore.getStore();
        if (store) {
          store.tenantId = activeTenantId;
          store.userId = user.id;
        }

        return true;
      }
      throw new UnauthorizedException('Invalid API key');
    }

    // --- Fall back to JWT ---
    try {
      const result = await (this.jwtGuard as any).canActivate(context);
      let hasJwtSession = false;

      if (result instanceof Observable) {
        hasJwtSession = await lastValueFrom(result);
      } else {
        hasJwtSession = result;
      }

      if (hasJwtSession) {
        const user = request.user;
        if (!user) return false;

        const requestedTenantId = request.headers['x-tenant-id'] as string;

        if (requestedTenantId && user.tenantId !== requestedTenantId) {
          const memberships = await this.userService.getUserTenants(user.id);
          const hasAccess = memberships.some(
            (m) => m.tenantId === requestedTenantId,
          );

          if (hasAccess || user.isAdmin) {
            user.tenantId = requestedTenantId;
          } else {
            throw new UnauthorizedException(
              'User does not belong to the requested tenant',
            );
          }
        }

        // Fetch the role for the active tenant
        const role = await this.tenantService.getUserRole(
          user.id,
          user.tenantId,
        );
        user.role = role;

        request.tenantId = user.tenantId;

        // Update tenant context store
        const store = tenantStore.getStore();
        if (store) {
          store.tenantId = user.tenantId;
          store.userId = user.id;
        }

        return true;
      }
      return false;
    } catch (e) {
      console.error(`[CombinedAuthGuard] JWT Auth Error:`, e);
      throw e instanceof UnauthorizedException
        ? e
        : new UnauthorizedException('Authentication required');
    }
  }

  private extractApiKey(request: Request): string | undefined {
    // Allow `Authorization: Bearer kb_...` form
    const authHeader = request.headers.authorization;
    if (authHeader?.startsWith('Bearer kb_')) {
      return authHeader.substring(7);
    }
    // Or a plain `x-api-key` header
    const headerKey = request.headers['x-api-key'] as string;
    if (headerKey) return headerKey;
    return undefined;
  }
}