From 04fd9e7741af51607dad5846a410f6b145c238c6 Mon Sep 17 00:00:00 2001
From: Jean-Laurent de Morlhon <jeanlaurent@morlhon.net>
Date: Thu, 14 Jan 2016 11:07:28 +0100
Subject: [PATCH] Fetching the default vpc id from account

Signed-off-by: Jean-Laurent de Morlhon <jeanlaurent@morlhon.net>
---
 drivers/amazonec2/amazonec2.go      | 24 ++++++++++-
 drivers/amazonec2/amazonec2_test.go | 66 +++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/drivers/amazonec2/amazonec2.go b/drivers/amazonec2/amazonec2.go
index 6033f486bb..33c0b0fffc 100644
--- a/drivers/amazonec2/amazonec2.go
+++ b/drivers/amazonec2/amazonec2.go
@@ -289,8 +289,15 @@ func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
 		return fmt.Errorf("amazonec2 driver requires the --amazonec2-secret-key option")
 	}
 
+	if d.VpcId == "" {
+		d.VpcId, err = d.getDefaultVPCId()
+		if err != nil {
+			log.Errorf("couldn't determine your account Default VPC ID : %q", err)
+		}
+	}
+
 	if d.SubnetId == "" && d.VpcId == "" {
-		return fmt.Errorf("amazonec2 driver requires either the --amazonec2-subnet-id or --amazonec2-vpc-id option")
+		return fmt.Errorf("amazonec2 driver requires either the --amazonec2-subnet-id or --amazonec2-vpc-id option or an AWS Account with a default vpc-id")
 	}
 
 	if d.SubnetId != "" && d.VpcId != "" {
@@ -952,6 +959,21 @@ func (d *Driver) deleteKeyPair() error {
 	return nil
 }
 
+func (d *Driver) getDefaultVPCId() (string, error) {
+	output, err := d.getClient().DescribeAccountAttributes(&ec2.DescribeAccountAttributesInput{})
+	if err != nil {
+		return "", err
+	}
+
+	for _, attribute := range output.AccountAttributes {
+		if *attribute.AttributeName == "default-vpc" {
+			return *attribute.AttributeValues[0].AttributeValue, nil
+		}
+	}
+
+	return "", errors.New("No default-vpc attribute")
+}
+
 func generateId() string {
 	rb := make([]byte, 10)
 	_, err := rand.Read(rb)
diff --git a/drivers/amazonec2/amazonec2_test.go b/drivers/amazonec2/amazonec2_test.go
index 98170f056e..6e9c0ee850 100644
--- a/drivers/amazonec2/amazonec2_test.go
+++ b/drivers/amazonec2/amazonec2_test.go
@@ -239,3 +239,69 @@ func TestSetConfigFromFlags(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Empty(t, checkFlags.InvalidFlags)
 }
+
+type fakeEC2WithDescribe struct {
+	*ec2.EC2
+	output *ec2.DescribeAccountAttributesOutput
+	err    error
+}
+
+func (f *fakeEC2WithDescribe) DescribeAccountAttributes(input *ec2.DescribeAccountAttributesInput) (*ec2.DescribeAccountAttributesOutput, error) {
+	return f.output, f.err
+}
+
+func TestFindDefaultVPC(t *testing.T) {
+	defaultVpc := "default-vpc"
+	vpcName := "vpc-9999"
+
+	driver := NewDriver("machineFoo", "path")
+	driver.clientFactory = func() Ec2Client {
+		return &fakeEC2WithDescribe{
+			output: &ec2.DescribeAccountAttributesOutput{
+				AccountAttributes: []*ec2.AccountAttribute{
+					{
+						AttributeName: &defaultVpc,
+						AttributeValues: []*ec2.AccountAttributeValue{
+							{AttributeValue: &vpcName},
+						},
+					},
+				},
+			},
+		}
+	}
+
+	vpc, err := driver.getDefaultVPCId()
+
+	assert.Equal(t, "vpc-9999", vpc)
+	assert.NoError(t, err)
+}
+
+func TestDefaultVPCIsMissing(t *testing.T) {
+	driver := NewDriver("machineFoo", "path")
+	driver.clientFactory = func() Ec2Client {
+		return &fakeEC2WithDescribe{
+			output: &ec2.DescribeAccountAttributesOutput{
+				AccountAttributes: []*ec2.AccountAttribute{},
+			},
+		}
+	}
+
+	vpc, err := driver.getDefaultVPCId()
+
+	assert.EqualError(t, err, "No default-vpc attribute")
+	assert.Empty(t, vpc)
+}
+
+func TestDescribeAccountAttributeFails(t *testing.T) {
+	driver := NewDriver("machineFoo", "path")
+	driver.clientFactory = func() Ec2Client {
+		return &fakeEC2WithDescribe{
+			err: errors.New("Not Found"),
+		}
+	}
+
+	vpc, err := driver.getDefaultVPCId()
+
+	assert.EqualError(t, err, "Not Found")
+	assert.Empty(t, vpc)
+}