From b0394c6d30308e48dd2d431b6f7a11339d11607f Mon Sep 17 00:00:00 2001
From: James Massara <james_massara@intuit.com>
Date: Thu, 8 Aug 2013 16:04:38 -0700
Subject: [PATCH] builder/amazon/all: Added support for setting attributes on
 the AMI

---
 .../amazon/common/step_modify_attributes.go   | 72 +++++++++++++++++++
 builder/amazon/ebs/builder.go                 | 10 +++
 builder/amazon/instance/builder.go            | 12 ++++
 3 files changed, 94 insertions(+)
 create mode 100644 builder/amazon/common/step_modify_attributes.go

diff --git a/builder/amazon/common/step_modify_attributes.go b/builder/amazon/common/step_modify_attributes.go
new file mode 100644
index 000000000..6833a96d8
--- /dev/null
+++ b/builder/amazon/common/step_modify_attributes.go
@@ -0,0 +1,72 @@
+package common
+
+import (
+	"fmt"
+	"github.com/mitchellh/goamz/ec2"
+	"github.com/mitchellh/multistep"
+	"github.com/mitchellh/packer/packer"
+)
+
+type StepModifyAttributes struct {
+	Users        []string
+	Groups       []string
+	ProductCodes []string
+	Description  string
+}
+
+func (s *StepModifyAttributes) Run(state map[string]interface{}) multistep.StepAction {
+	ec2conn := state["ec2"].(*ec2.EC2)
+	ui := state["ui"].(packer.Ui)
+	amis := state["amis"].(map[string]string)
+	ami := amis[ec2conn.Region.Name]
+
+	if s.Description != "" {
+		ui.Say(fmt.Sprintf("Setting Description of AMI (%s) to '%s'...", ami, s.Description))
+		_, err := ec2conn.ModifyImageAttribute(ami, &ec2.ModifyImageAttribute{
+			Attribute:   ec2.DescriptionAttribute,
+			Description: s.Description,
+		})
+		if err != nil {
+			err := fmt.Errorf("Error setting Description of AMI (%s): %s", ami, err)
+			state["error"] = err
+			ui.Error(err.Error())
+			return multistep.ActionHalt
+		}
+	}
+
+	if len(s.Users) > 0 || len(s.Groups) > 0 {
+		ui.Say(fmt.Sprintf("Setting Launch Permissions for AMI (%s)...", ami))
+		_, err := ec2conn.ModifyImageAttribute(ami, &ec2.ModifyImageAttribute{
+			Attribute: ec2.LaunchPermissionAttribute,
+			Operation: ec2.LaunchPermissionAdd,
+			Users:     s.Users,
+			Groups:    s.Groups,
+		})
+		if err != nil {
+			err := fmt.Errorf("Error setting Launch Permissions for AMI (%s): %s", ami, err)
+			state["error"] = err
+			ui.Error(err.Error())
+			return multistep.ActionHalt
+		}
+	}
+
+	if len(s.ProductCodes) > 0 {
+		ui.Say(fmt.Sprintf("Setting Product Code(s) for AMI (%s)...", ami))
+		_, err := ec2conn.ModifyImageAttribute(ami, &ec2.ModifyImageAttribute{
+			Attribute:    ec2.ProductCodeAttribute,
+			ProductCodes: s.ProductCodes,
+		})
+		if err != nil {
+			err := fmt.Errorf("Error setting Product Code(s) for AMI (%s): %s", ami, err)
+			state["error"] = err
+			ui.Error(err.Error())
+			return multistep.ActionHalt
+		}
+	}
+
+	return multistep.ActionContinue
+}
+
+func (s *StepModifyAttributes) Cleanup(state map[string]interface{}) {
+	// No cleanup...
+}
diff --git a/builder/amazon/ebs/builder.go b/builder/amazon/ebs/builder.go
index ac0707a99..860f05c08 100644
--- a/builder/amazon/ebs/builder.go
+++ b/builder/amazon/ebs/builder.go
@@ -27,6 +27,11 @@ type config struct {
 	// Configuration of the resulting AMI
 	AMIName string `mapstructure:"ami_name"`
 
+	// AMI attributes
+	AMIDescription string   `mapstructure:"ami_description"`
+	AMIUsers       []string `mapstructure:"ami_users"`
+	AMIGroups      []string `mapstructure:"ami_groups"`
+
 	// Tags for the AMI
 	Tags map[string]string
 
@@ -136,6 +141,11 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 		&stepStopInstance{},
 		&stepCreateAMI{},
 		&awscommon.StepCreateTags{Tags: b.config.Tags},
+		&awscommon.StepModifyAttributes{
+			Description: b.config.AMIDescription,
+			Users:       b.config.AMIUsers,
+			Groups:      b.config.AMIGroups,
+		},
 	}
 
 	// Run!
diff --git a/builder/amazon/instance/builder.go b/builder/amazon/instance/builder.go
index a3118260e..971dcfcfb 100644
--- a/builder/amazon/instance/builder.go
+++ b/builder/amazon/instance/builder.go
@@ -37,6 +37,12 @@ type Config struct {
 	X509KeyPath         string `mapstructure:"x509_key_path"`
 	X509UploadPath      string `mapstructure:"x509_upload_path"`
 
+	// AMI attributes
+	AMIDescription  string   `mapstructure:"ami_description"`
+	AMIUsers        []string `mapstructure:"ami_users"`
+	AMIGroups       []string `mapstructure:"ami_groups"`
+	AMIProductCodes []string `mapstructure:"ami_product_codes"`
+
 	tpl *common.Template
 }
 
@@ -210,6 +216,12 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 		&StepUploadBundle{},
 		&StepRegisterAMI{},
 		&awscommon.StepCreateTags{Tags: b.config.Tags},
+		&awscommon.StepModifyAttributes{
+			Description:  b.config.AMIDescription,
+			Users:        b.config.AMIUsers,
+			Groups:       b.config.AMIGroups,
+			ProductCodes: b.config.AMIProductCodes,
+		},
 	}
 
 	// Run!