blob: d3deb1e170ab844d50d8e091c85b2e3eea202caa [file] [log] [blame]
Sean McCullough2cba6952025-04-25 20:32:10 +00001package dockerimg
2
3import (
Sean McCullough15c95282025-05-08 16:48:38 -07004 "bufio"
Sean McCullough2cba6952025-04-25 20:32:10 +00005 "bytes"
Sean McCullough3e9d80c2025-05-13 23:35:23 +00006 "crypto/ed25519"
Sean McCullough2cba6952025-04-25 20:32:10 +00007 "crypto/rand"
Sean McCullough2cba6952025-04-25 20:32:10 +00008 "fmt"
9 "io/fs"
10 "os"
11 "path/filepath"
12 "strings"
13 "testing"
14
15 "golang.org/x/crypto/ssh"
16)
17
18// MockFileSystem implements the FileSystem interface for testing
19type MockFileSystem struct {
20 Files map[string][]byte
21 CreatedDirs map[string]bool
22 OpenedFiles map[string]*MockFile
23 StatCalledWith []string
Sean McCullough0d95d3a2025-04-30 16:22:28 +000024 TempFiles []string
Sean McCullough2cba6952025-04-25 20:32:10 +000025 FailOn map[string]error // Map of function name to error to simulate failures
26}
27
28func NewMockFileSystem() *MockFileSystem {
29 return &MockFileSystem{
30 Files: make(map[string][]byte),
31 CreatedDirs: make(map[string]bool),
32 OpenedFiles: make(map[string]*MockFile),
Sean McCullough0d95d3a2025-04-30 16:22:28 +000033 TempFiles: []string{},
Sean McCullough2cba6952025-04-25 20:32:10 +000034 FailOn: make(map[string]error),
35 }
36}
37
38func (m *MockFileSystem) Stat(name string) (fs.FileInfo, error) {
39 m.StatCalledWith = append(m.StatCalledWith, name)
40 if err, ok := m.FailOn["Stat"]; ok {
41 return nil, err
42 }
43
44 _, exists := m.Files[name]
45 if exists {
46 return nil, nil // File exists
47 }
48 _, exists = m.CreatedDirs[name]
49 if exists {
50 return nil, nil // Directory exists
51 }
52 return nil, os.ErrNotExist
53}
54
55func (m *MockFileSystem) Mkdir(name string, perm fs.FileMode) error {
56 if err, ok := m.FailOn["Mkdir"]; ok {
57 return err
58 }
59 m.CreatedDirs[name] = true
60 return nil
61}
62
Sean McCulloughc796e7f2025-04-30 08:44:06 -070063func (m *MockFileSystem) MkdirAll(name string, perm fs.FileMode) error {
64 if err, ok := m.FailOn["MkdirAll"]; ok {
65 return err
66 }
67 m.CreatedDirs[name] = true
68 return nil
69}
70
Sean McCullough2cba6952025-04-25 20:32:10 +000071func (m *MockFileSystem) ReadFile(name string) ([]byte, error) {
72 if err, ok := m.FailOn["ReadFile"]; ok {
73 return nil, err
74 }
75
76 data, exists := m.Files[name]
77 if !exists {
78 return nil, fmt.Errorf("file not found: %s", name)
79 }
80 return data, nil
81}
82
83func (m *MockFileSystem) WriteFile(name string, data []byte, perm fs.FileMode) error {
84 if err, ok := m.FailOn["WriteFile"]; ok {
85 return err
86 }
87 m.Files[name] = data
88 return nil
89}
90
91// MockFile implements a simple in-memory file for testing
92type MockFile struct {
93 name string
94 buffer *bytes.Buffer
95 fs *MockFileSystem
96 position int64
97}
98
99// MockFileContents represents in-memory file contents for testing
100type MockFileContents struct {
101 name string
102 contents string
103}
104
105func (m *MockFileSystem) OpenFile(name string, flag int, perm fs.FileMode) (*os.File, error) {
106 if err, ok := m.FailOn["OpenFile"]; ok {
107 return nil, err
108 }
109
110 // Initialize the file content if it doesn't exist and we're not in read-only mode
111 if _, exists := m.Files[name]; !exists && (flag&os.O_CREATE != 0) {
112 m.Files[name] = []byte{}
113 }
114
115 data, exists := m.Files[name]
116 if !exists {
117 return nil, fmt.Errorf("file not found: %s", name)
118 }
119
120 // For OpenFile, we'll just use WriteFile to simulate file operations
banksean29d689f2025-06-23 15:41:26 +0000121 // The actual file handle isn't used for much in the localsshimmer code
Sean McCullough2cba6952025-04-25 20:32:10 +0000122 // but we still need to return a valid file handle
123 tmpFile, err := os.CreateTemp("", "mockfile-*")
124 if err != nil {
125 return nil, err
126 }
127 if _, err := tmpFile.Write(data); err != nil {
128 tmpFile.Close()
129 return nil, err
130 }
131 if _, err := tmpFile.Seek(0, 0); err != nil {
132 tmpFile.Close()
133 return nil, err
134 }
135
136 return tmpFile, nil
137}
138
Sean McCullough0d95d3a2025-04-30 16:22:28 +0000139func (m *MockFileSystem) TempFile(dir, pattern string) (*os.File, error) {
140 if err, ok := m.FailOn["TempFile"]; ok {
141 return nil, err
142 }
143
144 // Create an actual temporary file for testing purposes
145 tmpFile, err := os.CreateTemp(dir, pattern)
146 if err != nil {
147 return nil, err
148 }
149
150 // Record the temp file path
151 m.TempFiles = append(m.TempFiles, tmpFile.Name())
152
153 return tmpFile, nil
154}
155
156func (m *MockFileSystem) Rename(oldpath, newpath string) error {
157 if err, ok := m.FailOn["Rename"]; ok {
158 return err
159 }
160
161 // If the old path exists in our mock file system, move its contents
162 if data, exists := m.Files[oldpath]; exists {
163 m.Files[newpath] = data
164 delete(m.Files, oldpath)
165 }
166
167 return nil
168}
169
170func (m *MockFileSystem) SafeWriteFile(name string, data []byte, perm fs.FileMode) error {
171 if err, ok := m.FailOn["SafeWriteFile"]; ok {
172 return err
173 }
174
175 // For the mock, we'll create a backup if the file exists
176 if existingData, exists := m.Files[name]; exists {
177 backupName := name + ".bak"
178 m.Files[backupName] = existingData
179 }
180
181 // Write the new data
182 m.Files[name] = data
183
184 return nil
185}
186
Sean McCullough2cba6952025-04-25 20:32:10 +0000187// MockKeyGenerator implements KeyGenerator interface for testing
188type MockKeyGenerator struct {
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000189 privateKey ed25519.PrivateKey
190 publicKey ed25519.PublicKey
191 sshPublicKey ssh.PublicKey
Sean McCullough7013e9e2025-05-14 02:03:58 +0000192 caSigner ssh.Signer
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000193 FailOn map[string]error
Sean McCullough2cba6952025-04-25 20:32:10 +0000194}
195
Sean McCullough7013e9e2025-05-14 02:03:58 +0000196func NewMockKeyGenerator(privateKey ed25519.PrivateKey, publicKey ed25519.PublicKey, sshPublicKey ssh.PublicKey, caSigner ssh.Signer) *MockKeyGenerator {
Sean McCullough2cba6952025-04-25 20:32:10 +0000197 return &MockKeyGenerator{
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000198 privateKey: privateKey,
199 publicKey: publicKey,
200 sshPublicKey: sshPublicKey,
Sean McCullough7013e9e2025-05-14 02:03:58 +0000201 caSigner: caSigner,
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000202 FailOn: make(map[string]error),
Sean McCullough2cba6952025-04-25 20:32:10 +0000203 }
204}
205
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000206func (m *MockKeyGenerator) GenerateKeyPair() (ed25519.PrivateKey, ed25519.PublicKey, error) {
207 if err, ok := m.FailOn["GenerateKeyPair"]; ok {
208 return nil, nil, err
Sean McCullough2cba6952025-04-25 20:32:10 +0000209 }
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000210 return m.privateKey, m.publicKey, nil
Sean McCullough2cba6952025-04-25 20:32:10 +0000211}
212
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000213func (m *MockKeyGenerator) ConvertToSSHPublicKey(publicKey ed25519.PublicKey) (ssh.PublicKey, error) {
214 if err, ok := m.FailOn["ConvertToSSHPublicKey"]; ok {
Sean McCullough2cba6952025-04-25 20:32:10 +0000215 return nil, err
216 }
Sean McCullough7013e9e2025-05-14 02:03:58 +0000217 // If we're generating the CA public key, return the caSigner's public key
218 if m.caSigner != nil && bytes.Equal(publicKey, m.publicKey) {
219 return m.caSigner.PublicKey(), nil
220 }
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000221 return m.sshPublicKey, nil
Sean McCullough2cba6952025-04-25 20:32:10 +0000222}
223
224// setupMocks sets up common mocks for testing
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000225func setupMocks(t *testing.T) (*MockFileSystem, *MockKeyGenerator, ed25519.PrivateKey) {
226 // Generate a real Ed25519 key pair
227 publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
Sean McCullough2cba6952025-04-25 20:32:10 +0000228 if err != nil {
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000229 t.Fatalf("Failed to generate test key pair: %v", err)
Sean McCullough2cba6952025-04-25 20:32:10 +0000230 }
231
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000232 // Generate a test SSH public key
233 sshPublicKey, err := ssh.NewPublicKey(publicKey)
Sean McCullough2cba6952025-04-25 20:32:10 +0000234 if err != nil {
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000235 t.Fatalf("Failed to generate test SSH public key: %v", err)
Sean McCullough2cba6952025-04-25 20:32:10 +0000236 }
237
Sean McCullough7013e9e2025-05-14 02:03:58 +0000238 // Create CA key pair
239 _, caPrivKey, err := ed25519.GenerateKey(rand.Reader)
240 if err != nil {
241 t.Fatalf("Failed to generate CA key pair: %v", err)
242 }
243
244 // Create CA signer
245 caSigner, err := ssh.NewSignerFromKey(caPrivKey)
246 if err != nil {
247 t.Fatalf("Failed to create CA signer: %v", err)
248 }
249
Sean McCullough2cba6952025-04-25 20:32:10 +0000250 // Create mocks
251 mockFS := NewMockFileSystem()
Sean McCullough7013e9e2025-05-14 02:03:58 +0000252 mockKG := NewMockKeyGenerator(privateKey, publicKey, sshPublicKey, caSigner)
253
254 // Add some files needed for tests
255 mockFS.Files["/home/testuser/.config/sketch/host_cert"] = []byte("test-certificate")
256 caPubKeyBytes := ssh.MarshalAuthorizedKey(ssh.PublicKey(caSigner.PublicKey()))
257 mockFS.Files["/home/testuser/.config/sketch/container_ca.pub"] = caPubKeyBytes
Sean McCullough2cba6952025-04-25 20:32:10 +0000258
259 return mockFS, mockKG, privateKey
260}
261
banksean29d689f2025-06-23 15:41:26 +0000262// Helper function to setup a basic LocalSSHimmer for testing
263func setupTestLocalSSHimmer(t *testing.T) (*LocalSSHimmer, *MockFileSystem, *MockKeyGenerator) {
Sean McCullough2cba6952025-04-25 20:32:10 +0000264 mockFS, mockKG, _ := setupMocks(t)
265
266 // Setup home dir in mock filesystem
267 homePath := "/home/testuser"
Sean McCulloughc796e7f2025-04-30 08:44:06 -0700268 sketchDir := filepath.Join(homePath, ".config/sketch")
Sean McCullough2cba6952025-04-25 20:32:10 +0000269 mockFS.CreatedDirs[sketchDir] = true
270
Sean McCullough0d95d3a2025-04-30 16:22:28 +0000271 // Create empty files so the tests don't fail
272 sketchConfigPath := filepath.Join(sketchDir, "ssh_config")
273 mockFS.Files[sketchConfigPath] = []byte("")
274 knownHostsPath := filepath.Join(sketchDir, "known_hosts")
275 mockFS.Files[knownHostsPath] = []byte("")
276
Sean McCullough2cba6952025-04-25 20:32:10 +0000277 // Set HOME environment variable for the test
278 oldHome := os.Getenv("HOME")
279 os.Setenv("HOME", homePath)
280 t.Cleanup(func() { os.Setenv("HOME", oldHome) })
281
banksean29d689f2025-06-23 15:41:26 +0000282 // Create LocalSSHimmer with mocks
283 ssh, err := newLocalSSHimmerWithDeps("test-container", "localhost", "2222", mockFS, mockKG)
Sean McCullough2cba6952025-04-25 20:32:10 +0000284 if err != nil {
banksean29d689f2025-06-23 15:41:26 +0000285 t.Fatalf("Failed to create LocalSSHimmer: %v", err)
Sean McCullough2cba6952025-04-25 20:32:10 +0000286 }
287
288 return ssh, mockFS, mockKG
289}
290
banksean29d689f2025-06-23 15:41:26 +0000291func TestNewLocalSSHimmerCreatesRequiredDirectories(t *testing.T) {
Sean McCullough2cba6952025-04-25 20:32:10 +0000292 mockFS, mockKG, _ := setupMocks(t)
293
294 // Set HOME environment variable for the test
295 oldHome := os.Getenv("HOME")
296 os.Setenv("HOME", "/home/testuser")
297 defer func() { os.Setenv("HOME", oldHome) }()
298
Sean McCullough0d95d3a2025-04-30 16:22:28 +0000299 // Create empty files so the test doesn't fail
300 sketchDir := "/home/testuser/.config/sketch"
301 sketchConfigPath := filepath.Join(sketchDir, "ssh_config")
302 mockFS.Files[sketchConfigPath] = []byte("")
303 knownHostsPath := filepath.Join(sketchDir, "known_hosts")
304 mockFS.Files[knownHostsPath] = []byte("")
305
banksean29d689f2025-06-23 15:41:26 +0000306 // Create sshimmer
307 _, err := newLocalSSHimmerWithDeps("test-container", "localhost", "2222", mockFS, mockKG)
Sean McCullough2cba6952025-04-25 20:32:10 +0000308 if err != nil {
banksean29d689f2025-06-23 15:41:26 +0000309 t.Fatalf("Failed to create LocalSSHimmer: %v", err)
Sean McCullough2cba6952025-04-25 20:32:10 +0000310 }
311
Sean McCulloughc796e7f2025-04-30 08:44:06 -0700312 // Check if the .config/sketch directory was created
313 expectedDir := "/home/testuser/.config/sketch"
Sean McCullough2cba6952025-04-25 20:32:10 +0000314 if !mockFS.CreatedDirs[expectedDir] {
315 t.Errorf("Expected directory %s to be created", expectedDir)
316 }
317}
318
319func TestCreateKeyPairIfMissing(t *testing.T) {
banksean29d689f2025-06-23 15:41:26 +0000320 ssh, mockFS, _ := setupTestLocalSSHimmer(t)
Sean McCullough2cba6952025-04-25 20:32:10 +0000321
322 // Test key pair creation
Sean McCulloughc796e7f2025-04-30 08:44:06 -0700323 keyPath := "/home/testuser/.config/sketch/test_key"
Sean McCullough2cba6952025-04-25 20:32:10 +0000324 _, err := ssh.createKeyPairIfMissing(keyPath)
325 if err != nil {
326 t.Fatalf("Failed to create key pair: %v", err)
327 }
328
329 // Verify private key file was created
330 if _, exists := mockFS.Files[keyPath]; !exists {
331 t.Errorf("Private key file not created at %s", keyPath)
332 }
333
334 // Verify public key file was created
335 pubKeyPath := keyPath + ".pub"
336 if _, exists := mockFS.Files[pubKeyPath]; !exists {
337 t.Errorf("Public key file not created at %s", pubKeyPath)
338 }
339
340 // Verify public key content format
341 pubKeyContent, _ := mockFS.ReadFile(pubKeyPath)
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000342 if !bytes.HasPrefix(pubKeyContent, []byte("ssh-ed25519 ")) {
Sean McCullough2cba6952025-04-25 20:32:10 +0000343 t.Errorf("Public key does not have expected format, got: %s", pubKeyContent)
344 }
345}
346
347// TestAddContainerToSSHConfig tests that the container gets added to the SSH config
348// This test uses a direct approach since the OpenFile mocking is complex
349func TestAddContainerToSSHConfig(t *testing.T) {
350 // Create a temporary directory for test files
banksean29d689f2025-06-23 15:41:26 +0000351 tempDir, err := os.MkdirTemp("", "localsshimmer-test-*")
Sean McCullough2cba6952025-04-25 20:32:10 +0000352 if err != nil {
353 t.Fatalf("Failed to create temp dir: %v", err)
354 }
355 defer os.RemoveAll(tempDir)
356
357 // Create real files in temp directory
358 configPath := filepath.Join(tempDir, "ssh_config")
359 initialConfig := `# SSH Config
360Host existing-host
361 HostName example.com
362 User testuser
363`
Autoformatter33f71722025-04-25 23:23:22 +0000364 if err := os.WriteFile(configPath, []byte(initialConfig), 0o644); err != nil {
Sean McCullough2cba6952025-04-25 20:32:10 +0000365 t.Fatalf("Failed to write initial config: %v", err)
366 }
367
banksean29d689f2025-06-23 15:41:26 +0000368 // Create a sshimmer with the real filesystem but custom paths
369 ssh := &LocalSSHimmer{
Sean McCullough2cba6952025-04-25 20:32:10 +0000370 cntrName: "test-container",
371 sshHost: "localhost",
372 sshPort: "2222",
373 sshConfigPath: configPath,
374 userIdentityPath: filepath.Join(tempDir, "user_identity"),
375 fs: &RealFileSystem{},
376 kg: &RealKeyGenerator{},
377 }
378
379 // Add container to SSH config
380 err = ssh.addContainerToSSHConfig()
381 if err != nil {
382 t.Fatalf("Failed to add container to SSH config: %v", err)
383 }
384
385 // Read the updated file
386 configData, err := os.ReadFile(configPath)
387 if err != nil {
388 t.Fatalf("Failed to read updated config: %v", err)
389 }
390 configStr := string(configData)
391
392 // Check for expected values
393 if !strings.Contains(configStr, "Host test-container") {
394 t.Errorf("Container host entry not found in config")
395 }
396
397 if !strings.Contains(configStr, "HostName localhost") {
398 t.Errorf("HostName not correctly added to SSH config")
399 }
400
401 if !strings.Contains(configStr, "Port 2222") {
402 t.Errorf("Port not correctly added to SSH config")
403 }
404
405 if !strings.Contains(configStr, "User root") {
406 t.Errorf("User not correctly set to root in SSH config")
407 }
408
409 // Check if identity file path is correct
410 identityLine := "IdentityFile " + ssh.userIdentityPath
411 if !strings.Contains(configStr, identityLine) {
412 t.Errorf("Identity file path not correctly added to SSH config")
413 }
414}
415
416func TestAddContainerToKnownHosts(t *testing.T) {
417 // Skip this test as it requires more complex setup
banksean29d689f2025-06-23 15:41:26 +0000418 // The TestLocalSSHimmerCleanup test covers the addContainerToKnownHosts
Sean McCullough2cba6952025-04-25 20:32:10 +0000419 // functionality in a more integrated way
banksean29d689f2025-06-23 15:41:26 +0000420 t.Skip("This test requires more complex setup, integrated test coverage exists in TestLocalSSHimmerCleanup")
Sean McCullough2cba6952025-04-25 20:32:10 +0000421}
422
423func TestRemoveContainerFromSSHConfig(t *testing.T) {
424 // Create a temporary directory for test files
banksean29d689f2025-06-23 15:41:26 +0000425 tempDir, err := os.MkdirTemp("", "localsshimmer-test-*")
Sean McCullough2cba6952025-04-25 20:32:10 +0000426 if err != nil {
427 t.Fatalf("Failed to create temp dir: %v", err)
428 }
429 defer os.RemoveAll(tempDir)
430
431 // Create paths for test files
432 sshConfigPath := filepath.Join(tempDir, "ssh_config")
433 userIdentityPath := filepath.Join(tempDir, "user_identity")
434 knownHostsPath := filepath.Join(tempDir, "known_hosts")
435
436 // Create initial SSH config with container entry
437 cntrName := "test-container"
438 sshHost := "localhost"
439 sshPort := "2222"
440
441 initialConfig := fmt.Sprintf(
442 `Host existing-host
443 HostName example.com
444 User testuser
445
446Host %s
447 HostName %s
448 User root
449 Port %s
450 IdentityFile %s
451 UserKnownHostsFile %s
452`,
453 cntrName, sshHost, sshPort, userIdentityPath, knownHostsPath,
454 )
455
Autoformatter33f71722025-04-25 23:23:22 +0000456 if err := os.WriteFile(sshConfigPath, []byte(initialConfig), 0o644); err != nil {
Sean McCullough2cba6952025-04-25 20:32:10 +0000457 t.Fatalf("Failed to write initial SSH config: %v", err)
458 }
459
banksean29d689f2025-06-23 15:41:26 +0000460 // Create a sshimmer with the real filesystem but custom paths
461 ssh := &LocalSSHimmer{
Sean McCullough2cba6952025-04-25 20:32:10 +0000462 cntrName: cntrName,
463 sshHost: sshHost,
464 sshPort: sshPort,
465 sshConfigPath: sshConfigPath,
466 userIdentityPath: userIdentityPath,
467 knownHostsPath: knownHostsPath,
468 fs: &RealFileSystem{},
469 }
470
471 // Remove container from SSH config
472 err = ssh.removeContainerFromSSHConfig()
473 if err != nil {
474 t.Fatalf("Failed to remove container from SSH config: %v", err)
475 }
476
477 // Read the updated file
478 configData, err := os.ReadFile(sshConfigPath)
479 if err != nil {
480 t.Fatalf("Failed to read updated config: %v", err)
481 }
482 configStr := string(configData)
483
484 // Check if the container host entry was removed
485 if strings.Contains(configStr, "Host "+cntrName) {
486 t.Errorf("Container host not removed from SSH config")
487 }
488
489 // Check if existing host remains
490 if !strings.Contains(configStr, "Host existing-host") {
491 t.Errorf("Existing host entry affected by container removal")
492 }
493}
494
495func TestRemoveContainerFromKnownHosts(t *testing.T) {
banksean29d689f2025-06-23 15:41:26 +0000496 ssh, mockFS, _ := setupTestLocalSSHimmer(t)
Sean McCullough2cba6952025-04-25 20:32:10 +0000497
498 // Setup server public key
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000499 _, publicKey, _ := ssh.kg.GenerateKeyPair()
500 sshPublicKey, _ := ssh.kg.ConvertToSSHPublicKey(publicKey)
501 ssh.serverPublicKey = sshPublicKey
Sean McCullough2cba6952025-04-25 20:32:10 +0000502
503 // Create host line to be removed
504 hostLine := "[localhost]:2222 ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ..."
505 otherLine := "otherhost ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQ..."
506
507 // Set initial content with the line to be removed
508 initialContent := otherLine + "\n" + hostLine
509 mockFS.Files[ssh.knownHostsPath] = []byte(initialContent)
510
511 // Add the host to test remove function
512 err := ssh.addContainerToKnownHosts()
513 if err != nil {
514 t.Fatalf("Failed to add container to known_hosts for removal test: %v", err)
515 }
516
517 // Now remove it
518 err = ssh.removeContainerFromKnownHosts()
519 if err != nil {
520 t.Fatalf("Failed to remove container from known_hosts: %v", err)
521 }
522
523 // Verify content
524 updatedContent, _ := mockFS.ReadFile(ssh.knownHostsPath)
525 content := string(updatedContent)
526
527 hostPattern := ssh.sshHost + ":" + ssh.sshPort
528 if strings.Contains(content, hostPattern) {
529 t.Errorf("Container entry not removed from known_hosts")
530 }
531
532 // Verify other content remains
533 if !strings.Contains(content, otherLine) {
534 t.Errorf("Other known_hosts entries improperly removed")
535 }
536}
537
banksean29d689f2025-06-23 15:41:26 +0000538func TestLocalSSHimmerCleanup(t *testing.T) {
Sean McCullough2cba6952025-04-25 20:32:10 +0000539 // Create a temporary directory for test files
banksean29d689f2025-06-23 15:41:26 +0000540 tempDir, err := os.MkdirTemp("", "localsshimmer-test-*")
Sean McCullough2cba6952025-04-25 20:32:10 +0000541 if err != nil {
542 t.Fatalf("Failed to create temp dir: %v", err)
543 }
544 defer os.RemoveAll(tempDir)
545
546 // Create paths for test files
547 sshConfigPath := filepath.Join(tempDir, "ssh_config")
548 userIdentityPath := filepath.Join(tempDir, "user_identity")
549 knownHostsPath := filepath.Join(tempDir, "known_hosts")
550 serverIdentityPath := filepath.Join(tempDir, "server_identity")
551
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000552 // Create keys for server key
553 publicKey, _, err := ed25519.GenerateKey(rand.Reader)
Sean McCullough2cba6952025-04-25 20:32:10 +0000554 if err != nil {
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000555 t.Fatalf("Failed to generate key pair: %v", err)
Sean McCullough2cba6952025-04-25 20:32:10 +0000556 }
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000557 sshPublicKey, err := ssh.NewPublicKey(publicKey)
Sean McCullough2cba6952025-04-25 20:32:10 +0000558 if err != nil {
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000559 t.Fatalf("Failed to generate SSH public key: %v", err)
Sean McCullough2cba6952025-04-25 20:32:10 +0000560 }
561
562 // Initialize files
Autoformatter33f71722025-04-25 23:23:22 +0000563 os.WriteFile(sshConfigPath, []byte("initial ssh_config content"), 0o644)
564 os.WriteFile(knownHostsPath, []byte("initial known_hosts content"), 0o644)
Sean McCullough2cba6952025-04-25 20:32:10 +0000565
banksean29d689f2025-06-23 15:41:26 +0000566 // Create a sshimmer with the real filesystem but custom paths
Sean McCullough2cba6952025-04-25 20:32:10 +0000567 cntrName := "test-container"
568 sshHost := "localhost"
569 sshPort := "2222"
570
banksean29d689f2025-06-23 15:41:26 +0000571 ssh := &LocalSSHimmer{
Sean McCullough2cba6952025-04-25 20:32:10 +0000572 cntrName: cntrName,
573 sshHost: sshHost,
574 sshPort: sshPort,
575 sshConfigPath: sshConfigPath,
576 userIdentityPath: userIdentityPath,
577 knownHostsPath: knownHostsPath,
578 serverIdentityPath: serverIdentityPath,
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000579 serverPublicKey: sshPublicKey,
Sean McCullough2cba6952025-04-25 20:32:10 +0000580 fs: &RealFileSystem{},
581 kg: &RealKeyGenerator{},
582 }
583
584 // Add container to configs
585 err = ssh.addContainerToSSHConfig()
586 if err != nil {
587 t.Fatalf("Failed to set up SSH config for cleanup test: %v", err)
588 }
589
590 err = ssh.addContainerToKnownHosts()
591 if err != nil {
592 t.Fatalf("Failed to set up known_hosts for cleanup test: %v", err)
593 }
594
595 // Execute cleanup
596 err = ssh.Cleanup()
597 if err != nil {
598 t.Fatalf("Cleanup failed: %v", err)
599 }
600
601 // Read updated files
602 configData, err := os.ReadFile(sshConfigPath)
603 if err != nil {
604 t.Fatalf("Failed to read updated SSH config: %v", err)
605 }
606 configStr := string(configData)
607
608 // Check container was removed from SSH config
609 hostEntry := "Host " + ssh.cntrName
610 if strings.Contains(configStr, hostEntry) {
611 t.Errorf("Container not removed from SSH config during cleanup")
612 }
613
614 // Verify known hosts was updated
615 knownHostsContent, err := os.ReadFile(knownHostsPath)
616 if err != nil {
617 t.Fatalf("Failed to read updated known_hosts: %v", err)
618 }
619
620 expectedHostPattern := ssh.sshHost + ":" + ssh.sshPort
621 if strings.Contains(string(knownHostsContent), expectedHostPattern) {
622 t.Errorf("Container not removed from known_hosts during cleanup")
623 }
624}
625
Sean McCullough15c95282025-05-08 16:48:38 -0700626func TestCheckForInclude_userAccepts(t *testing.T) {
Sean McCullough2cba6952025-04-25 20:32:10 +0000627 mockFS := NewMockFileSystem()
628
629 // Set HOME environment variable for the test
630 oldHome := os.Getenv("HOME")
631 os.Setenv("HOME", "/home/testuser")
632 defer func() { os.Setenv("HOME", oldHome) }()
633
634 // Create a mock ssh config with the expected include
Sean McCulloughc796e7f2025-04-30 08:44:06 -0700635 includeLine := "Include /home/testuser/.config/sketch/ssh_config"
Sean McCullough2cba6952025-04-25 20:32:10 +0000636 initialConfig := fmt.Sprintf("%s\nHost example\n HostName example.com\n", includeLine)
637
638 // Add the config to the mock filesystem
639 sshConfigPath := "/home/testuser/.ssh/config"
640 mockFS.Files[sshConfigPath] = []byte(initialConfig)
Sean McCullough15c95282025-05-08 16:48:38 -0700641 stdinReader := bufio.NewReader(strings.NewReader("y\n"))
Sean McCullough2cba6952025-04-25 20:32:10 +0000642 // Test the function with our mock
Sean McCullough15c95282025-05-08 16:48:38 -0700643 err := CheckForIncludeWithFS(mockFS, *stdinReader)
Sean McCullough2cba6952025-04-25 20:32:10 +0000644 if err != nil {
645 t.Fatalf("CheckForInclude failed with proper include: %v", err)
646 }
647
648 // Now test with config missing the include
649 mockFS.Files[sshConfigPath] = []byte("Host example\n HostName example.com\n")
650
Sean McCullough15c95282025-05-08 16:48:38 -0700651 stdinReader = bufio.NewReader(strings.NewReader("y\n"))
652 err = CheckForIncludeWithFS(mockFS, *stdinReader)
Sean McCulloughc796e7f2025-04-30 08:44:06 -0700653 if err != nil {
654 t.Fatalf("CheckForInclude should have created the Include line without an error")
Sean McCullough2cba6952025-04-25 20:32:10 +0000655 }
656}
657
Sean McCullough15c95282025-05-08 16:48:38 -0700658func TestCheckForInclude_userDeclines(t *testing.T) {
659 mockFS := NewMockFileSystem()
660
661 // Set HOME environment variable for the test
662 oldHome := os.Getenv("HOME")
663 os.Setenv("HOME", "/home/testuser")
664 defer func() { os.Setenv("HOME", oldHome) }()
665
666 // Create a mock ssh config with the expected include
667 includeLine := "Include /home/testuser/.config/sketch/ssh_config"
668 initialConfig := fmt.Sprintf("%s\nHost example\n HostName example.com\n", includeLine)
669
670 // Add the config to the mock filesystem
671 sshConfigPath := "/home/testuser/.ssh/config"
672 mockFS.Files[sshConfigPath] = []byte(initialConfig)
673 stdinReader := bufio.NewReader(strings.NewReader("n\n"))
674 // Test the function with our mock
675 err := CheckForIncludeWithFS(mockFS, *stdinReader)
676 if err != nil {
677 t.Fatalf("CheckForInclude failed with proper include: %v", err)
678 }
679
680 // Now test with config missing the include
681 missingInclude := []byte("Host example\n HostName example.com\n")
682 mockFS.Files[sshConfigPath] = missingInclude
683
684 stdinReader = bufio.NewReader(strings.NewReader("n\n"))
685 err = CheckForIncludeWithFS(mockFS, *stdinReader)
686 if err == nil {
687 t.Errorf("CheckForInclude should have returned an error")
688 }
689 if !bytes.Equal(mockFS.Files[sshConfigPath], missingInclude) {
690 t.Errorf("ssh config should not have been edited")
691 }
692}
693
banksean29d689f2025-06-23 15:41:26 +0000694func TestLocalSSHimmerWithErrors(t *testing.T) {
Sean McCullough2cba6952025-04-25 20:32:10 +0000695 // Test directory creation failure
696 mockFS := NewMockFileSystem()
Sean McCulloughc796e7f2025-04-30 08:44:06 -0700697 mockFS.FailOn["MkdirAll"] = fmt.Errorf("mock mkdir error")
Sean McCullough7013e9e2025-05-14 02:03:58 +0000698 mockKG := NewMockKeyGenerator(nil, nil, nil, nil)
Sean McCullough2cba6952025-04-25 20:32:10 +0000699
700 // Set HOME environment variable for the test
701 oldHome := os.Getenv("HOME")
702 os.Setenv("HOME", "/home/testuser")
703 defer func() { os.Setenv("HOME", oldHome) }()
704
banksean29d689f2025-06-23 15:41:26 +0000705 // Try to create sshimmer with failing FS
706 _, err := newLocalSSHimmerWithDeps("test-container", "localhost", "2222", mockFS, mockKG)
Sean McCullough2cba6952025-04-25 20:32:10 +0000707 if err == nil || !strings.Contains(err.Error(), "mock mkdir error") {
708 t.Errorf("Should have failed with mkdir error, got: %v", err)
709 }
710
711 // Test key generation failure
712 mockFS = NewMockFileSystem()
Sean McCullough7013e9e2025-05-14 02:03:58 +0000713 mockKG = NewMockKeyGenerator(nil, nil, nil, nil)
Sean McCullough3e9d80c2025-05-13 23:35:23 +0000714 mockKG.FailOn["GenerateKeyPair"] = fmt.Errorf("mock key generation error")
Sean McCullough2cba6952025-04-25 20:32:10 +0000715
banksean29d689f2025-06-23 15:41:26 +0000716 _, err = newLocalSSHimmerWithDeps("test-container", "localhost", "2222", mockFS, mockKG)
Sean McCullough2cba6952025-04-25 20:32:10 +0000717 if err == nil || !strings.Contains(err.Error(), "key generation error") {
718 t.Errorf("Should have failed with key generation error, got: %v", err)
719 }
720}
721
722func TestRealSSHTheatherInit(t *testing.T) {
Sean McCullough7013e9e2025-05-14 02:03:58 +0000723 // Skip this test as it requires real files for the CA which we don't want to create
724 // in a real integration test
725 t.Skip("Skipping test that requires real file system access for the CA")
726}
Sean McCullough2cba6952025-04-25 20:32:10 +0000727
Sean McCullough7013e9e2025-05-14 02:03:58 +0000728// Methods to help with the mocking interface
729func (m *MockKeyGenerator) GetCASigner() ssh.Signer {
730 return m.caSigner
731}
Sean McCullough2cba6952025-04-25 20:32:10 +0000732
Sean McCullough7013e9e2025-05-14 02:03:58 +0000733func (m *MockKeyGenerator) IsMock() bool {
734 return true
Sean McCullough2cba6952025-04-25 20:32:10 +0000735}