SSH client with identities support
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.

agent.go 3.3KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  1. package main
  2. import (
  3. "fmt"
  4. "io/ioutil"
  5. "os"
  6. "os/exec"
  7. "path"
  8. "path/filepath"
  9. "regexp"
  10. "sort"
  11. "strconv"
  12. "strings"
  13. "syscall"
  14. )
  15. type Agent struct {
  16. Identity Identity
  17. Path string
  18. Env []string
  19. }
  20. func (a *Agent) getEnv() []string {
  21. env := os.Environ()
  22. env = append(env, a.Env...)
  23. return env
  24. }
  25. func (a *Agent) loadEnv() {
  26. d, err := ioutil.ReadFile(a.envFile())
  27. fatal(err)
  28. properties := string(d)
  29. env := extractEnv(properties)
  30. a.Env = env
  31. }
  32. func (a *Agent) loadKeys() {
  33. path := path.Join(a.Identity.Path, "id_*")
  34. keys, err := filepath.Glob(path)
  35. fatal(err)
  36. sort.Strings(keys)
  37. fmt.Fprintf(os.Stderr, "Load private key:\n")
  38. privateKeys := []string{}
  39. for _, key := range keys {
  40. if strings.HasSuffix(key, ".pub") {
  41. continue
  42. }
  43. fmt.Fprintf(os.Stderr, " %s\n", key)
  44. privateKeys = append(privateKeys, key)
  45. }
  46. cmd := exec.Command("ssh-add", privateKeys...)
  47. cmd.Env = a.getEnv()
  48. _, e, err := capture3(cmd)
  49. if err != nil {
  50. os.Stderr.Write(e)
  51. fatal(err)
  52. }
  53. }
  54. func (a *Agent) envFile() string {
  55. return a.Path + ".env"
  56. }
  57. var propertyRegex = regexp.MustCompile(`([^=]+)=([^;]+); export .*;`)
  58. func extractEnv(str string) []string {
  59. properties := strings.Split(str, "\n")
  60. env := []string{}
  61. for _, property := range properties {
  62. match := propertyRegex.FindStringSubmatch(property)
  63. if match != nil {
  64. name := match[1]
  65. value := match[2]
  66. property = fmt.Sprintf("%s=%s", name, value)
  67. env = append(env, property)
  68. }
  69. }
  70. return env
  71. }
  72. func (a *Agent) start() {
  73. fmt.Fprintf(os.Stderr, "Start new agent for identity %s\n", a.Identity.Name)
  74. sock := a.Path + ".sock"
  75. cmd := exec.Command("ssh-agent", "-a", sock)
  76. o, e, err := capture3(cmd)
  77. if err != nil {
  78. os.Stderr.Write(e)
  79. fatal(err)
  80. }
  81. properties := string(o)
  82. env := extractEnv(properties)
  83. a.Env = env
  84. err = ioutil.WriteFile(a.envFile(), o, 0600)
  85. fatal(err)
  86. a.loadKeys()
  87. }
  88. func (a *Agent) getPid() int {
  89. d, err := ioutil.ReadFile(a.envFile())
  90. fatal(err)
  91. properties := string(d)
  92. env := strings.Split(properties, "\n")
  93. for _, property := range env {
  94. match := propertyRegex.FindStringSubmatch(property)
  95. if match != nil {
  96. name := match[1]
  97. if name == "SSH_AGENT_PID" {
  98. value := match[2]
  99. pid, err := strconv.Atoi(value)
  100. fatal(err)
  101. return pid
  102. }
  103. }
  104. }
  105. return -1
  106. }
  107. func (a *Agent) init() {
  108. if _, err := os.Stat(a.envFile()); os.IsNotExist(err) {
  109. a.start()
  110. return
  111. }
  112. pid := a.getPid()
  113. if pid <= 0 {
  114. a.start()
  115. return
  116. }
  117. proc, err := os.FindProcess(pid)
  118. if err == nil {
  119. err = proc.Signal(syscall.Signal(0))
  120. if err != nil {
  121. a.start()
  122. return
  123. }
  124. }
  125. a.loadEnv()
  126. }
  127. func NewAgent(config Config, identity Identity) Agent {
  128. p := path.Join(config.AgentsDir, identity.Name)
  129. agent := Agent{
  130. Identity: identity,
  131. Path: p,
  132. }
  133. agent.init()
  134. return agent
  135. }
  136. func (a *Agent) Run(config Config, prog string, args []string) {
  137. identity := a.Identity
  138. fmt.Fprintf(os.Stderr, "\033[1;41m[%s]\033[0m %s %s\n", identity.Name, prog, strings.Join(args, " "))
  139. exe := path.Join(config.BinDir, prog)
  140. if _, err := os.Stat(exe); os.IsNotExist(err) {
  141. fatal(fmt.Errorf("%s: no such file or directory", exe))
  142. }
  143. sshConfig := path.Join(identity.Path, "config")
  144. _, err := os.Stat(exe)
  145. if err == nil {
  146. args = append([]string{"-F", sshConfig}, args...)
  147. }
  148. args = append([]string{prog}, args...)
  149. env := a.getEnv()
  150. syscall.Exec(exe, args, env)
  151. }