From 519ed9963693ac7dbfc51d62851989d11c0812f9 Mon Sep 17 00:00:00 2001 From: aeris Date: Sun, 30 Sep 2018 00:43:41 +0200 Subject: [PATCH] Initial version --- .gitignore | 1 + agent.go | 176 ++++++++++++++++++++++++++++++++++++++++++++++++++++ go.mod | 3 + go.sum | 3 + identity.go | 72 +++++++++++++++++++++ main.go | 79 +++++++++++++++++++++++ utils.go | 45 ++++++++++++++ 7 files changed, 379 insertions(+) create mode 100644 .gitignore create mode 100644 agent.go create mode 100644 go.mod create mode 100644 go.sum create mode 100644 identity.go create mode 100644 main.go create mode 100644 utils.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3e318b3 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/ssh-ident diff --git a/agent.go b/agent.go new file mode 100644 index 0000000..194b8ea --- /dev/null +++ b/agent.go @@ -0,0 +1,176 @@ +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) +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..e635475 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module imirhil.fr/ssh-ident + +require gopkg.in/yaml.v2 v2.2.1 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..99fc7ca --- /dev/null +++ b/go.sum @@ -0,0 +1,3 @@ +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= diff --git a/identity.go b/identity.go new file mode 100644 index 0000000..a8d6dce --- /dev/null +++ b/identity.go @@ -0,0 +1,72 @@ +package main + +import ( + "os" + "path" + "regexp" + "strings" +) + +type Identity struct { + Name string + Path string +} + +func newIdentity(config Config, name string) Identity { + return Identity{ + Name: name, + Path: "", + } +} + +var hostRegexps = map[string]*regexp.Regexp{ + "ssh": regexp.MustCompile("(?P.*@)?(?P.*)"), + "scp": regexp.MustCompile("(?P.*@)?(?P.*):(?P.*)"), +} + +func extractHost(prog string, args []string) string { + args = removeOptions(prog, args) + switch prog { + case "ssh": + if len(args) == 0 { + return "" + } + re := hostRegexps["ssh"] + arg := args[0] + p := matchParams(re, arg) + return p["host"] + case "scp": + re := hostRegexps["scp"] + for _, arg := range args { + if p := matchParams(re, arg); p != nil { + return p["host"] + } + } + } + return "" +} + +func findIdentityName(config Config, prog string, args []string) string { + name := os.Getenv("SSH_IDENTITY") + if name != "" { + return name + } + identities := config.Identities + host := extractHost(prog, args) + for match, name := range identities { + if strings.Contains(host, match) { + return name + } + } + + return config.DefaultIdentity +} + +func FindIdentity(config Config, prog string, args []string) Identity { + name := findIdentityName(config, prog, args) + path := path.Join(config.IdentitiesDir, name) + return Identity{ + Name: name, + Path: path, + } +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..0c39097 --- /dev/null +++ b/main.go @@ -0,0 +1,79 @@ +package main + +import ( + "io/ioutil" + "os" + "os/user" + "path" + "strings" + + "gopkg.in/yaml.v2" +) + +type Config struct { + BinDir string `yaml:"bin_dir"` + AgentsDir string + IdentitiesDir string + DefaultIdentity string `yaml:"default_identity"` + Identities map[string]string +} + +var sshOptions = map[string]map[string]string{ + "ssh": { + "short": "1246AaconfiggKkMNnqsTtVvXxYy", + "long": "bcDEeFIiJLlmOopQRSWw", + }, + "scp": { + "short": "12346BCpqrv", + "long": "cFiloPS", + }, +} + +func removeOptions(prog string, args []string) []string { + notOptions := []string{} + longOptions := sshOptions[prog]["long"] + long := false + for _, arg := range args { + if long { + long = false + continue + } else if strings.HasPrefix(arg, "-") { + last := arg[len(arg)-1:] + long = strings.Contains(longOptions, last) + } else { + notOptions = append(notOptions, arg) + } + } + return notOptions +} + +func main() { + usr, err := user.Current() + fatal(err) + home := usr.HomeDir + sshDir := path.Join(home, ".ssh") + identitiesDir := path.Join(sshDir, "identities") + + configFile := path.Join(identitiesDir, "config.yml") + data, err := ioutil.ReadFile(configFile) + fatal(err) + config := Config{} + err = yaml.Unmarshal(data, &config) + fatal(err) + + config.IdentitiesDir = identitiesDir + + prog := os.Getenv("SSH_BINARY") + if prog == "" { + prog = os.Args[0] + prog = path.Base(prog) + } + args := os.Args[1:] + + identity := FindIdentity(config, prog, args) + agentsDir := path.Join(sshDir, "agents") + config.AgentsDir = agentsDir + + agent := NewAgent(config, identity) + agent.Run(config, prog, args) +} diff --git a/utils.go b/utils.go new file mode 100644 index 0000000..b319079 --- /dev/null +++ b/utils.go @@ -0,0 +1,45 @@ +package main + +import ( + "bytes" + "fmt" + "os" + "os/exec" + "regexp" +) + +func fatal(err error) { + if err == nil { + return + } + + fmt.Fprintf(os.Stderr, "%s\n", err.Error()) + os.Exit(-1) +} + +func matchParams(re *regexp.Regexp, str string) map[string]string { + match := re.FindStringSubmatch(str) + if match == nil { + return nil + } + params := make(map[string]string) + for i, name := range re.SubexpNames() { + params[name] = match[i] + } + return params +} + +func capture3(cmd *exec.Cmd) ([]byte, []byte, error) { + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + return stdout.Bytes(), stderr.Bytes(), err +} + +func dumpEnv(env []string) { + for _, e := range env { + fmt.Println(e) + } +}