Merge branch 'master' of https://github.com/LeeEirc/cocogo

Gopkg.lock

@@ -184,6 +184,22 @@
   revision = "477ce49ddf8f201350f40fdb5ed80a65d155cc33"
   version = "v1.4.0"
 
+[[projects]]
+  digest = "1:7c084e0e780596dd2a7e20d25803909a9a43689c153de953520dfbc0b0e51166"
+  name = "github.com/mattn/go-colorable"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "8029fb3788e5a4a9c00e415f586a6d033f5d38b3"
+  version = "v0.1.2"
+
+[[projects]]
+  digest = "1:9b90c7639a41697f3d4ad12d7d67dfacc9a7a4a6e0bbfae4fc72d0da57c28871"
+  name = "github.com/mattn/go-isatty"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "1311e847b0cb909da63b5fecfb5370aa66236465"
+  version = "v0.0.8"
+
 [[projects]]
   digest = "1:0356f3312c9bd1cbeda81505b7fd437501d8e778ab66998ef69f00d7f9b3a0d7"
   name = "github.com/mattn/go-runewidth"
@@ -192,6 +208,30 @@
   revision = "3ee7d812e62a0804a7d0a324e0249ca2db3476d3"
   version = "v0.0.4"
 
+[[projects]]
+  branch = "master"
+  digest = "1:2b32af4d2a529083275afc192d1067d8126b578c7a9613b26600e4df9c735155"
+  name = "github.com/mgutz/ansi"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "9520e82c474b0a04dd04f8a40959027271bab992"
+
+[[projects]]
+  branch = "master"
+  digest = "1:750abb588d8be68b8ad7e65e350fe2e79a060490ee2d24d437ed1a9f593d0dac"
+  name = "github.com/mlsquires/socketio"
+  packages = [
+    ".",
+    "engineio",
+    "engineio/message",
+    "engineio/parser",
+    "engineio/polling",
+    "engineio/transport",
+    "engineio/websocket",
+  ]
+  pruneopts = "UT"
+  revision = "169a6f09e62468882b9903aad586d2e1632d2f1c"
+
 [[projects]]
   digest = "1:abcdbf03ca6ca13d3697e2186edc1f33863bbdac2b3a44dfa39015e8903f7409"
   name = "github.com/olekukonko/tablewriter"
@@ -216,6 +256,30 @@
   revision = "a713b07e6d90e1831d7fefcb69f1310edb3783ae"
   version = "v1.10.0"
 
+[[projects]]
+  digest = "1:7a40ca559b74a39cb374acf9b2405bfbc1421c32c5d14843d97d3119714196b2"
+  name = "github.com/pschlump/MiscLib"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "a39843c365f40b726e9f30f3d3e8159a361fe1fe"
+  version = "v1.0.0"
+
+[[projects]]
+  digest = "1:74d7514f388723351b912975867519e6b368278d2c6e670aeb78d041b2215e71"
+  name = "github.com/pschlump/godebug"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "8c05f846059acd0b0ac57dc3eaa7db4c3f4dd61e"
+  version = "v0.9.1"
+
+[[projects]]
+  digest = "1:7e9ddd0dfe0afb9919d3470762b7e488a174206648a0c2412dfb1e6d63073392"
+  name = "github.com/pschlump/json"
+  packages = ["."]
+  pruneopts = "UT"
+  revision = "0d2e6a308e08fb871c1bcd9b56fb21a1be305e22"
+  version = "v1.12.0"
+
 [[projects]]
   digest = "1:274f67cb6fed9588ea2521ecdac05a6d62a8c51c074c1fccc6a49a40ba80e925"
   name = "github.com/satori/go.uuid"
@@ -308,6 +372,7 @@
     "github.com/jarcoal/httpmock",
     "github.com/kr/pty",
     "github.com/leonelquinteros/gotext",
+    "github.com/mlsquires/socketio",
     "github.com/olekukonko/tablewriter",
     "github.com/pkg/errors",
     "github.com/pkg/sftp",

Gopkg.toml

@@ -90,3 +90,7 @@
 [[constraint]]
   name = "github.com/elastic/go-elasticsearch"
   version = "0.0.0"
+
+[[constraint]]
+  branch = "master"
+  name = "github.com/mlsquires/socketio"

cmd/ws.go

+package main
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+
+	socketio "github.com/mlsquires/socketio"
+)
+
+func main() {
+	server, err := socketio.NewServer(nil)
+	if err != nil {
+		log.Fatal(err)
+	}
+	server.On("connection", func(s socketio.Socket) error {
+		fmt.Println("connected:")
+		return nil
+	})
+	server.On("/ssh", "host", func(s socketio.Conn, msg interface{}) {
+		fmt.Println("host:")
+	})
+	server.On("/ssh", "data", func(s socketio.Conn, msg interface{}) {
+		fmt.Println("On data")
+	})
+	server.OnEvent("/", "logout", func(s socketio.Conn) {
+		fmt.Println("logout: ")
+		last := s.Context().(string)
+		s.Emit("bye", last)
+	})
+	server.OnError("/ssh", func(e error) {
+		fmt.Println("meet error:", e)
+	})
+	server.OnDisconnect("/ssh", func(s socketio.Conn, msg string) {
+		fmt.Println("closed", msg)
+	})
+	go server.Serve()
+	defer server.Close()
+
+	http.Handle("/socket.io/", server)
+	http.Handle("/", http.FileServer(http.Dir("./asset")))
+	log.Println("Serving at localhost:5000...")
+	log.Fatal(http.ListenAndServe(":5000", nil))
+}

pkg/coco/coco.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"cocogo/pkg/config"
+	"cocogo/pkg/httpd"
 	"cocogo/pkg/logger"
 	"cocogo/pkg/service"
 	"cocogo/pkg/sshd"
@@ -23,6 +24,7 @@ func (c *Coco) Start() {
 	fmt.Printf("Coco version %s, more see https://www.jumpserver.org\n", version)
 	fmt.Println("Quit the server with CONTROL-C.")
 	go sshd.StartServer()
+	go httpd.StartHTTPServer()
 }
 
 func (c *Coco) Stop() {

pkg/webssh/client.go → pkg/httpd/client.go

-package webssh
+package httpd
 
 import (
 	"io"
@@ -19,6 +19,7 @@ type Client struct {
 	UserWrite io.WriteCloser
 	Conn      socketio.Conn
 	Closed    bool
+	pty       ssh.Pty
 }
 
 func (c *Client) Protocol() string {
@@ -50,9 +51,14 @@ func (c *Client) Write(p []byte) (n int, err error) {
 	return
 }
 
+func (c *Client) Pty() ssh.Pty {
+	return c.pty
+}
+
 func (c *Client) Close() (err error) {
-	if c.Closed {
-		return
-	}
+	//if c.Closed {
+	//	return
+	//}
+	c.Closed = true
 	return c.UserWrite.Close()
 }

pkg/webssh/connection.go → pkg/httpd/connection.go

-package webssh
+package httpd
 
 import (
 	"sync"
@@ -33,6 +33,11 @@ func (c *connections) AddWebConn(conID string, conn *WebConn) {
 	c.container[conID] = conn
 }
 
+func newWebConn(id string, sock socketio.Conn, addr string, user *model.User) *WebConn {
+	conn := &WebConn{Cid: id, Sock: sock, Addr: addr, User: user, mu: new(sync.RWMutex), Clients: make(map[string]*Client)}
+	return conn
+}
+
 type WebConn struct {
 	Cid     string
 	Sock    socketio.Conn

pkg/httpd/context.go

+package httpd
+
+import "cocogo/pkg/model"
+
+type WebContext struct {
+	User       *model.User
+	Connection *WebConn
+	Client     *Client
+}

pkg/webssh/data.go → pkg/httpd/data.go

-package webssh
+package httpd
 
 type HostMsg struct {
 	Uuid   string `json:"uuid"`
@@ -7,7 +7,7 @@ type HostMsg struct {
 	Size   []int  `json:"size"`
 }
 
-type ReSizeMsg struct {
+type ResizeMsg struct {
 	Height int `json:"rows"`
 	Width  int `json:"cols"`
 }

pkg/webssh/handler.go → pkg/httpd/handler.go

-package webssh
+package httpd
 
 import (
+	"cocogo/pkg/proxy"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -27,11 +29,9 @@ func AuthDecorator(handler http.HandlerFunc) http.HandlerFunc {
 				sessionid = strings.Split(line, "=")[1]
 			}
 		}
-		user := service.CheckUserCookie(sessionid, csrfToken)
-		if user.ID == "" {
-			// Todo: 构建login的url
+		_, err := service.CheckUserCookie(sessionid, csrfToken)
+		if err != nil {
 			http.Redirect(responseWriter, request, "", http.StatusFound)
-			return
 		}
 	}
 }
@@ -40,29 +40,32 @@ 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
+	var csrfToken, sessionID, 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]
+			sessionID = strings.Split(line, "=")[1]
 		}
 	}
-	user := service.CheckUserCookie(sessionid, csrfToken)
-	logger.Debug(user)
-	remoteAddrs := s.RemoteHeader().Get("X-Forwarded-For")
-	if remoteAddrs == "" {
+	user, err := service.CheckUserCookie(sessionID, csrfToken)
+	if err != nil {
+		return errors.New("user is not authenticated")
+	}
+	remoteAddr := s.RemoteHeader().Get("X-Forwarded-For")
+	if remoteAddr == "" {
 		remoteIP = s.RemoteAddr().String()
 	} else {
-		remoteIP = strings.Split(remoteAddrs, ",")[0]
+		remoteIP = strings.Split(remoteAddr, ",")[0]
 	}
-	conn := &WebConn{Cid: s.ID(), Sock: s, Addr: remoteIP, User: user}
-	cons.AddWebConn(s.ID(), conn)
+	logger.Infof("%s connect websocket from %s\n", user.Username, remoteIP)
+	conn := newWebConn(s.ID(), s, remoteIP, user)
+	ctx := WebContext{User: user, Connection: conn}
+	s.SetContext(ctx)
+	conns.AddWebConn(s.ID(), conn)
+	logger.Info("On Connect handler end")
 	return nil
-
 }
 
 func OnErrorHandler(e error) {
@@ -73,20 +76,20 @@ func OnErrorHandler(e error) {
 func OnHostHandler(s socketio.Conn, message HostMsg) {
 	// secret 	uuid string
 	logger.Debug("OnHost trigger")
-	winSiz := ssh.Window{Height: 24, Width: 80}
+	win := 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
+		win.Width = width
 	}
 	if height != 0 {
-		winSiz.Height = height
+		win.Height = height
 	}
 	clientID := uuid.NewV4().String()
-	emitMs := EmitRoomMsg{clientID, secret}
-	s.Emit("room", emitMs)
+	emitMsg := EmitRoomMsg{clientID, secret}
+	s.Emit("room", emitMsg)
 	asset := service.GetAsset(assetID)
 	systemUser := service.GetSystemUser(systemUserId)
 
@@ -94,16 +97,20 @@ func OnHostHandler(s socketio.Conn, message HostMsg) {
 		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
+	ctx := s.Context().(WebContext)
 
+	userR, userW := io.Pipe()
+	conn := conns.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,
+		pty: ssh.Pty{Term: "xterm", Window: win},
+	}
+	clientConn.WinChan <- win
+	conn.AddClient(clientID, clientConn)
+	proxySrv := proxy.ProxyServer{UserConn: clientConn, User: ctx.User, Asset: &asset, SystemUser: &systemUser}
+	go proxySrv.Proxy()
 }
 
 func OnTokenHandler(s socketio.Conn, message TokenMsg) {
@@ -140,7 +147,7 @@ func OnTokenHandler(s socketio.Conn, message TokenMsg) {
 	}
 
 	currentUser := service.GetUserProfile(tokenUser.UserId)
-	con := cons.GetWebConn(s.ID())
+	con := conns.GetWebConn(s.ID())
 	con.User = currentUser
 
 	asset := service.GetAsset(tokenUser.AssetId)
@@ -151,9 +158,12 @@ func OnTokenHandler(s socketio.Conn, message TokenMsg) {
 	}
 
 	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}
+	conn := conns.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, Closed: false,
+	}
 	clientConn.WinChan <- winSiz
 	conn.AddClient(clientID, &clientConn)
 
@@ -163,21 +173,33 @@ func OnTokenHandler(s socketio.Conn, message TokenMsg) {
 func OnDataHandler(s socketio.Conn, message DataMsg) {
 	logger.Debug("OnData trigger")
 	cid := message.Room
-	webconn := cons.GetWebConn(s.ID())
+	webconn := conns.GetWebConn(s.ID())
 	client := webconn.GetClient(cid)
+	if client == nil {
+		return
+	}
 	_, _ = client.UserWrite.Write([]byte(message.Data))
 }
 
-func OnResizeHandler(s socketio.Conn, message ReSizeMsg) {
+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)
+	logger.Debugf("On resize event trigger: %d*%d", message.Width, message.Height)
+	conn := conns.GetWebConn(s.ID())
+	conn.SetWinSize(winSize)
 }
 
 func OnLogoutHandler(s socketio.Conn, message string) {
 	logger.Debug("OnLogout trigger")
-	webConn := cons.GetWebConn(s.ID())
+	logger.Debugf("Msg: %s\n", message)
+	webConn := conns.GetWebConn(s.ID())
+	if webConn == nil {
+		logger.Error("No conn found")
+		return
+	}
 	client := webConn.GetClient(message)
+	if client == nil {
+		logger.Error("No client found")
+		return
+	}
 	_ = client.Close()
 }

pkg/webssh/server.go → pkg/httpd/server.go

-package webssh
+package httpd
 
 import (
 	"net/http"
@@ -13,7 +13,7 @@ import (
 
 var (
 	httpServer *http.Server
-	cons       = &connections{container: make(map[string]*WebConn), mu: new(sync.RWMutex)}
+	conns      = &connections{container: make(map[string]*WebConn), mu: new(sync.RWMutex)}
 )
 
 func StartHTTPServer() {
@@ -22,13 +22,13 @@ func StartHTTPServer() {
 	if err != nil {
 		logger.Fatal(err)
 	}
-	server.OnConnect("/ssh", OnConnectHandler)
+	server.OnConnect("/ssh", TestOnConnectHandler)
 	server.OnError("/ssh", OnErrorHandler)
-	server.OnEvent("/ssh", "host", OnHostHandler)
+	server.OnEvent("/ssh", "host", TestOnHostHandler)
 	server.OnEvent("/ssh", "token", OnTokenHandler)
-	server.OnEvent("/ssh", "data", OnDataHandler)
-	server.OnEvent("/ssh", "resize", OnResizeHandler)
-	server.OnEvent("/ssh", "logout", OnLogoutHandler)
+	server.OnEvent("/ssh", "data", TestOnDataHandler)
+	server.OnEvent("/ssh", "resize", TestOnResizeHandler)
+	server.OnEvent("/ssh", "logout", TestOnLogoutHandler)
 
 	go server.Serve()
 	defer server.Close()
@@ -37,5 +37,4 @@ func StartHTTPServer() {
 	logger.Debug("start HTTP Serving ", conf.HTTPPort)
 	httpServer = &http.Server{Addr: conf.BindHost + ":" + strconv.Itoa(conf.HTTPPort), Handler: nil}
 	logger.Fatal(httpServer.ListenAndServe())
-
 }

pkg/httpd/test.go

+package httpd
+
+import (
+	"fmt"
+	"github.com/googollee/go-socket.io"
+	"github.com/satori/go.uuid"
+)
+
+func TestOnConnectHandler(s socketio.Conn) error {
+	s.SetContext("")
+	fmt.Println("connected:", s.ID())
+	return nil
+}
+
+func TestOnHostHandler(s socketio.Conn, msg HostMsg) {
+	fmt.Println("On host")
+	secret := msg.Secret
+	clientID := uuid.NewV4().String()
+	emitMsg := EmitRoomMsg{clientID, secret}
+	s.Emit("room", emitMsg)
+	s.Emit("data", DataMsg{Room: clientID, Data: "Hello world"})
+}
+
+func TestOnDataHandler(s socketio.Conn, msg string) {
+	s.Emit("data", msg)
+}
+
+func TestOnResizeHandler(s socketio.Conn, msg ResizeMsg) {
+	fmt.Println("On Resize msg")
+}
+
+func TestOnLogoutHandler(s socketio.Conn, msg string) {
+	fmt.Println("On logout msg")
+}

pkg/httpd/wscontext.go

+package httpd
+
+import (
+	"cocogo/pkg/model"
+	"context"
+	"github.com/gliderlabs/ssh"
+	"github.com/googollee/go-socket.io"
+)
+
+type contextKey struct {
+	name string
+}
+
+var (
+	ContextKeyUser       = &contextKey{"user"}
+	ContextKeyAsset      = &contextKey{"asset"}
+	ContextKeySystemUser = &contextKey{"systemUser"}
+	ContextKeyLocalAddr  = &contextKey{"localAddr"}
+	ContextKeyConnection = &contextKey{"connection"}
+	ContextKeyClient     = &contextKey{"client"}
+)
+
+type WSContext interface {
+	context.Context
+	User() *model.User
+	Asset() *model.Asset
+	SystemUser() *model.SystemUser
+	SSHSession() *ssh.Session
+	SSHCtx() *ssh.Context
+	SetValue(key, value interface{})
+}
+
+type WebSocketContext struct {
+	context.Context
+}
+
+// user 返回当前连接的用户model
+func (ctx *WebSocketContext) User() *model.User {
+	return ctx.Value(ContextKeyUser).(*model.User)
+}
+
+func (ctx *WebSocketContext) Asset() *model.Asset {
+	return ctx.Value(ContextKeyAsset).(*model.Asset)
+}
+
+func (ctx *WebSocketContext) SystemUser() *model.SystemUser {
+	return ctx.Value(ContextKeySystemUser).(*model.SystemUser)
+}
+
+func (ctx *WebSocketContext) Connection() *WebConn {
+	return ctx.Value(ContextKeyConnection).(*WebConn)
+}
+
+func (ctx *WebSocketContext) Client() *Client {
+	return ctx.Value(ContextKeyClient).(*Client)
+}
+
+func (ctx *WebSocketContext) SetValue(key, value interface{}) {
+	ctx.Context = context.WithValue(ctx.Context, key, value)
+}
+
+func applySessionMetadata(ctx *WebSocketContext, sess ssh.Session) {
+	ctx.SetValue(ContextKeyLocalAddr, sess.LocalAddr())
+}
+
+func NewContext(s socketio.Conn) (*WebSocketContext, context.CancelFunc) {
+	parent, cancel := context.WithCancel(context.Background())
+	ctx := &WebSocketContext{parent}
+	return ctx, cancel
+}

pkg/proxy/telnetconn.go

@@ -100,7 +100,6 @@ func (tc *ServerTelnetConnection) optionNegotiate(data []byte) []byte {
 				buf.Write([]byte{WILL, TTYPE})
 			default:
 				buf.Write(bytes.ReplaceAll(item, []byte{DO}, []byte{WONT}))
-
 			}
 		case WILL:
 			switch item[1] {

pkg/service/users.go

@@ -72,12 +72,9 @@ func CheckUserOTP(seed, code string) (resp *AuthResp, err error) {
 	return
 }
 
-func CheckUserCookie(sessionID, csrfToken string) (user *model.User) {
+func CheckUserCookie(sessionID, csrfToken string) (user *model.User, err error) {
 	client.SetCookie("csrftoken", csrfToken)
 	client.SetCookie("sessionid", sessionID)
-	err := client.Get(UserProfileURL, &user)
-	if err != nil {
-		logger.Error(err)
-	}
+	err = client.Get(UserProfileURL, &user)
 	return
 }