[Update] 调整结构

cmd/coco.go

@@ -5,6 +5,5 @@ import (
 )
 
 func main() {
-	app := &coco.Coco{}
-	app.Start()
+	coco.RunForever()
 }

pkg/coco/coco.go

@@ -2,8 +2,15 @@ package coco
 
 import (
 	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
+	"cocogo/pkg/config"
+	"cocogo/pkg/logger"
+	"cocogo/pkg/proxy"
+	"cocogo/pkg/service"
 	"cocogo/pkg/sshd"
 )
 
@@ -16,9 +23,27 @@ func (c *Coco) Start() {
 	fmt.Println(time.Now().Format("2006-01-02 15:04:05"))
 	fmt.Printf("Coco version %s, more see https://www.jumpserver.org\n", version)
 	fmt.Println("Quit the server with CONTROL-C.")
-	sshd.StartServer()
+	go sshd.StartServer()
 }
 
 func (c *Coco) Stop() {
+	sshd.StopServer()
+	logger.Debug("Quit The Coco")
+}
+
+func RunForever() {
+	loadingBoot()
+	gracefulStop := make(chan os.Signal)
+	signal.Notify(gracefulStop, syscall.SIGTERM, syscall.SIGINT, syscall.SIGQUIT)
+	app := &Coco{}
+	app.Start()
+	<-gracefulStop
+	app.Stop()
+}
 
+func loadingBoot() {
+	config.Initial()
+	logger.Initial()
+	service.Initial()
+	proxy.Initial()
 }

pkg/coco/init.go

-package coco
-
-import (
-	"cocogo/pkg/config"
-	"cocogo/pkg/logger"
-	"cocogo/pkg/service"
-)
-
-func init() {
-	config.Initial()
-	logger.Initial()
-	service.Initial()
-}

pkg/config/config.go

@@ -41,6 +41,7 @@ type Config struct {
 	Comment             string            `yaml:"COMMENT"`
 	Language            string            `yaml:"LANG"`
 	LanguageCode        string            `yaml:"LANGUAGE_CODE"` // Abandon
+	UploadFailedReplay  bool              `yaml:"UPLOAD_FAILED_REPLAY_ON_START"`
 }
 
 func (c *Config) EnsureConfigValid() {
@@ -105,23 +106,25 @@ var lock = new(sync.RWMutex)
 var name, _ = os.Hostname()
 var rootPath, _ = os.Getwd()
 var Conf = &Config{
-	Name:           name,
-	CoreHost:       "http://localhost:8080",
-	BootstrapToken: "",
-	BindHost:       "0.0.0.0",
-	SSHPort:        2222,
-	SSHTimeout:     60,
-	HTTPPort:       5000,
-	AccessKey:      "",
-	AccessKeyFile:  "data/keys/.access_key",
-	LogLevel:       "DEBUG",
-	HostKeyFile:    "data/keys/host_key",
-	HostKey:        "",
-	RootPath:       rootPath,
-	Comment:        "Coco",
-	Language:       "zh",
-	ReplayStorage:  map[string]string{"TYPE": "server"},
-	CommandStorage: map[string]string{"TYPE": "server"},
+	Name:               name,
+	CoreHost:           "http://localhost:8080",
+	BootstrapToken:     "",
+	BindHost:           "0.0.0.0",
+	SSHPort:            2222,
+	SSHTimeout:         15,
+	HTTPPort:           5000,
+	HeartbeatDuration:  10,
+	AccessKey:          "",
+	AccessKeyFile:      "data/keys/.access_key",
+	LogLevel:           "DEBUG",
+	HostKeyFile:        "data/keys/host_key",
+	HostKey:            "",
+	RootPath:           rootPath,
+	Comment:            "Coco",
+	Language:           "zh",
+	ReplayStorage:      map[string]string{"TYPE": "server"},
+	CommandStorage:     map[string]string{"TYPE": "server"},
+	UploadFailedReplay: true,
 }
 
 func SetConf(conf *Config) {

pkg/handler/session.go

@@ -330,8 +330,10 @@ func (h *interactiveHandler) searchAsset(key string) (assets []model.Asset) {
 		searchData = h.searchResult
 	}
 
+	key = strings.ToLower(key)
 	for _, assetValue := range searchData {
-		contents := []string{assetValue.Hostname, assetValue.Ip, assetValue.Comment}
+		contents := []string{strings.ToLower(assetValue.Hostname),
+			strings.ToLower(assetValue.Ip), strings.ToLower(assetValue.Comment)}
 		if isSubstring(contents, key) {
 			assets = append(assets, assetValue)
 		}

pkg/model/assets.go

@@ -268,7 +268,6 @@ func (sf *SystemUserFilterRule) Pattern() *regexp.Regexp {
 
 func (sf *SystemUserFilterRule) Match(cmd string) (RuleAction, string) {
 	found := sf.Pattern().FindString(cmd)
-	fmt.Println(found)
 	if found == "" {
 		return ActionUnknown, ""
 	}

pkg/proxy/parser.go

@@ -2,10 +2,10 @@ package proxy
 
 import (
 	"bytes"
-	"cocogo/pkg/i18n"
 	"fmt"
 	"sync"
 
+	"cocogo/pkg/i18n"
 	"cocogo/pkg/logger"
 	"cocogo/pkg/model"
 	"cocogo/pkg/utils"
@@ -114,8 +114,8 @@ func (p *Parser) parseInputState(b []byte) []byte {
 		p.inputState = false
 		// 用户输入了Enter，开始结算命令
 		p.parseCmdInput()
-		if p.IsCommandForbidden() {
-			fbdMsg := utils.WrapperWarn(fmt.Sprintf(i18n.T("Command `%s` is forbidden"), p.command))
+		if cmd, ok := p.IsCommandForbidden(); !ok {
+			fbdMsg := utils.WrapperWarn(fmt.Sprintf(i18n.T("Command `%s` is forbidden"), cmd))
 			p.outputBuf.WriteString(fbdMsg)
 			p.srvOutputChan <- []byte("\r\n" + fbdMsg)
 			return []byte{utils.CharCleanLine, '\r'}
@@ -223,12 +223,20 @@ func (p *Parser) SetCMDFilterRules(rules []model.SystemUserFilterRule) {
 	p.cmdFilterRules = rules
 }
 
-func (p *Parser) IsCommandForbidden() bool {
+func (p *Parser) IsCommandForbidden() (string, bool) {
 	fmt.Println("Command is: ", p.command)
-	if p.command == "ls" {
-		return true
+	for _, rule := range p.cmdFilterRules {
+		allowed, cmd := rule.Match(p.command)
+		switch allowed {
+		case model.ActionAllow:
+			return "", true
+		case model.ActionDeny:
+			return cmd, false
+		default:
+
+		}
 	}
-	return false
+	return "", true
 }
 
 func (p *Parser) IsRecvState() bool {

pkg/proxy/proxy.go

@@ -148,8 +148,10 @@ func (p *ProxyServer) Proxy() {
 	}
 	cmdRules := p.GetFilterRules()
 	sw.SetFilterRules(cmdRules)
+	AddSession(sw)
 	_ = sw.Bridge(p.UserConn, srvConn)
 	p.finishSession(sw)
+	RemoveSession(sw)
 }
 
 func (p *ProxyServer) createSession(s *SwitchSession) bool {

pkg/proxy/recorder.go

@@ -14,14 +14,14 @@ import (
 	"cocogo/pkg/model"
 )
 
-func NewCommandRecorder(sess *SwitchSession) (recorder *CommandRecorder) {
-	recorder = &CommandRecorder{sessionID: sess.Id}
+func NewCommandRecorder(sid string) (recorder *CommandRecorder) {
+	recorder = &CommandRecorder{sessionID: sid}
 	recorder.initial()
 	return recorder
 }
 
-func NewReplyRecord(sess *SwitchSession) (recorder *ReplyRecorder) {
-	recorder = &ReplyRecorder{sessionID: sess.Id}
+func NewReplyRecord(sid string) (recorder *ReplyRecorder) {
+	recorder = &ReplyRecorder{sessionID: sid}
 	recorder.initial()
 	return recorder
 }

pkg/proxy/switch.go

@@ -47,8 +47,8 @@ func (s *SwitchSession) Initial() {
 	s.Id = uuid.NewV4().String()
 	s.DateStart = time.Now().UTC().Format("2006-01-02 15:04:05 +0000")
 	s.MaxIdleTime = config.GetConf().MaxIdleTime
-	s.cmdRecorder = NewCommandRecorder(s)
-	s.replayRecorder = NewReplyRecord(s)
+	s.cmdRecorder = NewCommandRecorder(s.Id)
+	s.replayRecorder = NewReplyRecord(s.Id)
 
 	s.parser = newParser()
 
@@ -83,13 +83,12 @@ func (s *SwitchSession) generateCommandResult(command [2]string) *model.Command
 		input = command[0]
 	}
 	i := strings.LastIndexByte(command[1], '\r')
-
-	if i > 1024 {
-		output = output[:1024]
-	} else if i > 0 {
+	if i < 0 {
+		output = command[1]
+	} else if i > 0 && i < 1024 {
 		output = command[1][:i]
 	} else {
-		output = command[1]
+		output = command[1][:1024]
 	}
 
 	return &model.Command{
@@ -134,13 +133,14 @@ func (s *SwitchSession) Bridge(userConn UserConnection, srvConn ServerConnection
 		select {
 		// 检测是否超过最大空闲时间
 		case <-time.After(time.Duration(s.MaxIdleTime) * time.Minute):
-			msg := i18n.T(fmt.Sprintf("Connect idle more than %d minutes, disconnect", s.MaxIdleTime))
+			msg := i18n.T(fmt.Sprintf("\n\nConnect idle more than %d minutes, disconnect", s.MaxIdleTime))
 			msg = utils.WrapperWarn(msg)
 			utils.IgnoreErrWriteString(s.userTran, msg)
 			return
 		// 手动结束
 		case <-s.ctx.Done():
-			msg := i18n.T("Terminated by administrator")
+			msg := i18n.T("\n\rTerminated by administrator")
+			msg = utils.WrapperWarn(msg)
 			utils.IgnoreErrWriteString(userConn, msg)
 			return
 		// 监控窗口大小变化

pkg/service/terminal.go

@@ -75,7 +75,7 @@ func FinishTask(tid string) bool {
 	var res map[string]interface{}
 	data := map[string]bool{"is_finished": true}
 	Url := fmt.Sprintf(FinishTaskURL, tid)
-	err := authClient.Patch(Url, data, res)
+	err := authClient.Patch(Url, data, &res)
 	if err != nil {
 		logger.Error(err)
 		return false

pkg/sshd/server.go

@@ -11,6 +11,8 @@ import (
 	"cocogo/pkg/logger"
 )
 
+var sshServer *ssh.Server
+
 func StartServer() {
 	conf := config.GetConf()
 	hostKey := HostKey{Value: conf.HostKey, Path: conf.HostKeyFile}
@@ -21,7 +23,7 @@ func StartServer() {
 	}
 
 	logger.Infof("Start ssh server at %s:%d", conf.BindHost, conf.SSHPort)
-	srv := ssh.Server{
+	sshServer = &ssh.Server{
 		Addr:                       conf.BindHost + ":" + strconv.Itoa(conf.SSHPort),
 		KeyboardInteractiveHandler: auth.CheckMFA,
 		PasswordHandler:            auth.CheckUserPassword,
@@ -32,6 +34,15 @@ func StartServer() {
 		SubsystemHandlers:          map[string]ssh.SubsystemHandler{},
 	}
 	// Set Auth Handler
-	srv.SetSubsystemHandler("sftp", handler.SftpHandler)
-	logger.Fatal(srv.ListenAndServe())
+	sshServer.SetSubsystemHandler("sftp", handler.SftpHandler)
+	logger.Fatal(sshServer.ListenAndServe())
+}
+
+func StopServer() {
+	err := sshServer.Close()
+	if err != nil {
+		logger.Debugf("ssh server close failed: %s", err.Error())
+	}
+	logger.Debug("Close ssh Server")
+
 }

pkg/utils/util.go

@@ -41,6 +41,6 @@ func WrapperTitle(text string) string {
 }
 
 func WrapperWarn(text string) string {
-	text += "\r\n"
+	text += "\n\r"
 	return WrapperString(text, Red)
 }