bandit-runner/ssh-proxy/server.ts

import express from 'express'
import { Client } from 'ssh2'
import cors from 'cors'

const app = express()
app.use(cors())
app.use(express.json())

// Store active connections
const connections = new Map<string, Client>()

// POST /ssh/connect
app.post('/ssh/connect', async (req, res) => {
  const { host, port, username, password, testOnly } = req.body
  
  // Security: Only allow connections to Bandit server
  if (host !== 'bandit.labs.overthewire.org' || port !== 2220) {
    return res.status(403).json({ 
      success: false, 
      message: 'Only connections to bandit.labs.overthewire.org:2220 are allowed' 
    })
  }

  const client = new Client()
  const connectionId = `conn-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`

  client.on('ready', () => {
    if (testOnly) {
      client.end()
      return res.json({ 
        connectionId: null,
        success: true, 
        message: 'Password validated successfully' 
      })
    }

    connections.set(connectionId, client)
    res.json({ 
      connectionId, 
      success: true, 
      message: 'Connected successfully' 
    })
  })

  client.on('error', (err) => {
    res.status(400).json({ 
      connectionId: null,
      success: false, 
      message: `Connection failed: ${err.message}` 
    })
  })

  client.connect({
    host,
    port,
    username,
    password,
    readyTimeout: 10000,
  })
})

// POST /ssh/exec - with PTY support for full terminal capture
app.post('/ssh/exec', async (req, res) => {
  const { connectionId, command, timeout = 30000, usePTY = true } = req.body
  const client = connections.get(connectionId)

  if (!client) {
    return res.status(404).json({ 
      success: false, 
      error: 'Connection not found' 
    })
  }

  let output = ''
  let stderr = ''
  
  const timeoutHandle = setTimeout(() => {
    res.json({
      output: output + '\n[Command timed out]',
      exitCode: 124,
      success: false,
      duration: timeout,
    })
  }, timeout)

  if (usePTY) {
    // Use PTY mode for full terminal emulation with ANSI codes
    client.exec(command, {
      pty: {
        term: 'xterm-256color',
        cols: 120,
        rows: 40,
      }
    }, (err, stream) => {
      if (err) {
        clearTimeout(timeoutHandle)
        return res.status(500).json({ 
          success: false, 
          error: err.message 
        })
      }

      stream.on('data', (data: Buffer) => {
        output += data.toString() // Includes ANSI codes and prompts
      })

      stream.on('close', (code: number) => {
        clearTimeout(timeoutHandle)
        res.json({
          output, // Full terminal output with ANSI
          exitCode: code || 0,
          success: (code || 0) === 0,
          duration: Date.now() % timeout,
        })
      })
    })
  } else {
    // Legacy mode without PTY
    client.exec(command, (err, stream) => {
      if (err) {
        clearTimeout(timeoutHandle)
        return res.status(500).json({ 
          success: false, 
          error: err.message 
        })
      }

      stream.on('data', (data: Buffer) => {
        output += data.toString()
      })

      stream.stderr.on('data', (data: Buffer) => {
        stderr += data.toString()
      })

      stream.on('close', (code: number) => {
        clearTimeout(timeoutHandle)
        res.json({
          output: output || stderr,
          exitCode: code,
          success: code === 0,
          duration: Date.now() % timeout,
        })
      })
    })
  }
})

// POST /ssh/disconnect
app.post('/ssh/disconnect', (req, res) => {
  const { connectionId } = req.body
  const client = connections.get(connectionId)

  if (client) {
    client.end()
    connections.delete(connectionId)
    res.json({ success: true, message: 'Disconnected' })
  } else {
    res.status(404).json({ success: false, message: 'Connection not found' })
  }
})

// GET /ssh/health
// POST /agent/run
app.post('/agent/run', async (req, res) => {
  const { runId, modelName, startLevel, endLevel, apiKey } = req.body

  if (!runId || !modelName || !apiKey) {
    return res.status(400).json({ error: 'Missing required parameters' })
  }

  try {
    // Set headers for Server-Sent Events / JSONL streaming
    res.setHeader('Content-Type', 'application/x-ndjson')
    res.setHeader('Cache-Control', 'no-cache')
    res.setHeader('Connection', 'keep-alive')

    // Import and create agent
    const { BanditAgent } = await import('./agent.js')
    
    const agent = new BanditAgent({
      runId,
      modelName,
      apiKey,
      startLevel: startLevel || 0,
      endLevel: endLevel || 33,
      responseSender: res,
    })

    // Run agent (it will stream events to response)
    await agent.run({
      runId,
      currentLevel: startLevel || 0,
      targetLevel: endLevel || 33,
      currentPassword: startLevel === 0 ? 'bandit0' : '',
      nextPassword: null,
      levelGoal: '', // Will be set by agent
      status: 'planning',
      retryCount: 0,
      maxRetries: 3,
      sshConnectionId: null,
      error: null,
    })
  } catch (error) {
    console.error('Agent run error:', error)
    if (!res.headersSent) {
      res.status(500).json({ error: error instanceof Error ? error.message : 'Unknown error' })
    }
  }
})

app.get('/ssh/health', (req, res) => {
  res.json({ 
    status: 'ok', 
    activeConnections: connections.size 
  })
})

const PORT = process.env.PORT || 3001
app.listen(PORT, () => {
  console.log(`SSH Proxy + LangGraph Agent running on port ${PORT}`)
})