Merge pull request #6980 from amydutta/amdut/snapshot

Adding options for Managed Image OS Disk and Data Disk(s) snapshot(s)

builder/azure/arm/artifact.go

@@ -33,18 +33,22 @@ type Artifact struct {
 	ManagedImageName                   string
 	ManagedImageLocation               string
 	ManagedImageId                     string
+	ManagedImageOSDiskSnapshotName     string
+	ManagedImageDataDiskSnapshotPrefix string
 
 	// Additional Disks
 	AdditionalDisks *[]AdditionalDiskArtifact
 }
 
-func NewManagedImageArtifact(osType, resourceGroup, name, location, id string) (*Artifact, error) {
+func NewManagedImageArtifact(osType, resourceGroup, name, location, id, osDiskSnapshotName, dataDiskSnapshotPrefix string) (*Artifact, error) {
 	return &Artifact{
 		ManagedImageResourceGroupName:      resourceGroup,
 		ManagedImageName:                   name,
 		ManagedImageLocation:               location,
 		ManagedImageId:                     id,
 		OSType:                             osType,
+		ManagedImageOSDiskSnapshotName:     osDiskSnapshotName,
+		ManagedImageDataDiskSnapshotPrefix: dataDiskSnapshotPrefix,
 	}, nil
 }
 
@@ -158,6 +162,12 @@ func (a *Artifact) String() string {
 		buf.WriteString(fmt.Sprintf("ManagedImageName: %s\n", a.ManagedImageName))
 		buf.WriteString(fmt.Sprintf("ManagedImageId: %s\n", a.ManagedImageId))
 		buf.WriteString(fmt.Sprintf("ManagedImageLocation: %s\n", a.ManagedImageLocation))
+		if a.ManagedImageOSDiskSnapshotName != "" {
+			buf.WriteString(fmt.Sprintf("ManagedImageOSDiskSnapshotName: %s\n", a.ManagedImageOSDiskSnapshotName))
+		}
+		if a.ManagedImageDataDiskSnapshotPrefix != "" {
+			buf.WriteString(fmt.Sprintf("ManagedImageDataDiskSnapshotPrefix: %s\n", a.ManagedImageDataDiskSnapshotPrefix))
+		}
 	} else {
 		buf.WriteString(fmt.Sprintf("StorageAccountLocation: %s\n", a.StorageAccountLocation))
 		buf.WriteString(fmt.Sprintf("OSDiskUri: %s\n", a.OSDiskUri))

builder/azure/arm/artifact_test.go

@@ -42,14 +42,67 @@ func TestArtifactIdVHD(t *testing.T) {
 }
 
 func TestArtifactIDManagedImage(t *testing.T) {
-	artifact, err := NewManagedImageArtifact("Linux", "fakeResourceGroup", "fakeName", "fakeLocation", "fakeID")
+	artifact, err := NewManagedImageArtifact("Linux", "fakeResourceGroup", "fakeName", "fakeLocation", "fakeID", "fakeOsDiskSnapshotName", "fakeDataDiskSnapshotPrefix")
 	if err != nil {
 		t.Fatalf("err=%s", err)
 	}
 
-	expected := "fakeID"
+	expected := `Azure.ResourceManagement.VMImage:
 
-	result := artifact.Id()
+OSType: Linux
+ManagedImageResourceGroupName: fakeResourceGroup
+ManagedImageName: fakeName
+ManagedImageId: fakeID
+ManagedImageLocation: fakeLocation
+ManagedImageOSDiskSnapshotName: fakeOsDiskSnapshotName
+ManagedImageDataDiskSnapshotPrefix: fakeDataDiskSnapshotPrefix
+`
+
+	result := artifact.String()
+	if result != expected {
+		t.Fatalf("bad: %s", result)
+	}
+}
+
+func TestArtifactIDManagedImageWithoutOSDiskSnapshotName(t *testing.T) {
+	artifact, err := NewManagedImageArtifact("Linux", "fakeResourceGroup", "fakeName", "fakeLocation", "fakeID", "", "fakeDataDiskSnapshotPrefix")
+	if err != nil {
+		t.Fatalf("err=%s", err)
+	}
+
+	expected := `Azure.ResourceManagement.VMImage:
+
+OSType: Linux
+ManagedImageResourceGroupName: fakeResourceGroup
+ManagedImageName: fakeName
+ManagedImageId: fakeID
+ManagedImageLocation: fakeLocation
+ManagedImageDataDiskSnapshotPrefix: fakeDataDiskSnapshotPrefix
+`
+
+	result := artifact.String()
+	if result != expected {
+		t.Fatalf("bad: %s", result)
+	}
+}
+
+func TestArtifactIDManagedImageWithoutDataDiskSnapshotPrefix(t *testing.T) {
+	artifact, err := NewManagedImageArtifact("Linux", "fakeResourceGroup", "fakeName", "fakeLocation", "fakeID", "fakeOsDiskSnapshotName", "")
+	if err != nil {
+		t.Fatalf("err=%s", err)
+	}
+
+	expected := `Azure.ResourceManagement.VMImage:
+
+OSType: Linux
+ManagedImageResourceGroupName: fakeResourceGroup
+ManagedImageName: fakeName
+ManagedImageId: fakeID
+ManagedImageLocation: fakeLocation
+ManagedImageOSDiskSnapshotName: fakeOsDiskSnapshotName
+`
+
+	result := artifact.String()
 	if result != expected {
 		t.Fatalf("bad: %s", result)
 	}

builder/azure/arm/azure_client.go

@@ -40,6 +40,7 @@ type AzureClient struct {
 	common.VaultClient
 	armStorage.AccountsClient
 	compute.DisksClient
+	compute.SnapshotsClient
 
 	InspectorMaxLength int
 	Template           *CaptureTemplate
@@ -189,6 +190,12 @@ func NewAzureClient(subscriptionID, resourceGroupName, storageAccountName string
 	azureClient.VirtualMachinesClient.ResponseInspector = byConcatDecorators(byInspecting(maxlen), templateCapture(azureClient), errorCapture(azureClient))
 	azureClient.VirtualMachinesClient.UserAgent = fmt.Sprintf("%s %s", useragent.String(), azureClient.VirtualMachinesClient.UserAgent)
 
+	azureClient.SnapshotsClient = compute.NewSnapshotsClientWithBaseURI(cloud.ResourceManagerEndpoint, subscriptionID)
+	azureClient.SnapshotsClient.Authorizer = autorest.NewBearerAuthorizer(servicePrincipalToken)
+	azureClient.SnapshotsClient.RequestInspector = withInspection(maxlen)
+	azureClient.SnapshotsClient.ResponseInspector = byConcatDecorators(byInspecting(maxlen), errorCapture(azureClient))
+	azureClient.SnapshotsClient.UserAgent = fmt.Sprintf("%s %s", useragent.String(), azureClient.SnapshotsClient.UserAgent)
+
 	azureClient.AccountsClient = armStorage.NewAccountsClientWithBaseURI(cloud.ResourceManagerEndpoint, subscriptionID)
 	azureClient.AccountsClient.Authorizer = autorest.NewBearerAuthorizer(servicePrincipalToken)
 	azureClient.AccountsClient.RequestInspector = withInspection(maxlen)

builder/azure/arm/builder.go

@@ -183,6 +183,8 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 			NewStepGetOSDisk(azureClient, ui),
 			NewStepGetAdditionalDisks(azureClient, ui),
 			NewStepPowerOffCompute(azureClient, ui),
+			NewStepSnapshotOSDisk(azureClient, ui, b.config.isManagedImage()),
+			NewStepSnapshotDataDisks(azureClient, ui, b.config.isManagedImage()),
 			NewStepCaptureImage(azureClient, ui),
 			NewStepDeleteResourceGroup(azureClient, ui),
 			NewStepDeleteOSDisk(azureClient, ui),
@@ -218,6 +220,8 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 			&packerCommon.StepProvision{},
 			NewStepGetOSDisk(azureClient, ui),
 			NewStepGetAdditionalDisks(azureClient, ui),
+			NewStepSnapshotOSDisk(azureClient, ui, b.config.isManagedImage()),
+			NewStepSnapshotDataDisks(azureClient, ui, b.config.isManagedImage()),
 			NewStepPowerOffCompute(azureClient, ui),
 			NewStepCaptureImage(azureClient, ui),
 			NewStepDeleteResourceGroup(azureClient, ui),
@@ -259,7 +263,7 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 
 	if b.config.isManagedImage() {
 		managedImageID := fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Compute/images/%s", b.config.SubscriptionID, b.config.ManagedImageResourceGroupName, b.config.ManagedImageName)
-		return NewManagedImageArtifact(b.config.OSType, b.config.ManagedImageResourceGroupName, b.config.ManagedImageName, b.config.manageImageLocation, managedImageID)
+		return NewManagedImageArtifact(b.config.OSType, b.config.ManagedImageResourceGroupName, b.config.ManagedImageName, b.config.manageImageLocation, managedImageID, b.config.ManagedImageOSDiskSnapshotName, b.config.ManagedImageDataDiskSnapshotPrefix)
 	} else if template, ok := b.stateBag.GetOk(constants.ArmCaptureTemplate); ok {
 		return NewArtifact(
 			template.(*CaptureTemplate),
@@ -361,6 +365,8 @@ func (b *Builder) configureStateBag(stateBag multistep.StateBag) {
 	stateBag.Put(constants.ArmIsManagedImage, b.config.isManagedImage())
 	stateBag.Put(constants.ArmManagedImageResourceGroupName, b.config.ManagedImageResourceGroupName)
 	stateBag.Put(constants.ArmManagedImageName, b.config.ManagedImageName)
+	stateBag.Put(constants.ArmManagedImageOSDiskSnapshotName, b.config.ManagedImageOSDiskSnapshotName)
+	stateBag.Put(constants.ArmManagedImageDataDiskSnapshotPrefix, b.config.ManagedImageDataDiskSnapshotPrefix)
 	stateBag.Put(constants.ArmAsyncResourceGroupDelete, b.config.AsyncResourceGroupDelete)
 }
 

builder/azure/arm/config.go

@@ -56,6 +56,8 @@ var (
 	reCaptureNamePrefix    = regexp.MustCompile("^[A-Za-z0-9][A-Za-z0-9_\\-\\.]{0,23}$")
 	reManagedDiskName      = regexp.MustCompile(validManagedDiskName)
 	reResourceGroupName    = regexp.MustCompile(validResourceGroupNameRe)
+	reSnapshotName         = regexp.MustCompile("^[A-Za-z0-9_]{10,79}$")
+	reSnapshotPrefix       = regexp.MustCompile("^[A-Za-z0-9_]{10,59}$")
 )
 
 type PlanInformation struct {
@@ -107,6 +109,8 @@ type Config struct {
 	ManagedImageResourceGroupName      string `mapstructure:"managed_image_resource_group_name"`
 	ManagedImageName                   string `mapstructure:"managed_image_name"`
 	ManagedImageStorageAccountType     string `mapstructure:"managed_image_storage_account_type"`
+	ManagedImageOSDiskSnapshotName     string `mapstructure:"managed_image_os_disk_snapshot_name"`
+	ManagedImageDataDiskSnapshotPrefix string `mapstructure:"managed_image_data_disk_snapshot_prefix"`
 	managedImageStorageAccountType     compute.StorageAccountTypes
 	manageImageLocation                string
 
@@ -685,6 +689,18 @@ func assertRequiredParametersSet(c *Config, errs *packer.MultiError) {
 		}
 	}
 
+	if c.ManagedImageOSDiskSnapshotName != "" {
+		if ok, err := assertManagedImageOSDiskSnapshotName(c.ManagedImageOSDiskSnapshotName, "managed_image_os_disk_snapshot_name"); !ok {
+			errs = packer.MultiErrorAppend(errs, err)
+		}
+	}
+
+	if c.ManagedImageDataDiskSnapshotPrefix != "" {
+		if ok, err := assertManagedImageDataDiskSnapshotName(c.ManagedImageDataDiskSnapshotPrefix, "managed_image_data_disk_snapshot_prefix"); !ok {
+			errs = packer.MultiErrorAppend(errs, err)
+		}
+	}
+
 	if c.VirtualNetworkName == "" && c.VirtualNetworkResourceGroupName != "" {
 		errs = packer.MultiErrorAppend(errs, fmt.Errorf("If virtual_network_resource_group_name is specified, so must virtual_network_name"))
 	}
@@ -738,6 +754,20 @@ func assertManagedImageName(name, setting string) (bool, error) {
 	return true, nil
 }
 
+func assertManagedImageOSDiskSnapshotName(name, setting string) (bool, error) {
+	if !isValidAzureName(reSnapshotName, name) {
+		return false, fmt.Errorf("The setting %s must only contain characters from a-z, A-Z, 0-9 and _ and the maximum length is 80 characters", setting)
+	}
+	return true, nil
+}
+
+func assertManagedImageDataDiskSnapshotName(name, setting string) (bool, error) {
+	if !isValidAzureName(reSnapshotPrefix, name) {
+		return false, fmt.Errorf("The setting %s must only contain characters from a-z, A-Z, 0-9 and _ and the maximum length (excluding the prefix) is 60 characters", setting)
+	}
+	return true, nil
+}
+
 func assertResourceGroupName(rgn, setting string) (bool, error) {
 	if !isValidAzureName(reResourceGroupName, rgn) {
 		return false, fmt.Errorf("The setting %s must match the regular expression %q, and not end with a '-' or '.'.", setting, validResourceGroupNameRe)

builder/azure/arm/config_test.go

@@ -628,6 +628,107 @@ func TestConfigShouldRejectMalformedCaptureContainerName(t *testing.T) {
 	}
 }
 
+func TestConfigShouldRejectMalformedManagedImageOSDiskSnapshotName(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                         "ignore",
+		"image_publisher":                     "ignore",
+		"image_sku":                           "ignore",
+		"location":                            "ignore",
+		"subscription_id":                     "ignore",
+		"communicator":                        "none",
+		"managed_image_resource_group_name":   "ignore",
+		"managed_image_name":                  "ignore",
+		"managed_image_os_disk_snapshot_name": "ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	wellFormedManagedImageOSDiskSnapshotName := []string{
+		"AbcdefghijklmnopqrstuvwX",
+		"underscore_underscore",
+		"0leading_number",
+		"really_loooooooooooooooooooooooooooooooooooooooooooooooooong",
+	}
+
+	for _, x := range wellFormedManagedImageOSDiskSnapshotName {
+		config["managed_image_os_disk_snapshot_name"] = x
+		_, _, err := newConfig(config, getPackerConfiguration())
+
+		if err != nil {
+			t.Errorf("Expected test to pass, but it failed with the well-formed managed_image_os_disk_snapshot_name set to %q.", x)
+		}
+	}
+
+	malformedManagedImageOSDiskSnapshotName := []string{
+		"min_ten",
+		"-leading-hyphen",
+		"trailing-hyphen-",
+		"trailing-period.",
+		"punc-!@#$%^&*()_+-=-punc",
+		"really_looooooooooooooooooooooooooooooooooooooooooooooooooooooong_exceeding_80_char_limit",
+	}
+
+	for _, x := range malformedManagedImageOSDiskSnapshotName {
+		config["managed_image_os_disk_snapshot_name"] = x
+		_, _, err := newConfig(config, getPackerConfiguration())
+
+		if err == nil {
+			t.Errorf("Expected test to fail, but it succeeded with the malformed managed_image_os_disk_snapshot_name set to %q.", x)
+		}
+	}
+}
+
+func TestConfigShouldRejectMalformedManagedImageDataDiskSnapshotPrefix(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                       "ignore",
+		"image_publisher":                   "ignore",
+		"image_sku":                         "ignore",
+		"location":                          "ignore",
+		"subscription_id":                   "ignore",
+		"communicator":                      "none",
+		"managed_image_resource_group_name": "ignore",
+		"managed_image_name":                "ignore",
+		"managed_image_data_disk_snapshot_prefix": "ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	wellFormedManagedImageDataDiskSnapshotPrefix := []string{
+		"min_ten_chars",
+		"AbcdefghijklmnopqrstuvwX",
+		"underscore_underscore",
+		"0leading_number",
+		"less_than_sixty_characters",
+	}
+
+	for _, x := range wellFormedManagedImageDataDiskSnapshotPrefix {
+		config["managed_image_data_disk_snapshot_prefix"] = x
+		_, _, err := newConfig(config, getPackerConfiguration())
+
+		if err != nil {
+			t.Errorf("Expected test to pass, but it failed with the well-formed managed_image_data_disk_snapshot_prefix set to %q.", x)
+		}
+	}
+
+	malformedManagedImageDataDiskSnapshotPrefix := []string{
+		"more_ten",
+		"-leading-hyphen",
+		"trailing-hyphen-",
+		"trailing-period.",
+		"punc-!@#$%^&*()_+-=-punc",
+		"really_looooooooooooooooooooooooooooooooooooooooooooooooooooooong_exceeding_60_char_limit",
+	}
+
+	for _, x := range malformedManagedImageDataDiskSnapshotPrefix {
+		config["managed_image_data_disk_snapshot_prefix"] = x
+		_, _, err := newConfig(config, getPackerConfiguration())
+
+		if err == nil {
+			t.Errorf("Expected test to fail, but it succeeded with the malformed managed_image_data_disk_snapshot_prefix set to %q.", x)
+		}
+	}
+}
+
 func TestConfigShouldAcceptTags(t *testing.T) {
 	config := map[string]interface{}{
 		"capture_name_prefix":    "ignore",
@@ -790,6 +891,150 @@ func TestConfigShouldRejectMissingCustomDataFile(t *testing.T) {
 	}
 }
 
+func TestConfigShouldRejectManagedImageOSDiskSnapshotNameWithoutManagedImageName(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                         "ignore",
+		"image_publisher":                     "ignore",
+		"image_sku":                           "ignore",
+		"location":                            "ignore",
+		"subscription_id":                     "ignore",
+		"communicator":                        "none",
+		"managed_image_resource_group_name":   "ignore",
+		"managed_image_os_disk_snapshot_name": "ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	_, _, err := newConfig(config, getPackerConfiguration())
+	if err == nil {
+		t.Fatal("expected config to reject Managed Image build with OS disk snapshot name but without managed image name")
+	}
+}
+
+func TestConfigShouldRejectManagedImageOSDiskSnapshotNameWithoutManagedImageResourceGroupName(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                         "ignore",
+		"image_publisher":                     "ignore",
+		"image_sku":                           "ignore",
+		"location":                            "ignore",
+		"subscription_id":                     "ignore",
+		"communicator":                        "none",
+		"managed_image_name":                  "ignore",
+		"managed_image_os_disk_snapshot_name": "ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	_, _, err := newConfig(config, getPackerConfiguration())
+	if err == nil {
+		t.Fatal("expected config to reject Managed Image build with OS disk snapshot name but without managed image resource group name")
+	}
+}
+
+func TestConfigShouldRejectImageDataDiskSnapshotPrefixWithoutManagedImageName(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                       "ignore",
+		"image_publisher":                   "ignore",
+		"image_sku":                         "ignore",
+		"location":                          "ignore",
+		"subscription_id":                   "ignore",
+		"communicator":                      "none",
+		"managed_image_resource_group_name": "ignore",
+		"managed_image_data_disk_snapshot_prefix": "ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	_, _, err := newConfig(config, getPackerConfiguration())
+	if err == nil {
+		t.Fatal("expected config to reject Managed Image build with data disk snapshot prefix but without managed image name")
+	}
+}
+
+func TestConfigShouldRejectImageDataDiskSnapshotPrefixWithoutManagedImageResourceGroupName(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":        "ignore",
+		"image_publisher":    "ignore",
+		"image_sku":          "ignore",
+		"location":           "ignore",
+		"subscription_id":    "ignore",
+		"communicator":       "none",
+		"managed_image_name": "ignore",
+		"managed_image_data_disk_snapshot_prefix": "ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	_, _, err := newConfig(config, getPackerConfiguration())
+	if err == nil {
+		t.Fatal("expected config to reject Managed Image build with data disk snapshot prefix but without managed image resource group name")
+	}
+}
+
+func TestConfigShouldAcceptManagedImageOSDiskSnapshotNameAndManagedImageDataDiskSnapshotPrefix(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                             "ignore",
+		"image_publisher":                         "ignore",
+		"image_sku":                               "ignore",
+		"location":                                "ignore",
+		"subscription_id":                         "ignore",
+		"communicator":                            "none",
+		"managed_image_resource_group_name":       "ignore",
+		"managed_image_name":                      "ignore",
+		"managed_image_os_disk_snapshot_name":     "ignore_ignore",
+		"managed_image_data_disk_snapshot_prefix": "ignore_ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	_, _, err := newConfig(config, getPackerConfiguration())
+	if err != nil {
+		t.Fatal("expected config to accept platform managed image build")
+	}
+}
+
+func TestConfigShouldRejectManagedImageOSDiskSnapshotNameAndManagedImageDataDiskSnapshotPrefixWithCaptureContainerName(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                         "ignore",
+		"image_publisher":                     "ignore",
+		"image_sku":                           "ignore",
+		"location":                            "ignore",
+		"subscription_id":                     "ignore",
+		"communicator":                        "none",
+		"capture_container_name":              "ignore",
+		"managed_image_os_disk_snapshot_name": "ignore_ignore",
+		"managed_image_data_disk_snapshot_prefix": "ignore_ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	_, _, err := newConfig(config, getPackerConfiguration())
+	if err == nil {
+		t.Fatal("expected config to reject Managed Image build with data disk snapshot prefix and OS disk snapshot name with capture container name")
+	}
+}
+
+func TestConfigShouldRejectManagedImageOSDiskSnapshotNameAndManagedImageDataDiskSnapshotPrefixWithCaptureNamePrefix(t *testing.T) {
+	config := map[string]interface{}{
+		"image_offer":                         "ignore",
+		"image_publisher":                     "ignore",
+		"image_sku":                           "ignore",
+		"location":                            "ignore",
+		"subscription_id":                     "ignore",
+		"communicator":                        "none",
+		"capture_name_prefix":                 "ignore",
+		"managed_image_os_disk_snapshot_name": "ignore_ignore",
+		"managed_image_data_disk_snapshot_prefix": "ignore_ignore",
+		// Does not matter for this test case, just pick one.
+		"os_type": constants.Target_Linux,
+	}
+
+	_, _, err := newConfig(config, getPackerConfiguration())
+	if err == nil {
+		t.Fatal("expected config to reject Managed Image build with data disk snapshot prefix and OS disk snapshot name with capture name prefix")
+	}
+}
+
 func TestConfigShouldAcceptPlatformManagedImageBuild(t *testing.T) {
 	config := map[string]interface{}{
 		"image_offer":                       "ignore",

builder/azure/arm/step_snapshot_data_disks.go

@@ -0,0 +1,102 @@
+package arm
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/go-autorest/autorest/to"
+	"github.com/hashicorp/packer/builder/azure/common/constants"
+	"github.com/hashicorp/packer/helper/multistep"
+	"github.com/hashicorp/packer/packer"
+)
+
+type StepSnapshotDataDisks struct {
+	client         *AzureClient
+	create         func(ctx context.Context, resourceGroupName string, srcUriVhd string, location string, tags map[string]*string, dstSnapshotName string) error
+	say            func(message string)
+	error          func(e error)
+	isManagedImage bool
+}
+
+func NewStepSnapshotDataDisks(client *AzureClient, ui packer.Ui, isManagedImage bool) *StepSnapshotDataDisks {
+	var step = &StepSnapshotDataDisks{
+		client:         client,
+		say:            func(message string) { ui.Say(message) },
+		error:          func(e error) { ui.Error(e.Error()) },
+		isManagedImage: isManagedImage,
+	}
+
+	step.create = step.createDataDiskSnapshot
+	return step
+}
+
+func (s *StepSnapshotDataDisks) createDataDiskSnapshot(ctx context.Context, resourceGroupName string, srcUriVhd string, location string, tags map[string]*string, dstSnapshotName string) error {
+
+	srcVhdToSnapshot := compute.Snapshot{
+		DiskProperties: &compute.DiskProperties{
+			CreationData: &compute.CreationData{
+				CreateOption:     compute.Copy,
+				SourceResourceID: to.StringPtr(srcUriVhd),
+			},
+		},
+		Location: to.StringPtr(location),
+		Tags:     tags,
+	}
+
+	f, err := s.client.SnapshotsClient.CreateOrUpdate(ctx, resourceGroupName, dstSnapshotName, srcVhdToSnapshot)
+
+	if err != nil {
+		s.say(s.client.LastError.Error())
+		return err
+	}
+
+	err = f.WaitForCompletion(ctx, s.client.SnapshotsClient.Client)
+
+	if err != nil {
+		s.say(s.client.LastError.Error())
+		return err
+	}
+
+	createdSnapshot, err := f.Result(s.client.SnapshotsClient)
+
+	if err != nil {
+		s.say(s.client.LastError.Error())
+		return err
+	}
+
+	s.say(fmt.Sprintf(" -> Managed Image Data Disk Snapshot		: '%s'", *(createdSnapshot.ID)))
+
+	return nil
+}
+
+func (s *StepSnapshotDataDisks) Run(ctx context.Context, stateBag multistep.StateBag) multistep.StepAction {
+	if s.isManagedImage {
+
+		s.say("Taking snapshot of data disk ...")
+
+		var resourceGroupName = stateBag.Get(constants.ArmManagedImageResourceGroupName).(string)
+		var location = stateBag.Get(constants.ArmLocation).(string)
+		var tags = stateBag.Get(constants.ArmTags).(map[string]*string)
+		var additionalDisks = stateBag.Get(constants.ArmAdditionalDiskVhds).([]string)
+		var dstSnapshotPrefix = stateBag.Get(constants.ArmManagedImageDataDiskSnapshotPrefix).(string)
+
+		for i, disk := range additionalDisks {
+			dstSnapshotName := dstSnapshotPrefix + strconv.Itoa(i)
+			err := s.create(ctx, resourceGroupName, disk, location, tags, dstSnapshotName)
+
+			if err != nil {
+				stateBag.Put(constants.Error, err)
+				s.error(err)
+
+				return multistep.ActionHalt
+			}
+		}
+	}
+
+	return multistep.ActionContinue
+}
+
+func (*StepSnapshotDataDisks) Cleanup(multistep.StateBag) {
+}

builder/azure/arm/step_snapshot_data_disks_test.go

@@ -0,0 +1,71 @@
+package arm
+
+import (
+	"context"
+	"fmt"
+	"github.com/hashicorp/packer/builder/azure/common/constants"
+	"github.com/hashicorp/packer/helper/multistep"
+	"testing"
+)
+
+func TestStepSnapshotDataDisksShouldFailIfSnapshotFails(t *testing.T) {
+	var testSubject = &StepSnapshotDataDisks{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return fmt.Errorf("!! Unit Test FAIL !!")
+		},
+		say:            func(message string) {},
+		error:          func(e error) {},
+		isManagedImage: true,
+	}
+
+	stateBag := createTestStateBagStepSnapshotDataDisks()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionHalt {
+		t.Fatalf("Expected the step to return 'ActionHalt', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == false {
+		t.Fatalf("Expected the step to set stateBag['%s'], but it was not.", constants.Error)
+	}
+}
+
+func TestStepSnapshotDataDisksShouldPassIfSnapshotPasses(t *testing.T) {
+	var testSubject = &StepSnapshotDataDisks{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return nil
+		},
+		say:            func(message string) {},
+		error:          func(e error) {},
+		isManagedImage: true,
+	}
+
+	stateBag := createTestStateBagStepSnapshotDataDisks()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionContinue {
+		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == true {
+		t.Fatalf("Expected the step to not set stateBag['%s'], but it was.", constants.Error)
+	}
+}
+
+func createTestStateBagStepSnapshotDataDisks() multistep.StateBag {
+	stateBag := new(multistep.BasicStateBag)
+
+	stateBag.Put(constants.ArmManagedImageResourceGroupName, "Unit Test: ResourceGroupName")
+	stateBag.Put(constants.ArmLocation, "Unit Test: Location")
+
+	value := "Unit Test: Tags"
+	tags := map[string]*string{
+		"tag01": &value,
+	}
+	stateBag.Put(constants.ArmTags, tags)
+
+	stateBag.Put(constants.ArmAdditionalDiskVhds, []string{"subscriptions/123-456-789/resourceGroups/existingresourcegroup/providers/Microsoft.Compute/disks/osdisk"})
+	stateBag.Put(constants.ArmManagedImageDataDiskSnapshotPrefix, "Unit Test: ManagedImageDataDiskSnapshotPrefix")
+
+	return stateBag
+}

builder/azure/arm/step_snapshot_os_disk.go

@@ -0,0 +1,97 @@
+package arm
+
+import (
+	"context"
+	"fmt"
+	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2018-04-01/compute"
+	"github.com/Azure/go-autorest/autorest/to"
+	"github.com/hashicorp/packer/builder/azure/common/constants"
+	"github.com/hashicorp/packer/helper/multistep"
+	"github.com/hashicorp/packer/packer"
+)
+
+type StepSnapshotOSDisk struct {
+	client         *AzureClient
+	create         func(ctx context.Context, resourceGroupName string, srcUriVhd string, location string, tags map[string]*string, dstSnapshotName string) error
+	say            func(message string)
+	error          func(e error)
+	isManagedImage bool
+}
+
+func NewStepSnapshotOSDisk(client *AzureClient, ui packer.Ui, isManagedImage bool) *StepSnapshotOSDisk {
+	var step = &StepSnapshotOSDisk{
+		client:         client,
+		say:            func(message string) { ui.Say(message) },
+		error:          func(e error) { ui.Error(e.Error()) },
+		isManagedImage: isManagedImage,
+	}
+
+	step.create = step.createSnapshot
+	return step
+}
+
+func (s *StepSnapshotOSDisk) createSnapshot(ctx context.Context, resourceGroupName string, srcUriVhd string, location string, tags map[string]*string, dstSnapshotName string) error {
+
+	srcVhdToSnapshot := compute.Snapshot{
+		DiskProperties: &compute.DiskProperties{
+			CreationData: &compute.CreationData{
+				CreateOption:     compute.Copy,
+				SourceResourceID: to.StringPtr(srcUriVhd),
+			},
+		},
+		Location: to.StringPtr(location),
+		Tags:     tags,
+	}
+
+	f, err := s.client.SnapshotsClient.CreateOrUpdate(ctx, resourceGroupName, dstSnapshotName, srcVhdToSnapshot)
+
+	if err != nil {
+		s.say(s.client.LastError.Error())
+		return err
+	}
+
+	err = f.WaitForCompletion(ctx, s.client.SnapshotsClient.Client)
+
+	if err != nil {
+		s.say(s.client.LastError.Error())
+		return err
+	}
+
+	createdSnapshot, err := f.Result(s.client.SnapshotsClient)
+
+	if err != nil {
+		s.say(s.client.LastError.Error())
+		return err
+	}
+
+	s.say(fmt.Sprintf(" -> Managed Image OS Disk Snapshot		: '%s'", *(createdSnapshot.ID)))
+
+	return nil
+}
+
+func (s *StepSnapshotOSDisk) Run(ctx context.Context, stateBag multistep.StateBag) multistep.StepAction {
+	if s.isManagedImage {
+
+		s.say("Taking snapshot of OS disk ...")
+
+		var resourceGroupName = stateBag.Get(constants.ArmManagedImageResourceGroupName).(string)
+		var location = stateBag.Get(constants.ArmLocation).(string)
+		var tags = stateBag.Get(constants.ArmTags).(map[string]*string)
+		var srcUriVhd = stateBag.Get(constants.ArmOSDiskVhd).(string)
+		var dstSnapshotName = stateBag.Get(constants.ArmManagedImageOSDiskSnapshotName).(string)
+
+		err := s.create(ctx, resourceGroupName, srcUriVhd, location, tags, dstSnapshotName)
+
+		if err != nil {
+			stateBag.Put(constants.Error, err)
+			s.error(err)
+
+			return multistep.ActionHalt
+		}
+	}
+
+	return multistep.ActionContinue
+}
+
+func (*StepSnapshotOSDisk) Cleanup(multistep.StateBag) {
+}

builder/azure/arm/step_snapshot_os_disk_test.go

@@ -0,0 +1,72 @@
+package arm
+
+import (
+	"context"
+	"fmt"
+	"github.com/hashicorp/packer/builder/azure/common/constants"
+	"github.com/hashicorp/packer/helper/multistep"
+	"testing"
+)
+
+func TestStepSnapshotOSDiskShouldFailIfSnapshotFails(t *testing.T) {
+	var testSubject = &StepSnapshotOSDisk{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return fmt.Errorf("!! Unit Test FAIL !!")
+		},
+		say:            func(message string) {},
+		error:          func(e error) {},
+		isManagedImage: true,
+	}
+
+	stateBag := createTestStateBagStepSnapshotOSDisk()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionHalt {
+		t.Fatalf("Expected the step to return 'ActionHalt', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == false {
+		t.Fatalf("Expected the step to set stateBag['%s'], but it was not.", constants.Error)
+	}
+}
+
+func TestStepSnapshotOSDiskShouldPassIfSnapshotPasses(t *testing.T) {
+	var testSubject = &StepSnapshotOSDisk{
+		create: func(context.Context, string, string, string, map[string]*string, string) error {
+			return nil
+		},
+		say:            func(message string) {},
+		error:          func(e error) {},
+		isManagedImage: true,
+	}
+
+	stateBag := createTestStateBagStepSnapshotOSDisk()
+
+	var result = testSubject.Run(context.Background(), stateBag)
+	if result != multistep.ActionContinue {
+		t.Fatalf("Expected the step to return 'ActionContinue', but got '%d'.", result)
+	}
+
+	if _, ok := stateBag.GetOk(constants.Error); ok == true {
+		t.Fatalf("Expected the step to not set stateBag['%s'], but it was.", constants.Error)
+	}
+}
+
+func createTestStateBagStepSnapshotOSDisk() multistep.StateBag {
+	stateBag := new(multistep.BasicStateBag)
+
+	stateBag.Put(constants.ArmManagedImageResourceGroupName, "Unit Test: ResourceGroupName")
+	stateBag.Put(constants.ArmLocation, "Unit Test: Location")
+
+	value := "Unit Test: Tags"
+	tags := map[string]*string{
+		"tag01": &value,
+	}
+
+	stateBag.Put(constants.ArmTags, tags)
+
+	stateBag.Put(constants.ArmOSDiskVhd, "subscriptions/123-456-789/resourceGroups/existingresourcegroup/providers/Microsoft.Compute/disks/osdisk")
+	stateBag.Put(constants.ArmManagedImageOSDiskSnapshotName, "Unit Test: ManagedImageOSDiskSnapshotName")
+
+	return stateBag
+}

builder/azure/common/constants/stateBag.go

@@ -35,4 +35,6 @@ const (
 	ArmManagedImageLocation               string = "arm.ManagedImageLocation"
 	ArmManagedImageName                   string = "arm.ManagedImageName"
 	ArmAsyncResourceGroupDelete           string = "arm.AsyncResourceGroupDelete"
+	ArmManagedImageOSDiskSnapshotName     string = "arm.ManagedImageOSDiskSnapshotName"
+	ArmManagedImageDataDiskSnapshotPrefix string = "arm.ManagedImageDataDiskSnapshotPrefix"
 )

website/source/docs/builders/azure.html.md

@@ -306,6 +306,14 @@ Providing `temp_resource_group_name` or `location` in combination with
     value and defaults to false. **Important** Setting this true means that
     your builds are faster, however any failed deletes are not reported.
 
+-   `managed_image_os_disk_snapshot_name` (string) If managed\_image\_os\_disk\_snapshot\_name 
+    is set, a snapshot of the OS disk is created with the same name as this value before the 
+    VM is captured.
+
+-   `managed_image_data_disk_snapshot_prefix` (string)  If managed\_image\_data\_disk\_snapshot\_prefix 
+    is set, snapshot of the data disk(s) is created with the same prefix as this value before the VM 
+    is captured.
+
 ## Basic Example
 
 Here is a basic example for Azure.