Skip to content

Commit f6ab002

Browse files
committed
Refactored, added DigitalOcean support
1 parent 5887d19 commit f6ab002

7 files changed

Lines changed: 217 additions & 106 deletions

File tree

Makefile

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
build-x64:
2+
GOOS=linux GOARCH=amd64 CGO_ENABLED=0 go build
3+
tar -czf cloud_ssh_x64.tar.gz cloud-ssh
4+
rm cloud-ssh
5+
6+
build-x86:
7+
GOOS=linux GOARCH=i386 CGO_ENABLED=0 go build
8+
tar -czf cloud_ssh_x86.tar.gz cloud-ssh
9+
rm cloud-ssh
10+
11+
build-macos:
12+
GOOS=darwin GOARCH=amd64 CGO_ENABLED=0 go build
13+
tar -czf cloud_ssh_macosx.tar.gz cloud-ssh
14+
rm cloud-ssh
15+
16+
all:
17+
build-macos
18+
build-x86
19+
build-x64

arguments_parser.go

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package main
2+
3+
import (
4+
"strings"
5+
)
6+
7+
func splitHostname(str string) (user string, hostname string) {
8+
if arr := strings.Split(str, "@"); len(arr) > 1 {
9+
return arr[0], arr[1]
10+
} else {
11+
return "", str
12+
}
13+
}
14+
15+
func joinHostname(user string, hostname string) string {
16+
if user != "" {
17+
return user + "@" + hostname
18+
} else {
19+
return hostname
20+
}
21+
}
22+
23+
// Go though arguments, and find one with host
24+
func getTargetHostname(args []string) (user string, hostname string, arg_idx int) {
25+
for idx, arg := range args {
26+
if !strings.HasPrefix(arg, "-") {
27+
if idx == 0 {
28+
hostname = arg
29+
arg_idx = idx
30+
break
31+
} else {
32+
if !strings.HasPrefix(args[idx-1], "-") {
33+
hostname = arg
34+
arg_idx = idx
35+
break
36+
}
37+
}
38+
}
39+
}
40+
41+
user, hostname = splitHostname(hostname)
42+
43+
return
44+
}

provider/aws.go renamed to aws.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
1-
package provider
1+
package main
22

33
import (
44
"log"
55
"launchpad.net/goamz/aws"
66
"launchpad.net/goamz/ec2"
77
)
88

9-
func GetEC2Instances(config map[string]string) (instances Instances) {
9+
func getEC2Instances(config map[string]string) (instances Instances) {
1010
instances = make(Instances)
1111

1212
if _,ok := config["access_key"]; !ok {

config.go

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
package main
2+
3+
import (
4+
"os"
5+
"fmt"
6+
"io/ioutil"
7+
"log"
8+
"github.com/go-yaml/yaml"
9+
"runtime"
10+
)
11+
12+
type Config map[string]StrMap
13+
type StrMap map[string]string
14+
15+
func userHomeDir() string {
16+
if runtime.GOOS == "windows" {
17+
home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
18+
if home == "" {
19+
home = os.Getenv("USERPROFILE")
20+
}
21+
return home
22+
}
23+
return os.Getenv("HOME")
24+
}
25+
26+
func readConfig() (config Config) {
27+
config = make(Config)
28+
29+
prefferedPaths := []string{
30+
"./cloud-ssh.yaml",
31+
userHomeDir() + "/.ssh/cloud-ssh.yaml",
32+
"/etc/cloud-ssh.yaml",
33+
}
34+
35+
var content []byte
36+
37+
for _, path := range prefferedPaths {
38+
if _, err := os.Stat(path); err == nil {
39+
fmt.Println("Found config:", path)
40+
content, err = ioutil.ReadFile(path)
41+
42+
if err != nil {
43+
log.Fatal("Error while reading config: ", err)
44+
}
45+
}
46+
}
47+
48+
if os.Getenv("AWS_ACCESS_KEY_ID") != "" && os.Getenv("AWS_SECRET_ACCESS_KEY") != "" {
49+
config["default"] = make(StrMap)
50+
config["default"]["access_key"] = os.Getenv("AWS_ACCESS_KEY_ID")
51+
config["default"]["secret_key"] = os.Getenv("AWS_SECRET_ACCESS_KEY")
52+
config["default"]["region"] = os.Getenv("AWS_REGION")
53+
config["default"]["provider"] = "aws"
54+
}
55+
56+
if len(content) == 0 {
57+
if len(config) == 0 {
58+
fmt.Println("Can't find any configuration or ENV variables. Check http://github.com/buger/cloud-ssh for documentation.")
59+
}
60+
return
61+
} else if err := yaml.Unmarshal(content, &config); err != nil {
62+
log.Fatal(err)
63+
}
64+
65+
return
66+
}

diginal_ocean.go

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package main
2+
3+
import (
4+
"log"
5+
"github.com/jmoiron/jsonq"
6+
"encoding/json"
7+
"net/http"
8+
)
9+
10+
func getDigitalOceanInstances(config map[string]string) (instances Instances) {
11+
instances = make(Instances)
12+
13+
if _,ok := config["client_id"]; !ok {
14+
log.Fatal("Missing client_id for ", config["name"], " DigitalOcean cloud")
15+
}
16+
17+
if _,ok := config["api_key"]; !ok {
18+
log.Fatal("Missing api_key for ", config["name"], " DigitalOcean cloud")
19+
}
20+
21+
22+
resp, err := http.Get("https://api.digitalocean.com/droplets/?client_id="+config["client_id"]+"&api_key="+config["api_key"])
23+
24+
if err != nil {
25+
log.Println("DigitalOcean API error:", err)
26+
return
27+
}
28+
29+
defer resp.Body.Close()
30+
31+
data := map[string]interface{}{}
32+
dec := json.NewDecoder(resp.Body)
33+
dec.Decode(&data)
34+
jq := jsonq.NewQuery(data)
35+
36+
status, err := jq.String("status")
37+
38+
if status == "ERROR" {
39+
err_msg, _ := jq.String("error_message")
40+
41+
log.Println("DigitalOcean API error: ", err_msg)
42+
return
43+
}
44+
45+
droplets, err := jq.ArrayOfObjects("droplets")
46+
47+
if err != nil {
48+
log.Println(err)
49+
return
50+
}
51+
52+
for _, droplet := range droplets {
53+
instances[droplet["ip_address"].(string)] = []Tag{
54+
Tag{"Name", droplet["name"].(string),},
55+
}
56+
}
57+
58+
return
59+
}

main.go

Lines changed: 27 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -2,137 +2,67 @@ package main
22

33
import (
44
"fmt"
5-
"github.com/buger/cloud-ssh/provider"
6-
"github.com/go-yaml/yaml"
7-
"io/ioutil"
85
"log"
96
"os"
107
"os/exec"
118
"regexp"
12-
"runtime"
139
"strconv"
1410
"strings"
11+
"sync"
1512
)
1613

17-
type CloudInstances map[string]provider.Instances
18-
type StrMap map[string]string
19-
type Config map[string]StrMap
20-
21-
func splitHostname(str string) (user string, hostname string) {
22-
if arr := strings.Split(str, "@"); len(arr) > 1 {
23-
return arr[0], arr[1]
24-
} else {
25-
return "", str
26-
}
14+
type Tag struct {
15+
Name, Value string
2716
}
2817

29-
func joinHostname(user string, hostname string) string {
30-
if user != "" {
31-
return user + "@" + hostname
32-
} else {
33-
return hostname
34-
}
35-
}
36-
37-
func getTargetHostname(args []string) (user string, hostname string, arg_idx int) {
38-
for idx, arg := range args {
39-
if !strings.HasPrefix(arg, "-") {
40-
if idx == 0 {
41-
hostname = arg
42-
arg_idx = idx
43-
break
44-
} else {
45-
if !strings.HasPrefix(args[idx-1], "-") {
46-
hostname = arg
47-
arg_idx = idx
48-
break
49-
}
50-
}
51-
}
52-
}
53-
54-
user, hostname = splitHostname(hostname)
18+
type Instances map[string][]Tag
19+
type CloudInstances map[string]Instances
5520

56-
return
57-
}
5821

5922
func getInstances(config Config) (clouds CloudInstances) {
60-
clouds = make(CloudInstances)
23+
clouds = make(CloudInstances)
24+
25+
var wg sync.WaitGroup
26+
var mux sync.RWMutex
6127

6228
for name, cfg := range config {
6329
for k, v := range cfg {
6430
cfg["name"] = name
31+
cfg["provider"] = k
6532

6633
if k == "provider" {
6734
switch v {
6835
case "aws":
69-
clouds[name] = provider.GetEC2Instances(cfg)
36+
wg.Add(1)
37+
go func(name string, cfg StrMap){
38+
mux.Lock()
39+
clouds[name] = getEC2Instances(cfg)
40+
mux.Unlock()
41+
wg.Done()
42+
}(name, cfg)
43+
case "digital_ocean":
44+
wg.Add(1)
45+
go func(name string, cfg StrMap){
46+
mux.Lock()
47+
clouds[name] = getDigitalOceanInstances(cfg)
48+
mux.Unlock()
49+
wg.Done()
50+
}(name, cfg)
7051
default:
7152
log.Println("Unknown provider: ", v)
7253
}
7354
}
7455
}
7556
}
7657

77-
return
78-
}
79-
80-
func userHomeDir() string {
81-
if runtime.GOOS == "windows" {
82-
home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
83-
if home == "" {
84-
home = os.Getenv("USERPROFILE")
85-
}
86-
return home
87-
}
88-
return os.Getenv("HOME")
89-
}
90-
91-
func readConfig() (config Config) {
92-
config = make(Config)
93-
94-
prefferedPaths := []string{
95-
"./cloud-ssh.yaml",
96-
userHomeDir() + "/.ssh/cloud-ssh.yaml",
97-
"/etc/cloud-ssh.yaml",
98-
}
99-
100-
var content []byte
101-
102-
for _, path := range prefferedPaths {
103-
if _, err := os.Stat(path); err == nil {
104-
fmt.Println("Found config:", path)
105-
content, err = ioutil.ReadFile(path)
106-
107-
if err != nil {
108-
log.Fatal("Error while reading config: ", err)
109-
}
110-
}
111-
}
112-
113-
if os.Getenv("AWS_ACCESS_KEY_ID") != "" && os.Getenv("AWS_SECRET_ACCESS_KEY") != "" {
114-
config["default"] = make(StrMap)
115-
config["default"]["access_key"] = os.Getenv("AWS_ACCESS_KEY_ID")
116-
config["default"]["secret_key"] = os.Getenv("AWS_SECRET_ACCESS_KEY")
117-
config["default"]["region"] = os.Getenv("AWS_REGION")
118-
config["default"]["provider"] = "aws"
119-
}
120-
121-
if len(content) == 0 {
122-
if len(config) == 0 {
123-
fmt.Println("Can't find any configuration or ENV variables. Check http://github.com/buger/cloud-ssh for documentation.")
124-
}
125-
return
126-
} else if err := yaml.Unmarshal(content, &config); err != nil {
127-
log.Fatal(err)
128-
}
58+
wg.Wait()
12959

13060
return
13161
}
13262

13363
func getMatchedInstances(clouds CloudInstances, filter string) (matched []StrMap) {
13464

135-
// Fuzzy matching (like SublimeText)
65+
// Fuzzy matching, like SublimeText
13666
filter = strings.Join(strings.Split(filter, ""), ".*?")
13767

13868
rHost := regexp.MustCompile(filter)

provider/types.go

Lines changed: 0 additions & 7 deletions
This file was deleted.

0 commit comments

Comments
 (0)