package ssh

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"strconv"
	"io/ioutil"
	"time"
)

type SSHConfig struct {
	User string
	Auth []ssh.AuthMethod
	Host string
	Port int
}

type SSHClient struct {
	session *ssh.Session
	conf    *SSHConfig
}

func New(ip, port, user, pw, keyFilePath string) (*SSHClient, error) {
	p, _ := strconv.Atoi(port)

	auth := make([]ssh.AuthMethod, 0)

	if keyFilePath != "" {
		key, err := parsePrivateKey(keyFilePath, pw)
		if err != nil {
			return nil, err
		}
		auth = append(auth, ssh.PublicKeys(key))
	}
	auth = append(auth, ssh.Password(pw))

	conf := &SSHConfig{
		User: user,
		Auth: auth,
		Host: ip,
		Port: p,
	}

	return &SSHClient{conf: conf}, nil
}

func (cli *SSHClient) Session() error {
	sshConfig := &ssh.ClientConfig{
		User:            cli.conf.User,
		Auth:            cli.conf.Auth,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout: time.Second * 10,
	}

	connection, err := ssh.Dial("tcp", fmt.Sprintf("%s:%d", cli.conf.Host, cli.conf.Port), sshConfig)
	if err != nil {
		return err
	}
	session, err := connection.NewSession()
	if err != nil {
		return err
	}

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	if err := session.RequestPty("xterm", 80, 40, modes); err != nil {
		session.Close()
		return err
	}

	cli.session = session

	return nil
}

func (cli *SSHClient) RunCommand(command string) ([]byte, error) {

	if err := cli.Session(); err != nil {
		return nil, err
	}

	var b bytes.Buffer
	cli.session.Stdout = &b

	err := cli.session.Run(command)
	if err != nil {
		return nil, err
	}

	cli.session.Close()

	return b.Bytes(), nil
}


func parsePrivateKey(keyPath, pw string) (ssh.Signer, error) {
	buff, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, err
	}
	if pw == "" {
		return ssh.ParsePrivateKey(buff)
	}
	return ssh.ParsePrivateKeyWithPassphrase(buff, []byte(pw))
}