Go进阶52:开发扩展SSH的使用领域和功能

package main

import (
	"golang.org/x/crypto/ssh"
	"gorm.io/gorm"
	"log"
	"net"
)

var privateBytes = []byte(`
-----BEGIN OPENSSH PRIVATE KEY-----
...
-----END OPENSSH PRIVATE KEY-----
`) // sshd-server 的私钥正式

// startSshSvrListen 启动ssh-server服务 db来管理 用户登录 和 用户登录日志 完成一些聊天的功能
// 这个方法将被 main.go 的 main方法调用,ssh-server的启动入口
func startSshSvrListen(addr string, db *gorm.DB) {
	//初始话 ssh-server 客户端的配置, 用户认证
	config := &ssh.ServerConfig{
		NoClientAuth: false, //如果是true ssh-sever不需要用户认证
		MaxAuthTries: 6,     //用户认证重试次数
		//PasswordCallback:            authUserMfa(db),       //ssh用户名密码认证 可以扩展成  LDAP ... google MFA
		PublicKeyCallback:           authPublicKeysOfGithub(db), //github用户ssh公钥登录 https://github.com/${githubUserName}.keys  https://github.com/mojocn.keys
		KeyboardInteractiveCallback: authKeyboard(db),           //键盘问答输入用户认证 可以扩展成  LDAP ... google MFA
		AuthLogCallback:             nil,                        //记录用户登录认证日志的callback
		ServerVersion:               "",                         //"ssh可以扩展的更多功能的聊天服务", ascii string only  //自定义服务端的版本信息
	}
	// 解析需要给服务端设置ssh 私钥
	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		log.Fatal("无效私钥证书:", err)
	}
	config.AddHostKey(private)

	//监听socket
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal("启动socket失败::", err)
	}
	for {
		// 处理连接
		conn, err := listener.Accept()
		if err != nil {
			// handle error
			log.Println(err)
			continue
		}
		// 用户认证协议SSH User Authentication Protocol
		// 开始工作
		// 开始 handshake 用户登录之前这里用户身份认证, ssh.NewServerConn 会调用上面 PasswordCallback  PublicKeyCallback KeyboardInteractiveCallback ...的callback
		// 用户登录成之后需要向后传递的参数可以 从 sConn.Permissions 中获取
		sConn, chans, reqs, err := ssh.NewServerConn(conn, config)
		if err != nil {
			// handle error
			log.Print(err)
			continue
		}
		//处理 连接协议SSH Connection Protocol
		// 用户handshake 认证成功
		// 强制必须 丢弃服务的request,防止被攻击
		go ssh.DiscardRequests(reqs)
		// 核心/VIP/MVP 处理连接协议SSH Connection Protocol
		go handleChannels(chans, sConn)
	}
}

package main

import (
	"bytes"
	"context"
	"crypto/md5"
	"errors"
	"fmt"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"gorm.io/gorm"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)
//authPublicKeysOfGithub github.com 公钥身份authentication
func authPublicKeysOfGithub(db *gorm.DB) func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
	return func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
		userName := conn.User()
		err := sshPublicKeysAuthByGithub(userName, key)//用过github.com api 获取用户名的公钥 校验
		if err != nil {
			return nil, err
		}
		one := new(User)
		err = db.Where("name = ?", userName).Take(one).Error
		if errors.Is(err, gorm.ErrRecordNotFound) {
			one.Name = userName
			err := db.Save(one).Error
			if err != nil {
				logrus.Error(err)
			}
		}
		return setPermission(one, Fingerprint(key), "github"), nil
	}
}
//authKeyboard 用户普通匿名登录 记录用户信息
func  authKeyboard(db *gorm.DB) func(conn ssh.ConnMetadata, challenge ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
	return func(conn ssh.ConnMetadata, challenge ssh.KeyboardInteractiveChallenge) (*ssh.Permissions, error) {
		userName := conn.User()
		one := new(User)
		err := db.Where("name = ?", userName).Take(one).Error
		if errors.Is(err, gorm.ErrRecordNotFound) {
			one.Name = userName
			err := db.Save(one).Error
			if err != nil {
				logrus.Error(err)
			}
		}
		return setPermission(one, "", "anon"), nil
	}
}


//Fingerprint 计算公钥指纹
func Fingerprint(k ssh.PublicKey) string {
	hash := md5.Sum(k.Marshal())
	r := fmt.Sprintf("% x", hash)
	return strings.Replace(r, " ", ":", -1)
}
//sshPublicKeysAuthByGithub 比较github的公钥
func sshPublicKeysAuthByGithub(user string, key ssh.PublicKey) error {
	publicKeys, err := fetchGithubPublicKeys(user)
	if err != nil {
		return err
	}
	for _, pbk := range publicKeys {
		if bytes.Equal(key.Marshal(), pbk.Marshal()) {
			return nil
		}
	}
	return fmt.Errorf("the key is not match any https://github.com/%s.keys", user)
}
//fetchGithubPublicKeys 获取当用户名的公钥
func fetchGithubPublicKeys(githubUser string) ([]ssh.PublicKey, error) {
	keyURL := fmt.Sprintf("https://github.com/%s.keys", githubUser)
	ctx, cancelFunc := context.WithTimeout(context.Background(), time.Second*15)
	defer cancelFunc()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, keyURL, nil)
	if err != nil {
		return nil, err
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		return nil, errors.New("invalid response from github")
	}
	authorizedKeysBytes, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("reading body:%v", err)
	}
	var keys []ssh.PublicKey
	for len(authorizedKeysBytes) > 0 {
		pubKey, _, _, rest, err := ssh.ParseAuthorizedKey(authorizedKeysBytes)
		if err != nil {
			return nil, fmt.Errorf("parsing key: %v",err)//errors.Wrap(err, "parsing key")
		}
		keys = append(keys, pubKey)
		authorizedKeysBytes = rest
	}
	return keys, nil
}

package main

import (
	"fmt"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
	"log"
)

func  handleChannels(channels <-chan ssh.NewChannel, sshConn *ssh.ServerConn) {
	user, err := getPermissionUser(sshConn)// 获取  SSH User Authentication Protocol 传递过来的认证用户信息
	if err != nil {
		log.Println(err)
		return
	}
	//创建session pty client
	c := NewClient(user, sshConn, postManVar)
	promptString := fmt.Sprintf("[%s] ", user.Name)

	hasShell := false

	for ch := range channels {
		//交互Session
		//一个Session就是一个远程的程序的执行.这个程序或许是shell,应用程序,系统调用或者内建的子系统.它可能没有绑定到虚拟终端上,又或者有或没有涉及到X11转发.同时间,可以有多个Session正在被运行.
		if t := ch.ChannelType(); t != "session" {
			ch.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
			continue
		}

		channel, requests, err := ch.Accept()
		if err != nil {
			log.Printf("Could not accept channel: %v", err)
			continue
		}
		defer channel.Close()

		c.Term = term.NewTerminal(channel, promptString)

		for req := range requests {
			var width, height int
			var ok bool
			switch req.Type {
			case "shell":
				// 开启交互 tty shell
				//一旦一个Session被设置完毕,在远端就会有一个程序被启动.这个程序可以是一个Shell,也可以时一个应用程序或者是一个有着独立域名的子系统.
				if c.Term != nil && !hasShell {
					go c.HandleShell(channel)
					ok = true
					hasShell = true
				}
			case "pty-req"://通过如下消息可以让服务器为Session分配一个虚拟终端
				//当客户端的终端窗口大小被改变时,或许需要发送这个消息给服务器.
				width, height, ok = parsePtyRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			case "window-change":
				//客户terminal size 改变 client.Term (pty) 的size也需要改变,否则console输出会出现排版错误
				width, height, ok = parseWinchRequest(req.Payload)
				if ok {
					err := c.Resize(width, height)
					ok = err == nil
				}
			case "exec":
				// ssh root@mojotv.cn whoami
				//一旦一个Session被设置完毕,在远端就会有一个程序被启动.这个程序可以是一个Shell,也可以时一个应用程序或者是一个有着独立域名的子系统.
				command, err := c.ParseCommandLine(req)// 协议 req.Payload 里面的用户命令输出
				if err != nil {
					logrus.Printf("error parsing ssh execMsg: %s\n", err)
					return
				} else {
					ok = true
				}
				//开始执行从 whoami 远程shell 命令
				// 执行完成 结果直接返回
				go c.HandleExec(command, channel)
			case "env":
				//在shell或command被开始时之后,或许有环境变量需要被传递过去.然而在特权程序里不受控制的设置环境变量是一个很有风险的事情,
				//所以规范推荐实现维护一个允许被设置的环境变量列表或者只有当sshd丢弃权限后设置环境变量.
				//todo set language i18n
				logrus.Info(string(req.Payload))
			case "subsystem":
				//一旦一个Session被设置完毕,在远端就会有一个程序被启动.这个程序可以是一个Shell,也可以时一个应用程序或者是一个有着独立域名的子系统.
				// 实现一下功能可以实现 sftp功能
				//fmt.Fprintf(debugStream, "Subsystem: %s\n", req.Payload[4:])
				if string(req.Payload[4:]) == "sftp" {
					ok = true
					go c.HandleSftp(channel)
				}

			default:
				log.Println(req.Type, string(req.Payload))
			}
			if req.WantReply {
				req.Reply(ok, nil)
			}
		}
	}
}

package main

import (
	"fmt"
	"github.com/fatih/color"
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
	"time"
)

type Client struct {
	DeviceSessionID string //设备uuid
	Conn            *ssh.ServerConn
	postman         *PostMam //处理 用户关系和 消息msg MQ
	User            *User//当前用户
	selectedFriend  *User//当前窗口选择的对话的好友 可以是nil
	selectedGroup   *Group//当前窗口选择的对话的群组 可以是nil
	Term            *term.Terminal
	termWidth       int
	termHeight      int
}

// NewClient constructs a new client
// 1.记录client terminal的状态
// 2.当前用户的状态
// 3.消息发送接收
// 4.好友群组关系管理
// 5.读取客户段terminal的输入
func NewClient(user *User, conn *ssh.ServerConn, pm *PostMam) *Client {

	return &Client{
		DeviceSessionID: string(conn.SessionID()),
		Conn:            conn,
		postman:         pm,
		User:            user,
		selectedFriend:  nil,
		selectedGroup:   nil,
		Term:            nil,//pty
		termWidth:       0,
		termHeight:      0,
	}
}

// TermWrite 写入消息到当强用户ssh 客户端
func (c *Client) writeBack(msg string) {
	c.Term.Write([]byte(msg))
}


func (c *Client) PromptHome() {
	c.Term.SetPrompt(fmt.Sprintf("[%s]", "🌏"))
}
func (c *Client) SetPrompt(s string) {
	c.Term.SetPrompt(fmt.Sprintf("[%s]", s))
}

func (c *Client) Danger(msg string) {
	content := color.RedString("🔴  %s\r\n", msg)
	c.writeBack(content)
}
func (c *Client) Warning(msg string) {
	content := color.YellowString("🟠  %s\r\n", msg)
	c.writeBack(content)
}

func (c *Client) Success(msg string) {
	content := color.GreenString("🟢  %s\n", msg)
	c.writeBack(content)
}

func (c *Client) Primary(msg string) {
	content := color.BlueString("🔵  %s\r\n", msg)
	c.writeBack(content)
}

func (c *Client) MsgPrivate(msg string) {
	content := color.HiCyanString("💬  %s\r\n", msg)
	c.writeBack(content)
}

func (c *Client) MsgGroup(msg string) {
	content := color.HiYellowString("📻 %s\r\n", msg)
	c.writeBack(content)
}
func (c *Client) WritePigeonMsg(msg Msg) {
	if msg.GroupID > 0 {
		c.MsgGroup(msg.Content + "\r\n")
	} else {
		c.MsgPrivate(msg.Content + "\r\n")
		return
	}
}


// Resize resizes the client to the given width and height
func (c *Client) Resize(width, height int) error {
	width = 1000000 //
	err := c.Term.SetSize(width, height)
	if err != nil {
		logrus.Errorf("Resize failed: %dx%d", width, height)
		return err
	}
	c.termWidth, c.termHeight = width, height
	return nil
}

func (c *Client) setSessionPrompt() {
	prompt := fmt.Sprintf("[%s]", c.User.Name)
	if c.selectedFriend != nil {
		prompt = fmt.Sprintf("[%s -> %s]", c.User.Name, c.selectedFriend.Name)
	}
	if c.selectedGroup != nil {
		prompt = fmt.Sprintf("[%s IN %s]", c.User.Name, c.selectedGroup.Name)
	}
	c.Term.SetPrompt(prompt)
}

package main

import (
	"github.com/sirupsen/logrus"
	"golang.org/x/crypto/ssh"
	"strings"
)

//HandleShell 这里将真这的处理用户 ssh shell 输入
func (c *Client) HandleShell(channel ssh.Channel) {
	defer channel.Close()
	exitChan := make(chan bool, 1)
	//c.Server.Add(c)
	// 用户进入聊天界面 开始注册用户在线状态, 聊天消息的队列
	err := c.postman.RegisterClientDevice(c.User, c.DeviceSessionID)
	if err != nil {
		logrus.Println(err)
		return
	}
	go func() {
		// Block until done, then remove.
		c.Conn.Wait()
		c.closed = true
		//c.Server.Remove(c)
		//close(c.Messages)
		//c.postman.UserOffline(c.User)  // 用户离线
	}()

	go func() {
		//todo:: send history msg

		// 接受其他用户发送给你的消息 或 广播消息
		c.postman.ReceiveMsgLoop(c.DeviceSessionID, c, exitChan)
	}()
	new(actionHelp).Exec(c,nil) //输出帮助信息
	for {
		line, err := c.Term.ReadLine()
		if err != nil {
			break
		}
		// 使用 默认的 hook action 来处理 交互shell的键盘输入(聊天 或者 指令)
		var doer ActionDoer = new(ActionDefault)
		// choose action
		isCmd, action, args := parseInputLine(line) // 解析用户输入 return 是否是指令 or 是发送聊天消息
		if isCmd {
			v, ok := ActionMap[action] //开始匹配指令 有点类似与gin中的路由匹配
			if ok {
				doer = v // 匹配指令的 hook
			} else {
				c.Danger("未知动作指令: " + line)
				continue
			}
		}
		//
		if hint := doer.Hint(args); hint != "" { //参数检查
			c.Warning("Invalid command: " + line)
			continue
		}
		err = doer.Exec(c, args) // 执行自定义的action hook
		if err != nil {
			c.Warning(err.Error())
			logrus.Error(err)
			//c.TermWrite(err.Error())
		}
	}

}
//parseInputLine 解析input
func parseInputLine(line string) (isCmd bool, action string, args []string) {
	parts := strings.Split(line, " ")
	if len(parts) > 0 && strings.HasPrefix(parts[0], "/") {
		args = []string{}
		for _, p := range parts {
			if t := strings.TrimSpace(p); t != "" {
				args = append(args, t)
			}
		}
		return true, strings.TrimPrefix(args[0], "/"), args
	}
	return false, "", []string{line}
}

package main

import "log"


var ActionMap = map[string]ActionDoer{}

//ActionDoer 编写插件hook 来扩展更多的功能
type ActionDoer interface {
	Help() (alias, long string) //命令扩展的帮助信息
	Hint(args []string) string //exec 之前的参数检查
	Exec(c *Client, args []string) error //扩展的执行逻辑
}

//registerAction 注册编写的action hook 扩展功能,这个方法建议在 init 函数中调用
func registerAction(name string, doer ActionDoer) {
	_, ok := ActionMap[name]
	if ok {
		log.Fatal("action has already existed: ", name)
	} else {
		ActionMap[name] = doer
	}
}

package main

import (
	"bytes"
	"fmt"
	"github.com/olekukonko/tablewriter" //在terminal中打印出漂亮的表格
	"golang.org/x/text/encoding/simplifiedchinese"
	"golang.org/x/text/transform"
	"io/ioutil"
	"net/http"
	"strings"
)

func init() {
	registerAction("stock", new(ActionStock))
}

type ActionStock struct{}

func (a ActionStock) Help() (alias, log string) {
	return "stock", "输入股票代码查询股票价格 eg: /stock sh688111 sh600036 sz002594"
}
func (a ActionStock) Exec(c *Client, args []string) error {
	var headers []string
	for _, v := range template {
		headers = append(headers, v.Desc)
	}
	table := tablewriter.NewWriter(c.Term)
	//设置表格样式
	table.SetHeader(headers)
	table.SetBorder(true)
	table.SetHeaderColor(
		tablewriter.Colors{tablewriter.FgHiBlueColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgWhiteColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgCyanColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgHiRedColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgMagentaColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgGreenColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgYellowColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgYellowColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgYellowColor, tablewriter.Bold},
		tablewriter.Colors{tablewriter.FgYellowColor, tablewriter.Bold},

	)
	table.SetColumnColor(
		tablewriter.Colors{tablewriter.Bold, tablewriter.FgHiBlueColor},
		tablewriter.Colors{tablewriter.Normal, tablewriter.FgWhiteColor},

		tablewriter.Colors{tablewriter.Normal, tablewriter.FgCyanColor},
		tablewriter.Colors{tablewriter.Bold, tablewriter.FgHiRedColor},
		tablewriter.Colors{tablewriter.Normal, tablewriter.FgMagentaColor},
		tablewriter.Colors{tablewriter.Normal, tablewriter.FgGreenColor},
		tablewriter.Colors{tablewriter.Normal, tablewriter.FgYellowColor},
		tablewriter.Colors{tablewriter.Normal, tablewriter.FgYellowColor},
		tablewriter.Colors{tablewriter.Normal, tablewriter.FgYellowColor},
		tablewriter.Colors{tablewriter.Normal, tablewriter.FgYellowColor},
	)
	var err error
	for _,code :=range args[1:]{
		row, err2 := stockPrice(code)//sina api 获取股票价格
		if err2 != nil {
			err = err2
		}
		table.Append(row)// 插入股票信息到表格
	}
	table.Render()
	table = nil
	return err
}

func (a ActionStock) Hint(args []string) string {
	return ""
}

var template = []StockItem{
	{
		Idx:   0,
		Desc:  "Name",
		Value: "",
	},
	{
		Idx:   1,
		Desc:  "Today_Start_Price",
		Value: "",
	},
	{
		Idx:   2,
		Desc:  "Yesterday_End_Price",
		Value: "",
	},
	{
		Idx:   3,
		Desc:  "Current_Price",
		Value: "",
	},
	{
		Idx:   4,
		Desc:  "Today_Top",
		Value: "",
	},
	{
		Idx:   5,
		Desc:  "Today_Bottom",
		Value: "",
	},
	{
		Idx:   6,
		Desc:  "Buy_One",
		Value: "",
	},
	{
		Idx:   7,
		Desc:  "Sell_One",
		Value: "",
	},
	{
		Idx:   8,
		Desc:  "Deal_Amount",
		Value: "",
	},
	{
		Idx:   9,
		Desc:  "Deal_Money",
		Value: "",
	},
}

type StockItem struct {
	Idx   int
	Desc  string
	Value string
}


func GbkToUtf8(s []byte) ([]byte, error) {
	reader := transform.NewReader(bytes.NewReader(s), simplifiedchinese.GBK.NewDecoder())
	d, e := ioutil.ReadAll(reader)
	if e != nil {
		return nil, e
	}
	return d, nil
}

func Utf8ToGbk(s []byte) ([]byte, error) {
	reader := transform.NewReader(bytes.NewReader(s), simplifiedchinese.GBK.NewEncoder())
	d, e := ioutil.ReadAll(reader)
	if e != nil {
		return nil, e
	}
	return d, nil
}


func stockPrice(stockCode string) (list []string, err error) {
	url := fmt.Sprintf("http://hq.sinajs.cn/list=%s", stockCode)
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	bs, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	utf8, err := GbkToUtf8(bs)
	if err != nil {
		return nil, err
	}
	parts := strings.Split(string(utf8), `="`)
	body := parts[1]

	ps := strings.Split(body, ",")

	for _, v := range template {
		list = append(list, ps[v.Idx])
	}

	return

}