From d23ad907c03f351dd672436cf154a638d59ac2e0 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 30 Jul 2013 18:28:21 -0700
Subject: [PATCH] builder/amazon/chroot: register AMI

---
 builder/amazon/chroot/builder.go           |  1 +
 builder/amazon/chroot/step_register_ami.go | 67 ++++++++++++++++++++++
 2 files changed, 68 insertions(+)
 create mode 100644 builder/amazon/chroot/step_register_ami.go

diff --git a/builder/amazon/chroot/builder.go b/builder/amazon/chroot/builder.go
index 82338cb2b..75dcaed3b 100644
--- a/builder/amazon/chroot/builder.go
+++ b/builder/amazon/chroot/builder.go
@@ -140,6 +140,7 @@ func (b *Builder) Run(ui packer.Ui, hook packer.Hook, cache packer.Cache) (packe
 		&StepChrootProvision{},
 		&StepEarlyCleanup{},
 		&StepSnapshot{},
+		&StepRegisterAMI{},
 	}
 
 	// Run!
diff --git a/builder/amazon/chroot/step_register_ami.go b/builder/amazon/chroot/step_register_ami.go
new file mode 100644
index 000000000..9ec70f038
--- /dev/null
+++ b/builder/amazon/chroot/step_register_ami.go
@@ -0,0 +1,67 @@
+package chroot
+
+import (
+	"fmt"
+	"github.com/mitchellh/goamz/ec2"
+	"github.com/mitchellh/multistep"
+	awscommon "github.com/mitchellh/packer/builder/amazon/common"
+	"github.com/mitchellh/packer/packer"
+)
+
+// StepRegisterAMI creates the AMI.
+type StepRegisterAMI struct{}
+
+func (s *StepRegisterAMI) Run(state map[string]interface{}) multistep.StepAction {
+	ec2conn := state["ec2"].(*ec2.EC2)
+	image := state["source_image"].(*ec2.Image)
+	snapshotId := state["snapshot_id"].(string)
+	ui := state["ui"].(packer.Ui)
+
+	amiName := "foo"
+
+	ui.Say("Registering the AMI...")
+	blockDevices := make([]ec2.BlockDeviceMapping, len(image.BlockDevices))
+	for i, device := range image.BlockDevices {
+		newDevice := device
+		if newDevice.DeviceName == image.RootDeviceName {
+			newDevice.SnapshotId = snapshotId
+		}
+
+		blockDevices[i] = newDevice
+	}
+
+	registerOpts := &ec2.RegisterImage{
+		Name:           amiName,
+		Architecture:   image.Architecture,
+		KernelId:       image.KernelId,
+		RamdiskId:      image.RamdiskId,
+		RootDeviceName: image.RootDeviceName,
+		BlockDevices:   blockDevices,
+	}
+
+	registerResp, err := ec2conn.RegisterImage(registerOpts)
+	if err != nil {
+		state["error"] = fmt.Errorf("Error registering AMI: %s", err)
+		ui.Error(state["error"].(error).Error())
+		return multistep.ActionHalt
+	}
+
+	// Set the AMI ID in the state
+	ui.Say(fmt.Sprintf("AMI: %s", registerResp.ImageId))
+	amis := make(map[string]string)
+	amis[ec2conn.Region.Name] = registerResp.ImageId
+	state["amis"] = amis
+
+	// Wait for the image to become ready
+	ui.Say("Waiting for AMI to become ready...")
+	if err := awscommon.WaitForAMI(ec2conn, registerResp.ImageId); err != nil {
+		err := fmt.Errorf("Error waiting for AMI: %s", err)
+		state["error"] = err
+		ui.Error(err.Error())
+		return multistep.ActionHalt
+	}
+
+	return multistep.ActionContinue
+}
+
+func (s *StepRegisterAMI) Cleanup(state map[string]interface{}) {}