package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/binary"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"golang.org/x/crypto/ssh"
	"github.com/creack/pty/v2"
)

func main() {
	if os.Getenv("COMMAND") == "" {
		log.Fatal("COMMAND environment variable must be set")
	}
	config := &ssh.ServerConfig{
		NoClientAuth: true,
	}
	_, key, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatal(err)
	}
	signer, err := ssh.NewSignerFromKey(key)
	if err != nil {
		log.Fatal(err)
	}
	config.AddHostKey(signer)
	listener, err := net.Listen("tcp", ":22")
	if err != nil {
		log.Fatal(err)
	}
	log.Println("SSH server listening on :22")
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Accept error:", err)
			continue
		}
		go handleConn(conn, config)
	}
}

func handleConn(conn net.Conn, config *ssh.ServerConfig) {
	sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
	if err != nil {
		log.Println("ServerConn error:", err)
		conn.Close()
		return
	}
	log.Println("New connection from", sshConn.RemoteAddr(), "user", sshConn.User())
	go ssh.DiscardRequests(reqs)
	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}
		channel, requests, err := newChannel.Accept()
		if err != nil {
			log.Println("Channel accept error:", err)
			continue
		}
		go handleChannel(channel, requests)
	}
	sshConn.Wait()
}

func handleChannel(channel ssh.Channel, requests <-chan *ssh.Request) {
	defer channel.Close()
	var ptmx *os.File
	for req := range requests {
		switch req.Type {
		case "pty-req":
			req.Reply(true, nil)
		case "window-change":
			if ptmx != nil {
				width := binary.BigEndian.Uint32(req.Payload)
				height := binary.BigEndian.Uint32(req.Payload[4:])
				pty.Setsize(ptmx, &pty.Winsize{Cols: uint16(width), Rows: uint16(height)})
			}
			req.Reply(true, nil)
		case "shell":
			req.Reply(true, nil)
			command := os.Getenv("COMMAND")
			if command == "" {
				command = "/app/tui"
			}
			cmd := exec.Command(command)
			cmd.Env = []string{"PATH=/bin"}
			cmd.Dir = "/"
			var err error
			ptmx, err = pty.Start(cmd)
			if err != nil {
				log.Println("PTY start error:", err)
				return
			}
			go func() {
				defer ptmx.Close()
				go io.Copy(channel, ptmx)
				go io.Copy(ptmx, channel)
				cmd.Wait()
				channel.Close()
			}()
		case "exec":
			req.Reply(true, nil)
			command := string(req.Payload[4:])
			runCommand(channel, command)
			return
		default:
			req.Reply(false, nil)
		}
	}
}

func runCommand(channel ssh.Channel, command string) {
	cmd := exec.Command("/bin/bash", "-c", command)
	cmd.Env = []string{"PATH=/bin"}
	cmd.Dir = "/"
	stdin, err := cmd.StdinPipe()
	if err != nil {
		log.Println("StdinPipe error:", err)
		return
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Println("StdoutPipe error:", err)
		return
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		log.Println("StderrPipe error:", err)
		return
	}
	if err := cmd.Start(); err != nil {
		log.Println("Start error:", err)
		return
	}
	go io.Copy(stdin, channel)
	go io.Copy(channel, stdout)
	go io.Copy(channel, stderr)
	cmd.Wait()
}