HostKeeper/v1/test/knownhosts/knownhosts_test.go

package knownhosts_test

import (
	"crypto/ed25519"
	"crypto/rand"
	"os"
	"path/filepath"
	"testing"

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

	"git.tukangketik.id/swanadiva/hostkeeper/pkg/knownhosts"
)

func tempDir(t *testing.T) string {
	t.Helper()
	dir, err := os.MkdirTemp("", "knownhosts-test-*")
	if err != nil {
		t.Fatalf("Failed to create temp dir: %v", err)
	}
	t.Cleanup(func() { os.RemoveAll(dir) })
	return dir
}

func generateTestKey(t *testing.T) cryptossh.PublicKey {
	t.Helper()
	pubKey, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	sshPubKey, err := cryptossh.NewPublicKey(pubKey)
	if err != nil {
		t.Fatalf("Failed to create SSH public key: %v", err)
	}
	return sshPubKey
}

func generateTestKey2(t *testing.T) cryptossh.PublicKey {
	t.Helper()
	pubKey, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}
	sshPubKey, err := cryptossh.NewPublicKey(pubKey)
	if err != nil {
		t.Fatalf("Failed to create SSH public key: %v", err)
	}
	return sshPubKey
}

// 2.1 New creates file on first call
func TestNewCreatesFile(t *testing.T) {
	dir := tempDir(t)
	kh, err := knownhosts.New(dir)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	if kh == nil {
		t.Fatal("New should return non-nil KnownHosts")
	}
	// File should be created after first Add+Save
}

// 2.2 New loads existing
func TestNewLoadsExisting(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	// Create and add a host
	kh1, _ := knownhosts.New(dir)
	_ = kh1.Add("example.com", 22, key)

	// Load again
	kh2, err := knownhosts.New(dir)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	stored := kh2.Get("example.com", 22)
	if stored == nil {
		t.Fatal("Should load existing host from file")
	}
}

// 2.3 Add new host
func TestAddNewHost(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh, _ := knownhosts.New(dir)
	err := kh.Add("example.com", 22, key)
	if err != nil {
		t.Fatalf("Add failed: %v", err)
	}

	stored := kh.Get("example.com", 22)
	if stored == nil {
		t.Fatal("Get should return the added host")
	}
	if stored.Hostname != "example.com" {
		t.Errorf("Hostname = %q, want %q", stored.Hostname, "example.com")
	}
	if stored.Port != 22 {
		t.Errorf("Port = %d, want 22", stored.Port)
	}
}

// 2.4 Add duplicate — no error, no duplicate
func TestAddDuplicate(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh, _ := knownhosts.New(dir)
	_ = kh.Add("example.com", 22, key)
	err := kh.Add("example.com", 22, key)
	if err != nil {
		t.Fatalf("Add duplicate should not error: %v", err)
	}
}

// 2.5 Get existing host
func TestGetExisting(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh, _ := knownhosts.New(dir)
	_ = kh.Add("example.com", 22, key)

	stored := kh.Get("example.com", 22)
	if stored == nil {
		t.Fatal("Get should return existing host")
	}
	if stored.Hostname != "example.com" {
		t.Errorf("Hostname = %q, want %q", stored.Hostname, "example.com")
	}
}

// 2.6 Get non-existent host
func TestGetNonExistent(t *testing.T) {
	dir := tempDir(t)
	kh, _ := knownhosts.New(dir)

	stored := kh.Get("unknown.com", 22)
	if stored != nil {
		t.Error("Get should return nil for non-existent host")
	}
}

// 2.7 Remove existing host
func TestRemoveExisting(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh, _ := knownhosts.New(dir)
	_ = kh.Add("example.com", 22, key)

	err := kh.Remove("example.com", 22)
	if err != nil {
		t.Fatalf("Remove failed: %v", err)
	}

	stored := kh.Get("example.com", 22)
	if stored != nil {
		t.Error("Get should return nil after Remove")
	}
}

// 2.8 Remove non-existent host — no error
func TestRemoveNonExistent(t *testing.T) {
	dir := tempDir(t)
	kh, _ := knownhosts.New(dir)

	err := kh.Remove("unknown.com", 22)
	if err != nil {
		t.Fatalf("Remove non-existent should not error: %v", err)
	}
}

// 2.9 Verify unknown host — (false, nil) TOFU
func TestVerifyUnknown(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh, _ := knownhosts.New(dir)
	matches, stored := kh.Verify("unknown.com", 22, key)
	if matches {
		t.Error("Verify should return false for unknown host")
	}
	if stored != nil {
		t.Error("Verify should return nil HostKey for unknown host")
	}
}

// 2.10 Verify known host, matching key — (true, hostKey)
func TestVerifyKnownMatch(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh, _ := knownhosts.New(dir)
	_ = kh.Add("example.com", 22, key)

	matches, stored := kh.Verify("example.com", 22, key)
	if !matches {
		t.Error("Verify should return true for matching key")
	}
	if stored == nil {
		t.Error("Verify should return stored HostKey")
	}
}

// 2.11 Verify known host, mismatched key — (false, hostKey) MITM
func TestVerifyKnownMismatch(t *testing.T) {
	dir := tempDir(t)
	key1 := generateTestKey(t)
	key2 := generateTestKey2(t)

	kh, _ := knownhosts.New(dir)
	_ = kh.Add("example.com", 22, key1)

	matches, stored := kh.Verify("example.com", 22, key2)
	if matches {
		t.Error("Verify should return false for mismatched key (MITM)")
	}
	if stored == nil {
		t.Error("Verify should return stored HostKey for mismatch")
	}
}

// 2.12 HostKeyCallback — autoAdd=true adds unknown hosts
func TestHostKeyCallbackAutoAdd(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh, _ := knownhosts.New(dir)
	callback := kh.HostKeyCallback(true)

	// Simulate host key check via callback
	// HostKeyCallback expects net.Addr, so we create a fake one
	addr := &fakeAddr{addr: "192.168.1.1:22"}
	err := callback("example.com", addr, key)
	if err != nil {
		t.Fatalf("HostKeyCallback with autoAdd failed: %v", err)
	}

	// Verify host was added
	stored := kh.Get("example.com", 22)
	if stored == nil {
		t.Error("HostKeyCallback should add unknown host when autoAdd=true")
	}
}

// 2.13 Persistence — Add → Save → New → Get
func TestPersistence(t *testing.T) {
	dir := tempDir(t)
	key := generateTestKey(t)

	kh1, _ := knownhosts.New(dir)
	_ = kh1.Add("example.com", 22, key)

	// Create new instance from same directory
	kh2, err := knownhosts.New(dir)
	if err != nil {
		t.Fatalf("New failed: %v", err)
	}
	stored := kh2.Get("example.com", 22)
	if stored == nil {
		t.Fatal("Host should persist across New() calls")
	}
}

// 2.14 Corrupted file → error
func TestCorruptedFile(t *testing.T) {
	dir := tempDir(t)
	path := filepath.Join(dir, "known_hosts")
	_ = os.WriteFile(path, []byte("not valid json {{{"), 0600)

	_, err := knownhosts.New(dir)
	if err == nil {
		t.Error("New should return error for corrupted file")
	}
}

// fakeAddr implements net.Addr for testing
type fakeAddr struct {
	addr string
}

func (f *fakeAddr) Network() string { return "tcp" }
func (f *fakeAddr) String() string  { return f.addr }