package kvm2

import (
	"bytes"
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"errors"
	"io"
	"net"
	"os"
	"path/filepath"
	"sync"
	"testing"

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

// testSSHServer is an in-process x/crypto/ssh server that accepts one
// authorised public key and forwards direct-streamlocal@openssh.com channels
// for socketPath to that unix socket, rejecting every other channel.
type testSSHServer struct {
	addr       string
	hostKey    ssh.PublicKey
	socketPath string

	listener net.Listener
	mu       sync.Mutex
	conns    []net.Conn
	wg       sync.WaitGroup
	stopOnce sync.Once

	reqMu     sync.Mutex
	requested []string
}

func startSSHServer(t *testing.T, authorized ssh.PublicKey, socketPath string) *testSSHServer {
	t.Helper()
	_, hostPriv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("generating host key: %v", err)
	}
	hostSigner, err := ssh.NewSignerFromKey(hostPriv)
	if err != nil {
		t.Fatalf("host signer: %v", err)
	}
	cfg := &ssh.ServerConfig{
		PublicKeyCallback: func(_ ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			if bytes.Equal(key.Marshal(), authorized.Marshal()) {
				return &ssh.Permissions{}, nil
			}
			return nil, errors.New("key not authorised")
		},
	}
	cfg.AddHostKey(hostSigner)
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	s := &testSSHServer{addr: ln.Addr().String(), hostKey: hostSigner.PublicKey(), socketPath: socketPath, listener: ln}
	s.wg.Add(1)
	go s.serve(cfg)
	t.Cleanup(s.stop)
	return s
}

func (s *testSSHServer) serve(cfg *ssh.ServerConfig) {
	defer s.wg.Done()
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			return
		}
		s.mu.Lock()
		s.conns = append(s.conns, conn)
		s.mu.Unlock()
		s.wg.Add(1)
		go s.handle(conn, cfg)
	}
}

func (s *testSSHServer) handle(conn net.Conn, cfg *ssh.ServerConfig) {
	defer s.wg.Done()
	sc, chans, reqs, err := ssh.NewServerConn(conn, cfg)
	if err != nil {
		_ = conn.Close()
		return
	}
	defer sc.Close()
	go ssh.DiscardRequests(reqs)
	for nc := range chans {
		if nc.ChannelType() != "direct-streamlocal@openssh.com" {
			_ = nc.Reject(ssh.UnknownChannelType, "only direct-streamlocal@openssh.com is served")
			continue
		}
		var open struct {
			SocketPath string
			Reserved0  string
			Reserved1  uint32
		}
		if err := ssh.Unmarshal(nc.ExtraData(), &open); err != nil {
			_ = nc.Reject(ssh.ConnectionFailed, "bad direct-streamlocal payload")
			continue
		}
		s.reqMu.Lock()
		s.requested = append(s.requested, open.SocketPath)
		s.reqMu.Unlock()
		if open.SocketPath != s.socketPath {
			_ = nc.Reject(ssh.ConnectionFailed, "no such socket")
			continue
		}
		target, err := net.Dial("unix", s.socketPath)
		if err != nil {
			_ = nc.Reject(ssh.ConnectionFailed, "socket unavailable")
			continue
		}
		ch, chReqs, err := nc.Accept()
		if err != nil {
			_ = target.Close()
			continue
		}
		go ssh.DiscardRequests(chReqs)
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			proxy(ch, target)
		}()
	}
}

func proxy(ch ssh.Channel, target net.Conn) {
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		_, _ = io.Copy(target, ch)
		_ = target.Close()
	}()
	go func() {
		defer wg.Done()
		_, _ = io.Copy(ch, target)
		_ = ch.Close()
	}()
	wg.Wait()
}

func (s *testSSHServer) stop() {
	s.stopOnce.Do(func() {
		_ = s.listener.Close()
		s.mu.Lock()
		for _, c := range s.conns {
			_ = c.Close()
		}
		s.mu.Unlock()
		s.wg.Wait()
	})
}

func (s *testSSHServer) requestedPaths() []string {
	s.reqMu.Lock()
	defer s.reqMu.Unlock()
	return append([]string(nil), s.requested...)
}

// startEchoSocket listens on a unix socket that echoes what it reads. The
// path is kept short: unix socket paths are limited to about 100 bytes and
// t.TempDir() names are long.
func startEchoSocket(t *testing.T) string {
	t.Helper()
	dir, err := os.MkdirTemp("", "kvm2")
	if err != nil {
		t.Fatalf("temp dir: %v", err)
	}
	t.Cleanup(func() { _ = os.RemoveAll(dir) })
	path := filepath.Join(dir, "libvirt-sock")
	ln, err := net.Listen("unix", path)
	if err != nil {
		t.Fatalf("listen unix: %v", err)
	}
	t.Cleanup(func() { _ = ln.Close() })
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go func() {
				_, _ = io.Copy(c, c)
				_ = c.Close()
			}()
		}
	}()
	return path
}

// silentListener accepts TCP connections and never speaks, for handshake
// timeouts. It returns the address.
func silentListener(t *testing.T) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	var conns []net.Conn
	var mu sync.Mutex
	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			mu.Lock()
			conns = append(conns, c)
			mu.Unlock()
		}
	}()
	t.Cleanup(func() {
		_ = ln.Close()
		mu.Lock()
		defer mu.Unlock()
		for _, c := range conns {
			_ = c.Close()
		}
	})
	return ln.Addr().String()
}

// closedPort returns an address nothing listens on.
func closedPort(t *testing.T) string {
	t.Helper()
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr := ln.Addr().String()
	_ = ln.Close()
	return addr
}

// writeKeyFile generates an ed25519 key, writes it in OpenSSH format and
// returns the path and the public half.
func writeKeyFile(t *testing.T, dir, name string) (string, ssh.PublicKey) {
	t.Helper()
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("generating key: %v", err)
	}
	block, err := ssh.MarshalPrivateKey(priv, "kvm2 test")
	if err != nil {
		t.Fatalf("marshalling key: %v", err)
	}
	path := writeFile(t, dir, name, string(pem.EncodeToMemory(block)))
	sshPub, err := ssh.NewPublicKey(pub)
	if err != nil {
		t.Fatalf("public key: %v", err)
	}
	return path, sshPub
}

// writeKnownHosts writes one known_hosts line for addr and key.
func writeKnownHosts(t *testing.T, dir, name, addr string, key ssh.PublicKey) string {
	t.Helper()
	return writeFile(t, dir, name, knownhosts.Line([]string{addr}, key)+"\n")
}

func writeFile(t *testing.T, dir, name, content string) string {
	t.Helper()
	path := filepath.Join(dir, name)
	if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
		t.Fatalf("writing %s: %v", path, err)
	}
	return path
}

// otherPublicKey returns a public key unrelated to any file, for mismatch
// cases.
func otherPublicKey(t *testing.T) ssh.PublicKey {
	t.Helper()
	pub, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("generating key: %v", err)
	}
	sshPub, err := ssh.NewPublicKey(pub)
	if err != nil {
		t.Fatalf("public key: %v", err)
	}
	return sshPub
}