package ai import ( "encoding/json" "fmt" ) // ToolRouter dispatches tool calls to the appropriate service. // Services are stored as interface{} to avoid circular imports between packages. type ToolRouter struct { ssh interface{} // *ssh.SSHService sftp interface{} // *sftp.SFTPService rdp interface{} // *rdp.RDPService sessions interface{} // *session.Manager connections interface{} // *connections.ConnectionService aiService interface{} // *AIService — for terminal buffer access } // NewToolRouter creates an empty ToolRouter. Call SetServices to wire in backends. func NewToolRouter() *ToolRouter { return &ToolRouter{} } // SetServices wires the router to actual service implementations. func (r *ToolRouter) SetServices(ssh, sftp, rdp, sessions, connections interface{}) { r.ssh = ssh r.sftp = sftp r.rdp = rdp r.sessions = sessions r.connections = connections } // SetAIService wires the router to the AI service for terminal buffer access. func (r *ToolRouter) SetAIService(aiService interface{}) { r.aiService = aiService } // sshWriter is the interface we need from SSHService for terminal_write. type sshWriter interface { Write(sessionID string, data string) error } // sshSessionLister is the interface for listing SSH sessions. type sshSessionLister interface { ListSessions() interface{} } // sftpLister is the interface for SFTP list operations. type sftpLister interface { List(sessionID string, path string) (interface{}, error) } // sftpReader is the interface for SFTP read operations. type sftpReader interface { ReadFile(sessionID string, path string) (string, error) } // sftpWriter is the interface for SFTP write operations. type sftpWriter interface { WriteFile(sessionID string, path string, content string) error } // rdpFrameGetter is the interface for getting RDP screenshots. type rdpFrameGetter interface { GetFrame(sessionID string) ([]byte, error) GetSessionInfo(sessionID string) (interface{}, error) } // rdpMouseSender is the interface for RDP mouse events. type rdpMouseSender interface { SendMouse(sessionID string, x, y int, flags uint32) error } // rdpKeySender is the interface for RDP key events. type rdpKeySender interface { SendKey(sessionID string, scancode uint32, pressed bool) error } // rdpClipboardSender is the interface for RDP clipboard. type rdpClipboardSender interface { SendClipboard(sessionID string, data string) error } // rdpSessionLister is the interface for listing RDP sessions. type rdpSessionLister interface { ListSessions() interface{} } // sessionLister is the interface for listing sessions. type sessionLister interface { List() interface{} } // bufferProvider provides terminal output buffers. type bufferProvider interface { GetBuffer(sessionId string) *TerminalBuffer } // Dispatch routes a tool call to the appropriate handler. func (r *ToolRouter) Dispatch(toolName string, input json.RawMessage) (interface{}, error) { switch toolName { case "terminal_write": return r.handleTerminalWrite(input) case "terminal_read": return r.handleTerminalRead(input) case "terminal_cwd": return r.handleTerminalCwd(input) case "sftp_list": return r.handleSFTPList(input) case "sftp_read": return r.handleSFTPRead(input) case "sftp_write": return r.handleSFTPWrite(input) case "rdp_screenshot": return r.handleRDPScreenshot(input) case "rdp_click": return r.handleRDPClick(input) case "rdp_doubleclick": return r.handleRDPDoubleClick(input) case "rdp_type": return r.handleRDPType(input) case "rdp_keypress": return r.handleRDPKeypress(input) case "rdp_scroll": return r.handleRDPScroll(input) case "rdp_move": return r.handleRDPMove(input) case "list_sessions": return r.handleListSessions(input) case "connect_ssh": return r.handleConnectSSH(input) case "disconnect": return r.handleDisconnect(input) default: return nil, fmt.Errorf("unknown tool: %s", toolName) } } func (r *ToolRouter) handleTerminalWrite(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` Text string `json:"text"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } w, ok := r.ssh.(sshWriter) if !ok || r.ssh == nil { return nil, fmt.Errorf("SSH service not available") } if err := w.Write(params.SessionID, params.Text); err != nil { return nil, err } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleTerminalRead(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` Lines int `json:"lines"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } if params.Lines <= 0 { params.Lines = 50 } bp, ok := r.aiService.(bufferProvider) if !ok || r.aiService == nil { return nil, fmt.Errorf("terminal buffer not available") } buf := bp.GetBuffer(params.SessionID) lines := buf.ReadLast(params.Lines) return map[string]interface{}{ "lines": lines, "count": len(lines), }, nil } func (r *ToolRouter) handleTerminalCwd(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } // terminal_cwd works by writing "pwd" to the terminal and reading output w, ok := r.ssh.(sshWriter) if !ok || r.ssh == nil { return nil, fmt.Errorf("SSH service not available") } if err := w.Write(params.SessionID, "pwd\n"); err != nil { return nil, fmt.Errorf("send pwd: %w", err) } return map[string]string{ "status": "pwd command sent — read terminal output for result", }, nil } func (r *ToolRouter) handleSFTPList(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` Path string `json:"path"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } l, ok := r.sftp.(sftpLister) if !ok || r.sftp == nil { return nil, fmt.Errorf("SFTP service not available") } return l.List(params.SessionID, params.Path) } func (r *ToolRouter) handleSFTPRead(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` Path string `json:"path"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } reader, ok := r.sftp.(sftpReader) if !ok || r.sftp == nil { return nil, fmt.Errorf("SFTP service not available") } content, err := reader.ReadFile(params.SessionID, params.Path) if err != nil { return nil, err } return map[string]string{"content": content}, nil } func (r *ToolRouter) handleSFTPWrite(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` Path string `json:"path"` Content string `json:"content"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } w, ok := r.sftp.(sftpWriter) if !ok || r.sftp == nil { return nil, fmt.Errorf("SFTP service not available") } if err := w.WriteFile(params.SessionID, params.Path, params.Content); err != nil { return nil, err } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleRDPScreenshot(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } fg, ok := r.rdp.(rdpFrameGetter) if !ok || r.rdp == nil { return nil, fmt.Errorf("RDP service not available") } frame, err := fg.GetFrame(params.SessionID) if err != nil { return nil, err } // Get session info for dimensions info, err := fg.GetSessionInfo(params.SessionID) if err != nil { return nil, err } // Try to get dimensions from session config type configGetter interface { GetConfig() (int, int) } width, height := 1920, 1080 if cg, ok := info.(configGetter); ok { width, height = cg.GetConfig() } // Encode as JPEG jpeg, err := EncodeScreenshot(frame, width, height, 1280, 720, 75) if err != nil { return nil, fmt.Errorf("encode screenshot: %w", err) } return map[string]interface{}{ "image": jpeg, "width": width, "height": height, }, nil } func (r *ToolRouter) handleRDPClick(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` X int `json:"x"` Y int `json:"y"` Button string `json:"button"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } ms, ok := r.rdp.(rdpMouseSender) if !ok || r.rdp == nil { return nil, fmt.Errorf("RDP service not available") } var buttonFlag uint32 = 0x1000 // left switch params.Button { case "right": buttonFlag = 0x2000 case "middle": buttonFlag = 0x4000 } // Press if err := ms.SendMouse(params.SessionID, params.X, params.Y, buttonFlag|0x8000); err != nil { return nil, err } // Release if err := ms.SendMouse(params.SessionID, params.X, params.Y, buttonFlag); err != nil { return nil, err } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleRDPDoubleClick(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` X int `json:"x"` Y int `json:"y"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } ms, ok := r.rdp.(rdpMouseSender) if !ok || r.rdp == nil { return nil, fmt.Errorf("RDP service not available") } // Two clicks for i := 0; i < 2; i++ { if err := ms.SendMouse(params.SessionID, params.X, params.Y, 0x1000|0x8000); err != nil { return nil, err } if err := ms.SendMouse(params.SessionID, params.X, params.Y, 0x1000); err != nil { return nil, err } } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleRDPType(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` Text string `json:"text"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } cs, ok := r.rdp.(rdpClipboardSender) if !ok || r.rdp == nil { return nil, fmt.Errorf("RDP service not available") } // Send text via clipboard, then simulate Ctrl+V if err := cs.SendClipboard(params.SessionID, params.Text); err != nil { return nil, err } ks, ok := r.rdp.(rdpKeySender) if !ok { return nil, fmt.Errorf("RDP key service not available") } // Ctrl down, V down, V up, Ctrl up if err := ks.SendKey(params.SessionID, 0x001D, true); err != nil { return nil, err } if err := ks.SendKey(params.SessionID, 0x002F, true); err != nil { return nil, err } if err := ks.SendKey(params.SessionID, 0x002F, false); err != nil { return nil, err } if err := ks.SendKey(params.SessionID, 0x001D, false); err != nil { return nil, err } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleRDPKeypress(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` Key string `json:"key"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } ks, ok := r.rdp.(rdpKeySender) if !ok || r.rdp == nil { return nil, fmt.Errorf("RDP service not available") } // Simple key name to scancode mapping for common keys keyMap := map[string]uint32{ "Enter": 0x001C, "Tab": 0x000F, "Escape": 0x0001, "Backspace": 0x000E, "Delete": 0xE053, "Space": 0x0039, "Up": 0xE048, "Down": 0xE050, "Left": 0xE04B, "Right": 0xE04D, } scancode, ok := keyMap[params.Key] if !ok { return nil, fmt.Errorf("unknown key: %s", params.Key) } if err := ks.SendKey(params.SessionID, scancode, true); err != nil { return nil, err } if err := ks.SendKey(params.SessionID, scancode, false); err != nil { return nil, err } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleRDPScroll(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` X int `json:"x"` Y int `json:"y"` Direction string `json:"direction"` Clicks int `json:"clicks"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } if params.Clicks <= 0 { params.Clicks = 3 } ms, ok := r.rdp.(rdpMouseSender) if !ok || r.rdp == nil { return nil, fmt.Errorf("RDP service not available") } var flags uint32 = 0x0200 // wheel flag if params.Direction == "down" { flags |= 0x0100 // negative flag } for i := 0; i < params.Clicks; i++ { if err := ms.SendMouse(params.SessionID, params.X, params.Y, flags); err != nil { return nil, err } } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleRDPMove(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` X int `json:"x"` Y int `json:"y"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } ms, ok := r.rdp.(rdpMouseSender) if !ok || r.rdp == nil { return nil, fmt.Errorf("RDP service not available") } if err := ms.SendMouse(params.SessionID, params.X, params.Y, 0x0800); err != nil { return nil, err } return map[string]string{"status": "ok"}, nil } func (r *ToolRouter) handleListSessions(_ json.RawMessage) (interface{}, error) { result := map[string]interface{}{ "ssh": []interface{}{}, "rdp": []interface{}{}, } if r.sessions != nil { if sl, ok := r.sessions.(sessionLister); ok { result["all"] = sl.List() } } return result, nil } func (r *ToolRouter) handleConnectSSH(input json.RawMessage) (interface{}, error) { var params struct { ConnectionID int64 `json:"connectionId"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } // This will be wired to the app-level connect logic later return nil, fmt.Errorf("connect_ssh requires app-level wiring — not yet available via tool dispatch") } func (r *ToolRouter) handleDisconnect(input json.RawMessage) (interface{}, error) { var params struct { SessionID string `json:"sessionId"` } if err := json.Unmarshal(input, ¶ms); err != nil { return nil, fmt.Errorf("parse input: %w", err) } // Try SSH first type disconnecter interface { Disconnect(sessionID string) error } if d, ok := r.ssh.(disconnecter); ok { if err := d.Disconnect(params.SessionID); err == nil { return map[string]string{"status": "disconnected", "protocol": "ssh"}, nil } } if d, ok := r.rdp.(disconnecter); ok { if err := d.Disconnect(params.SessionID); err == nil { return map[string]string{"status": "disconnected", "protocol": "rdp"}, nil } } return nil, fmt.Errorf("session %s not found in SSH or RDP", params.SessionID) }