package main

import (
	"bytes"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/pem"
	"errors"
	"io"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

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

// genClientKey makes a client key via keygen, writes it to a file for RSH_KEY,
// and returns the private key path and the corresponding public key.
func genClientKey(t *testing.T) (keyPath string, pub ssh.PublicKey) {
	t.Helper()
	var buf bytes.Buffer
	if err := keygen(&buf); err != nil {
		t.Fatal(err)
	}
	keyPath = filepath.Join(t.TempDir(), "id_ed25519")
	if err := os.WriteFile(keyPath, buf.Bytes(), 0o600); err != nil {
		t.Fatal(err)
	}
	signer, err := ssh.ParsePrivateKey(buf.Bytes())
	if err != nil {
		t.Fatal(err)
	}
	return keyPath, signer.PublicKey()
}

func writeKnownHosts(t *testing.T, host string, port int, hostKey ssh.PublicKey) string {
	t.Helper()
	addr := knownhosts.Normalize(net.JoinHostPort(host, strconv.Itoa(port)))
	line := knownhosts.Line([]string{addr}, hostKey)
	p := filepath.Join(t.TempDir(), "known_hosts")
	if err := os.WriteFile(p, []byte(line+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	return p
}

func setEnv(t *testing.T, key, kh string, port int) {
	t.Setenv("RSH_KEY", key)
	t.Setenv("RSH_KNOWN_HOSTS", kh)
	t.Setenv("RSH_PORT", strconv.Itoa(port))
}

func TestTransportBridgesStdio(t *testing.T) {
	keyPath, pub := genClientKey(t)
	echo := func(_ string, stdin io.Reader, stdout, _ io.Writer) int {
		io.Copy(stdout, stdin)
		return 0
	}
	srv := newTestServer(t, pub, echo)
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), srv.hostKey.PublicKey())
	setEnv(t, keyPath, kh, srv.port())

	want := []byte("the quick brown fox\x00\x01\x02 binary tail")
	var out bytes.Buffer
	if err := transport([]string{"-l", "u", "127.0.0.1", "cat"}, bytes.NewReader(want), &out, io.Discard); err != nil {
		t.Fatalf("transport: %v", err)
	}
	if !bytes.Equal(out.Bytes(), want) {
		t.Fatalf("bridged data mismatch:\n got %q\nwant %q", out.Bytes(), want)
	}
}

func TestTransportKeyFromData(t *testing.T) {
	keyPath, pub := genClientKey(t)
	echo := func(_ string, stdin io.Reader, stdout, _ io.Writer) int {
		io.Copy(stdout, stdin)
		return 0
	}
	srv := newTestServer(t, pub, echo)
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), srv.hostKey.PublicKey())
	keyData, err := os.ReadFile(keyPath)
	if err != nil {
		t.Fatal(err)
	}
	// RSH_KEY points nowhere; RSH_KEY_DATA must take precedence and work.
	t.Setenv("RSH_KEY", "/does/not/exist")
	t.Setenv("RSH_KEY_DATA", string(keyData))
	t.Setenv("RSH_KNOWN_HOSTS", kh)
	t.Setenv("RSH_PORT", strconv.Itoa(srv.port()))

	want := []byte("via key data")
	var out bytes.Buffer
	if err := transport([]string{"-l", "u", "127.0.0.1", "cat"}, bytes.NewReader(want), &out, io.Discard); err != nil {
		t.Fatalf("transport: %v", err)
	}
	if !bytes.Equal(out.Bytes(), want) {
		t.Fatalf("bridged data mismatch: %q", out.Bytes())
	}
}

func TestTransportRejectsUnknownHostKey(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 })

	// Pin a DIFFERENT key, so the server's real host key must be rejected.
	_, wrong := genClientKey(t)
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), wrong)
	setEnv(t, keyPath, kh, srv.port())

	if err := transport([]string{"-l", "u", "127.0.0.1", "true"}, bytes.NewReader(nil), io.Discard, io.Discard); err == nil {
		t.Fatal("expected host-key mismatch error, got nil")
	}
}

func TestTransportPropagatesExitCode(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 7 })
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), srv.hostKey.PublicKey())
	setEnv(t, keyPath, kh, srv.port())

	err := transport([]string{"-l", "u", "127.0.0.1", "false"}, bytes.NewReader(nil), io.Discard, io.Discard)
	var ee *ssh.ExitError
	if !errors.As(err, &ee) || ee.ExitStatus() != 7 {
		t.Fatalf("expected exit status 7, got %v", err)
	}
}

func TestDialTimeoutFromEnv(t *testing.T) {
	t.Setenv("RSH_CONNECT_TIMEOUT", "")
	if d, err := dialTimeout(); err != nil || d != defaultDialTimeout {
		t.Fatalf("unset: got %v, %v; want %v", d, err, defaultDialTimeout)
	}
	t.Setenv("RSH_CONNECT_TIMEOUT", "3")
	if d, err := dialTimeout(); err != nil || d != 3*time.Second {
		t.Fatalf("set: got %v, %v; want 3s", d, err)
	}
	for _, bad := range []string{"0", "-1", "abc", "3.5", "99999"} {
		t.Setenv("RSH_CONNECT_TIMEOUT", bad)
		if _, err := dialTimeout(); err == nil {
			t.Fatalf("RSH_CONNECT_TIMEOUT=%q: expected error, got nil", bad)
		}
	}
}

// A host we cannot reach must be reported with the marker the app greps for,
// so it can skip the rest of the folders bound to that remote.
func TestUnreachableHostIsMarked(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 })

	// A port with nothing behind it: the connect fails before any handshake.
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	dead := l.Addr().(*net.TCPAddr).Port
	l.Close()

	// Pin the dead port so the setup is the realistic one: a host rsend has
	// synced with before and now cannot reach.
	kh := writeKnownHosts(t, "127.0.0.1", dead, srv.hostKey.PublicKey())
	setEnv(t, keyPath, kh, dead)
	err = transport([]string{"-l", "u", "127.0.0.1", "true"}, bytes.NewReader(nil), io.Discard, io.Discard)
	if err == nil {
		t.Fatal("expected an error dialing a closed port, got nil")
	}
	if !strings.Contains(err.Error(), unreachablePrefix) {
		t.Fatalf("error %q does not carry %q", err, unreachablePrefix)
	}
}

// An auth or host-key failure is not "unreachable": the host answered. Marking
// it would make the app skip folders that a retry could still fix.
func TestHostKeyMismatchIsNotMarkedUnreachable(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 })
	_, wrong := genClientKey(t)
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), wrong)
	setEnv(t, keyPath, kh, srv.port())

	err := transport([]string{"-l", "u", "127.0.0.1", "true"}, bytes.NewReader(nil), io.Discard, io.Discard)
	if err == nil {
		t.Fatal("expected a host-key mismatch error, got nil")
	}
	if strings.Contains(err.Error(), unreachablePrefix) {
		t.Fatalf("host-key mismatch wrongly marked unreachable: %v", err)
	}
}

func TestPubkeyMatchesKeygen(t *testing.T) {
	keyPath, pub := genClientKey(t)
	keyData, err := os.ReadFile(keyPath)
	if err != nil {
		t.Fatal(err)
	}
	t.Setenv("RSH_KEY", "")
	t.Setenv("RSH_KEY_DATA", string(keyData))

	var out bytes.Buffer
	if err := pubkey(&out); err != nil {
		t.Fatalf("pubkey: %v", err)
	}
	got, _, _, _, err := ssh.ParseAuthorizedKey(out.Bytes())
	if err != nil {
		t.Fatalf("parse pubkey output: %v", err)
	}
	if !bytes.Equal(got.Marshal(), pub.Marshal()) {
		t.Fatal("pubkey output does not match the keygen public key")
	}
}

func TestPubkeyRejectsMissingKey(t *testing.T) {
	t.Setenv("RSH_KEY", "")
	t.Setenv("RSH_KEY_DATA", "")
	if err := pubkey(io.Discard); err == nil {
		t.Fatal("expected error when no key is configured")
	}
}

func TestPubkeyRejectsNonEd25519Key(t *testing.T) {
	private, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	block, err := ssh.MarshalPrivateKey(private, "test")
	if err != nil {
		t.Fatal(err)
	}
	t.Setenv("RSH_KEY", "")
	t.Setenv("RSH_KEY_DATA", string(pem.EncodeToMemory(block)))
	if err := pubkey(io.Discard); err == nil || !strings.Contains(err.Error(), "ssh-ed25519") {
		t.Fatalf("expected ed25519-only error, got %v", err)
	}
}

func TestScanPrintsFingerprintAndLine(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 })
	t.Setenv("RSH_KEY", keyPath)
	t.Setenv("RSH_PORT", strconv.Itoa(srv.port()))

	var out bytes.Buffer
	if err := scan("u@127.0.0.1", &out); err != nil {
		t.Fatalf("scan: %v", err)
	}
	lines := strings.SplitN(strings.TrimSpace(out.String()), "\n", 2)
	if len(lines) != 2 {
		t.Fatalf("want fingerprint + known_hosts line, got %q", out.String())
	}
	// Line 0 is "<keytype> SHA256:<base64>": the type is what the user compares
	// against a specific key file on the server.
	if lines[0] != "ssh-ed25519 "+ssh.FingerprintSHA256(srv.hostKey.PublicKey()) {
		t.Errorf("type + fingerprint line: %q", lines[0])
	}
	// The printed known_hosts line must validate the real host key.
	khPath := filepath.Join(t.TempDir(), "kh")
	if err := os.WriteFile(khPath, []byte(lines[1]+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	cb, err := knownhosts.New(khPath)
	if err != nil {
		t.Fatal(err)
	}
	addr := net.JoinHostPort("127.0.0.1", strconv.Itoa(srv.port()))
	remote := &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: srv.port()}
	if err := cb(addr, remote, srv.hostKey.PublicKey()); err != nil {
		t.Errorf("pinned line did not validate the host key: %v", err)
	}
}

// A stock sshd offers several host key types. x/crypto's default client
// preference puts ssh-ed25519 last, behind ecdsa and rsa, so without an
// explicit policy rsend would pin the ecdsa key while the README tells the user
// to verify the ed25519 fingerprint. Scanning must pick ed25519.
func TestScanPrefersEd25519(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServerKeys(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 }, true)
	t.Setenv("RSH_KEY", keyPath)
	t.Setenv("RSH_PORT", strconv.Itoa(srv.port()))

	var out bytes.Buffer
	if err := scan("u@127.0.0.1", &out); err != nil {
		t.Fatalf("scan: %v", err)
	}
	lines := strings.SplitN(strings.TrimSpace(out.String()), "\n", 2)
	want := "ssh-ed25519 " + ssh.FingerprintSHA256(srv.hostKey.PublicKey())
	if lines[0] != want {
		t.Fatalf("scan pinned the wrong host key:\n got %q\nwant %q", lines[0], want)
	}
}

// Once a type is pinned, only that type may be offered, so a server that also
// has an ecdsa key cannot be negotiated onto it.
func TestTransportOffersOnlyPinnedKeyType(t *testing.T) {
	keyPath, pub := genClientKey(t)
	echo := func(_ string, stdin io.Reader, stdout, _ io.Writer) int {
		io.Copy(stdout, stdin)
		return 0
	}
	srv := newTestServerKeys(t, pub, echo, true)
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), srv.hostKey.PublicKey())
	setEnv(t, keyPath, kh, srv.port())

	want := []byte("pinned to ed25519")
	var out bytes.Buffer
	if err := transport([]string{"-l", "u", "127.0.0.1", "cat"}, bytes.NewReader(want), &out, io.Discard); err != nil {
		t.Fatalf("transport: %v", err)
	}
	if !bytes.Equal(out.Bytes(), want) {
		t.Fatalf("bridged data mismatch: %q", out.Bytes())
	}
}

// A pin for a key type the server no longer has must fail loud at negotiation
// rather than fall through to some other key the user never verified.
func TestTransportRejectsWhenPinnedTypeAbsent(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 })

	// Pin an ecdsa key; the server only has ed25519.
	epriv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	epub, err := ssh.NewPublicKey(&epriv.PublicKey)
	if err != nil {
		t.Fatal(err)
	}
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), epub)
	setEnv(t, keyPath, kh, srv.port())

	err = transport([]string{"-l", "u", "127.0.0.1", "true"}, bytes.NewReader(nil), io.Discard, io.Discard)
	if err == nil {
		t.Fatal("expected a failure when the pinned key type is unavailable, got nil")
	}
	if strings.Contains(err.Error(), unreachablePrefix) {
		t.Fatalf("a reachable host must not be marked unreachable: %v", err)
	}
}

// With nothing pinned for this host:port the algorithm policy has nothing to
// go on, so it falls back to the modern preference and lets the known_hosts
// callback decide. The host must still be rejected, loudly, and not be mistaken
// for unreachable.
func TestTransportRejectsUnpinnedHost(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 })
	// Pin a different port, so the file parses but holds nothing for this one.
	kh := writeKnownHosts(t, "127.0.0.1", srv.port()+1, srv.hostKey.PublicKey())
	setEnv(t, keyPath, kh, srv.port())

	err := transport([]string{"-l", "u", "127.0.0.1", "true"}, bytes.NewReader(nil), io.Discard, io.Discard)
	if err == nil {
		t.Fatal("expected an unknown-host error, got nil")
	}
	if strings.Contains(err.Error(), unreachablePrefix) {
		t.Fatalf("a reachable host must not be marked unreachable: %v", err)
	}
}

// Re-scanning a host that is already pinned must reproduce the pinned key type,
// not the scan-time preference. Otherwise a config pinned by an older rsend
// (which negotiated ecdsa) would scan ed25519 and look like a key change on a
// server nobody touched.
func TestScanReproducesThePinnedKeyType(t *testing.T) {
	keyPath, pub := genClientKey(t)
	srv := newTestServerKeys(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 }, true)

	// Pin the ecdsa key, as an older rsend would have.
	ecdsaPub := srv.extraKey.PublicKey()
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), ecdsaPub)
	t.Setenv("RSH_KEY", keyPath)
	t.Setenv("RSH_KNOWN_HOSTS", kh)
	t.Setenv("RSH_PORT", strconv.Itoa(srv.port()))

	var out bytes.Buffer
	if err := scan("u@127.0.0.1", &out); err != nil {
		t.Fatalf("scan: %v", err)
	}
	lines := strings.SplitN(strings.TrimSpace(out.String()), "\n", 2)
	want := ecdsaPub.Type() + " " + ssh.FingerprintSHA256(ecdsaPub)
	if lines[0] != want {
		t.Fatalf("scan did not reproduce the pinned key:\n got %q\nwant %q", lines[0], want)
	}
	// The reproduced known_hosts line must equal the stored pin exactly, or the
	// app's pin comparison reports a change.
	stored, err := os.ReadFile(kh)
	if err != nil {
		t.Fatal(err)
	}
	if lines[1] != strings.TrimSpace(string(stored)) {
		t.Fatalf("reproduced line differs from the pin:\n got %q\nwant %q", lines[1], strings.TrimSpace(string(stored)))
	}
}

// A host whose key type was genuinely rotated must still be scannable, or the
// app can never show its "host key changed" prompt and the user has to delete
// and recreate the remote. Preferring the pinned type is right; requiring it is
// not.
func TestScanStillWorksAfterAKeyTypeRotation(t *testing.T) {
	keyPath, pub := genClientKey(t)
	// The server now has only ed25519, as if its ecdsa key had been removed.
	srv := newTestServer(t, pub, func(_ string, _ io.Reader, _, _ io.Writer) int { return 0 })

	// The stored pin is an ecdsa key the server no longer offers.
	epriv, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	epub, err := ssh.NewPublicKey(&epriv.PublicKey)
	if err != nil {
		t.Fatal(err)
	}
	kh := writeKnownHosts(t, "127.0.0.1", srv.port(), epub)
	t.Setenv("RSH_KEY", keyPath)
	t.Setenv("RSH_KNOWN_HOSTS", kh)
	t.Setenv("RSH_PORT", strconv.Itoa(srv.port()))

	var out bytes.Buffer
	if err := scan("u@127.0.0.1", &out); err != nil {
		t.Fatalf("scan must reach the new key so the caller can warn about it: %v", err)
	}
	lines := strings.SplitN(strings.TrimSpace(out.String()), "\n", 2)
	want := "ssh-ed25519 " + ssh.FingerprintSHA256(srv.hostKey.PublicKey())
	if lines[0] != want {
		t.Fatalf("scan did not report the rotated key:\n got %q\nwant %q", lines[0], want)
	}
}

func TestPinnedAlgos(t *testing.T) {
	_, pub := genClientKey(t)
	kh := writeKnownHosts(t, "10.0.0.1", 2222, pub)
	addr := knownhosts.Normalize(net.JoinHostPort("10.0.0.1", "2222"))

	got, err := pinnedAlgos(kh, addr)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 1 || got[0] != ssh.KeyAlgoED25519 {
		t.Fatalf("pinned algos: got %v want [%s]", got, ssh.KeyAlgoED25519)
	}
	// A host with no pin yields nothing rather than a default.
	other, err := pinnedAlgos(kh, knownhosts.Normalize(net.JoinHostPort("10.0.0.2", "22")))
	if err != nil {
		t.Fatal(err)
	}
	if len(other) != 0 {
		t.Fatalf("unpinned host: got %v want none", other)
	}
	// An ssh-rsa pin allows the SHA-2 signature algorithms and not SHA-1.
	got = algosForKeyType(ssh.KeyAlgoRSA)
	if len(got) != 2 || got[0] != ssh.KeyAlgoRSASHA512 || got[1] != ssh.KeyAlgoRSASHA256 {
		t.Fatalf("rsa expansion: %v", got)
	}
	for _, a := range got {
		if a == ssh.KeyAlgoRSA {
			t.Fatal("plain ssh-rsa (SHA-1) must not be offered")
		}
	}
}

// A @revoked entry names a key that must never be trusted, and @cert-authority
// names a signing key, not a host key. Neither may contribute an algorithm.
func TestPinnedAlgosIgnoresMarkers(t *testing.T) {
	_, pub := genClientKey(t)
	addr := knownhosts.Normalize(net.JoinHostPort("10.0.0.9", "22"))
	line := knownhosts.Line([]string{addr}, pub)
	p := filepath.Join(t.TempDir(), "kh")
	if err := os.WriteFile(p, []byte("@revoked "+line+"\n"), 0o600); err != nil {
		t.Fatal(err)
	}
	got, err := pinnedAlgos(p, addr)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 0 {
		t.Fatalf("a revoked entry must contribute nothing, got %v", got)
	}
}

func TestParseTransport(t *testing.T) {
	cases := []struct {
		args       []string
		user, host string
		cmd        string
		wantErr    bool
	}{
		{[]string{"-l", "alice", "host", "rsync", "--server"}, "alice", "host", "rsync --server", false},
		{[]string{"bob@host", "echo", "hi"}, "bob", "host", "echo hi", false},
		{[]string{"-l", "alice", "carol@host", "x"}, "alice", "host", "x", false},
		{[]string{"host"}, "", "host", "", false},
		{[]string{"-l"}, "", "", "", true},
		{[]string{"-p", "22", "host", "x"}, "", "", "", true},
		// A bare IPv6 literal already works; a bracketed one must not survive
		// into net.JoinHostPort, which would double-bracket it.
		{[]string{"-l", "u", "2001:db8::1", "x"}, "u", "2001:db8::1", "x", false},
		{[]string{"-l", "u", "[2001:db8::1]", "x"}, "u", "2001:db8::1", "x", false},
		{[]string{"u@[2001:db8::1]", "x"}, "u", "2001:db8::1", "x", false},
		{[]string{"[::1]"}, "", "::1", "", false},
	}
	for _, c := range cases {
		u, h, cmd, err := parseTransport(c.args)
		if c.wantErr {
			if err == nil {
				t.Errorf("%v: expected error", c.args)
			}
			continue
		}
		if err != nil {
			t.Errorf("%v: %v", c.args, err)
			continue
		}
		if u != c.user || h != c.host || strings.Join(cmd, " ") != c.cmd {
			t.Errorf("%v: got (%q,%q,%q) want (%q,%q,%q)",
				c.args, u, h, strings.Join(cmd, " "), c.user, c.host, c.cmd)
		}
	}
}