fix conflicts

Gopkg.lock

@@ -182,6 +182,14 @@
   pruneopts = "UT"
   revision = "a5b02f93d862f065920dd6a40dddc66b60d0dec4"
 
+[[projects]]
+  digest = "1:c805e517269b0ba4c21ded5836019ed7d16953d4026cb7d00041d039c7906be9"
+  name = "gopkg.in/natefinch/lumberjack.v2"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "a96e63847dc3c67d17befa69c303767e2f84e54f"
+  version = "v2.1"
+
 [[projects]]
   digest = "1:4d2e5a73dc1500038e504a8d78b986630e3626dc027bc030ba5c75da257cdb96"
   name = "gopkg.in/yaml.v2"
@@ -207,6 +215,7 @@
     "github.com/xlab/treeprint",
     "golang.org/x/crypto/ssh",
     "golang.org/x/crypto/ssh/terminal",
+    "gopkg.in/natefinch/lumberjack.v2",
     "gopkg.in/yaml.v2",
   ]
   solver-name = "gps-cdcl"

Gopkg.toml

@@ -70,3 +70,7 @@
   go-tests = true
   unused-packages = true
 
+
+[[constraint]]
+  name = "gopkg.in/natefinch/lumberjack.v2"
+  version = "2.1.0"

cmd/coco.go

 package main
 
 import (
-	"cocogo/pkg/config"
-	"cocogo/pkg/service"
-	"cocogo/pkg/sshd"
+	"cocogo/pkg/coco"
 )
 
-func init() {
-	config.Initial()
-}
-
 func main() {
-	service.Initial()
-	sshd.StartServer()
+	app := &coco.Coco{}
+	app.Start()
 }

cmd/config_example.yml

 NAME: coco2
 BOOTSTRAP_TOKEN: PleaseChangeMe
 CORE_HOST: http://127.0.0.1:8080
+LOG_LEVEL: INFO

pkg/coco/coco.go

+package coco
+
+import (
+	"fmt"
+	"time"
+
+	"cocogo/pkg/sshd"
+)
+
+const version = "1.4.0"
+
+type Coco struct {
+}
+
+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()
+}
+
+func (c *Coco) Stop() {
+
+}

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

 package config
 
 import (
-	"cocogo/pkg/logger"
 	"encoding/json"
 	"io/ioutil"
+	"log"
 	"os"
 	"strings"
 	"sync"
 	"time"
 
-	log "github.com/sirupsen/logrus"
 	"gopkg.in/yaml.v2"
 )
 
@@ -54,7 +53,7 @@ func (c *Config) LoadFromYAML(body []byte) error {
 	defer c.mux.Unlock()
 	err := yaml.Unmarshal(body, c)
 	if err != nil {
-		log.Errorf("Load yaml error: %v", err)
+		log.Printf("Load yaml error: %v", err)
 	}
 	return err
 }
@@ -62,8 +61,7 @@ func (c *Config) LoadFromYAML(body []byte) error {
 func (c *Config) LoadFromYAMLPath(filepath string) error {
 	body, err := ioutil.ReadFile(filepath)
 	if err != nil {
-		log.Errorf("Not found file: %s", filepath)
-		os.Exit(1)
+		log.Printf("Not found file: %s", filepath)
 	}
 	return c.LoadFromYAML(body)
 }
@@ -73,8 +71,7 @@ func (c *Config) LoadFromJSON(body []byte) error {
 	defer c.mux.Unlock()
 	err := json.Unmarshal(body, c)
 	if err != nil {
-		logger.Error("Config load yaml error")
-		os.Exit(1)
+		log.Printf("Config load yaml error")
 	}
 	return nil
 }
@@ -88,7 +85,7 @@ func (c *Config) LoadFromEnv() error {
 	}
 	envYAML, err := yaml.Marshal(envMap)
 	if err != nil {
-		log.Errorf("Error occur: %v", err)
+		log.Fatalf("Error occur: %v", err)
 	}
 	return c.LoadFromYAML(envYAML)
 }

pkg/handler/session.go

@@ -36,14 +36,14 @@ func SessionHandler(sess ssh.Session) {
 	}
 }
 
-func newInteractiveHandler(sess ssh.Session, user *model.User) *InteractiveHandler {
+func newInteractiveHandler(sess ssh.Session, user *model.User) *interactiveHandler {
 	term := terminal.NewTerminal(sess, "Opt> ")
-	handler := &InteractiveHandler{sess: sess, user: user, term: term}
+	handler := &interactiveHandler{sess: sess, user: user, term: term}
 	handler.Initial()
 	return handler
 }
 
-type InteractiveHandler struct {
+type interactiveHandler struct {
 	sess ssh.Session
 	user *model.User
 	term *terminal.Terminal
@@ -56,18 +56,18 @@ type InteractiveHandler struct {
 	mu               *sync.RWMutex
 }
 
-func (h *InteractiveHandler) Initial() {
+func (h *interactiveHandler) Initial() {
 	h.displayBanner()
 	h.loadUserAssets()
 	h.loadUserAssetNodes()
 	h.searchResult = h.assets
 }
 
-func (h *InteractiveHandler) displayBanner() {
+func (h *interactiveHandler) displayBanner() {
 	displayBanner(h.sess, h.user.Name)
 }
 
-func (h *InteractiveHandler) watchWinSizeChange(winCh <-chan ssh.Window, done <-chan struct{}) {
+func (h *interactiveHandler) watchWinSizeChange(winCh <-chan ssh.Window, done <-chan struct{}) {
 	for {
 		select {
 		case <-done:
@@ -77,13 +77,13 @@ func (h *InteractiveHandler) watchWinSizeChange(winCh <-chan ssh.Window, done <-
 			if !ok {
 				return
 			}
-			logger.Debugf("Term change: %d*%d", win.Height, win.Width)
+			logger.Debugf("Term window size change: %d*%d", win.Height, win.Width)
 			_ = h.term.SetSize(win.Width, win.Height)
 		}
 	}
 }
 
-func (h *InteractiveHandler) Dispatch(ctx cctx.Context) {
+func (h *interactiveHandler) Dispatch(ctx cctx.Context) {
 	_, winCh, _ := h.sess.Pty()
 	for {
 		doneChan := make(chan struct{})
@@ -147,7 +147,7 @@ func (h *InteractiveHandler) Dispatch(ctx cctx.Context) {
 	}
 }
 
-func (h *InteractiveHandler) chooseSystemUser(systemUsers []model.SystemUser) model.SystemUser {
+func (h *interactiveHandler) chooseSystemUser(systemUsers []model.SystemUser) model.SystemUser {
 	length := len(systemUsers)
 	switch length {
 	case 0:
@@ -195,7 +195,7 @@ func (h *InteractiveHandler) chooseSystemUser(systemUsers []model.SystemUser) mo
 }
 
 // 当资产的数量为1的时候，就进行代理转化
-func (h *InteractiveHandler) displayAssetsOrProxy(assets []model.Asset) {
+func (h *interactiveHandler) displayAssetsOrProxy(assets []model.Asset) {
 	if len(assets) == 1 {
 		logger.Debug(assets[0].SystemUsers)
 		systemUser := h.chooseSystemUser(assets[0].SystemUsers)
@@ -207,7 +207,7 @@ func (h *InteractiveHandler) displayAssetsOrProxy(assets []model.Asset) {
 	}
 }
 
-func (h *InteractiveHandler) displayAssets(assets model.AssetList) {
+func (h *interactiveHandler) displayAssets(assets model.AssetList) {
 	if len(assets) == 0 {
 		_, _ = io.WriteString(h.term, "\r\n No Assets\r\n\r")
 	} else {
@@ -228,7 +228,7 @@ func (h *InteractiveHandler) displayAssets(assets model.AssetList) {
 
 }
 
-func (h *InteractiveHandler) displayNodes(nodes []model.Node) {
+func (h *interactiveHandler) displayNodes(nodes []model.Node) {
 	tree := ConstructAssetNodeTree(nodes)
 	tipHeaderMsg := "\r\nNode: [ ID.Name(Asset amount) ]"
 	tipEndMsg := "Tips: Enter g+NodeID to display the host under the node, such as g1\r\n\r"
@@ -242,22 +242,22 @@ func (h *InteractiveHandler) displayNodes(nodes []model.Node) {
 
 }
 
-func (h *InteractiveHandler) refreshAssetsAndNodesData() {
+func (h *interactiveHandler) refreshAssetsAndNodesData() {
 	_, err := io.WriteString(h.sess, "Refresh done\r\n")
 	if err != nil {
 		logger.Error("refresh Assets  Nodes err:", err)
 	}
 }
 
-func (h *InteractiveHandler) loadUserAssets() {
+func (h *interactiveHandler) loadUserAssets() {
 	h.assets = service.GetUserAssets(h.user.ID, "1")
 }
 
-func (h *InteractiveHandler) loadUserAssetNodes() {
+func (h *interactiveHandler) loadUserAssetNodes() {
 	h.nodes = service.GetUserNodes(h.user.ID, "1")
 }
 
-func (h *InteractiveHandler) searchAsset(key string) (assets []model.Asset) {
+func (h *interactiveHandler) searchAsset(key string) (assets []model.Asset) {
 	if indexNum, err := strconv.Atoi(key); err == nil {
 		if indexNum > 0 && indexNum <= len(h.searchResult) {
 			assets = []model.Asset{h.searchResult[indexNum-1]}
@@ -281,7 +281,7 @@ func (h *InteractiveHandler) searchAsset(key string) (assets []model.Asset) {
 	return assets
 }
 
-func (h *InteractiveHandler) searchNodeAssets(num int) (assets []model.Asset) {
+func (h *interactiveHandler) searchNodeAssets(num int) (assets []model.Asset) {
 	if num > len(h.nodes) || num == 0 {
 		return assets
 	}
@@ -289,7 +289,7 @@ func (h *InteractiveHandler) searchNodeAssets(num int) (assets []model.Asset) {
 
 }
 
-func (h *InteractiveHandler) Proxy(ctx context.Context) {
+func (h *interactiveHandler) Proxy(ctx context.Context) {
 	//h.assetSelect = &model.Asset{Hostname: "centos", Port: 22, Ip: "192.168.244.185", Protocol: "ssh"}
 	//h.systemUserSelect = &model.SystemUser{Id: "5dd8b5a0-8cdb-4857-8629-faf811c525e1", Name: "web", Username: "root", Password: "redhat", Protocol: "telnet"}
 

pkg/logger/formatter.go

+/*
+Using https://github.com/t-tomalak/logrus-easy-formatter/ as formatter
+*/
+
+package logger
+
+import (
+	"strings"
+	"time"
+
+	"github.com/sirupsen/logrus"
+)
+
+const (
+	// Default log format will output [INFO]: 2006-01-02T15:04:05Z07:00 - Log message
+	defaultLogFormat       = "[%lvl%]: %time% - %msg%"
+	defaultTimestampFormat = time.RFC3339
+)
+
+// Formatter implements logrus.Formatter interface.
+type Formatter struct {
+	// Timestamp format
+	TimestampFormat string
+	// Available standard keys: time, msg, lvl
+	// Also can include custom fields but limited to strings.
+	// All of fields need to be wrapped inside %% i.e %time% %msg%
+	LogFormat string
+
+	// Disables the truncation of the level text to 4 characters.
+	DisableLevelTruncation bool
+}
+
+// Format building log message.
+func (f *Formatter) Format(entry *logrus.Entry) ([]byte, error) {
+	output := f.LogFormat
+	if output == "" {
+		output = defaultLogFormat
+	}
+
+	timestampFormat := f.TimestampFormat
+	if timestampFormat == "" {
+		timestampFormat = defaultTimestampFormat
+	}
+
+	output = strings.Replace(output, "%time%", entry.Time.Format(timestampFormat), 1)
+	output = strings.Replace(output, "%msg%", entry.Message, 1)
+	level := strings.ToUpper(entry.Level.String())
+	if !f.DisableLevelTruncation {
+		level = level[:4]
+	}
+	output = strings.Replace(output, "%lvl%", level, 1)
+
+	for k, v := range entry.Data {
+		if s, ok := v.(string); ok {
+			output = strings.Replace(output, "%"+k+"%", s, 1)
+		}
+	}
+	output += "\n"
+
+	return []byte(output), nil
+}

pkg/logger/logger.go

 package logger
 
 import (
+	"cocogo/pkg/common"
+	"fmt"
 	"os"
+	"path"
+	"strings"
 
 	"github.com/sirupsen/logrus"
+
+	"cocogo/pkg/config"
 )
 
 var logger = logrus.New()
+var logLevels = map[string]logrus.Level{
+	"DEBUG": logrus.DebugLevel,
+	"INFO":  logrus.InfoLevel,
+	"WARN":  logrus.WarnLevel,
+	"ERROR": logrus.ErrorLevel,
+}
 
-func init() {
-	customFormatter := &logrus.TextFormatter{
-		DisableColors:          false,
-		FullTimestamp:          true,
-		DisableLevelTruncation: false,
+func Initial() {
+	formatter := &Formatter{
+		LogFormat:       "%time% [%lvl%] %msg%",
+		TimestampFormat: "2006-01-02 15:04:05",
+	}
+	level, ok := logLevels[strings.ToUpper(config.Conf.LogLevel)]
+	if !ok {
+		level = logrus.InfoLevel
 	}
-	customFormatter.TimestampFormat = "2006-01-02 15:04:05"
-	logger.SetFormatter(customFormatter)
 
 	// Output to stdout instead of the default stderr
 	// Can be any io.Writer, see below for File example
+	logger.SetFormatter(formatter)
 	logger.SetOutput(os.Stdout)
+	logger.SetLevel(level)
+
+	// Output to file
+	logFilePath := path.Join(config.Conf.RootPath, "logs", "coco.log")
+	logDirPath := path.Dir(logFilePath)
+	if common.FileExists(logDirPath) {
+		err := os.MkdirAll(logDirPath, os.ModePerm)
+		if err != nil {
+			fmt.Printf("Create log dir %s error: %s\n", logDirPath, err)
+			return
+		}
+	}
 
-	// Only logger the warning severity or above.
-	logger.SetLevel(logrus.DebugLevel)
+	rotateFileHook, err := NewRotateFileHook(RotateFileConfig{
+		Filename:   logFilePath,
+		MaxSize:    50,
+		MaxBackups: 7,
+		MaxAge:     7,
+		LocalTime:  true,
+		Level:      level,
+		Formatter:  formatter,
+	})
+	if err != nil {
+		fmt.Printf("Create log rotate hook error: %s\n", err)
+		return
+	}
+	logger.AddHook(rotateFileHook)
 }
 
 func Debug(args ...interface{}) {

pkg/logger/rotate.go

+package logger
+
+import (
+	"io"
+
+	"github.com/sirupsen/logrus"
+	"gopkg.in/natefinch/lumberjack.v2"
+)
+
+type RotateFileConfig struct {
+	Filename   string
+	MaxSize    int
+	MaxBackups int
+	MaxAge     int
+	Level      logrus.Level
+	LocalTime  bool
+	Formatter  logrus.Formatter
+}
+
+type RotateFileHook struct {
+	Config    RotateFileConfig
+	logWriter io.Writer
+}
+
+func NewRotateFileHook(config RotateFileConfig) (logrus.Hook, error) {
+	hook := RotateFileHook{
+		Config: config,
+	}
+	hook.logWriter = &lumberjack.Logger{
+		Filename:   config.Filename,
+		MaxSize:    config.MaxSize,
+		MaxBackups: config.MaxBackups,
+		MaxAge:     config.MaxAge,
+		LocalTime:  config.LocalTime,
+	}
+
+	return &hook, nil
+}
+
+func (hook *RotateFileHook) Levels() []logrus.Level {
+	return logrus.AllLevels[:hook.Config.Level+1]
+}
+
+func (hook *RotateFileHook) Fire(entry *logrus.Entry) (err error) {
+	b, err := hook.Config.Formatter.Format(entry)
+	if err != nil {
+		return err
+	}
+	hook.logWriter.Write(b)
+	return nil
+}

pkg/proxy/parser.go

@@ -26,7 +26,7 @@ var (
 
 // Parse 解析用户输入输出, 拦截过滤用户输入输出
 type Parser struct {
-	session   *Session
+	session   *SwitchSession
 	inputBuf  *bytes.Buffer
 	cmdBuf    *bytes.Buffer
 	outputBuf *bytes.Buffer

pkg/proxy/proxy.go

@@ -132,7 +132,7 @@ func (p *ProxyServer) Proxy() {
 		return
 	}
 
-	sw := NewSwitch(p.UserConn, srvConn)
+	sw := NewSwitchSession(p.UserConn, srvConn)
 	cmdRules, err := service.GetSystemUserFilterRules(p.SystemUser.Id)
 	if err != nil {
 		logger.Error("Get system user filter rule error: ", err)

pkg/proxy/recorder.go

@@ -16,10 +16,10 @@ import (
 )
 
 type CommandRecorder struct {
-	Session *Session
+	Session *SwitchSession
 }
 
-func NewCommandRecorder(sess *Session) (recorder *CommandRecorder) {
+func NewCommandRecorder(sess *SwitchSession) (recorder *CommandRecorder) {
 	return &CommandRecorder{Session: sess}
 }
 

pkg/proxy/switch.go

@@ -10,15 +10,15 @@ import (
 	"cocogo/pkg/logger"
 )
 
-func NewSwitch(userConn UserConnection, serverConn ServerConnection) (sw *Session) {
+func NewSwitchSession(userConn UserConnection, serverConn ServerConnection) (sw *SwitchSession) {
 	parser := new(Parser)
 	parser.Initial()
-	sw = &Session{userConn: userConn, serverConn: serverConn, parser: parser}
+	sw = &SwitchSession{userConn: userConn, serverConn: serverConn, parser: parser}
 	parser.session = sw
 	return sw
 }
 
-type Session struct {
+type SwitchSession struct {
 	Id         string
 	User       string    `json:"user"`
 	Server     string    `json:"asset"`
@@ -40,7 +40,7 @@ type Session struct {
 	cancelFunc context.CancelFunc
 }
 
-func (s *Session) Initial() {
+func (s *SwitchSession) Initial() {
 	s.Id = uuid.NewV4().String()
 	s.User = s.userConn.User()
 	s.Server = s.serverConn.Name()
@@ -50,15 +50,15 @@ func (s *Session) Initial() {
 	s.DateStart = time.Now()
 }
 
-func (s *Session) preBridge() {
+func (s *SwitchSession) preBridge() {
 
 }
 
-func (s *Session) postBridge() {
+func (s *SwitchSession) postBridge() {
 
 }
 
-func (s *Session) watchWindowChange(ctx context.Context, winCh <-chan ssh.Window) {
+func (s *SwitchSession) watchWindowChange(ctx context.Context, winCh <-chan ssh.Window) {
 	defer func() {
 		logger.Debug("Watch window change routine end")
 	}()
@@ -80,7 +80,7 @@ func (s *Session) watchWindowChange(ctx context.Context, winCh <-chan ssh.Window
 	}
 }
 
-func (s *Session) readUserToServer(ctx context.Context) {
+func (s *SwitchSession) readUserToServer(ctx context.Context) {
 	defer func() {
 		logger.Debug("Read user to server end")
 	}()
@@ -102,7 +102,7 @@ func (s *Session) readUserToServer(ctx context.Context) {
 	}
 }
 
-func (s *Session) readServerToUser(ctx context.Context) {
+func (s *SwitchSession) readServerToUser(ctx context.Context) {
 	defer func() {
 		logger.Debug("Read server to user end")
 	}()
@@ -124,7 +124,7 @@ func (s *Session) readServerToUser(ctx context.Context) {
 	}
 }
 
-func (s *Session) Bridge() (err error) {
+func (s *SwitchSession) Bridge() (err error) {
 	winCh := s.userConn.WinCh()
 	ctx, cancel := context.WithCancel(context.Background())
 	s.cancelFunc = cancel

pkg/service/init.go

@@ -25,10 +25,10 @@ func Initial() {
 	ak := AccessKey{Value: config.Conf.AccessKey, Path: keyPath}
 	_ = ak.Load()
 	authClient.Auth = ak
-	ValidateAccessAuth()
+	validateAccessAuth()
 }
 
-func ValidateAccessAuth() {
+func validateAccessAuth() {
 	maxTry := 30
 	count := 0
 	for count < maxTry {

pkg/sshd/server.go

 package sshd
 
 import (
-	"fmt"
 	"strconv"
-	"time"
 
 	"github.com/gliderlabs/ssh"
 
@@ -13,34 +11,36 @@ import (
 	"cocogo/pkg/logger"
 )
 
-const version = "v1.4.0"
-
-var (
-	conf = config.Conf
-)
+var conf = config.Conf
 
 func StartServer() {
-	logger.Debug("Load host key")
 	hostKey := HostKey{Value: conf.HostKey, Path: conf.HostKeyFile}
+
+	logger.Debug("Loading host key")
 	signer, err := hostKey.Load()
 	if err != nil {
 		logger.Fatal("Load host key error: ", err)
 	}
-	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.Printf("Start ssh server at %s:%d\n", conf.BindHost, conf.SSHPort)
-	fmt.Println("Quit the server with CONTROL-C.")
 
+	logger.Infof("Start ssh server at %s:%d", conf.BindHost, conf.SSHPort)
 	srv := ssh.Server{
 		Addr:                       conf.BindHost + ":" + strconv.Itoa(conf.SSHPort),
-		PasswordHandler:            auth.CheckUserPassword,
-		PublicKeyHandler:           auth.CheckUserPublicKey,
 		KeyboardInteractiveHandler: auth.CheckMFA,
 		NextAuthMethodsHandler:     auth.CheckUserNeedMFA,
 		HostSigners:                []ssh.Signer{signer},
 		Handler:                    handler.SessionHandler,
 		SubsystemHandlers:          map[string]ssh.SubsystemHandler{},
 	}
+	// Set Auth Handler
+	if conf.PasswordAuth {
+		srv.PasswordHandler = auth.CheckUserPassword
+	}
+	if conf.PublicKeyAuth {
+		srv.PublicKeyHandler = auth.CheckUserPublicKey
+	}
+	if !conf.PasswordAuth && !conf.PublicKeyAuth {
+		srv.PasswordHandler = auth.CheckUserPassword
+	}
 	srv.SetSubsystemHandler("sftp", handler.SftpHandler)
 	logger.Fatal(srv.ListenAndServe())
 }