Save challenge username to config file

README.md

@@ -19,11 +19,10 @@ Weep searches for a configuration in the following locations:
 
 - embedded configuration (see below)
 - `/etc/weep/weep.yaml`
-- `~/.config/weep/weep.yaml`
-- `~/.weep.yaml`
+- `~/.weep/weep.yaml`
 - `./weep.yaml`
 
-Multiple configurations in these locations **will be merged**.
+Multiple configurations in these locations **will be merged** in the order listed above (e.g. entries in `./weep.yaml` will take precedence over `~/.weep/weep.yaml`.
 
 You can also specify a config file as a CLI arg. This configuration will be used exclusively and will not be merged with other configurations:
 

challenge/challenge.go

@@ -26,12 +26,17 @@ import (
 	"net/url"
 	"os"
 	"os/exec"
-	"os/user"
-	"path/filepath"
+	"os/signal"
+	"path"
 	"runtime"
 	"strings"
+	"syscall"
 	"time"
 
+	"github.com/mitchellh/go-homedir"
+
+	"github.com/netflix/weep/config"
+
 	"github.com/manifoldco/promptui"
 
 	"github.com/spf13/viper"
@@ -102,6 +107,8 @@ func isWSL() bool {
 func poll(pollingUrl string) (*ConsolemeChallengeResponse, error) {
 	timeout := time.After(2 * time.Minute)
 	tick := time.Tick(3 * time.Second)
+	interrupt := make(chan os.Signal, 1)
+	signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
 	req, err := http.NewRequest("GET", pollingUrl, nil)
 	if err != nil {
 		return nil, err
@@ -121,6 +128,8 @@ func poll(pollingUrl string) (*ConsolemeChallengeResponse, error) {
 			if pollResponse.Status == "success" {
 				return pollResponse, nil
 			}
+		case <-interrupt:
+			return nil, errors.New("interrupt received")
 		}
 	}
 }
@@ -139,18 +148,18 @@ func pollRequest(c *http.Client, r *http.Request) (*ConsolemeChallengeResponse,
 }
 
 func getCredentialsPath() (string, error) {
-	currentUser, err := user.Current()
+	home, err := homedir.Dir()
 	if err != nil {
 		return "", err
 	}
-	weepDir := filepath.Join(currentUser.HomeDir, ".weep")
+	weepDir := path.Join(home, ".weep")
 	// Setup the directories where we will be writing credentials
 	if _, err := os.Stat(weepDir); os.IsNotExist(err) {
 		_ = os.Mkdir(weepDir, 0700)
 	} else {
 		_ = os.Chmod(weepDir, 0700)
 	}
-	credentialsPath := filepath.Join(weepDir, ".credentials")
+	credentialsPath := path.Join(weepDir, "credentials")
 	return credentialsPath, nil
 }
 
@@ -177,11 +186,14 @@ func promptUser() (string, error) {
 	}
 
 	result, err := prompt.Run()
-
 	if err != nil {
 		return "", err
 	}
 
+	if err := config.SetUser(result); err != nil {
+		return "", err
+	}
+
 	return result, nil
 }
 
@@ -199,7 +211,7 @@ func HasValidJwt(challenge *ConsolemeChallengeResponse) bool {
 
 func RefreshChallenge() error {
 	existingChallengeBody, err := getChallenge()
-	var user = viper.GetString("challenge_settings.user")
+	var userName = viper.GetString("challenge_settings.user")
 	if err != nil {
 		log.Debugf("unable to read existing challenge file: %v", err)
 
@@ -210,23 +222,23 @@ func RefreshChallenge() error {
 	}
 	// Step 1: Make unauthed request to ConsoleMe challenge endpoint and get a challenge challenge
 	// Check Config for username
-	if user == "" && existingChallengeBody != nil {
+	if userName == "" && existingChallengeBody != nil {
 		// Find user from old jwt
-		user = existingChallengeBody.User
+		userName = existingChallengeBody.User
 	}
-	if user == "" {
-		user, err = promptUser()
+	if userName == "" {
+		userName, err = promptUser()
 		if err != nil {
 			return err
 		}
 	}
-	if user == "" {
+	if userName == "" {
 		return fmt.Errorf("invalid configuration: challenge_settings.user must be set")
 	}
 	var consoleMeChallengeGeneratorEndpoint = fmt.Sprintf(
 		"%s/noauth/v1/challenge_generator/%s",
 		viper.GetString("consoleme_url"),
-		user,
+		userName,
 	)
 	var challenge ConsolemeChallenge
 	req, err := http.NewRequest("GET", consoleMeChallengeGeneratorEndpoint, nil)

cmd/root.go

@@ -23,10 +23,10 @@ import (
 	"path/filepath"
 	"syscall"
 
+	"github.com/netflix/weep/config"
+
 	"github.com/kardianos/service"
 
-	"github.com/mitchellh/go-homedir"
-	"github.com/netflix/weep/config"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -62,72 +62,9 @@ func Execute() {
 	}
 }
 
-// initConfig reads in configs by precedence, with later configs overriding earlier:
-//   - embedded
-//   - /etc/weep/weep.yaml
-//   - ~/.config/weep/weep.yaml
-//   - ~/.weep.yaml
-//   - ./weep.yaml
-// If a config file is specified via CLI arg, it will be read exclusively and not merged with other
-// configuration.
 func initConfig() {
-	home, err := homedir.Dir()
-	if err != nil {
-		log.Fatal(err)
-	}
-	viper.SetConfigType("yaml")
-
-	// Read in explicitly defined config file
-	if cfgFile != "" {
-		viper.SetConfigFile(cfgFile)
-		err = viper.ReadInConfig()
-		if err != nil {
-			log.Fatalf("could not open config file %s: %v", cfgFile, err)
-		}
-		return
-	}
-
-	// Read embedded config if available
-	if err := config.ReadEmbeddedConfig(); err != nil {
-		log.Debugf("unable to read embedded config: %v; falling back to config file", err)
-	}
-
-	// Read in config from etc
-	viper.SetConfigName("weep")
-	viper.AddConfigPath("/etc/weep/")
-	_ = viper.MergeInConfig()
-
-	// Read in config from config dir
-	viper.SetConfigName("weep")
-	viper.AddConfigPath(home + "/.config/weep/")
-	_ = viper.MergeInConfig()
-
-	// Read in config from home dir
-	viper.SetConfigName(".weep")
-	viper.AddConfigPath(home)
-	_ = viper.MergeInConfig()
-
-	// Read in config from current directory
-	viper.SetConfigName("weep")
-	viper.AddConfigPath(".")
-	_ = viper.MergeInConfig()
-
-	// TODO: revisit first-run setup
-	//if err := viper.MergeInConfig(); err != nil {
-	//	if _, ok := err.(viper.ConfigFileNotFoundError); ok && config.EmbeddedConfigFile != "" {
-	//		log.Debugf("no config file found, trying to use embedded config")
-	//	} else if isatty.IsTerminal(os.Stdout.Fd()) {
-	//		err = util.FirstRunPrompt()
-	//		if err != nil {
-	//			log.Fatalf("config bootstrap failed: %v", err)
-	//		}
-	//	} else {
-	//		log.Debugf("unable to read config file: %v", err)
-	//	}
-	//}
-
-	if err := viper.Unmarshal(&config.Config); err != nil {
-		log.Fatalf("unable to decode config into struct: %v", err)
+	if err := config.InitConfig(cfgFile); err != nil {
+		log.Fatalf("failed to initialize config: %v", err)
 	}
 }
 

config/config.go

@@ -17,9 +17,14 @@
 package config
 
 import (
+	"path"
 	"path/filepath"
 	"runtime"
 
+	"github.com/mitchellh/go-homedir"
+	"github.com/pkg/errors"
+	log "github.com/sirupsen/logrus"
+
 	"github.com/spf13/viper"
 )
 
@@ -49,6 +54,98 @@ func getDefaultLogFile() string {
 	}
 }
 
+// initConfig reads in configs by precedence, with later configs overriding earlier:
+//   - embedded
+//   - /etc/weep/weep.yaml
+//   - ~/.weep/weep.yaml
+//   - ./weep.yaml
+// If a config file is specified via CLI arg, it will be read exclusively and not merged with other
+// configuration.
+func InitConfig(filename string) error {
+	viper.SetConfigType("yaml")
+
+	// Read in explicitly defined config file
+	if filename != "" {
+		viper.SetConfigFile(filename)
+		if err := viper.ReadInConfig(); err != nil {
+			log.Errorf("could not open config file %s: %v", filename, err)
+			return err
+		}
+		return nil
+	}
+
+	// Read embedded config if available
+	if err := ReadEmbeddedConfig(); err != nil {
+		log.Debugf("unable to read embedded config: %v", err)
+	}
+
+	configLocations := []string{
+		"/etc/weep",
+		"$HOME/.weep",
+		".",
+	}
+
+	for _, dir := range configLocations {
+		viper.SetConfigName("weep")
+		viper.AddConfigPath(dir)
+		if err := viper.MergeInConfig(); err != nil {
+			return errors.Wrapf(err, "failed to read config from %s", viper.ConfigFileUsed())
+		}
+	}
+
+	// TODO: revisit first-run setup
+	//if err := viper.MergeInConfig(); err != nil {
+	//	if _, ok := err.(viper.ConfigFileNotFoundError); ok && config.EmbeddedConfigFile != "" {
+	//		log.Debugf("no config file found, trying to use embedded config")
+	//	} else if isatty.IsTerminal(os.Stdout.Fd()) {
+	//		err = util.FirstRunPrompt()
+	//		if err != nil {
+	//			log.Fatalf("config bootstrap failed: %v", err)
+	//		}
+	//	} else {
+	//		log.Debugf("unable to read config file: %v", err)
+	//	}
+	//}
+
+	if err := viper.Unmarshal(&Config); err != nil {
+		return errors.Wrap(err, "unable to decode config into struct")
+	}
+	return nil
+}
+
+// SetUser saves the provided username to ~/.weep/weep.yaml
+func SetUser(user string) error {
+	// Create a temporary viper instance to isolate from main config
+	v := viper.New()
+	v.SetConfigType("yaml")
+	v.SetConfigName("weep")
+	v.AddConfigPath("$HOME/.weep")
+
+	// Read in existing config if there is one so we don't overwrite it
+	if err := v.ReadInConfig(); err != nil {
+		switch err.(type) {
+		case viper.ConfigFileNotFoundError:
+			break
+		default:
+			return err
+		}
+	}
+
+	// Set user in the temp config then write to file
+	v.Set("challenge_settings.user", user)
+
+	home, err := homedir.Dir()
+	if err != nil {
+		return err
+	}
+	configPath := path.Join(home, ".weep/weep.yaml")
+
+	if err := v.WriteConfigAs(configPath); err != nil {
+		return err
+	}
+	return nil
+}
+
 var (
 	Config WeepConfig
 )