adobe · Thushara67 · Jun 24, 2020 · Jun 25, 2020 · Jun 25, 2020 · Jun 25, 2020
diff --git a/README.md b/README.md
@@ -7,6 +7,8 @@ Current supported:
   - EC2
   - RDS
   - ELB
+- Azure
+  - Virtual Networks
 
 (PRs welcome for more!)
 
@@ -32,6 +34,7 @@ Use "cloudinventory [command] --help" for more information about a command.
 ```
 
 The subcommands all have additional options which be found by doing a `-h`
+
 Example for AWS:
 
 ```bash
@@ -53,10 +56,32 @@ Global Flags:
   -p, --path string     file path to dump the inventory in (default "cloudinventory.json")
 ```
 
+Example for Azure :
+
+```bash
+cloudinventory dump azurevnet -h
+Dump Azure inventory. Currently supports Virtual networks
+
+Usage:
+  cloudinventory dump azurevnet [flags]
+
+Flags:
+  -h, --help                help for azurevnet
+  -i, --input_Path string   file path to take subscriptionIDs as input
+  -m, --maxGoRoutines int   customize maximum no.of Goroutines  (default -1)
+  -s, --stats               dumps the stats of different resources for subscriptions
+
+Global Flags:
+  -f, --filter string   limit dump to a particular cloud service, e.g ec2/rds/route53/loadbalancer in aws and virtual networks in azure
+  -p, --path string     file path to dump the inventory in (default "cloudinventory.json")
+```
+
 The tool reads credentials from your environment.
 
 For AWS see: <https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-envvars.html>
 
+For Azure see: <https://github.com/Azure/azure-sdk-for-go#more-authentication-details>
+
 ## Library Use
 
 The packages with helping wrappers can be imported individually.
@@ -73,4 +98,4 @@ Contributions are very welcome. Please see [Contributing Guide](CONTRIBUTING.md)
 
 ## Licensing
 
-The project is licensed under the Apache V2 License. See [License](LICENSE) for more information
+The project is licensed under the Apache V2 License. See [License](LICENSE) for more information
diff --git a/azurenetwork/base.go b/azurenetwork/base.go
@@ -0,0 +1,31 @@
+package azurenetwork
+
+import (
+        "context"
+        "github.com/Azure/azure-sdk-for-go/profiles/2019-03-01/resources/mgmt/subscriptions"
+        "github.com/Azure/go-autorest/autorest/azure/auth"
+)
+
+// GetAllSubscriptionIDs function returns a map of subscription name and subscription ID at account level
+func GetAllSubscriptionIDs(ctx context.Context) (map[string]string, error) {
+        subscriptionMap := make(map[string]string)
+        authorizer, err := auth.NewAuthorizerFromEnvironment()
+        if err != nil {
+                return subscriptionMap, err
+        }
+        newClient := subscriptions.NewClient()
+        newClient.Authorizer = authorizer
+        result, err := newClient.ListComplete(ctx)
+        if err != nil {
+                return subscriptionMap, err
+        }
+        for result.NotDone() {
+                subscription := result.Value()
+                subscriptionMap[*subscription.DisplayName] = *subscription.SubscriptionID
+                if err := result.Next(); err != nil {
+                        return subscriptionMap, err
+                }
+        }
+        return subscriptionMap, nil
+
+}
diff --git a/azurenetwork/base_test.go b/azurenetwork/base_test.go
@@ -0,0 +1,21 @@
+package azurenetwork
+
+import (
+        "context"
+        "testing"
+        "time"
+)
+
+// TestGetAllSubscriptionIDs tests function GetAllSubscriptionIDs
+func TestGetAllSubscriptionIDs(t *testing.T) {
+
+        ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
+        defer cancel()
+        subscriptionMap, err := GetAllSubscriptionIDs(ctx)
+        if err != nil {
+                t.Errorf("Failed to  Get all SubscriptionIDs: %v", err)
+        } else {
+                t.Logf("GetallSubscriptionIDs successful.Found %d subscription IDs", len(subscriptionMap))
+
+        }
+}
diff --git a/azurenetwork/vnet.go b/azurenetwork/vnet.go
@@ -0,0 +1,36 @@
+package azurenetwork
+
+import (
+        "context"
+        "github.com/Azure/azure-sdk-for-go/profiles/2019-03-01/network/mgmt/network"
+        "github.com/Azure/go-autorest/autorest/azure/auth"
+        "time"
+)
+
+// GetAllVNet function gathers all virtual networks for a given subscriptionID
+func GetAllVNet(subscriptionID string) ([]*network.VirtualNetwork, error) {
+        var vnetList []*network.VirtualNetwork
+        authorizer, err := auth.NewAuthorizerFromEnvironment()
+        if err != nil {
+                return vnetList, err
+        }
+        vnetClient := network.NewVirtualNetworksClient(subscriptionID)
+        vnetClient.Authorizer = authorizer
+        ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
+        defer cancel()
+        results, err := vnetClient.ListAllComplete(ctx)
+        if err != nil {
+                return vnetList, err
+        }
+        for results.NotDone() {
+                vnet := results.Value()
+                vnetList = append(vnetList, &vnet)
+                if err := results.Next(); err != nil {
+                        return vnetList, err
+                }
+
+        }
+
+        return vnetList, nil
+
+}
diff --git a/azurenetwork/vnet_test.go b/azurenetwork/vnet_test.go
@@ -0,0 +1,27 @@
+package azurenetwork
+
+import (
+        "context"
+        "testing"
+        "time"
+)
+
+// TestGetAllVNet tests the function GetAllVNet
+func TestGetAllVNet(t *testing.T) {
+        if testing.Short() {
+                t.Skip("Skipping test in short mode")
+        }
+        ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
+        defer cancel()
+        subscriptions, err := GetAllSubscriptionIDs(ctx)
+        if err != nil {
+                t.Errorf("Unable to get subscriptionIDs: %v", err)
+        }
+        for key, subsID := range subscriptions {
+                Dbs, err := GetAllVNet(subsID)
+                if err != nil {
+                        t.Errorf("Failed to get virtual networks for subscription: %s because %v", key, err)
+                }
+                t.Logf("Found %d virtual networks in %s", len(Dbs), key)
+        }
+}
diff --git a/azurevnetcollector/vnetcollect.go b/azurevnetcollector/vnetcollect.go
@@ -0,0 +1,122 @@
+package azurevnetcollector
+
+import (
+        "context"
+        "fmt"
+        "github.com/adobe/cloudinventory/azurenetwork"
+        "github.com/Azure/azure-sdk-for-go/profiles/2019-03-01/network/mgmt/network"
+        "sync"
+        "time"
+)
+
+// AzureCollector is a struct that contains a map of subscription name and its subscriptionID
+type AzureCollector struct {
+        SubscriptionMap map[string]string
+}
+
+// NewAzureCollector returns an AzureCollector with subscription info set in subscriptionMap.
+func NewAzureCollector() (AzureCollector, error) {
+        var col AzureCollector
+        ctx, cancel := context.WithTimeout(context.Background(), 100*time.Second)
+        defer cancel()
+        err := col.GetSubscription(ctx)
+        if err != nil {
+                return col, err
+        }
+        return col, nil
+}
+
+// NewAzureCollectorUserDefined returns an AzureCollector with subscription info given by user in subscriptionMap.
+func NewAzureCollectorUserDefined(subscriptionID []string) (AzureCollector, error) {
+        var col AzureCollector
+        subID := make(map[string]string)
+        for i := 0; i < len(subscriptionID); i++ {
+                subID["subscription_id :"+subscriptionID[i]] = subscriptionID[i]
+        }
+        col.SubscriptionMap = subID
+        return col, nil
+}
+
+// GetSubscription adds map of subscription name with subscriptionID to the AzureCollector
+func (col *AzureCollector) GetSubscription(ctx context.Context) error {
+        subscription := make(map[string]string)
+        var err error
+        subscription, err = azurenetwork.GetAllSubscriptionIDs(ctx)
+        if err != nil {
+                return err
+        }
+        col.SubscriptionMap = subscription
+        return nil
+
+}
+
+// CollectVirtualNetworks gathers Virtual Networks stats and data for each subscriptionID in an account level
+// Function takes no.of goroutines to be created from user as input
+func (col *AzureCollector) CollectVirtualNetworks(maxGoRoutines int) (map[string][]*network.VirtualNetwork, map[string]int, error) {
+        VNet := make(map[string][]*network.VirtualNetwork)
+        VNetCount := make(map[string]int)
+        type VirtualNetworksPerSubscriptionID struct {
+                SubscriptionName string
+                VnetList         []*network.VirtualNetwork
+        }
+        var chanCapacity int
+
+        if maxGoRoutines >= len(col.SubscriptionMap) || maxGoRoutines < 0 {
+                chanCapacity = len(col.SubscriptionMap)
+        } else {
+                chanCapacity = maxGoRoutines
+        }
+        subscriptionCount := 0
+        vnetsChan := make(chan VirtualNetworksPerSubscriptionID, chanCapacity)
+        errChan := make(chan error, chanCapacity)
+        var wg sync.WaitGroup
+        for subscriptionName, subID := range col.SubscriptionMap {
+                if subscriptionCount < chanCapacity {
+                        wg.Add(1)
+                        go func(subID string, subscriptionName string, vnetsChan chan VirtualNetworksPerSubscriptionID, errChan chan error) {
+                                defer wg.Done()
+                                vnets, err := CollectVirtualNetworksPerSubscriptionID(subID)
+                                if err != nil {
+                                        errChan <- fmt.Errorf(fmt.Sprintf("Error while gathering %s: %v", subscriptionName, err))
+                                        return
+                                }
+                                if vnets == nil {
+                                        return
+                                }
+                                vnetsChan <- VirtualNetworksPerSubscriptionID{subscriptionName, vnets}
+                        }(subID, subscriptionName, vnetsChan, errChan)
+                        if subscriptionCount == chanCapacity-1 {
+                                wg.Wait()
+                                close(vnetsChan)
+                                close(errChan)
+                                if len(errChan) > 0 {
+                                        return nil, nil, fmt.Errorf(fmt.Sprintf("Failed to gather virtual network data: %v", <-errChan))
+                                }
+                                for subscriptionVnets := range vnetsChan {
+                                        VNet[subscriptionVnets.SubscriptionName] = subscriptionVnets.VnetList
+                                        VNetCount[subscriptionVnets.SubscriptionName] = len(subscriptionVnets.VnetList)
+                                }
+                        }
+                } else {
+                        vnets, err := CollectVirtualNetworksPerSubscriptionID(subID)
+                        if err != nil {
+                                return nil, nil, fmt.Errorf(fmt.Sprintf("Failed to gather virtual network data: %v", err))
+                        }
+                        if vnets == nil {
+                                subscriptionCount++
+                                continue
+                        }
+                        VNet[subscriptionName] = vnets
+                        VNetCount[subscriptionName] = len(vnets)
+                }
+                subscriptionCount++
+        }
+
+        return VNet, VNetCount, nil
+}
+
+// CollectVirtualNetworksPerSubscriptionID returns a slice of virtual networks for a given subscriptionID
+func CollectVirtualNetworksPerSubscriptionID(subscriptionID string) ([]*network.VirtualNetwork, error) {
+        vnetlist, err := azurenetwork.GetAllVNet(subscriptionID)
+        return vnetlist, err
+}
diff --git a/azurevnetcollector/vnetcollect_test.go b/azurevnetcollector/vnetcollect_test.go
@@ -0,0 +1,19 @@
+package azurevnetcollector
+
+import "testing"
+
+// TestCollectVirtualNetworks test the function CollectVirtualNetworks
+func TestCollectVirtualNetworks(t *testing.T) {
+        if testing.Short() {
+                t.Skip("Skipping test in short mode")
+        }
+        col, err := NewAzureCollector()
+        if err != nil {
+                t.Errorf("Failed to create collector: %v", err)
+        }
+        maxGoRoutines := len(col.SubscriptionMap)
+        _, _, err = col.CollectVirtualNetworks(maxGoRoutines)
+        if err != nil {
+                t.Errorf("Failed to collect Virtual networks: %v", err)
+        }
+}