package nebula

import (
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"encoding/pem"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/slackhq/nebula/commands"
	"github.com/slackhq/nebula/config"
	"github.com/slackhq/nebula/sshd"
	"github.com/slackhq/nebula/test"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/crypto/ssh"
)

type sshdReloadHarness struct {
	c       *config.C
	reg     *commands.Registry
	server  *sshd.SSHServer
	keyFile string
	signer  ssh.Signer
	listen  string
	users   []string
	extra   string
}

func newSSHDReloadHarness(t *testing.T) *sshdReloadHarness {
	t.Helper()
	l := test.NewLogger()
	h := &sshdReloadHarness{keyFile: filepath.Join(t.TempDir(), "host_key"), listen: "127.0.0.1:0", users: []string{"test"}}
	h.writeHostKey(t, newEd25519Key(t))

	_, userKey, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	h.signer, err = ssh.NewSignerFromKey(userKey)
	require.NoError(t, err)

	h.c = config.NewC(l)
	require.NoError(t, h.c.LoadString(h.config()))
	h.reg = commands.New()
	h.server, err = sshd.NewSSHServer(t.Context(), l, h.reg)
	require.NoError(t, err)
	wireSSHReload(l, h.server, h.c)
	run, err := configSSH(l, h.server, h.c, nil)
	require.NoError(t, err)
	go run()
	t.Cleanup(h.server.Stop)
	require.Eventually(t, h.server.Running, 5*time.Second, 10*time.Millisecond)
	return h
}

func newEd25519Key(t *testing.T) crypto.Signer {
	_, key, err := ed25519.GenerateKey(rand.Reader)
	require.NoError(t, err)
	return key
}

// writeHostKey replaces the host key file and returns the public key a client should now see
func (h *sshdReloadHarness) writeHostKey(t *testing.T, key crypto.Signer) ssh.PublicKey {
	t.Helper()
	block, err := ssh.MarshalPrivateKey(key, "")
	require.NoError(t, err)
	require.NoError(t, os.WriteFile(h.keyFile, pem.EncodeToMemory(block), 0o600))
	pub, err := ssh.NewPublicKey(key.Public())
	require.NoError(t, err)
	return pub
}

func (h *sshdReloadHarness) config() string {
	users := ""
	for _, u := range h.users {
		users += fmt.Sprintf("    - user: %s\n      keys: [%q]\n", u, string(ssh.MarshalAuthorizedKey(h.signer.PublicKey())))
	}
	return fmt.Sprintf("sshd:\n  enabled: true\n  listen: %s\n  host_key: %s\n  authorized_users:\n%s%s",
		h.listen, h.keyFile, users, h.extra)
}

// dial connects to wherever the server listens now, a restart on :0 moves it. hostKeyAlgorithms limits which host keys
// the client accepts, none means any
func (h *sshdReloadHarness) dial(hostKeyAlgorithms ...string) (*ssh.Client, ssh.PublicKey, error) {
	addr := h.server.Addr()
	if addr == nil {
		return nil, nil, fmt.Errorf("not listening")
	}
	var hostKey ssh.PublicKey
	c, err := ssh.Dial("tcp", addr.String(), &ssh.ClientConfig{
		User: "test",
		Auth: []ssh.AuthMethod{ssh.PublicKeys(h.signer)},
		HostKeyCallback: func(_ string, _ net.Addr, key ssh.PublicKey) error {
			hostKey = key
			return nil
		},
		HostKeyAlgorithms: hostKeyAlgorithms,
		Timeout:           time.Second,
	})
	return c, hostKey, err
}

func (h *sshdReloadHarness) connect(t *testing.T) (*ssh.Client, ssh.PublicKey) {
	t.Helper()
	var client *ssh.Client
	var hostKey ssh.PublicKey
	require.Eventually(t, func() bool {
		var err error
		client, hostKey, err = h.dial()
		return err == nil
	}, 5*time.Second, 20*time.Millisecond)
	t.Cleanup(func() { _ = client.Close() })
	return client, hostKey
}

// dropped is true when the server closed the client's connection within wait
func dropped(client *ssh.Client, wait time.Duration) bool {
	done := make(chan struct{})
	go func() {
		_ = client.Wait()
		close(done)
	}()
	select {
	case <-done:
		return true
	case <-time.After(wait):
		return false
	}
}

// Restarting sshd drops every open session, the one that asked for the reload included. A reload that touches nothing
// sshd reads leaves it alone
func TestSSHDReload_UnchangedKeepsSessions(t *testing.T) {
	h := newSSHDReloadHarness(t)
	client, _ := h.connect(t)

	require.NoError(t, h.c.ReloadConfigString(h.config()+"punchy:\n  punch: true\n").Err)
	assert.False(t, dropped(client, 500*time.Millisecond), "an unchanged sshd was restarted")
}

// Every key the server reads restarts it when it changes
func TestSSHDReload_ChangedConfigRestarts(t *testing.T) {
	cases := map[string]func(t *testing.T, h *sshdReloadHarness){
		"listen":           func(t *testing.T, h *sshdReloadHarness) { h.listen = "localhost:0" },
		"authorized_users": func(t *testing.T, h *sshdReloadHarness) { h.users = append(h.users, "other") },
		"trusted_cas":      func(t *testing.T, h *sshdReloadHarness) { h.extra = "  trusted_cas: []\n" },
		"host_key": func(t *testing.T, h *sshdReloadHarness) {
			h.keyFile = filepath.Join(t.TempDir(), "other_host_key")
			h.writeHostKey(t, newEd25519Key(t))
		},
	}
	for name, change := range cases {
		t.Run(name, func(t *testing.T) {
			h := newSSHDReloadHarness(t)
			client, _ := h.connect(t)

			change(t, h)
			require.NoError(t, h.c.ReloadConfigString(h.config()).Err)
			assert.True(t, dropped(client, 5*time.Second), "a changed sshd config was not applied")
			require.Eventually(t, h.server.Running, 5*time.Second, 10*time.Millisecond)
			h.connect(t)
		})
	}
}

// The config can stay the same while the host key file it names changes. New connections get the new key, open
// sessions are left alone
func TestSSHDReload_ChangedHostKeyFile(t *testing.T) {
	h := newSSHDReloadHarness(t)
	client, _ := h.connect(t)

	want := h.writeHostKey(t, newEd25519Key(t))
	require.NoError(t, h.c.ReloadConfigString(h.config()).Err)
	_, got := h.connect(t)
	assert.Equal(t, want.Marshal(), got.Marshal(), "a new host key was not applied")
	assert.False(t, dropped(client, 500*time.Millisecond), "a new host key restarted sshd")
}

// A host key of another algorithm replaces the old one rather than being offered next to it
func TestSSHDReload_HostKeyAlgorithmChange(t *testing.T) {
	h := newSSHDReloadHarness(t)
	h.connect(t)

	rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
	require.NoError(t, err)
	want := h.writeHostKey(t, rsaKey)
	require.NoError(t, h.c.ReloadConfigString(h.config()).Err)

	_, got := h.connect(t)
	assert.Equal(t, want.Marshal(), got.Marshal())
	_, _, err = h.dial(ssh.KeyAlgoED25519)
	assert.Error(t, err, "the old ed25519 host key is still offered")
}

// A server that isn't running, say its port was busy, is tried again on every reload as before
func TestSSHDReload_NotRunningRestarts(t *testing.T) {
	h := newSSHDReloadHarness(t)
	h.server.Stop()
	require.Eventually(t, func() bool { return !h.server.Running() }, 5*time.Second, 10*time.Millisecond)

	require.NoError(t, h.c.ReloadConfigString(h.config()).Err)
	require.Eventually(t, h.server.Running, 5*time.Second, 10*time.Millisecond)
	h.connect(t)
}

// A host key that no longer loads stops sshd, the same as a config that doesn't
func TestSSHDReload_BadHostKeyFileStops(t *testing.T) {
	h := newSSHDReloadHarness(t)
	h.connect(t)

	require.NoError(t, os.WriteFile(h.keyFile, []byte("not a key"), 0o600))
	assert.True(t, h.c.ReloadConfigString(h.config()).Failed)
	require.Eventually(t, func() bool { return !h.server.Running() }, 5*time.Second, 10*time.Millisecond)
}

// A trusted CA or authorized user that is skipped with a warning fails the reload, the valid ones are still used
func TestSSHDReload_SkippedEntriesFail(t *testing.T) {
	cases := map[string]string{
		"bad trusted ca":             "  trusted_cas: [nope]\n",
		"user entry not a map":       "    - nope\n",
		"user missing":               "    - keys: [nope]\n",
		"bad key":                    "    - user: other\n      keys: nope\n",
		"key in a list not a string": "    - user: other\n      keys: [5]\n",
		"bad key in a list":          "    - user: other\n      keys: [nope]\n",
		"keys missing":               "    - user: other\n",
	}
	for name, extra := range cases {
		t.Run(name, func(t *testing.T) {
			h := newSSHDReloadHarness(t)
			h.extra = extra
			res := h.c.ReloadConfigString(h.config())
			require.NoError(t, res.Err)
			assert.True(t, res.Failed)
			require.Eventually(t, h.server.Running, 5*time.Second, 10*time.Millisecond)
			h.connect(t)
		})
	}
}

// TestSSHExecExitStatusReflectsTheCommand: an exec request exits 1 when the command
// returned an error and 0 otherwise, with the command's output either way.
func TestSSHExecExitStatusReflectsTheCommand(t *testing.T) {
	h := newSSHDReloadHarness(t)
	h.reg.Register(&commands.Command{
		Name: "fails",
		Callback: func(fs any, a []string, w commands.StringWriter) error {
			_ = w.WriteLine("nope")
			return errors.New("failed")
		},
	})
	client, _ := h.connect(t)

	sess, err := client.NewSession()
	require.NoError(t, err)
	out, err := sess.CombinedOutput("fails")
	var exit *ssh.ExitError
	require.ErrorAs(t, err, &exit)
	assert.Equal(t, 1, exit.ExitStatus())
	assert.Equal(t, "nope\n", string(out))

	sess, err = client.NewSession()
	require.NoError(t, err)
	out, err = sess.CombinedOutput("help fails")
	require.NoError(t, err)
	assert.Equal(t, "fails - \n", string(out))

	// A word that is not a command is a failure too, with the command list as before.
	sess, err = client.NewSession()
	require.NoError(t, err)
	out, err = sess.CombinedOutput("nope")
	require.ErrorAs(t, err, &exit)
	assert.Equal(t, 1, exit.ExitStatus())
	assert.True(t, strings.HasPrefix(string(out), "did not understand: nope\nAvailable commands:\n"), string(out))
}