opsxshell/electron/services/zmodem-service.ts

import Zmodem from 'zmodem-ts'
import fs from 'fs'
import path from 'path'
import { app } from 'electron'
import { createLogger } from '../utils/logger'

const logger = createLogger('ZmodemService')

// Type definitions for Zmodem sessions (since zmodem-ts doesn't export them properly)
interface ZmodemSendSession {
  type: 'send'
  set_sender(func: (octets: number[]) => void): void
  send_offer(params: { name: string; size: number; mtime?: number }): Promise<any>
  on(event: string, handler: (data: any) => void): void
  abort(): void
  close(): Promise<void>
}

interface ZmodemReceiveSession {
  type: 'receive'
  set_sender(func: (octets: number[]) => void): void
  start(): Promise<any>
  on(event: string, handler: (data: any) => void): void
  abort(): void
}

interface ZmodemSentry {
  consume(octets: number[]): void
}

export interface ZmodemProgress {
  percentage: number
  transferred: number
  total: number
  fileName?: string
}

interface ZmodemSessionInfo {
  id: string
  shellId: string
  type: 'send' | 'receive'
  session: ZmodemSendSession | ZmodemReceiveSession
  fileName?: string
  filePath?: string
  totalSize?: number
  transferred?: number
  sentry?: ZmodemSentry
  onComplete?: (success: boolean, error?: string) => void
}

export class ZmodemService {
  private sessions: Map<string, ZmodemSessionInfo> = new Map()
  private shellDataCallbacks: Map<string, (data: Buffer) => void> = new Map()

  constructor() {
    app.on('before-quit', () => {
      this.closeAll()
    })
  }

  /**
   * Send file to server via Zmodem (rz)
   *
   * Flow: Server types "rz" -> sends ZRQINIT/ZRINIT -> we send file
   * We use Sentry to detect when server is ready (ZRINIT received)
   */
  sendFile(
    shellId: string,
    fileName: string,
    fileData: Buffer,
    writeFunc: (data: Uint8Array) => void,
    onProgress?: (progress: ZmodemProgress) => void,
    onComplete?: (success: boolean, error?: string) => void
  ): string {
    const sessionId = `zmodem-send-${Date.now()}`
    const totalSize = fileData.length

    // Create session info
    const sessionInfo: ZmodemSessionInfo = {
      id: sessionId,
      shellId,
      type: 'send',
      session: null as any,
      fileName,
      totalSize,
      transferred: 0
    }

    this.sessions.set(sessionId, sessionInfo)

    // Store completion callback
    if (onComplete) {
      sessionInfo.onComplete = onComplete
    }

    // Create Sentry to detect ZRINIT from server
    // @ts-ignore - zmodem-ts types are incomplete
    const sentry = new Zmodem.Sentry({
      to_terminal: (octets: number[]) => {
        // Echo to terminal
        writeFunc(new Uint8Array(octets))
      },
      on_detect: (detection: any) => {
        const detectedSession = detection.get_session()
        if (detectedSession && detectedSession.type === 'send') {
          // Server is ready for file transfer (received ZRINIT)
          const sendSession = detectedSession as ZmodemSendSession

          // Update session reference
          sessionInfo.session = sendSession

          // Set up sender to write to shell
          sendSession.set_sender((octets: number[]) => {
            writeFunc(new Uint8Array(octets))
          })

          // Handle session end
          sendSession.on('session_end', () => {
            this.sessions.delete(sessionId)
          })

          sendSession.on('error', (err: Error) => {
            logger.error('Zmodem send error:', err)
            onComplete?.(false, err.message)
          })

          // Send the file
          this.doSendFile(sessionId, sendSession, fileName, fileData, onProgress, onComplete)
        }
      },
      on_retract: () => {
        // Session was retracted
        logger.info('Zmodem send session retracted')
        this.sessions.delete(sessionId)
        onComplete?.(false, 'Session retracted')
      },
      sender: (octets: number[]) => {
        writeFunc(new Uint8Array(octets))
      }
    })

    sessionInfo.sentry = sentry

    // Store shell data callback to feed to Sentry
    this.shellDataCallbacks.set(shellId, (data: Buffer) => {
      sentry.consume(Array.from(data))
    })

    return sessionId
  }

  /**
   * Internal method to send file after ZRINIT is received
   */
  private async doSendFile(
    sessionId: string,
    session: ZmodemSendSession,
    fileName: string,
    fileData: Buffer,
    onProgress?: (progress: ZmodemProgress) => void,
    onComplete?: (success: boolean, error?: string) => void
  ): Promise<void> {
    try {
      // Send file offer
      const result = await session.send_offer({
        name: fileName,
        size: fileData.length,
        mtime: Date.now()
      })

      if (result) {
        // Transfer completed successfully
        onProgress?.({
          percentage: 100,
          transferred: fileData.length,
          total: fileData.length,
          fileName
        })
        onComplete?.(true)
      } else {
        // Transfer was skipped or rejected
        onComplete?.(false, 'Transfer skipped or rejected')
      }
    } catch (err: any) {
      logger.error('Zmodem send error:', err)
      onComplete?.(false, err.message)
    }
  }

  /**
   * Feed shell data to Zmodem service for protocol detection
   * This should be called when shell receives data
   */
  feedShellData(shellId: string, data: Buffer): void {
    const callback = this.shellDataCallbacks.get(shellId)
    if (callback) {
      callback(data)
    }

    // Also check all sessions for matching shellId
    for (const [sessionId, session] of this.sessions) {
      if (session.shellId === shellId && session.sentry) {
        // Feed to sentry for detection
      }
    }
  }

  /**
   * Receive file from server via Zmodem (sz)
   *
   * Flow: User clicks "sz" in UI -> we open receive session
   * Server sends ZRQINIT -> we accept -> server sends file
   */
  receiveFile(
    shellId: string,
    filePath: string,
    writeFunc: (data: Uint8Array) => void,
    onProgress?: (progress: ZmodemProgress) => void,
    onComplete?: (success: boolean, error?: string) => void
  ): string {
    const sessionId = `zmodem-receive-${Date.now()}`
    const fileName = path.basename(filePath)

    // Create receive session
    // @ts-ignore - zmodem-ts types are incomplete
    const session: ZmodemReceiveSession = new Zmodem.ReceiveSession()

    const sessionInfo: ZmodemSessionInfo = {
      id: sessionId,
      shellId,
      type: 'receive',
      session,
      filePath,
      fileName,
      totalSize: 0,
      transferred: 0
    }

    this.sessions.set(sessionId, sessionInfo)

    if (onComplete) {
      sessionInfo.onComplete = onComplete
    }

    // Set up sender to write responses to shell
    session.set_sender((octets: number[]) => {
      writeFunc(new Uint8Array(octets))
    })

    // Handle file offer
    session.on('offer', (offer: any) => {
      const details = offer.getDetails()
      logger.info('Zmodem offer received:', details)

      sessionInfo.fileName = details.name || fileName
      sessionInfo.totalSize = details.size || 0

      // Accept the offer
      offer.accept({
        on_input: 'spool_array'
      }).then((transfer: any) => {
        // Handle incoming data
        transfer.on('input', (data: Uint8Array) => {
          sessionInfo.transferred = (sessionInfo.transferred || 0) + data.length

          // Write to file
          try {
            fs.writeFileSync(filePath, data, { flag: 'a' })
          } catch (err: any) {
            logger.error('Error writing Zmodem file:', err)
          }

          const total = sessionInfo.totalSize || sessionInfo.transferred || 0
          onProgress?.({
            percentage: total > 0
              ? Math.round((sessionInfo.transferred! / total) * 100)
              : 0,
            transferred: sessionInfo.transferred!,
            total: total,
            fileName: sessionInfo.fileName
          })
        })

        transfer.on('complete', () => {
          logger.info('Zmodem receive complete')
          onProgress?.({
            percentage: 100,
            transferred: sessionInfo.totalSize || sessionInfo.transferred || 0,
            total: sessionInfo.totalSize || sessionInfo.transferred || 0,
            fileName: sessionInfo.fileName
          })
          onComplete?.(true)
        })
      }).catch((err: Error) => {
        logger.error('Zmodem accept error:', err)
        onComplete?.(false, err.message)
      })
    })

    // Handle session end
    session.on('session_end', () => {
      this.sessions.delete(sessionId)
    })

    session.on('error', (err: Error) => {
      logger.error('Zmodem receive error:', err)
      onComplete?.(false, err.message)
    })

    // Start receiving - this sends ZRINIT to signal we're ready
    session.start().catch((err: Error) => {
      logger.error('Zmodem start error:', err)
      this.sessions.delete(sessionId)
      onComplete?.(false, err.message)
    })

    return sessionId
  }

  /**
   * Cancel a Zmodem session
   */
  cancelSession(sessionId: string): void {
    const session = this.sessions.get(sessionId)
    if (session) {
      try {
        const s = session.session as any
        s.abort?.()
        s.close?.()
      } catch (err: any) {
        logger.error('Error canceling Zmodem session:', err)
      }

      // Clean up shell data callback
      this.shellDataCallbacks.delete(session.shellId)

      this.sessions.delete(sessionId)
    }
  }

  /**
   * Close all active sessions
   */
  closeAll(): void {
    for (const sessionId of this.sessions.keys()) {
      this.cancelSession(sessionId)
    }
  }

  /**
   * Get active session IDs
   */
  getActiveSessions(): string[] {
    return Array.from(this.sessions.keys())
  }

  /**
   * Get session by ID
   */
  getSession(sessionId: string): ZmodemSessionInfo | undefined {
    return this.sessions.get(sessionId)
  }
}

// Singleton instance
export const zmodemService = new ZmodemService()