From 838941e3eb5eed1b9de9d23e5ca1c1409c5726f7 Mon Sep 17 00:00:00 2001
From: Gu1llaum-3 <archambault.guillaume06@gmail.com>
Date: Mon, 23 Feb 2026 23:04:40 +0100
Subject: [PATCH] fix: allow env vars and SSH tokens in IdentityFile validation
 (issue #33)

ValidateIdentityFile now accepts $VAR/${VAR} (expanded via os.Expand, undefined vars accepted as-is) and SSH tokens like %d, %h before falling back to os.Stat.
The raw value is preserved when writing to ssh_config.
---
 internal/validation/ssh.go      | 19 +++++++++++++++++++
 internal/validation/ssh_test.go | 14 ++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/internal/validation/ssh.go b/internal/validation/ssh.go
index 93bee63..dea0ab8 100644
--- a/internal/validation/ssh.go
+++ b/internal/validation/ssh.go
@@ -66,6 +66,25 @@ func ValidateIdentityFile(path string) bool {
 	if path == "" {
 		return true // Optional field
 	}
+	// SSH tokens (e.g. %d, %h, %r, %u) are resolved by SSH at connection time
+	sshTokenRegex := regexp.MustCompile(`%[hprunCdiklLT]`)
+	if sshTokenRegex.MatchString(path) {
+		return true
+	}
+	// Expand environment variables ($VAR and ${VAR}); track undefined ones
+	hasUndefined := false
+	path = os.Expand(path, func(key string) string {
+		val, ok := os.LookupEnv(key)
+		if !ok {
+			hasUndefined = true
+			return "$" + key
+		}
+		return val
+	})
+	// If any variable was undefined, accept the path (SSH will report the error)
+	if hasUndefined {
+		return true
+	}
 	// Expand ~ to home directory
 	if strings.HasPrefix(path, "~/") {
 		homeDir, err := os.UserHomeDir()
diff --git a/internal/validation/ssh_test.go b/internal/validation/ssh_test.go
index 32dca0c..17564d1 100644
--- a/internal/validation/ssh_test.go
+++ b/internal/validation/ssh_test.go
@@ -133,6 +133,9 @@ func TestValidateIdentityFile(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	// Set up an env var pointing to the valid file's directory for env var tests
+	t.Setenv("TEST_SSHM_DIR", tmpDir)
+
 	tests := []struct {
 		name string
 		path string
@@ -143,6 +146,13 @@ func TestValidateIdentityFile(t *testing.T) {
 		{"non-existent file", "/path/to/nonexistent", false},
 		// Skip tilde path test in CI environments where ~/.ssh/id_rsa may not exist
 		// {"tilde path", "~/.ssh/id_rsa", true}, // Will pass if file exists
+		// Environment variable expansion (issue #33)
+		{"env var $VAR/key defined", "$TEST_SSHM_DIR/test_key", true},
+		{"env var ${VAR}/key defined", "${TEST_SSHM_DIR}/test_key", true},
+		{"env var undefined", "$UNDEFINED_SSHM_VAR_XYZ/key", true},
+		// SSH tokens
+		{"SSH token %d", "%d/.ssh/id_rsa", true},
+		{"SSH token %h", "%h-key", true},
 	}
 
 	for _, tt := range tests {
@@ -170,6 +180,7 @@ func TestValidateHost(t *testing.T) {
 	if err := os.WriteFile(validIdentity, []byte("test"), 0600); err != nil {
 		t.Fatal(err)
 	}
+	t.Setenv("TEST_SSHM_HOST_DIR", tmpDir)
 
 	tests := []struct {
 		name     string
@@ -187,6 +198,9 @@ func TestValidateHost(t *testing.T) {
 		{"invalid hostname", "myserver", "invalid..hostname", "22", "", true},
 		{"invalid port", "myserver", "example.com", "99999", "", true},
 		{"invalid identity", "myserver", "example.com", "22", "/nonexistent", true},
+		// Environment variables and SSH tokens in identity (issue #33)
+		{"identity with env var", "myserver", "example.com", "22", "$TEST_SSHM_HOST_DIR/test_key", false},
+		{"identity with SSH token", "myserver", "example.com", "22", "%d/.ssh/id_rsa", false},
 	}
 
 	for _, tt := range tests {