You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
176 lines
3.3 KiB
176 lines
3.3 KiB
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"
|
|
|
|
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)
|
|
}
|
|
|