LocalNetMsg/src/main/remote/terminal.ts

// 远程终端: 控制端发起 open, 被控端 spawn 本地 PTY, 双向字节流 (base64) 走现有 WS
import { EventEmitter } from 'node:events'
import { randomUUID } from 'node:crypto'
// @ts-ignore - @lydell/node-pty 没有自带 .d.ts, 这里只使用其运行时 API
import nodePty from '@lydell/node-pty'
import type { DeviceInfo, TerminalFrame, WsFrame } from '../protocol'
import type { ChatClient } from '../chat-client'
import { defaultShell, defaultCwd } from './shell'
import { shouldAsk, markApproved, enqueueApproval, resetApproval } from './approval'
import { recordAudit } from '../db'

// @lydell/node-pty 1.x 导出: const pty = require('@lydell/node-pty'); pty.spawn(...)
const pty: { spawn: (file: string, args: string[] | string, opts?: any) => any } = (nodePty as any).default ?? nodePty

export interface TerminalManagerEvents {
  // UI 收到这些事件做弹窗 / 显示状态
  approvalRequested: (req: { peerId: string; peerName: string; kind: 'terminal'; detail: string; requestId: string }) => void
  sessionOpened: (sessionId: string, peerId: string, info: { shell: string; rows: number; cols: number }) => void
  sessionClosed: (sessionId: string, peerId: string, reason?: string) => void
  output: (sessionId: string, bytes: Buffer, peerId: string) => void
}

interface Session {
  sessionId: string
  ownerId: string          // 谁开的 (控制端 deviceId)
  peerId: string           // 被控端 deviceId (= self 本机, 当本机是 server)
  proc: any                // IPty from @lydell/node-pty
  readOnly: boolean
  shell: string
  cols: number
  rows: number
  createdAt: number
}

type OpenFrame   = Extract<TerminalFrame, { type: 'open' }>
type InputFrame  = Extract<TerminalFrame, { type: 'input' }>
type ResizeFrame = Extract<TerminalFrame, { type: 'resize' }>
type CloseFrame  = Extract<TerminalFrame, { type: 'close' }>
type AckFrame    = Extract<TerminalFrame, { type: 'ack' }>

export class TerminalManager extends EventEmitter {
  // 作为被控端 (server): 接收对端的 open
  private sessions = new Map<string, Session>()
  // 作为控制端 (client): 自己发起的 session
  private outgoing = new Map<string, { peerId: string; createdAt: number; rows: number; cols: number; shell: string; readOnly: boolean }>()
  private chatClient: ChatClient
  private selfId: string

  constructor(chatClient: ChatClient, selfId: string) {
    super()
    this.chatClient = chatClient
    this.selfId = selfId
  }

  // ===== 控制端 =====

  // 发起一次 open, 等被控端 ack; 被控端同意后立刻开始收 output
  async openOnPeer(peerId: string, opts: { rows: number; cols: number; shell?: string; readOnly?: boolean } = { rows: 24, cols: 80 }): Promise<{ ok: boolean; reason?: string; sessionId?: string }> {
    const sessionId = randomUUID()
    const frame: OpenFrame = { type: 'open', sessionId, rows: opts.rows, cols: opts.cols, ...(opts.shell ? { shell: opts.shell } : {}), readOnly: !!opts.readOnly }
    const sent = this.chatClient.send({ type: 'terminal', payload: frame } as WsFrame, peerId)
    if (!sent) {
      recordAudit({ action: 'terminal.open', target: peerId, result: 'error', payload: { reason: 'ws closed' } })
      return { ok: false, reason: '对方离线' }
    }
    this.outgoing.set(sessionId, {
      peerId,
      createdAt: Date.now(),
      rows: opts.rows,
      cols: opts.cols,
      shell: opts.shell || '',
      readOnly: !!opts.readOnly,
    })
    // 等待 ack: 由 handleIncoming 的 'ack' 分支 resolve
    return new Promise((resolve) => {
      const timer = setTimeout(() => {
        if (this.outgoing.has(sessionId)) {
          this.outgoing.delete(sessionId)
          resolve({ ok: false, reason: '超时 (对方 30s 内未响应)' })
        }
      }, 30_000)
      const onResolve = (ok: boolean, reason?: string) => {
        clearTimeout(timer)
        this.off(`ack:${sessionId}`, onResolve as any)
        resolve({ ok, reason, sessionId })
      }
      this.on(`ack:${sessionId}`, onResolve as any)
    })
  }

  sendInput(sessionId: string, bytes: Buffer | string) {
    const info = this.outgoing.get(sessionId)
    if (!info) return false
    if (info.readOnly) return false
    const b = Buffer.isBuffer(bytes) ? bytes : Buffer.from(bytes, 'utf8')
    const frame: InputFrame = { type: 'input', sessionId, data: b.toString('base64') }
    return this.chatClient.send({ type: 'terminal', payload: frame } as WsFrame, info.peerId)
  }

  resize(sessionId: string, rows: number, cols: number) {
    const info = this.outgoing.get(sessionId)
    if (!info) return false
    info.rows = rows; info.cols = cols
    const frame: ResizeFrame = { type: 'resize', sessionId, rows, cols }
    return this.chatClient.send({ type: 'terminal', payload: frame } as WsFrame, info.peerId)
  }

  close(sessionId: string, reason = 'user') {
    const info = this.outgoing.get(sessionId)
    if (!info) return
    this.outgoing.delete(sessionId)
    const frame: CloseFrame = { type: 'close', sessionId, reason }
    this.chatClient.send({ type: 'terminal', payload: frame } as WsFrame, info.peerId)
  }

  // ===== 被控端 =====

  // 收到来自对端的 terminal 帧
  async handleIncoming(peer: DeviceInfo, f: TerminalFrame) {
    switch (f.type) {
      case 'open':
        return this.handleOpen(peer, f)
      case 'input':
        return this.handleInput(peer, f)
      case 'resize':
        return this.handleResize(peer, f)
      case 'close':
        return this.handleClose(peer, f)
      case 'ack':
        // 控制端收到我们的 open 后, ack 给对端
        this.emit(`ack:${f.sessionId}`, f.ok, f.reason)
        return
      case 'output': {
        // 控制端: 被控端发来的终端输出 -> 交给渲染层
        // (ipc.ts 监听 TerminalManager 的 'output' 事件, 广播 'remote:terminal:output' 给 UI)
        // 之前这里直接 return, 导致控制端永远收不到任何输出 —— 终端窗口一片空白
        if (!f.data) return
        const buf = Buffer.from(f.data, 'base64')
        if (buf.length === 0) return
        this.emit('output', f.sessionId, buf, peer.deviceId)
        return
      }
    }
  }

  private async handleOpen(peer: DeviceInfo, f: OpenFrame) {
    // 1. 授权闸门
    if (shouldAsk(peer.deviceId, 'terminal')) {
      const reqId = randomUUID()
      const reqPromise = enqueueApproval({
        requestId: reqId,
        peerId: peer.deviceId,
        peerName: peer.name,
        kind: 'terminal',
        detail: `对方请求打开终端 (rows=${f.rows} cols=${f.cols}${f.readOnly ? ', 只读' : ''})`,
        ts: Date.now(),
      })
      this.emit('approvalRequested', {
        requestId: reqId, peerId: peer.deviceId, peerName: peer.name, kind: 'terminal',
        detail: `请求打开终端 (rows=${f.rows} cols=${f.cols}${f.readOnly ? ', 只读' : ''})`,
      })
      const verdict = await reqPromise
      if (!verdict.ok) {
        this.sendAck(peer.deviceId, f.sessionId, false, '用户拒绝')
        recordAudit({ action: 'terminal.open', source: peer.deviceId, sessionId: f.sessionId, result: 'denied', payload: { rows: f.rows, cols: f.cols } })
        return
      }
      if (verdict.remember) markApproved(peer.deviceId, 'terminal')
    }

    // 2. spawn pty
    try {
      const shell = defaultShell()
      const proc = pty.spawn(shell.file, shell.args, {
        name: 'xterm-256color',
        cols: f.cols || 80,
        rows: f.rows || 24,
        cwd: defaultCwd(),
        env: { ...process.env, TERM: 'xterm-256color', LANG: process.env.LANG || 'en_US.UTF-8', COLORTERM: 'truecolor' } as any,
        encoding: null,
      })
      const session: Session = {
        sessionId: f.sessionId,
        ownerId: peer.deviceId,
        peerId: peer.deviceId,
        proc,
        readOnly: !!f.readOnly,
        shell: shell.label,
        cols: f.cols || 80,
        rows: f.rows || 24,
        createdAt: Date.now(),
      }
      this.sessions.set(f.sessionId, session)
      this.sendAck(peer.deviceId, f.sessionId, true)
      this.emit('sessionOpened', f.sessionId, peer.deviceId, { shell: shell.label, rows: session.rows, cols: session.cols })

      proc.onData((data: string | Buffer) => {
        // encoding:null 时 data 是 Buffer; encoding:'utf8' 时是 string
        // 我们以 Buffer 形式发给对端, 对端 xterm.write 直接接 Uint8Array
        const buf = Buffer.isBuffer(data) ? data : Buffer.from(data, 'utf8')
        const out: Extract<TerminalFrame, { type: 'output' }> = { type: 'output', sessionId: f.sessionId, data: buf.toString('base64') }
        this.chatClient.send({ type: 'terminal', payload: out } as WsFrame, peer.deviceId)
        this.emit('output', f.sessionId, buf, peer.deviceId)
      })
      proc.onExit(({ exitCode, signal }: { exitCode: number; signal?: number }) => {
        const reason = `退出 (code=${exitCode}${signal ? `, signal=${signal}` : ''})`
        this.sessions.delete(f.sessionId)
        const out: CloseFrame = { type: 'close', sessionId: f.sessionId, reason }
        this.chatClient.send({ type: 'terminal', payload: out } as WsFrame, peer.deviceId)
        this.emit('sessionClosed', f.sessionId, peer.deviceId, reason)
        recordAudit({ action: 'terminal.exit', source: peer.deviceId, sessionId: f.sessionId, result: 'closed', payload: { exitCode, signal } })
      })

      recordAudit({ action: 'terminal.open', source: peer.deviceId, sessionId: f.sessionId, result: 'ok', payload: { shell: shell.label, rows: session.rows, cols: session.cols } })
    } catch (e: any) {
      this.sendAck(peer.deviceId, f.sessionId, false, `spawn 失败: ${e?.message || e}`)
      recordAudit({ action: 'terminal.open', source: peer.deviceId, sessionId: f.sessionId, result: 'error', payload: { error: String(e?.message || e) } })
    }
  }

  private handleInput(_peer: DeviceInfo, f: InputFrame) {
    const s = this.sessions.get(f.sessionId)
    if (!s) return
    if (s.readOnly) return
    try {
      const buf = Buffer.from(f.data, 'base64')
      s.proc.write(buf)
    } catch (e: any) {
      console.warn('[terminal] write failed:', e?.message)
    }
  }

  private handleResize(_peer: DeviceInfo, f: ResizeFrame) {
    const s = this.sessions.get(f.sessionId)
    if (!s) return
    try {
      s.proc.resize(Math.max(2, f.cols), Math.max(2, f.rows))
      s.rows = f.rows; s.cols = f.cols
    } catch (e: any) {
      console.warn('[terminal] resize failed:', e?.message)
    }
  }

  private handleClose(peer: DeviceInfo, f: CloseFrame) {
    const s = this.sessions.get(f.sessionId)
    if (!s) return
    try { s.proc.kill() } catch {}
    this.sessions.delete(f.sessionId)
    this.emit('sessionClosed', f.sessionId, peer.deviceId, f.reason)
    recordAudit({ action: 'terminal.close', source: peer.deviceId, sessionId: f.sessionId, result: 'closed', payload: { reason: f.reason } })
  }

  // 控制端收到被控端的 output -> 转发给 UI
  onRemoteOutput(cb: (sessionId: string, bytes: Buffer, peerId: string) => void) {
    this.on('output', cb)
  }

  private sendAck(peerId: string, sessionId: string, ok: boolean, reason?: string) {
    const frame: AckFrame = { type: 'ack', sessionId, ok, reason }
    this.chatClient.send({ type: 'terminal', payload: frame } as WsFrame, peerId)
  }

  // ===== 维护 =====

  stop() {
    for (const s of this.sessions.values()) {
      try { s.proc.kill() } catch {}
    }
    this.sessions.clear()
    this.outgoing.clear()
  }

  // 被控端 ws 断线 -> 关掉所有属于这个对端的 session
  onPeerDisconnected(peerId: string) {
    for (const [sid, s] of this.sessions) {
      if (s.peerId === peerId) {
        try { s.proc.kill() } catch {}
        this.sessions.delete(sid)
        this.emit('sessionClosed', sid, peerId, 'peer-disconnected')
      }
    }
    // 控制端: 告知 UI (renderer 通过 device:wsState 也会刷新)
    for (const [sid, info] of this.outgoing) {
      if (info.peerId === peerId) {
        this.outgoing.delete(sid)
      }
    }
  }

  // 列出本机作为被控端当前所有 session (UI: 谁正在用我)
  listServerSessions() {
    return Array.from(this.sessions.values()).map(s => ({
      sessionId: s.sessionId,
      peerId: s.peerId,
      shell: s.shell,
      rows: s.rows,
      cols: s.cols,
      readOnly: s.readOnly,
      createdAt: s.createdAt,
    }))
  }

  listClientSessions() {
    return Array.from(this.outgoing.entries()).map(([sid, info]) => ({
      sessionId: sid,
      peerId: info.peerId,
      shell: info.shell,
      rows: info.rows,
      cols: info.cols,
      readOnly: info.readOnly,
      createdAt: info.createdAt,
    }))
  }

  hasClientSession(sessionId: string) {
    return this.outgoing.has(sessionId)
  }

  // 用户在 UI 上点 "拒绝所有未决请求"
  resetPeer(peerId: string, kind?: 'terminal' | 'forward' | 'usb') {
    resetApproval(peerId, kind)
  }
}