feat: add ssh server

backend/pkg/proto/ssh/handler/interactive.go

@@ -0,0 +1,795 @@
+// Package handler
+/**
+Copyright (c) The Authors.
+* @Author: feng.xiang
+* @Date: 2023/12/13 09:50
+* @Desc:
+*/
+package handler
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/BurntSushi/toml"
+	"github.com/c-bata/go-prompt"
+	"github.com/chzyer/readline"
+	gossh "github.com/gliderlabs/ssh"
+	"github.com/mattn/go-runewidth"
+	"github.com/nicksnyder/go-i18n/v2/i18n"
+	"github.com/olekukonko/tablewriter"
+	"github.com/patrickmn/go-cache"
+	"github.com/spf13/cast"
+	"github.com/veops/go-ansiterm"
+	"go.uber.org/zap"
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/text/language"
+
+	myi18n "github.com/veops/oneterm/pkg/i18n"
+	"github.com/veops/oneterm/pkg/logger"
+	"github.com/veops/oneterm/pkg/proto/ssh/client"
+	"github.com/veops/oneterm/pkg/proto/ssh/config"
+	"github.com/veops/oneterm/pkg/server/model"
+	"github.com/veops/oneterm/pkg/util"
+)
+
+type InteractiveHandler struct {
+	Locker *sync.RWMutex
+
+	Session gossh.Session
+	//Term      *term.Terminal
+	Term      *readline.Instance
+	Prompt    *prompt.Prompt
+	Localizer *i18n.Localizer
+	SshType   int
+	pty       *gossh.Pty
+
+	Sshd     *sshdServer
+	Pty      gossh.Pty
+	Language int
+
+	Assets       []*model.Asset
+	Accounts     map[int]*model.Account
+	Commands     map[int]*model.Command
+	HistoryInput []string
+
+	SshClient        *ssh.Client
+	SshSession       map[string]*client.Connection
+	GatewayCloseChan chan struct{}
+
+	SelectedAsset *model.Asset
+	SessionReq    *model.SshReq
+
+	AccountInfo *model.Account
+	NeedAccount bool
+
+	Parser *Parser
+
+	GatewayListener   net.Listener
+	MessageChan       chan string
+	AccountsForSelect []*model.Account
+	Cache             *cache.Cache
+}
+
+type Parser struct {
+	Input      *ansiterm.ByteStream
+	Output     *ansiterm.ByteStream
+	InputData  []byte
+	OutputData []byte
+	Ps1        string
+	Ps2        string
+}
+
+var (
+	Bundle       = i18n.NewBundle(language.Chinese)
+	TotalSession = map[string]*client.Connection{}
+)
+
+func I18nInit(path string) {
+	Bundle.RegisterUnmarshalFunc("toml", toml.Unmarshal)
+	files, err := util.ListFiles(path)
+	if err != nil {
+		logger.L.Error(err.Error(), zap.String("module", "i18n"))
+	}
+	for _, f := range files {
+		_, err = Bundle.LoadMessageFile(f)
+		if err != nil {
+			logger.L.Warn(err.Error(), zap.String("module", "i18n"))
+		}
+	}
+}
+
+func NewInteractiveHandler(s gossh.Session, ss *sshdServer, pty gossh.Pty) *InteractiveHandler {
+	//t := term.NewTerminal(s, "> ")
+
+	t, err := readline.NewEx(&readline.Config{
+		Stdin:  s,
+		Stdout: s,
+		Prompt: ">",
+	})
+	if err != nil {
+		logger.L.Error(err.Error())
+	}
+
+	ih := &InteractiveHandler{
+		Locker:  new(sync.RWMutex),
+		Term:    t,
+		Session: s,
+		Sshd:    ss,
+
+		SessionReq:  &model.SshReq{},
+		SshSession:  map[string]*client.Connection{},
+		Pty:         pty,
+		MessageChan: make(chan string, 128),
+		Cache:       cache.New(time.Minute, time.Minute*5),
+	}
+	ih.Language = 1
+	ih.Localizer = i18n.NewLocalizer(Bundle)
+	width := 120
+	height := 40
+	if pty.Window.Width != 0 {
+		width = pty.Window.Width
+	}
+	if pty.Window.Height != 0 {
+		height = pty.Window.Height
+	}
+	ih.Parser = &Parser{
+		Input:  NewParser(width, height),
+		Output: NewParser(width, height),
+	}
+
+	return ih
+}
+
+func completer(d prompt.Document) []prompt.Suggest {
+	// 这里可以根据用户的实时输入来动态生成建议
+	suggestions := []prompt.Suggest{
+		{Text: "users", Description: "Store the username"},
+		{Text: "articles", Description: "Store the article text posted by user"},
+	}
+
+	// 只有当用户输入为空，或者以 'u' 开始时，才显示建议
+	if d.TextBeforeCursor() == "" || d.GetWordBeforeCursorWithSpace() == "u" {
+		return prompt.FilterHasPrefix(suggestions, d.GetWordBeforeCursor(), true)
+	}
+
+	// 其他情况不显示任何建议
+	return []prompt.Suggest{}
+}
+
+func NewParser(width, height int) *ansiterm.ByteStream {
+	screen := ansiterm.NewScreen(width, height)
+	stream := ansiterm.InitByteStream(screen, false)
+	stream.Attach(screen)
+	return stream
+}
+
+func (i *InteractiveHandler) WatchWinSize(winChan <-chan gossh.Window) {
+	for {
+		select {
+		case <-i.Session.Context().Done():
+			return
+		case win, ok := <-winChan:
+			if !ok {
+				return
+			}
+			for _, v := range i.SshSession {
+				client.ResizeSshClient(v.Session, win.Height, win.Width)
+				_ = v.Record.Resize(win.Height, win.Width)
+			}
+		}
+	}
+}
+
+func (i *InteractiveHandler) SwitchLanguage(lang string) {
+	languages := []string{"zh", "en"}
+
+	switch len(lang) {
+	case 0:
+		length := len(languages)
+		if length <= 1 {
+			return
+		}
+		if i.Language >= length {
+			i.Language = 1
+		} else {
+			i.Language += 1
+		}
+	default:
+		for index, v := range languages {
+			if v == lang {
+				i.Language = index + 1
+			}
+		}
+	}
+	i.Localizer = i18n.NewLocalizer(Bundle, languages[i.Language-1])
+
+}
+
+func (i *InteractiveHandler) SwitchLang(lang string) {
+	languages := []string{"zh", "en"}
+
+	switch len(lang) {
+	case 0:
+		length := len(languages)
+		if length <= 1 {
+			return
+		}
+		if i.Language >= length {
+			i.Language = 1
+		} else {
+			i.Language += 1
+		}
+	default:
+		for index, v := range languages {
+			if v == lang {
+				i.Language = index + 1
+			}
+		}
+	}
+	i.Localizer = i18n.NewLocalizer(Bundle, languages[i.Language-1])
+	i.PrintMessage(myi18n.MsgSshWelcome, map[string]any{"User": i.Session.User()})
+}
+
+func (i *InteractiveHandler) output(msg string) {
+	_, _ = io.WriteString(i.Session, msg)
+}
+
+func (i *InteractiveHandler) HostInfo(id int) (asset *model.Asset, err error) {
+	if id < 0 {
+		return
+	}
+
+	cookie, ok := i.Session.Context().Value("cookie").(string)
+	if !ok {
+		err = fmt.Errorf("no cookie")
+		return
+	}
+	res, er := i.Sshd.Core.Asset.Lists(cookie, "", id)
+	if er != nil {
+		err = er
+		return
+	}
+	if res.Count != 1 {
+		er = fmt.Errorf("found %d hosts: not unique", res.Count)
+		return
+	}
+	bs, er := json.Marshal(res.List[0])
+	if er != nil {
+		err = er
+		return
+	}
+	err = json.Unmarshal(bs, &asset)
+	if err != nil {
+		return
+	}
+	return asset, nil
+}
+
+func (i *InteractiveHandler) Check(id int, host *model.Asset) (asset *model.Asset, state bool, err error) {
+	assets, er := i.AcquireAssets("", id)
+	if er != nil {
+		err = err
+		return
+	}
+	if len(assets) == 0 {
+		return
+	}
+	asset = assets[0]
+	state = i.Sshd.Core.Asset.HasPermission(asset.AccessAuth)
+	return
+}
+
+func (i *InteractiveHandler) generateSessionRecord(conn *client.Connection, status int) (res *model.Session, err error) {
+	res = &model.Session{
+		SessionType: cast.ToInt(i.Session.Context().Value("sshType")),
+	}
+	if i.SessionReq != nil && i.SessionReq.SessionId != "" {
+		err = util.DecodeStruct(&res, i.SessionReq)
+		if err != nil {
+			return
+		}
+		res.Uid = i.SessionReq.Uid
+	} else {
+		res.ClientIp = i.Session.RemoteAddr().String()
+		res.UserName = i.Session.Context().User()
+		res.AccountInfo = fmt.Sprintf("%s(%s)", i.AccountInfo.Name, i.AccountInfo.Account)
+		res.Protocol = i.SessionReq.Protocol
+
+	}
+	s, er := i.Sshd.Core.Auth.AclInfo(i.Session.Context().Value("cookie").(string))
+	if er != nil {
+		logger.L.Warn(er.Error(), zap.String("session", "add"))
+	} else if s != nil {
+		res.Uid = s.Uid
+		res.UserName = s.UserName
+	}
+	res.Status = status
+	res.AssetInfo = fmt.Sprintf("%s(%s)", i.SelectedAsset.Name, i.SelectedAsset.Ip)
+	res.SessionId = conn.SessionId
+	res.GatewayId = i.SelectedAsset.GatewayId
+	if conn.Gateway != nil {
+		res.GatewayInfo = fmt.Sprintf("%s:%d", conn.Gateway.Host, conn.Gateway.Port)
+	}
+	res.AssetId = i.SelectedAsset.Id
+	res.AccountId = i.AccountInfo.Id
+	if status == model.SESSIONSTATUS_OFFLINE {
+		t := time.Now()
+		res.ClosedAt = &t
+	}
+	return
+}
+func readLine(s gossh.Session) string {
+	buf := make([]byte, 1)
+	var in []byte
+	for {
+		_, _ = s.Read(buf)
+		switch buf[0] {
+		case []byte("\r")[0], []byte("\r\n")[0]:
+			return string(in)
+		default:
+			in = append(in, buf[0])
+		}
+	}
+}
+
+func (i *InteractiveHandler) Schedule(pty *gossh.Pty) {
+	i.pty = pty
+	var err error
+	var line string
+	if st, ok := i.Session.Context().Value("sshType").(int); ok && st == model.SESSIONTYPE_WEB {
+		//line, err = i.Term.ReadLine()
+		line = readLine(i.Session)
+		if err != nil {
+			logger.L.Debug("connection closed", zap.String("msg", err.Error()))
+			return
+		}
+		var r *model.SshReq
+		err = json.Unmarshal([]byte(line), &r)
+		if err != nil {
+			logger.L.Warn(err.Error())
+			return
+		}
+		// "Accept-Language")
+		//i.Localizer = i18n.NewLocalizer(conf.Bundle, lang, accept)
+		i.Session.Context().SetValue("cookie", r.Cookie)
+		i.SessionReq = r
+
+		// monitor
+		{
+			if i.SessionReq.SessionId != "" {
+				switch i.SessionReq.Action {
+				case model.SESSIONACTION_MONITOR:
+					i.wrapJsonResponse(i.SessionReq.SessionId, 0, "success")
+					RegisterMonitorSession(i.SessionReq.SessionId, i.Session)
+					return
+				case model.SESSIONACTION_CLOSE:
+					if v, ok := config.TotalHostSession.Load(i.SessionReq.SessionId); ok {
+						err = v.(*client.Connection).Session.Close()
+						if err != nil {
+							logger.L.Warn(err.Error())
+							i.wrapJsonResponse(i.SessionReq.SessionId, 1, "failed")
+							return
+						}
+						close(v.(*client.Connection).Exit)
+					}
+					i.wrapJsonResponse(i.SessionReq.SessionId, 0, "success")
+					return
+				}
+			}
+		}
+
+		host, ok, err := i.Check(r.AssetId, nil)
+		if err != nil {
+			logger.L.Warn(err.Error())
+			i.wrapJsonResponse("", 1, err.Error())
+			return
+		}
+
+		if !ok {
+			i.wrapJsonResponse("", 1, fmt.Sprintf("invalid status for %v", r.AssetId))
+			return
+		}
+		i.SelectedAsset = host
+
+		commands, er := i.AcquireCommands()
+		if er != nil {
+			return
+		}
+		i.Commands = commands
+		_, err = i.Proxy(host.Ip, r.AccountId)
+		if err != nil {
+			logger.L.Error(err.Error(), zap.String("module", "proxy"))
+		}
+		return
+	} else {
+		if config.SSHConfig.PlainMode {
+			i.SwitchLang("zh")
+			for {
+				//line, err = i.Term.ReadLine()
+				line, err = i.Term.Readline()
+
+				if err != nil {
+					logger.L.Debug("connection closed", zap.String("msg", err.Error()))
+					break
+				}
+				if strings.TrimSpace(line) == "" {
+					continue
+				}
+				if i.HandleInput(strings.TrimSpace(line)) {
+					break
+				}
+			}
+		} else {
+			tm := InitAndRunTerm(i)
+			_, err := tm.Run()
+			if err != nil {
+				logger.L.Error(err.Error(), zap.String("module", "schedule"))
+			}
+		}
+	}
+}
+
+func (i *InteractiveHandler) HandleInput(line string) (exit bool) {
+
+	switch strings.TrimSpace(line) {
+	case "/*":
+		i.SelectedAsset = nil
+		assets, er := i.AcquireAssets("", 0)
+		if er != nil {
+			return
+		}
+		accounts, er := i.AcquireAccounts()
+		if er != nil {
+			return
+		}
+		commands, er := i.AcquireCommands()
+		if er != nil {
+			return
+		}
+		i.Locker.Lock()
+		i.Assets = assets
+		i.Accounts = accounts
+		i.Commands = commands
+		i.Locker.Unlock()
+
+		i.showResult(assets)
+		return
+	case "/?", "/？":
+		i.PrintMessage(myi18n.MsgSshWelcome, map[string]any{"User": i.Session.User()})
+		return
+	case "/s":
+		i.SwitchLang("")
+		return
+	case "/q":
+		i.Session.Close()
+		return
+	default:
+		switch {
+		case line == "exit":
+			logger.L.Info("exit", zap.String("user", i.Session.User()), zap.String("input", line))
+			i.Session.Close()
+			return
+		}
+	}
+	_, er := i.Proxy(line, -1)
+	if er != nil {
+		logger.L.Info(er.Error())
+	}
+	if st, ok := i.Session.Context().Value("sshType").(int); ok && st == model.SESSIONTYPE_WEB {
+		exit = true
+	}
+	return
+}
+
+func (i *InteractiveHandler) AcquireAndStoreAssets(search string, id int) (selectedHosts, likeHosts []*model.Asset, err error) {
+	i.Locker.RLock()
+	count := len(i.Assets)
+	i.Locker.RUnlock()
+	var find = func(assets []*model.Asset) (selectedHosts, likeHosts []*model.Asset) {
+		if search == "" {
+			return
+		}
+		for _, v := range assets {
+			if v.Ip == search || v.Name == search {
+				selectedHosts = append(selectedHosts, v)
+			} else if strings.Contains(v.Ip, search) || strings.Contains(v.Name, search) {
+				likeHosts = append(likeHosts, v)
+			}
+		}
+		return
+	}
+	if count == 0 {
+		res, er := i.AcquireAssets(search, id)
+		if er != nil {
+			err = er
+			return
+		}
+		selectedHosts, likeHosts = find(res)
+		i.Locker.Lock()
+		i.Assets = res
+		i.Locker.Unlock()
+		return
+	} else {
+		i.Locker.Lock()
+		selectedHosts, likeHosts = find(i.Assets)
+		i.Locker.Unlock()
+		return
+	}
+}
+
+func (i *InteractiveHandler) AcquireAssets(search string, id int) (assets []*model.Asset, err error) {
+	if search == "" && id <= 0 {
+		if v, ok := i.Cache.Get("assets"); ok {
+			return v.([]*model.Asset), nil
+		} else {
+			defer func() {
+				if err == nil {
+					i.Cache.Set("assets", assets, 0)
+				}
+			}()
+		}
+	}
+	if totalAssets, ok := i.Cache.Get("assets"); ok {
+		for _, v := range totalAssets.([]*model.Asset) {
+			if id > 0 {
+				if id == v.Id {
+					assets = append(assets, v)
+					return
+				} else {
+					continue
+				}
+			} else {
+				if strings.Contains(v.Name, search) {
+					assets = append(assets, v)
+				}
+			}
+		}
+	} else {
+		cookie, ok := i.Session.Context().Value("cookie").(string)
+		if ok {
+			res, er := i.Sshd.Core.Asset.Lists(cookie, search, id)
+			if er != nil {
+				err = er
+				return
+			}
+			if res != nil {
+				for _, v := range res.List {
+					var v1 model.Asset
+					_ = util.DecodeStruct(&v1, v)
+					bs, _ := json.Marshal(v.(map[string]interface{})["authorization"])
+					er = json.Unmarshal(bs, &v1.Authorization)
+					if er != nil {
+						logger.L.Warn(er.Error())
+					}
+					assets = append(assets, &v1)
+				}
+			}
+		} else {
+			err = fmt.Errorf("no cookies")
+		}
+	}
+	return
+}
+
+func (i *InteractiveHandler) AcquireAccounts() (accounts map[int]*model.Account, err error) {
+	accounts = map[int]*model.Account{}
+	cookie, ok := i.Session.Context().Value("cookie").(string)
+	if ok {
+		res, er := i.Sshd.Core.Auth.Accounts(cookie)
+		if er != nil {
+			err = er
+			return
+		}
+		for _, v := range res {
+			var v1 model.Account
+			_ = util.DecodeStruct(&v1, v)
+			accounts[v1.Id] = &v1
+		}
+	} else {
+		err = fmt.Errorf("no cookies")
+	}
+	return
+}
+
+func (i *InteractiveHandler) AcquireAccountInfo(id int, name string) (res *model.Account, err error) {
+	cookie, ok := i.Session.Context().Value("cookie").(string)
+	if ok {
+		return i.Sshd.Core.Auth.AccountInfo(cookie, id, name)
+	} else {
+		err = fmt.Errorf("no cookies")
+	}
+	return
+}
+
+func (i *InteractiveHandler) AcquireCommands() (commands map[int]*model.Command, err error) {
+	commands = map[int]*model.Command{}
+	cookie, ok := i.Session.Context().Value("cookie").(string)
+	if ok {
+		res, er := i.Sshd.Core.Asset.Commands(cookie)
+		if er != nil {
+			err = er
+			return
+		}
+		for _, v := range res {
+			var v1 model.Command
+			_ = util.DecodeStruct(&v1, v)
+			commands[v1.Id] = &v1
+		}
+	} else {
+		err = fmt.Errorf("no cookies")
+	}
+	return
+}
+
+func (i *InteractiveHandler) AcquireConfig() (config *model.Config, err error) {
+	config = &model.Config{}
+	cookie, ok := i.Session.Context().Value("cookie").(string)
+	if ok {
+		res, er := i.Sshd.Core.Asset.Config(cookie)
+		if er != nil {
+			err = er
+			return
+		}
+		config = res
+	} else {
+		err = fmt.Errorf("no cookies")
+	}
+	return
+}
+
+func (i *InteractiveHandler) showResult(data []*model.Asset) {
+	i.Term.SetPrompt("host> ")
+	var hosts []string
+	for _, d := range data {
+		hosts = append(hosts, d.Name)
+	}
+
+	var templateData = map[string]interface{}{
+		"Count": len(data),
+		"Msg":   "",
+	}
+
+	if data != nil {
+		templateData["Msg"] = i.tableData(hosts)
+	}
+	i.PrintMessage(myi18n.MsgSshShowAssetResults, templateData)
+}
+
+func (i *InteractiveHandler) tableData(data []string) string {
+	chunkData := i.chunkData(data)
+	buf := &bytes.Buffer{}
+	tw := tablewriter.NewWriter(buf)
+	tw.SetAutoWrapText(false)
+	tw.SetColumnSeparator(" ")
+	tw.SetNoWhiteSpace(false)
+	tw.SetBorder(false)
+	tw.SetAlignment(tablewriter.ALIGN_LEFT)
+	tw.AppendBulk(chunkData)
+	tw.Render()
+	return buf.String()
+}
+
+func (i *InteractiveHandler) chunkData(data []string) (res [][]string) {
+	width := 80
+	if i.pty != nil {
+		width = i.pty.Window.Width
+	}
+	n := len(data)
+	chunk := n
+	for ; chunk >= 1; chunk -= 1 {
+		ok := true
+		for i := 0; i < n && ok; i += chunk {
+			w := chunk*3 + 4
+			r := i + chunk
+			if r > n {
+				r = n
+			}
+			for _, s := range data[i:r] {
+				w += runewidth.StringWidth(s)
+			}
+			ok = ok && w <= width
+		}
+		if ok {
+			t := i.getChunk(data, chunk)
+			maxLen := make(map[int]int)
+			for _, c := range t {
+				for i, v := range c {
+					l := runewidth.StringWidth(v)
+					if l > maxLen[i] {
+						maxLen[i] = l
+					}
+				}
+			}
+			for _, row := range t {
+				w := chunk*3 + 4
+				for i := range row {
+					w += maxLen[i]
+				}
+				ok = ok && w <= width
+			}
+
+		}
+		if ok {
+			break
+		}
+	}
+	if chunk < 1 {
+		chunk = 1
+	}
+	res = i.getChunk(data, chunk)
+
+	return
+}
+
+func (i *InteractiveHandler) getChunk(data []string, chunk int) (res [][]string) {
+	n := len(data)
+	for i := 0; i < n; i += chunk {
+		r := i + chunk
+		if r > n {
+			r = n
+		}
+		res = append(res, data[i:r])
+	}
+	return
+}
+
+func (i *InteractiveHandler) wrapJsonResponse(sessionId string, code int, message string) {
+	if st, ok := i.Session.Context().Value("sshType").(int); ok && st != model.SESSIONTYPE_WEB {
+		return
+	}
+	res, er := json.Marshal(model.SshResp{
+		Code:      code,
+		Message:   message,
+		SessionId: sessionId,
+		Uid:       i.SessionReq.Uid,
+		UserName:  i.SessionReq.UserName,
+	})
+
+	if er != nil {
+		logger.L.Error(er.Error())
+	}
+	i.output(string(append(res, []byte("\r")...)))
+}
+
+func (i *InteractiveHandler) NewSession(account *model.Account, gateway *model.Gateway) (conn *client.Connection, err error) {
+	i.Locker.Lock()
+	defer i.Locker.Unlock()
+	if i.SshClient == nil {
+		con, ch, er := client.NewSShClient(strings.ReplaceAll(i.SessionReq.Protocol, "ssh", i.SelectedAsset.Ip), account, gateway)
+		if er != nil {
+			err = er
+			return
+		}
+		i.SshClient = con
+		i.GatewayCloseChan = ch
+	}
+	i.AccountInfo = account
+
+	conn, err = client.NewSShSession(i.SshClient, i.Pty, i.GatewayCloseChan)
+
+	if err != nil {
+		return
+	}
+	conn.AssetId = i.SelectedAsset.Id
+	conn.AccountId = account.Id
+	conn.Gateway = gateway
+	i.SshSession[conn.SessionId] = conn
+
+	return
+}
+
+func (i *InteractiveHandler) UpsertSession(conn *client.Connection, status int) error {
+	resp, err := i.generateSessionRecord(conn, status)
+	if err != nil {
+		return err
+	}
+
+	return i.Sshd.Core.Audit.NewSession(resp)
+}