package main

import (
	"bufio"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"os"
	"sync"
)

func MuxShell(w io.Writer, r io.Reader) (chan<- string, <-chan string) {
	in := make(chan string, 1)
	out := make(chan string, 1)
	var wg sync.WaitGroup
	wg.Add(1) //for the shell itself
	go func() {
		for cmd := range in {
			wg.Add(1)
			w.Write([]byte(cmd + "\n"))
			wg.Wait()
		}
	}()
	go func() {
		var (
			buf [65 * 1024]byte
			t   int
		)
		for {
			n, err := r.Read(buf[t:])
			if err != nil {
				close(in)
				close(out)
				return
			}
			t += n
			if buf[t-2] == '$' { //assuming the $PS1 == 'sh-4.3$ '
				out <- string(buf[:t])
				wg.Done()
			}
		}
	}()
	return in, out
}

func main() {
	config1 := &ssh.ClientConfig {
		User: "lab_42_2018",
		Auth: []ssh.AuthMethod{
			ssh.Password("DFfdffuy934"),
		},
	}
	config1.HostKeyCallback = ssh.InsecureIgnoreHostKey()

	config2 := &ssh.ClientConfig{
		User: "lab2_42_2018",
		Auth: []ssh.AuthMethod{
			ssh.Password("DFfdffuy934"),
		},
	}
	config2.HostKeyCallback = ssh.InsecureIgnoreHostKey()

	hosts := [2]string{"lab.posevin.com:22", "lab2.posevin.com:22"}

	var(
		configs [2]*ssh.ClientConfig
		clients [2]*ssh.Client
		sessions [2]*ssh.Session
		writers [2]io.WriteCloser
		readers [2]io.Reader
		ins [2]chan<-string
		outs [2]<-chan string
		err error
	)

	configs[0]=config1
	configs[1]=config2

	for i:=0; i<len(hosts); i++ {
		clients[i], err = ssh.Dial("tcp", hosts[i], configs[i])
		if err != nil {
			panic(err)
		}

		defer clients[i].Close()
		sessions[i], err = clients[i].NewSession()

		if err != nil {
			log.Fatalf("unable to create session: %s", err)
		}
		defer sessions[i].Close()

		modes := ssh.TerminalModes{
			ssh.ECHO: 0,
		}

		if err := sessions[i].RequestPty("xterm", 80, 40, modes); err != nil {
			log.Fatal(err)
		}


		writers[i], err = sessions[i].StdinPipe()
		if err != nil {
			panic(err)
		}
		readers[i], err = sessions[i].StdoutPipe()
		if err != nil {
			panic(err)
		}

		ins[i], outs[i] = MuxShell(writers[i], readers[i])
		if err := sessions[i].Start("/bin/sh"); err != nil {
			log.Fatal(err)
		}
		fmt.Print(<-outs[i]) //ignore the shell output
	}

	for {
		reader := bufio.NewReader(os.Stdin)
		str, _ := reader.ReadString('\n')

		for i:=0; i<len(hosts); i++ {
			fmt.Print(len(hosts))
			ins[i] <- str
		}
		for i:=0; i<len(hosts); i++{
			fmt.Print(<-outs[i])
			fmt.Print("------------------------------" + "\n")
		}
	}
}