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.

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