Merge pull request #13 from ernoaapa/add-support-for-usernames

Added --user flag to support fetching users public keys

Author: ernoaapa

README.md

@@ -14,7 +14,11 @@ fetch-ssh-keys <source name> <parameters> <output file>
 
 For example fetch users public SSH keys from GitHub `my-lovely-team` team in `my-awesome-company` organization and output in SSH authorized_keys format
 ```shell
+# Fetch 'my-lovely-team' keys in 'my-awesome-company' organisation
 fetch-ssh-keys github --organization my-awesome-company --team my-lovely-team --token YOUR-TOKEN-HERE ./the-keys
+
+# Fetch 'ernoaapa' and 'arnested' public keys
+fetch-ssh-keys github --user ernoaapa --user arnested  --token YOUR-TOKEN-HERE ./the-keys
 ```
 
 Tool can be used for example to automatically update `.ssh/authorized_keys` file by giving path to `.ssh/authorized_keys` as last argument and adding the script to cron job.
@@ -30,12 +34,15 @@ Tool can be used for example to automatically update `.ssh/authorized_keys` file
 | --file-mode    | No (default 0600) | File permissions when writing to a file                                                                   |
 
 #### GitHub
-| Parameter      | Required | Description                                                                                               |
-|----------------|----------|-----------------------------------------------------------------------------------------------------------|
-| --organization | Yes      | Name of the organization which members keys to pick                                                       |
-| --team         | No       | Name of the team which members keys to pick                                                               |
-| --token        | No       | GitHub API token to use for communication. Without token you get only public members of the organization. |
-| --public-only  | No       | Return only members what are publicly members of the given organization                                   |
+| Parameter      | Description                                                                                               |
+|----------------|-----------------------------------------------------------------------------------------------------------|
+| --organization | Name of the organization which members keys to pick                                                       |
+| --team         | Name of the team which members keys to pick                                                               |
+| --user         | Name of the user which keys to pick                                                                       |
+| --token        | GitHub API token to use for communication. Without token you get only public members of the organization. |
+| --public-only  | Return only members what are publicly members of the given organization                                   |
+
+You can give `--organisation` (optionally combined with `--team` flag) and/or one or more `--user` flags.
 
 ## Development
 ### Get dependencies

fetch/github.go

@@ -17,39 +17,33 @@ type GithubFetchParams struct {
 	PublicMembersOnly bool
 }
 
-// GitHubKeys fetches organization users public SSH key from GitHub
-func GitHubKeys(organizationName string, params GithubFetchParams) (map[string][]string, error) {
-	ctx := context.Background()
-
-	client := getClient(params)
+// GitHubOrganisationKeys fetches organization users public SSH key from GitHub
+func GitHubOrganisationKeys(organizationName string, params GithubFetchParams) (map[string][]string, error) {
+	client := getClient(params.Token)
 	users, err := fetchUsers(client, organizationName, params)
 	if err != nil {
 		return map[string][]string{}, err
 	}
 	log.Debugf("Users found: %d", len(users))
 
-	result := map[string][]string{}
+	usernames := []string{}
 	for _, user := range users {
-		username := *user.Login
-		keys, _, err := client.Users.ListKeys(ctx, username, &github.ListOptions{})
-		if err != nil {
-			return map[string][]string{}, err
-		}
-
-		result[username] = make([]string, len(keys))
-
-		for index, key := range keys {
-			result[username][index] = *key.Key
-		}
+		usernames = append(usernames, *user.Login)
 	}
 
-	return result, nil
+	return fetchUserKeys(client, usernames, params.Token)
+}
+
+// GitHubUsers fetches users public SSH keys from GitHub
+func GitHubUsers(usernames []string, token string) (map[string][]string, error) {
+	client := getClient(token)
+	return fetchUserKeys(client, usernames, token)
 }
 
-func getClient(params GithubFetchParams) *github.Client {
-	if len(params.Token) > 0 {
+func getClient(token string) *github.Client {
+	if len(token) > 0 {
 		ts := oauth2.StaticTokenSource(
-			&oauth2.Token{AccessToken: params.Token},
+			&oauth2.Token{AccessToken: token},
 		)
 		return github.NewClient(oauth2.NewClient(oauth2.NoContext, ts))
 	}
@@ -100,3 +94,23 @@ func resolveTeamID(client *github.Client, organizationName, teamName string) (in
 
 	return -1, fmt.Errorf("Unable to find team [%s] from organization [%s]", teamName, organizationName)
 }
+
+func fetchUserKeys(client *github.Client, usernames []string, token string) (map[string][]string, error) {
+	ctx := context.Background()
+
+	result := map[string][]string{}
+	for _, username := range usernames {
+		keys, _, err := client.Users.ListKeys(ctx, username, &github.ListOptions{})
+		if err != nil {
+			return map[string][]string{}, err
+		}
+
+		result[username] = make([]string, len(keys))
+
+		for index, key := range keys {
+			result[username][index] = *key.Key
+		}
+	}
+
+	return result, nil
+}

fetch/github_test.go

@@ -9,10 +9,10 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-func TestFetchPublicKeys(t *testing.T) {
+func TestFetchOrganisationKeys(t *testing.T) {
 	log.SetLevel(log.DebugLevel)
 
-	keys, err := GitHubKeys("devopsfinland", GithubFetchParams{
+	keys, err := GitHubOrganisationKeys("devopsfinland", GithubFetchParams{
 		// Use token if it's available to avoid hitting API rate limits with the tests...
 		Token:             os.Getenv("GITHUB_TOKEN"),
 		PublicMembersOnly: true,
@@ -23,3 +23,14 @@ func TestFetchPublicKeys(t *testing.T) {
 	assert.True(t, len(keys["ernoaapa"]) > 0, "should return ernoaapa public SSH key")
 	assert.True(t, len(keys["ernoaapa"][0]) > 0, "should not return empty key for ernoaapa")
 }
+
+func TestFetchUserKeys(t *testing.T) {
+	log.SetLevel(log.DebugLevel)
+
+	keys, err := GitHubUsers([]string{"ernoaapa", "arnested"}, os.Getenv("GITHUB_TOKEN"))
+
+	assert.NoError(t, err, "Fetch GitHub keys returned error")
+	assert.Equal(t, 2, len(keys), "should return SSH keys for both users")
+	assert.True(t, len(keys["ernoaapa"]) > 0, "should return ernoaapa public SSH key")
+	assert.True(t, len(keys["ernoaapa"][0]) > 0, "should not return empty key for ernoaapa")
+}

main.go

@@ -1,12 +1,15 @@
 package main
 
 import (
+	"fmt"
 	"os"
 
 	log "github.com/Sirupsen/logrus"
+	"github.com/pkg/errors"
 
 	"github.com/ernoaapa/fetch-ssh-keys/fetch"
 	"github.com/ernoaapa/fetch-ssh-keys/output"
+	"github.com/ernoaapa/fetch-ssh-keys/utils"
 	"github.com/urfave/cli"
 )
 
@@ -55,22 +58,52 @@ func main() {
 					Name:  "team",
 					Usage: "Return only members of `TEAM` (this option can be used multiple times for multiple teams)",
 				},
+				cli.StringSliceFlag{
+					Name:  "user",
+					Usage: "Return given user public ssh keys (this option can be used multiple times for multiple users)",
+				},
 			},
 			Action: func(c *cli.Context) error {
-				if c.String("organization") == "" {
-					log.Fatalln("You must give --organization value")
+				var (
+					token        = c.String("token")
+					organisation = c.String("organization")
+					teams        = c.StringSlice("team")
+					users        = c.StringSlice("user")
+					publicOnly   = c.Bool("public-only")
+
+					orgKeys  map[string][]string
+					userKeys map[string][]string
+
+					target   = c.Args().Get(0)
+					fileMode = os.FileMode(c.GlobalInt("file-mode"))
+					format   = c.GlobalString("format")
+
+					err error
+				)
+
+				if organisation == "" && len(users) == 0 {
+					return fmt.Errorf("You must give either --organisation or --user parameter")
+				}
+
+				if c.IsSet("organization") {
+					orgKeys, err = fetch.GitHubOrganisationKeys(organisation, fetch.GithubFetchParams{
+						Token:             token,
+						TeamNames:         teams,
+						PublicMembersOnly: publicOnly,
+					})
+					if err != nil {
+						return errors.Wrapf(err, "Failed to fetch keys from organisation %s", organisation)
+					}
 				}
 
-				keys, err := fetch.GitHubKeys(c.String("organization"), fetch.GithubFetchParams{
-					Token:             c.String("token"),
-					TeamNames:         c.StringSlice("team"),
-					PublicMembersOnly: c.Bool("public-only"),
-				})
-				if err != nil {
-					log.Fatalln("Failed to fetch keys", err)
+				if c.IsSet("user") {
+					userKeys, err = fetch.GitHubUsers(users, token)
+					if err != nil {
+						return errors.Wrap(err, "Failed to fetch GitHub user(s) keys")
+					}
 				}
 
-				return output.Write(c.GlobalString("format"), c.Args().Get(0), os.FileMode(c.GlobalInt("file-mode")), keys)
+				return output.Write(format, target, fileMode, utils.MergeKeys(orgKeys, userKeys))
 			},
 		},
 	}

utils/merge.go

@@ -0,0 +1,18 @@
+package utils
+
+// MergeKeys merges key maps together to single map
+func MergeKeys(keySets ...map[string][]string) map[string][]string {
+	result := make(map[string][]string)
+
+	for _, userKeys := range keySets {
+		for username, keys := range userKeys {
+			if _, ok := result[username]; !ok {
+				result[username] = []string{}
+			}
+
+			result[username] = append(result[username], keys...)
+		}
+	}
+
+	return result
+}

utils/merge_test.go

@@ -0,0 +1,23 @@
+package utils
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMergeKeys(t *testing.T) {
+	result := MergeKeys(map[string][]string{
+		"user-1": []string{"key1"},
+		"user-2": []string{"key1"},
+	}, map[string][]string{
+		"user-2": []string{"key2"},
+		"user-3": []string{"key1"},
+	})
+
+	assert.Equal(t, map[string][]string{
+		"user-1": []string{"key1"},
+		"user-2": []string{"key1", "key2"},
+		"user-3": []string{"key1"},
+	}, result)
+}