add feature license management
This commit is contained in:
Warp Agent
328
backend/internal/system/terminal.go
Normal file
328
backend/internal/system/terminal.go
Normal file
@@ -0,0 +1,328 @@
|
||||
package system
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"os/user"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/atlasos/calypso/internal/common/logger"
|
||||
"github.com/creack/pty"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
const (
|
||||
// WebSocket timeouts
|
||||
writeWait = 10 * time.Second
|
||||
pongWait = 60 * time.Second
|
||||
pingPeriod = (pongWait * 9) / 10
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
ReadBufferSize: 4096,
|
||||
WriteBufferSize: 4096,
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
// Allow all origins - in production, validate against allowed domains
|
||||
return true
|
||||
},
|
||||
}
|
||||
|
||||
// TerminalSession manages a single terminal session
|
||||
type TerminalSession struct {
|
||||
conn *websocket.Conn
|
||||
pty *os.File
|
||||
cmd *exec.Cmd
|
||||
logger *logger.Logger
|
||||
mu sync.RWMutex
|
||||
closed bool
|
||||
username string
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// HandleTerminalWebSocket handles WebSocket connection for terminal
|
||||
func HandleTerminalWebSocket(c *gin.Context, log *logger.Logger) {
|
||||
// Verify authentication
|
||||
userID, exists := c.Get("user_id")
|
||||
if !exists {
|
||||
log.Warn("Terminal WebSocket: unauthorized access", "ip", c.ClientIP())
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "unauthorized"})
|
||||
return
|
||||
}
|
||||
|
||||
username, _ := c.Get("username")
|
||||
if username == nil {
|
||||
username = userID
|
||||
}
|
||||
|
||||
log.Info("Terminal WebSocket: connection attempt", "username", username, "ip", c.ClientIP())
|
||||
|
||||
// Upgrade connection
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
log.Error("Terminal WebSocket: upgrade failed", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
log.Info("Terminal WebSocket: connection upgraded", "username", username)
|
||||
|
||||
// Create session
|
||||
session := &TerminalSession{
|
||||
conn: conn,
|
||||
logger: log,
|
||||
username: username.(string),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
// Start terminal
|
||||
if err := session.startPTY(); err != nil {
|
||||
log.Error("Terminal WebSocket: failed to start PTY", "error", err, "username", username)
|
||||
session.sendError(err.Error())
|
||||
session.close()
|
||||
return
|
||||
}
|
||||
|
||||
// Handle messages and PTY output
|
||||
go session.handleRead()
|
||||
go session.handleWrite()
|
||||
}
|
||||
|
||||
// startPTY starts the PTY session
|
||||
func (s *TerminalSession) startPTY() error {
|
||||
// Get user info
|
||||
currentUser, err := user.Lookup(s.username)
|
||||
if err != nil {
|
||||
// Fallback to current user
|
||||
currentUser, err = user.Current()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Determine shell
|
||||
shell := os.Getenv("SHELL")
|
||||
if shell == "" {
|
||||
shell = "/bin/bash"
|
||||
}
|
||||
|
||||
// Create command
|
||||
s.cmd = exec.Command(shell)
|
||||
s.cmd.Env = append(os.Environ(),
|
||||
"TERM=xterm-256color",
|
||||
"HOME="+currentUser.HomeDir,
|
||||
"USER="+currentUser.Username,
|
||||
"USERNAME="+currentUser.Username,
|
||||
)
|
||||
s.cmd.Dir = currentUser.HomeDir
|
||||
|
||||
// Start PTY
|
||||
ptyFile, err := pty.Start(s.cmd)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.pty = ptyFile
|
||||
|
||||
// Set initial size
|
||||
pty.Setsize(ptyFile, &pty.Winsize{
|
||||
Rows: 24,
|
||||
Cols: 80,
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleRead handles incoming WebSocket messages
|
||||
func (s *TerminalSession) handleRead() {
|
||||
defer s.close()
|
||||
|
||||
// Set read deadline and pong handler
|
||||
s.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
s.conn.SetPongHandler(func(string) error {
|
||||
s.conn.SetReadDeadline(time.Now().Add(pongWait))
|
||||
return nil
|
||||
})
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-s.done:
|
||||
return
|
||||
default:
|
||||
messageType, data, err := s.conn.ReadMessage()
|
||||
if err != nil {
|
||||
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
|
||||
s.logger.Error("Terminal WebSocket: read error", "error", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Handle binary messages (raw input)
|
||||
if messageType == websocket.BinaryMessage {
|
||||
s.writeToPTY(data)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle text messages (JSON commands)
|
||||
if messageType == websocket.TextMessage {
|
||||
var msg map[string]interface{}
|
||||
if err := json.Unmarshal(data, &msg); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
switch msg["type"] {
|
||||
case "input":
|
||||
if data, ok := msg["data"].(string); ok {
|
||||
s.writeToPTY([]byte(data))
|
||||
}
|
||||
|
||||
case "resize":
|
||||
if cols, ok1 := msg["cols"].(float64); ok1 {
|
||||
if rows, ok2 := msg["rows"].(float64); ok2 {
|
||||
s.resizePTY(uint16(cols), uint16(rows))
|
||||
}
|
||||
}
|
||||
|
||||
case "ping":
|
||||
s.writeWS(websocket.TextMessage, []byte(`{"type":"pong"}`))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleWrite handles PTY output to WebSocket
|
||||
func (s *TerminalSession) handleWrite() {
|
||||
defer s.close()
|
||||
|
||||
ticker := time.NewTicker(pingPeriod)
|
||||
defer ticker.Stop()
|
||||
|
||||
// Read from PTY and write to WebSocket
|
||||
buffer := make([]byte, 4096)
|
||||
for {
|
||||
select {
|
||||
case <-s.done:
|
||||
return
|
||||
case <-ticker.C:
|
||||
// Send ping
|
||||
if err := s.writeWS(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
default:
|
||||
// Read from PTY
|
||||
if s.pty != nil {
|
||||
n, err := s.pty.Read(buffer)
|
||||
if err != nil {
|
||||
if err != io.EOF {
|
||||
s.logger.Error("Terminal WebSocket: PTY read error", "error", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if n > 0 {
|
||||
// Write binary data to WebSocket
|
||||
if err := s.writeWS(websocket.BinaryMessage, buffer[:n]); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// writeToPTY writes data to PTY
|
||||
func (s *TerminalSession) writeToPTY(data []byte) {
|
||||
s.mu.RLock()
|
||||
closed := s.closed
|
||||
pty := s.pty
|
||||
s.mu.RUnlock()
|
||||
|
||||
if closed || pty == nil {
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := pty.Write(data); err != nil {
|
||||
s.logger.Error("Terminal WebSocket: PTY write error", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// resizePTY resizes the PTY
|
||||
func (s *TerminalSession) resizePTY(cols, rows uint16) {
|
||||
s.mu.RLock()
|
||||
closed := s.closed
|
||||
ptyFile := s.pty
|
||||
s.mu.RUnlock()
|
||||
|
||||
if closed || ptyFile == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// Use pty.Setsize from package, not method from variable
|
||||
pty.Setsize(ptyFile, &pty.Winsize{
|
||||
Cols: cols,
|
||||
Rows: rows,
|
||||
})
|
||||
}
|
||||
|
||||
// writeWS writes message to WebSocket
|
||||
func (s *TerminalSession) writeWS(messageType int, data []byte) error {
|
||||
s.mu.RLock()
|
||||
closed := s.closed
|
||||
conn := s.conn
|
||||
s.mu.RUnlock()
|
||||
|
||||
if closed || conn == nil {
|
||||
return io.ErrClosedPipe
|
||||
}
|
||||
|
||||
conn.SetWriteDeadline(time.Now().Add(writeWait))
|
||||
return conn.WriteMessage(messageType, data)
|
||||
}
|
||||
|
||||
// sendError sends error message
|
||||
func (s *TerminalSession) sendError(errMsg string) {
|
||||
msg := map[string]interface{}{
|
||||
"type": "error",
|
||||
"error": errMsg,
|
||||
}
|
||||
data, _ := json.Marshal(msg)
|
||||
s.writeWS(websocket.TextMessage, data)
|
||||
}
|
||||
|
||||
// close closes the terminal session
|
||||
func (s *TerminalSession) close() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
if s.closed {
|
||||
return
|
||||
}
|
||||
|
||||
s.closed = true
|
||||
close(s.done)
|
||||
|
||||
// Close PTY
|
||||
if s.pty != nil {
|
||||
s.pty.Close()
|
||||
}
|
||||
|
||||
// Kill process
|
||||
if s.cmd != nil && s.cmd.Process != nil {
|
||||
s.cmd.Process.Signal(syscall.SIGTERM)
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
if s.cmd.ProcessState == nil || !s.cmd.ProcessState.Exited() {
|
||||
s.cmd.Process.Kill()
|
||||
}
|
||||
}
|
||||
|
||||
// Close WebSocket
|
||||
if s.conn != nil {
|
||||
s.conn.Close()
|
||||
}
|
||||
|
||||
s.logger.Info("Terminal WebSocket: session closed", "username", s.username)
|
||||
}
|
||||
Reference in New Issue
Block a user