package sshd

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"net"
	"os"
	"path/filepath"
	"testing"
	"time"

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

func newSigner(t *testing.T) ssh.Signer {
	t.Helper()
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	s, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		t.Fatal(err)
	}
	return s
}

func authLine(s ssh.Signer) string { return string(ssh.MarshalAuthorizedKey(s.PublicKey())) }

// start runs a server that accepts the given signers and returns its address.
func start(t *testing.T, allowed ...ssh.Signer) (addr, keysFile string, srv *Server) {
	t.Helper()
	dir := t.TempDir()
	keysFile = filepath.Join(dir, "authorized_keys")
	var lines string
	for _, s := range allowed {
		lines += authLine(s)
	}
	if err := os.WriteFile(keysFile, []byte("# who may sign in\n"+lines), 0o600); err != nil {
		t.Fatal(err)
	}
	srv, err := New(Config{AuthorizedKeys: keysFile, HostKey: filepath.Join(dir, "host_key"), Logf: t.Logf})
	if err != nil {
		t.Fatal(err)
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { srv.Serve(ctx, ln); close(done) }()
	t.Cleanup(func() { cancel(); <-done })
	return ln.Addr().String(), keysFile, srv
}

func dial(addr string, auth ...ssh.AuthMethod) (*ssh.Client, error) {
	return ssh.Dial("tcp", addr, &ssh.ClientConfig{User: "tester", Auth: auth, HostKeyCallback: ssh.InsecureIgnoreHostKey(), Timeout: 5 * time.Second})
}

func TestNewNeedsKeys(t *testing.T) {
	dir := t.TempDir()
	cfg := Config{AuthorizedKeys: filepath.Join(dir, "authorized_keys"), HostKey: filepath.Join(dir, "host_key")}
	if _, err := New(cfg); err == nil {
		t.Fatal("a missing authorized_keys file should be an error")
	}
	os.WriteFile(cfg.AuthorizedKeys, []byte("# nothing yet\n\n"), 0o600)
	if _, err := New(cfg); err == nil {
		t.Fatal("an authorized_keys file without keys should be an error")
	}
}

func TestOptionsLinesAreSkipped(t *testing.T) {
	a, b := newSigner(t), newSigner(t)
	path := filepath.Join(t.TempDir(), "authorized_keys")
	os.WriteFile(path, []byte(`command="/bin/false" `+authLine(a)+authLine(b)), 0o600)
	keys, err := loadAuthorizedKeys(path)
	if err != nil || len(keys) != 1 || ssh.FingerprintSHA256(keys[0]) != ssh.FingerprintSHA256(b.PublicKey()) {
		t.Fatalf("only the key without options should load, got %d keys, err %v", len(keys), err)
	}
}

func TestHostKeyIsStable(t *testing.T) {
	dir := t.TempDir()
	keys := filepath.Join(dir, "authorized_keys")
	os.WriteFile(keys, []byte(authLine(newSigner(t))), 0o600)
	cfg := Config{AuthorizedKeys: keys, HostKey: filepath.Join(dir, "sub", "host_key")}
	a, err := New(cfg)
	if err != nil {
		t.Fatal(err)
	}
	b, err := New(cfg)
	if err != nil {
		t.Fatal(err)
	}
	if a.Fingerprint() != b.Fingerprint() {
		t.Fatal("the host key changed between starts")
	}
}

func TestKeyAuth(t *testing.T) {
	good, bad := newSigner(t), newSigner(t)
	addr, _, _ := start(t, good)

	c, err := dial(addr, ssh.PublicKeys(good))
	if err != nil {
		t.Fatalf("authorized key refused: %v", err)
	}
	c.Close()
	if _, err := dial(addr, ssh.PublicKeys(bad)); err == nil {
		t.Fatal("unknown key was accepted")
	}
	if _, err := dial(addr, ssh.Password("hunter2")); err == nil {
		t.Fatal("a password was accepted")
	}
	if _, err := dial(addr, ssh.KeyboardInteractive(func(string, string, []string, []bool) ([]string, error) { return nil, nil })); err == nil {
		t.Fatal("keyboard-interactive was accepted")
	}
}

func TestKeysCanChangeWhileRunning(t *testing.T) {
	first, second := newSigner(t), newSigner(t)
	addr, keysFile, _ := start(t, first)
	if _, err := dial(addr, ssh.PublicKeys(second)); err == nil {
		t.Fatal("second key accepted before it was added")
	}
	f, _ := os.OpenFile(keysFile, os.O_APPEND|os.O_WRONLY, 0o600)
	f.WriteString(authLine(second))
	f.Close()
	c, err := dial(addr, ssh.PublicKeys(second))
	if err != nil {
		t.Fatalf("key added to the file was not picked up: %v", err)
	}
	c.Close()
}

func TestNoForwarding(t *testing.T) {
	good := newSigner(t)
	addr, _, _ := start(t, good)
	c, err := dial(addr, ssh.PublicKeys(good))
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	if conn, err := c.Dial("tcp", "127.0.0.1:1"); err == nil {
		conn.Close()
		t.Fatal("port forwarding was allowed")
	}
}