banger/internal/guest/known_hosts_test.go

package guest

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

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

// makeTestHostKey generates a fresh ed25519 key and returns the
// ssh.PublicKey the server would present during a handshake.
func makeTestHostKey(t *testing.T) ssh.PublicKey {
	t.Helper()
	pub, _, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKey: %v", err)
	}
	sshPub, err := ssh.NewPublicKey(pub)
	if err != nil {
		t.Fatalf("NewPublicKey: %v", err)
	}
	return sshPub
}

func TestTOFUHostKeyCallbackPinsOnFirstUse(t *testing.T) {
	t.Parallel()
	path := filepath.Join(t.TempDir(), "known_hosts")
	cb := TOFUHostKeyCallback(path)

	key := makeTestHostKey(t)
	addr := &net.TCPAddr{IP: net.ParseIP("172.16.0.5"), Port: 22}

	if err := cb("172.16.0.5:22", addr, key); err != nil {
		t.Fatalf("first-use callback: %v", err)
	}

	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("ReadFile: %v", err)
	}
	content := string(data)
	if !strings.Contains(content, "172.16.0.5") {
		t.Errorf("known_hosts missing host:\n%s", content)
	}
	if !strings.Contains(content, key.Type()) {
		t.Errorf("known_hosts missing key type:\n%s", content)
	}
}

func TestTOFUHostKeyCallbackAcceptsMatch(t *testing.T) {
	t.Parallel()
	path := filepath.Join(t.TempDir(), "known_hosts")
	cb := TOFUHostKeyCallback(path)
	key := makeTestHostKey(t)
	addr := &net.TCPAddr{IP: net.ParseIP("172.16.0.6"), Port: 22}

	if err := cb("172.16.0.6:22", addr, key); err != nil {
		t.Fatalf("first-use: %v", err)
	}
	// Same key, second dial: must succeed.
	if err := cb("172.16.0.6:22", addr, key); err != nil {
		t.Fatalf("second dial with matching key: %v", err)
	}
}

func TestTOFUHostKeyCallbackRejectsMismatch(t *testing.T) {
	t.Parallel()
	path := filepath.Join(t.TempDir(), "known_hosts")
	cb := TOFUHostKeyCallback(path)
	addr := &net.TCPAddr{IP: net.ParseIP("172.16.0.7"), Port: 22}

	original := makeTestHostKey(t)
	if err := cb("172.16.0.7:22", addr, original); err != nil {
		t.Fatalf("pin original: %v", err)
	}

	impostor := makeTestHostKey(t)
	err := cb("172.16.0.7:22", addr, impostor)
	if err == nil {
		t.Fatal("expected mismatch error, got nil")
	}
	if !strings.Contains(err.Error(), "does not match") {
		t.Errorf("error = %v, want message about mismatch", err)
	}
}

func TestTOFUEmptyPathDisablesVerification(t *testing.T) {
	t.Parallel()
	// Empty path returns an Insecure callback — useful for tests /
	// throwaway tools. Document behaviour so the fallback doesn't
	// silently regress to "always verify but without a file".
	cb := TOFUHostKeyCallback("")
	addr := &net.TCPAddr{IP: net.ParseIP("127.0.0.1"), Port: 22}
	if err := cb("127.0.0.1:22", addr, makeTestHostKey(t)); err != nil {
		t.Fatalf("empty-path callback should accept: %v", err)
	}
}

func TestRemoveKnownHostsDropsEntry(t *testing.T) {
	t.Parallel()
	path := filepath.Join(t.TempDir(), "known_hosts")
	cb := TOFUHostKeyCallback(path)
	keep := makeTestHostKey(t)
	drop := makeTestHostKey(t)

	if err := cb("172.16.0.10:22", &net.TCPAddr{IP: net.ParseIP("172.16.0.10"), Port: 22}, keep); err != nil {
		t.Fatalf("pin keep: %v", err)
	}
	if err := cb("172.16.0.11:22", &net.TCPAddr{IP: net.ParseIP("172.16.0.11"), Port: 22}, drop); err != nil {
		t.Fatalf("pin drop: %v", err)
	}

	if err := RemoveKnownHosts(path, "172.16.0.11"); err != nil {
		t.Fatalf("RemoveKnownHosts: %v", err)
	}

	data, _ := os.ReadFile(path)
	content := string(data)
	if !strings.Contains(content, "172.16.0.10") {
		t.Errorf("kept entry missing:\n%s", content)
	}
	if strings.Contains(content, "172.16.0.11") {
		t.Errorf("dropped entry still present:\n%s", content)
	}
}

func TestRemoveKnownHostsMissingFileIsNoOp(t *testing.T) {
	t.Parallel()
	missing := filepath.Join(t.TempDir(), "absent")
	if err := RemoveKnownHosts(missing, "any"); err != nil {
		t.Fatalf("RemoveKnownHosts on missing: %v", err)
	}
}

func TestRemoveKnownHostsEmptyPathIsNoOp(t *testing.T) {
	t.Parallel()
	if err := RemoveKnownHosts("", "any"); err != nil {
		t.Fatalf("RemoveKnownHosts(empty): %v", err)
	}
}

// TestTOFURewritesAllowsReuseAfterRemove: after a VM is deleted and
// its pin is cleared, a future VM reusing the same IP (with a fresh
// host key) should re-pin cleanly, not fail the mismatch branch.
func TestTOFURewritesAllowsReuseAfterRemove(t *testing.T) {
	t.Parallel()
	path := filepath.Join(t.TempDir(), "known_hosts")
	cb := TOFUHostKeyCallback(path)
	addr := &net.TCPAddr{IP: net.ParseIP("172.16.0.15"), Port: 22}

	original := makeTestHostKey(t)
	if err := cb("172.16.0.15:22", addr, original); err != nil {
		t.Fatalf("pin original: %v", err)
	}

	// VM deleted → pin removed.
	if err := RemoveKnownHosts(path, "172.16.0.15"); err != nil {
		t.Fatalf("RemoveKnownHosts: %v", err)
	}

	// New VM, same IP, new host key. Must re-pin without error.
	replacement := makeTestHostKey(t)
	if err := cb("172.16.0.15:22", addr, replacement); err != nil {
		t.Fatalf("re-pin after remove: %v", err)
	}
}

func TestHostLookupKeyStripsPort(t *testing.T) {
	t.Parallel()
	if got := hostLookupKey("10.0.0.1:22", nil); got != "10.0.0.1" {
		t.Errorf("got %q, want 10.0.0.1", got)
	}
	if got := hostLookupKey("host.vm", nil); got != "host.vm" {
		t.Errorf("got %q, want host.vm", got)
	}
	addr := &net.TCPAddr{IP: net.ParseIP("1.2.3.4"), Port: 22}
	if got := hostLookupKey("", addr); got != "1.2.3.4" {
		t.Errorf("fallback: got %q, want 1.2.3.4", got)
	}
}