package ssh_client

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

type SSHClient struct {
	client *ssh.Client
}

func NewSSHClient(host, username, password string) (*SSHClient, error) {
	config := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return nil
		},
	}

	client, err := ssh.Dial("tcp", host, config)
	if err != nil {
		return nil, err
	}

	return &SSHClient{
		client: client,
	}, nil
}

func (c *SSHClient) newSession() (*ssh.Session, error) {
	session, err := c.client.NewSession()
	if err != nil {
		return nil, err
	}

	modes := ssh.TerminalModes{
		ssh.ECHO: 0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

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

	return session, nil
}

// Start a command that will run continuously, calling Close() method of the session to end the command
func (c *SSHClient) Start(cmd string) (*ssh.Session, error) {
	session, err := c.newSession()
	if err != nil {
		return session, err
	}

	go func() {
		_ = session.Start(cmd)
	}()

	return session, err
}

// Start a command and receive all output from the command
func (c *SSHClient) CombinedOutput(cmd string) ([]byte, error) {
	session, err := c.newSession()
	if err != nil {
		return nil, err
	}

	return session.CombinedOutput(cmd)
}

// Start a command and receive standard output from the command
func (c *SSHClient) Output(cmd string) ([]byte, error) {
	session, err := c.newSession()
	if err != nil {
		return nil, err
	}

	return session.Output(cmd)
}

// Start a command that does not require output
func (c *SSHClient) Run(cmd string) error {
	session, err := c.newSession()
	if err != nil {
		return err
	}

	return session.Run(cmd)
}

// Close the ssh client
func (c *SSHClient) Close() error {
	return c.client.Close()
}