From d1f8b8545dad133ea772ef44f370a09b866e4a15 Mon Sep 17 00:00:00 2001
From: Paul Meyer <paulmey@microsoft.com>
Date: Wed, 29 May 2019 18:00:04 +0000
Subject: [PATCH] Add StepAttachDisk

---
 builder/azure/chroot/builder.go           |  6 ++-
 builder/azure/chroot/diskattacher_test.go |  2 +-
 builder/azure/chroot/step_attach_disk.go  | 65 +++++++++++++++++++++++
 3 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 builder/azure/chroot/step_attach_disk.go

diff --git a/builder/azure/chroot/builder.go b/builder/azure/chroot/builder.go
index a91fb3035..4454ffd55 100644
--- a/builder/azure/chroot/builder.go
+++ b/builder/azure/chroot/builder.go
@@ -131,7 +131,11 @@ func (b *Builder) Run(ctx context.Context, ui packer.Ui, hook packer.Hook) (pack
 	}
 
 	steps = append(steps,
-		//&StepAttachDisk{},
+		&StepAttachDisk{ // sets 'device' in stateBag
+			SubscriptionID: info.SubscriptionID,
+			ResourceGroup:  info.ResourceGroupName,
+			DiskName:       osDiskName,
+		},
 		&amznchroot.StepPreMountCommands{
 			Commands: b.config.PreMountCommands,
 		},
diff --git a/builder/azure/chroot/diskattacher_test.go b/builder/azure/chroot/diskattacher_test.go
index fbdc359ed..088dbd147 100644
--- a/builder/azure/chroot/diskattacher_test.go
+++ b/builder/azure/chroot/diskattacher_test.go
@@ -6,7 +6,7 @@ import (
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2019-03-01/compute"
-	"github.com/hashicorp/packer/builder/azure/chroot/client"
+	"github.com/hashicorp/packer/builder/azure/common/client"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
diff --git a/builder/azure/chroot/step_attach_disk.go b/builder/azure/chroot/step_attach_disk.go
new file mode 100644
index 000000000..9f05592b2
--- /dev/null
+++ b/builder/azure/chroot/step_attach_disk.go
@@ -0,0 +1,65 @@
+package chroot
+
+import (
+	"context"
+	"fmt"
+	"github.com/hashicorp/packer/builder/azure/common/client"
+	"github.com/hashicorp/packer/helper/multistep"
+	"github.com/hashicorp/packer/packer"
+	"log"
+	"time"
+)
+
+var _ multistep.Step = &StepAttachDisk{}
+
+type StepAttachDisk struct {
+	SubscriptionID, ResourceGroup, DiskName string
+}
+
+func (s StepAttachDisk) Run(ctx context.Context, state multistep.StateBag) multistep.StepAction {
+	azcli := state.Get("azureclient").(client.AzureClientSet)
+	ui := state.Get("ui").(packer.Ui)
+
+	diskResourceID := fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/disks/%s",
+		s.SubscriptionID,
+		s.ResourceGroup,
+		s.DiskName)
+	ui.Say(fmt.Sprintf("Attaching disk '%s'", diskResourceID))
+
+	da := NewDiskAttacher(azcli)
+	lun, err := da.AttachDisk(ctx, diskResourceID)
+	if err != nil {
+		log.Printf("StepAttachDisk.Run: error: %+v", err)
+		err := fmt.Errorf(
+			"error attaching disk '%s': %v", diskResourceID, err)
+		state.Put("error", err)
+		ui.Error(err.Error())
+		return multistep.ActionHalt
+	}
+
+	ui.Say("Disk attached, waiting for device to show up")
+	ctx, cancel := context.WithTimeout(ctx, time.Minute*3) // in case is not configured correctly
+	defer cancel()
+	device, err := da.WaitForDevice(ctx, lun)
+
+	ui.Say(fmt.Sprintf("Disk available at %q", device))
+	state.Put("device", device)
+	return multistep.ActionContinue
+}
+
+func (s StepAttachDisk) Cleanup(state multistep.StateBag) {
+	azcli := state.Get("azureclient").(client.AzureClientSet)
+	ui := state.Get("ui").(packer.Ui)
+
+	diskResourceID := fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/disks/%s",
+		s.SubscriptionID,
+		s.ResourceGroup,
+		s.DiskName)
+	ui.Say(fmt.Sprintf("Detaching disk '%s'", diskResourceID))
+
+	da := NewDiskAttacher(azcli)
+	err := da.DetachDisk(context.Background(), diskResourceID)
+	if err != nil {
+		ui.Error(fmt.Sprintf("error detaching %q: %v", diskResourceID, err))
+	}
+}