package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"

	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
)

func transfer(dst io.Writer, src io.Reader) {
	buf := make([]byte, 4096)
	for {
		n, err := src.Read(buf)
		switch {
		case n == 0:
			continue
		case err == io.EOF:
			return
		case err != nil:
			log.Println("transfer read:", err)
		}
		out, err := dst.Write(buf[:n])
		switch {
		case out != n:
			log.Println("transfer write: short write")
		case err != nil:
			log.Println("transfer write:", err)
		}
	}
}

func shellHandler(s ssh.Session) {
	cmd := exec.Command("/bin/rc")
	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Fatal(err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Fatal(err)
	}
	err = cmd.Start()
	if err != nil {
		log.Fatal(err)
	}
	//go io.Copy(stdin, s)
	go transfer(stdin, s)
	//go io.Copy(s, stderr)
	go transfer(s, stderr)
	//io.Copy(s, stdout)
	transfer(s, stdout)
}

func main() {
	priv, err := ioutil.ReadFile("./id_rsa")
	if err != nil {
		log.Fatal(err)
	}
	signer, err := gossh.ParsePrivateKey(priv)
	if err != nil {
		log.Fatal(err)
	}
	s := &ssh.Server{
		Addr:    ":2222",
		Handler: shellHandler,
	}
	s.AddHostKey(signer)
	err = s.SetOption(ssh.NoPty())
	if err != nil {
		log.Fatal(err)
	}
	if len(os.Args) > 1 {
		pubAuth := ssh.PublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
			data, err := ioutil.ReadFile(os.Args[1])
			if err != nil {
				log.Fatal(err)
			}
			allowed, _, _, _, _ := ssh.ParseAuthorizedKey(data)
			return ssh.KeysEqual(key, allowed)
		})
		err = s.SetOption(pubAuth)
		if err != nil {
			log.Fatal(err)
		}
	}

	log.Println("starting ssh server on port 2222...")
	log.Fatal(s.ListenAndServe())
}