跳板机管理平台
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

318 lines
6.3 KiB

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)
}
}
}