terminal/server/api/term.go

package api

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"next-terminal/server/common"
	"next-terminal/server/common/nt"
	"path"
	"strconv"

	"next-terminal/server/common/guacamole"
	"next-terminal/server/common/term"
	"next-terminal/server/config"
	"next-terminal/server/dto"
	"next-terminal/server/global/session"
	"next-terminal/server/log"
	"next-terminal/server/model"
	"next-terminal/server/repository"
	"next-terminal/server/service"
	"next-terminal/server/utils"

	"github.com/gorilla/websocket"
	"github.com/labstack/echo/v4"
)

const (
	Closed = 0
	Data   = 1
	Resize = 2
	Ping   = 9
)

type WebTerminalApi struct {
}

func (api WebTerminalApi) SshEndpoint(c echo.Context) error {
	ws, err := UpGrader.Upgrade(c.Response().Writer, c.Request(), nil)
	if err != nil {
		return err
	}

	defer func() {
		_ = ws.Close()
	}()
	ctx := context.TODO()

	sessionId := c.Param("id")
	cols, _ := strconv.Atoi(c.QueryParam("cols"))
	rows, _ := strconv.Atoi(c.QueryParam("rows"))

	s, err := service.SessionService.FindByIdAndDecrypt(ctx, sessionId)
	if err != nil {
		return WriteMessage(ws, dto.NewMessage(Closed, "获取会话或解密数据失败"))
	}

	if err := api.permissionCheck(c, s.AssetId); err != nil {
		return WriteMessage(ws, dto.NewMessage(Closed, err.Error()))
	}

	var (
		username   = s.Username
		password   = s.Password
		privateKey = s.PrivateKey
		passphrase = s.Passphrase
		ip         = s.IP
		port       = s.Port
	)

	if s.AccessGatewayId != "" && s.AccessGatewayId != "-" {
		g, err := service.GatewayService.GetGatewayById(s.AccessGatewayId)
		if err != nil {
			return WriteMessage(ws, dto.NewMessage(Closed, "获取接入网关失败："+err.Error()))
		}

		defer g.CloseSshTunnel(s.ID)
		exposedIP, exposedPort, err := g.OpenSshTunnel(s.ID, ip, port)
		if err != nil {
			return WriteMessage(ws, dto.NewMessage(Closed, "创建隧道失败："+err.Error()))
		}
		ip = exposedIP
		port = exposedPort
	}

	recording := ""
	var isRecording = false
	property, err := repository.PropertyRepository.FindByName(ctx, guacamole.EnableRecording)
	if err == nil && property.Value == "true" {
		isRecording = true
	}

	if isRecording {
		recording = path.Join(config.GlobalCfg.Guacd.Recording, sessionId, "recording.cast")
	}

	attributes, err := repository.AssetRepository.FindAssetAttrMapByAssetId(ctx, s.AssetId)
	if err != nil {
		return WriteMessage(ws, dto.NewMessage(Closed, "获取资产属性失败："+err.Error()))
	}

	var xterm = "xterm-256color"
	var nextTerminal *term.NextTerminal
	if attributes[nt.SocksProxyEnable] == "true" {
		nextTerminal, err = term.NewNextTerminalUseSocks(ip, port, username, password, privateKey, passphrase, rows, cols, recording, xterm, true, attributes[nt.SocksProxyHost], attributes[nt.SocksProxyPort], attributes[nt.SocksProxyUsername], attributes[nt.SocksProxyPassword])
	} else {
		nextTerminal, err = term.NewNextTerminal(ip, port, username, password, privateKey, passphrase, rows, cols, recording, xterm, true)
	}

	if err != nil {
		return WriteMessage(ws, dto.NewMessage(Closed, "创建SSH客户端失败："+err.Error()))
	}

	if err := nextTerminal.RequestPty(xterm, rows, cols); err != nil {
		return err
	}

	if err := nextTerminal.Shell(); err != nil {
		return err
	}

	sessionForUpdate := model.Session{
		ConnectionId:  sessionId,
		Width:         cols,
		Height:        rows,
		Status:        nt.Connected,
		Recording:     recording,
		ConnectedTime: common.NowJsonTime(),
	}
	if sessionForUpdate.Recording == "" {
		// 未录屏时无需审计
		sessionForUpdate.Reviewed = true
	}
	// 更新会话状态
	if err := repository.SessionRepository.UpdateById(ctx, &sessionForUpdate, sessionId); err != nil {
		return err
	}

	nextSession := &session.Session{
		ID:           s.ID,
		Protocol:     s.Protocol,
		Mode:         s.Mode,
		WebSocket:    ws,
		GuacdTunnel:  nil,
		NextTerminal: nextTerminal,
		Observer:     session.NewObserver(s.ID),
	}
	session.GlobalSessionManager.Add(nextSession)

	termHandler := NewTermHandler(s.Creator, s.AssetId, sessionId, isRecording, ws, nextTerminal)
	termHandler.Start()
	defer termHandler.Stop()

	for {
		_, message, err := ws.ReadMessage()
		if err != nil {
			// web socket会话关闭后主动关闭ssh会话
			service.SessionService.CloseSessionById(sessionId, Normal, "用户正常退出")
			break
		}

		msg, err := dto.ParseMessage(string(message))
		if err != nil {
			continue
		}

		switch msg.Type {
		case Resize:
			decodeString, err := base64.StdEncoding.DecodeString(msg.Content)
			if err != nil {
				continue
			}
			var winSize dto.WindowSize
			err = json.Unmarshal(decodeString, &winSize)
			if err != nil {
				continue
			}
			_ = termHandler.WindowChange(winSize.Rows, winSize.Cols)
			_ = repository.SessionRepository.UpdateWindowSizeById(ctx, winSize.Rows, winSize.Cols, sessionId)
		case Data:
			input := []byte(msg.Content)
			err := termHandler.Write(input)
			if err != nil {
				service.SessionService.CloseSessionById(sessionId, TunnelClosed, "远程连接已关闭")
			}
		case Ping:
			err := termHandler.SendRequest()
			if err != nil {
				service.SessionService.CloseSessionById(sessionId, TunnelClosed, "远程连接已关闭")
			} else {
				_ = termHandler.SendMessageToWebSocket(dto.NewMessage(Ping, msg.Content))
			}

		}
	}
	return nil
}

func (api WebTerminalApi) SshMonitorEndpoint(c echo.Context) error {
	ws, err := UpGrader.Upgrade(c.Response().Writer, c.Request(), nil)
	if err != nil {
		return err
	}

	defer func() {
		_ = ws.Close()
	}()
	ctx := context.TODO()

	sessionId := c.Param("id")
	s, err := repository.SessionRepository.FindById(ctx, sessionId)
	if err != nil {
		return WriteMessage(ws, dto.NewMessage(Closed, "获取会话失败"))
	}

	nextSession := session.GlobalSessionManager.GetById(sessionId)
	if nextSession == nil {
		return WriteMessage(ws, dto.NewMessage(Closed, "会话已离线"))
	}

	obId := utils.UUID()
	obSession := &session.Session{
		ID:        obId,
		Protocol:  s.Protocol,
		Mode:      s.Mode,
		WebSocket: ws,
	}
	nextSession.Observer.Add(obSession)

	for {
		_, _, err := ws.ReadMessage()
		if err != nil {
			nextSession.Observer.Del(obId)
			break
		}
	}
	return nil
}

func (api WebTerminalApi) AccessTerminalEndpoint(c echo.Context) error {
	ws, err := UpGrader.Upgrade(c.Response().Writer, c.Request(), nil)
	if err != nil {
		return err
	}
	defer func() {
		_ = ws.Close()
	}()

	sessionId := c.QueryParam("sessionId")
	log.Debug("AccessTerminal: WebSocket connected, sessionId=" + sessionId + ", cols=" + c.QueryParam("cols") + ", rows=" + c.QueryParam("rows"))
	if sessionId == "" {
		return WriteMessage(ws, dto.NewMessage(Closed, "sessionId is required"))
	}

	ctx := context.TODO()

	s, err := service.SessionService.FindByIdAndDecrypt(ctx, sessionId)
	if err != nil {
		log.Debug("AccessTerminal: session not found, err=" + err.Error())
		return WriteMessage(ws, dto.NewMessage(Closed, "session not found: "+err.Error()))
	}
	log.Debug("AccessTerminal: session found, protocol=" + s.Protocol + " ip=" + s.IP + ":" + strconv.Itoa(s.Port) + " user=" + s.Username + " pwdLen=" + strconv.Itoa(len(s.Password)) + " keyLen=" + strconv.Itoa(len(s.PrivateKey)))

	xterm := "xterm-256color"
	cols, _ := strconv.Atoi(c.QueryParam("cols"))
	rows, _ := strconv.Atoi(c.QueryParam("rows"))
	if cols <= 0 {
		cols = 80
	}
	if rows <= 0 {
		rows = 24
	}

	// 检查是否启用录制
	recording := ""
	isRecording := false
	property, err := repository.PropertyRepository.FindByName(ctx, guacamole.EnableRecording)
	if err == nil && property.Value == "true" {
		isRecording = true
		recording = path.Join(config.GlobalCfg.Guacd.Recording, sessionId, "recording.cast")
	}

	var nextTerminal *term.NextTerminal
	if isRecording {
		nextTerminal, err = term.NewNextTerminal(s.IP, s.Port, s.Username, s.Password, s.PrivateKey, s.Passphrase, rows, cols, recording, xterm, true)
	} else {
		nextTerminal, err = CreateNextTerminalBySession(s)
	}
	if err != nil {
		log.Debug("AccessTerminal: CreateNextTerminalBySession failed, err=" + err.Error())
		return WriteMessage(ws, dto.NewMessage(Closed, "创建SSH客户端失败："+err.Error()))
	}
	log.Debug("AccessTerminal: SSH client created successfully")

	if err := nextTerminal.RequestPty(xterm, rows, cols); err != nil {
		log.Debug("AccessTerminal: RequestPty failed, err=" + err.Error())
		return WriteMessage(ws, dto.NewMessage(Closed, "请求PTY失败："+err.Error()))
	}
	log.Debug("AccessTerminal: RequestPty OK")

	if err := nextTerminal.Shell(); err != nil {
		log.Debug("AccessTerminal: Shell failed, err=" + err.Error())
		return WriteMessage(ws, dto.NewMessage(Closed, "启动Shell失败："+err.Error()))
	}
	log.Debug("AccessTerminal: Shell OK, starting TermHandler...")

	sessionForUpdate := model.Session{
		ConnectionId:  sessionId,
		Width:         cols,
		Height:        rows,
		Status:        nt.Connected,
		Recording:     recording,
		ConnectedTime: common.NowJsonTime(),
	}
	if sessionForUpdate.Recording == "" {
		sessionForUpdate.Reviewed = true
	}
	if err := repository.SessionRepository.UpdateById(ctx, &sessionForUpdate, sessionId); err != nil {
		return err
	}

	nextSession := &session.Session{
		ID:           s.ID,
		Protocol:     s.Protocol,
		Mode:         s.Mode,
		WebSocket:    ws,
		GuacdTunnel:  nil,
		NextTerminal: nextTerminal,
		Observer:     session.NewObserver(s.ID),
	}
	session.GlobalSessionManager.Add(nextSession)

	termHandler := NewTermHandler(s.Creator, s.AssetId, sessionId, isRecording, ws, nextTerminal)
	termHandler.Start()
	defer termHandler.Stop()

	log.Debug("AccessTerminal: TermHandler started, entering message loop")

	for {
		_, message, err := ws.ReadMessage()
		if err != nil {
			log.Debug("AccessTerminal: WebSocket read error, closing session. err=" + err.Error())
			service.SessionService.CloseSessionById(sessionId, Normal, "用户正常退出")
			break
		}
		msg, err := dto.ParseMessage(string(message))
		if err != nil {
			continue
		}
		switch msg.Type {
		case Resize:
			decodeString, _ := base64.StdEncoding.DecodeString(msg.Content)
			var winSize dto.WindowSize
			_ = json.Unmarshal(decodeString, &winSize)
			_ = termHandler.WindowChange(winSize.Rows, winSize.Cols)
			_ = repository.SessionRepository.UpdateWindowSizeById(ctx, winSize.Rows, winSize.Cols, sessionId)
		case Data:
			_ = termHandler.Write([]byte(msg.Content))
		case Ping:
			_ = termHandler.SendRequest()
			_ = termHandler.SendMessageToWebSocket(dto.NewMessage(Ping, msg.Content))
		}
	}
	return nil
}

func (api WebTerminalApi) permissionCheck(c echo.Context, assetId string) error {
	user, _ := GetCurrentAccount(c)
	_ = user
	return nil
}

func WriteMessage(ws *websocket.Conn, msg dto.Message) error {
	message := []byte(msg.ToString())
	return ws.WriteMessage(websocket.TextMessage, message)
}

func CreateNextTerminalBySession(session model.Session) (*term.NextTerminal, error) {
	var (
		username   = session.Username
		password   = session.Password
		privateKey = session.PrivateKey
		passphrase = session.Passphrase
		ip         = session.IP
		port       = session.Port
	)
	return term.NewNextTerminal(ip, port, username, password, privateKey, passphrase,
		10, 10, "", "", true)
}