From bba93792d59d24d0799a579591a6aee9ff50646e Mon Sep 17 00:00:00 2001
From: Josh Curl <josh@curl.me>
Date: Wed, 16 Nov 2016 13:02:35 -0800
Subject: [PATCH] Rewrite update-ssh-keys in Go

---
 cmd/cloudinitexecute/authorize_ssh_keys.go | 108 +++++++++++++++++++--
 cmd/cloudinitexecute/cloudinitexecute.go   |   8 +-
 images/02-console/Dockerfile               |   1 -
 images/02-console/update-ssh-keys          |  20 ----
 4 files changed, 105 insertions(+), 32 deletions(-)
 delete mode 100755 images/02-console/update-ssh-keys

diff --git a/cmd/cloudinitexecute/authorize_ssh_keys.go b/cmd/cloudinitexecute/authorize_ssh_keys.go
index 79016135..98d14b24 100644
--- a/cmd/cloudinitexecute/authorize_ssh_keys.go
+++ b/cmd/cloudinitexecute/authorize_ssh_keys.go
@@ -1,20 +1,110 @@
 package cloudinitexecute
 
 import (
+	"io/ioutil"
 	"os"
-	"os/exec"
+	"path"
+	"strconv"
+	"strings"
 
 	log "github.com/Sirupsen/logrus"
+	"github.com/rancher/os/util"
 )
 
-func authorizeSSHKeys(user string, authorizedKeys []string, name string) {
-	for _, authorizedKey := range authorizedKeys {
-		cmd := exec.Command("update-ssh-keys", user, authorizedKey)
-		cmd.Stdout = os.Stdout
-		cmd.Stderr = os.Stderr
-		err := cmd.Run()
-		if err != nil {
-			log.WithFields(log.Fields{"err": err, "user": user, "auth_key": authorizedKey}).Error("Error updating SSH authorized_keys")
+var (
+	sshDir             = ".ssh"
+	authorizedKeysFile = path.Join(sshDir, "authorized_keys")
+)
+
+func authorizeSSHKeys(username string, authorizedKeys []string, name string) error {
+	var uid int
+	var gid int
+	var homeDir string
+
+	bytes, err := ioutil.ReadFile("/etc/passwd")
+	if err != nil {
+		return err
+	}
+
+	for _, line := range strings.Split(string(bytes), "\n") {
+		if strings.HasPrefix(line, username) {
+			split := strings.Split(line, ":")
+			if len(split) < 6 {
+				break
+			}
+			uid, err = strconv.Atoi(split[2])
+			if err != nil {
+				return err
+			}
+			gid, err = strconv.Atoi(split[3])
+			if err != nil {
+				return err
+			}
+			homeDir = split[5]
 		}
 	}
+
+	sshDir = path.Join(homeDir, sshDir)
+	authorizedKeysFile = path.Join(homeDir, authorizedKeysFile)
+
+	for _, authorizedKey := range authorizedKeys {
+		if err = authorizeSSHKey(authorizedKey, uid, gid, homeDir); err != nil {
+			log.Errorf("Failed to authorize SSH key %s: %v", authorizedKey, err)
+		}
+	}
+
+	return nil
+}
+
+func authorizeSSHKey(authorizedKey string, uid, gid int, homeDir string) error {
+	if _, err := os.Stat(sshDir); os.IsNotExist(err) {
+		if err = os.Mkdir(sshDir, 0700); err != nil {
+			return err
+		}
+	} else if err != nil {
+		return err
+	}
+
+	authorizedKeysFileInfo, err := os.Stat(authorizedKeysFile)
+	if os.IsNotExist(err) {
+		keysFile, err := os.Create(authorizedKeysFile)
+		if err != nil {
+			return err
+		}
+		if err = keysFile.Chmod(0600); err != nil {
+			return err
+		}
+		if err = keysFile.Close(); err != nil {
+			return err
+		}
+		authorizedKeysFileInfo, err = os.Stat(authorizedKeysFile)
+		if err != nil {
+			return err
+		}
+	} else if err != nil {
+		return err
+	}
+
+	bytes, err := ioutil.ReadFile(authorizedKeysFile)
+	if err != nil {
+		return err
+	}
+
+	if !strings.Contains(string(bytes), authorizedKey) {
+		bytes = append(bytes, []byte(authorizedKey)...)
+		bytes = append(bytes, '\n')
+	}
+
+	perm := authorizedKeysFileInfo.Mode().Perm()
+	if err = util.WriteFileAtomic(authorizedKeysFile, bytes, perm); err != nil {
+		return err
+	}
+	if err = os.Chown(sshDir, uid, gid); err != nil {
+		return err
+	}
+	if err = os.Chown(authorizedKeysFile, uid, gid); err != nil {
+		return err
+	}
+
+	return nil
 }
diff --git a/cmd/cloudinitexecute/cloudinitexecute.go b/cmd/cloudinitexecute/cloudinitexecute.go
index 5a557028..ea50d1b4 100644
--- a/cmd/cloudinitexecute/cloudinitexecute.go
+++ b/cmd/cloudinitexecute/cloudinitexecute.go
@@ -56,8 +56,12 @@ func Main() {
 
 func ApplyConsole(cfg *rancherConfig.CloudConfig) {
 	if len(cfg.SSHAuthorizedKeys) > 0 {
-		authorizeSSHKeys("rancher", cfg.SSHAuthorizedKeys, sshKeyName)
-		authorizeSSHKeys("docker", cfg.SSHAuthorizedKeys, sshKeyName)
+		if err := authorizeSSHKeys("rancher", cfg.SSHAuthorizedKeys, sshKeyName); err != nil {
+			log.Error(err)
+		}
+		if err := authorizeSSHKeys("docker", cfg.SSHAuthorizedKeys, sshKeyName); err != nil {
+			log.Error(err)
+		}
 	}
 
 	WriteFiles(cfg, "console")
diff --git a/images/02-console/Dockerfile b/images/02-console/Dockerfile
index bd114aac..a933b8cc 100644
--- a/images/02-console/Dockerfile
+++ b/images/02-console/Dockerfile
@@ -1,5 +1,4 @@
 FROM rancher/os-base
-COPY update-ssh-keys /usr/sbin/
 COPY build/lsb-release /etc/
 RUN sed -i 's/rancher:!/rancher:*/g' /etc/shadow && \
     sed -i 's/docker:!/docker:*/g' /etc/shadow && \
diff --git a/images/02-console/update-ssh-keys b/images/02-console/update-ssh-keys
deleted file mode 100755
index f2a0c637..00000000
--- a/images/02-console/update-ssh-keys
+++ /dev/null
@@ -1,20 +0,0 @@
-#!/bin/bash
-
-USERNAME=$1
-HOME_DIR=$(grep ^$USERNAME /etc/passwd | cut -f6 -d:)
-
-if [ ! -d $HOME_DIR/.ssh ]; then
-	mkdir -p $HOME_DIR/.ssh
-    chmod 0700 $HOME_DIR/.ssh
-fi
-
-if [ ! -e $HOME_DIR/.ssh/authorized_keys ]; then
-    touch $HOME_DIR/.ssh/authorized_keys
-    chmod 0600 $HOME_DIR/.ssh/authorized_keys
-fi
-
-if ! grep -q "$2" $HOME_DIR/.ssh/authorized_keys; then
-    echo "$2" >> $HOME_DIR/.ssh/authorized_keys
-fi
-
-chown -R $USERNAME $HOME_DIR/.ssh