[Update] 解决MFA的问题

cmd/hello.go

-package main
-
-import (
-	"bytes"
-	"fmt"
-	"golang.org/x/crypto/ssh/terminal"
-	"io"
-	"io/ioutil"
-	"os"
-)
-
-var f, _ = os.Create("/tmp/new.txt")
-
-var buf, _ = ioutil.ReadFile("/tmp/cmd.text")
-
-type CmdRwParser struct {
-	content []byte
-}
-
-func (c *CmdRwParser) Read(b []byte) (int, error) {
-	for i, v := range c.content {
-		b[i] = v
-	}
-	fmt.Printf("Read %s\n", b)
-	return len(c.content), io.EOF
-}
-
-func (c *CmdRwParser) Write(b []byte) (int, error) {
-	fmt.Printf("Write %s\n", b)
-	return len(b), nil
-}
-
-func main() {
-	nb := new(bytes.Buffer)
-	term := terminal.NewTerminal(nb, ">")
-	nb.Write(buf)
-	nb.Write([]byte("\r"))
-	fmt.Printf("Buf: %s\n", buf)
-
-	line, _ := term.ReadLine()
-	f.WriteString(line)
-	fmt.Printf("Line: %s\n", []byte(line))
-	fmt.Println(".......................")
-	fmt.Printf(nb.String())
-	f.Close()
-}

pkg/auth/server.go

 package auth
 
 import (
+	"cocogo/pkg/model"
+	"fmt"
 	"strings"
 
 	"github.com/ibuler/ssh"
@@ -12,36 +14,45 @@ import (
 	"cocogo/pkg/service"
 )
 
-func checkAuth(ctx ssh.Context, password, publicKey string) (ok bool) {
+func checkAuth(ctx ssh.Context, password, publicKey string) (res ssh.AuthResult) {
 	username := ctx.User()
 	remoteAddr := strings.Split(ctx.RemoteAddr().String(), ":")[0]
-	user := service.Authenticate(username, password, publicKey, remoteAddr, "T")
+	user, err := service.Authenticate(username, password, publicKey, remoteAddr, "T")
 	authMethod := "publickey"
 	action := "Accepted"
+	res = ssh.AuthFailed
 	if password != "" {
 		authMethod = "password"
 	}
-	if user.Id == "" {
+	if err != nil {
 		action = "Failed"
 	} else {
 		ctx.SetValue(cctx.ContextKeyUser, user)
-		ok = true
+		res = ssh.AuthPartiallySuccessful
 	}
 	logger.Infof("%s %s for %s from %s", action, authMethod, username, remoteAddr)
-	return ok
+	return res
 }
 
-func CheckUserPassword(ctx ssh.Context, password string) bool {
-	ok := checkAuth(ctx, password, "")
-	return ok
+func CheckUserPassword(ctx ssh.Context, password string) ssh.AuthResult {
+	res := checkAuth(ctx, password, "")
+	return res
 }
 
-func CheckUserPublicKey(ctx ssh.Context, key ssh.PublicKey) bool {
+func CheckUserPublicKey(ctx ssh.Context, key ssh.PublicKey) ssh.AuthResult {
 	b := key.Marshal()
 	publicKey := common.Base64Encode(string(b))
 	return checkAuth(ctx, "", publicKey)
 }
 
-func CheckMFA(ctx ssh.Context, challenger gossh.KeyboardInteractiveChallenge) bool {
-	return false
+func CheckMFA(ctx ssh.Context, challenger gossh.KeyboardInteractiveChallenge) ssh.AuthResult {
+	answers, err := challenger("admin", "> ", []string{"MFA"}, []bool{true})
+	if err != nil {
+		return ssh.AuthFailed
+	}
+	fmt.Println(answers)
+
+	//ok := checkAuth(ctx, "admin", "")
+	ctx.SetValue(cctx.ContextKeyUser, &model.User{Username: "admin", Name: "admin"})
+	return ssh.AuthSuccessful
 }

pkg/cctx/context.go

@@ -2,6 +2,7 @@ package cctx
 
 import (
 	"context"
+
 	"github.com/ibuler/ssh"
 
 	"cocogo/pkg/model"

pkg/common/utils.go

 package common
+
+import "os"
+
+func FileExists(name string) bool {
+	if _, err := os.Stat(name); err != nil {
+		if os.IsNotExist(err) {
+			return false
+		}
+	}
+	return true
+}

pkg/config/config.go

 package config
 
 import (
+	"cocogo/pkg/logger"
 	"encoding/json"
-	"fmt"
 	"io/ioutil"
 	"os"
 	"strings"
@@ -73,7 +73,7 @@ func (c *Config) LoadFromJSON(body []byte) error {
 	defer c.mux.Unlock()
 	err := json.Unmarshal(body, c)
 	if err != nil {
-		fmt.Println("Load yaml err")
+		logger.Error("Config load yaml error")
 		os.Exit(1)
 	}
 	return nil
@@ -113,9 +113,9 @@ var Conf = &Config{
 	SSHTimeout:     60,
 	HTTPPort:       5000,
 	AccessKey:      "",
-	AccessKeyFile:  "access_key",
+	AccessKeyFile:  "data/keys/.access_key",
 	LogLevel:       "DEBUG",
-	HostKeyFile:    "host_key",
+	HostKeyFile:    "data/keys/host_key",
 	HostKey:        "",
 	RootPath:       rootPath,
 	Comment:        "Coco",

pkg/proxy/parser.go

@@ -2,7 +2,6 @@ package proxy
 
 import (
 	"bytes"
-	"cocogo/pkg/recorder"
 	"sync"
 	"time"
 
@@ -46,8 +45,8 @@ type Parser struct {
 	counter         int
 
 	cmdFilterRules  []model.SystemUserFilterRule
-	commandRecorder *recorder.CommandRecorder
-	replayRecorder  *recorder.ReplyRecorder
+	commandRecorder *CommandRecorder
+	replayRecorder  *ReplyRecorder
 }
 
 func (p *Parser) Initial() {
@@ -174,13 +173,13 @@ func (p *Parser) SetCMDFilterRules(rules []model.SystemUserFilterRule) {
 	p.cmdFilterRules = rules
 }
 
-func (p *Parser) SetReplayRecorder(recorder *recorder.ReplyRecorder) {
+func (p *Parser) SetReplayRecorder(recorder *ReplyRecorder) {
 	p.replayRecorder = recorder
 
 }
 
 func (p *Parser) recordCommand() {
-	cmd := &recorder.Command{
+	cmd := &Command{
 		SessionId:  p.session.Id,
 		OrgId:      p.session.Org,
 		Input:      p.command,
@@ -193,7 +192,7 @@ func (p *Parser) recordCommand() {
 	p.commandRecorder.Record(cmd)
 }
 
-func (p *Parser) SetCommandRecorder(recorder *recorder.CommandRecorder) {
+func (p *Parser) SetCommandRecorder(recorder *CommandRecorder) {
 	p.commandRecorder = recorder
 }
 

pkg/proxy/proxy.go

 package proxy
 
 import (
-	"cocogo/pkg/recorder"
 	"fmt"
 	"io"
 	"strings"
@@ -47,10 +46,11 @@ func (p *ProxyServer) validatePermission() bool {
 
 func (p *ProxyServer) getServerConn() (srvConn ServerConnection, err error) {
 	srvConn = &ServerSSHConnection{
-		host:     "192.168.244.145",
+		host:     "192.168.244.185",
 		port:     "22",
-		user:     "root",
+		user:     "web",
 		password: "redhat",
+		timeout:  config.Conf.SSHTimeout,
 	}
 	pty, _, ok := p.Session.Pty()
 	if !ok {
@@ -61,7 +61,7 @@ func (p *ProxyServer) getServerConn() (srvConn ServerConnection, err error) {
 	go p.sendConnectingMsg(done)
 	err = srvConn.Connect(pty.Window.Height, pty.Window.Width, pty.Term)
 	_, _ = io.WriteString(p.Session, "\r\n")
-	done <- struct{}{}
+	close(done)
 	return
 }
 
@@ -100,10 +100,8 @@ func (p *ProxyServer) Proxy() {
 		logger.Error("Get system user filter rule error: ", err)
 	}
 	sw.parser.SetCMDFilterRules(cmdRules)
-	replayRecorder := recorder.NewReplyRecord(sw.Id)
+	replayRecorder := NewReplyRecord(sw.Id)
 	sw.parser.SetReplayRecorder(replayRecorder)
-	cmdR
-	sw.parser.SetCommandRecorder()
 	_ = sw.Bridge()
 	_ = srvConn.Close()
 }

pkg/proxy/srvconn.go

@@ -29,7 +29,7 @@ type ServerSSHConnection struct {
 	password       string
 	privateKey     string
 	privateKeyPath string
-	timeout        time.Duration
+	timeout        int
 	Proxy          *ServerSSHConnection
 
 	client    *gossh.Client
@@ -61,7 +61,7 @@ func (sc *ServerSSHConnection) Port() string {
 }
 
 func (sc *ServerSSHConnection) Timeout() time.Duration {
-	return sc.timeout
+	return time.Duration(sc.timeout) * time.Second
 }
 
 func (sc *ServerSSHConnection) String() string {
@@ -93,7 +93,7 @@ func (sc *ServerSSHConnection) Config() (config *gossh.ClientConfig, err error)
 		User:            sc.user,
 		Auth:            authMethods,
 		HostKeyCallback: gossh.InsecureIgnoreHostKey(),
-		Timeout:         sc.timeout,
+		Timeout:         sc.Timeout(),
 	}
 	return config, nil
 }
@@ -158,6 +158,7 @@ func (sc *ServerSSHConnection) invokeShell(h, w int, term string) (err error) {
 
 func (sc *ServerSSHConnection) Connect(h, w int, term string) (err error) {
 	_, err = sc.connect()
+	fmt.Println("error")
 	if err != nil {
 		return
 	}

pkg/proxy/switch.go

 package proxy
 
 import (
-	"cocogo/pkg/logger"
 	"context"
+	"time"
+
 	"github.com/ibuler/ssh"
 	"github.com/satori/go.uuid"
-	"time"
+
+	"cocogo/pkg/logger"
 )
 
 func NewSwitch(userConn UserConnection, serverConn ServerConnection) (sw *Session) {

pkg/service/accesskey.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"path"
 	"strings"
 
 	"cocogo/pkg/common"
@@ -61,6 +62,13 @@ func (ak *AccessKey) LoadAccessKeyFromFile(keyPath string) error {
 }
 
 func (ak *AccessKey) SaveToFile() error {
+	keyDir := path.Dir(ak.Path)
+	if !common.FileExists(keyDir) {
+		err := os.MkdirAll(keyDir, os.ModePerm)
+		if err != nil {
+			return err
+		}
+	}
 	f, err := os.Create(ak.Path)
 	defer f.Close()
 	if err != nil {
@@ -74,7 +82,6 @@ func (ak *AccessKey) SaveToFile() error {
 }
 
 func (ak *AccessKey) Register(times int) error {
-	fmt.Println(config.Conf.Name)
 	name := config.Conf.Name
 	token := config.Conf.BootstrapToken
 	comment := "Coco"

pkg/service/users.go

@@ -7,7 +7,7 @@ import (
 	"cocogo/pkg/model"
 )
 
-func Authenticate(username, password, publicKey, remoteAddr, loginType string) (user *model.User) {
+func Authenticate(username, password, publicKey, remoteAddr, loginType string) (user *model.User, err error) {
 	data := map[string]string{
 		"username":    username,
 		"password":    password,
@@ -20,11 +20,13 @@ func Authenticate(username, password, publicKey, remoteAddr, loginType string) (
 		User  *model.User `json:"user"`
 	}
 	Url := client.ParseUrlQuery(UserAuthURL, nil)
-	err := client.Post(Url, data, &resp)
+	err = client.Post(Url, data, &resp)
 	if err != nil {
 		logger.Error(err)
+		return
 	}
-	return resp.User
+	user = resp.User
+	return
 }
 
 func GetUserProfile(userId string) (user *model.User) {

pkg/sshd/hostkey.go

@@ -3,6 +3,7 @@ package sshd
 import (
 	"io/ioutil"
 	"os"
+	"path"
 
 	"golang.org/x/crypto/ssh"
 
@@ -37,6 +38,13 @@ func (hk *HostKey) Gen() (signer ssh.Signer, err error) {
 		return
 	}
 	keyBytes := common.EncodePrivateKeyToPEM(key)
+	keyDir := path.Dir(hk.Path)
+	if !common.FileExists(keyDir) {
+		err := os.MkdirAll(keyDir, os.ModePerm)
+		if err != nil {
+			return signer, err
+		}
+	}
 	err = common.WriteKeyToFile(keyBytes, hk.Path)
 	if err != nil {
 		return

pkg/sshd/server.go

 package sshd
 
 import (
+	"fmt"
 	"strconv"
+	"time"
 
 	"github.com/ibuler/ssh"
+	gossh "golang.org/x/crypto/ssh"
 
 	"cocogo/pkg/auth"
 	"cocogo/pkg/config"
@@ -11,27 +14,43 @@ import (
 	"cocogo/pkg/logger"
 )
 
-const version = "coco-v1.4"
+const version = "v1.4.0"
+
+func defaultConfig(ctx ssh.Context) (conf *gossh.ServerConfig) {
+	conf = new(gossh.ServerConfig)
+	conf.NextAuthMethodsCallback = func(conn gossh.ConnMetadata) (methods []string) {
+		fmt.Println("Username: ", conn.User())
+		if conn.User() == "ibuler" {
+			return []string{"keyboard-interactive"}
+		}
+		return
+	}
+	return conf
+}
 
 var (
 	conf = config.Conf
 )
 
 func StartServer() {
-	logger.Debug("Load host access key")
+	logger.Debug("Load host key")
 	hostKey := HostKey{Value: conf.HostKey, Path: conf.HostKeyFile}
 	signer, err := hostKey.Load()
 	if err != nil {
-		logger.Fatal("Load access key error: %s", err)
+		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.")
 
 	srv := ssh.Server{
 		Addr:                        conf.BindHost + ":" + strconv.Itoa(conf.SSHPort),
 		PasswordHandler:             auth.CheckUserPassword,
 		PublicKeyHandler:            auth.CheckUserPublicKey,
 		KeyboardInteractiveHandler:  auth.CheckMFA,
+		DefaultServerConfigCallback: defaultConfig,
 		HostSigners:                 []ssh.Signer{signer},
-		Version:                    version,
 		Handler:                     handler.SessionHandler,
 		SubsystemHandlers:           map[string]ssh.SubsystemHandler{},
 	}