Prefer public IP if available

AWS instances in a private subnet are generally accessed via a VPN, but
many EC2 classic instances still have a public IP. We should use that
where possible, and fall back to a private IP if not.

parser.go

@@ -11,6 +11,19 @@ type state struct {
 	Modules []moduleState `json:"modules"`
 }
 
+// keyNames contains the names of the keys to check for in each resource in the
+// state file. This allows us to support multiple types of resource without too
+// much fuss.
+var keyNames []string
+
+func init() {
+	keyNames = []string{
+		"ipv4_address", // DO
+		"public_ip", // AWS
+		"private_ip", // AWS
+	}
+}
+
 // read populates the state object from a statefile.
 func (s *state) read(stateFile io.Reader) error {
 
@@ -57,26 +70,23 @@ type resourceState struct {
 }
 
 // isSupported returns true if terraform-inventory supports this resource.
-func (s *resourceState) isSupported() bool {
+func (s resourceState) isSupported() bool {
 	return s.Address() != ""
 }
 
 // Address returns the IP address of this resource.
-func (s *resourceState) Address() string {
-	switch s.Type {
-	case "aws_instance":
-		return s.Primary.Attributes["private_ip"]
-
-	case "digitalocean_droplet":
-		return s.Primary.Attributes["ipv4_address"]
-
-	default:
-		return ""
+func (s resourceState) Address() string {
+	for _, key := range(keyNames) {
+		if ip := s.Primary.Attributes[key]; ip != "" {
+			return ip
+		}
 	}
+
+	return ""
 }
 
 // Attributes returns a map containing everything we know about this resource.
-func (s *resourceState) Attributes() map[string]string {
+func (s resourceState) Attributes() map[string]string {
 	return s.Primary.Attributes
 }
 

parser_test.go

@@ -33,7 +33,8 @@ const exampleStateFile = `
 						"id": "i-bbbbbbbb",
 						"attributes": {
 							"id": "i-bbbbbbbb",
-							"private_ip": "10.0.0.2"
+							"private_ip": "10.0.0.2",
+							"public_ip": "50.0.0.1"
 						}
 					}
 				},
@@ -91,6 +92,7 @@ func TestStateRead(t *testing.T) {
 							Attributes: map[string]string{
 								"id":         "i-bbbbbbbb",
 								"private_ip": "10.0.0.2",
+								"public_ip": "50.0.0.1",
 							},
 						},
 					},
@@ -136,8 +138,21 @@ func TestResources(t *testing.T) {
 	assert.Equal(t, "digitalocean_droplet", inst["three"].Type)
 }
 
-func TestIsSupported(t *testing.T) {
+func TestAddress(t *testing.T) {
+	r := strings.NewReader(exampleStateFile)
 
+	var s state
+	err := s.read(r)
+	assert.Nil(t, err)
+
+	inst := s.resources()
+	assert.Equal(t, 3, len(inst))
+	assert.Equal(t, "10.0.0.1", inst["one"].Address())
+	assert.Equal(t, "50.0.0.1", inst["two"].Address())
+	assert.Equal(t, "192.168.0.3", inst["three"].Address())
+}
+
+func TestIsSupported(t *testing.T) {
 	r := resourceState{
 		Type: "something",
 	}