[Update] 修改结构

pkg/cctx/context.go

@@ -12,7 +12,7 @@ type contextKey struct {
 }
 
 var (
-	ContextKeyUser       = &contextKey{"User"}
+	ContextKeyUser       = &contextKey{"user"}
 	ContextKeyAsset      = &contextKey{"asset"}
 	ContextKeySystemUser = &contextKey{"systemUser"}
 	ContextKeySSHSession = &contextKey{"sshSession"}
@@ -35,7 +35,7 @@ type CocoContext struct {
 	context.Context
 }
 
-// User 返回当前连接的用户model
+// user 返回当前连接的用户model
 func (ctx *CocoContext) User() *model.User {
 	return ctx.Value(ContextKeyUser).(*model.User)
 }

pkg/common/client.go

@@ -117,7 +117,7 @@ func (c *Client) SetReqHeaders(req *http.Request, params ...map[string]string) {
 	if req.Header.Get("Content-Type") == "" {
 		req.Header.Set("Content-Type", "application/json")
 	}
-	req.Header.Set("User-Agent", "coco-client")
+	req.Header.Set("user-Agent", "coco-client")
 	c.SetAuthHeader(req)
 }
 

pkg/common/client_test.go

@@ -38,7 +38,7 @@ func TestClient_Do(t *testing.T) {
 		t.Errorf("Failed Do() error: %s", err.Error())
 	}
 	if len(res) != 2 {
-		t.Errorf("User not equal 2: %d", len(res))
+		t.Errorf("user not equal 2: %d", len(res))
 	}
 }
 

pkg/handler/session.go

@@ -97,7 +97,7 @@ func (i *InteractiveHandler) Dispatch(ctx cctx.Context) {
 
 		if err != nil {
 			if err != io.EOF {
-				logger.Debug("User disconnected")
+				logger.Debug("user disconnected")
 			} else {
 				logger.Error("Read from user err: ", err)
 			}
@@ -342,10 +342,10 @@ func (i *InteractiveHandler) Proxy(ctx context.Context) {
 //	serverAuth := transport.ServerAuth{
 //		SessionID: uuid.NewV4().String(),
 //		IP:        asset.Ip,
-//		Port:      asset.Port,
+//		port:      asset.port,
 //		UserName:  systemUser.UserName,
-//		Password:  systemUser.Password,
-//		PublicKey: parsePrivateKey(systemUser.PrivateKey)}
+//		password:  systemUser.password,
+//		PublicKey: parsePrivateKey(systemUser.privateKey)}
 //
 //	nodeConn, err := transport.NewNodeConn(i.sess.Context(), serverAuth, ptyReq, winChan)
 //	if err != nil {

pkg/proxy/proxy.go

@@ -49,11 +49,11 @@ func (p *ProxyServer) Proxy() {
 	if !p.checkProtocol() {
 		return
 	}
-	conn := SSHConnection{
-		Host:     "192.168.244.185",
-		Port:     "22",
-		User:     "root",
-		Password: "redhat",
+	conn := ServerSSHConnection{
+		host:     "192.168.244.185",
+		port:     "22",
+		user:     "root",
+		password: "redhat",
 	}
 	ptyReq, _, ok := p.Session.Pty()
 	if !ok {
@@ -66,9 +66,9 @@ func (p *ProxyServer) Proxy() {
 	}
 
 	sw := Switch{
-		userSession: p.Session,
-		serverConn:  &conn,
-		parser:      parser,
+		userConn:   p.Session,
+		serverConn: &conn,
+		parser:     parser,
 	}
 	_ = sw.Bridge()
 	_ = conn.Close()

pkg/proxy/connection.go → pkg/proxy/srvconn.go

@@ -11,20 +11,26 @@ import (
 
 type ServerConnection interface {
 	io.ReadWriteCloser
+	Name() string
+	Host() string
+	Port() string
+	User() string
+	Timeout() time.Duration
 	Protocol() string
 	Connect(h, w int, term string) error
 	SetWinSize(w, h int) error
 }
 
-type SSHConnection struct {
-	Host           string
-	Port           string
-	User           string
-	Password       string
-	PrivateKey     string
-	PrivateKeyPath string
-	Timeout        time.Duration
-	Proxy          *SSHConnection
+type ServerSSHConnection struct {
+	name           string
+	host           string
+	port           string
+	user           string
+	password       string
+	privateKey     string
+	privateKeyPath string
+	timeout        time.Duration
+	Proxy          *ServerSSHConnection
 
 	client    *gossh.Client
 	Session   *gossh.Session
@@ -34,25 +40,49 @@ type SSHConnection struct {
 	closed    bool
 }
 
-func (sc *SSHConnection) Protocol() string {
+func (sc *ServerSSHConnection) Protocol() string {
 	return "ssh"
 }
 
-func (sc *SSHConnection) Config() (config *gossh.ClientConfig, err error) {
+func (sc *ServerSSHConnection) User() string {
+	return sc.user
+}
+
+func (sc *ServerSSHConnection) Host() string {
+	return sc.host
+}
+
+func (sc *ServerSSHConnection) Name() string {
+	return sc.name
+}
+
+func (sc *ServerSSHConnection) Port() string {
+	return sc.port
+}
+
+func (sc *ServerSSHConnection) Timeout() time.Duration {
+	return sc.timeout
+}
+
+func (sc *ServerSSHConnection) String() string {
+	return fmt.Sprintf("%s@%s:%s", sc.user, sc.host, sc.port)
+}
+
+func (sc *ServerSSHConnection) Config() (config *gossh.ClientConfig, err error) {
 	authMethods := make([]gossh.AuthMethod, 0)
-	if sc.Password != "" {
-		authMethods = append(authMethods, gossh.Password(sc.Password))
+	if sc.password != "" {
+		authMethods = append(authMethods, gossh.Password(sc.password))
 	}
-	if sc.PrivateKeyPath != "" {
-		if pubkey, err := GetPubKeyFromFile(sc.PrivateKeyPath); err != nil {
+	if sc.privateKeyPath != "" {
+		if pubkey, err := GetPubKeyFromFile(sc.privateKeyPath); err != nil {
 			err = fmt.Errorf("parse private key from file error: %sc", err)
 			return config, err
 		} else {
 			authMethods = append(authMethods, gossh.PublicKeys(pubkey))
 		}
 	}
-	if sc.PrivateKey != "" {
-		if signer, err := gossh.ParsePrivateKey([]byte(sc.PrivateKey)); err != nil {
+	if sc.privateKey != "" {
+		if signer, err := gossh.ParsePrivateKey([]byte(sc.privateKey)); err != nil {
 			err = fmt.Errorf("parse private key error: %sc", err)
 			return config, err
 		} else {
@@ -60,15 +90,15 @@ func (sc *SSHConnection) Config() (config *gossh.ClientConfig, err error) {
 		}
 	}
 	config = &gossh.ClientConfig{
-		User:            sc.User,
+		User:            sc.user,
 		Auth:            authMethods,
 		HostKeyCallback: gossh.InsecureIgnoreHostKey(),
-		Timeout:         sc.Timeout,
+		Timeout:         sc.timeout,
 	}
 	return config, nil
 }
 
-func (sc *SSHConnection) connect() (client *gossh.Client, err error) {
+func (sc *ServerSSHConnection) connect() (client *gossh.Client, err error) {
 	config, err := sc.Config()
 	if err != nil {
 		return
@@ -78,20 +108,20 @@ func (sc *SSHConnection) connect() (client *gossh.Client, err error) {
 		if err != nil {
 			return client, err
 		}
-		proxySock, err := proxyClient.Dial("tcp", net.JoinHostPort(sc.Host, sc.Port))
+		proxySock, err := proxyClient.Dial("tcp", net.JoinHostPort(sc.host, sc.port))
 		if err != nil {
 			return client, err
 		}
-		proxyConn, chans, reqs, err := gossh.NewClientConn(proxySock, net.JoinHostPort(sc.Host, sc.Port), config)
+		proxyConn, chans, reqs, err := gossh.NewClientConn(proxySock, net.JoinHostPort(sc.host, sc.port), config)
 		if err != nil {
 			return client, err
 		}
 		sc.proxyConn = proxyConn
 		client = gossh.NewClient(proxyConn, chans, reqs)
 	} else {
-		client, err = gossh.Dial("tcp", net.JoinHostPort(sc.Host, sc.Port), config)
+		client, err = gossh.Dial("tcp", net.JoinHostPort(sc.host, sc.port), config)
 		if err != nil {
-			err = fmt.Errorf("connect host %sc error: %sc", sc.Host, err)
+			err = fmt.Errorf("connect host %sc error: %sc", sc.host, err)
 			return
 		}
 	}
@@ -99,7 +129,7 @@ func (sc *SSHConnection) connect() (client *gossh.Client, err error) {
 	return client, nil
 }
 
-func (sc *SSHConnection) invokeShell(h, w int, term string) (err error) {
+func (sc *ServerSSHConnection) invokeShell(h, w int, term string) (err error) {
 	sess, err := sc.client.NewSession()
 	if err != nil {
 		return
@@ -126,7 +156,7 @@ func (sc *SSHConnection) invokeShell(h, w int, term string) (err error) {
 	return err
 }
 
-func (sc *SSHConnection) Connect(h, w int, term string) (err error) {
+func (sc *ServerSSHConnection) Connect(h, w int, term string) (err error) {
 	_, err = sc.connect()
 	if err != nil {
 		return
@@ -138,19 +168,19 @@ func (sc *SSHConnection) Connect(h, w int, term string) (err error) {
 	return nil
 }
 
-func (sc *SSHConnection) SetWinSize(h, w int) error {
+func (sc *ServerSSHConnection) SetWinSize(h, w int) error {
 	return sc.Session.WindowChange(h, w)
 }
 
-func (sc *SSHConnection) Read(p []byte) (n int, err error) {
+func (sc *ServerSSHConnection) Read(p []byte) (n int, err error) {
 	return sc.stdout.Read(p)
 }
 
-func (sc *SSHConnection) Write(p []byte) (n int, err error) {
+func (sc *ServerSSHConnection) Write(p []byte) (n int, err error) {
 	return sc.stdin.Write(p)
 }
 
-func (sc *SSHConnection) Close() (err error) {
+func (sc *ServerSSHConnection) Close() (err error) {
 	if sc.closed {
 		return
 	}

pkg/proxy/connection_test.go → pkg/proxy/srvconn_test.go

@@ -5,12 +5,12 @@ import (
 	"testing"
 )
 
-var testConnection = SSHConnection{
-	Host:     "127.0.0.1",
-	Port:     "22",
-	User:     "root",
-	Password: "redhat",
-	Proxy:    &SSHConnection{Host: "192.168.244.185", Port: "22", User: "root", Password: "redhat"},
+var testConnection = ServerSSHConnection{
+	host:     "127.0.0.1",
+	port:     "22",
+	user:     "root",
+	password: "redhat",
+	Proxy:    &ServerSSHConnection{host: "192.168.244.185", port: "22", user: "root", password: "redhat"},
 }
 
 func TestSSHConnection_Config(t *testing.T) {

pkg/proxy/switch.go

 package proxy
 
 import (
+	"cocogo/pkg/logger"
 	"cocogo/pkg/service"
 	"context"
 	"github.com/ibuler/ssh"
 	"github.com/satori/go.uuid"
 	"time"
-
-	"cocogo/pkg/logger"
 )
 
-func NewSwitch(userSess ssh.Session, serverConn ServerConnection) (sw *Switch) {
+func NewSwitch(userConn UserConnection, serverConn ServerConnection) (sw *Switch) {
 	rules, err := service.GetSystemUserFilterRules("")
 	if err != nil {
 		logger.Error("Get system user filter rule error: ", err)
@@ -19,14 +18,14 @@ func NewSwitch(userSess ssh.Session, serverConn ServerConnection) (sw *Switch) {
 		cmdFilterRules: rules,
 	}
 	parser.Initial()
-	sw = &Switch{userSession: userSess, serverConn: serverConn, parser: parser}
+	sw = &Switch{userConn: userConn, serverConn: serverConn, parser: parser}
 	return sw
 }
 
-type SwitchInfo struct {
-	Id         string    `json:"id"`
+type Switch struct {
+	Id         string
 	User       string    `json:"user"`
-	Asset      string    `json:"asset"`
+	Server     string    `json:"asset"`
 	SystemUser string    `json:"system_user"`
 	Org        string    `json:"org_id"`
 	LoginFrom  string    `json:"login_from"`
@@ -36,20 +35,23 @@ type SwitchInfo struct {
 	DateActive time.Time `json:"date_last_active"`
 	Finished   bool      `json:"is_finished"`
 	Closed     bool
-}
 
-type Switch struct {
-	Info        *SwitchInfo
-	parser      *Parser
-	userSession ssh.Session
-	serverConn  ServerConnection
-	userTran    Transport
-	serverTran  Transport
-	cancelFunc  context.CancelFunc
+	parser     *Parser
+	userConn   UserConnection
+	serverConn ServerConnection
+	userTran   Transport
+	serverTran Transport
+	cancelFunc context.CancelFunc
 }
 
 func (s *Switch) Initial() {
 	s.Id = uuid.NewV4().String()
+	s.User = s.userConn.User()
+	s.Server = s.serverConn.Name()
+	s.SystemUser = s.serverConn.User()
+	s.LoginFrom = s.userConn.LoginFrom()
+	s.RemoteAddr = s.userConn.RemoteAddr()
+	s.DateStart = time.Now()
 }
 
 func (s *Switch) preBridge() {
@@ -128,11 +130,11 @@ func (s *Switch) readServerToUser(ctx context.Context) {
 }
 
 func (s *Switch) Bridge() (err error) {
-	_, winCh, _ := s.userSession.Pty()
+	winCh := s.userConn.WinCh()
 	ctx, cancel := context.WithCancel(context.Background())
 	s.cancelFunc = cancel
 
-	s.userTran = NewDirectTransport("", s.userSession)
+	s.userTran = NewDirectTransport("", s.userConn)
 	s.serverTran = NewDirectTransport("", s.serverConn)
 	go s.watchWindowChange(ctx, winCh)
 	go s.readServerToUser(ctx)

pkg/proxy/userconn.go

+package proxy
+
+import (
+	"io"
+	"strings"
+
+	"github.com/ibuler/ssh"
+)
+
+type UserConnection interface {
+	io.ReadWriteCloser
+	Protocol() string
+	WinCh() <-chan ssh.Window
+	User() string
+	Name() string
+	LoginFrom() string
+	RemoteAddr() string
+}
+
+type SSHUserConnection struct {
+	ssh.Session
+	winch <-chan ssh.Window
+}
+
+func (uc *SSHUserConnection) Protocol() string {
+	return "ssh"
+}
+
+func (uc *SSHUserConnection) User() string {
+	return uc.Session.User()
+}
+
+func (uc *SSHUserConnection) WinCh() (winch <-chan ssh.Window) {
+	_, winch, ok := uc.Pty()
+	if ok {
+		return
+	}
+	return nil
+}
+
+func (uc *SSHUserConnection) LoginFrom() string {
+	return "T"
+}
+
+func (uc *SSHUserConnection) RemoteAddr() string {
+	return strings.Split(uc.Session.RemoteAddr().String(), ":")[0]
+}