// 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 type + 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, and
// only the already-pinned key type is offered, so a server that switches key
// type fails at negotiation rather than presenting an unverified key.
// 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)
//
// The only non-standard Go code linked here is golang.org/x/crypto. The
// Android binaries are built with cgo against Bionic (see rsh/build.sh); the
// host binary is not.
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"
)

// defaultDialTimeout bounds the TCP connect and the SSH handshake. A firewall
// that drops packets instead of refusing them makes a connect hang for the
// whole budget, so keep it short: a reachable host completes the handshake
// well inside this even on poor mobile networks.
const defaultDialTimeout = 10 * time.Second

// unreachablePrefix marks a failure to reach the host at all, as opposed to a
// rejected key or a mismatched host key. The app greps rsync's merged output
// for it so it can skip the remaining folders on the same dead remote instead
// of paying the timeout once per folder.
const unreachablePrefix = "unreachable: "

// scanAlgos is the host-key preference used when first pinning a host. The
// x/crypto default puts ssh-ed25519 last, behind ecdsa, rsa and dss, so a stock
// server would get its ecdsa key pinned while the README tells the user to
// verify the ed25519 one. Prefer ed25519, accept the other modern types if that
// is all the server has, and never ssh-rsa (SHA-1) or ssh-dss.
var scanAlgos = []string{
	ssh.KeyAlgoED25519,
	ssh.KeyAlgoECDSA256,
	ssh.KeyAlgoECDSA384,
	ssh.KeyAlgoECDSA521,
	ssh.KeyAlgoRSASHA512,
	ssh.KeyAlgoRSASHA256,
}

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)
	}
	port, err := sshPort()
	if err != nil {
		return err
	}
	algos, err := pinnedAlgos(kh, knownhosts.Normalize(net.JoinHostPort(host, strconv.Itoa(port))))
	if err != nil {
		return err
	}
	if len(algos) == 0 {
		// No literal pin for this address. The file may still authorise the
		// host through a hashed or wildcard entry, which the literal match
		// deliberately does not understand, so fall back to the modern
		// preference and let the callback decide. Verification stays strict:
		// an unknown or changed key still fails loud, just after the dial.
		algos = scanAlgos
	}
	client, err := dial(user, host, cb, algos)
	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")
	}
	port, err := sshPort()
	if err != nil {
		return err
	}
	addr := knownhosts.Normalize(net.JoinHostPort(host, strconv.Itoa(port)))
	// Prefer the key type already pinned for this host, then fall back to the
	// general preference. A server usually offers several types; picking a
	// different one than last time would reproduce a different line and read as
	// a host-key change on a server nobody touched. Preferring rather than
	// requiring matters: negotiation walks the client list in order, so an
	// untouched server still reproduces its pinned line, while one whose key
	// type was genuinely rotated presents the new key and reaches the "host key
	// changed" prompt instead of failing with no common algorithm and no way to
	// accept the new key short of deleting the remote.
	algos := scanAlgos
	if kh := os.Getenv("RSH_KNOWN_HOSTS"); kh != "" {
		if pinned, perr := pinnedAlgos(kh, addr); perr == nil && len(pinned) > 0 {
			algos = append(append([]string{}, pinned...), scanAlgos...)
		}
	}
	var hostKey ssh.PublicKey
	capture := func(_ string, _ net.Addr, key ssh.PublicKey) error {
		hostKey = key
		return nil
	}
	client, err := dial(user, host, capture, algos)
	if err != nil {
		return err
	}
	client.Close()

	line := knownhosts.Line([]string{addr}, hostKey)
	// The type belongs next to the fingerprint: the user is told to compare it
	// against a specific key file on the server, and "SHA256:..." alone does not
	// say which one.
	fmt.Fprintf(out, "%s %s\n%s\n", hostKey.Type(), ssh.FingerprintSHA256(hostKey), line)
	return nil
}

// pinnedAlgos returns the host-key algorithms to offer for addr: exactly the
// key types already pinned for it. Offering only the pinned type turns a server
// that switches key type into a loud negotiation failure instead of a silent
// prompt to trust a key the user never verified.
//
// rsend writes this file itself from its own pins, one exact unhashed address
// per line (knownhosts.Line), so matching the address literally is enough.
func pinnedAlgos(khPath, addr string) ([]string, error) {
	data, err := os.ReadFile(khPath)
	if err != nil {
		return nil, fmt.Errorf("known_hosts: %w", err)
	}
	var algos []string
	seen := make(map[string]bool)
	for rest := data; len(rest) > 0; {
		marker, hosts, key, _, next, perr := ssh.ParseKnownHosts(rest)
		if perr == io.EOF {
			break
		}
		if perr != nil {
			return nil, fmt.Errorf("known_hosts: %w", perr)
		}
		rest = next
		// A @revoked line names a key that must never be accepted, and a
		// @cert-authority line names a signing key rather than a host key.
		// Neither says anything about what this host may present.
		if marker != "" {
			continue
		}
		for _, h := range hosts {
			if h != addr {
				continue
			}
			for _, a := range algosForKeyType(key.Type()) {
				if !seen[a] {
					seen[a] = true
					algos = append(algos, a)
				}
			}
		}
	}
	return algos, nil
}

// algosForKeyType maps a pinned known_hosts key type to the signature
// algorithms a server may use with it. Only RSA differs: an "ssh-rsa" pin
// names the same key as an rsa-sha2-* signature. Plain ssh-rsa is SHA-1 and is
// left out, matching scanAlgos and OpenSSH's own default.
func algosForKeyType(t string) []string {
	if t == ssh.KeyAlgoRSA {
		return []string{ssh.KeyAlgoRSASHA512, ssh.KeyAlgoRSASHA256}
	}
	return []string{t}
}

// 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) {
	var signer ssh.Signer
	var err error
	if data := os.Getenv("RSH_KEY_DATA"); data != "" {
		signer, err = ssh.ParsePrivateKey([]byte(data))
	} else {
		keyPath := os.Getenv("RSH_KEY")
		if keyPath == "" {
			return nil, errors.New("RSH_KEY or RSH_KEY_DATA not set")
		}
		pemBytes, readErr := os.ReadFile(keyPath)
		if readErr != nil {
			return nil, fmt.Errorf("read key: %w", readErr)
		}
		signer, err = ssh.ParsePrivateKey(pemBytes)
	}
	if err != nil {
		return nil, err
	}
	if signer.PublicKey().Type() != ssh.KeyAlgoED25519 {
		return nil, fmt.Errorf("client key must be %s", ssh.KeyAlgoED25519)
	}
	return signer, nil
}

// dial opens an SSH connection authenticated with the configured key, verifying
// the host key with hostKey and offering only the host-key algorithms in algos.
func dial(user, host string, hostKey ssh.HostKeyCallback, algos []string) (*ssh.Client, error) {
	signer, err := loadSigner()
	if err != nil {
		return nil, err
	}
	port, err := sshPort()
	if err != nil {
		return nil, err
	}
	timeout, err := dialTimeout()
	if err != nil {
		return nil, err
	}
	cfg := &ssh.ClientConfig{
		User:              user,
		Auth:              []ssh.AuthMethod{ssh.PublicKeys(signer)},
		HostKeyCallback:   hostKey,
		HostKeyAlgorithms: algos,
		Timeout:           timeout,
	}
	addr := net.JoinHostPort(host, strconv.Itoa(port))
	conn, err := net.DialTimeout("tcp", addr, timeout)
	if err != nil {
		return nil, fmt.Errorf("%s%w", unreachablePrefix, err)
	}
	if err := conn.SetDeadline(time.Now().Add(timeout)); err != nil {
		conn.Close()
		return nil, err
	}
	clientConn, chans, reqs, err := ssh.NewClientConn(conn, addr, cfg)
	if err != nil {
		conn.Close()
		return nil, err
	}
	if err := conn.SetDeadline(time.Time{}); err != nil {
		clientConn.Close()
		return nil, err
	}
	return ssh.NewClient(clientConn, chans, reqs), nil
}

// 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")
	}
	cmd = args[i+1:]
	u, h := splitUserHost(args[i])
	host = h
	if u != "" && 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], unbracket(s[i+1:])
	}
	return "", unbracket(s)
}

// unbracket strips the brackets people habitually put around an IPv6 literal.
// A bare literal already works, because net.JoinHostPort adds the brackets and
// knownhosts.Normalize takes them back off; a pre-bracketed one would be
// double-bracketed into an address that can be neither dialled nor matched
// against a pin.
func unbracket(h string) string {
	if len(h) > 1 && h[0] == '[' && h[len(h)-1] == ']' {
		return h[1 : len(h)-1]
	}
	return h
}

// dialTimeout returns the connect budget, overridable in whole seconds with
// RSH_CONNECT_TIMEOUT for running rsh by hand against a slow or filtered host.
func dialTimeout() (time.Duration, error) {
	s := os.Getenv("RSH_CONNECT_TIMEOUT")
	if s == "" {
		return defaultDialTimeout, nil
	}
	n, err := strconv.Atoi(s)
	if err != nil || n < 1 || n > 3600 {
		return 0, fmt.Errorf("invalid RSH_CONNECT_TIMEOUT %q", s)
	}
	return time.Duration(n) * time.Second, nil
}

func sshPort() (int, error) {
	s := os.Getenv("RSH_PORT")
	if s == "" {
		return 22, nil
	}
	p, err := strconv.Atoi(s)
	if err != nil || p < 1 || p > 65535 {
		return 0, fmt.Errorf("invalid RSH_PORT %q", s)
	}
	return p, nil
}

// 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
		}
	}
}