From c3131ec3faa91d66aa5251386f234e2f0f0a1870 Mon Sep 17 00:00:00 2001
From: Jeffrey Ellin <jeff@ellin.com>
Date: Fri, 15 Jan 2016 15:10:03 -0500
Subject: [PATCH] Fixes #1898. Add parameter to pass private key path to use
 when creating a machine.

It is assumed that the corresponding public certificate will be the same file name + .pub

Signed-off-by: Jeffrey Ellin <jeff@ellin.com>
---
 docs/drivers/aws.md                           |  2 ++
 drivers/amazonec2/amazonec2.go                | 29 +++++++++++++++--
 .../amazonec2/createwithkeypair.bats          | 32 +++++++++++++++++++
 3 files changed, 60 insertions(+), 3 deletions(-)
 create mode 100644 test/integration/amazonec2/createwithkeypair.bats

diff --git a/docs/drivers/aws.md b/docs/drivers/aws.md
index aae8303dfd..04ddd624d1 100644
--- a/docs/drivers/aws.md
+++ b/docs/drivers/aws.md
@@ -68,6 +68,7 @@ You can use environment variables:
 -   `--amazonec2-private-address-only`: Use the private IP address only.
 -   `--amazonec2-monitoring`: Enable CloudWatch Monitoring.
 -   `--amazonec2-use-ebs-optimized-instance`: Create an EBS Optimized Instance, instance type must support it.
+-   `--amazonec2-ssh-keypath`: Path to Private Key file to use for instance. Matching public key with .pub extension should exist
 
 By default, the Amazon EC2 driver will use a daily image of Ubuntu 15.10 LTS.
 
@@ -111,6 +112,7 @@ Environment variables and default values:
 | `--amazonec2-private-address-only`       | -                       | `false`          |
 | `--amazonec2-monitoring`                 | -                       | `false`          |
 | `--amazonec2-use-ebs-optimized-instance` | -                       | `false`          |
+| `--amazonec2-ssh-keypath`                |`AWS_SSH_KEYPATH`        | -                |
 
 ## Security Group
 Note that a security group will be created and associated to the host. This security group will have the following ports opened inbound:
diff --git a/drivers/amazonec2/amazonec2.go b/drivers/amazonec2/amazonec2.go
index 8f2dbc5e01..89fd809a35 100644
--- a/drivers/amazonec2/amazonec2.go
+++ b/drivers/amazonec2/amazonec2.go
@@ -86,6 +86,7 @@ type Driver struct {
 	UsePrivateIP            bool
 	UseEbsOptimizedInstance bool
 	Monitoring              bool
+	SSHPrivateKeyPath       string
 }
 
 type clientFactory interface {
@@ -207,6 +208,11 @@ func (d *Driver) GetCreateFlags() []mcnflag.Flag {
 			Name:  "amazonec2-use-ebs-optimized-instance",
 			Usage: "Create an EBS optimized instance",
 		},
+		mcnflag.StringFlag{
+			Name:   "amazonec2-ssh-keypath",
+			Usage:  "SSH Key for Instance",
+			EnvVar: "AWS_SSH_KEYPATH",
+		},
 	}
 }
 
@@ -283,6 +289,7 @@ func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
 	d.UsePrivateIP = flags.Bool("amazonec2-use-private-address")
 	d.Monitoring = flags.Bool("amazonec2-monitoring")
 	d.UseEbsOptimizedInstance = flags.Bool("amazonec2-use-ebs-optimized-instance")
+	d.SSHPrivateKeyPath = flags.String("amazonec2-ssh-keypath")
 	d.SetSwarmConfigFromFlags(flags)
 
 	if d.AccessKey == "" && d.SecretKey == "" {
@@ -738,11 +745,27 @@ func (d *Driver) waitForInstance() error {
 }
 
 func (d *Driver) createKeyPair() error {
-	if err := ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
-		return err
+
+	keyPath := ""
+
+	if d.SSHPrivateKeyPath == "" {
+		log.Debugf("Creating New SSH Key")
+		if err := ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
+			return err
+		}
+		keyPath = d.GetSSHKeyPath()
+	} else {
+		log.Debugf("Using ExistingKeyPair: %s", d.SSHPrivateKeyPath)
+		if err := mcnutils.CopyFile(d.SSHPrivateKeyPath, d.GetSSHKeyPath()); err != nil {
+			return err
+		}
+		if err := mcnutils.CopyFile(d.SSHPrivateKeyPath+".pub", d.GetSSHKeyPath()+".pub"); err != nil {
+			return err
+		}
+		keyPath = d.SSHPrivateKeyPath
 	}
 
-	publicKey, err := ioutil.ReadFile(d.GetSSHKeyPath() + ".pub")
+	publicKey, err := ioutil.ReadFile(keyPath + ".pub")
 	if err != nil {
 		return err
 	}
diff --git a/test/integration/amazonec2/createwithkeypair.bats b/test/integration/amazonec2/createwithkeypair.bats
new file mode 100644
index 0000000000..c62486c235
--- /dev/null
+++ b/test/integration/amazonec2/createwithkeypair.bats
@@ -0,0 +1,32 @@
+#!/usr/bin/env bats
+
+load ${BASE_TEST_DIR}/helpers.bash
+
+only_if_env DRIVER amazonec2
+
+use_disposable_machine
+
+require_env AWS_ACCESS_KEY_ID
+
+require_env AWS_SECRET_ACCESS_KEY
+
+export AWS_SSH_DIR="$MACHINE_STORAGE_PATH/mcnkeys"
+
+export AWS_SSH_KEYPATH=$AWS_SSH_DIR/id_rsa
+
+@test "$DRIVER: Should Create Instance with Pre existing SSH Key" {
+
+	mkdir -p $AWS_SSH_DIR
+
+	run ssh-keygen -f $AWS_SSH_KEYPATH -t rsa -N ''
+
+	machine create -d amazonec2 $NAME
+	
+	run diff $AWS_SSH_KEYPATH $MACHINE_STORAGE_PATH/machines/$NAME/id_rsa
+	[[ $output == "" ]]
+
+	run diff $AWS_SSH_KEYPATH.pub $MACHINE_STORAGE_PATH/machines/$NAME/id_rsa.pub
+	[[ $output == "" ]]
+
+
+}
\ No newline at end of file