From 9a46f080dfddb1c94fa8681fb7fe5ba64b1ce6ff Mon Sep 17 00:00:00 2001
From: AvineshTripathi <avineshtripathi1@gmail.com>
Date: Thu, 20 Jun 2024 03:34:26 +0530
Subject: [PATCH] feat: added dependency support to azure

---
 providers/azure/compute/snapshots.go        | 27 +++++++++++++++++++++
 providers/azure/compute/virtual_machines.go | 27 +++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/providers/azure/compute/snapshots.go b/providers/azure/compute/snapshots.go
index 8dd86895f..cab0ff7ca 100644
--- a/providers/azure/compute/snapshots.go
+++ b/providers/azure/compute/snapshots.go
@@ -46,6 +46,7 @@ func Snapshots(ctx context.Context, client providers.ProviderClient) ([]models.R
 				ResourceId: *snapshot.ID,
 				Cost:       0,
 				Name:       *snapshot.Name,
+				Relations:  getSnapshotRelation(snapshot),
 				FetchedAt:  time.Now(),
 				Tags:       tags,
 				CreatedAt:  *snapshot.Properties.TimeCreated,
@@ -62,3 +63,29 @@ func Snapshots(ctx context.Context, client providers.ProviderClient) ([]models.R
 	}).Info("Fetched resources")
 	return resources, nil
 }
+
+func getSnapshotRelation(snp *armcompute.Snapshot) []models.Link {
+	var rel []models.Link
+
+	if snp.Properties.CreationData != nil {
+		if snp.Properties.CreationData.ImageReference != nil && snp.Properties.CreationData.ImageReference.ID != nil {
+			rel = append(rel, models.Link{
+				ResourceID: *snp.Properties.CreationData.ImageReference.ID,
+				Type:       "Image",
+				Name:       *snp.Properties.CreationData.ImageReference.ID,
+				Relation:   "USES",
+			})
+		}
+
+		if snp.Properties.CreationData.SourceResourceID != nil {
+			rel = append(rel, models.Link{
+				ResourceID: *snp.Properties.CreationData.SourceResourceID,
+				Type:       "Disk",
+				Name:       *snp.Properties.CreationData.SourceResourceID,
+				Relation:   "USES",
+			})
+		}
+
+	}
+	return rel
+}
diff --git a/providers/azure/compute/virtual_machines.go b/providers/azure/compute/virtual_machines.go
index 48f308449..273fbd6c3 100644
--- a/providers/azure/compute/virtual_machines.go
+++ b/providers/azure/compute/virtual_machines.go
@@ -73,6 +73,7 @@ func VirtualMachines(ctx context.Context, client providers.ProviderClient) ([]mo
 				Cost:       cost,
 				Name:       *vm.Name,
 				FetchedAt:  time.Now(),
+				Relations:  getVmRelation(vm),
 				Tags:       tags,
 				Link:       fmt.Sprintf("https://portal.azure.com/#resource%s", *vm.ID),
 			})
@@ -87,3 +88,29 @@ func VirtualMachines(ctx context.Context, client providers.ProviderClient) ([]mo
 	}).Info("Fetched resources")
 	return resources, nil
 }
+
+func getVmRelation(vm *armcompute.VirtualMachine) []models.Link {
+
+	var rel []models.Link
+
+	if vm.Properties.StorageProfile != nil && vm.Properties.StorageProfile.DataDisks != nil {
+		for _, disk := range vm.Properties.StorageProfile.DataDisks {
+			rel = append(rel, models.Link{
+				ResourceID: *disk.ManagedDisk.ID,
+				Type:       "Disk",
+				Name:       *disk.Name,
+				Relation:   "USES",
+			})
+		}
+
+		if vm.Properties.StorageProfile.ImageReference != nil {
+			rel = append(rel, models.Link{
+				ResourceID: *vm.Properties.StorageProfile.ImageReference.ID,
+				Type:       "Image",
+				Name:       *vm.Properties.StorageProfile.ImageReference.ID,
+				Relation:   "USES",
+			})
+		}
+	}
+	return rel
+}