package main import ( "database/sql" "fmt" "io" "io/ioutil" "os" "os/signal" "os/user" "runtime" "strconv" "strings" "syscall" "time" "github.com/manifoldco/promptui" _ "github.com/mattn/go-sqlite3" "golang.org/x/crypto/ssh" "golang.org/x/crypto/ssh/terminal" ) var ( DarwinAuthFileFmt = "/Users/%s/.ssh/id_rsa" LinuxAuthFileFmt = "/home/%s/.ssh/id_rsa" AuthFileFmt = LinuxAuthFileFmt DatabaseFile = "/usr/local/jumpserver/jumpserver.db" QuitSelect = "quit" ) type SSHTerminal struct { Session *ssh.Session exitMsg string stdout io.Reader stdin io.Writer stderr io.Reader } // ssh server type HostServer struct { Name string Host string Port int } // SSHUser ssh user type SSHUser struct { Username string Password string IdentityFile string } func init() { if runtime.GOOS == "linux" { AuthFileFmt = LinuxAuthFileFmt } else if runtime.GOOS == "darwin" { AuthFileFmt = DarwinAuthFileFmt } else { fmt.Printf("program not support os(%s)\n", runtime.GOOS) } } func (t *SSHTerminal) updateTerminalSize() { go func() { // 窗口大小改变 sigwinchCh := make(chan os.Signal, 1) signal.Notify(sigwinchCh, syscall.SIGWINCH) fd := int(os.Stdin.Fd()) termWidth, termHeight, err := terminal.GetSize(fd) if err != nil { fmt.Println(err) } for { select { case sigwinch := <-sigwinchCh: if sigwinch == nil { return } currTermWidth, currTermHeight, err := terminal.GetSize(fd) // 窗口没有发生改变 if currTermHeight == termHeight && currTermWidth == termWidth { continue } // 调整窗体大小 t.Session.WindowChange(currTermHeight, currTermWidth) if err != nil { fmt.Printf("Unable to send window-change reqest: %s.", err) continue } termWidth, termHeight = currTermWidth, currTermHeight } } }() } func (t *SSHTerminal) interactiveSession() error { defer func() { if t.exitMsg == "" { fmt.Fprintln(os.Stdout, "the connection was closed on the remote side on ", time.Now().Format(time.RFC822)) } else { fmt.Fprintln(os.Stdout, t.exitMsg) } }() fd := int(os.Stdin.Fd()) // raw模式 cbreak state, err := terminal.MakeRaw(fd) if err != nil { return err } defer terminal.Restore(fd, state) // 获取当前窗体大小 termWidth, termHeight, err := terminal.GetSize(fd) if err != nil { return err } termType := os.Getenv("TERM") if termType == "" { termType = "xterm-256color" } // 设置tty err = t.Session.RequestPty(termType, termHeight, termWidth, ssh.TerminalModes{ ssh.ECHO: 1, ssh.TTY_OP_ISPEED: 14400, ssh.TTY_OP_OSPEED: 14400, }) if err != nil { return err } // 动态调整窗体大小 t.updateTerminalSize() t.stdin, err = t.Session.StdinPipe() if err != nil { return err } t.stdout, err = t.Session.StdoutPipe() if err != nil { return err } t.stderr, err = t.Session.StderrPipe() go io.Copy(os.Stderr, t.stderr) go io.Copy(os.Stdout, t.stdout) go func() { buf := make([]byte, 128) for { n, err := os.Stdin.Read(buf) if err != nil { fmt.Println(err) return } if n > 0 { _, err = t.stdin.Write(buf[:n]) if err != nil { fmt.Println(err) t.exitMsg = err.Error() return } } } }() err = t.Session.Shell() if err != nil { return err } err = t.Session.Wait() if err != nil { return err } return nil } func NewTerminal(server *HostServer, sshUser *SSHUser) (err error) { key, err := ioutil.ReadFile(fmt.Sprintf(AuthFileFmt, sshUser.Username)) if err != nil { return } signer, err := ssh.ParsePrivateKey(key) if err != nil { return } sshConfig := &ssh.ClientConfig{ User: sshUser.Username, Auth: []ssh.AuthMethod{ ssh.Password(sshUser.Password), ssh.PublicKeys(signer), }, HostKeyCallback: ssh.InsecureIgnoreHostKey(), } client, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", server.Host, server.Port), sshConfig) if err != nil { fmt.Println(err) } defer client.Close() session, err := client.NewSession() if err != nil { return err } defer session.Close() s := SSHTerminal{ Session: session, } return s.interactiveSession() } func handlerSignal() { c := make(chan os.Signal, 1) signal.Notify(c, syscall.SIGHUP, syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGINT) for { s := <-c switch s { case syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGINT: return case syscall.SIGHUP: default: return } } } func main() { // 获取当前用户 userInfo, err := user.Current() if err != nil { fmt.Println("user.Current failed!") return } // 连接数据库 db, err := sql.Open("sqlite3", DatabaseFile) if err != nil { fmt.Printf("open database failed! err=%v", err) return } defer db.Close() // 查询当前用户有访问权限的主机 rows, err := db.Query(fmt.Sprintf("SELECT name, ip, port FROM hosts where isdelete=0 and name in(select hostname from hostuser where username='%s');", userInfo.Username)) if err != nil { fmt.Printf("db.Query failed! err=%v", err) return } defer rows.Close() // ui展示列表 menuLabels := make([]string, 0) menuLabels = append(menuLabels, QuitSelect) // 退出 for rows.Next() { var ( name string ip string port int ) err = rows.Scan(&name, &ip, &port) if err != nil { fmt.Printf("rows.Scan failed! err=%v", err) continue } menuLabels = append(menuLabels, fmt.Sprintf("%s:%s:%d", name, ip, port)) } // 选项列表 prompt := promptui.Select{ Label: "Select Host (quit for ctrl+c or select quit)", Items: menuLabels, Size: len(menuLabels), } for { _, selectLabel, err := prompt.Run() if err != nil { fmt.Printf("Prompt failed %v\n", err) return } fmt.Printf("You choose %q\n", selectLabel) if selectLabel == QuitSelect { return } // 获取选中信息 hostItems := strings.Split(selectLabel, ":") if len(hostItems) != 3 { fmt.Printf("invalid ssh host: %s", selectLabel) continue } port, err := strconv.ParseInt(hostItems[2], 10, 32) if err != nil { fmt.Printf("invalid ssh port: %s", selectLabel[2]) continue } server := &HostServer{Name: hostItems[0], Host: hostItems[1], Port: int(port)} user := &SSHUser{Username: userInfo.Username, IdentityFile: fmt.Sprintf(AuthFileFmt, userInfo.Username)} // 新建一个终端 err = NewTerminal(server, user) if err != nil { fmt.Printf("NewTerminal err=%v\n", err) } } }