From 95a00070c310e140802f058b63f0ac85c47c693d Mon Sep 17 00:00:00 2001
From: Brian Waldon <bcwaldon@gmail.com>
Date: Wed, 5 Mar 2014 14:30:38 -0800
Subject: [PATCH] feat(ssh): Set custom name for ssh identity

---
 README.md                 | 3 +++
 cloudinit/cloud_config.go | 6 ++++--
 cloudinit/ssh_key.go      | 4 ++--
 coreos-cloudinit.go       | 5 ++++-
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 131e108..41fffa6 100644
--- a/README.md
+++ b/README.md
@@ -15,6 +15,9 @@ Only a subset of [cloud-config functionality][cloud-config] is implemented. A se
 
 Provided public SSH keys will be authorized for the `core` user.
 
+The keys will be named "coreos-cloudinit" by default.
+Override this with the `--ssh-key-name` flag when calling `coreos-cloudinit`.
+
 ### Custom cloud-config Parameters
 
 #### coreos.etcd.discovery_url
diff --git a/cloudinit/cloud_config.go b/cloudinit/cloud_config.go
index 815bbad..6e35b56 100644
--- a/cloudinit/cloud_config.go
+++ b/cloudinit/cloud_config.go
@@ -6,6 +6,8 @@ import (
 	"launchpad.net/goyaml"
 )
 
+const DefaultSSHKeyName = "coreos-cloudinit"
+
 type CloudConfig struct {
 	SSH_Authorized_Keys []string
 	Coreos struct{Etcd struct{ Discovery_URL string }; Fleet struct{ Autostart bool } }
@@ -26,9 +28,9 @@ func (cc CloudConfig) String() string {
 	}
 }
 
-func ResolveCloudConfig(cfg CloudConfig) error {
+func ResolveCloudConfig(cfg CloudConfig, sshKeyName string) error {
 	if len(cfg.SSH_Authorized_Keys) > 0 {
-		err := AuthorizeSSHKeys(cfg.SSH_Authorized_Keys)
+		err := AuthorizeSSHKeys(sshKeyName, cfg.SSH_Authorized_Keys)
 		if err == nil {
 			log.Printf("Authorized SSH keys for core user")
 		} else {
diff --git a/cloudinit/ssh_key.go b/cloudinit/ssh_key.go
index 598bd01..7a3c502 100644
--- a/cloudinit/ssh_key.go
+++ b/cloudinit/ssh_key.go
@@ -10,7 +10,7 @@ import (
 
 // Add the provide SSH public key to the core user's list of
 // authorized keys
-func AuthorizeSSHKeys(keys []string) error {
+func AuthorizeSSHKeys(name string, keys []string) error {
 	for i, key := range keys {
 		keys[i] = strings.TrimSpace(key)
 	}
@@ -19,7 +19,7 @@ func AuthorizeSSHKeys(keys []string) error {
 	// also ends with a newline
 	joined := fmt.Sprintf("%s\n", strings.Join(keys, "\n"))
 
-	cmd := exec.Command("update-ssh-keys", "-u", "core", "-a", "coreos-cloudinit")
+	cmd := exec.Command("update-ssh-keys", "-u", "core", "-a", name)
 	stdin, err := cmd.StdinPipe()
 	if err != nil {
 		return err
diff --git a/coreos-cloudinit.go b/coreos-cloudinit.go
index a986aef..cdf3af1 100644
--- a/coreos-cloudinit.go
+++ b/coreos-cloudinit.go
@@ -28,6 +28,9 @@ func main() {
 	var workspace string
 	flag.StringVar(&workspace, "workspace", "/var/lib/coreos-cloudinit", "Base directory coreos-cloudinit should use to store data")
 
+	var sshKeyName string
+	flag.StringVar(&sshKeyName, "ssh-key-name", cloudinit.DefaultSSHKeyName, "Add SSH keys to the system with the given name")
+
 	flag.Parse()
 
 	if printVersion == true {
@@ -70,7 +73,7 @@ func main() {
 
 	switch t := parsed.(type) {
 	case cloudinit.CloudConfig:
-		err = cloudinit.ResolveCloudConfig(t)
+		err = cloudinit.ResolveCloudConfig(t, sshKeyName)
 	case cloudinit.Script:
 		var path string
 		path, err = cloudinit.PersistScriptInWorkspace(t, workspace)