// rsh is rsend's SSH transport. It is invoked three ways:
//
//	rsh [-l USER] [USER@]HOST CMD...   rsync remote shell (rsync -e), strict
//	rsh -keygen                        generate an ed25519 key, print the private key (PEM)
//	rsh -pubkey                        print the pubkey for RSH_KEY_DATA/RSH_KEY
//	rsh -scan USER@HOST                connect, print host-key fingerprint + line
//
// Transport mode mirrors what ssh does for rsync: it dials the host, runs the
// remote command, and bridges stdin/stdout/stderr. Host keys are verified
// strictly against RSH_KNOWN_HOSTS; an unknown or changed key fails loud.
// Key, known_hosts path, and port come from the environment because rsync owns
// the argument vector:
//
//	RSH_KEY           path to the ed25519 private key
//	RSH_KEY_DATA      the private key itself (PEM), preferred over RSH_KEY so the
//	                  app need never write the plaintext key to disk
//	RSH_KNOWN_HOSTS   path to the known_hosts file (required in transport mode)
//	RSH_PORT          TCP port (optional, default 22)
//
// Pure Go, no cgo.
package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"
)

const dialTimeout = 30 * time.Second

func main() {
	if err := run(os.Args[1:], os.Stdin, os.Stdout, os.Stderr); err != nil {
		var ee *ssh.ExitError
		if errors.As(err, &ee) {
			os.Exit(ee.ExitStatus())
		}
		fmt.Fprintf(os.Stderr, "rsh: %v\n", err)
		os.Exit(1)
	}
}

func run(args []string, in io.Reader, out, errw io.Writer) error {
	if len(args) >= 1 {
		switch args[0] {
		case "-keygen":
			if len(args) != 1 {
				return errors.New("usage: rsh -keygen")
			}
			return keygen(out)
		case "-pubkey":
			if len(args) != 1 {
				return errors.New("usage: rsh -pubkey")
			}
			return pubkey(out)
		case "-scan":
			if len(args) != 2 {
				return errors.New("usage: rsh -scan USER@HOST")
			}
			return scan(args[1], out)
		}
	}
	return transport(args, in, out, errw)
}

// transport runs the remote command and bridges I/O, verifying the host key
// strictly against RSH_KNOWN_HOSTS.
func transport(args []string, in io.Reader, out, errw io.Writer) error {
	user, host, cmd, err := parseTransport(args)
	if err != nil {
		return err
	}
	if user == "" {
		return errors.New("no remote user (expected -l USER or USER@HOST)")
	}
	if len(cmd) == 0 {
		return errors.New("no remote command")
	}
	kh := os.Getenv("RSH_KNOWN_HOSTS")
	if kh == "" {
		return errors.New("RSH_KNOWN_HOSTS not set")
	}
	cb, err := knownhosts.New(kh)
	if err != nil {
		return fmt.Errorf("known_hosts: %w (run Test connection first)", err)
	}
	client, err := dial(user, host, cb)
	if err != nil {
		return err
	}
	defer client.Close()
	go keepAlive(client)

	session, err := client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()
	session.Stdin = in
	session.Stdout = out
	session.Stderr = errw
	return session.Run(strings.Join(cmd, " "))
}

// scan connects (authenticating, so it also proves the key is installed),
// captures the presented host key, and prints its SHA256 fingerprint and a
// known_hosts line for the caller to pin.
func scan(target string, out io.Writer) error {
	user, host := splitUserHost(target)
	if user == "" {
		return errors.New("scan target must be USER@HOST")
	}
	var hostKey ssh.PublicKey
	capture := func(_ string, _ net.Addr, key ssh.PublicKey) error {
		hostKey = key
		return nil
	}
	client, err := dial(user, host, capture)
	if err != nil {
		return err
	}
	client.Close()

	addr := knownhosts.Normalize(net.JoinHostPort(host, strconv.Itoa(sshPort())))
	line := knownhosts.Line([]string{addr}, hostKey)
	fmt.Fprintf(out, "%s\n%s\n", ssh.FingerprintSHA256(hostKey), line)
	return nil
}

// keygen generates an ed25519 key and prints the private key in PEM form to
// out. Nothing touches disk: the caller owns persistence (the app stores it
// encrypted) and derives the public key with -pubkey.
func keygen(out io.Writer) error {
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return err
	}
	block, err := ssh.MarshalPrivateKey(priv, "rsend")
	if err != nil {
		return err
	}
	return pem.Encode(out, block)
}

// pubkey loads the private key (RSH_KEY_DATA or RSH_KEY) and prints its public
// key in authorized_keys format. The app uses it to validate and display a key
// the user imported: if loadSigner accepts it here, the transport will too.
func pubkey(out io.Writer) error {
	signer, err := loadSigner()
	if err != nil {
		return err
	}
	_, err = out.Write(ssh.MarshalAuthorizedKey(signer.PublicKey()))
	return err
}

// loadSigner reads the private key from RSH_KEY_DATA (the key itself) or, if
// that is unset, the file named by RSH_KEY.
func loadSigner() (ssh.Signer, error) {
	if data := os.Getenv("RSH_KEY_DATA"); data != "" {
		return ssh.ParsePrivateKey([]byte(data))
	}
	keyPath := os.Getenv("RSH_KEY")
	if keyPath == "" {
		return nil, errors.New("RSH_KEY or RSH_KEY_DATA not set")
	}
	pemBytes, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("read key: %w", err)
	}
	return ssh.ParsePrivateKey(pemBytes)
}

// dial opens an SSH connection authenticated with the configured key, verifying
// the host key with hostKey.
func dial(user, host string, hostKey ssh.HostKeyCallback) (*ssh.Client, error) {
	signer, err := loadSigner()
	if err != nil {
		return nil, err
	}
	cfg := &ssh.ClientConfig{
		User:            user,
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(signer)},
		HostKeyCallback: hostKey,
		Timeout:         dialTimeout,
	}
	addr := net.JoinHostPort(host, strconv.Itoa(sshPort()))
	return ssh.Dial("tcp", addr, cfg)
}

// parseTransport extracts the user, host, and remote command from the argument
// vector rsync hands to its remote shell: an optional "-l USER", the host
// (possibly USER@HOST), then the remote command.
func parseTransport(args []string) (user, host string, cmd []string, err error) {
	i := 0
	for i < len(args) {
		a := args[i]
		if a == "-l" {
			if i+1 >= len(args) {
				return "", "", nil, errors.New("-l requires an argument")
			}
			user = args[i+1]
			i += 2
			continue
		}
		if strings.HasPrefix(a, "-") {
			return "", "", nil, fmt.Errorf("unexpected option %q", a)
		}
		break
	}
	if i >= len(args) {
		return "", "", nil, errors.New("missing host")
	}
	host = args[i]
	cmd = args[i+1:]
	if u, h := splitUserHost(host); u != "" {
		host = h
		if user == "" {
			user = u
		}
	}
	return user, host, cmd, nil
}

func splitUserHost(s string) (user, host string) {
	if i := strings.Index(s, "@"); i >= 0 {
		return s[:i], s[i+1:]
	}
	return "", s
}

func sshPort() int {
	if s := os.Getenv("RSH_PORT"); s != "" {
		if p, err := strconv.Atoi(s); err == nil && p > 0 {
			return p
		}
	}
	return 22
}

// keepAlive pings the server so a long transfer is not dropped by an idle NAT.
// It returns once the connection is closed.
func keepAlive(client *ssh.Client) {
	t := time.NewTicker(60 * time.Second)
	defer t.Stop()
	for range t.C {
		if _, _, err := client.SendRequest("keepalive@openssh.com", true, nil); err != nil {
			return
		}
	}
}