[Update] 提交1

cmd/coco.go

 package main
 
 import (
-	"cocogo/pkg/auth"
 	"cocogo/pkg/config"
 	"cocogo/pkg/sshd"
 )
 
 func init() {
 	config.Initial()
-	auth.Initial()
-	sshd.Initial()
 }
 
 func main() {

pkg/auth/server.go

 package auth
 
 import (
+	gossh "golang.org/x/crypto/ssh"
+
 	"github.com/gliderlabs/ssh"
 
 	"cocogo/pkg/common"
 	"cocogo/pkg/service"
 )
 
+func CheckUserPassword(ctx ssh.Context, password string) bool {
+	return true
+}
+
 func CheckUserPublicKey(ctx ssh.Context, key ssh.PublicKey) bool {
 	username := ctx.User()
 	b := key.Marshal()
@@ -18,5 +24,8 @@ func CheckUserPublicKey(ctx ssh.Context, key ssh.PublicKey) bool {
 	}
 	ctx.SetValue("LoginUser", authUser)
 	return true
+}
 
+func CheckMFA(ctx ssh.Context, challenger gossh.KeyboardInteractiveChallenge) bool {
+	return true
 }

pkg/common/sshutil.go

+package common

pkg/config/config.go

@@ -26,6 +26,7 @@ type Config struct {
 	TelnetRegex         string            `json:"TERMINAL_TELNET_REGEX"`
 	MaxIdleTime         time.Duration     `json:"SECURITY_MAX_IDLE_TIME"`
 	Name                string            `yaml:"NAME"`
+	HostKeyFile         string            `yaml:"HOST_KEY_FILE"`
 	CoreHost            string            `yaml:"CORE_HOST"`
 	BootstrapToken      string            `yaml:"BOOTSTRAP_TOKEN"`
 	BindHost            string            `yaml:"BIND_HOST"`

pkg/logger/logger.go

@@ -40,3 +40,7 @@ func Error(args ...interface{}) {
 func Panic(args ...interface{}) {
 	logrus.Panic(args...)
 }
+
+func Fatal(args ...interface{}) {
+	logrus.Fatal(args...)
+}

pkg/record/replay.go

 package record
 
 import (
-	"cocogo/pkg/auth"
-	"cocogo/pkg/config"
-	"cocogo/pkg/storage"
 	"compress/gzip"
 	"context"
 	"encoding/json"
@@ -13,10 +10,15 @@ import (
 	"path/filepath"
 	"strings"
 	"time"
+
+	"cocogo/pkg/config"
+	"cocogo/pkg/storage"
 )
 
+var conf = config.Conf
+
 func NewReplyRecord(sessionID string) *Reply {
-	rootPath := config.GetGlobalConfig().RootPath
+	rootPath := conf.RootPath
 	currentData := time.Now().UTC().Format("2006-01-02")
 	gzFileName := sessionID + ".replay.gz"
 	absFilePath := filepath.Join(rootPath, "data", "replays", currentData, sessionID)
@@ -52,9 +54,9 @@ func (r *Reply) Record(b []byte) {
 }
 
 func (r *Reply) StartRecord() {
-	auth.MakeSureDirExit(r.absFilePath)
-	r.WriteF, _ = os.Create(r.absFilePath)
-	_, _ = r.WriteF.Write([]byte("{"))
+	//auth.MakeSureDirExit(r.absFilePath)
+	//r.WriteF, _ = os.Create(r.absFilePath)
+	//_, _ = r.WriteF.Write([]byte("{"))
 }
 
 func (r *Reply) EndRecord(ctx context.Context) {

pkg/service/assets.go

 package service
 
 import (
-	"fmt"
-
 	"cocogo/pkg/logger"
 	"cocogo/pkg/model"
 )
@@ -20,12 +18,12 @@ import (
 //
 //}
 //
-func GetSystemUserAuthInfo(systemUserID string) {
+func GetSystemUserAssetAuthInfo(systemUserID, assetID string) (info model.SystemUserAuthInfo, err error) {
 	var authInfo model.SystemUserAuthInfo
-	err := client.Get("systemUserAuthInfo", nil, &authInfo)
+	err = Client.Get("systemUserAuthInfo", nil, &authInfo)
 	if err != nil {
 		logger.Info("get User Assets Groups err:", err)
 		return
 	}
-	fmt.Println(authInfo)
+	return
 }

pkg/service/client.go

 package service
 
 import (
-	"cocogo/pkg/model"
 	"net/http"
 	"path"
 	"path/filepath"
@@ -61,10 +60,10 @@ func (c *WrapperClient) LoadAuth() error {
 }
 
 func (c *WrapperClient) CheckAuth() error {
-	var user model.User
-	err := c.Get("UserProfileUrl", &user)
-	if err != nil {
-		return err
-	}
+	//var user model.User
+	//err := c.Get("UserProfileUrl", &user)
+	//if err != nil {
+	//	return err
+	//}
 	return nil
 }

pkg/service/init.go

@@ -6,15 +6,15 @@ import (
 	"cocogo/pkg/logger"
 )
 
-var client = WrapperClient{}
+var Client = WrapperClient{}
 
 func init() {
-	err := client.LoadAuth()
+	err := Client.LoadAuth()
 	if err != nil {
 		logger.Error("Load client access key error: %s", err)
 		os.Exit(10)
 	}
-	err = client.CheckAuth()
+	err = Client.CheckAuth()
 	if err != nil {
 		logger.Error("Check client auth error: %s", err)
 		os.Exit(11)

pkg/service/users.go

 package service
 
+import "cocogo/pkg/model"
+
+func CheckAuth(username, password, publicKey, remoteAddr, loginType string) (user model.User, err error) {
+	return user, nil
+}
+
 //
 //func (s *Service) CheckAuth(username, password, publicKey, remoteAddr, loginType string) (model.User, error) {
 //	/*

pkg/sshd/template.go → pkg/sshd/handlers/banner.go

-package sshd
+package handlers
 
-const welcomeTemplate = `
+import (
+	"text/template"
+
+	"github.com/gliderlabs/ssh"
+
+	"cocogo/pkg/logger"
+)
+
+const bannerTemplate = `
 		{{.UserName}}	Welcome to use Jumpserver open source fortress system{{.EndLine}}
 {{.Tab}}1) Enter {{.ColorCode}}ID{{.ColorEnd}} directly login or enter {{.ColorCode}}part IP, Hostname, Comment{{.ColorEnd}} to search login(if unique). {{.EndLine}}
 {{.Tab}}2) Enter {{.ColorCode}}/{{.ColorEnd}} + {{.ColorCode}}IP, Hostname{{.ColorEnd}} or {{.ColorCode}}Comment{{.ColorEnd}} search, such as: /ip. {{.EndLine}}
@@ -12,3 +20,30 @@ const welcomeTemplate = `
 {{.Tab}}8) Enter {{.ColorCode}}r{{.ColorEnd}} to refresh your assets and nodes.{{.EndLine}}
 {{.Tab}}0) Enter {{.ColorCode}}q{{.ColorEnd}} exit.{{.EndLine}}
 `
+
+var displayTemplate = template.Must(template.New("display").Parse(bannerTemplate))
+
+type Banner struct {
+	UserName  string
+	ColorCode string
+	ColorEnd  string
+	Tab       string
+	EndLine   string
+}
+
+func (h *Banner) display(sess ssh.Session) {
+	e := displayTemplate.Execute(sess, h)
+	if e != nil {
+		logger.Warn("Display help info failed")
+	}
+}
+
+func NewBanner(userName string) *Banner {
+	return &Banner{
+		UserName:  userName,
+		ColorCode: GreenColorCode,
+		ColorEnd:  ColorEnd,
+		Tab:       Tab,
+		EndLine:   EndLine,
+	}
+}

pkg/sshd/const.go → pkg/sshd/handlers/const.go

-package sshd
+package handlers
 
 const (
 	GreenColorCode = "\033[32m"

pkg/sshd/hostkey.go

+package sshd
+
+import (
+	"golang.org/x/crypto/ssh"
+	"io/ioutil"
+	"os"
+)
+
+type HostKey struct {
+	Value string
+	Path  string
+}
+
+func (hk *HostKey) loadHostKeyFromFile(keyPath string) (signer ssh.Signer, err error) {
+	_, err = os.Stat(conf.HostKeyFile)
+	if err != nil {
+		return
+	}
+	buf, err := ioutil.ReadFile(conf.HostKeyFile)
+	if err != nil {
+		return
+	}
+	return hk.loadHostKeyFromString(string(buf))
+}
+
+func (hk *HostKey) loadHostKeyFromString(value string) (signer ssh.Signer, err error) {
+	signer, err = ssh.ParsePrivateKey([]byte(value))
+	return
+}
+
+func (hk *HostKey) Gen() (signer ssh.Signer, err error) {
+	return
+}
+
+func (hk *HostKey) SaveToFile(signer ssh.Signer) (err error) {
+	return
+}
+
+func (hk *HostKey) Load() (signer ssh.Signer, err error) {
+	if hk.Value != "" {
+		signer, err = hk.loadHostKeyFromString(hk.Value)
+		if err == nil {
+			return
+		}
+	}
+	if hk.Path != "" {
+		signer, err = hk.loadHostKeyFromFile(hk.Path)
+		if err == nil {
+			return
+		}
+	}
+	signer, err = hk.Gen()
+	if err != nil {
+		return
+	}
+	err = hk.SaveToFile(signer)
+	if err != nil {
+		return
+	}
+	return
+}

pkg/sshd/server.go

 package sshd
 
 import (
-	"cocogo/pkg/auth"
-	"cocogo/pkg/config"
-	"cocogo/pkg/model"
-	"io"
 	"strconv"
-	"sync"
-	"text/template"
-
-	"golang.org/x/crypto/ssh/terminal"
 
 	"github.com/gliderlabs/ssh"
-	"github.com/sirupsen/logrus"
+
+	"cocogo/pkg/auth"
+	"cocogo/pkg/config"
+	"cocogo/pkg/logger"
+	"cocogo/pkg/sshd/handlers"
 )
 
 var (
-	conf            *config.Config
-	appService      *auth.Service
-	serverSig       ssh.Signer
-	displayTemplate *template.Template
-	log             *logrus.Logger
-
-	Cached *sync.Map
+	conf = config.Conf
 )
 
-func Initial() {
-	displayTemplate = template.Must(template.New("display").Parse(welcomeTemplate))
-	Cached = new(sync.Map)
-	conf = config.GetGlobalConfig()
-	appService = auth.GetGlobalService()
-	serverSig = parsePrivateKey(conf.TermConfig.HostKey)
-
-	log = logrus.New()
-
-	if level, err := logrus.ParseLevel(conf.LogLevel); err != nil {
-		log.SetLevel(logrus.InfoLevel)
-	} else {
-		log.SetLevel(level)
-	}
-
-}
-
 func StartServer() {
-	srv := ssh.Server{
-		Addr:             conf.BindHost + ":" + strconv.Itoa(conf.SSHPort),
-		PasswordHandler:  appService.CheckSSHPassword,
-		PublicKeyHandler: appService.CheckSSHPublicKey,
-		HostSigners:      []ssh.Signer{serverSig},
-		Version:          "coco-v1.4",
-		Handler:          connectHandler,
+	hostKey := HostKey{Value: conf.HostKey, Path: conf.HostKeyFile}
+	signer, err := hostKey.Load()
+	if err != nil {
+		logger.Fatal("Load access key error: %s", err)
 	}
-	log.Fatal(srv.ListenAndServe())
-}
 
-func connectHandler(sess ssh.Session) {
-	_, _, ptyOk := sess.Pty()
-	if ptyOk {
-		user, ok := sess.Context().Value("LoginUser").(model.User)
-		if !ok {
-			log.Info("Get current User failed")
-			return
-		}
-
-		userInteractive := &sshInteractive{
-			sess:      sess,
-			term:      terminal.NewTerminal(sess, "Opt>"),
-			user:      user,
-			assetData: new(sync.Map),
-			helpInfo: HelpInfo{UserName: sess.User(),
-				ColorCode: GreenColorCode,
-				ColorEnd:  ColorEnd,
-				Tab:       Tab,
-				EndLine:   EndLine}}
-
-		log.Info("accept one session")
-		userInteractive.displayHelpInfo()
-		userInteractive.StartDispatch()
-
-	} else {
-		_, err := io.WriteString(sess, "No PTY requested.\n")
-		if err != nil {
-			return
-		}
+	srv := ssh.Server{
+		Addr:                       conf.BindHost + ":" + strconv.Itoa(conf.SSHPort),
+		PasswordHandler:            auth.CheckUserPassword,
+		PublicKeyHandler:           auth.CheckUserPublicKey,
+		KeyboardInteractiveHandler: auth.CheckMFA,
+		HostSigners:                []ssh.Signer{signer},
+		Version:                    "coco-v1.4",
+		Handler:                    handlers.SessionHandler,
 	}
-
+	logger.Fatal(srv.ListenAndServe())
 }

pkg/sshd/sshutil.go

-package sshd
-
-import (
-	gossh "golang.org/x/crypto/ssh"
-)
-
-func parsePrivateKey(privateKey string) gossh.Signer {
-	private, err := gossh.ParsePrivateKey([]byte(privateKey))
-	if err != nil {
-		log.Error("Failed to parse private key: ", err)
-	}
-	return private
-}

pkg/storage/interface.go

 package storage
 
-import "cocogo/pkg/config"
-
 type Storage interface {
 	Upload(gZipFile, target string)
 }
 
 func NewStorageServer() Storage {
-	conf := config.GetGlobalConfig()
-
-	switch conf.TermConfig.RePlayStorage["TYPE"] {
-	case "server":
-		return NewJmsStorage()
-	}
+	//conf := config.GetGlobalConfig()
+	//
+	//switch conf.TermConfig.RePlayStorage["TYPE"] {
+	//case "server":
+	//	return NewJmsStorage()
+	//}
 	return nil
 }

pkg/storage/jms.go

 package storage
 
-import (
-	"cocogo/pkg/auth"
-	"path/filepath"
-	"strings"
-)
+//var client = service.Client
 
 func NewJmsStorage() Storage {
-	appService := auth.GetGlobalService()
-	return &Server{
-		StorageType: "jms",
-		service:     appService,
-	}
+	//appService := auth.GetGlobalService()
+	//return &Server{
+	//	StorageType: "jms",
+	//	service:     appService,
+	//}
+	return &Server{}
 }
 
 type Server struct {
 	StorageType string
-	service     *auth.Service
 }
 
 func (s *Server) Upload(gZipFilePath, target string) {
-	sessionID := strings.Split(filepath.Base(gZipFilePath), ".")[0]
-	_ = s.service.PushSessionReplay(gZipFilePath, sessionID)
+	//sessionID := strings.Split(filepath.Base(gZipFilePath), ".")[0]
+	//_ = client.PushSessionReplay(gZipFilePath, sessionID)
 }