dockerimg: rename ssh_theater to local_sshimmer

Co-Authored-By: sketch <hello@sketch.dev>
Change-ID: s650e5e56560011fak
diff --git a/dockerimg/local_sshimmer.go b/dockerimg/local_sshimmer.go
new file mode 100644
index 0000000..3f0fb11
--- /dev/null
+++ b/dockerimg/local_sshimmer.go
@@ -0,0 +1,780 @@
+package dockerimg
+
+import (
+	"bufio"
+	"bytes"
+	"crypto/ed25519"
+	"crypto/rand"
+	"encoding/pem"
+	"fmt"
+	"io/fs"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/kevinburke/ssh_config"
+	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/knownhosts"
+)
+
+// Ed25519 has a fixed key size, no bit size constant needed
+
+// LocalSSHimmer does the necessary key pair generation, known_hosts updates, ssh_config file updates etc steps
+// so that ssh can connect to a locally running sketch container to other local processes like vscode without
+// the user having to run the usual ssh obstacle course.
+//
+// LocalSSHimmer does not modify your default .ssh/config, or known_hosts files.  However, in order for you
+// to be able to use it properly you will have to make a one-time edit to your ~/.ssh/config file.
+//
+// In your ~/.ssh/config file, add the following line:
+//
+// Include $HOME/.config/sketch/ssh_config
+//
+// where $HOME is your home directory.
+//
+// LocalSSHimmer uses Ed25519 keys for improved security and performance.
+type LocalSSHimmer struct {
+	cntrName string
+	sshHost  string
+	sshPort  string
+
+	knownHostsPath     string
+	userIdentityPath   string
+	sshConfigPath      string
+	serverIdentityPath string
+	containerCAPath    string
+	hostCertPath       string
+
+	serverPublicKey      ssh.PublicKey
+	serverIdentity       []byte
+	userIdentity         []byte
+	hostCertificate      []byte
+	containerCA          ssh.Signer
+	containerCAPublicKey ssh.PublicKey
+
+	fs FileSystem
+	kg KeyGenerator
+}
+
+// NewLocalSSHimmer will set up everything so that you can use ssh on localhost to connect to
+// the sketch container.  Call #Clean when you are done with the container to remove the
+// various entries it created in its known_hosts and ssh_config files. Also note that
+// this will generate key pairs for both the ssh server identity and the user identity, if
+// these files do not already exist.  These key pair files are not deleted by #Cleanup,
+// so they can be re-used across invocations of sketch. This means every sketch container
+// that runs on this host will use the same ssh server identity.
+// The system uses Ed25519 keys for better security and performance.
+//
+// If this doesn't return an error, you should be able to run "ssh <cntrName>"
+// in a terminal on your host machine to open a shell into the container without having
+// to manually accept changes to your known_hosts file etc.
+func NewLocalSSHimmer(cntrName, sshHost, sshPort string) (*LocalSSHimmer, error) {
+	return newLocalSSHimmerWithDeps(cntrName, sshHost, sshPort, &RealFileSystem{}, &RealKeyGenerator{})
+}
+
+// newLocalSSHimmerWithDeps creates a new LocalSSHimmer with the specified dependencies
+func newLocalSSHimmerWithDeps(cntrName, sshHost, sshPort string, fs FileSystem, kg KeyGenerator) (*LocalSSHimmer, error) {
+	base := filepath.Join(os.Getenv("HOME"), ".config", "sketch")
+	if _, err := fs.Stat(base); err != nil {
+		if err := fs.MkdirAll(base, 0o777); err != nil {
+			return nil, fmt.Errorf("couldn't create %s: %w", base, err)
+		}
+	}
+
+	cst := &LocalSSHimmer{
+		cntrName:           cntrName,
+		sshHost:            sshHost,
+		sshPort:            sshPort,
+		knownHostsPath:     filepath.Join(base, "known_hosts"),
+		userIdentityPath:   filepath.Join(base, "container_user_identity"),
+		serverIdentityPath: filepath.Join(base, "container_server_identity"),
+		containerCAPath:    filepath.Join(base, "container_ca"),
+		hostCertPath:       filepath.Join(base, "host_cert"),
+		sshConfigPath:      filepath.Join(base, "ssh_config"),
+		fs:                 fs,
+		kg:                 kg,
+	}
+
+	// Step 1: Create regular server identity for the container SSH server
+	if _, err := cst.createKeyPairIfMissing(cst.serverIdentityPath); err != nil {
+		return nil, fmt.Errorf("couldn't create server identity: %w", err)
+	}
+
+	// Step 2: Create user identity that will be used to connect to the container
+	if _, err := cst.createKeyPairIfMissing(cst.userIdentityPath); err != nil {
+		return nil, fmt.Errorf("couldn't create user identity: %w", err)
+	}
+
+	// Step 3: Generate host certificate and CA for mutual authentication
+	// This now handles both CA creation and certificate signing in one step
+	if err := cst.createHostCertificate(cst.userIdentityPath); err != nil {
+		return nil, fmt.Errorf("couldn't create host certificate: %w", err)
+	}
+
+	// Step 5: Load all necessary key materials
+	serverIdentity, err := fs.ReadFile(cst.serverIdentityPath)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't read container's ssh server identity: %w", err)
+	}
+	cst.serverIdentity = serverIdentity
+
+	serverPubKeyBytes, err := fs.ReadFile(cst.serverIdentityPath + ".pub")
+	if err != nil {
+		return nil, fmt.Errorf("couldn't read ssh server public key file: %w", err)
+	}
+	serverPubKey, _, _, _, err := ssh.ParseAuthorizedKey(serverPubKeyBytes)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't parse ssh server public key: %w", err)
+	}
+	cst.serverPublicKey = serverPubKey
+
+	userIdentity, err := fs.ReadFile(cst.userIdentityPath + ".pub")
+	if err != nil {
+		return nil, fmt.Errorf("couldn't read ssh user identity: %w", err)
+	}
+	cst.userIdentity = userIdentity
+
+	hostCert, err := fs.ReadFile(cst.hostCertPath)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't read host certificate: %w", err)
+	}
+	cst.hostCertificate = hostCert
+
+	// Step 6: Configure SSH settings
+	if err := cst.addContainerToSSHConfig(); err != nil {
+		return nil, fmt.Errorf("couldn't add container to ssh_config: %w", err)
+	}
+
+	if err := cst.addContainerToKnownHosts(); err != nil {
+		return nil, fmt.Errorf("couldn't update known hosts: %w", err)
+	}
+
+	return cst, nil
+}
+
+func checkSSHResolve(hostname string) error {
+	cmd := exec.Command("ssh", "-T", hostname)
+	out, err := cmd.CombinedOutput()
+	if strings.HasPrefix(string(out), "ssh: Could not resolve hostname") {
+		return err
+	}
+	return nil
+}
+
+func CheckForIncludeWithFS(fs FileSystem, stdinReader bufio.Reader) error {
+	sketchSSHPathInclude := "Include " + filepath.Join(os.Getenv("HOME"), ".config", "sketch", "ssh_config")
+	defaultSSHPath := filepath.Join(os.Getenv("HOME"), ".ssh", "config")
+
+	// Read the existing SSH config file
+	existingContent, err := fs.ReadFile(defaultSSHPath)
+	if err != nil {
+		// If the file doesn't exist, create a new one with just the include line
+		if os.IsNotExist(err) {
+			return fs.SafeWriteFile(defaultSSHPath, []byte(sketchSSHPathInclude+"\n"), 0o644)
+		}
+		return fmt.Errorf("⚠️  SSH connections are disabled. cannot open SSH config file: %s: %w", defaultSSHPath, err)
+	}
+
+	// Parse the config file
+	cfg, err := ssh_config.Decode(bytes.NewReader(existingContent))
+	if err != nil {
+		return fmt.Errorf("couldn't decode ssh_config: %w", err)
+	}
+
+	var sketchInludePos *ssh_config.Position
+	var firstNonIncludePos *ssh_config.Position
+	for _, host := range cfg.Hosts {
+		for _, node := range host.Nodes {
+			inc, ok := node.(*ssh_config.Include)
+			if ok {
+				if strings.TrimSpace(inc.String()) == sketchSSHPathInclude {
+					pos := inc.Pos()
+					sketchInludePos = &pos
+				}
+			} else if firstNonIncludePos == nil && !strings.HasPrefix(strings.TrimSpace(node.String()), "#") {
+				pos := node.Pos()
+				firstNonIncludePos = &pos
+			}
+		}
+	}
+
+	if sketchInludePos == nil {
+		fmt.Printf("\nTo enable you to use ssh to connect to local sketch containers: \nAdd %q to the top of %s [y/N]? ", sketchSSHPathInclude, defaultSSHPath)
+		char, _, err := stdinReader.ReadRune()
+		if err != nil {
+			return fmt.Errorf("couldn't read from stdin: %w", err)
+		}
+		if char != 'y' && char != 'Y' {
+			return fmt.Errorf("User declined to edit ssh config file")
+		}
+		// Include line not found, add it to the top of the file
+		cfgBytes, err := cfg.MarshalText()
+		if err != nil {
+			return fmt.Errorf("couldn't marshal ssh_config: %w", err)
+		}
+
+		// Add the include line to the beginning
+		cfgBytes = append([]byte(sketchSSHPathInclude+"\n"), cfgBytes...)
+
+		// Safely write the updated config back to the file
+		if err := fs.SafeWriteFile(defaultSSHPath, cfgBytes, 0o644); err != nil {
+			return fmt.Errorf("couldn't safely write ssh_config: %w", err)
+		}
+		return nil
+	}
+
+	if firstNonIncludePos != nil && firstNonIncludePos.Line < sketchInludePos.Line {
+		fmt.Printf("⚠️  SSH confg warning: the location of the Include statement for sketch's ssh config on line %d of %s may prevent ssh from working with sketch containers. try moving it to the top of the file (before any 'Host' lines) if ssh isn't working for you.\n", sketchInludePos.Line, defaultSSHPath)
+	}
+	return nil
+}
+
+func removeFromHosts(cntrName string, cfgHosts []*ssh_config.Host) []*ssh_config.Host {
+	hosts := []*ssh_config.Host{}
+	for _, host := range cfgHosts {
+		if host.Matches(cntrName) || strings.Contains(host.String(), cntrName) {
+			continue
+		}
+		patMatch := false
+		for _, pat := range host.Patterns {
+			if strings.Contains(pat.String(), cntrName) {
+				patMatch = true
+			}
+		}
+		if patMatch {
+			continue
+		}
+
+		hosts = append(hosts, host)
+	}
+	return hosts
+}
+
+// encodePrivateKeyToPEM encodes an Ed25519 private key for storage
+func encodePrivateKeyToPEM(privateKey ed25519.PrivateKey) []byte {
+	// No need to create a signer first, we can directly marshal the key
+
+	// Format and encode as a binary private key format
+	pkBytes, err := ssh.MarshalPrivateKey(privateKey, "sketch key")
+	if err != nil {
+		panic(fmt.Sprintf("failed to marshal private key: %v", err))
+	}
+
+	// Return PEM encoded bytes
+	return pem.EncodeToMemory(pkBytes)
+}
+
+func (c *LocalSSHimmer) writeKeyToFile(keyBytes []byte, filename string) error {
+	err := c.fs.WriteFile(filename, keyBytes, 0o600)
+	return err
+}
+
+func (c *LocalSSHimmer) createKeyPairIfMissing(idPath string) (ssh.PublicKey, error) {
+	if _, err := c.fs.Stat(idPath); err == nil {
+		return nil, nil
+	}
+
+	privateKey, publicKey, err := c.kg.GenerateKeyPair()
+	if err != nil {
+		return nil, fmt.Errorf("error generating key pair: %w", err)
+	}
+
+	sshPublicKey, err := c.kg.ConvertToSSHPublicKey(publicKey)
+	if err != nil {
+		return nil, fmt.Errorf("error converting to SSH public key: %w", err)
+	}
+
+	privateKeyPEM := encodePrivateKeyToPEM(privateKey)
+
+	err = c.writeKeyToFile(privateKeyPEM, idPath)
+	if err != nil {
+		return nil, fmt.Errorf("error writing private key to file %w", err)
+	}
+	pubKeyBytes := ssh.MarshalAuthorizedKey(sshPublicKey)
+
+	err = c.writeKeyToFile([]byte(pubKeyBytes), idPath+".pub")
+	if err != nil {
+		return nil, fmt.Errorf("error writing public key to file %w", err)
+	}
+	return sshPublicKey, nil
+}
+
+func (c *LocalSSHimmer) addSketchHostMatchIfMissing(cfg *ssh_config.Config) error {
+	found := false
+	for _, host := range cfg.Hosts {
+		if strings.Contains(host.String(), "host=\"sketch-*\"") {
+			found = true
+			break
+		}
+	}
+	if !found {
+		hostPattern, err := ssh_config.NewPattern("host=\"sketch-*\"")
+		if err != nil {
+			return fmt.Errorf("couldn't add pattern to ssh_config: %w", err)
+		}
+
+		hostCfg := &ssh_config.Host{Patterns: []*ssh_config.Pattern{hostPattern}}
+		hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "UserKnownHostsFile", Value: c.knownHostsPath})
+
+		hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "IdentityFile", Value: c.userIdentityPath})
+		hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.Empty{})
+
+		cfg.Hosts = append([]*ssh_config.Host{hostCfg}, cfg.Hosts...)
+	}
+	return nil
+}
+
+func (c *LocalSSHimmer) addContainerToSSHConfig() error {
+	// Read the existing file contents or start with an empty config if file doesn't exist
+	var configData []byte
+	var cfg *ssh_config.Config
+	var err error
+
+	configData, err = c.fs.ReadFile(c.sshConfigPath)
+	if err != nil {
+		// If the file doesn't exist, create an empty config
+		if os.IsNotExist(err) {
+			cfg = &ssh_config.Config{}
+		} else {
+			return fmt.Errorf("couldn't read ssh_config: %w", err)
+		}
+	} else {
+		// Parse the existing config
+		cfg, err = ssh_config.Decode(bytes.NewReader(configData))
+		if err != nil {
+			return fmt.Errorf("couldn't decode ssh_config: %w", err)
+		}
+	}
+
+	cntrPattern, err := ssh_config.NewPattern(c.cntrName)
+	if err != nil {
+		return fmt.Errorf("couldn't add pattern to ssh_config: %w", err)
+	}
+
+	// Remove any matches for this container if they already exist.
+	cfg.Hosts = removeFromHosts(c.cntrName, cfg.Hosts)
+
+	hostCfg := &ssh_config.Host{Patterns: []*ssh_config.Pattern{cntrPattern}}
+	hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "HostName", Value: c.sshHost})
+	hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "User", Value: "root"})
+	hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "Port", Value: c.sshPort})
+	hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "IdentityFile", Value: c.userIdentityPath})
+	hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "CertificateFile", Value: c.hostCertPath})
+	hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.KV{Key: "UserKnownHostsFile", Value: c.knownHostsPath})
+
+	hostCfg.Nodes = append(hostCfg.Nodes, &ssh_config.Empty{})
+	cfg.Hosts = append(cfg.Hosts, hostCfg)
+
+	if err := c.addSketchHostMatchIfMissing(cfg); err != nil {
+		return fmt.Errorf("couldn't add missing host match: %w", err)
+	}
+
+	cfgBytes, err := cfg.MarshalText()
+	if err != nil {
+		return fmt.Errorf("couldn't marshal ssh_config: %w", err)
+	}
+
+	// Safely write the updated configuration to file
+	if err := c.fs.SafeWriteFile(c.sshConfigPath, cfgBytes, 0o644); err != nil {
+		return fmt.Errorf("couldn't safely write ssh_config: %w", err)
+	}
+
+	return nil
+}
+
+func (c *LocalSSHimmer) addContainerToKnownHosts() error {
+	// Instead of adding individual host entries, we'll use a CA-based approach
+	// by adding a single "@cert-authority" entry
+
+	// Format the CA public key line for the known_hosts file
+	var caPublicKeyLine string
+	if c.containerCAPublicKey != nil {
+		// Create a line that trusts only localhost hosts with a certificate signed by our CA
+		// This restricts the CA authority to only localhost addresses for security
+		caLine := "@cert-authority localhost,127.0.0.1,[::1] " + string(ssh.MarshalAuthorizedKey(c.containerCAPublicKey))
+		caPublicKeyLine = strings.TrimSpace(caLine)
+	}
+
+	// For backward compatibility, also add the host key itself
+	pkBytes := c.serverPublicKey.Marshal()
+	if len(pkBytes) == 0 {
+		return fmt.Errorf("empty serverPublicKey, this is a bug")
+	}
+	hostKeyLine := knownhosts.Line([]string{c.sshHost + ":" + c.sshPort}, c.serverPublicKey)
+
+	// Read existing known_hosts content or start with empty if the file doesn't exist
+	outputLines := []string{}
+	existingContent, err := c.fs.ReadFile(c.knownHostsPath)
+	if err == nil {
+		scanner := bufio.NewScanner(bytes.NewReader(existingContent))
+		for scanner.Scan() {
+			line := scanner.Text()
+			// Skip existing CA lines to avoid duplicates
+			if caPublicKeyLine != "" && strings.HasPrefix(line, "@cert-authority * ") {
+				continue
+			}
+			// Skip existing host key lines for this host:port
+			if strings.Contains(line, c.sshHost+":"+c.sshPort) {
+				continue
+			}
+			outputLines = append(outputLines, line)
+		}
+	} else if !os.IsNotExist(err) {
+		return fmt.Errorf("couldn't read known_hosts file: %w", err)
+	}
+
+	// Add the CA public key line if available
+	if caPublicKeyLine != "" {
+		outputLines = append(outputLines, caPublicKeyLine)
+	}
+
+	// Also add the host key line for backward compatibility
+	outputLines = append(outputLines, hostKeyLine)
+
+	// Safely write the updated content to the file
+	if err := c.fs.SafeWriteFile(c.knownHostsPath, []byte(strings.Join(outputLines, "\n")), 0o644); err != nil {
+		return fmt.Errorf("couldn't safely write updated known_hosts to %s: %w", c.knownHostsPath, err)
+	}
+
+	return nil
+}
+
+func (c *LocalSSHimmer) removeContainerFromKnownHosts() error {
+	// Read the existing known_hosts file
+	existingContent, err := c.fs.ReadFile(c.knownHostsPath)
+	if err != nil {
+		// If the file doesn't exist, there's nothing to do
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return fmt.Errorf("couldn't read known_hosts file: %w", err)
+	}
+
+	// Line we want to remove for specific host
+	lineToRemove := knownhosts.Line([]string{c.sshHost + ":" + c.sshPort}, c.serverPublicKey)
+
+	// We don't need to track cert-authority lines anymore as we always preserve them
+
+	// Filter out the line we want to remove
+	outputLines := []string{}
+	scanner := bufio.NewScanner(bytes.NewReader(existingContent))
+	for scanner.Scan() {
+		line := scanner.Text()
+
+		// Remove specific host entry
+		if line == lineToRemove {
+			continue
+		}
+
+		// We will preserve all lines, including certificate authority lines
+		// because they might be used by other containers
+
+		// Keep all lines, including CA entries which might be used by other containers
+		outputLines = append(outputLines, line)
+	}
+
+	// Safely write the updated content back to the file
+	if err := c.fs.SafeWriteFile(c.knownHostsPath, []byte(strings.Join(outputLines, "\n")), 0o644); err != nil {
+		return fmt.Errorf("couldn't safely write updated known_hosts to %s: %w", c.knownHostsPath, err)
+	}
+
+	return nil
+}
+
+// Cleanup removes the container-specific entries from the SSH configuration and known_hosts files.
+// It preserves the certificate authority entries that might be used by other containers.
+func (c *LocalSSHimmer) Cleanup() error {
+	if err := c.removeContainerFromSSHConfig(); err != nil {
+		return fmt.Errorf("couldn't remove container from ssh_config: %v\n", err)
+	}
+	if err := c.removeContainerFromKnownHosts(); err != nil {
+		return fmt.Errorf("couldn't remove container from known_hosts: %v\n", err)
+	}
+
+	return nil
+}
+
+func (c *LocalSSHimmer) removeContainerFromSSHConfig() error {
+	// Read the existing file contents
+	configData, err := c.fs.ReadFile(c.sshConfigPath)
+	if err != nil {
+		return fmt.Errorf("couldn't read ssh_config: %w", err)
+	}
+
+	cfg, err := ssh_config.Decode(bytes.NewReader(configData))
+	if err != nil {
+		return fmt.Errorf("couldn't decode ssh_config: %w", err)
+	}
+	cfg.Hosts = removeFromHosts(c.cntrName, cfg.Hosts)
+
+	if err := c.addSketchHostMatchIfMissing(cfg); err != nil {
+		return fmt.Errorf("couldn't add missing host match: %w", err)
+	}
+
+	cfgBytes, err := cfg.MarshalText()
+	if err != nil {
+		return fmt.Errorf("couldn't marshal ssh_config: %w", err)
+	}
+
+	// Safely write the updated configuration to file
+	if err := c.fs.SafeWriteFile(c.sshConfigPath, cfgBytes, 0o644); err != nil {
+		return fmt.Errorf("couldn't safely write ssh_config: %w", err)
+	}
+	return nil
+}
+
+// FileSystem represents a filesystem interface for testability
+type FileSystem interface {
+	Stat(name string) (fs.FileInfo, error)
+	Mkdir(name string, perm fs.FileMode) error
+	MkdirAll(name string, perm fs.FileMode) error
+	ReadFile(name string) ([]byte, error)
+	WriteFile(name string, data []byte, perm fs.FileMode) error
+	OpenFile(name string, flag int, perm fs.FileMode) (*os.File, error)
+	TempFile(dir, pattern string) (*os.File, error)
+	Rename(oldpath, newpath string) error
+	SafeWriteFile(name string, data []byte, perm fs.FileMode) error
+}
+
+func (fs *RealFileSystem) MkdirAll(name string, perm fs.FileMode) error {
+	return os.MkdirAll(name, perm)
+}
+
+// RealFileSystem is the default implementation of FileSystem that uses the OS
+type RealFileSystem struct{}
+
+func (fs *RealFileSystem) Stat(name string) (fs.FileInfo, error) {
+	return os.Stat(name)
+}
+
+func (fs *RealFileSystem) Mkdir(name string, perm fs.FileMode) error {
+	return os.Mkdir(name, perm)
+}
+
+func (fs *RealFileSystem) ReadFile(name string) ([]byte, error) {
+	return os.ReadFile(name)
+}
+
+func (fs *RealFileSystem) WriteFile(name string, data []byte, perm fs.FileMode) error {
+	return os.WriteFile(name, data, perm)
+}
+
+func (fs *RealFileSystem) OpenFile(name string, flag int, perm fs.FileMode) (*os.File, error) {
+	return os.OpenFile(name, flag, perm)
+}
+
+func (fs *RealFileSystem) TempFile(dir, pattern string) (*os.File, error) {
+	return os.CreateTemp(dir, pattern)
+}
+
+func (fs *RealFileSystem) Rename(oldpath, newpath string) error {
+	return os.Rename(oldpath, newpath)
+}
+
+// SafeWriteFile writes data to a temporary file, syncs to disk, creates a backup of the existing file if it exists,
+// and then renames the temporary file to the target file name.
+func (fs *RealFileSystem) SafeWriteFile(name string, data []byte, perm fs.FileMode) error {
+	// Get the directory from the target filename
+	dir := filepath.Dir(name)
+
+	// Create a temporary file in the same directory
+	tmpFile, err := fs.TempFile(dir, filepath.Base(name)+".*.tmp")
+	if err != nil {
+		return fmt.Errorf("couldn't create temporary file: %w", err)
+	}
+	tmpFilename := tmpFile.Name()
+	defer os.Remove(tmpFilename) // Clean up if we fail
+
+	// Write data to the temporary file
+	if _, err := tmpFile.Write(data); err != nil {
+		tmpFile.Close()
+		return fmt.Errorf("couldn't write to temporary file: %w", err)
+	}
+
+	// Sync to disk to ensure data is written
+	if err := tmpFile.Sync(); err != nil {
+		tmpFile.Close()
+		return fmt.Errorf("couldn't sync temporary file: %w", err)
+	}
+
+	// Close the temporary file
+	if err := tmpFile.Close(); err != nil {
+		return fmt.Errorf("couldn't close temporary file: %w", err)
+	}
+
+	// If the original file exists, create a backup
+	if _, err := fs.Stat(name); err == nil {
+		backupName := name + ".bak"
+		// Remove any existing backup
+		_ = os.Remove(backupName) // Ignore errors if the backup doesn't exist
+
+		// Create the backup
+		if err := fs.Rename(name, backupName); err != nil {
+			return fmt.Errorf("couldn't create backup file: %w", err)
+		}
+	}
+
+	// Rename the temporary file to the target file
+	if err := fs.Rename(tmpFilename, name); err != nil {
+		return fmt.Errorf("couldn't rename temporary file to target: %w", err)
+	}
+
+	// Set permissions on the new file
+	if err := os.Chmod(name, perm); err != nil {
+		return fmt.Errorf("couldn't set permissions on file: %w", err)
+	}
+
+	return nil
+}
+
+// KeyGenerator represents an interface for generating SSH keys for testability
+type KeyGenerator interface {
+	GenerateKeyPair() (ed25519.PrivateKey, ed25519.PublicKey, error)
+	ConvertToSSHPublicKey(publicKey ed25519.PublicKey) (ssh.PublicKey, error)
+}
+
+// RealKeyGenerator is the default implementation of KeyGenerator
+type RealKeyGenerator struct{}
+
+func (kg *RealKeyGenerator) GenerateKeyPair() (ed25519.PrivateKey, ed25519.PublicKey, error) {
+	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
+	return privateKey, publicKey, err
+}
+
+func (kg *RealKeyGenerator) ConvertToSSHPublicKey(publicKey ed25519.PublicKey) (ssh.PublicKey, error) {
+	return ssh.NewPublicKey(publicKey)
+}
+
+// CheckSSHReachability checks if the user's SSH config includes the Sketch SSH config file
+func CheckSSHReachability(cntrName string) error {
+	if err := checkSSHResolve(cntrName); err != nil {
+		return CheckForIncludeWithFS(&RealFileSystem{}, *bufio.NewReader(os.Stdin))
+	}
+	return nil
+}
+
+// setupContainerCA creates or loads the Container CA keys
+// Note: The setupContainerCA functionality has been incorporated directly into createHostCertificate
+// to simplify the certificate and CA creation process and avoid key format issues.
+
+// createHostCertificate creates a certificate for the host to authenticate to the container
+func (c *LocalSSHimmer) createHostCertificate(identityPath string) error {
+	// For testing purposes, create a minimal empty certificate
+	// This check will only be true in tests
+	if _, ok := c.kg.(interface{ IsMock() bool }); ok {
+		c.hostCertificate = []byte("test-host-certificate")
+		return nil
+	}
+
+	// Check if certificate already exists
+	if _, err := c.fs.Stat(c.hostCertPath); err == nil {
+		// Certificate exists, verify it's still valid
+		certBytes, err := c.fs.ReadFile(c.hostCertPath)
+		if err != nil {
+			return fmt.Errorf("error reading host certificate: %w", err)
+		}
+
+		// Parse certificate to check validity
+		pk, _, _, _, err := ssh.ParseAuthorizedKey(certBytes)
+		if err != nil {
+			// Invalid certificate, will regenerate
+		} else if cert, ok := pk.(*ssh.Certificate); ok {
+			// Check if certificate is still valid
+			if time.Now().Before(time.Unix(int64(cert.ValidBefore), 0)) &&
+				time.Now().After(time.Unix(int64(cert.ValidAfter), 0)) {
+				// Certificate is still valid
+				c.hostCertificate = certBytes // Store the valid certificate
+				return nil
+			}
+		}
+		// Otherwise, certificate is invalid or expired, regenerate it
+	}
+
+	// Load the private key to sign
+	privKeyBytes, err := c.fs.ReadFile(identityPath)
+	if err != nil {
+		return fmt.Errorf("error reading private key: %w", err)
+	}
+
+	// Parse the private key
+	signer, err := ssh.ParsePrivateKey(privKeyBytes)
+	if err != nil {
+		return fmt.Errorf("error parsing private key: %w", err)
+	}
+
+	// Create a new certificate
+	cert := &ssh.Certificate{
+		Key:             signer.PublicKey(),
+		Serial:          1,
+		CertType:        ssh.UserCert,
+		KeyId:           "sketch-host",
+		ValidPrincipals: []string{"root"},                               // Only valid for root user in container
+		ValidAfter:      uint64(time.Now().Add(-1 * time.Hour).Unix()),  // Valid from 1 hour ago
+		ValidBefore:     uint64(time.Now().Add(720 * time.Hour).Unix()), // Valid for 30 days
+		Permissions: ssh.Permissions{
+			CriticalOptions: map[string]string{
+				"source-address": "127.0.0.1,::1", // Only valid from localhost
+			},
+			Extensions: map[string]string{
+				"permit-pty":              "",
+				"permit-agent-forwarding": "",
+				"permit-port-forwarding":  "",
+			},
+		},
+	}
+
+	// Create a signer from the CA key for certificate signing
+	// The containerCA should already be a valid signer, but we'll create a fresh one for robustness
+	// Generate a fresh ed25519 key pair for the CA
+	caPrivate, caPublic, err := c.kg.GenerateKeyPair()
+	if err != nil {
+		return fmt.Errorf("error generating temporary CA key pair: %w", err)
+	}
+
+	// Create a signer from the private key
+	caSigner, err := ssh.NewSignerFromKey(caPrivate)
+	if err != nil {
+		return fmt.Errorf("error creating temporary CA signer: %w", err)
+	}
+
+	// Sign the certificate with the temporary CA
+	if err := cert.SignCert(rand.Reader, caSigner); err != nil {
+		return fmt.Errorf("error signing host certificate: %w", err)
+	}
+
+	// Marshal the certificate
+	certBytes := ssh.MarshalAuthorizedKey(cert)
+
+	// Store the certificate in memory
+	c.hostCertificate = certBytes
+
+	// Also update the CA public key for the known_hosts file
+	c.containerCAPublicKey, err = c.kg.ConvertToSSHPublicKey(caPublic)
+	if err != nil {
+		return fmt.Errorf("error converting temporary CA to SSH public key: %w", err)
+	}
+
+	// Write the certificate to file
+	if err := c.writeKeyToFile(certBytes, c.hostCertPath); err != nil {
+		return fmt.Errorf("error writing host certificate to file: %w", err)
+	}
+
+	// Also write the new CA public key
+	caPubKeyBytes := ssh.MarshalAuthorizedKey(c.containerCAPublicKey)
+	if err := c.writeKeyToFile(caPubKeyBytes, c.containerCAPath+".pub"); err != nil {
+		return fmt.Errorf("error writing CA public key to file: %w", err)
+	}
+
+	// And the CA private key
+	caPrivKeyPEM := encodePrivateKeyToPEM(caPrivate)
+	if err := c.writeKeyToFile(caPrivKeyPEM, c.containerCAPath); err != nil {
+		return fmt.Errorf("error writing CA private key to file: %w", err)
+	}
+
+	// Update the in-memory CA signer
+	c.containerCA = caSigner
+
+	return nil
+}