From 50cc9e51f9dd13f3420c21d2e4ea349ec0be7981 Mon Sep 17 00:00:00 2001
From: v4n <andre@vyanes.com>
Date: Tue, 1 Nov 2016 15:15:25 -0700
Subject: [PATCH] Digital Ocean] Add ability to speicify the private SSH key
 path
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: ΛNDΓΕ <andre@vyanes.com>
---
 drivers/digitalocean/digitalocean.go | 52 +++++++++++++++++++++++++---
 1 file changed, 48 insertions(+), 4 deletions(-)

diff --git a/drivers/digitalocean/digitalocean.go b/drivers/digitalocean/digitalocean.go
index 899aa65c..0a8c32a2 100644
--- a/drivers/digitalocean/digitalocean.go
+++ b/drivers/digitalocean/digitalocean.go
@@ -5,12 +5,14 @@ import (
 	"io/ioutil"
 	"net"
 	"os"
+	"path"
 	"time"
 
 	"github.com/digitalocean/godo"
 	"github.com/docker/machine/libmachine/drivers"
 	"github.com/docker/machine/libmachine/log"
 	"github.com/docker/machine/libmachine/mcnflag"
+	"github.com/docker/machine/libmachine/mcnutils"
 	"github.com/docker/machine/libmachine/ssh"
 	"github.com/docker/machine/libmachine/state"
 	"golang.org/x/oauth2"
@@ -25,6 +27,7 @@ type Driver struct {
 	Region            string
 	SSHKeyID          int
 	SSHKeyFingerprint string
+	SSHKey            string
 	Size              string
 	IPv6              bool
 	Backups           bool
@@ -60,6 +63,11 @@ func (d *Driver) GetCreateFlags() []mcnflag.Flag {
 			Name:   "digitalocean-ssh-key-fingerprint",
 			Usage:  "SSH key fingerprint",
 		},
+		mcnflag.StringFlag{
+			EnvVar: "DIGITALOCEAN_SSH_KEY_PATH",
+			Name:   "digitalocean-ssh-key-path",
+			Usage:  "SSH private key path ",
+		},
 		mcnflag.IntFlag{
 			EnvVar: "DIGITALOCEAN_SSH_PORT",
 			Name:   "digitalocean-ssh-port",
@@ -140,6 +148,8 @@ func (d *Driver) SetConfigFromFlags(flags drivers.DriverOptions) error {
 	d.SSHUser = flags.String("digitalocean-ssh-user")
 	d.SSHPort = flags.Int("digitalocean-ssh-port")
 	d.SSHKeyFingerprint = flags.String("digitalocean-ssh-key-fingerprint")
+	d.SSHKey = flags.String("digitalocean-ssh-key-path")
+
 	d.SetSwarmConfigFromFlags(flags)
 
 	if d.AccessToken == "" {
@@ -156,6 +166,16 @@ func (d *Driver) PreCreateCheck() error {
 		}
 	}
 
+	if d.SSHKey != "" {
+		if d.SSHKeyFingerprint == "" {
+			return fmt.Errorf("ssh-key-fingerpint needs to be provided for %q", d.SSHKey)
+		}
+
+		if _, err := os.Stat(d.SSHKey); os.IsNotExist(err) {
+			return fmt.Errorf("SSH key does not exist: %q", d.SSHKey)
+		}
+	}
+
 	client := d.getClient()
 	regions, _, err := client.Regions.List(nil)
 	if err != nil {
@@ -239,15 +259,26 @@ func (d *Driver) Create() error {
 }
 
 func (d *Driver) createSSHKey() (*godo.Key, error) {
+	d.SSHKeyPath = d.GetSSHKeyPath()
+
 	if d.SSHKeyFingerprint != "" {
 		key, resp, err := d.getClient().Keys.GetByFingerprint(d.SSHKeyFingerprint)
 		if err != nil && resp.StatusCode == 404 {
 			return nil, fmt.Errorf("Digital Ocean SSH key with fingerprint %s doesn't exist", d.SSHKeyFingerprint)
 		}
-		return key, err
+
+		if d.SSHKey == "" {
+			log.Infof("Assuming Digital Ocean private SSH is located at ~/.ssh/id_rsa")
+			return key, nil
+		}
+
+		if err := copySSHKey(d.SSHKey, d.SSHKeyPath); err != nil {
+			return nil, err
+		}
+		return key, nil
 	}
 
-	if err := ssh.GenerateSSHKey(d.GetSSHKeyPath()); err != nil {
+	if err := ssh.GenerateSSHKey(d.SSHKeyPath); err != nil {
 		return nil, err
 	}
 
@@ -348,8 +379,9 @@ func (d *Driver) getClient() *godo.Client {
 }
 
 func (d *Driver) GetSSHKeyPath() string {
-	// don't set SSHKeyPath when using an existing key fingerprint
-	if d.SSHKeyPath == "" && d.SSHKeyFingerprint == "" {
+	if d.SSHKey != "" {
+		d.SSHKeyPath = d.ResolveStorePath(path.Base(d.SSHKey))
+	} else if d.SSHKeyPath == "" && d.SSHKeyFingerprint == "" {
 		d.SSHKeyPath = d.ResolveStorePath("id_rsa")
 	}
 	return d.SSHKeyPath
@@ -358,3 +390,15 @@ func (d *Driver) GetSSHKeyPath() string {
 func (d *Driver) publicSSHKeyPath() string {
 	return d.GetSSHKeyPath() + ".pub"
 }
+
+func copySSHKey(src, dst string) error {
+	if err := mcnutils.CopyFile(src, dst); err != nil {
+		return fmt.Errorf("unable to copy ssh key: %s", err)
+	}
+
+	if err := os.Chmod(dst, 0600); err != nil {
+		return fmt.Errorf("unable to set permissions on the ssh key: %s", err)
+	}
+
+	return nil
+}