it0/packages/services/inventory-service/src/interfaces/rest/controllers/pool-server.controller.ts

import {
  Controller,
  Get,
  Post,
  Put,
  Delete,
  Body,
  Param,
  Query,
  Headers,
  NotFoundException,
  BadRequestException,
  UnauthorizedException,
} from '@nestjs/common';
import * as crypto from 'crypto';
import { ConfigService } from '@nestjs/config';
import { PoolServerRepository } from '../../../infrastructure/repositories/pool-server.repository';
import { CredentialVaultService } from '../../../infrastructure/crypto/credential-vault.service';
import { PoolServer } from '../../../domain/entities/pool-server.entity';

@Controller('api/v1/inventory/pool-servers')
export class PoolServerController {
  private readonly internalApiKey: string;

  constructor(
    private readonly poolServerRepository: PoolServerRepository,
    private readonly vault: CredentialVaultService,
    private readonly configService: ConfigService,
  ) {
    this.internalApiKey = this.configService.get<string>('INTERNAL_API_KEY', '');
  }

  private serialize(server: PoolServer) {
    // Never return raw encrypted key
    const { sshKeyEncrypted, sshKeyIv, ...safe } = server;
    return {
      ...safe,
      hasKey: !!sshKeyEncrypted,
      available: server.status === 'active' && server.currentInstances < server.maxInstances,
      freeSlots: Math.max(0, server.maxInstances - server.currentInstances),
    };
  }

  @Get()
  async list(@Query('available') available?: string) {
    const servers = available === 'true'
      ? await this.poolServerRepository.findAvailable()
      : await this.poolServerRepository.findAll();
    return servers.map((s) => this.serialize(s));
  }

  @Get(':id')
  async getOne(@Param('id') id: string) {
    const server = await this.poolServerRepository.findById(id);
    if (!server) throw new NotFoundException(`Pool server ${id} not found`);
    return this.serialize(server);
  }

  @Post()
  async create(@Body() body: {
    name: string;
    host: string;
    sshPort?: number;
    sshUser: string;
    sshKey: string;        // plaintext private key, encrypted immediately
    maxInstances?: number;
    region?: string;
    notes?: string;
  }) {
    if (!body.name || !body.host || !body.sshUser || !body.sshKey) {
      throw new BadRequestException('name, host, sshUser, sshKey are required');
    }

    const { encrypted, iv } = this.vault.encrypt(body.sshKey);

    const server = new PoolServer();
    server.id = crypto.randomUUID();
    server.name = body.name;
    server.host = body.host;
    server.sshPort = body.sshPort ?? 22;
    server.sshUser = body.sshUser;
    server.sshKeyEncrypted = encrypted.toString('base64');
    server.sshKeyIv = iv.toString('base64');
    server.maxInstances = body.maxInstances ?? 10;
    server.currentInstances = 0;
    server.status = 'active';
    server.region = body.region;
    server.notes = body.notes;

    return this.serialize(await this.poolServerRepository.save(server));
  }

  @Put(':id')
  async update(@Param('id') id: string, @Body() body: {
    name?: string;
    maxInstances?: number;
    status?: 'active' | 'maintenance' | 'offline';
    region?: string;
    notes?: string;
    sshKey?: string;       // optional re-key
  }) {
    const server = await this.poolServerRepository.findById(id);
    if (!server) throw new NotFoundException(`Pool server ${id} not found`);

    if (body.name !== undefined) server.name = body.name;
    if (body.maxInstances !== undefined) server.maxInstances = body.maxInstances;
    if (body.status !== undefined) server.status = body.status;
    if (body.region !== undefined) server.region = body.region;
    if (body.notes !== undefined) server.notes = body.notes;
    if (body.sshKey) {
      const { encrypted, iv } = this.vault.encrypt(body.sshKey);
      server.sshKeyEncrypted = encrypted.toString('base64');
      server.sshKeyIv = iv.toString('base64');
    }

    return this.serialize(await this.poolServerRepository.save(server));
  }

  @Delete(':id')
  async remove(@Param('id') id: string) {
    const server = await this.poolServerRepository.findById(id);
    if (!server) throw new NotFoundException(`Pool server ${id} not found`);
    if (server.currentInstances > 0) {
      throw new BadRequestException(
        `Cannot delete server with ${server.currentInstances} active instance(s). Stop all instances first.`,
      );
    }
    await this.poolServerRepository.remove(server);
    return { message: 'Pool server deleted', id };
  }

  // Internal endpoint: called by agent-service to get decrypted SSH credentials for deployment
  @Get(':id/deploy-creds')
  async deployCreds(
    @Param('id') id: string,
    @Headers('x-internal-api-key') key: string,
  ) {
    if (!this.internalApiKey || key !== this.internalApiKey) {
      throw new UnauthorizedException('Invalid internal API key');
    }
    const server = await this.poolServerRepository.findById(id);
    if (!server) throw new NotFoundException(`Pool server ${id} not found`);

    const sshKey = this.vault.decrypt(
      Buffer.from(server.sshKeyEncrypted, 'base64'),
      Buffer.from(server.sshKeyIv, 'base64'),
    );

    return {
      id: server.id,
      host: server.host,
      sshPort: server.sshPort,
      sshUser: server.sshUser,
      sshKey,
    };
  }

  // Internal endpoint: called by agent-service after deploying / removing an OpenClaw instance
  @Post(':id/increment')
  async increment(@Param('id') id: string) {
    const server = await this.poolServerRepository.findById(id);
    if (!server) throw new NotFoundException(`Pool server ${id} not found`);
    await this.poolServerRepository.incrementInstances(id);
    return { ok: true };
  }

  @Post(':id/decrement')
  async decrement(@Param('id') id: string) {
    const server = await this.poolServerRepository.findById(id);
    if (!server) throw new NotFoundException(`Pool server ${id} not found`);
    await this.poolServerRepository.decrementInstances(id);
    return { ok: true };
  }
}