wraith/internal/ssh/service_test.go

package ssh

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
	"testing"
	"time"

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

func TestNewSSHService(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)
	if svc == nil {
		t.Fatal("NewSSHService returned nil")
	}
	if len(svc.ListSessions()) != 0 {
		t.Error("new service should have no sessions")
	}
}

func TestBuildPasswordAuth(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)
	auth := svc.BuildPasswordAuth("mypassword")
	if auth == nil {
		t.Error("BuildPasswordAuth returned nil")
	}
}

func TestBuildKeyAuth(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)

	// Generate a test Ed25519 key
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("GenerateKey error: %v", err)
	}
	pemBlock, err := ssh.MarshalPrivateKey(priv, "")
	if err != nil {
		t.Fatalf("MarshalPrivateKey error: %v", err)
	}
	keyBytes := pem.EncodeToMemory(pemBlock)

	auth, err := svc.BuildKeyAuth(keyBytes, "")
	if err != nil {
		t.Fatalf("BuildKeyAuth error: %v", err)
	}
	if auth == nil {
		t.Error("BuildKeyAuth returned nil")
	}
}

func TestBuildKeyAuthInvalidKey(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)
	_, err := svc.BuildKeyAuth([]byte("not a key"), "")
	if err == nil {
		t.Error("BuildKeyAuth should fail with invalid key")
	}
}

func TestSessionTracking(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)

	// Manually add a session to test tracking
	svc.mu.Lock()
	svc.sessions["test-123"] = &SSHSession{
		ID:        "test-123",
		Hostname:  "192.168.1.4",
		Port:      22,
		Username:  "vstockwell",
		Connected: time.Now(),
	}
	svc.mu.Unlock()

	s, ok := svc.GetSession("test-123")
	if !ok {
		t.Fatal("session not found")
	}
	if s.Hostname != "192.168.1.4" {
		t.Errorf("Hostname = %q, want %q", s.Hostname, "192.168.1.4")
	}

	sessions := svc.ListSessions()
	if len(sessions) != 1 {
		t.Errorf("ListSessions() = %d, want 1", len(sessions))
	}
}

func TestGetSessionNotFound(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)
	_, ok := svc.GetSession("nonexistent")
	if ok {
		t.Error("GetSession should return false for nonexistent session")
	}
}

func TestWriteNotFound(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)
	err := svc.Write("nonexistent", "data")
	if err == nil {
		t.Error("Write should fail for nonexistent session")
	}
}

func TestResizeNotFound(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)
	err := svc.Resize("nonexistent", 80, 24)
	if err == nil {
		t.Error("Resize should fail for nonexistent session")
	}
}

func TestDisconnectNotFound(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)
	err := svc.Disconnect("nonexistent")
	if err == nil {
		t.Error("Disconnect should fail for nonexistent session")
	}
}

func TestDisconnectRemovesSession(t *testing.T) {
	svc := NewSSHService(nil, nil, nil, nil)

	// Manually add a session with nil Client/Session/Stdin (no real connection)
	svc.mu.Lock()
	svc.sessions["test-dc"] = &SSHSession{
		ID:        "test-dc",
		Hostname:  "10.0.0.1",
		Port:      22,
		Username:  "admin",
		Connected: time.Now(),
	}
	svc.mu.Unlock()

	if err := svc.Disconnect("test-dc"); err != nil {
		t.Fatalf("Disconnect error: %v", err)
	}

	_, ok := svc.GetSession("test-dc")
	if ok {
		t.Error("session should be removed after Disconnect")
	}

	if len(svc.ListSessions()) != 0 {
		t.Error("ListSessions should be empty after Disconnect")
	}
}