wraith/backend/src/terminal/ssh-connection.service.ts

import { Injectable, Logger } from '@nestjs/common';
import { Client, ClientChannel, utils } from 'ssh2';
import { createHash } from 'crypto';
import { CredentialsService } from '../vault/credentials.service';
import { HostsService } from '../connections/hosts.service';
import { PrismaService } from '../prisma/prisma.service';
import { v4 as uuid } from 'uuid';

export interface SshSession {
  id: string;
  hostId: number;
  client: Client;
  stream: ClientChannel | null;
}

@Injectable()
export class SshConnectionService {
  private readonly logger = new Logger(SshConnectionService.name);
  private sessions = new Map<string, SshSession>();

  constructor(
    private credentials: CredentialsService,
    private hosts: HostsService,
    private prisma: PrismaService,
  ) {}

  async connect(
    hostId: number,
    onData: (data: string) => void,
    onClose: (reason: string) => void,
    onHostKeyVerify: (fingerprint: string, isNew: boolean) => Promise<boolean>,
  ): Promise<string> {
    const host = await this.hosts.findOne(hostId);
    const cred = host.credentialId
      ? await this.credentials.decryptForConnection(host.credentialId)
      : null;

    const sessionId = uuid();
    const client = new Client();

    return new Promise((resolve, reject) => {
      client.on('ready', () => {
        client.shell({ term: 'xterm-256color' }, (err, stream) => {
          if (err) {
            client.end();
            return reject(err);
          }
          const session: SshSession = { id: sessionId, hostId, client, stream };
          this.sessions.set(sessionId, session);

          stream.on('data', (data: Buffer) => onData(data.toString('utf-8')));
          stream.on('close', () => {
            this.disconnect(sessionId);
            onClose('Session ended');
          });

          // Update lastConnectedAt and create connection log
          this.hosts.touchLastConnected(hostId);
          this.prisma.connectionLog.create({
            data: { hostId, protocol: host.protocol },
          }).catch(() => {});

          resolve(sessionId);
        });
      });

      client.on('error', (err) => {
        this.logger.error(`SSH error for host ${hostId}: ${err.message}`);
        this.disconnect(sessionId);
        onClose(err.message);
        reject(err);
      });

      const connectConfig: any = {
        host: host.hostname,
        port: host.port,
        username: cred?.username || 'root',
        debug: (msg: string) => {
          if (msg.includes('auth') || msg.includes('Auth') || msg.includes('key') || msg.includes('Key')) {
            this.logger.log(`[SSH-DEBUG] ${msg}`);
          }
        },
        hostVerifier: (key: Buffer, verify: (accept: boolean) => void) => {
          const fingerprint = createHash('sha256').update(key).digest('base64');
          const fp = `SHA256:${fingerprint}`;

          if (host.hostFingerprint === fp) {
            verify(true); // known host — accept silently
            return;
          }

          // Unknown or changed fingerprint — ask the user via WebSocket
          const isNew = !host.hostFingerprint;
          onHostKeyVerify(fp, isNew).then((accepted) => {
            if (accepted) {
              // Persist fingerprint so future connections auto-accept
              this.prisma.host.update({
                where: { id: hostId },
                data: { hostFingerprint: fp },
              }).catch(() => {});
            }
            verify(accepted);
          });
        },
      };

      if (cred?.sshKey) {
        connectConfig.privateKey = cred.sshKey.privateKey;
        if (cred.sshKey.passphrase) {
          connectConfig.passphrase = cred.sshKey.passphrase;
        }
        this.logger.log(`[SSH] Using key auth for ${connectConfig.username}@${connectConfig.host}:${connectConfig.port}`);
        this.logger.log(`[SSH] Key starts with: ${cred.sshKey.privateKey.substring(0, 40)}...`);
        this.logger.log(`[SSH] Key length: ${cred.sshKey.privateKey.length}, has passphrase: ${!!cred.sshKey.passphrase}`);

        // Verify ssh2 can parse the key
        try {
          const parsed = utils.parseKey(cred.sshKey.privateKey, cred.sshKey.passphrase || undefined);
          if (parsed instanceof Error) {
            this.logger.error(`[SSH] Key parse FAILED: ${parsed.message}`);
          } else {
            const keyInfo = Array.isArray(parsed) ? parsed[0] : parsed;
            this.logger.log(`[SSH] Key parsed OK — type: ${keyInfo.type}, comment: ${keyInfo.comment || 'none'}`);
            this.logger.log(`[SSH] Public key fingerprint: ${keyInfo.getPublicSSH?.()?.toString('base64')?.substring(0, 40) || 'N/A'}`);
          }
        } catch (e: any) {
          this.logger.error(`[SSH] Key parse threw: ${e.message}`);
        }
      } else if (cred?.password) {
        connectConfig.password = cred.password;
        this.logger.log(`[SSH] Using password auth for ${connectConfig.username}@${connectConfig.host}:${connectConfig.port}`);
      } else {
        this.logger.warn(`[SSH] No auth method available for host ${hostId}`);
      }

      client.connect(connectConfig);
    });
  }

  write(sessionId: string, data: string) {
    const session = this.sessions.get(sessionId);
    if (session?.stream) {
      session.stream.write(data);
    }
  }

  resize(sessionId: string, cols: number, rows: number) {
    const session = this.sessions.get(sessionId);
    if (session?.stream) {
      session.stream.setWindow(rows, cols, 0, 0);
    }
  }

  disconnect(sessionId: string) {
    const session = this.sessions.get(sessionId);
    if (session) {
      session.stream?.close();
      session.client.end();
      this.sessions.delete(sessionId);
      this.sftpChannels.delete(sessionId);

      // Update connection log with disconnect time
      this.prisma.connectionLog.updateMany({
        where: { hostId: session.hostId, disconnectedAt: null },
        data: { disconnectedAt: new Date() },
      }).catch(() => {});
    }
  }

  getSession(sessionId: string): SshSession | undefined {
    return this.sessions.get(sessionId);
  }

  private sftpChannels = new Map<string, any>();

  getSftpChannel(sessionId: string): Promise<any> {
    const logger = this.logger;
    const cached = this.sftpChannels.get(sessionId);
    if (cached) {
      return Promise.resolve(cached);
    }
    return new Promise((resolve, reject) => {
      const session = this.sessions.get(sessionId);
      if (!session) {
        logger.error(`[SFTP] Session ${sessionId} not found in sessions map (${this.sessions.size} active sessions)`);
        return reject(new Error('Session not found'));
      }
      logger.log(`[SFTP] Requesting SFTP subsystem on session ${sessionId}`);
      session.client.sftp((err, sftp) => {
        if (err) {
          logger.error(`[SFTP] client.sftp() callback error: ${err.message}`);
          return reject(err);
        }
        logger.log(`[SFTP] SFTP subsystem opened successfully for session ${sessionId}`);
        sftp.on('close', () => {
          this.sftpChannels.delete(sessionId);
          logger.log(`[SFTP] Channel closed for session ${sessionId}`);
        });
        this.sftpChannels.set(sessionId, sftp);
        resolve(sftp);
      });
    });
  }
}