main_test.go

package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/hex"
	"net"
	"os"
	"path/filepath"
	"testing"

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

// testBackend holds a test agent backend with its socket path
type testBackend struct {
	agent    agent.Agent
	listener net.Listener
	sockPath string
	tmpDir   string // directory to clean up (only if we created it)
}

// createShortTempDir creates a temporary directory with a short path in /tmp
// to avoid Unix socket path length limits (108 chars on Linux).
// t.TempDir() paths can be too long in nested nix shells.
func createShortTempDir(t *testing.T) string {
	t.Helper()

	randBytes := make([]byte, 6)
	if _, err := rand.Read(randBytes); err != nil {
		t.Fatalf("failed to generate random bytes: %v", err)
	}

	dir := filepath.Join("/tmp", "ua-"+hex.EncodeToString(randBytes))
	if err := os.MkdirAll(dir, 0700); err != nil {
		t.Fatalf("failed to create temp dir: %v", err)
	}

	return dir
}

// newTestBackend creates a new in-memory agent backend served over a Unix socket
func newTestBackend(t *testing.T, tmpDir string, name string) *testBackend {
	t.Helper()

	// Use short socket names to stay under Unix socket path limit
	sockPath := filepath.Join(tmpDir, name+".s")
	keyring := agent.NewKeyring()

	listener, err := net.Listen("unix", sockPath)
	if err != nil {
		t.Fatalf("failed to create listener for %s: %v", name, err)
	}

	// Serve agent in background
	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				return // listener closed
			}
			go agent.ServeAgent(keyring, conn)
		}
	}()

	return &testBackend{
		agent:    keyring,
		listener: listener,
		sockPath: sockPath,
	}
}

func (tb *testBackend) close() {
	tb.listener.Close()
}

// testEnv manages a test environment with short temp directory
type testEnv struct {
	dir string
	t   *testing.T
}

func newTestEnv(t *testing.T) *testEnv {
	t.Helper()
	return &testEnv{
		dir: createShortTempDir(t),
		t:   t,
	}
}

func (e *testEnv) cleanup() {
	os.RemoveAll(e.dir)
}

// generateTestKey creates a new ed25519 key pair for testing
func generateTestKey(t *testing.T) (ed25519.PrivateKey, ssh.PublicKey) {
	t.Helper()

	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("failed to generate test key: %v", err)
	}

	sshPub, err := ssh.NewPublicKey(pub)
	if err != nil {
		t.Fatalf("failed to convert public key: %v", err)
	}

	return priv, sshPub
}

// addKeyToBackend adds a key to a backend agent
func addKeyToBackend(t *testing.T, backend *testBackend, priv ed25519.PrivateKey, comment string) {
	t.Helper()

	err := backend.agent.Add(agent.AddedKey{
		PrivateKey: priv,
		Comment:    comment,
	})
	if err != nil {
		t.Fatalf("failed to add key to backend: %v", err)
	}
}

func TestList_SingleBackend(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend := newTestBackend(t, env.dir, "a1")
	defer backend.close()

	priv1, _ := generateTestKey(t)
	priv2, _ := generateTestKey(t)
	addKeyToBackend(t, backend, priv1, "key1")
	addKeyToBackend(t, backend, priv2, "key2")

	union := NewUnionAgent([]string{backend.sockPath})

	keys, err := union.List()
	if err != nil {
		t.Fatalf("List() error: %v", err)
	}

	if len(keys) != 2 {
		t.Errorf("expected 2 keys, got %d", len(keys))
	}
}

func TestList_MultipleBackends(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	priv1, _ := generateTestKey(t)
	priv2, _ := generateTestKey(t)
	priv3, _ := generateTestKey(t)

	addKeyToBackend(t, backend1, priv1, "key1-backend1")
	addKeyToBackend(t, backend1, priv2, "key2-backend1")
	addKeyToBackend(t, backend2, priv3, "key3-backend2")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	keys, err := union.List()
	if err != nil {
		t.Fatalf("List() error: %v", err)
	}

	if len(keys) != 3 {
		t.Errorf("expected 3 keys, got %d", len(keys))
	}
}

func TestList_Deduplication(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	// Add the same key to both backends
	priv, _ := generateTestKey(t)
	addKeyToBackend(t, backend1, priv, "shared-key-backend1")
	addKeyToBackend(t, backend2, priv, "shared-key-backend2")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	keys, err := union.List()
	if err != nil {
		t.Fatalf("List() error: %v", err)
	}

	if len(keys) != 1 {
		t.Errorf("expected 1 key (deduplicated), got %d", len(keys))
	}
}

func TestList_UnavailableBackend(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend := newTestBackend(t, env.dir, "a1")
	defer backend.close()

	priv, _ := generateTestKey(t)
	addKeyToBackend(t, backend, priv, "key1")

	// Include a non-existent socket path
	nonExistent := filepath.Join(env.dir, "bad.s")
	union := NewUnionAgent([]string{nonExistent, backend.sockPath})

	keys, err := union.List()
	if err != nil {
		t.Fatalf("List() error: %v", err)
	}

	// Should still return keys from the available backend
	if len(keys) != 1 {
		t.Errorf("expected 1 key from available backend, got %d", len(keys))
	}
}

func TestSign_SingleBackend(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend := newTestBackend(t, env.dir, "a1")
	defer backend.close()

	priv, pub := generateTestKey(t)
	addKeyToBackend(t, backend, priv, "key1")

	union := NewUnionAgent([]string{backend.sockPath})

	data := []byte("test data to sign")
	sig, err := union.Sign(pub, data)
	if err != nil {
		t.Fatalf("Sign() error: %v", err)
	}

	// Verify signature
	if err := pub.Verify(data, sig); err != nil {
		t.Errorf("signature verification failed: %v", err)
	}
}

func TestSign_MultipleBackends_KeyInFirst(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	priv, pub := generateTestKey(t)
	addKeyToBackend(t, backend1, priv, "key1")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	data := []byte("test data to sign")
	sig, err := union.Sign(pub, data)
	if err != nil {
		t.Fatalf("Sign() error: %v", err)
	}

	if err := pub.Verify(data, sig); err != nil {
		t.Errorf("signature verification failed: %v", err)
	}
}

func TestSign_MultipleBackends_KeyInSecond(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	priv, pub := generateTestKey(t)
	addKeyToBackend(t, backend2, priv, "key1") // Key only in second backend

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	data := []byte("test data to sign")
	sig, err := union.Sign(pub, data)
	if err != nil {
		t.Fatalf("Sign() error: %v", err)
	}

	if err := pub.Verify(data, sig); err != nil {
		t.Errorf("signature verification failed: %v", err)
	}
}

func TestSign_KeyNotFound(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend := newTestBackend(t, env.dir, "a1")
	defer backend.close()

	// Generate a key but don't add it to the backend
	_, pub := generateTestKey(t)

	union := NewUnionAgent([]string{backend.sockPath})

	data := []byte("test data to sign")
	_, err := union.Sign(pub, data)
	if err == nil {
		t.Error("expected error when signing with unknown key, got nil")
	}
}

func TestSignWithFlags(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend := newTestBackend(t, env.dir, "a1")
	defer backend.close()

	priv, pub := generateTestKey(t)
	addKeyToBackend(t, backend, priv, "key1")

	union := NewUnionAgent([]string{backend.sockPath})

	data := []byte("test data to sign")
	sig, err := union.SignWithFlags(pub, data, 0)
	if err != nil {
		t.Fatalf("SignWithFlags() error: %v", err)
	}

	if err := pub.Verify(data, sig); err != nil {
		t.Errorf("signature verification failed: %v", err)
	}
}

func TestAdd_ForwardsToFirstBackend(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	priv, _ := generateTestKey(t)
	err := union.Add(agent.AddedKey{
		PrivateKey: priv,
		Comment:    "added-via-union",
	})
	if err != nil {
		t.Fatalf("Add() error: %v", err)
	}

	// Key should be in first backend
	keys1, _ := backend1.agent.List()
	if len(keys1) != 1 {
		t.Errorf("expected 1 key in first backend, got %d", len(keys1))
	}

	// Key should NOT be in second backend
	keys2, _ := backend2.agent.List()
	if len(keys2) != 0 {
		t.Errorf("expected 0 keys in second backend, got %d", len(keys2))
	}
}

func TestAdd_NoBackends(t *testing.T) {
	union := NewUnionAgent([]string{})

	priv, _ := generateTestKey(t)
	err := union.Add(agent.AddedKey{
		PrivateKey: priv,
		Comment:    "test",
	})
	if err == nil {
		t.Error("expected error when adding to union with no backends, got nil")
	}
}

func TestRemove_FromAllBackends(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	// Add the same key to both backends
	priv, pub := generateTestKey(t)
	addKeyToBackend(t, backend1, priv, "shared-key")
	addKeyToBackend(t, backend2, priv, "shared-key")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	err := union.Remove(pub)
	if err != nil {
		t.Fatalf("Remove() error: %v", err)
	}

	// Key should be removed from both backends
	keys1, _ := backend1.agent.List()
	if len(keys1) != 0 {
		t.Errorf("expected 0 keys in first backend after remove, got %d", len(keys1))
	}

	keys2, _ := backend2.agent.List()
	if len(keys2) != 0 {
		t.Errorf("expected 0 keys in second backend after remove, got %d", len(keys2))
	}
}

func TestRemoveAll_FromAllBackends(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	priv1, _ := generateTestKey(t)
	priv2, _ := generateTestKey(t)
	priv3, _ := generateTestKey(t)

	addKeyToBackend(t, backend1, priv1, "key1")
	addKeyToBackend(t, backend1, priv2, "key2")
	addKeyToBackend(t, backend2, priv3, "key3")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	err := union.RemoveAll()
	if err != nil {
		t.Fatalf("RemoveAll() error: %v", err)
	}

	// All keys should be removed from both backends
	keys1, _ := backend1.agent.List()
	if len(keys1) != 0 {
		t.Errorf("expected 0 keys in first backend after RemoveAll, got %d", len(keys1))
	}

	keys2, _ := backend2.agent.List()
	if len(keys2) != 0 {
		t.Errorf("expected 0 keys in second backend after RemoveAll, got %d", len(keys2))
	}
}

func TestLock_AllBackends(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	priv1, _ := generateTestKey(t)
	priv2, _ := generateTestKey(t)
	addKeyToBackend(t, backend1, priv1, "key1")
	addKeyToBackend(t, backend2, priv2, "key2")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	passphrase := []byte("test-passphrase")
	err := union.Lock(passphrase)
	if err != nil {
		t.Fatalf("Lock() error: %v", err)
	}

	// Both backends should be locked (return 0 keys when locked)
	keys1, _ := backend1.agent.List()
	if len(keys1) != 0 {
		t.Errorf("expected 0 keys from locked first backend, got %d", len(keys1))
	}

	keys2, _ := backend2.agent.List()
	if len(keys2) != 0 {
		t.Errorf("expected 0 keys from locked second backend, got %d", len(keys2))
	}
}

func TestUnlock_AllBackends(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	priv1, _ := generateTestKey(t)
	priv2, _ := generateTestKey(t)
	addKeyToBackend(t, backend1, priv1, "key1")
	addKeyToBackend(t, backend2, priv2, "key2")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	passphrase := []byte("test-passphrase")

	// Lock first
	err := union.Lock(passphrase)
	if err != nil {
		t.Fatalf("Lock() error: %v", err)
	}

	// Then unlock
	err = union.Unlock(passphrase)
	if err != nil {
		t.Fatalf("Unlock() error: %v", err)
	}

	// Both backends should be unlocked and show their keys
	keys1, _ := backend1.agent.List()
	if len(keys1) != 1 {
		t.Errorf("expected 1 key from unlocked first backend, got %d", len(keys1))
	}

	keys2, _ := backend2.agent.List()
	if len(keys2) != 1 {
		t.Errorf("expected 1 key from unlocked second backend, got %d", len(keys2))
	}
}

func TestUnlock_WrongPassphrase(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend := newTestBackend(t, env.dir, "a1")
	defer backend.close()

	priv, _ := generateTestKey(t)
	addKeyToBackend(t, backend, priv, "key1")

	union := NewUnionAgent([]string{backend.sockPath})

	// Lock with one passphrase
	err := union.Lock([]byte("correct-passphrase"))
	if err != nil {
		t.Fatalf("Lock() error: %v", err)
	}

	// Try to unlock with wrong passphrase
	err = union.Unlock([]byte("wrong-passphrase"))
	if err == nil {
		t.Error("expected error when unlocking with wrong passphrase, got nil")
	}

	// Backend should still be locked
	keys, _ := backend.agent.List()
	if len(keys) != 0 {
		t.Errorf("expected 0 keys from still-locked backend, got %d", len(keys))
	}
}

func TestSigners_SingleBackend(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend := newTestBackend(t, env.dir, "a1")
	defer backend.close()

	priv1, _ := generateTestKey(t)
	priv2, _ := generateTestKey(t)
	addKeyToBackend(t, backend, priv1, "key1")
	addKeyToBackend(t, backend, priv2, "key2")

	union := NewUnionAgent([]string{backend.sockPath})

	signers, err := union.Signers()
	if err != nil {
		t.Fatalf("Signers() error: %v", err)
	}

	if len(signers) != 2 {
		t.Errorf("expected 2 signers, got %d", len(signers))
	}
}

func TestSigners_MultipleBackends_Deduplication(t *testing.T) {
	env := newTestEnv(t)
	defer env.cleanup()

	backend1 := newTestBackend(t, env.dir, "a1")
	defer backend1.close()
	backend2 := newTestBackend(t, env.dir, "a2")
	defer backend2.close()

	// Add the same key to both backends
	priv, _ := generateTestKey(t)
	addKeyToBackend(t, backend1, priv, "shared-key")
	addKeyToBackend(t, backend2, priv, "shared-key")

	union := NewUnionAgent([]string{backend1.sockPath, backend2.sockPath})

	signers, err := union.Signers()
	if err != nil {
		t.Fatalf("Signers() error: %v", err)
	}

	if len(signers) != 1 {
		t.Errorf("expected 1 signer (deduplicated), got %d", len(signers))
	}
}

func TestGenerateSocketPath(t *testing.T) {
	// Save and restore TMPDIR
	origTmpdir := os.Getenv("TMPDIR")
	defer os.Setenv("TMPDIR", origTmpdir)

	tmpDir := t.TempDir()
	os.Setenv("TMPDIR", tmpDir)

	path, err := generateSocketPath()
	if err != nil {
		t.Fatalf("generateSocketPath() error: %v", err)
	}

	// Check that the directory was created
	dir := filepath.Dir(path)
	info, err := os.Stat(dir)
	if err != nil {
		t.Fatalf("socket directory not created: %v", err)
	}

	if !info.IsDir() {
		t.Error("expected socket path parent to be a directory")
	}

	// Check permissions (should be 0700)
	if info.Mode().Perm() != 0700 {
		t.Errorf("expected directory permissions 0700, got %o", info.Mode().Perm())
	}

	// Check path format
	if filepath.Base(dir)[:4] != "ssh-" {
		t.Errorf("expected directory name to start with 'ssh-', got %s", filepath.Base(dir))
	}
}