Browse Source

Initial version

aeris 3 months ago
commit
519ed99636
7 changed files with 379 additions and 0 deletions
  1. 1
    0
      .gitignore
  2. 176
    0
      agent.go
  3. 3
    0
      go.mod
  4. 3
    0
      go.sum
  5. 72
    0
      identity.go
  6. 79
    0
      main.go
  7. 45
    0
      utils.go

+ 1
- 0
.gitignore View File

@@ -0,0 +1 @@
1
+/ssh-ident

+ 176
- 0
agent.go View File

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

+ 3
- 0
go.mod View File

@@ -0,0 +1,3 @@
1
+module imirhil.fr/ssh-ident
2
+
3
+require gopkg.in/yaml.v2 v2.2.1

+ 3
- 0
go.sum View File

@@ -0,0 +1,3 @@
1
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
2
+gopkg.in/yaml.v2 v2.2.1 h1:mUhvW9EsL+naU5Q3cakzfE91YhliOondGd6ZrsDBHQE=
3
+gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

+ 72
- 0
identity.go View File

@@ -0,0 +1,72 @@
1
+package main
2
+
3
+import (
4
+	"os"
5
+	"path"
6
+	"regexp"
7
+	"strings"
8
+)
9
+
10
+type Identity struct {
11
+	Name string
12
+	Path string
13
+}
14
+
15
+func newIdentity(config Config, name string) Identity {
16
+	return Identity{
17
+		Name: name,
18
+		Path: "",
19
+	}
20
+}
21
+
22
+var hostRegexps = map[string]*regexp.Regexp{
23
+	"ssh": regexp.MustCompile("(?P<user>.*@)?(?P<host>.*)"),
24
+	"scp": regexp.MustCompile("(?P<user>.*@)?(?P<host>.*):(?P<file>.*)"),
25
+}
26
+
27
+func extractHost(prog string, args []string) string {
28
+	args = removeOptions(prog, args)
29
+	switch prog {
30
+	case "ssh":
31
+		if len(args) == 0 {
32
+			return ""
33
+		}
34
+		re := hostRegexps["ssh"]
35
+		arg := args[0]
36
+		p := matchParams(re, arg)
37
+		return p["host"]
38
+	case "scp":
39
+		re := hostRegexps["scp"]
40
+		for _, arg := range args {
41
+			if p := matchParams(re, arg); p != nil {
42
+				return p["host"]
43
+			}
44
+		}
45
+	}
46
+	return ""
47
+}
48
+
49
+func findIdentityName(config Config, prog string, args []string) string {
50
+	name := os.Getenv("SSH_IDENTITY")
51
+	if name != "" {
52
+		return name
53
+	}
54
+	identities := config.Identities
55
+	host := extractHost(prog, args)
56
+	for match, name := range identities {
57
+		if strings.Contains(host, match) {
58
+			return name
59
+		}
60
+	}
61
+
62
+	return config.DefaultIdentity
63
+}
64
+
65
+func FindIdentity(config Config, prog string, args []string) Identity {
66
+	name := findIdentityName(config, prog, args)
67
+	path := path.Join(config.IdentitiesDir, name)
68
+	return Identity{
69
+		Name: name,
70
+		Path: path,
71
+	}
72
+}

+ 79
- 0
main.go View File

@@ -0,0 +1,79 @@
1
+package main
2
+
3
+import (
4
+	"io/ioutil"
5
+	"os"
6
+	"os/user"
7
+	"path"
8
+	"strings"
9
+
10
+	"gopkg.in/yaml.v2"
11
+)
12
+
13
+type Config struct {
14
+	BinDir          string `yaml:"bin_dir"`
15
+	AgentsDir       string
16
+	IdentitiesDir   string
17
+	DefaultIdentity string `yaml:"default_identity"`
18
+	Identities      map[string]string
19
+}
20
+
21
+var sshOptions = map[string]map[string]string{
22
+	"ssh": {
23
+		"short": "1246AaconfiggKkMNnqsTtVvXxYy",
24
+		"long":  "bcDEeFIiJLlmOopQRSWw",
25
+	},
26
+	"scp": {
27
+		"short": "12346BCpqrv",
28
+		"long":  "cFiloPS",
29
+	},
30
+}
31
+
32
+func removeOptions(prog string, args []string) []string {
33
+	notOptions := []string{}
34
+	longOptions := sshOptions[prog]["long"]
35
+	long := false
36
+	for _, arg := range args {
37
+		if long {
38
+			long = false
39
+			continue
40
+		} else if strings.HasPrefix(arg, "-") {
41
+			last := arg[len(arg)-1:]
42
+			long = strings.Contains(longOptions, last)
43
+		} else {
44
+			notOptions = append(notOptions, arg)
45
+		}
46
+	}
47
+	return notOptions
48
+}
49
+
50
+func main() {
51
+	usr, err := user.Current()
52
+	fatal(err)
53
+	home := usr.HomeDir
54
+	sshDir := path.Join(home, ".ssh")
55
+	identitiesDir := path.Join(sshDir, "identities")
56
+
57
+	configFile := path.Join(identitiesDir, "config.yml")
58
+	data, err := ioutil.ReadFile(configFile)
59
+	fatal(err)
60
+	config := Config{}
61
+	err = yaml.Unmarshal(data, &config)
62
+	fatal(err)
63
+
64
+	config.IdentitiesDir = identitiesDir
65
+
66
+	prog := os.Getenv("SSH_BINARY")
67
+	if prog == "" {
68
+		prog = os.Args[0]
69
+		prog = path.Base(prog)
70
+	}
71
+	args := os.Args[1:]
72
+
73
+	identity := FindIdentity(config, prog, args)
74
+	agentsDir := path.Join(sshDir, "agents")
75
+	config.AgentsDir = agentsDir
76
+
77
+	agent := NewAgent(config, identity)
78
+	agent.Run(config, prog, args)
79
+}

+ 45
- 0
utils.go View File

@@ -0,0 +1,45 @@
1
+package main
2
+
3
+import (
4
+	"bytes"
5
+	"fmt"
6
+	"os"
7
+	"os/exec"
8
+	"regexp"
9
+)
10
+
11
+func fatal(err error) {
12
+	if err == nil {
13
+		return
14
+	}
15
+
16
+	fmt.Fprintf(os.Stderr, "%s\n", err.Error())
17
+	os.Exit(-1)
18
+}
19
+
20
+func matchParams(re *regexp.Regexp, str string) map[string]string {
21
+	match := re.FindStringSubmatch(str)
22
+	if match == nil {
23
+		return nil
24
+	}
25
+	params := make(map[string]string)
26
+	for i, name := range re.SubexpNames() {
27
+		params[name] = match[i]
28
+	}
29
+	return params
30
+}
31
+
32
+func capture3(cmd *exec.Cmd) ([]byte, []byte, error) {
33
+	var stdout, stderr bytes.Buffer
34
+	cmd.Stdout = &stdout
35
+	cmd.Stderr = &stderr
36
+
37
+	err := cmd.Run()
38
+	return stdout.Bytes(), stderr.Bytes(), err
39
+}
40
+
41
+func dumpEnv(env []string) {
42
+	for _, e := range env {
43
+		fmt.Println(e)
44
+	}
45
+}

Loading…
Cancel
Save