/*
@Time : 2019/8/30 11:22
@Author : louis
@File : ssh
@Software: GoLand
*/

package main

import (
	"errors"
	"fmt"
	flag "github.com/spf13/pflag"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
	"io/ioutil"
	"log"
	"net"
	"os"
	"time"
)

type Conn struct {
	user string
	path string
	auth []ssh.AuthMethod
	addr string
}

var (
	client   *ssh.Client
	session  *ssh.Session
	password string
	host     string
	path     string
	port     int
	user     string
	help     bool
)

func usage() {
	_, _ = fmt.Fprintf(os.Stderr, `sshgo version:1.0.0
Usage sshgo [-h host] [-P port] [-u user] [--pkPath path] [-p password]

example: 1) sshgo -u root -P 9527 -h 1.1.1.1 --pkPath /path/to/id_rsa -p 123456
         2) sshgo --user root --port 9527 --host 1.1.1.1 --pkPath /path/to/id_rsa --password 123456 

if use key to login, the password is the key password; if use user & pasword; the password is for user.

Options:
`)
	flag.PrintDefaults()
}

func init() {
	flag.StringVarP(&password, "password", "p", "", "" +
		"if use key to login, the password is the key password; " +
		"if use user & password to login; the password is for user.")
	flag.StringVarP(&path, "pkPath", "", "", "private key path")
	flag.StringVarP(&host, "host", "h", "fenghong.tech", "remote host addr ip")
	flag.IntVarP(&port, "port", "P", 9527, "remote host port")
	flag.StringVarP(&user, "user", "u", "root", "remote host user")
	flag.BoolVarP(&help, "help", "", false, "this help")
	flag.Usage = usage
}

func (c *Conn) SetConf() (err error) {
	c.addr = fmt.Sprintf("%s:%d", host, port)
	c.user = user
	c.path = path
	c.auth = make([]ssh.AuthMethod, 0)
	var method ssh.AuthMethod
	// use privatakey to login,defualt is ""
	if path != "" {
		c.auth = make([]ssh.AuthMethod, 0)
		method, err = PublicFile(path, password)
		if err != nil {
			return err
		}
        // use password & user to login
	} else {
		method = ssh.Password(password)
	}
	c.auth = append(c.auth, method)
	return nil
}

func (c *Conn) SetSession() (session *ssh.Session, err error) {
	client, err = ssh.Dial("tcp", c.addr, &ssh.ClientConfig{
		User: c.user,
		Auth: c.auth,
		//需要验证服务端，不做验证返回nil就可以，点击HostKeyCallback看源码就知道了
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
		Timeout: time.Second * 2,
	})
	if err != nil {
		fmt.Println(err)
		return nil, err
	}

	// create session
	if session, err = client.NewSession(); err != nil {
		return nil, err
	}

	return session, nil
}

func isFile(path string) bool {
	_, err := os.Stat(path)
	return err == nil || os.IsExist(err)
}

//采用公钥验证,这里封装了一下,使用秘钥+密码验证
func PublicFile(privateKeyPath, password string) (method ssh.AuthMethod, err error) {
	if !isFile(privateKeyPath) {
		return nil, errors.New("file not exist")
	}
	bufKey, err := ioutil.ReadFile(privateKeyPath)
	if err != nil {
		return nil, err
	}
	var key ssh.Signer
	if password == "" {
		key, err = ssh.ParsePrivateKey(bufKey)
		if err != nil {
			return nil, err
		}
	} else {
		bufPwd := []byte(password)
		key, err = ssh.ParsePrivateKeyWithPassphrase(bufKey, bufPwd)
		if err != nil {
			return nil, err
		}
	}
	return ssh.PublicKeys(key), nil
}




func main() {
	flag.Parse()
	if help {
		flag.Usage()
		return
	}
	c := Conn{}
	err := c.SetConf()
	if err != nil {
		log.Fatalln(err)
	}
	session, err = c.SetSession()
	if err != nil {
		log.Fatalln(err)
	}
	defer session.Close()
	//当ssh连接建立过后, 我们就可以通过这个连接建立一个回话, 在回话上和远程主机通信。
	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO:1,
		ssh.ECHOCTL:0,
		ssh.TTY_OP_ISPEED:14400,
		ssh.TTY_OP_OSPEED:14400,
	}
	termFd := int(os.Stdin.Fd())
	w,h,_ := terminal.GetSize(termFd)
	termState, _ := terminal.MakeRaw(termFd)
	// TODO 主动从服务器exit,会导致空指针.
	defer terminal.Restore(termFd,termState)

	err = session.RequestPty("xterm-256color",h,w,modes)
	if err != nil {
		log.Fatalln(err)
	}
	err = session.Shell()
	if err != nil {
		log.Fatalln(err)
	}
	err = session.Wait()
	if err != nil {
		log.Fatalln(err)
	}
}

var cmd string 
// init 函数加上命令行参数
flag.StringVarP(&cmd, "cmd", "c", "pwd", "execute cmd in server")

func main() {
	flag.Parse()
	if help {
		flag.Usage()
		return
	}
	c := Conn{}
	err := c.SetConf()
	if err != nil {
		log.Fatalln(err)
	}
	session, err = c.SetSession()
	if err != nil {
		log.Fatalln(err)
	}
	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	_ = session.Run(cmd)
}