implement compgen and cd in mshell server

pkg/packet/packet.go

@@ -26,6 +26,8 @@ import (
 //                   >cd, >getcmd, >untailcmd, >input, <resp
 // all             : <>error, <>message, <>ping, <raw
 
+const MaxCompGenValues = 100
+
 var GlobalDebug = false
 
 const (
@@ -46,7 +48,8 @@ const (
 	CdPacketStr        = "cd"        // rpc
 	CmdDataPacketStr   = "cmddata"   // rpc-response
 	RawPacketStr       = "raw"
-	InputPacketStr     = "input" // command
+	InputPacketStr     = "input"   // command
+	CompGenPacketStr   = "compgen" // rpc
 )
 
 const PacketSenderQueueSize = 20
@@ -73,11 +76,13 @@ func init() {
 	TypeStrToFactory[DataPacketStr] = reflect.TypeOf(DataPacketType{})
 	TypeStrToFactory[DataAckPacketStr] = reflect.TypeOf(DataAckPacketType{})
 	TypeStrToFactory[DataEndPacketStr] = reflect.TypeOf(DataEndPacketType{})
+	TypeStrToFactory[CompGenPacketStr] = reflect.TypeOf(CompGenPacketType{})
 
 	var _ RpcPacketType = (*RunPacketType)(nil)
 	var _ RpcPacketType = (*GetCmdPacketType)(nil)
 	var _ RpcPacketType = (*UntailCmdPacketType)(nil)
 	var _ RpcPacketType = (*CdPacketType)(nil)
+	var _ RpcPacketType = (*CompGenPacketType)(nil)
 
 	var _ RpcResponsePacketType = (*CmdStartPacketType)(nil)
 	var _ RpcResponsePacketType = (*ResponsePacketType)(nil)
@@ -315,6 +320,30 @@ func MakeCdPacket() *CdPacketType {
 	return &CdPacketType{Type: CdPacketStr}
 }
 
+type CompGenPacketType struct {
+	Type     string `json:"type"`
+	ReqId    string `json:"reqid"`
+	Prefix   string `json:"prefix"`
+	CompType string `json:"comptype"`
+	Cwd      string `json:"cwd"`
+}
+
+func IsValidCompGenType(t string) bool {
+	return (t == "file" || t == "command" || t == "directory")
+}
+
+func (*CompGenPacketType) GetType() string {
+	return CompGenPacketStr
+}
+
+func (p *CompGenPacketType) GetReqId() string {
+	return p.ReqId
+}
+
+func MakeCompGenPacket() *CompGenPacketType {
+	return &CompGenPacketType{Type: CompGenPacketStr}
+}
+
 type ResponsePacketType struct {
 	Type    string      `json:"type"`
 	RespId  string      `json:"respid"`

pkg/server/server.go

@@ -10,8 +10,11 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"os/exec"
+	"strings"
 	"sync"
 
+	"github.com/alessio/shellescape"
 	"github.com/scripthaus-dev/mshell/pkg/base"
 	"github.com/scripthaus-dev/mshell/pkg/packet"
 	"github.com/scripthaus-dev/mshell/pkg/shexec"
@@ -48,6 +51,47 @@ func (m *MServer) ProcessCommandPacket(pk packet.CommandPacketType) {
 	return
 }
 
+func (m *MServer) runCompGen(compPk *packet.CompGenPacketType) {
+	reqId := compPk.GetReqId()
+	if !packet.IsValidCompGenType(compPk.CompType) {
+		m.Sender.SendErrorResponse(reqId, fmt.Errorf("invalid compgen type '%s'", compPk.CompType))
+		return
+	}
+	compGenCmdStr := fmt.Sprintf("cd %s; compgen -A %s -- %s | head -n %d", shellescape.Quote(compPk.Cwd), shellescape.Quote(compPk.CompType), shellescape.Quote(compPk.Prefix), packet.MaxCompGenValues)
+	ecmd := exec.Command("bash", "-c", compGenCmdStr)
+	outputBytes, err := ecmd.Output()
+	if err != nil {
+		m.Sender.SendErrorResponse(reqId, fmt.Errorf("compgen error: %w", err))
+		return
+	}
+	outputStr := string(outputBytes)
+	parts := strings.Split(outputStr, "\n")
+	if len(parts) > 0 && parts[len(parts)-1] == "" {
+		parts = parts[0 : len(parts)-1]
+	}
+	m.Sender.SendResponse(reqId, map[string]interface{}{"comps": parts})
+	return
+}
+
+func (m *MServer) ProcessRpcPacket(pk packet.RpcPacketType) {
+	reqId := pk.GetReqId()
+	if cdPk, ok := pk.(*packet.CdPacketType); ok {
+		err := os.Chdir(cdPk.Dir)
+		if err != nil {
+			m.Sender.SendErrorResponse(reqId, fmt.Errorf("cannot change directory: %w", err))
+			return
+		}
+		m.Sender.SendResponse(reqId, true)
+		return
+	}
+	if compPk, ok := pk.(*packet.CompGenPacketType); ok {
+		go m.runCompGen(compPk)
+		return
+	}
+	m.Sender.SendErrorResponse(reqId, fmt.Errorf("invalid rpc type '%s'", pk.GetType()))
+	return
+}
+
 func (m *MServer) runCommand(runPacket *packet.RunPacketType) {
 	if err := runPacket.CK.Validate("packet"); err != nil {
 		m.Sender.SendErrorResponse(runPacket.ReqId, fmt.Errorf("server run packets require valid ck: %s", err))
@@ -117,6 +161,10 @@ func RunServer() (int, error) {
 			server.ProcessCommandPacket(cmdPk)
 			continue
 		}
+		if rpcPk, ok := pk.(packet.RpcPacketType); ok {
+			server.ProcessRpcPacket(rpcPk)
+			continue
+		}
 		server.Sender.SendMessage(fmt.Sprintf("invalid packet '%s' sent to mshell server", packet.AsString(pk)))
 		continue
 	}