gistfile1.txt

package protocols

import (
	"fmt"
	"io"
	"net"
	"runtime/debug"
	"strings"
	"time"
	// "log"

	"sync"

	"golang.org/x/crypto/ssh"

	"github.com/felixge/tcpkeepalive"
	"github.com/pkg/errors"

	"github.com/DanielRenne/GoCore/core/atomicTypes"
	"github.com/atlonaeng/studio/sessionFunctions"
)

const (
	SSH_BUFF          = 1024
	RECONNECTION_TIME = 3000
)

//OnData callback when data is received by a protocol.
type OnData func(client *TerminalSSH, data []byte)

//OnError callback when an error occurs.
type OnError func(client *TerminalSSH, err error, message string)

//OnConnection callback when a connection is established.
type OnConnection func(client *TerminalSSH)

type OnLog func(message string)

//TerminalSSH provides an object that can exec a new ssh terminal for mac and linux.
type TerminalSSH struct {
	OnConnection           OnConnection
	OnData                 OnData
	OnError                OnError
	OnLog                  OnLog
	processWriter          stdInSync
	processReader          stdOutSync
	processError           stdErrorSync
	endConnection          chan int
	endConnectionWaiting   atomicTypes.AtomicBool
	password               string
	explicitlyDisconnected atomicTypes.AtomicBool
	KillReadErr            atomicTypes.AtomicBool
	KillRead               atomicTypes.AtomicBool
}

type stdInSync struct {
	sync.RWMutex
	writer io.Writer
}

type stdOutSync struct {
	sync.RWMutex
	reader io.Reader
}

type stdErrorSync struct {
	sync.RWMutex
	reader io.Reader
}

//Connect will start the ssh terminal launch and connection.
func (terminal *TerminalSSH) Connect(userName string, password string, address string, port string) {
	go terminal.connectTCPKeepAlive(userName, password, address, port)
}

func (terminal *TerminalSSH) sshInteractive(user, instruction, password string, questions []string, echos []bool) (answers []string, err error) {

	answers = make([]string, len(questions))

	for n, _ := range questions {
		answers[n] = password
	}

	return answers, nil
}

func (terminal *TerminalSSH) connectTCPKeepAlive(userName string, password string, address string, port string) {
	defer func() {
		if r := recover(); r != nil {
			session_functions.Print("\n\nPanic Stack: " + string(debug.Stack()))
			session_functions.VelocityLog("Panic Recovered", fmt.Sprintf("%+v", r))
			return
		}
	}()

	//Default to Non-KeyboardInteractive
	keyboardInteractiveCheck := false

	for {

		if terminal.OnLog != nil {
			terminal.OnLog("Connecting to SSH " + address)
		}

		if terminal.explicitlyDisconnected.Get() {
			return
		}

		terminal.endConnectionWaiting.Set(false)
		terminal.endConnection = make(chan int)

		keyboardInteractive := func(user, instruction string, questions []string, echos []bool) (answers []string, err error) {
			answers, err = terminal.sshInteractive(user, instruction, password, questions, echos)
			return
		}

		var config *ssh.ClientConfig

		if keyboardInteractiveCheck == false {
			config = &ssh.ClientConfig{
				User:            userName,
				HostKeyCallback: ssh.InsecureIgnoreHostKey(),
				Auth: []ssh.AuthMethod{
					ssh.Password(password),
				},
			}
		} else {
			config = &ssh.ClientConfig{
				User:            userName,
				HostKeyCallback: ssh.InsecureIgnoreHostKey(),
				Auth: []ssh.AuthMethod{
					ssh.KeyboardInteractive(keyboardInteractive),
				},
			}
		}

		// modes := ssh.TerminalModes{
		// 	ssh.ECHO: 0, // disable echoing
		// }

		conn, err := net.DialTimeout("tcp", address+":"+port, time.Duration(3000)*time.Millisecond)

		if err != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "TCP Dial Error")
			}
			time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
			continue
		}

		keepAliveTCPConn, _ := conn.(*net.TCPConn)

		kaConn, errKeepAlive := tcpkeepalive.EnableKeepAlive(keepAliveTCPConn)
		if errKeepAlive == nil {
			kaConn.SetKeepAliveIdle(2 * time.Second)
			kaConn.SetKeepAliveCount(2)
			kaConn.SetKeepAliveInterval(2 * time.Second)
		} else {
			session_functions.VelocityLog("Error->protocols->tcp.Connect", "Failed to set Keep Alive on SSH Connection: "+errKeepAlive.Error())
		}

		sshConn, chans, reqs, err := ssh.NewClientConn(conn, address, config)
		if err != nil {
			if terminal.OnError != nil {

				if strings.Contains(err.Error(), "no supported methods remain") {
					if keyboardInteractiveCheck {
						keyboardInteractiveCheck = false
					} else {
						keyboardInteractiveCheck = true
					}
				}

				terminal.OnError(terminal, err, "Failed to Create New Client Conn")
			}
			conn.Close()
			time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
			continue
		}

		if terminal.OnLog != nil {
			terminal.OnLog("Created New Client Conn")
		}

		sshClient := ssh.NewClient(sshConn, chans, reqs)

		session, err := sshClient.NewSession()
		if err != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "Failed to Establish SSH Session")
			}
			conn.Close()
			time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
			continue
		}

		// err = session.RequestPty("xterm", 80, 40, modes)
		// if err != nil {
		// 	if terminal.OnError != nil {
		// 		terminal.OnError(terminal, err, "Failed to Establish Psuedo Terminal")
		// 		conn.Close()
		// 	}
		// 	time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
		// 	continue
		// }

		// if terminal.OnLog != nil {
		// 	terminal.OnLog("Created Psuedo Terminal")
		// }

		writer, err := session.StdinPipe()
		if err != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "Failed to set Stdin Pipe")
			}
			time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
			continue
		}

		terminal.setProcessWriter(writer)

		reader, err := session.StdoutPipe()
		if err != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "Failed to set Stdout Pipe")
			}
			time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
			continue
		}

		terminal.setProcessReader(reader)

		errReader, err := session.StderrPipe()
		if err != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "Failed to set Stderr Pipe")
			}
			time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
			continue
		}

		terminal.setProcessError(errReader)

		go terminal.readStdOut()
		go terminal.readStdErr()

		// go func() {
		err = session.Shell()
		if err != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "Shell Exited")
			}
			conn.Close()
			time.Sleep(time.Millisecond * time.Duration(RECONNECTION_TIME))
			continue
		}

		if terminal.OnConnection != nil {
			go terminal.OnConnection(terminal)
		}

		if terminal.OnLog != nil {
			terminal.OnLog("Started SSH Shell Successfully.")
		}

		terminal.endConnectionWaiting.Set(true)
		<-terminal.endConnection

		session.Close()

		if terminal.OnLog != nil {
			terminal.OnLog("SSH Connection Ended:  IP:  " + address + " Port:  " + port)
		}

		time.Sleep(time.Second * 3)
	}
}

func (terminal *TerminalSSH) Disconnect() {
	if terminal.endConnectionWaiting.Get() {
		terminal.endConnection <- 0
	}
}

func (terminal *TerminalSSH) DisconnectExplicitlyFromTerminal() {
	terminal.explicitlyDisconnected.Set(true)
	if terminal.endConnectionWaiting.Get() {
		terminal.endConnection <- 0
	}
}

func (terminal *TerminalSSH) Write(data []byte) {
	defer func() {
		if r := recover(); r != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, fmt.Errorf("%+v", r), "Write: Line 130 with string ("+string(data)+")")
			}
			return
		}
	}()

	writer := terminal.getProcessWriter()
	_, err := writer.Write(data)

	if err != nil {
		if terminal.endConnectionWaiting.Get() {
			terminal.endConnection <- 0
		}
		if terminal.OnError != nil {
			terminal.OnError(terminal, err, "Write: Line 141")
		}
		return
	}

	return
}

func (terminal *TerminalSSH) readStdOut() {

	defer func() {
		if r := recover(); r != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, fmt.Errorf("%+v", r), "readStdOut: Panic")
			}
			return
		}
	}()

	result := make([]byte, SSH_BUFF)
	for {

		reader := terminal.getProcessReader()
		length, err := reader.Read(result)

		if err != nil {
			if terminal.endConnectionWaiting.Get() {
				terminal.endConnection <- 0
			}
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "readStdOutGeneralError:  "+err.Error())
				return
			}
		}

		if terminal.OnData != nil {
			terminal.OnData(terminal, result[:length])
		}
	}
}

func (terminal *TerminalSSH) readStdErr() {

	defer func() {
		if r := recover(); r != nil {
			if terminal.OnError != nil {
				terminal.OnError(terminal, fmt.Errorf("%+v", r), "readStdErr:  Panic")
			}
			return
		}
	}()

	result := make([]byte, SSH_BUFF)
	for {

		reader := terminal.getProcessError()
		length, err := reader.Read(result)

		if err != nil {
			if terminal.endConnectionWaiting.Get() {
				terminal.endConnection <- 0
			}
			if terminal.OnError != nil {
				terminal.OnError(terminal, err, "readStdErrGeneralError:  "+err.Error())
				return
			}
			return
		}

		if terminal.endConnectionWaiting.Get() {
			terminal.endConnection <- 0
		}

		if terminal.OnError != nil {
			terminal.OnError(terminal, errors.New("StdErr Recieved"), string(result[:length]))
		}
	}
}

func (terminal *TerminalSSH) setProcessWriter(writer io.WriteCloser) {
	terminal.processWriter.Lock()
	terminal.processWriter.writer = writer
	terminal.processWriter.Unlock()
}

func (terminal *TerminalSSH) getProcessWriter() (writer io.Writer) {
	terminal.processWriter.RLock()
	writer = terminal.processWriter.writer
	terminal.processWriter.RUnlock()
	return
}

func (terminal *TerminalSSH) setProcessReader(reader io.Reader) {
	terminal.processReader.Lock()
	terminal.processReader.reader = reader
	terminal.processReader.Unlock()
}

func (terminal *TerminalSSH) getProcessReader() (reader io.Reader) {
	terminal.processReader.RLock()
	reader = terminal.processReader.reader
	terminal.processReader.RUnlock()
	return
}

func (terminal *TerminalSSH) setProcessError(reader io.Reader) {
	terminal.processError.Lock()
	terminal.processError.reader = reader
	terminal.processError.Unlock()
}

func (terminal *TerminalSSH) getProcessError() (reader io.Reader) {
	terminal.processError.RLock()
	reader = terminal.processError.reader
	terminal.processError.RUnlock()
	return
}