Factor into separate files

duration.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+	"time"
+)
+
+func fmtDuration(duration time.Duration) string {
+	result := []string{}
+	started := false
+	if duration > 7*24*time.Hour || started {
+		started = true
+		weeks := int(duration.Hours() / 24 / 7)
+		duration %= 7 * 24 * time.Hour
+		result = append(result, fmt.Sprintf("%2dw", weeks))
+	}
+	if duration > 24*time.Hour || started {
+		started = true
+		days := int(duration.Hours() / 24)
+		duration %= 24 * time.Hour
+		result = append(result, fmt.Sprintf("%2dd", days))
+	}
+	if duration > time.Hour || started {
+		started = true
+		hours := int(duration.Hours())
+		duration %= time.Hour
+		result = append(result, fmt.Sprintf("%2dh", hours))
+	}
+	if duration > time.Minute || started {
+		started = true
+		minutes := int(duration.Minutes())
+		duration %= time.Minute
+		result = append(result, fmt.Sprintf("%2dm", minutes))
+	}
+	seconds := int(duration.Seconds())
+	result = append(result, fmt.Sprintf("%2ds", seconds))
+	return strings.TrimSpace(strings.Join(result, ""))
+}

instance.go

@@ -0,0 +1,94 @@
+package main
+
+import (
+	"fmt"
+	"sort"
+	"time"
+
+	"github.com/stripe/aws-go/gen/ec2"
+)
+
+type Instance struct {
+	InstanceID, PrivateIP, State           string
+	Up                                     time.Duration
+	Tags                                   map[string]string
+	ICMPPing, SSHPing, HTTPPing, HTTPSPing <-chan PingResponse
+}
+
+func NewInstance(i ec2.Instance) *Instance {
+	return &Instance{
+		*i.InstanceID, *i.PrivateIPAddress, *i.State.Name,
+		time.Since(i.LaunchTime),
+		TagMap(i.Tags),
+		ICMPPing(*i.PrivateIPAddress),
+		SSHPing(*i.PrivateIPAddress),
+		HTTPPing(*i.PrivateIPAddress),
+		HTTPSPing(*i.PrivateIPAddress),
+	}
+}
+
+func TagMap(ts []ec2.Tag) map[string]string {
+	m := map[string]string{}
+	for _, t := range ts {
+		m[*t.Key] = *t.Value
+	}
+	return m
+}
+
+func (i *Instance) Name() string {
+	return i.Tags["Name"]
+}
+
+func (i *Instance) String() string {
+	return fmt.Sprint(i.InstanceID, " ", i.Name(), " ", i.PrivateIP)
+}
+
+func (i *Instance) PrettyState() string {
+	var (
+		s     = ""
+		color = ""
+	)
+	switch i.State {
+	default:
+		s = "U"
+	case "running":
+		s = "R"
+		color = "32" // Green
+	case "rebooting":
+		s = "B"
+		color = "34" // Blue
+	case "pending":
+		s = "P"
+		color = "33" // Yellow
+	case "stopping":
+		s = "-"
+		color = "33" // Yellow
+	case "shutting-down":
+		s = "G"
+		color = "33" // Yellow
+	case "stopped":
+		s = "."
+		color = "31" // Red
+	case "terminated":
+		s = "T"
+		color = "31" // Red
+	}
+	return fmt.Sprint("[" + color + "m" + s + "")
+}
+
+func InstancesFromEC2Result(in *ec2.DescribeInstancesResult) []*Instance {
+	out := []*Instance{}
+	for _, r := range in.Reservations {
+		for _, oi := range r.Instances {
+			out = append(out, NewInstance(oi))
+		}
+	}
+	sort.Sort(InstancesByName(out))
+	return out
+}
+
+type InstancesByName []*Instance
+
+func (a InstancesByName) Len() int           { return len(a) }
+func (a InstancesByName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
+func (a InstancesByName) Less(i, j int) bool { return a[i].Name() < a[j].Name() }

main.go

@@ -2,18 +2,13 @@ package main
 
 import (
 	"bufio"
-	"crypto/tls"
 	"fmt"
 	"log"
-	"net"
-	"net/http"
 	"os"
-	"os/exec"
-	"sort"
 	"syscall"
-	"time"
 
 	"github.com/olekukonko/tablewriter"
+
 	"github.com/stripe/aws-go/aws"
 	"github.com/stripe/aws-go/gen/ec2"
 )
@@ -27,159 +22,18 @@ func init() {
 	}
 }
 
-type PingResponse struct {
-	OK       bool
-	Duration time.Duration
-}
-
-func (p PingResponse) String() string {
-	var (
-		result = "-"
-		color  = "31"
-	)
-
-	if p.OK {
-		result = fmt.Sprintf("%.01fms", p.Duration.Seconds()*1000)
-		if p.Duration < 100*time.Millisecond {
-			color = "32"
-		}
-	}
-
-	return "\033[" + color + "m" + result + "\033[m"
-}
-
-func SSHPing(where string) <-chan PingResponse {
-	return DoPing(func() error {
-		d := &net.Dialer{Timeout: time.Second}
-		c, err := d.Dial("tcp", where+":22")
-		if err == nil {
-			c.Close()
-		}
-		return err
-	})
-}
-
-func HTTPPing(where string) <-chan PingResponse {
-	return DoPing(func() error {
-		_, err := http.Get(fmt.Sprint("http://", where))
-		return err
-	})
-}
-
-func HTTPSPing(where string) <-chan PingResponse {
-	return DoPing(func() error {
-		_, err := http.Get(fmt.Sprint("https://", where))
-		return err
-	})
-}
-
-func Ping(where string) <-chan PingResponse {
-	// -c: Count, -W timeout in seconds
-	return DoPing(exec.Command("ping", "-W1", "-c1", where).Run)
-}
-
-func DoPing(f func() error) <-chan PingResponse {
-	respChan := make(chan PingResponse)
-	go func() {
-		start := time.Now()
-		err := f()
-		responseTime := time.Since(start)
-		respChan <- PingResponse{err == nil, responseTime}
-	}()
-	return respChan
-}
-
-type Instance struct {
-	InstanceID, PrivateIP              string
-	Tags                               map[string]string
-	Ping, SSHPing, HTTPPing, HTTPSPing <-chan PingResponse
-}
-
-func TagMap(ts []ec2.Tag) map[string]string {
-	m := map[string]string{}
-	for _, t := range ts {
-		m[*t.Key] = *t.Value
-	}
-	return m
-}
-
-func NewInstance(i ec2.Instance) *Instance {
-	ping := Ping(*i.PrivateIPAddress)
-	sshPing := SSHPing(*i.PrivateIPAddress)
-	httpPing := HTTPPing(*i.PrivateIPAddress)
-	httpsPing := HTTPSPing(*i.PrivateIPAddress)
-	return &Instance{
-		*i.InstanceID, *i.PrivateIPAddress, TagMap(i.Tags),
-		ping, sshPing, httpPing, httpsPing}
-}
-
-func (i *Instance) Name() string {
-	return i.Tags["Name"]
-}
-
-func (i *Instance) String() string {
-	return fmt.Sprint(i.InstanceID, " ", i.Name(), " ", i.PrivateIP)
-}
-
-func InstancesFromEC2Result(in *ec2.DescribeInstancesResult) []*Instance {
-	out := []*Instance{}
-	for _, r := range in.Reservations {
-		for _, oi := range r.Instances {
-			out = append(out, NewInstance(oi))
-		}
-	}
-	sort.Sort(InstancesByName(out))
-	return out
-}
-
-type InstancesByName []*Instance
-
-func (a InstancesByName) Len() int           { return len(a) }
-func (a InstancesByName) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
-func (a InstancesByName) Less(i, j int) bool { return a[i].Name() < a[j].Name() }
-
-// Configure HTTP for 1s timeout and HTTPS to ignore SSL CA errors.
-func ConfigureHTTP() {
-
-	http.DefaultClient.Timeout = 1 * time.Second
-
-	t := http.DefaultTransport.(*http.Transport)
-
-	// Ignore SSL certificate errors
-	t.TLSClientConfig = &tls.Config{
-		InsecureSkipVerify: true,
-	}
-}
-
-func main() {
-
-	if os.Getenv("SSH_AUTH_SOCK") == "" {
-		fmt.Fprintln(os.Stderr, "Warning: agent forwarding not enabled")
-	}
-
-	creds := aws.IAMCreds()
-
-	c := ec2.New(creds, AWS_REGION, nil)
-
-	resp, err := c.DescribeInstances(&ec2.DescribeInstancesRequest{})
-	if err != nil {
-		log.Fatal("DescribeInstances error:", err)
-	}
-
-	// Do this after querying the AWS endpoint (otherwise vulnerable to MITM.)
-	ConfigureHTTP()
-
-	instances := InstancesFromEC2Result(resp)
-
+func ShowInstances(instances []*Instance) {
 	table := tablewriter.NewWriter(os.Stderr)
 	table.SetAlignment(tablewriter.ALIGN_RIGHT)
-	table.SetHeader([]string{"N", "InstanceID", "Name", "PrivateIP",
+	table.SetHeader([]string{
+		"N", "ID", "Name", "S", "Private", "Launch",
 		"ICMP", "SSH", "HTTP", "HTTPS"})
 
 	for n, i := range instances {
 		row := []string{
-			fmt.Sprint(n), i.InstanceID, i.Name(), i.PrivateIP,
-			(<-i.Ping).String(),
+			fmt.Sprint(n), i.InstanceID[2:], i.Name(), i.PrettyState(),
+			i.PrivateIP, fmtDuration(i.Up),
+			(<-i.ICMPPing).String(),
 			(<-i.SSHPing).String(),
 			(<-i.HTTPPing).String(),
 			(<-i.HTTPSPing).String(),
@@ -188,30 +42,63 @@ func main() {
 	}
 
 	table.Render()
+}
 
+func UserSelectInstance(max int) int {
 	s := bufio.NewScanner(os.Stdin)
-	if s.Err() != nil || !s.Scan() {
+	if !s.Scan() {
+		// User closed stdin before we read anything
+		os.Exit(1)
+	}
+	if s.Err() != nil {
 		log.Fatalln("Error reading stdin:", s.Err())
 	}
-
 	var n int
-	_, err = fmt.Sscan(s.Text(), &n)
+	_, err := fmt.Sscan(s.Text(), &n)
 	if err != nil {
 		log.Fatalln("Unrecognised input:", s.Text())
 	}
-	if n >= len(instances) {
+	if n >= max {
 		log.Fatalln("%d is not a valid instance", n)
 	}
+	return n
+}
 
-	instance := instances[n]
+func InvokeSSH(instance *Instance) {
 	log.Println("Sending you to:", instance)
 
 	args := []string{"/usr/bin/ssh"}
-	args = append(args, os.Args[1:]...)
+	// Enable the user to specify arguments to the left and right of the host.
+	left, right := BreakArgsBySeparator()
+	args = append(args, left...)
 	args = append(args, instance.PrivateIP)
+	args = append(args, right...)
 
-	err = syscall.Exec("/usr/bin/ssh", args, os.Environ())
+	err := syscall.Exec("/usr/bin/ssh", args, os.Environ())
 	if err != nil {
 		log.Fatalln("Failed to exec:", err)
 	}
 }
+
+func main() {
+
+	if os.Getenv("SSH_AUTH_SOCK") == "" {
+		fmt.Fprintln(os.Stderr, "Warning: agent forwarding not enabled")
+	}
+
+	creds := aws.IAMCreds()
+	c := ec2.New(creds, AWS_REGION, nil)
+	ec2Instances, err := c.DescribeInstances(&ec2.DescribeInstancesRequest{})
+	if err != nil {
+		log.Fatal("DescribeInstances error:", err)
+	}
+
+	// Do this after querying the AWS endpoint (otherwise vulnerable to MITM.)
+	ConfigureHTTP()
+
+	instances := InstancesFromEC2Result(ec2Instances)
+	ShowInstances(instances)
+
+	n := UserSelectInstance(len(instances))
+	InvokeSSH(instances[n])
+}