ssh-ident/agent.go

package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"syscall"
)

type Agent struct {
	Identity Identity
	Path     string
	Env      []string
}

func (a *Agent) getEnv() []string {
	env := os.Environ()
	env = append(env, a.Env...)
	return env
}

func (a *Agent) loadEnv() {
	d, err := ioutil.ReadFile(a.envFile())
	fatal(err)
	properties := string(d)
	env := extractEnv(properties)
	a.Env = env
}

func (a *Agent) loadKeys() {
	path := path.Join(a.Identity.Path, "id_*")
	keys, err := filepath.Glob(path)
	fatal(err)
	sort.Strings(keys)

	fmt.Fprintf(os.Stderr, "Load private key:\n")
	privateKeys := []string{}
	for _, key := range keys {
		if strings.HasSuffix(key, ".pub") {
			continue
		}
		fmt.Fprintf(os.Stderr, "  %s\n", key)
		privateKeys = append(privateKeys, key)
	}

	cmd := exec.Command("ssh-add", privateKeys...)
	cmd.Env = a.getEnv()

	_, e, err := capture3(cmd)
	if err != nil {
		os.Stderr.Write(e)
		fatal(err)
	}
}

func (a *Agent) envFile() string {
	return a.Path + ".env"
}

var propertyRegex = regexp.MustCompile(`([^=]+)=([^;]+); export .*;`)

func extractEnv(str string) []string {
	properties := strings.Split(str, "\n")
	env := []string{}
	for _, property := range properties {
		match := propertyRegex.FindStringSubmatch(property)
		if match != nil {
			name := match[1]
			value := match[2]
			property = fmt.Sprintf("%s=%s", name, value)
			env = append(env, property)
		}
	}
	return env
}

func (a *Agent) start() {
	fmt.Fprintf(os.Stderr, "Start new agent for identity %s\n", a.Identity.Name)
	sock := a.Path + ".sock"

	syscall.Unlink(sock)
	cmd := exec.Command("ssh-agent", "-a", sock)
	o, e, err := capture3(cmd)
	if err != nil {
		os.Stderr.Write(e)
		fatal(err)
	}

	properties := string(o)
	env := extractEnv(properties)
	a.Env = env

	err = ioutil.WriteFile(a.envFile(), o, 0600)
	fatal(err)

	a.loadKeys()
}

func (a *Agent) getPid() int {
	d, err := ioutil.ReadFile(a.envFile())
	fatal(err)
	properties := string(d)
	env := strings.Split(properties, "\n")
	for _, property := range env {
		match := propertyRegex.FindStringSubmatch(property)
		if match != nil {
			name := match[1]
			if name == "SSH_AGENT_PID" {
				value := match[2]
				pid, err := strconv.Atoi(value)
				fatal(err)
				return pid
			}
		}
	}
	return -1
}

func (a *Agent) init() {
	if _, err := os.Stat(a.envFile()); os.IsNotExist(err) {
		a.start()
		return
	}

	pid := a.getPid()
	if pid <= 0 {
		a.start()
		return
	}

	proc, err := os.FindProcess(pid)
	if err == nil {
		err = proc.Signal(syscall.Signal(0))
		if err != nil {
			a.start()
			return
		}
	}

	a.loadEnv()
}

func NewAgent(config Config, identity Identity) Agent {
	p := path.Join(config.AgentsDir, identity.Name)
	agent := Agent{
		Identity: identity,
		Path:     p,
	}
	agent.init()
	return agent
}

func (a *Agent) Run(config Config, prog string, args []string) {
	identity := a.Identity
	fmt.Fprintf(os.Stderr, "\033[1;41m[%s]\033[0m %s %s\n", identity.Name, prog, strings.Join(args, " "))
	exe := path.Join(config.BinDir, prog)
	if _, err := os.Stat(exe); os.IsNotExist(err) {
		fatal(fmt.Errorf("%s: no such file or directory", exe))
	}

	sshConfig := path.Join(identity.Path, "config")
	_, err := os.Stat(exe)
	if err == nil {
		args = append([]string{"-F", sshConfig}, args...)
	}
	args = append([]string{prog}, args...)

	env := a.getEnv()
	syscall.Exec(exe, args, env)
}