Add better user home dir handling

main.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"os/user"
+	"path/filepath"
 	"strings"
 )
 
@@ -57,12 +59,19 @@ func getUsernamesKeys(config config, client *githubClient, singleUsername string
 	return usernamesKeys
 }
 
+func getAuthorizedKeysFilename(userHomeDir string) string {
+	authorizedKeysFilename := filepath.Join(userHomeDir, ".ssh", "authorized_keys")
+	return authorizedKeysFilename
+}
+
 func getKeysOutput(keys []string) string {
-	return strings.Join(keys, "\n")
+	keysOutput := strings.Join(keys, "\n")
+	return keysOutput
 }
 
-func writeKeysToFile(keys []string, filename string) error {
-	authorizedKeysFile, err := os.Create(filename)
+func writeKeysToUserAuthorizedKeysFile(keys []string, userHomeDir string) error {
+	authorizedKeysFilename := getAuthorizedKeysFilename(userHomeDir)
+	authorizedKeysFile, err := os.Create(authorizedKeysFilename)
 	if err != nil {
 		return err
 	}
@@ -84,17 +93,20 @@ func main() {
 
 	if *debug {
 		if rate, _, err := client.RateLimit(); err == nil {
-			fmt.Printf("GitHub API Rate Limit: %#v\n", rate)
+			log.Printf("GitHub API Rate Limit: %#v\n", rate)
 		}
 	}
 
 	usernamesKeys := getUsernamesKeys(config, client, flag.Arg(0))
 
 	for username, keys := range usernamesKeys {
 		if *writeToFile {
-			// TODO: Is this always the path?
-			err := writeKeysToFile(keys, fmt.Sprintf("/home/%s/.ssh/authorized_keys", username))
-			check(err)
+			if user, err := user.Lookup(username); err == nil {
+				err = writeKeysToUserAuthorizedKeysFile(keys, user.HomeDir)
+				check(err)
+			} else {
+				log.Printf("User '%s' not found, no keys written.\n", username)
+			}
 		} else {
 			fmt.Println(getKeysOutput(keys))
 		}

main_test.go

@@ -7,6 +7,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -188,16 +189,27 @@ github_user_1_key_3`
 	assert.Equal(t, expected, actual)
 }
 
-func TestWriteKeysToFile(t *testing.T) {
+func TestWriteKeysToUserAuthorizedKeysFile(t *testing.T) {
 	// Test writing test keys to a test file
 	keys := []string{"github_user_1_key_1", "github_user_1_key_2", "github_user_1_key_3"}
-	authorizedKeysFile, _ := ioutil.TempFile("", "ghkeys-authorized_keys")
-	authorizedKeysFilename := authorizedKeysFile.Name()
-	defer os.RemoveAll(authorizedKeysFilename)
-	err := writeKeysToFile(keys, authorizedKeysFilename)
+	userHomeDir, _ := ioutil.TempDir("", "ghkeys-userHomeDir")
+	defer os.RemoveAll(userHomeDir)
+
+	// No .ssh directory exists, we expect an error
+	err := writeKeysToUserAuthorizedKeysFile(keys, userHomeDir)
+	assert.Error(t, err)
+
+	// Create a temp authorized_keys file with junk data
+	authorizedKeysFilename := getAuthorizedKeysFilename(userHomeDir)
+	os.MkdirAll(filepath.Dir(authorizedKeysFilename), os.ModePerm)
+	authorizedKeysFile, _ := os.Create(authorizedKeysFilename)
+	authorizedKeysFile.WriteString("junk")
+
+	// Test writing to an authorized_keys file
+	err = writeKeysToUserAuthorizedKeysFile(keys, userHomeDir)
 	assert.Nil(t, err)
 
-	// Compare test file to what we expect it to have
+	// Compare test authorized_keys file to what we expect it to have
 	expectedKeysFileContent, _ := ioutil.ReadFile(testAuthorizedKeysFilename)
 	actualKeysFileContent, _ := ioutil.ReadFile(authorizedKeysFilename)
 	assert.Equal(t, expectedKeysFileContent, actualKeysFileContent)