package main

import (
	"bufio"
	"fmt"
	"golang.org/x/crypto/ssh"
	"io"
	"log"
	"os"
	"strconv"
	"strings"
	"syscall"
	"time"
)

/**
Compile: go build remotessh.go
Call as client: ./remotessh
Call as target: ./remotessh <exit code>
*/


const privateKey = `
-----BEGIN OPENSSH PRIVATE KEY-----
<< private key >>
-----END OPENSSH PRIVATE KEY-----
`

// Dummy logger for stdout and stderr
type MyReader struct {
	prefix string
}

func (r MyReader) Write(src []byte) (int, error) {
	fmt.Println(r.prefix, string(src))
	return len(src), nil
}

type MyWriter struct {
	data string
	offset int
}

func (w *MyWriter) Read(dst []byte) (int, error) {
	if w.offset >= len(w.data) {
		return 0, io.EOF
	}

	length := copy(dst, w.data[w.offset:])
	w.offset += length
	return length, nil
}

func main() {
	if len(os.Args) > 1 {
		// Act as basic programmable command
		fmt.Fprintf(os.Stdout, "Some output to stdout\n")
		fmt.Fprintf(os.Stderr, "Some other output to stderr\n")
		
		// Read from stdin until end marker is reached
		scanner := bufio.NewScanner(os.Stdin)
		buf := ""
		for scanner.Scan() {
			buf += scanner.Text()
			if strings.Contains(buf, "<<end>>") {
				break
			}
		}

		if err := scanner.Err(); err != nil {
			fmt.Fprintf(os.Stderr,"Error reading from stdin: %v", err)
		} else {
			fmt.Fprintf(os.Stdout, "Received on stdin: %s", buf)
		}

		code, _ := strconv.Atoi(os.Args[1])
		syscall.Exit(code)
	}


	key, err := ssh.ParsePrivateKey([]byte(privateKey))
	if err != nil {
		log.Fatalf("Parse key: %v", err)
	}

	config := &ssh.ClientConfig{
		User: "tobias",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(key),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}
	client, err := ssh.Dial("tcp", "localhost:22", config)
	if err != nil {
		log.Fatal("Failed to dial: ", err)
	}

	// Each ClientConn can support multiple interactive sessions,
	// represented by a Session.
	session, err := client.NewSession()
	if err != nil {
		log.Fatal("Failed to create session: ", err)
	}
	defer session.Close()

	// Output and input
	session.Stdout = MyReader{prefix: "stdout:"}
	session.Stderr = MyReader{prefix: "stderr:"}
	session.Stdin = &MyWriter{data: "Some input, some more input, the <<end>>"}

	t0 := time.Now()
	err = session.Start("<path to remotessh>/remotessh 1")
	if err != nil {
		log.Fatalf("Error starting command: %v", err)
	}

	err = session.Wait()
	if err != nil {
		if e, ok := err.(*ssh.ExitError); ok {
			fmt.Printf("ExitError, status: %d, signal: %s, message: %s\n", e.ExitStatus(), e.Signal(), e.Msg())
		} else {
			fmt.Printf("Unknown error, %v", e)
		}
	}
	fmt.Println("Time taken", time.Since(t0).Seconds(), "s")
}