ssh_config_tool.go

package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"sort"
	"strings"

	"github.com/google/shlex"
	"github.com/pmezard/go-difflib/difflib"
)

type SSHCommand struct {
	Hostname string
	User     string
	Options  map[string]string
}

func parseSSHCommand(sshCommand string) (*SSHCommand, error) {
	tokens, err := shlex.Split(sshCommand)
	if err != nil {
		return nil, err
	}

	hostname := ""
	user := ""
	options := make(map[string]string)

	for i := 0; i < len(tokens); i++ {
		switch tokens[i] {
		case "ssh":
			// skip
		case "-p":
			i++
			options["Port"] = tokens[i]
		case "-i":
			i++
			options["IdentityFile"] = tokens[i]
		case "-o":
			i++
			opt := strings.SplitN(tokens[i], "=", 2)
			options[opt[0]] = opt[1]
		case "-X":
			options["ForwardX11"] = "yes"
		case "-A":
			options["ForwardAgent"] = "yes"
		case "-L":
			i++
			localForwardParts := strings.SplitN(tokens[i], ":", 3)
			options["LocalForward"] = fmt.Sprintf(
				"%s %s:%s",
				localForwardParts[0],
				localForwardParts[1],
				localForwardParts[2])
		case "-J":
			i++
			options["ProxyJump"] = tokens[i]
		default:
			if strings.Contains(tokens[i], "@") {
				parts := strings.Split(tokens[i], "@")
				user = parts[0]
				hostname = parts[1]
			} else {
				hostname = tokens[i]
			}
		}
	}

	return &SSHCommand{
		Hostname: hostname,
		User:     user,
		Options:  options,
	}, nil
}

func sshCommandToConfigEntry(sshCommand string) string {

	sshCmd, _ := parseSSHCommand(sshCommand)

	var b strings.Builder
	fmt.Fprintf(&b, "Host %s\n", sshCmd.Hostname)
	fmt.Fprintf(&b, "  HostName %s\n", sshCmd.Hostname)
	fmt.Fprintf(&b, "  User %s\n", sshCmd.User)

	// Get sorted keys
	keys := make([]string, 0, len(sshCmd.Options))
	for k := range sshCmd.Options {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	// Iterate through sorted keys for deterministic output
	for _, k := range keys {
		fmt.Fprintf(&b, "  %s %s\n", k, sshCmd.Options[k])
	}

	return b.String()
}

func calculateDiff(existingConfig, newConfig string) string {
	diff := difflib.UnifiedDiff{
		A:        difflib.SplitLines(existingConfig),
		B:        difflib.SplitLines(newConfig),
		FromFile: "Existing Config",
		ToFile:   "New Config",
		Context:  3,
	}
	diffStr, _ := difflib.GetUnifiedDiffString(diff)
	return diffStr
}

func splitSSHConfig(force bool, configPath string, dryRun bool) {
	usr, _ := user.Current()
	if configPath == "" {
		configPath = filepath.Join(usr.HomeDir, ".ssh", "config")
	}

	configData, err := ioutil.ReadFile(configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error reading SSH config file: %v\n", err)
		os.Exit(1)
	}
	hosts := parseSSHConfig(string(configData))

	sshConfigD := filepath.Join(usr.HomeDir, ".ssh", "config.d")
	if !dryRun {
		_ = os.MkdirAll(sshConfigD, 0755)
	}

	createNewDotSSHConfig := true
	for hostName, hostLines := range hosts {
		if hostName == "" {
			continue
		}
		hostConfigFile := filepath.Join(sshConfigD, hostName+".conf")
		hostConfig := strings.Join(hostLines, "\n")

		overwrite := force
		if !force {
			if _, err := os.Stat(hostConfigFile); err == nil {
				existingData, _ := ioutil.ReadFile(hostConfigFile)
				existingConfig := string(existingData)

				if existingConfig != hostConfig {
					diffStr := calculateDiff(existingConfig, hostConfig)
					fmt.Fprintf(os.Stderr, "Differences in host %s configuration:\n%s\n", hostName, diffStr)
					fmt.Fprintf(os.Stderr, "Not modifying %s - use --force to overwrite\n", hostConfigFile)
					createNewDotSSHConfig = false
				}
				continue
			} else {
				overwrite = true
			}
		}

		if overwrite && !dryRun {
			err = ioutil.WriteFile(hostConfigFile, []byte(hostConfig), 0644)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error writing host config file: %v\n", err)
				os.Exit(1)
			}
		} else if dryRun {
			fmt.Fprintf(os.Stderr, "Dry run: would write file %s with content:\n%s\n", hostConfigFile, hostConfig)
		}
	}

	if !dryRun {
		backupSSHConfig()
	}

	if ((createNewDotSSHConfig && len(hosts) > 0) || force) && !dryRun {
		newConfig := "# This file is generated by sshcfgtool - see ~/.ssh/config.d for host configurations.\n"
		newConfig += "Include " + filepath.Join(usr.HomeDir, ".ssh", "config.d", "*")
		err = ioutil.WriteFile(configPath, []byte(newConfig), 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error writing new SSH config file: %v\n", err)
			os.Exit(1)
		}
	} else if dryRun {
		fmt.Fprintf(os.Stderr, "Dry run: would update SSH config file %s\n", configPath)
	}
}

func copyFile(srcPath, dstPath string) error {
	srcFile, err := os.Open(srcPath)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	dstFile, err := os.Create(dstPath)
	if err != nil {
		return err
	}
	defer dstFile.Close()

	_, err = io.Copy(dstFile, srcFile)
	if err != nil {
		return err
	}

	srcInfo, err := srcFile.Stat()
	if err != nil {
		return err
	}

	return os.Chmod(dstPath, srcInfo.Mode())
}

func backupSSHConfig() {
	usr, _ := user.Current()
	sshConfigPath := filepath.Join(usr.HomeDir, ".ssh", "config")
	backupBase := filepath.Join(usr.HomeDir, ".ssh", "config.backup")

	var backupPath string
	for i := 0; ; i++ {
		backupPath = backupBase
		if i > 0 {
			backupPath = fmt.Sprintf("%s.%d", backupBase, i)
		}
		if _, err := os.Stat(backupPath); os.IsNotExist(err) {
			break
		}
	}

	//err := exec.Command("cp", "-p", sshConfigPath, backupPath).Run()
	err := copyFile(sshConfigPath, backupPath)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error creating backup: %v\n", err)
		os.Exit(1)
	}
}

func parseSSHConfig(configData string) map[string][]string {
	hosts := make(map[string][]string)

	scanner := bufio.NewScanner(strings.NewReader(configData))
	hostName := ""

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())

		if line == "" {
			continue
		}

		if strings.HasPrefix(line, "Host ") {
			hostName = strings.TrimSpace(strings.TrimPrefix(line, "Host"))
		} else if hostName != "" {
			hosts[hostName] = append(hosts[hostName], "  "+line)
		}
	}

	return hosts
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "Usage: %s [subcommand] [options]\n\n", os.Args[0])
		fmt.Fprintln(flag.CommandLine.Output(), "Subcommands:")
		fmt.Fprintln(flag.CommandLine.Output(), "  split\tSplit SSH config into separate files")
		fmt.Fprintln(flag.CommandLine.Output(), "  translate\tTranslate SSH command to SSH config entry")
		flag.PrintDefaults()
	}

	force := flag.Bool("force", false, "Force overwriting existing files in split subcommand")

	var dryRun bool
	flag.BoolVar(&dryRun, "dry-run", false, "Print intentions but don't actually change any files")
	flag.BoolVar(&dryRun, "n", false, "Print intentions but don't actually change any files (shorthand)")

	var configPath string
	flag.StringVar(&configPath, "config", "", "Path to the SSH config file")
	flag.StringVar(&configPath, "c", "", "Path to the SSH config file (shorthand)")

	flag.Parse()

	if flag.NArg() < 1 {
		flag.Usage()
		os.Exit(1)
	}

	subcommand := flag.Arg(0)
	switch subcommand {
	case "split":
		splitSSHConfig(*force, configPath, dryRun)
	case "translate":
		sshCommand := flag.Args()[1:]
		configEntry := sshCommandToConfigEntry(strings.Join(sshCommand, " "))
		fmt.Println(configEntry)
	default:
		flag.Usage()
		os.Exit(1)
	}
}