[update] websocket handler相关

pkg/model/users.go

@@ -28,3 +28,12 @@ type User struct {
 	IsValid  bool   `json:"is_valid"`
 	IsActive bool   `json:"is_active"`
 }
+
+type TokenUser struct {
+	UserId         string `json:"user"`
+	UserName       string `json:"username"`
+	AssetId        string `json:"asset"`
+	Hostname       string `json:"hostname"`
+	SystemUserId   string `json:"system_user"`
+	SystemUserName string `json:"system_user_name"`
+}

pkg/service/assets.go

@@ -93,3 +93,12 @@ func GetAsset(assetID string) (asset model.Asset) {
 	}
 	return
 }
+
+func GetTokenAsset(token string) (tokenUser model.TokenUser) {
+	Url := authClient.ParseUrlQuery(fmt.Sprintf(TokenAsset, token), nil)
+	err := authClient.Get(Url, &tokenUser)
+	if err != nil {
+		logger.Error("Get Token Asset info failed")
+	}
+	return
+}

pkg/service/urls.go

@@ -10,6 +10,7 @@ const (
 	SystemUserCmdFilterRules = "/api/assets/v1/system-user/%s/cmd-filter-rules/"   // 过滤规则url
 	SystemUser               = "/api/assets/v1/system-user/%s"                     //	某个系统用户的信息
 	Asset                    = "/api/assets/v1/assets/%s/"                         // 某一个资产信息
+	TokenAsset               = "/api/users/v1/connection-token/?token=%s"          // Token name
 
 	TerminalRegisterURL  = "/api/terminal/v2/terminal-registrations/" // 注册当前coco
 	TerminalConfigURL    = "/api/terminal/v1/terminal/config/"        // 从jumpserver获取coco的配置

pkg/webssh/client.go

 package webssh
 
 import (
+	"io"
+
 	socketio "github.com/googollee/go-socket.io"
 	"github.com/ibuler/ssh"
+
+	"cocogo/pkg/model"
 )
 
 type Client struct {
-	Uuid  string
-	Cid   string
-	WinCh chan ssh.Window
+	Uuid      string
+	Cid       string
+	user      *model.User
+	addr      string
+	WinChan   chan ssh.Window
+	UserRead  io.Reader
+	UserWrite io.WriteCloser
+	Conn      socketio.Conn
+	Closed    bool
+}
+
+func (c *Client) Protocol() string {
+	return "ws"
+}
+
+func (c *Client) WinCh() <-chan ssh.Window {
+	return c.WinChan
+}
+
+func (c *Client) User() string {
+	return c.user.Username
+}
+func (c *Client) LoginFrom() string {
+	return "WT"
+}
+func (c *Client) RemoteAddr() string {
+	return c.addr
+}
+
+func (c *Client) Read(p []byte) (n int, err error) {
+	return c.UserRead.Read(p)
+}
+
+func (c *Client) Write(p []byte) (n int, err error) {
+	data := DataMsg{Data: string(p), Room: c.Uuid}
+	n = len(p)
+	c.Conn.Emit("data", data)
+	return
+}
 
-	Conn socketio.Conn
+func (c *Client) Close() (err error) {
+	if c.Closed {
+		return
+	}
+	return c.UserWrite.Close()
 }

pkg/webssh/connection.go

@@ -3,9 +3,10 @@ package webssh
 import (
 	"sync"
 
-	"cocogo/pkg/model"
-
 	socketio "github.com/googollee/go-socket.io"
+	"github.com/ibuler/ssh"
+
+	"cocogo/pkg/model"
 )
 
 type connections struct {
@@ -58,3 +59,11 @@ func (w *WebConn) AddClient(clientID string, conn *Client) {
 	defer w.mu.Unlock()
 	w.Clients[clientID] = conn
 }
+
+func (w *WebConn) SetWinSize(winSize ssh.Window) {
+	w.mu.RLock()
+	defer w.mu.RUnlock()
+	for _, client := range w.Clients {
+		client.WinChan <- winSize
+	}
+}

pkg/webssh/data.go

@@ -22,3 +22,20 @@ type DataMsg struct {
 	Data string `json:"data"`
 	Room string `json:"room"`
 }
+
+type EmitRoomMsg struct {
+	Room   string `json:"room"`
+	Secret string `json:"secret"`
+}
+
+type EmitDataMsg struct {
+	Room string `json:"room"`
+	Data string `json:"data"`
+}
+
+type EmitLogoutMsg struct {
+	Room string `json:"room"`
+}
+
+type EmitDisconnectMsg struct {
+}

pkg/webssh/handler.go

+package webssh
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	socketio "github.com/googollee/go-socket.io"
+	"github.com/ibuler/ssh"
+	uuid "github.com/satori/go.uuid"
+
+	"cocogo/pkg/logger"
+	"cocogo/pkg/service"
+)
+
+func AuthDecorator(handler http.HandlerFunc) http.HandlerFunc {
+	return func(responseWriter http.ResponseWriter, request *http.Request) {
+		cookies := strings.Split(request.Header.Get("Cookie"), ";")
+		var csrfToken string
+		var sessionid string
+		for _, line := range cookies {
+			if strings.Contains(line, "csrftoken") {
+				csrfToken = strings.Split(line, "=")[1]
+			}
+			if strings.Contains(line, "sessionid") {
+				sessionid = strings.Split(line, "=")[1]
+			}
+		}
+		user := service.CheckUserCookie(sessionid, csrfToken)
+		if user.Id == "" {
+			// Todo: 构建login的url
+			http.Redirect(responseWriter, request, "", http.StatusFound)
+			return
+		}
+	}
+}
+
+func OnConnectHandler(s socketio.Conn) error {
+	// 首次连接 1.获取当前用户的信息
+	logger.Debug("OnConnectHandler")
+	cookies := strings.Split(s.RemoteHeader().Get("Cookie"), ";")
+	var csrfToken string
+	var sessionid string
+	var remoteIP string
+	for _, line := range cookies {
+		if strings.Contains(line, "csrftoken") {
+			csrfToken = strings.Split(line, "=")[1]
+		}
+		if strings.Contains(line, "sessionid") {
+			sessionid = strings.Split(line, "=")[1]
+		}
+	}
+	user := service.CheckUserCookie(sessionid, csrfToken)
+	logger.Debug(user)
+	remoteAddrs := s.RemoteHeader().Get("X-Forwarded-For")
+	if remoteAddrs == "" {
+		remoteIP = s.RemoteAddr().String()
+	} else {
+		remoteIP = strings.Split(remoteAddrs, ",")[0]
+	}
+	conn := &WebConn{Cid: s.ID(), Sock: s, Addr: remoteIP, User: user}
+	cons.AddWebConn(s.ID(), conn)
+	return nil
+
+}
+
+func OnErrorHandler(e error) {
+	logger.Debug("OnError trigger")
+	logger.Debug(e)
+}
+
+func OnHostHandler(s socketio.Conn, message HostMsg) {
+	// secret 	uuid string
+	logger.Debug("OnHost trigger")
+	winSiz := ssh.Window{Height: 24, Width: 80}
+	assetID := message.Uuid
+	systemUserId := message.UserID
+	secret := message.Secret
+	width, height := message.Size[0], message.Size[1]
+	if width != 0 {
+		winSiz.Width = width
+	}
+	if height != 0 {
+		winSiz.Height = height
+	}
+	clientID := uuid.NewV4().String()
+	emitMs := EmitRoomMsg{clientID, secret}
+	s.Emit("room", emitMs)
+	asset := service.GetAsset(assetID)
+	systemUser := service.GetSystemUser(systemUserId)
+
+	if asset.Id == "" || systemUser.Id == "" {
+		return
+	}
+
+	userR, userW := io.Pipe()
+
+	conn := cons.GetWebConn(s.ID())
+	clientConn := Client{Uuid: clientID, Cid: conn.Cid, user: conn.User,
+		WinChan: make(chan ssh.Window, 100), Conn: s, UserRead: userR, UserWrite: userW}
+	clientConn.WinChan <- winSiz
+	conn.AddClient(clientID, &clientConn)
+
+	// Todo: 构建proxy server 启动goroutine
+
+}
+
+func OnTokenHandler(s socketio.Conn, message TokenMsg) {
+	logger.Debug("OnToken trigger")
+	winSiz := ssh.Window{Height: 24, Width: 80}
+	token := message.Token
+	secret := message.Secret
+	width, height := message.Size[0], message.Size[1]
+	if width != 0 {
+		winSiz.Width = width
+	}
+	if height != 0 {
+		winSiz.Height = height
+	}
+	clientID := uuid.NewV4().String()
+	emitMs := EmitRoomMsg{clientID, secret}
+	s.Emit("room", emitMs)
+
+	// check token
+
+	if token == "" || secret == "" {
+		msg := fmt.Sprintf("Token or secret is None: %s %s", token, secret)
+		dataMsg := EmitDataMsg{Data: msg, Room: clientID}
+		s.Emit("data", dataMsg)
+		s.Emit("disconnect")
+	}
+	tokenUser := service.GetTokenAsset(token)
+	logger.Debug(tokenUser)
+	if tokenUser.UserId == "" {
+		msg := "Token info is none, maybe token expired"
+		dataMsg := EmitDataMsg{Data: msg, Room: clientID}
+		s.Emit("data", dataMsg)
+		s.Emit("disconnect")
+	}
+
+	currentUser := service.GetUserProfile(tokenUser.UserId)
+	con := cons.GetWebConn(s.ID())
+	con.User = currentUser
+
+	asset := service.GetAsset(tokenUser.AssetId)
+	systemUser := service.GetSystemUser(tokenUser.SystemUserId)
+
+	if asset.Id == "" || systemUser.Id == "" {
+		return
+	}
+
+	userR, userW := io.Pipe()
+	conn := cons.GetWebConn(s.ID())
+	clientConn := Client{Uuid: clientID, Cid: conn.Cid, user: conn.User,
+		WinChan: make(chan ssh.Window, 100), Conn: s, UserRead: userR, UserWrite: userW}
+	clientConn.WinChan <- winSiz
+	conn.AddClient(clientID, &clientConn)
+
+	// Todo: 构建proxy server 启动goroutine
+}
+
+func OnDataHandler(s socketio.Conn, message DataMsg) {
+	logger.Debug("OnData trigger")
+	cid := message.Room
+	webconn := cons.GetWebConn(s.ID())
+	client := webconn.GetClient(cid)
+	_, _ = client.UserWrite.Write([]byte(message.Data))
+}
+
+func OnResizeHandler(s socketio.Conn, message ReSizeMsg) {
+	winSize := ssh.Window{Height: message.Height, Width: message.Width}
+	logger.Debugf("On resize event trigger: %s*%s", message.Width, message.Height)
+	con := cons.GetWebConn(s.ID())
+	con.SetWinSize(winSize)
+}
+
+func OnLogoutHandler(s socketio.Conn, message string) {
+	logger.Debug("OnLogout trigger")
+	webConn := cons.GetWebConn(s.ID())
+	client := webConn.GetClient(message)
+	_ = client.Close()
+}

pkg/webssh/server.go

@@ -3,16 +3,12 @@ package webssh
 import (
 	"net/http"
 	"strconv"
+	"sync"
 
 	socketio "github.com/googollee/go-socket.io"
-	uuid "github.com/satori/go.uuid"
 
 	"cocogo/pkg/config"
 	"cocogo/pkg/logger"
-
-	"cocogo/pkg/service"
-	"strings"
-	"sync"
 )
 
 var (
@@ -38,68 +34,8 @@ func StartHTTPServer() {
 	defer server.Close()
 
 	http.Handle("/socket.io/", server)
-	logger.Debug("start HTTP Serving")
+	logger.Debug("start HTTP Serving ", conf.HTTPPort)
 	httpServer = &http.Server{Addr: conf.BindHost + ":" + strconv.Itoa(conf.HTTPPort), Handler: nil}
 	logger.Fatal(httpServer.ListenAndServe())
 
 }
-
-func OnConnectHandler(s socketio.Conn) error {
-	// 首次连接 1.获取当前用户的信息
-	cookies := strings.Split(s.RemoteHeader().Get("Cookie"), ";")
-	var csrfToken string
-	var sessionid string
-	var remoteIP string
-	for _, line := range cookies {
-		if strings.Contains(line, "csrftoken") {
-			csrfToken = strings.Split(line, "=")[1]
-		}
-		if strings.Contains(line, "sessionid") {
-			sessionid = strings.Split(line, "=")[1]
-		}
-	}
-	user := service.CheckUserCookie(sessionid, csrfToken)
-	remoteAddrs := s.RemoteHeader().Get("X-Forwarded-For")
-	if remoteAddrs == "" {
-		remoteIP = s.RemoteAddr().String()
-	} else {
-		remoteIP = strings.Split(remoteAddrs, ",")[0]
-	}
-	conn := &WebConn{Cid: s.ID(), Sock: s, Addr: remoteIP, User: user}
-	cons.AddWebConn(s.ID(), conn)
-	return nil
-
-}
-func OnErrorHandler(e error) {
-
-}
-
-func OnHostHandler(s socketio.Conn, message HostMsg) {
-	// secret 	uuid string
-	//assetID := message.Uuid
-	//systemUserId := message.UserID
-	secret := message.Secret
-	//width, height := message.Size[0], message.Size[1]
-
-	clientID := uuid.NewV4().String()
-	//asset := service.GetAsset(assetID)
-	//systemUser := service.GetSystemUser(systemUserId)
-	s.Emit("room", map[string]string{"room": clientID, "secret": secret})
-}
-
-func OnTokenHandler(s socketio.Conn, message TokenMsg) {
-
-}
-
-func OnDataHandler(s socketio.Conn, message DataMsg) {
-
-}
-
-func OnResizeHandler(s socketio.Conn, message ReSizeMsg) {
-
-}
-
-func OnLogoutHandler(s socketio.Conn, message string) {
-	// message: room
-
-}