calypso/backend/internal/system/terminal.go

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)
}