-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.go
155 lines (120 loc) · 2.85 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
package main
import (
"bytes"
"encoding/json"
"fmt"
"io/ioutil"
"net/http"
"os"
"os/user"
"strings"
"text/template"
)
// Server struct holds Droplet data for simpler syntax.
type Server struct {
Hostname string
IP string
}
type Droplets struct {
Droplets []Droplet
}
type Droplet struct {
Name string `json:"name"`
Networks Network
}
type Network struct {
Version []ProtocolStruct `json:"v4"`
}
type ProtocolStruct struct {
IP string `json:"ip_address"`
}
const templ = `
#START dossh{{range .}}
Host {{.Hostname}}
Hostname {{.IP}}
Port 22
User root
IdentitiesOnly yes
{{end}}
#END dossh
`
// getServers connects to DO API and fetches droplet list.
func getServers(token string) []byte {
req, err := http.NewRequest("GET", "https://api.digitalocean.com/v2/droplets", nil)
req.Header.Set("Authorization", "Bearer "+token)
if err != nil {
panic(err)
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
panic(err)
}
defer resp.Body.Close()
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
panic(err)
}
return body
}
// parseServers unmarshals the JSON data from DO API and returns []Server struct
func parseServers(body []byte) []Server {
var d Droplets
err := json.Unmarshal(body, &d)
if err != nil {
panic(err)
}
var servers []Server
for _, droplet := range d.Droplets {
servers = append(servers, Server{droplet.Name, droplet.Networks.Version[0].IP})
}
return servers
}
// templateFrom creates a new text template from given []Server struct. It also fetches the old .ssh/config file
// and parses it so that the old dossh values are removed. It then inserts the new values to the old file template
// and finally returns the completed byte array.
func templateFrom(servers []Server) []byte {
t, err := template.New("config").Parse(templ)
if err != nil {
panic(err)
}
var doc bytes.Buffer
err = t.Execute(&doc, servers)
if err != nil {
panic(err)
}
usr, err := user.Current()
if err != nil {
panic(err)
}
file, err := ioutil.ReadFile(usr.HomeDir + "/.ssh/config")
if err != nil {
panic(err)
}
firstRun := strings.Contains(string(file), "#START dossh")
if !firstRun {
return []byte(strings.TrimSpace(string(file) + doc.String()))
}
config := string(file)
e := strings.Split(config, "#START dossh")
h := strings.Split(config, "#END dossh")
newconfig := e[0] + h[1]
return []byte(strings.TrimSpace(newconfig) + doc.String())
}
// saveconfig saves the new .ssh/config file with file permissions 0612.
func saveconfig(document []byte) {
usr, err := user.Current()
if err != nil {
panic(err)
}
err = ioutil.WriteFile(usr.HomeDir+"/.ssh/config", document, 0612)
if err != nil {
panic(err)
}
}
func main() {
body := getServers(os.Args[1])
servers := parseServers(body)
document := templateFrom(servers)
saveconfig(document)
fmt.Println("SSH config file updated successfully.")
}