[Update] 暂存

pkg/common/client.go

@@ -154,6 +154,11 @@ func (c *Client) Do(method, url string, data, res interface{}, params ...map[str
 		return
 	}
 
+	// If is buffer return the raw response body
+	if buf, ok := res.(*bytes.Buffer); ok {
+		buf.Write(body)
+		return
+	}
 	// Unmarshal response body to result struct
 	if res != nil {
 		err = json.Unmarshal(body, res)

pkg/config/config.go

@@ -2,14 +2,12 @@ package config
 
 import (
 	"encoding/json"
+	"gopkg.in/yaml.v2"
 	"io/ioutil"
 	"log"
 	"os"
 	"strings"
 	"sync"
-	"time"
-
-	"gopkg.in/yaml.v2"
 )
 
 type Config struct {
@@ -23,7 +21,7 @@ type Config struct {
 	ReplayStorage       map[string]string `json:"TERMINAL_REPLAY_STORAGE" yaml:"REPLAY_STORAGE"`
 	SessionKeepDuration int               `json:"TERMINAL_SESSION_KEEP_DURATION"`
 	TelnetRegex         string            `json:"TERMINAL_TELNET_REGEX"`
-	MaxIdleTime         time.Duration     `json:"SECURITY_MAX_IDLE_TIME"`
+	MaxIdleTime         int               `json:"SECURITY_MAX_IDLE_TIME"`
 	SftpRoot            string            `json:"TERMINAL_SFTP_ROOT" yaml:"SFTP_ROOT"`
 	Name                string            `yaml:"NAME"`
 	SecretKey           string            `yaml:"SECRET_KEY"`
@@ -37,20 +35,29 @@ type Config struct {
 	AccessKey           string            `yaml:"ACCESS_KEY"`
 	AccessKeyFile       string            `yaml:"ACCESS_KEY_FILE"`
 	LogLevel            string            `yaml:"LOG_LEVEL"`
-	HeartbeatDuration   time.Duration     `yaml:"HEARTBEAT_INTERVAL"`
+	HeartbeatDuration   int               `yaml:"HEARTBEAT_INTERVAL"`
 	RootPath            string            `yaml:"ROOT_PATH"`
 	Comment             string            `yaml:"COMMENT"`
 	Language            string            `yaml:"LANG"`
+	LanguageCode        string            `yaml:"LANGUAGE_CODE"` // Abandon
 
-	mux sync.RWMutex
+	Mux sync.RWMutex
 }
 
 func (c *Config) EnsureConfigValid() {
+	// 兼容原来config
+	if c.LanguageCode != "" && c.Language == "" {
+		c.Language = c.LanguageCode
+	}
+	// 确保至少有一个认证
+	if !c.PublicKeyAuth && !c.PasswordAuth {
+		c.PasswordAuth = true
+	}
 }
 
 func (c *Config) LoadFromYAML(body []byte) error {
-	c.mux.Lock()
-	defer c.mux.Unlock()
+	c.Mux.Lock()
+	defer c.Mux.Unlock()
 	err := yaml.Unmarshal(body, c)
 	if err != nil {
 		log.Printf("Load yaml error: %v", err)
@@ -67,8 +74,8 @@ func (c *Config) LoadFromYAMLPath(filepath string) error {
 }
 
 func (c *Config) LoadFromJSON(body []byte) error {
-	c.mux.Lock()
-	defer c.mux.Unlock()
+	c.Mux.Lock()
+	defer c.Mux.Unlock()
 	err := json.Unmarshal(body, c)
 	if err != nil {
 		log.Printf("Config load yaml error")
@@ -117,6 +124,6 @@ var Conf = &Config{
 	RootPath:       rootPath,
 	Comment:        "Coco",
 	Language:       "zh",
-	ReplayStorage:  map[string]string{},
-	CommandStorage: map[string]string{},
+	ReplayStorage:  map[string]string{"TYPE": "server"},
+	CommandStorage: map[string]string{"TYPE": "server"},
 }

pkg/handler/session.go

@@ -9,9 +9,6 @@ import (
 	"sync"
 
 	"github.com/gliderlabs/ssh"
-	//"github.com/olekukonko/tablewriter"
-	//"github.com/satori/go.uuid"
-	//"github.com/xlab/treeprint"
 	"github.com/olekukonko/tablewriter"
 	"github.com/xlab/treeprint"
 	"golang.org/x/crypto/ssh/terminal"
@@ -21,23 +18,20 @@ import (
 	"cocogo/pkg/model"
 	"cocogo/pkg/proxy"
 	"cocogo/pkg/service"
-	//"cocogo/pkg/transport"
-	//"cocogo/pkg/userhome"
+	"cocogo/pkg/utils"
 )
 
 func SessionHandler(sess ssh.Session) {
-	_, _, ptyOk := sess.Pty()
-	if ptyOk {
+	pty, _, ok := sess.Pty()
+	if ok {
 		ctx, cancel := cctx.NewContext(sess)
+		defer cancel()
 		handler := newInteractiveHandler(sess, ctx.User())
-		logger.Infof("New connection from: %s %s", sess.User(), sess.RemoteAddr().String())
+		logger.Debugf("User Request pty: %s %s", sess.User(), pty.Term)
 		handler.Dispatch(ctx)
-		cancel()
 	} else {
-		_, err := io.WriteString(sess, "No PTY requested.\n")
-		if err != nil {
-			return
-		}
+		utils.IgnoreErrWriteString(sess, "No PTY requested.\n")
+		return
 	}
 }
 

pkg/model/terminal.go

@@ -13,19 +13,6 @@ type Terminal struct {
 	} `json:"service_account"`
 }
 
-type TerminalConf struct {
-	AssetListPageSize   string            `json:"TERMINAL_ASSET_LIST_PAGE_SIZE"`
-	AssetListSortBy     string            `json:"TERMINAL_ASSET_LIST_SORT_BY"`
-	HeaderTitle         string            `json:"TERMINAL_HEADER_TITLE"`
-	HostKey             string            `json:"TERMINAL_HOST_KEY" yaml:"HOST_KEY"`
-	PasswordAuth        bool              `json:"TERMINAL_PASSWORD_AUTH" yaml:"PASSWORD_AUTH"`
-	PublicKeyAuth       bool              `json:"TERMINAL_PUBLIC_KEY_AUTH" yaml:"PUBLIC_KEY_AUTH"`
-	CommandStorage      map[string]string `json:"TERMINAL_COMMAND_STORAGE"`
-	ReplayStorage       map[string]string `json:"TERMINAL_REPLAY_STORAGE" yaml:"REPLAY_STORAGE"`
-	SessionKeepDuration int               `json:"TERMINAL_SESSION_KEEP_DURATION"`
-	TelnetRegex         string            `json:"TERMINAL_TELNET_REGEX"`
-}
-
 type TerminalTask struct {
 	Id         string `json:"id"`
 	Name       string `json:"name"`

pkg/proxy/parser.go

@@ -2,11 +2,9 @@ package proxy
 
 import (
 	"bytes"
-	"sync"
-	"time"
-
 	"cocogo/pkg/logger"
 	"cocogo/pkg/model"
+	"sync"
 )
 
 var (
@@ -26,11 +24,12 @@ var (
 
 // Parse 解析用户输入输出, 拦截过滤用户输入输出
 type Parser struct {
-	session   *SwitchSession
 	inputBuf  *bytes.Buffer
 	cmdBuf    *bytes.Buffer
 	outputBuf *bytes.Buffer
 
+	cmdCh chan *[2]string
+
 	inputInitial  bool
 	inputPreState bool
 	inputState    bool
@@ -44,9 +43,7 @@ type Parser struct {
 	cmdOutputParser *CmdParser
 	counter         int
 
-	cmdFilterRules  []model.SystemUserFilterRule
-	commandRecorder *CommandRecorder
-	replayRecorder  *ReplyRecorder
+	cmdFilterRules []model.SystemUserFilterRule
 }
 
 func (p *Parser) Initial() {
@@ -60,8 +57,6 @@ func (p *Parser) Initial() {
 	p.cmdOutputParser = &CmdParser{}
 	p.cmdInputParser.Initial()
 	p.cmdOutputParser.Initial()
-	record := NewCommandRecorder(p.session)
-	p.SetCommandRecorder(record)
 }
 
 // Todo: parseMultipleInput 依然存在问题
@@ -81,8 +76,7 @@ func (p *Parser) parseInputState(b []byte) {
 		// 用户又开始输入，并上次不处于输入状态，开始结算上次命令的结果
 		if !p.inputPreState {
 			p.parseCmdOutput()
-			// 开始记录命令
-			p.recordCommand()
+			p.cmdCh <- &[2]string{p.command, p.output}
 		}
 	}
 }
@@ -92,7 +86,6 @@ func (p *Parser) parseCmdInput() {
 	p.command = p.cmdInputParser.Parse(data)
 	p.cmdBuf.Reset()
 	p.inputBuf.Reset()
-	p.counter += 1
 }
 
 func (p *Parser) parseCmdOutput() {
@@ -174,30 +167,3 @@ func (p *Parser) ParseServerOutput(b []byte) []byte {
 func (p *Parser) SetCMDFilterRules(rules []model.SystemUserFilterRule) {
 	p.cmdFilterRules = rules
 }
-
-func (p *Parser) SetReplayRecorder(recorder *ReplyRecorder) {
-	p.replayRecorder = recorder
-
-}
-
-func (p *Parser) recordCommand() {
-	cmd := &Command{
-		SessionId:  p.session.Id,
-		OrgId:      p.session.Org,
-		Input:      p.command,
-		Output:     p.output,
-		User:       p.session.User,
-		Server:     p.session.Server,
-		SystemUser: p.session.SystemUser,
-		Timestamp:  time.Now(),
-	}
-	p.commandRecorder.Record(cmd)
-}
-
-func (p *Parser) SetCommandRecorder(recorder *CommandRecorder) {
-	p.commandRecorder = recorder
-}
-
-func (p *Parser) recordReplay(b []byte) {
-	p.replayRecorder.Record(b)
-}

pkg/proxy/switch.go

 package proxy
 
 import (
+	"cocogo/pkg/model"
 	"context"
 	"time"
 
@@ -14,7 +15,6 @@ func NewSwitchSession(userConn UserConnection, serverConn ServerConnection) (sw
 	parser := new(Parser)
 	parser.Initial()
 	sw = &SwitchSession{userConn: userConn, serverConn: serverConn, parser: parser}
-	parser.session = sw
 	return sw
 }
 
@@ -32,7 +32,14 @@ type SwitchSession struct {
 	Finished   bool      `json:"is_finished"`
 	Closed     bool
 
-	parser     *Parser
+	srvChan  chan []byte
+	userChan chan []byte
+
+	cmdFilterRules []model.SystemUserFilterRule
+	cmdRecorder    *CommandRecorder
+	replayRecorder *ReplayStorage
+	parser         *Parser
+
 	userConn   UserConnection
 	serverConn ServerConnection
 	userTran   Transport

pkg/service/init.go

@@ -26,6 +26,7 @@ func Initial() {
 	_ = ak.Load()
 	authClient.Auth = ak
 	validateAccessAuth()
+	go KeepSyncConfigWithServer()
 }
 
 func validateAccessAuth() {
@@ -37,7 +38,7 @@ func validateAccessAuth() {
 			break
 		}
 		if err != nil {
-			msg := `Connect server error or access key is invalid, remove %s run again`
+			msg := "Connect server error or access key is invalid, remove %s run again"
 			logger.Errorf(msg, config.Conf.AccessKeyFile)
 		}
 		if user.Role != "App" {
@@ -49,5 +50,28 @@ func validateAccessAuth() {
 			os.Exit(1)
 		}
 	}
+}
+
+func MustLoadServerConfigOnce() {
+
+}
+
+func LoadConfigFromServer(conf *config.Config) (err error) {
+	conf.Mux.Lock()
+	defer conf.Mux.Unlock()
+	err = authClient.Get(TerminalConfigURL, conf)
+	if err != nil {
+		logger.Warn("Sync config with server error: ", err)
+	}
+	return err
+}
 
+func KeepSyncConfigWithServer() {
+	for {
+		err := LoadConfigFromServer(config.Conf)
+		if err != nil {
+			logger.Warn("Sync config with server error: ", err)
+		}
+		time.Sleep(60 * time.Second)
+	}
 }

pkg/service/terminal.go

@@ -79,14 +79,6 @@ func FinishTask(tid string) bool {
 	return true
 }
 
-func LoadConfigFromServer() (res model.TerminalConf) {
-	err := authClient.Get(TerminalConfigURL, &res)
-	if err != nil {
-		logger.Error(err)
-	}
-	return
-}
-
 func PushSessionReplay(sessionID, gZipFile string) {
 
 }

pkg/sshd/server.go

@@ -15,7 +15,6 @@ var conf = config.Conf
 
 func StartServer() {
 	hostKey := HostKey{Value: conf.HostKey, Path: conf.HostKeyFile}
-
 	logger.Debug("Loading host key")
 	signer, err := hostKey.Load()
 	if err != nil {