Added argument for the Azure driver to created or use an already exis…

…ting network security group.

Added Network Security Group name for machine to saved state file. Will only remove NSG if there are no more NICs attached to it anymore which means the current VM being deleted was the last.

Signed-off-by: Mitchell Maler <mitchell.maler@live.com>

drivers/azure/azure.go

@@ -34,28 +34,29 @@ const (
 )
 
 const (
-	flAzureEnvironment     = "azure-environment"
-	flAzureSubscriptionID  = "azure-subscription-id"
-	flAzureResourceGroup   = "azure-resource-group"
-	flAzureSSHUser         = "azure-ssh-user"
-	flAzureDockerPort      = "azure-docker-port"
-	flAzureLocation        = "azure-location"
-	flAzureSize            = "azure-size"
-	flAzureImage           = "azure-image"
-	flAzureVNet            = "azure-vnet"
-	flAzureSubnet          = "azure-subnet"
-	flAzureSubnetPrefix    = "azure-subnet-prefix"
-	flAzureAvailabilitySet = "azure-availability-set"
-	flAzurePorts           = "azure-open-port"
-	flAzurePrivateIPAddr   = "azure-private-ip-address"
-	flAzureUsePrivateIP    = "azure-use-private-ip"
-	flAzureStaticPublicIP  = "azure-static-public-ip"
-	flAzureNoPublicIP      = "azure-no-public-ip"
-	flAzureDNSLabel        = "azure-dns"
-	flAzureStorageType     = "azure-storage-type"
-	flAzureCustomData      = "azure-custom-data"
-	flAzureClientID        = "azure-client-id"
-	flAzureClientSecret    = "azure-client-secret"
+	flAzureEnvironment           = "azure-environment"
+	flAzureSubscriptionID        = "azure-subscription-id"
+	flAzureResourceGroup         = "azure-resource-group"
+	flAzureSSHUser               = "azure-ssh-user"
+	flAzureDockerPort            = "azure-docker-port"
+	flAzureLocation              = "azure-location"
+	flAzureSize                  = "azure-size"
+	flAzureImage                 = "azure-image"
+	flAzureVNet                  = "azure-vnet"
+	flAzureSubnet                = "azure-subnet"
+	flAzureSubnetPrefix          = "azure-subnet-prefix"
+	flAzureNetworkSecurityGroup  = "azure-network-security-group"
+	flAzureAvailabilitySet       = "azure-availability-set"
+	flAzurePorts                 = "azure-open-port"
+	flAzurePrivateIPAddr         = "azure-private-ip-address"
+	flAzureUsePrivateIP          = "azure-use-private-ip"
+	flAzureStaticPublicIP        = "azure-static-public-ip"
+	flAzureNoPublicIP            = "azure-no-public-ip"
+	flAzureDNSLabel              = "azure-dns"
+	flAzureStorageType           = "azure-storage-type"
+	flAzureCustomData            = "azure-custom-data"
+	flAzureClientID              = "azure-client-id"
+	flAzureClientSecret          = "azure-client-secret"
 )
 
 const (
@@ -74,15 +75,16 @@ type Driver struct {
 	SubscriptionID string
 	ResourceGroup  string
 
-	DockerPort      int
-	Location        string
-	Size            string
-	Image           string
-	VirtualNetwork  string
-	SubnetName      string
-	SubnetPrefix    string
-	AvailabilitySet string
-	StorageType     string
+	DockerPort           int
+	Location             string
+	Size                 string
+	Image                string
+	VirtualNetwork       string
+	SubnetName           string
+	SubnetPrefix         string
+	NetworkSecurityGroup string
+	AvailabilitySet      string
+	StorageType          string
 
 	OpenPorts      []string
 	PrivateIPAddr  string
@@ -216,7 +218,7 @@ func (d *Driver) GetCreateFlags() []mcnflag.Flag {
 		},
 		mcnflag.StringFlag{
 			Name:   flAzureDNSLabel,
-			Usage:  "A unique DNS label for the public IP adddress",
+			Usage:  "A unique DNS label for the public IP address",
 			EnvVar: "AZURE_DNS_LABEL",
 		},
 		mcnflag.StringSliceFlag{
@@ -233,6 +235,11 @@ func (d *Driver) GetCreateFlags() []mcnflag.Flag {
 			Usage:  "Azure Service Principal Account password (optional, browser auth is used if not specified)",
 			EnvVar: "AZURE_CLIENT_SECRET",
 		},
+		mcnflag.StringFlag{
+			Name:   flAzureNetworkSecurityGroup,
+			Usage:  "Azure Network Security Group to attach to the virtual machine.",
+			EnvVar: "AZURE_NETWORK_SECURITY_GROUP",
+		},
 	}
 }
 
@@ -276,6 +283,7 @@ func (d *Driver) SetConfigFromFlags(fl drivers.DriverOptions) error {
 	d.DockerPort = fl.Int(flAzureDockerPort)
 	d.DNSLabel = fl.String(flAzureDNSLabel)
 	d.CustomDataFile = fl.String(flAzureCustomData)
+	d.NetworkSecurityGroup = fl.String(flAzureNetworkSecurityGroup)
 
 	d.ClientID = fl.String(flAzureClientID)
 	d.ClientSecret = fl.String(flAzureClientSecret)
@@ -363,7 +371,11 @@ func (d *Driver) Create() error {
 	if err := c.CreateAvailabilitySetIfNotExists(d.ctx, d.ResourceGroup, d.AvailabilitySet, d.Location); err != nil {
 		return err
 	}
-	if err := c.CreateNetworkSecurityGroup(d.ctx, d.ResourceGroup, d.naming().NSG(), d.Location, d.ctx.FirewallRules); err != nil {
+	// Check if network security group was set by optional param and if not set it from naming.
+	if d.NetworkSecurityGroup == "" {
+		d.NetworkSecurityGroup = d.naming().NSG()
+	}
+	if err := c.CreateNetworkSecurityGroup(d.ctx, d.ResourceGroup, d.NetworkSecurityGroup, d.Location, d.ctx.FirewallRules); err != nil {
 		return err
 	}
 	vnetResourceGroup, vNetName := parseVirtualNetwork(d.VirtualNetwork, d.ResourceGroup)
@@ -422,7 +434,7 @@ func (d *Driver) Remove() error {
 	if err := c.DeletePublicIPAddressIfExists(d.ResourceGroup, d.naming().IP()); err != nil {
 		return err
 	}
-	if err := c.DeleteNetworkSecurityGroupIfExists(d.ResourceGroup, d.naming().NSG()); err != nil {
+	if err := c.CleanupNetworkSecurityGroupIfExists(d.ResourceGroup, d.NetworkSecurityGroup); err != nil {
 		return err
 	}
 	if err := c.CleanupAvailabilitySetIfExists(d.ResourceGroup, d.AvailabilitySet); err != nil {

drivers/azure/azureutil/azureutil.go

@@ -142,6 +142,10 @@ func (a AzureClient) DeleteNetworkSecurityGroupIfExists(resourceGroup, name stri
 		func() (autorest.Response, error) { return a.securityGroupsClient().Delete(resourceGroup, name, nil) })
 }
 
+func (a AzureClient) CleanupNetworkSecurityGroupIfExists(resourceGroup, name string) error {
+	return a.cleanupResourceIfExists(&nsgCleanup{rg: resourceGroup, name: name})
+}
+
 func (a AzureClient) CreatePublicIPAddress(ctx *DeploymentContext, resourceGroup, name, location string, isStatic bool, dnsLabel string) error {
 	log.Info("Creating public IP address.", logutil.Fields{
 		"name":   name,

drivers/azure/azureutil/cleanup.go

@@ -124,3 +124,27 @@ func (c *avSetCleanup) LogFields() logutil.Fields { return logutil.Fields{"name"
 func (c *avSetCleanup) HasAttachedResources() bool {
 	return c.ref.Properties.VirtualMachines != nil && len(*c.ref.Properties.VirtualMachines) > 0
 }
+
+// nsgCleanup manages cleanup of Network Security Group resources.
+type nsgCleanup struct {
+	rg, name string
+	ref      network.SecurityGroup
+}
+
+func (c *nsgCleanup) Get(a AzureClient) (err error) {
+	c.ref, err = a.securityGroupsClient().Get(c.rg, c.name, "")
+	return err
+}
+
+func (c *nsgCleanup) Delete(a AzureClient) error {
+	_, err := a.securityGroupsClient().Delete(c.rg, c.name, nil)
+	return err
+}
+
+func (c *nsgCleanup) ResourceType() string { return "Network Security Group" }
+
+func (c *nsgCleanup) LogFields() logutil.Fields { return logutil.Fields{"name": c.name} }
+
+func (c *nsgCleanup) HasAttachedResources() bool {
+	return c.ref.Properties.NetworkInterfaces != nil && len(*c.ref.Properties.NetworkInterfaces) > 0
+}