bandit-runner/ssh-proxy/agent.ts

/**
 * LangGraph agent integrated with SSH proxy
 * Runs in Node.js environment with full dependency support
 * Based on context7 best practices for streaming and config passing
 */

import { StateGraph, END, START, Annotation } from "@langchain/langgraph"
import { HumanMessage, SystemMessage } from "@langchain/core/messages"
import { ChatOpenAI } from "@langchain/openai"
import type { RunnableConfig } from "@langchain/core/runnables"
import type { Client } from 'ssh2'
import type { Response } from 'express'

// Define state using Annotation for proper LangGraph typing
const BanditState = Annotation.Root({
  runId: Annotation<string>,
  currentLevel: Annotation<number>,
  targetLevel: Annotation<number>,
  currentPassword: Annotation<string>,
  nextPassword: Annotation<string | null>,
  levelGoal: Annotation<string>,
  commandHistory: Annotation<Array<{
    command: string
    output: string
    exitCode: number
    timestamp: string
    level: number
  }>>({
    reducer: (left, right) => left.concat(right),
    default: () => [],
  }),
  thoughts: Annotation<Array<{
    type: 'plan' | 'observation' | 'reasoning' | 'decision'
    content: string
    timestamp: string
    level: number
  }>>({
    reducer: (left, right) => left.concat(right),
    default: () => [],
  }),
  status: Annotation<'planning' | 'executing' | 'validating' | 'advancing' | 'paused' | 'complete' | 'failed'>,
  retryCount: Annotation<number>,
  maxRetries: Annotation<number>,
  sshConnectionId: Annotation<string | null>,
  error: Annotation<string | null>,
})

type BanditAgentState = typeof BanditState.State

const LEVEL_GOALS: Record<number, string> = {
  0: "Read 'readme' file in home directory",
  1: "Read '-' file (use 'cat ./-' or 'cat < -')",
  2: "Find and read hidden file with spaces in name",
  3: "Find file with specific permissions",
  4: "Find file in inhere directory that is human-readable",
  5: "Find file owned by bandit7, group bandit6, 33 bytes",
  // Add more as needed
}

const SYSTEM_PROMPT = `You are BanditRunner, an autonomous operator solving the OverTheWire Bandit wargame.

RULES:
1. Only use safe commands: ls, cat, grep, find, base64, etc.
2. Think step-by-step
3. Extract passwords (32-char alphanumeric strings)
4. Validate before advancing

WORKFLOW:
1. Plan - analyze level goal
2. Execute - run command
3. Validate - check for password
4. Advance - move to next level`

/**
 * Create planning node - LLM decides next command
 * Following context7 pattern: pass RunnableConfig for proper streaming
 */
async function planLevel(
  state: BanditAgentState,
  config?: RunnableConfig
): Promise<Partial<BanditAgentState>> {
  const { currentLevel, levelGoal, commandHistory, sshConnectionId, currentPassword } = state
  
  // Establish SSH connection if needed
  if (!sshConnectionId) {
    const sshProxyUrl = process.env.SSH_PROXY_URL || 'http://localhost:3001'
    const connectResponse = await fetch(`${sshProxyUrl}/ssh/connect`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        host: 'bandit.labs.overthewire.org',
        port: 2220,
        username: `bandit${currentLevel}`,
        password: currentPassword,
        testOnly: false,
      }),
    })

    const connectData = await connectResponse.json() as { connectionId?: string; success?: boolean; message?: string }
    
    if (!connectData.success || !connectData.connectionId) {
      return {
        status: 'failed',
        error: `SSH connection failed: ${connectData.message || 'Unknown error'}`,
      }
    }

    // Update state with connection ID
    return {
      sshConnectionId: connectData.connectionId,
      status: 'planning',
    }
  }
  
  // Get LLM from config (injected by agent)
  const llm = (config?.configurable?.llm) as ChatOpenAI
  
  // Build context from recent commands
  const recentCommands = commandHistory.slice(-3).map(cmd => 
    `Command: ${cmd.command}\nOutput: ${cmd.output.slice(0, 300)}\nExit: ${cmd.exitCode}`
  ).join('\n\n')

  const messages = [
    new SystemMessage(SYSTEM_PROMPT),
    new HumanMessage(`Level ${currentLevel}: ${levelGoal}

Recent Commands:
${recentCommands || 'No commands yet'}

What command should I run next? Provide ONLY the exact command to execute.`),
  ]

  const response = await llm.invoke(messages, config)
  const thought = response.content as string

  return {
    thoughts: [{
      type: 'plan',
      content: thought,
      timestamp: new Date().toISOString(),
      level: currentLevel,
    }],
    status: 'executing',
  }
}

/**
 * Execute SSH command via proxy with PTY
 */
async function executeCommand(
  state: BanditAgentState,
  config?: RunnableConfig
): Promise<Partial<BanditAgentState>> {
  const { thoughts, currentLevel, sshConnectionId } = state
  
  // Extract command from latest thought
  const latestThought = thoughts[thoughts.length - 1]
  const commandMatch = latestThought.content.match(/```(?:bash|sh)?\s*\n?(.+?)\n?```/s) ||
                       latestThought.content.match(/^(.+)$/m)
  
  if (!commandMatch) {
    return {
      status: 'failed',
      error: 'Could not extract command from LLM response',
    }
  }

  const command = commandMatch[1].trim()

  // Execute via SSH with PTY enabled
  try {
    const sshProxyUrl = process.env.SSH_PROXY_URL || 'http://localhost:3001'
    const response = await fetch(`${sshProxyUrl}/ssh/exec`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        connectionId: sshConnectionId,
        command,
        usePTY: true, // Enable PTY for full terminal capture
        timeout: 30000,
      }),
    })

    const data = await response.json() as { output?: string; exitCode?: number; success?: boolean }
    
    const result = {
      command,
      output: data.output || '',
      exitCode: data.exitCode || 1,
      timestamp: new Date().toISOString(),
      level: currentLevel,
    }

    return {
      commandHistory: [result],
      status: 'validating',
    }
  } catch (error) {
    return {
      status: 'failed',
      error: `SSH execution failed: ${error instanceof Error ? error.message : String(error)}`,
    }
  }
}

/**
 * Validate if password was found
 */
async function validateResult(
  state: BanditAgentState,
  config?: RunnableConfig
): Promise<Partial<BanditAgentState>> {
  const { commandHistory } = state
  const lastCommand = commandHistory[commandHistory.length - 1]
  
  // Simple password extraction (32-char alphanumeric)
  const passwordMatch = lastCommand.output.match(/([A-Za-z0-9]{32,})/)
  
  if (passwordMatch) {
    return {
      nextPassword: passwordMatch[1],
      status: 'advancing',
    }
  }

  // Retry if under limit
  if (state.retryCount < state.maxRetries) {
    return {
      retryCount: state.retryCount + 1,
      status: 'planning',
    }
  }

  return {
    status: 'failed',
    error: `Max retries reached for level ${state.currentLevel}`,
  }
}

/**
 * Advance to next level
 */
async function advanceLevel(
  state: BanditAgentState,
  config?: RunnableConfig
): Promise<Partial<BanditAgentState>> {
  const nextLevel = state.currentLevel + 1
  
  if (nextLevel > state.targetLevel) {
    return {
      status: 'complete',
      currentLevel: nextLevel,
      currentPassword: state.nextPassword || '',
    }
  }

  return {
    currentLevel: nextLevel,
    currentPassword: state.nextPassword || '',
    nextPassword: null,
    levelGoal: LEVEL_GOALS[nextLevel] || 'Unknown',
    retryCount: 0,
    status: 'planning',
  }
}

/**
 * Conditional routing function
 */
function shouldContinue(state: BanditAgentState): string {
  if (state.status === 'complete' || state.status === 'failed') return END
  if (state.status === 'paused') return END
  if (state.status === 'planning') return 'plan_level'
  if (state.status === 'executing') return 'execute_command'
  if (state.status === 'validating') return 'validate_result'
  if (state.status === 'advancing') return 'advance_level'
  return END
}

/**
 * Agent executor that can run in SSH proxy
 */
export class BanditAgent {
  private llm: ChatOpenAI
  private graph: ReturnType<typeof StateGraph.prototype.compile>
  private responseSender?: Response

  constructor(config: {
    runId: string
    modelName: string
    apiKey: string
    startLevel: number
    endLevel: number
    responseSender?: Response
  }) {
    this.llm = new ChatOpenAI({
      model: config.modelName,
      apiKey: config.apiKey,
      temperature: 0.7,
      configuration: {
        baseURL: 'https://openrouter.ai/api/v1',
      },
    })

    this.responseSender = config.responseSender
    this.graph = this.createGraph()
  }

  private createGraph() {
    const workflow = new StateGraph(BanditState)
      .addNode('plan_level', planLevel)
      .addNode('execute_command', executeCommand)
      .addNode('validate_result', validateResult)
      .addNode('advance_level', advanceLevel)
      .addEdge(START, 'plan_level')
      .addConditionalEdges('plan_level', shouldContinue)
      .addConditionalEdges('execute_command', shouldContinue)
      .addConditionalEdges('validate_result', shouldContinue)
      .addConditionalEdges('advance_level', shouldContinue)

    return workflow.compile()
  }

  private emit(event: any) {
    if (this.responseSender && !this.responseSender.writableEnded) {
      // Send as JSONL (newline-delimited JSON)
      this.responseSender.write(JSON.stringify(event) + '\n')
    }
  }

  async run(initialState: Partial<BanditAgentState>): Promise<void> {
    try {
      // Stream updates using context7 recommended pattern
      const stream = await this.graph.stream(
        initialState,
        {
          streamMode: "updates", // Per context7: emit after each step
          configurable: { llm: this.llm }, // Pass LLM through config
        }
      )

      for await (const update of stream) {
        // Emit each update as JSONL event
        const [nodeName, nodeOutput] = Object.entries(update)[0]
        
        this.emit({
          type: 'node_update',
          node: nodeName,
          data: nodeOutput,
          timestamp: new Date().toISOString(),
        })

        // Send specific event types based on node
        if (nodeName === 'plan_level' && nodeOutput.thoughts) {
          const thought = nodeOutput.thoughts[nodeOutput.thoughts.length - 1]
          
          // Emit as 'thinking' event for UI
          this.emit({
            type: 'thinking',
            data: {
              content: thought.content,
              level: thought.level,
            },
            timestamp: new Date().toISOString(),
          })
          
          // Also emit as 'agent_message' for chat panel
          this.emit({
            type: 'agent_message',
            data: {
              content: `Planning: ${thought.content}`,
              level: thought.level,
              metadata: {
                thoughtType: thought.type,
              },
            },
            timestamp: new Date().toISOString(),
          })
        }

        if (nodeName === 'execute_command' && nodeOutput.commandHistory) {
          const cmd = nodeOutput.commandHistory[nodeOutput.commandHistory.length - 1]
          
          // Emit tool call event
          this.emit({
            type: 'tool_call',
            data: {
              content: `ssh_exec: ${cmd.command}`,
              level: cmd.level,
              metadata: {
                tool: 'ssh_exec',
                command: cmd.command,
              },
            },
            timestamp: new Date().toISOString(),
          })
          
          // Emit terminal output with prompt
          this.emit({
            type: 'terminal_output',
            data: {
              content: `$ ${cmd.command}`,
              command: cmd.command,
              level: cmd.level,
            },
            timestamp: new Date().toISOString(),
          })
          
          // Emit command result (includes ANSI codes from PTY)
          this.emit({
            type: 'terminal_output',
            data: {
              content: cmd.output,
              level: cmd.level,
            },
            timestamp: new Date().toISOString(),
          })
        }

        if (nodeName === 'validate_result' && nodeOutput.nextPassword) {
          this.emit({
            type: 'agent_message',
            data: {
              content: `Password found: ${nodeOutput.nextPassword}`,
              level: nodeOutput.currentLevel,
            },
            timestamp: new Date().toISOString(),
          })
        }

        if (nodeName === 'advance_level' && nodeOutput.currentLevel !== undefined) {
          this.emit({
            type: 'level_complete',
            data: {
              content: `Level ${nodeOutput.currentLevel - 1} completed successfully`,
              level: nodeOutput.currentLevel - 1,
            },
            timestamp: new Date().toISOString(),
          })
          
          this.emit({
            type: 'agent_message',
            data: {
              content: `Advancing to Level ${nodeOutput.currentLevel}`,
              level: nodeOutput.currentLevel,
            },
            timestamp: new Date().toISOString(),
          })
        }

        if (nodeOutput.error) {
          this.emit({
            type: 'error',
            data: {
              content: nodeOutput.error,
              level: nodeOutput.currentLevel,
            },
            timestamp: new Date().toISOString(),
          })
        }
      }

      // Final completion event
      this.emit({
        type: 'run_complete',
        data: { content: 'Agent run completed successfully' },
        timestamp: new Date().toISOString(),
      })
    } catch (error) {
      this.emit({
        type: 'error',
        data: { content: error instanceof Error ? error.message : String(error) },
        timestamp: new Date().toISOString(),
      })
    } finally {
      if (this.responseSender && !this.responseSender.writableEnded) {
        this.responseSender.end()
      }
    }
  }
}