123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177 |
- 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)
- }
|