package e2e

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

	"git.veillax.com/Veillax/bamboo/sshd"
	"git.veillax.com/Veillax/blacktip"
	"git.veillax.com/Veillax/reef/model"
	"golang.org/x/crypto/ssh"
)

func newSSHSigner(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
}

// A TCP port on the edge whose target is Bamboo's SSH server at the site reaches that server
// with an ordinary ssh client: the path is client, edge, WireGuard tunnel, Bamboo, sshd.
func TestSSHThroughTheEdgeToBamboosServer(t *testing.T) {
	w := setup(t, blacktip.TraefikConfig{})
	dir := t.TempDir()
	signer := newSSHSigner(t)
	keys := filepath.Join(dir, "authorized_keys")
	if err := os.WriteFile(keys, ssh.MarshalAuthorizedKey(signer.PublicKey()), 0o600); err != nil {
		t.Fatal(err)
	}
	srv, err := sshd.New(sshd.Config{AuthorizedKeys: keys, HostKey: filepath.Join(dir, "ssh_host_ed25519_key")})
	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())
	t.Cleanup(cancel)
	go srv.Serve(ctx, ln)

	edgePort := freePort(t)
	w.open(model.Mapping{Proto: model.ProtoTCP, Port: edgePort, Name: "ssh", Targets: tgt(w.site, ln.Addr().(*net.TCPAddr).Port)})
	addr := "127.0.0.1:" + strconv.Itoa(edgePort)

	connect := func(auth ssh.AuthMethod) (*ssh.Client, error) {
		return ssh.Dial("tcp", addr, &ssh.ClientConfig{User: "alex", Auth: []ssh.AuthMethod{auth}, Timeout: 5 * time.Second,
			// the key the client sees through the tunnel must be Bamboo's, not something in between
			HostKeyCallback: func(_ string, _ net.Addr, k ssh.PublicKey) error {
				if got := ssh.FingerprintSHA256(k); got != srv.Fingerprint() {
					return fmt.Errorf("host key %s is not the Bamboo server's %s", got, srv.Fingerprint())
				}
				return nil
			}})
	}
	var client *ssh.Client
	var last error
	for i := 0; i < 60 && client == nil; i++ {
		if client, last = connect(ssh.PublicKeys(signer)); client == nil {
			time.Sleep(250 * time.Millisecond)
		}
	}
	if client == nil {
		t.Fatalf("could not sign in through the edge: %v", last)
	}
	defer client.Close()

	if sshd.Supported {
		sess, err := client.NewSession()
		if err != nil {
			t.Fatal(err)
		}
		out, err := sess.Output("echo through-the-tunnel")
		sess.Close()
		if err != nil || strings.TrimSpace(string(out)) != "through-the-tunnel" {
			t.Fatalf("running a command: %q, %v", out, err)
		}
	}

	// the same narrow rules hold through the tunnel
	if c, err := connect(ssh.PublicKeys(newSSHSigner(t))); err == nil {
		c.Close()
		t.Error("a key that is not in authorized_keys signed in")
	}
	if c, err := connect(ssh.Password("hunter2")); err == nil {
		c.Close()
		t.Error("a password signed in")
	}
	if c, err := client.Dial("tcp", "127.0.0.1:1"); err == nil {
		c.Close()
		t.Error("port forwarding through the SSH server is allowed")
	}
}