<commit_before>package cli

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	osexec "os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/v5/server"
	"github.com/99designs/aws-vault/v5/vault"
	"github.com/99designs/keyring"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"gopkg.in/alecthomas/kingpin.v2"
)

type ExecCommandInput struct {
	ProfileName      string
	Command          string
	Args             []string
	StartEc2Server   bool
	StartEcsServer   bool
	CredentialHelper bool
	Config           vault.Config
	SessionDuration  time.Duration
	NoSession        bool
}

// AwsCredentialHelperData is metadata for AWS CLI credential process
// See https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes
type AwsCredentialHelperData struct {
	Version         int    `json:"Version"`
	AccessKeyID     string `json:"AccessKeyId"`
	SecretAccessKey string `json:"SecretAccessKey"`
	SessionToken    string `json:"SessionToken,omitempty"`
	Expiration      string `json:"Expiration,omitempty"`
}

func ConfigureExecCommand(app *kingpin.Application, a *AwsVault) {
	input := ExecCommandInput{}

	cmd := app.Command("exec", "Executes a command with AWS credentials in the environment")

	cmd.Flag("duration", "Duration of the temporary or assume-role session. Defaults to 1h").
		Short('d').
		DurationVar(&input.SessionDuration)

	cmd.Flag("no-session", "Skip creating STS session with GetSessionToken").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Flag("region", "The AWS region").
		StringVar(&input.Config.Region)

	cmd.Flag("mfa-token", "The MFA token to use").
		Short('t').
		StringVar(&input.Config.MfaToken)

	cmd.Flag("json", "AWS credential helper. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes").
		Short('j').
		BoolVar(&input.CredentialHelper)

	cmd.Flag("server", "Run a server in the background for credentials").
		Short('s').
		BoolVar(&input.StartEc2Server)

	cmd.Flag("ec2-server", "Run a EC2 metadata server in the background for credentials").
		Hidden().
		BoolVar(&input.StartEc2Server)

	cmd.Flag("ecs-server", "Run a ECS credential server in the background for credentials").
		Hidden().
		BoolVar(&input.StartEcsServer)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(a.MustGetProfileNames).
		StringVar(&input.ProfileName)

	cmd.Arg("cmd", "Command to execute, defaults to $SHELL").
		StringVar(&input.Command)

	cmd.Arg("args", "Command arguments").
		StringsVar(&input.Args)

	cmd.Action(func(c *kingpin.ParseContext) (err error) {
		input.Config.MfaPromptMethod = a.PromptDriver
		input.Config.NonChainedGetSessionTokenDuration = input.SessionDuration
		input.Config.AssumeRoleDuration = input.SessionDuration
		if input.Command == "" {
			input.Command, input.Args = getDefaultShellCmd()
		}
		if input.Command == "" {
			app.Fatalf("Argument 'cmd' not provided, and SHELL not present, try --help")
		}

		cl, err := a.ConfigLoader()
		if err != nil {
			return err
		}
		keyring, err := a.Keyring()
		if err != nil {
			return err
		}

		err = ExecCommand(input, cl, keyring)
		app.FatalIfError(err, "exec")
		return nil
	})
}

func getDefaultShellCmd() (string, []string) {
	shellCmd := os.Getenv("SHELL")
	s := strings.ToLower(shellCmd)
	s = strings.TrimSuffix(s, ".exe")
	s = filepath.Base(s)

	// for shells that support it start an interactive login shell
	shellArgs := []string{}
	if s == "sh" ||
		s == "bash" ||
		s == "zsh" ||
		s == "csh" ||
		s == "fish" {
		shellArgs = []string{"-l"}
	}

	return shellCmd, shellArgs
}

func ExecCommand(input ExecCommandInput, configLoader *vault.ConfigLoader, keyring keyring.Keyring) error {
	if os.Getenv("AWS_VAULT") != "" {
		return fmt.Errorf("aws-vault sessions should be nested with care, unset $AWS_VAULT to force")
	}

	if input.StartEc2Server && input.StartEcsServer {
		return fmt.Errorf("Can't use --server with --ecs-server")
	}
	if input.StartEc2Server && input.CredentialHelper {
		return fmt.Errorf("Can't use --server with --json")
	}
	if input.StartEc2Server && input.NoSession {
		return fmt.Errorf("Can't use --server with --no-session")
	}
	if input.StartEcsServer && input.CredentialHelper {
		return fmt.Errorf("Can't use --ecs-server with --json")
	}
	if input.StartEcsServer && input.NoSession {
		return fmt.Errorf("Can't use --ecs-server with --no-session")
	}

	vault.UseSession = !input.NoSession

	configLoader.BaseConfig = input.Config
	configLoader.ActiveProfile = input.ProfileName
	config, err := configLoader.LoadFromProfile(input.ProfileName)
	if err != nil {
		return err
	}

	ckr := &vault.CredentialKeyring{Keyring: keyring}
	creds, err := vault.NewTempCredentials(config, ckr)
	if err != nil {
		return fmt.Errorf("Error getting temporary credentials: %w", err)
	}

	if input.StartEc2Server {
		return execEc2Server(input, config, creds)
	}

	if input.StartEcsServer {
		return execEcsServer(input, config, creds)
	}

	if input.CredentialHelper {
		return execCredentialHelper(input, config, creds)
	}

	return execEnvironment(input, config, creds)
}

func updateEnvForAwsVault(env environ, profileName string, region string) environ {
	env.Unset("AWS_ACCESS_KEY_ID")
	env.Unset("AWS_SECRET_ACCESS_KEY")
	env.Unset("AWS_SESSION_TOKEN")
	env.Unset("AWS_SECURITY_TOKEN")
	env.Unset("AWS_CREDENTIAL_FILE")
	env.Unset("AWS_DEFAULT_PROFILE")
	env.Unset("AWS_PROFILE")
	env.Unset("AWS_SDK_LOAD_CONFIG")

	env.Set("AWS_VAULT", profileName)

	if region != "" {
		log.Printf("Setting subprocess env: AWS_DEFAULT_REGION=%s, AWS_REGION=%s", region, region)
		env.Set("AWS_DEFAULT_REGION", region)
		env.Set("AWS_REGION", region)
	}

	return env
}

func execEc2Server(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	if err := server.StartEc2CredentialsServer(creds, config.Region); err != nil {
		return fmt.Errorf("Failed to start credential server: %w", err)
	}

	env := environ(os.Environ())
	env = updateEnvForAwsVault(env, input.ProfileName, config.Region)

	return execCmd(input.Command, input.Args, env)
}

func execEcsServer(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	uri, token, err := server.StartEcsCredentialServer(creds)
	if err != nil {
		return fmt.Errorf("Failed to start credential server: %w", err)
	}

	env := environ(os.Environ())
	env = updateEnvForAwsVault(env, input.ProfileName, config.Region)

	log.Println("Setting subprocess env AWS_CONTAINER_CREDENTIALS_FULL_URI, AWS_CONTAINER_AUTHORIZATION_TOKEN")
	env.Set("AWS_CONTAINER_CREDENTIALS_FULL_URI", uri)
	env.Set("AWS_CONTAINER_AUTHORIZATION_TOKEN", token)

	return execCmd(input.Command, input.Args, env)
}

func execCredentialHelper(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	val, err := creds.Get()
	if err != nil {
		return fmt.Errorf("Failed to get credentials for %s: %w", input.ProfileName, err)
	}

	credentialData := AwsCredentialHelperData{
		Version:         1,
		AccessKeyID:     val.AccessKeyID,
		SecretAccessKey: val.SecretAccessKey,
	}
	if val.SessionToken != "" {
		credentialData.SessionToken = val.SessionToken
	}
	if credsExpiresAt, err := creds.ExpiresAt(); err == nil {
		credentialData.Expiration = credsExpiresAt.Format("2006-01-02T15:04:05Z")
	}

	json, err := json.Marshal(&credentialData)
	if err != nil {
		return fmt.Errorf("Error creating credential json: %w", err)
	}

	fmt.Print(string(json))

	return nil
}

func execEnvironment(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	val, err := creds.Get()
	if err != nil {
		return fmt.Errorf("Failed to get credentials for %s: %w", input.ProfileName, err)
	}

	env := environ(os.Environ())
	env = updateEnvForAwsVault(env, input.ProfileName, config.Region)

	log.Println("Setting subprocess env: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY")
	env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
	env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

	if val.SessionToken != "" {
		log.Println("Setting subprocess env: AWS_SESSION_TOKEN, AWS_SECURITY_TOKEN")
		env.Set("AWS_SESSION_TOKEN", val.SessionToken)
		env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
	}
	if expiration, err := creds.ExpiresAt(); err == nil {
		log.Println("Setting subprocess env: AWS_SESSION_EXPIRATION")
		env.Set("AWS_SESSION_EXPIRATION", expiration.Format(time.RFC3339))
	}

	if !supportsExecSyscall() {
		return execCmd(input.Command, input.Args, env)
	}

	return execSyscall(input.Command, input.Args, env)
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}

func execCmd(command string, args []string, env []string) error {
	log.Printf("Starting child process: %s %s", command, strings.Join(args, " "))

	cmd := osexec.Command(command, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = env

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan)

	if err := cmd.Start(); err != nil {
		return err
	}

	go func() {
		for {
			sig := <-sigChan
			cmd.Process.Signal(sig)
		}
	}()

	if err := cmd.Wait(); err != nil {
		cmd.Process.Signal(os.Kill)
		return fmt.Errorf("Failed to wait for command termination: %v", err)
	}

	waitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus)
	os.Exit(waitStatus.ExitStatus())
	return nil
}

func supportsExecSyscall() bool {
	return runtime.GOOS == "linux" || runtime.GOOS == "darwin" || runtime.GOOS == "freebsd"
}

func execSyscall(command string, args []string, env []string) error {
	log.Printf("Exec command %s %s", command, strings.Join(args, " "))

	argv0, err := osexec.LookPath(command)
	if err != nil {
		return err
	}

	argv := make([]string, 0, 1+len(args))
	argv = append(argv, command)
	argv = append(argv, args...)

	return syscall.Exec(argv0, argv, env)
}
<commit_msg>Revert "Make the exec command default to a login shell. Fixes #546"<commit_after>package cli

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	osexec "os/exec"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/v5/server"
	"github.com/99designs/aws-vault/v5/vault"
	"github.com/99designs/keyring"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"gopkg.in/alecthomas/kingpin.v2"
)

type ExecCommandInput struct {
	ProfileName      string
	Command          string
	Args             []string
	StartEc2Server   bool
	StartEcsServer   bool
	CredentialHelper bool
	Config           vault.Config
	SessionDuration  time.Duration
	NoSession        bool
}

// AwsCredentialHelperData is metadata for AWS CLI credential process
// See https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes
type AwsCredentialHelperData struct {
	Version         int    `json:"Version"`
	AccessKeyID     string `json:"AccessKeyId"`
	SecretAccessKey string `json:"SecretAccessKey"`
	SessionToken    string `json:"SessionToken,omitempty"`
	Expiration      string `json:"Expiration,omitempty"`
}

func ConfigureExecCommand(app *kingpin.Application, a *AwsVault) {
	input := ExecCommandInput{}

	cmd := app.Command("exec", "Executes a command with AWS credentials in the environment")

	cmd.Flag("duration", "Duration of the temporary or assume-role session. Defaults to 1h").
		Short('d').
		DurationVar(&input.SessionDuration)

	cmd.Flag("no-session", "Skip creating STS session with GetSessionToken").
		Short('n').
		BoolVar(&input.NoSession)

	cmd.Flag("region", "The AWS region").
		StringVar(&input.Config.Region)

	cmd.Flag("mfa-token", "The MFA token to use").
		Short('t').
		StringVar(&input.Config.MfaToken)

	cmd.Flag("json", "AWS credential helper. Ref: https://docs.aws.amazon.com/cli/latest/topic/config-vars.html#sourcing-credentials-from-external-processes").
		Short('j').
		BoolVar(&input.CredentialHelper)

	cmd.Flag("server", "Run a server in the background for credentials").
		Short('s').
		BoolVar(&input.StartEc2Server)

	cmd.Flag("ec2-server", "Run a EC2 metadata server in the background for credentials").
		Hidden().
		BoolVar(&input.StartEc2Server)

	cmd.Flag("ecs-server", "Run a ECS credential server in the background for credentials").
		Hidden().
		BoolVar(&input.StartEcsServer)

	cmd.Arg("profile", "Name of the profile").
		Required().
		HintAction(a.MustGetProfileNames).
		StringVar(&input.ProfileName)

	cmd.Arg("cmd", "Command to execute, defaults to $SHELL").
		Default(os.Getenv("SHELL")).
		StringVar(&input.Command)

	cmd.Arg("args", "Command arguments").
		StringsVar(&input.Args)

	cmd.Action(func(c *kingpin.ParseContext) (err error) {
		input.Config.MfaPromptMethod = a.PromptDriver
		input.Config.NonChainedGetSessionTokenDuration = input.SessionDuration
		input.Config.AssumeRoleDuration = input.SessionDuration

		cl, err := a.ConfigLoader()
		if err != nil {
			return err
		}
		keyring, err := a.Keyring()
		if err != nil {
			return err
		}

		err = ExecCommand(input, cl, keyring)
		app.FatalIfError(err, "exec")
		return nil
	})
}

func ExecCommand(input ExecCommandInput, configLoader *vault.ConfigLoader, keyring keyring.Keyring) error {
	if os.Getenv("AWS_VAULT") != "" {
		return fmt.Errorf("aws-vault sessions should be nested with care, unset $AWS_VAULT to force")
	}

	if input.StartEc2Server && input.StartEcsServer {
		return fmt.Errorf("Can't use --server with --ecs-server")
	}
	if input.StartEc2Server && input.CredentialHelper {
		return fmt.Errorf("Can't use --server with --json")
	}
	if input.StartEc2Server && input.NoSession {
		return fmt.Errorf("Can't use --server with --no-session")
	}
	if input.StartEcsServer && input.CredentialHelper {
		return fmt.Errorf("Can't use --ecs-server with --json")
	}
	if input.StartEcsServer && input.NoSession {
		return fmt.Errorf("Can't use --ecs-server with --no-session")
	}

	vault.UseSession = !input.NoSession

	configLoader.BaseConfig = input.Config
	configLoader.ActiveProfile = input.ProfileName
	config, err := configLoader.LoadFromProfile(input.ProfileName)
	if err != nil {
		return err
	}

	ckr := &vault.CredentialKeyring{Keyring: keyring}
	creds, err := vault.NewTempCredentials(config, ckr)
	if err != nil {
		return fmt.Errorf("Error getting temporary credentials: %w", err)
	}

	if input.StartEc2Server {
		return execEc2Server(input, config, creds)
	}

	if input.StartEcsServer {
		return execEcsServer(input, config, creds)
	}

	if input.CredentialHelper {
		return execCredentialHelper(input, config, creds)
	}

	return execEnvironment(input, config, creds)
}

func updateEnvForAwsVault(env environ, profileName string, region string) environ {
	env.Unset("AWS_ACCESS_KEY_ID")
	env.Unset("AWS_SECRET_ACCESS_KEY")
	env.Unset("AWS_SESSION_TOKEN")
	env.Unset("AWS_SECURITY_TOKEN")
	env.Unset("AWS_CREDENTIAL_FILE")
	env.Unset("AWS_DEFAULT_PROFILE")
	env.Unset("AWS_PROFILE")
	env.Unset("AWS_SDK_LOAD_CONFIG")

	env.Set("AWS_VAULT", profileName)

	if region != "" {
		log.Printf("Setting subprocess env: AWS_DEFAULT_REGION=%s, AWS_REGION=%s", region, region)
		env.Set("AWS_DEFAULT_REGION", region)
		env.Set("AWS_REGION", region)
	}

	return env
}

func execEc2Server(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	if err := server.StartEc2CredentialsServer(creds, config.Region); err != nil {
		return fmt.Errorf("Failed to start credential server: %w", err)
	}

	env := environ(os.Environ())
	env = updateEnvForAwsVault(env, input.ProfileName, config.Region)

	return execCmd(input.Command, input.Args, env)
}

func execEcsServer(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	uri, token, err := server.StartEcsCredentialServer(creds)
	if err != nil {
		return fmt.Errorf("Failed to start credential server: %w", err)
	}

	env := environ(os.Environ())
	env = updateEnvForAwsVault(env, input.ProfileName, config.Region)

	log.Println("Setting subprocess env AWS_CONTAINER_CREDENTIALS_FULL_URI, AWS_CONTAINER_AUTHORIZATION_TOKEN")
	env.Set("AWS_CONTAINER_CREDENTIALS_FULL_URI", uri)
	env.Set("AWS_CONTAINER_AUTHORIZATION_TOKEN", token)

	return execCmd(input.Command, input.Args, env)
}

func execCredentialHelper(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	val, err := creds.Get()
	if err != nil {
		return fmt.Errorf("Failed to get credentials for %s: %w", input.ProfileName, err)
	}

	credentialData := AwsCredentialHelperData{
		Version:         1,
		AccessKeyID:     val.AccessKeyID,
		SecretAccessKey: val.SecretAccessKey,
	}
	if val.SessionToken != "" {
		credentialData.SessionToken = val.SessionToken
	}
	if credsExpiresAt, err := creds.ExpiresAt(); err == nil {
		credentialData.Expiration = credsExpiresAt.Format("2006-01-02T15:04:05Z")
	}

	json, err := json.Marshal(&credentialData)
	if err != nil {
		return fmt.Errorf("Error creating credential json: %w", err)
	}

	fmt.Print(string(json))

	return nil
}

func execEnvironment(input ExecCommandInput, config *vault.Config, creds *credentials.Credentials) error {
	val, err := creds.Get()
	if err != nil {
		return fmt.Errorf("Failed to get credentials for %s: %w", input.ProfileName, err)
	}

	env := environ(os.Environ())
	env = updateEnvForAwsVault(env, input.ProfileName, config.Region)

	log.Println("Setting subprocess env: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY")
	env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
	env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

	if val.SessionToken != "" {
		log.Println("Setting subprocess env: AWS_SESSION_TOKEN, AWS_SECURITY_TOKEN")
		env.Set("AWS_SESSION_TOKEN", val.SessionToken)
		env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
	}
	if expiration, err := creds.ExpiresAt(); err == nil {
		log.Println("Setting subprocess env: AWS_SESSION_EXPIRATION")
		env.Set("AWS_SESSION_EXPIRATION", expiration.Format(time.RFC3339))
	}

	if !supportsExecSyscall() {
		return execCmd(input.Command, input.Args, env)
	}

	return execSyscall(input.Command, input.Args, env)
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}

func execCmd(command string, args []string, env []string) error {
	log.Printf("Starting child process: %s %s", command, strings.Join(args, " "))

	cmd := osexec.Command(command, args...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = env

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan)

	if err := cmd.Start(); err != nil {
		return err
	}

	go func() {
		for {
			sig := <-sigChan
			cmd.Process.Signal(sig)
		}
	}()

	if err := cmd.Wait(); err != nil {
		cmd.Process.Signal(os.Kill)
		return fmt.Errorf("Failed to wait for command termination: %v", err)
	}

	waitStatus := cmd.ProcessState.Sys().(syscall.WaitStatus)
	os.Exit(waitStatus.ExitStatus())
	return nil
}

func supportsExecSyscall() bool {
	return runtime.GOOS == "linux" || runtime.GOOS == "darwin" || runtime.GOOS == "freebsd"
}

func execSyscall(command string, args []string, env []string) error {
	log.Printf("Exec command %s %s", command, strings.Join(args, " "))

	argv0, err := osexec.LookPath(command)
	if err != nil {
		return err
	}

	argv := make([]string, 0, 1+len(args))
	argv = append(argv, command)
	argv = append(argv, args...)

	return syscall.Exec(argv0, argv, env)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"math"
	"strconv"
	"strings"

	"github.com/chzyer/readline"

	"github.com/AsynkronIT/protoactor-go/actor"
	"github.com/AsynkronIT/protoactor-go/remote"
	proto "github.com/gogo/protobuf/proto"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal
var _ = fmt.Errorf
var _ = math.Inf

// Function constructor - constructs new function for listing given directory
var completer = readline.NewPrefixCompleter(
	readline.PcItem("tell"),
	readline.PcItem("exit"),
)

func filterInput(r rune) (rune, bool) {
	switch r {
	// block CtrlZ feature
	case readline.CharCtrlZ:
		return r, false
	}
	return r, true
}

func main() {
	logo := `
     ___         _         ___ _    ___
    | _ \_ _ ___| |_ ___  / __| |  |_ _|
    |  _/ '_/ _ \  _/ _ \| (__| |__ | |
    |_| |_| \___/\__\___(_)___|____|___|
`
	fmt.Println(logo)

	remote.DefaultSerializerID = 1
	remote.Start("127.0.0.1:0")
	actor.SpawnNamed(actor.FromFunc(func(ctx actor.Context) {
		log.Printf("ECHO: %+v", ctx.Message())
	}), "echo")

	vars := make(map[string]string)
	vars["%address%"] = actor.ProcessRegistry.Address
	vars["%echo%"] = fmt.Sprintf(`{"Address":"%v", "Id":"echo"}`, actor.ProcessRegistry.Address)

	l, err := readline.NewEx(&readline.Config{
		Prompt:          "\033[31m»\033[0m ",
		HistoryFile:     "/tmp/readline.tmp",
		AutoComplete:    completer,
		InterruptPrompt: "^C",
		EOFPrompt:       "exit",

		HistorySearchFold:   true,
		FuncFilterInputRune: filterInput,
	})

	if err != nil {
		panic(err)
	}
	defer l.Close()

	log.SetOutput(l.Stderr())
	for {
		line, err := l.Readline()
		if err == readline.ErrInterrupt {
			if len(line) == 0 {
				break
			} else {
				continue
			}
		} else if err == io.EOF {
			break
		}

		line = strings.TrimSpace(line)
		for k, v := range vars {
			line = strings.Replace(line, k, v, 1000)
		}
		log.Println(line)
		switch {

		case strings.HasPrefix(line, "tell "):
			line = tell(line)

		case line == "exit":
			goto exit
		case line == "":
		default:
			log.Println("Unknown command :", strconv.Quote(line))
		}
	}
exit:
}
func tell(line string) string {
	parts := strings.SplitN(line, " ", 4)
	i := parts[1]
	x := strings.SplitN(i, "/", 2)
	address := x[0]
	id := x[1]
	m := &remote.JsonMessage{
		Json:     parts[3],
		TypeName: parts[2],
	}
	pid := actor.NewPID(address, id)
	remote.SendMessage(pid, m, nil, 1)
	return line
}
<commit_msg>cli error handling<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"math"
	"strconv"
	"strings"

	"github.com/chzyer/readline"

	"encoding/json"
	"github.com/AsynkronIT/protoactor-go/actor"
	"github.com/AsynkronIT/protoactor-go/remote"
	proto "github.com/gogo/protobuf/proto"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal
var _ = fmt.Errorf
var _ = math.Inf

// Function constructor - constructs new function for listing given directory
var completer = readline.NewPrefixCompleter(
	readline.PcItem("tell"),
	readline.PcItem("exit"),
)

func filterInput(r rune) (rune, bool) {
	switch r {
	// block CtrlZ feature
	case readline.CharCtrlZ:
		return r, false
	}
	return r, true
}

func main() {
	logo := `
     ___         _         ___ _    ___
    | _ \_ _ ___| |_ ___  / __| |  |_ _|
    |  _/ '_/ _ \  _/ _ \| (__| |__ | |
    |_| |_| \___/\__\___(_)___|____|___|
`
	fmt.Println(logo)

	remote.DefaultSerializerID = 1
	remote.Start("127.0.0.1:0")
	actor.SpawnNamed(actor.FromFunc(func(ctx actor.Context) {
		fmt.Printf("ECHO: %+v\n", ctx.Message())
	}), "echo")

	vars := make(map[string]string)
	vars["%address%"] = actor.ProcessRegistry.Address
	vars["%echo%"] = fmt.Sprintf(`{"Address":"%v", "Id":"echo"}`, actor.ProcessRegistry.Address)

	l, err := readline.NewEx(&readline.Config{
		Prompt:          "\033[31m»\033[0m ",
		HistoryFile:     "/tmp/readline.tmp",
		AutoComplete:    completer,
		InterruptPrompt: "^C",
		EOFPrompt:       "exit",

		HistorySearchFold:   true,
		FuncFilterInputRune: filterInput,
	})

	if err != nil {
		panic(err)
	}
	defer l.Close()

	log.SetOutput(l.Stderr())
	for {
		line, err := l.Readline()
		if err == readline.ErrInterrupt {
			if len(line) == 0 {
				break
			} else {
				continue
			}
		} else if err == io.EOF {
			break
		}

		line = strings.TrimSpace(line)
		for k, v := range vars {
			line = strings.Replace(line, k, v, 1000)
		}
		switch {

		case strings.HasPrefix(line, "tell "):
			tell(line)

		case line == "exit":
			goto exit
		case line == "":
		default:
			log.Println("Unknown command :", strconv.Quote(line))
		}
	}
exit:
}
func tell(line string) {
	parts := strings.SplitN(line, " ", 4)

	if len(parts) != 4 {
		fmt.Printf("Wrong number of arguments for `tell`. expected: pid type-name json\n")
	} else {

		pidStr := parts[1]
		typeNameStr := parts[2]
		jsonStr := parts[3]

		x := strings.SplitN(pidStr, "/", 2)
		address := x[0]
		id := x[1]

		err := parseJson(jsonStr)
		if err == nil {
			m := &remote.JsonMessage{
				Json:     jsonStr,
				TypeName: typeNameStr,
			}
			pid := actor.NewPID(address, id)
			remote.SendMessage(pid, m, nil, 1)
		} else {
			fmt.Printf("Invalid JSON payload: %v\n", err)
		}
	}
}

func parseJson(s string) error {
	var js map[string]interface{}
	return json.Unmarshal([]byte(s), &js)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/rfjakob/gocryptfs/internal/configfile"
	"github.com/rfjakob/gocryptfs/internal/prefer_openssl"
	"github.com/rfjakob/gocryptfs/internal/tlog"
)

// argContainer stores the parsed CLI options and arguments
type argContainer struct {
	debug, init, zerokey, fusedebug, openssl, passwd, foreground, version,
	plaintextnames, quiet, nosyslog, wpanic,
	longnames, allow_other, ro, reverse, aessiv, nonempty bool
	masterkey, mountpoint, cipherdir, cpuprofile, extpass,
	memprofile, ko, passfile string
	// Configuration file name override
	config             string
	notifypid, scryptn int
	// _configCustom is true when the user sets a custom config file name.
	// This is not a CLI option.
	_configCustom bool
}

var flagSet *flag.FlagSet

// prefixOArgs transform options passed via "-o foo,bar" into regular options
// like "-foo -bar" and prefixes them to the command line.
func prefixOArgs(osArgs []string) []string {
	// Need at least 3, example: gocryptfs -o foo,bar
	if len(osArgs) < 3 {
		return osArgs
	}
	// Find and extract "-o foo,bar"
	var otherArgs, oOpts []string
	for i := 1; i < len(osArgs); i++ {
		if osArgs[i] == "-o" {
			// Last argument?
			if i+1 >= len(osArgs) {
				tlog.Fatal.Printf("The \"-o\" option requires an argument")
				os.Exit(ErrExitUsage)
			}
			oOpts = strings.Split(osArgs[i+1], ",")
			// Skip over the arguments to "-o"
			i++
		} else if strings.HasPrefix(osArgs[i], "-o=") {
			oOpts = strings.Split(osArgs[i][3:], ",")
		} else {
			otherArgs = append(otherArgs, osArgs[i])
		}
	}
	// Start with program name
	newArgs := []string{osArgs[0]}
	// Add options from "-o"
	for _, o := range oOpts {
		if o == "" {
			continue
		}
		if o == "o" || o == "-o" {
			tlog.Fatal.Printf("You can't pass \"-o\" to \"-o\"")
			os.Exit(ErrExitUsage)
		}
		newArgs = append(newArgs, "-"+o)
	}
	// Add other arguments
	newArgs = append(newArgs, otherArgs...)
	return newArgs
}

// parseCliOpts - parse command line options (i.e. arguments that start with "-")
func parseCliOpts() (args argContainer) {
	os.Args = prefixOArgs(os.Args)

	var err error
	var opensslAuto string

	flagSet = flag.NewFlagSet(tlog.ProgramName, flag.ContinueOnError)
	flagSet.Usage = usageText
	flagSet.BoolVar(&args.debug, "d", false, "")
	flagSet.BoolVar(&args.debug, "debug", false, "Enable debug output")
	flagSet.BoolVar(&args.fusedebug, "fusedebug", false, "Enable fuse library debug output")
	flagSet.BoolVar(&args.init, "init", false, "Initialize encrypted directory")
	flagSet.BoolVar(&args.zerokey, "zerokey", false, "Use all-zero dummy master key")
	// Tri-state true/false/auto
	flagSet.StringVar(&opensslAuto, "openssl", "auto", "Use OpenSSL instead of built-in Go crypto")
	flagSet.BoolVar(&args.passwd, "passwd", false, "Change password")
	flagSet.BoolVar(&args.foreground, "f", false, "Stay in the foreground")
	flagSet.BoolVar(&args.version, "version", false, "Print version and exit")
	flagSet.BoolVar(&args.plaintextnames, "plaintextnames", false, "Do not encrypt file names")
	flagSet.BoolVar(&args.quiet, "q", false, "")
	flagSet.BoolVar(&args.quiet, "quiet", false, "Quiet - silence informational messages")
	flagSet.BoolVar(&args.nosyslog, "nosyslog", false, "Do not redirect output to syslog when running in the background")
	flagSet.BoolVar(&args.wpanic, "wpanic", false, "When encountering a warning, panic and exit immediately")
	flagSet.BoolVar(&args.longnames, "longnames", true, "Store names longer than 176 bytes in extra files")
	flagSet.BoolVar(&args.allow_other, "allow_other", false, "Allow other users to access the filesystem. "+
		"Only works if user_allow_other is set in /etc/fuse.conf.")
	flagSet.BoolVar(&args.ro, "ro", false, "Mount the filesystem read-only")
	flagSet.BoolVar(&args.reverse, "reverse", false, "Reverse mode")
	flagSet.BoolVar(&args.aessiv, "aessiv", false, "AES-SIV encryption")
	flagSet.BoolVar(&args.nonempty, "nonempty", false, "Allow mounting over non-empty directories")
	flagSet.StringVar(&args.masterkey, "masterkey", "", "Mount with explicit master key")
	flagSet.StringVar(&args.cpuprofile, "cpuprofile", "", "Write cpu profile to specified file")
	flagSet.StringVar(&args.memprofile, "memprofile", "", "Write memory profile to specified file")
	flagSet.StringVar(&args.config, "config", "", "Use specified config file instead of CIPHERDIR/gocryptfs.conf")
	flagSet.StringVar(&args.extpass, "extpass", "", "Use external program for the password prompt")
	flagSet.StringVar(&args.passfile, "passfile", "", "Read password from file")
	flagSet.StringVar(&args.ko, "ko", "", "Pass additional options directly to the kernel, comma-separated list")
	flagSet.IntVar(&args.notifypid, "notifypid", 0, "Send USR1 to the specified process after "+
		"successful mount - used internally for daemonization")
	flagSet.IntVar(&args.scryptn, "scryptn", configfile.ScryptDefaultLogN, "scrypt cost parameter logN. "+
		"Setting this to a lower value speeds up mounting but makes the password susceptible to brute-force attacks")
	// Ignored otions
	var dummyBool bool
	ignoreText := "(ignored for compatability)"
	flagSet.BoolVar(&dummyBool, "rw", false, ignoreText)
	flagSet.BoolVar(&dummyBool, "nosuid", false, ignoreText)
	flagSet.BoolVar(&dummyBool, "nodev", false, ignoreText)
	var dummyString string
	flagSet.StringVar(&dummyString, "o", "", "For compatability, all options can be also passed as a comma-separated list to -o.")
	// Actual parsing
	err = flagSet.Parse(os.Args[1:])
	if err == flag.ErrHelp {
		os.Exit(0)
	}
	if err != nil {
		tlog.Warn.Printf("You passed: %s", prettyArgs())
		tlog.Fatal.Printf("%v", err)
		os.Exit(ErrExitUsage)
	}
	// "-openssl" needs some post-processing
	if opensslAuto == "auto" {
		args.openssl = prefer_openssl.PreferOpenSSL()
	} else {
		args.openssl, err = strconv.ParseBool(opensslAuto)
		if err != nil {
			tlog.Fatal.Printf("Invalid \"-openssl\" setting: %v", err)
			os.Exit(ErrExitUsage)
		}
	}
	// "-passfile FILE" is a shortcut for "-extpass=/bin/cat FILE"
	if args.passfile != "" {
		args.extpass = "/bin/cat " + args.passfile
	}
	return args
}

// prettyArgs pretty-prints the command-line arguments.
func prettyArgs() string {
	pa := fmt.Sprintf("%q", os.Args[1:])
	// Get rid of "[" and "]"
	pa = pa[1 : len(pa)-1]
	return pa
}
<commit_msg>main: catch "-extpass" AND "-masterkey" usage early<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/rfjakob/gocryptfs/internal/configfile"
	"github.com/rfjakob/gocryptfs/internal/prefer_openssl"
	"github.com/rfjakob/gocryptfs/internal/tlog"
)

// argContainer stores the parsed CLI options and arguments
type argContainer struct {
	debug, init, zerokey, fusedebug, openssl, passwd, foreground, version,
	plaintextnames, quiet, nosyslog, wpanic,
	longnames, allow_other, ro, reverse, aessiv, nonempty bool
	masterkey, mountpoint, cipherdir, cpuprofile, extpass,
	memprofile, ko, passfile string
	// Configuration file name override
	config             string
	notifypid, scryptn int
	// _configCustom is true when the user sets a custom config file name.
	// This is not a CLI option.
	_configCustom bool
}

var flagSet *flag.FlagSet

// prefixOArgs transform options passed via "-o foo,bar" into regular options
// like "-foo -bar" and prefixes them to the command line.
func prefixOArgs(osArgs []string) []string {
	// Need at least 3, example: gocryptfs -o foo,bar
	if len(osArgs) < 3 {
		return osArgs
	}
	// Find and extract "-o foo,bar"
	var otherArgs, oOpts []string
	for i := 1; i < len(osArgs); i++ {
		if osArgs[i] == "-o" {
			// Last argument?
			if i+1 >= len(osArgs) {
				tlog.Fatal.Printf("The \"-o\" option requires an argument")
				os.Exit(ErrExitUsage)
			}
			oOpts = strings.Split(osArgs[i+1], ",")
			// Skip over the arguments to "-o"
			i++
		} else if strings.HasPrefix(osArgs[i], "-o=") {
			oOpts = strings.Split(osArgs[i][3:], ",")
		} else {
			otherArgs = append(otherArgs, osArgs[i])
		}
	}
	// Start with program name
	newArgs := []string{osArgs[0]}
	// Add options from "-o"
	for _, o := range oOpts {
		if o == "" {
			continue
		}
		if o == "o" || o == "-o" {
			tlog.Fatal.Printf("You can't pass \"-o\" to \"-o\"")
			os.Exit(ErrExitUsage)
		}
		newArgs = append(newArgs, "-"+o)
	}
	// Add other arguments
	newArgs = append(newArgs, otherArgs...)
	return newArgs
}

// parseCliOpts - parse command line options (i.e. arguments that start with "-")
func parseCliOpts() (args argContainer) {
	os.Args = prefixOArgs(os.Args)

	var err error
	var opensslAuto string

	flagSet = flag.NewFlagSet(tlog.ProgramName, flag.ContinueOnError)
	flagSet.Usage = usageText
	flagSet.BoolVar(&args.debug, "d", false, "")
	flagSet.BoolVar(&args.debug, "debug", false, "Enable debug output")
	flagSet.BoolVar(&args.fusedebug, "fusedebug", false, "Enable fuse library debug output")
	flagSet.BoolVar(&args.init, "init", false, "Initialize encrypted directory")
	flagSet.BoolVar(&args.zerokey, "zerokey", false, "Use all-zero dummy master key")
	// Tri-state true/false/auto
	flagSet.StringVar(&opensslAuto, "openssl", "auto", "Use OpenSSL instead of built-in Go crypto")
	flagSet.BoolVar(&args.passwd, "passwd", false, "Change password")
	flagSet.BoolVar(&args.foreground, "f", false, "Stay in the foreground")
	flagSet.BoolVar(&args.version, "version", false, "Print version and exit")
	flagSet.BoolVar(&args.plaintextnames, "plaintextnames", false, "Do not encrypt file names")
	flagSet.BoolVar(&args.quiet, "q", false, "")
	flagSet.BoolVar(&args.quiet, "quiet", false, "Quiet - silence informational messages")
	flagSet.BoolVar(&args.nosyslog, "nosyslog", false, "Do not redirect output to syslog when running in the background")
	flagSet.BoolVar(&args.wpanic, "wpanic", false, "When encountering a warning, panic and exit immediately")
	flagSet.BoolVar(&args.longnames, "longnames", true, "Store names longer than 176 bytes in extra files")
	flagSet.BoolVar(&args.allow_other, "allow_other", false, "Allow other users to access the filesystem. "+
		"Only works if user_allow_other is set in /etc/fuse.conf.")
	flagSet.BoolVar(&args.ro, "ro", false, "Mount the filesystem read-only")
	flagSet.BoolVar(&args.reverse, "reverse", false, "Reverse mode")
	flagSet.BoolVar(&args.aessiv, "aessiv", false, "AES-SIV encryption")
	flagSet.BoolVar(&args.nonempty, "nonempty", false, "Allow mounting over non-empty directories")
	flagSet.StringVar(&args.masterkey, "masterkey", "", "Mount with explicit master key")
	flagSet.StringVar(&args.cpuprofile, "cpuprofile", "", "Write cpu profile to specified file")
	flagSet.StringVar(&args.memprofile, "memprofile", "", "Write memory profile to specified file")
	flagSet.StringVar(&args.config, "config", "", "Use specified config file instead of CIPHERDIR/gocryptfs.conf")
	flagSet.StringVar(&args.extpass, "extpass", "", "Use external program for the password prompt")
	flagSet.StringVar(&args.passfile, "passfile", "", "Read password from file")
	flagSet.StringVar(&args.ko, "ko", "", "Pass additional options directly to the kernel, comma-separated list")
	flagSet.IntVar(&args.notifypid, "notifypid", 0, "Send USR1 to the specified process after "+
		"successful mount - used internally for daemonization")
	flagSet.IntVar(&args.scryptn, "scryptn", configfile.ScryptDefaultLogN, "scrypt cost parameter logN. "+
		"Setting this to a lower value speeds up mounting but makes the password susceptible to brute-force attacks")
	// Ignored otions
	var dummyBool bool
	ignoreText := "(ignored for compatability)"
	flagSet.BoolVar(&dummyBool, "rw", false, ignoreText)
	flagSet.BoolVar(&dummyBool, "nosuid", false, ignoreText)
	flagSet.BoolVar(&dummyBool, "nodev", false, ignoreText)
	var dummyString string
	flagSet.StringVar(&dummyString, "o", "", "For compatability, all options can be also passed as a comma-separated list to -o.")
	// Actual parsing
	err = flagSet.Parse(os.Args[1:])
	if err == flag.ErrHelp {
		os.Exit(0)
	}
	if err != nil {
		tlog.Warn.Printf("You passed: %s", prettyArgs())
		tlog.Fatal.Printf("%v", err)
		os.Exit(ErrExitUsage)
	}
	// "-openssl" needs some post-processing
	if opensslAuto == "auto" {
		args.openssl = prefer_openssl.PreferOpenSSL()
	} else {
		args.openssl, err = strconv.ParseBool(opensslAuto)
		if err != nil {
			tlog.Fatal.Printf("Invalid \"-openssl\" setting: %v", err)
			os.Exit(ErrExitUsage)
		}
	}
	// "-passfile FILE" is a shortcut for "-extpass=/bin/cat FILE"
	if args.passfile != "" {
		args.extpass = "/bin/cat " + args.passfile
	}
	if args.extpass != "" && args.masterkey != "" {
		tlog.Fatal.Printf("The options -extpass and -masterkey cannot be used at the same time")
		os.Exit(ErrExitUsage)
	}
	return args
}

// prettyArgs pretty-prints the command-line arguments.
func prettyArgs() string {
	pa := fmt.Sprintf("%q", os.Args[1:])
	// Get rid of "[" and "]"
	pa = pa[1 : len(pa)-1]
	return pa
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package httplog

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
	"runtime"
	"time"

	"github.com/golang/glog"
)

// StacktracePred returns true if a stacktrace should be logged for this status.
type StacktracePred func(httpStatus int) (logStacktrace bool)

type logger interface {
	Addf(format string, data ...interface{})
}

// Add a layer on top of ResponseWriter, so we can track latency and error
// message sources.
//
// TODO now that we're using go-restful, we shouldn't need to be wrapping
// the http.ResponseWriter. We can recover panics from go-restful, and
// the logging value is questionable.
type respLogger struct {
	hijacked       bool
	statusRecorded bool
	status         int
	statusStack    string
	addedInfo      string
	startTime      time.Time

	captureErrorOutput bool

	req *http.Request
	w   http.ResponseWriter

	logStacktracePred StacktracePred
}

// Simple logger that logs immediately when Addf is called
type passthroughLogger struct{}

// Addf logs info immediately.
func (passthroughLogger) Addf(format string, data ...interface{}) {
	glog.V(2).Info(fmt.Sprintf(format, data...))
}

// DefaultStacktracePred is the default implementation of StacktracePred.
func DefaultStacktracePred(status int) bool {
	return (status < http.StatusOK || status >= http.StatusInternalServerError) && status != http.StatusSwitchingProtocols
}

// NewLogged turns a normal response writer into a logged response writer.
//
// Usage:
//
// defer NewLogged(req, &w).StacktraceWhen(StatusIsNot(200, 202)).Log()
//
// (Only the call to Log() is deferred, so you can set everything up in one line!)
//
// Note that this *changes* your writer, to route response writing actions
// through the logger.
//
// Use LogOf(w).Addf(...) to log something along with the response result.
func NewLogged(req *http.Request, w *http.ResponseWriter) *respLogger {
	if _, ok := (*w).(*respLogger); ok {
		// Don't double-wrap!
		panic("multiple NewLogged calls!")
	}
	rl := &respLogger{
		startTime:         time.Now(),
		req:               req,
		w:                 *w,
		logStacktracePred: DefaultStacktracePred,
	}
	*w = rl // hijack caller's writer!
	return rl
}

// LogOf returns the logger hiding in w. If there is not an existing logger
// then a passthroughLogger will be created which will log to stdout immediately
// when Addf is called.
func LogOf(req *http.Request, w http.ResponseWriter) logger {
	if _, exists := w.(*respLogger); !exists {
		pl := &passthroughLogger{}
		return pl
	}
	if rl, ok := w.(*respLogger); ok {
		return rl
	}
	panic("Unable to find or create the logger!")
}

// Unlogged returns the original ResponseWriter, or w if it is not our inserted logger.
func Unlogged(w http.ResponseWriter) http.ResponseWriter {
	if rl, ok := w.(*respLogger); ok {
		return rl.w
	}
	return w
}

// StacktraceWhen sets the stacktrace logging predicate, which decides when to log a stacktrace.
// There's a default, so you don't need to call this unless you don't like the default.
func (rl *respLogger) StacktraceWhen(pred StacktracePred) *respLogger {
	rl.logStacktracePred = pred
	return rl
}

// StatusIsNot returns a StacktracePred which will cause stacktraces to be logged
// for any status *not* in the given list.
func StatusIsNot(statuses ...int) StacktracePred {
	return func(status int) bool {
		for _, s := range statuses {
			if status == s {
				return false
			}
		}
		return true
	}
}

// Addf adds additional data to be logged with this request.
func (rl *respLogger) Addf(format string, data ...interface{}) {
	rl.addedInfo += "\n" + fmt.Sprintf(format, data...)
}

// Log is intended to be called once at the end of your request handler, via defer
func (rl *respLogger) Log() {
	latency := time.Since(rl.startTime)
	if glog.V(2) {
		if !rl.hijacked {
			glog.InfoDepth(1, fmt.Sprintf("%s %s: (%v) %v%v%v [%s %s]", rl.req.Method, rl.req.RequestURI, latency, rl.status, rl.statusStack, rl.addedInfo, rl.req.Header["User-Agent"], rl.req.RemoteAddr))
		} else {
			glog.InfoDepth(1, fmt.Sprintf("%s %s: (%v) hijacked [%s %s]", rl.req.Method, rl.req.RequestURI, latency, rl.req.Header["User-Agent"], rl.req.RemoteAddr))
		}
	}
}

// Header implements http.ResponseWriter.
func (rl *respLogger) Header() http.Header {
	return rl.w.Header()
}

// Write implements http.ResponseWriter.
func (rl *respLogger) Write(b []byte) (int, error) {
	if !rl.statusRecorded {
		rl.recordStatus(http.StatusOK) // Default if WriteHeader hasn't been called
	}
	if rl.captureErrorOutput {
		rl.Addf("logging error output: %q\n", string(b))
	}
	return rl.w.Write(b)
}

// Flush implements http.Flusher even if the underlying http.Writer doesn't implement it.
// Flush is used for streaming purposes and allows to flush buffered data to the client.
func (rl *respLogger) Flush() {
	if flusher, ok := rl.w.(http.Flusher); ok {
		flusher.Flush()
	} else if glog.V(2) {
		glog.InfoDepth(1, fmt.Sprintf("Unable to convert %+v into http.Flusher", rl.w))
	}
}

// WriteHeader implements http.ResponseWriter.
func (rl *respLogger) WriteHeader(status int) {
	rl.recordStatus(status)
	rl.w.WriteHeader(status)
}

// Hijack implements http.Hijacker.
func (rl *respLogger) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	rl.hijacked = true
	return rl.w.(http.Hijacker).Hijack()
}

// CloseNotify implements http.CloseNotifier
func (rl *respLogger) CloseNotify() <-chan bool {
	return rl.w.(http.CloseNotifier).CloseNotify()
}

func (rl *respLogger) recordStatus(status int) {
	rl.status = status
	rl.statusRecorded = true
	if rl.logStacktracePred(status) {
		// Only log stacks for errors
		stack := make([]byte, 50*1024)
		stack = stack[:runtime.Stack(stack, false)]
		rl.statusStack = "\n" + string(stack)
		rl.captureErrorOutput = true
	} else {
		rl.statusStack = ""
	}
}
<commit_msg>UPSTREAM: <carry>: Fix to avoid REST API calls at log level 2.<commit_after>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package httplog

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
	"runtime"
	"time"

	"github.com/golang/glog"
)

// StacktracePred returns true if a stacktrace should be logged for this status.
type StacktracePred func(httpStatus int) (logStacktrace bool)

type logger interface {
	Addf(format string, data ...interface{})
}

// Add a layer on top of ResponseWriter, so we can track latency and error
// message sources.
//
// TODO now that we're using go-restful, we shouldn't need to be wrapping
// the http.ResponseWriter. We can recover panics from go-restful, and
// the logging value is questionable.
type respLogger struct {
	hijacked       bool
	statusRecorded bool
	status         int
	statusStack    string
	addedInfo      string
	startTime      time.Time

	captureErrorOutput bool

	req *http.Request
	w   http.ResponseWriter

	logStacktracePred StacktracePred
}

// Simple logger that logs immediately when Addf is called
type passthroughLogger struct{}

// Addf logs info immediately.
func (passthroughLogger) Addf(format string, data ...interface{}) {
	glog.V(2).Info(fmt.Sprintf(format, data...))
}

// DefaultStacktracePred is the default implementation of StacktracePred.
func DefaultStacktracePred(status int) bool {
	return (status < http.StatusOK || status >= http.StatusInternalServerError) && status != http.StatusSwitchingProtocols
}

// NewLogged turns a normal response writer into a logged response writer.
//
// Usage:
//
// defer NewLogged(req, &w).StacktraceWhen(StatusIsNot(200, 202)).Log()
//
// (Only the call to Log() is deferred, so you can set everything up in one line!)
//
// Note that this *changes* your writer, to route response writing actions
// through the logger.
//
// Use LogOf(w).Addf(...) to log something along with the response result.
func NewLogged(req *http.Request, w *http.ResponseWriter) *respLogger {
	if _, ok := (*w).(*respLogger); ok {
		// Don't double-wrap!
		panic("multiple NewLogged calls!")
	}
	rl := &respLogger{
		startTime:         time.Now(),
		req:               req,
		w:                 *w,
		logStacktracePred: DefaultStacktracePred,
	}
	*w = rl // hijack caller's writer!
	return rl
}

// LogOf returns the logger hiding in w. If there is not an existing logger
// then a passthroughLogger will be created which will log to stdout immediately
// when Addf is called.
func LogOf(req *http.Request, w http.ResponseWriter) logger {
	if _, exists := w.(*respLogger); !exists {
		pl := &passthroughLogger{}
		return pl
	}
	if rl, ok := w.(*respLogger); ok {
		return rl
	}
	panic("Unable to find or create the logger!")
}

// Unlogged returns the original ResponseWriter, or w if it is not our inserted logger.
func Unlogged(w http.ResponseWriter) http.ResponseWriter {
	if rl, ok := w.(*respLogger); ok {
		return rl.w
	}
	return w
}

// StacktraceWhen sets the stacktrace logging predicate, which decides when to log a stacktrace.
// There's a default, so you don't need to call this unless you don't like the default.
func (rl *respLogger) StacktraceWhen(pred StacktracePred) *respLogger {
	rl.logStacktracePred = pred
	return rl
}

// StatusIsNot returns a StacktracePred which will cause stacktraces to be logged
// for any status *not* in the given list.
func StatusIsNot(statuses ...int) StacktracePred {
	return func(status int) bool {
		for _, s := range statuses {
			if status == s {
				return false
			}
		}
		return true
	}
}

// Addf adds additional data to be logged with this request.
func (rl *respLogger) Addf(format string, data ...interface{}) {
	rl.addedInfo += "\n" + fmt.Sprintf(format, data...)
}

// Log is intended to be called once at the end of your request handler, via defer
func (rl *respLogger) Log() {
	latency := time.Since(rl.startTime)
	if glog.V(3) {
		if !rl.hijacked {
			glog.InfoDepth(1, fmt.Sprintf("%s %s: (%v) %v%v%v [%s %s]", rl.req.Method, rl.req.RequestURI, latency, rl.status, rl.statusStack, rl.addedInfo, rl.req.Header["User-Agent"], rl.req.RemoteAddr))
		} else {
			glog.InfoDepth(1, fmt.Sprintf("%s %s: (%v) hijacked [%s %s]", rl.req.Method, rl.req.RequestURI, latency, rl.req.Header["User-Agent"], rl.req.RemoteAddr))
		}
	}
}

// Header implements http.ResponseWriter.
func (rl *respLogger) Header() http.Header {
	return rl.w.Header()
}

// Write implements http.ResponseWriter.
func (rl *respLogger) Write(b []byte) (int, error) {
	if !rl.statusRecorded {
		rl.recordStatus(http.StatusOK) // Default if WriteHeader hasn't been called
	}
	if rl.captureErrorOutput {
		rl.Addf("logging error output: %q\n", string(b))
	}
	return rl.w.Write(b)
}

// Flush implements http.Flusher even if the underlying http.Writer doesn't implement it.
// Flush is used for streaming purposes and allows to flush buffered data to the client.
func (rl *respLogger) Flush() {
	if flusher, ok := rl.w.(http.Flusher); ok {
		flusher.Flush()
	} else if glog.V(2) {
		glog.InfoDepth(1, fmt.Sprintf("Unable to convert %+v into http.Flusher", rl.w))
	}
}

// WriteHeader implements http.ResponseWriter.
func (rl *respLogger) WriteHeader(status int) {
	rl.recordStatus(status)
	rl.w.WriteHeader(status)
}

// Hijack implements http.Hijacker.
func (rl *respLogger) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	rl.hijacked = true
	return rl.w.(http.Hijacker).Hijack()
}

// CloseNotify implements http.CloseNotifier
func (rl *respLogger) CloseNotify() <-chan bool {
	return rl.w.(http.CloseNotifier).CloseNotify()
}

func (rl *respLogger) recordStatus(status int) {
	rl.status = status
	rl.statusRecorded = true
	if rl.logStacktracePred(status) {
		// Only log stacks for errors
		stack := make([]byte, 50*1024)
		stack = stack[:runtime.Stack(stack, false)]
		rl.statusStack = "\n" + string(stack)
		rl.captureErrorOutput = true
	} else {
		rl.statusStack = ""
	}
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2016 Alexey Derbyshev
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package ugo

// ChainWrapper is the special struct,
// containing resulting and middleware data
type ChainWrapper struct {
	Mid Seq    // Mid is for middleware calculations
	Res Object // Res if for resulting data
}

// Each is a chaining wrapper for #Each
func (wrapper *ChainWrapper) Each(cb Action) *ChainWrapper {
	Each(wrapper.Mid, cb)
	return wrapper
}

// ForEach is a chaining wrapper for #ForEach
func (wrapper *ChainWrapper) ForEach(cb Action) *ChainWrapper {
	return wrapper.Each(cb)
}

// Map is a chaining wrapper for #Map
func (wrapper *ChainWrapper) Map(cb Callback) *ChainWrapper {
	wrapper.Mid = Map(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Collect is a chaining wrapper for #Collect
func (wrapper *ChainWrapper) Collect(cb Callback) *ChainWrapper {
	return wrapper.Map(cb)
}

// Filter is a chaining wrapper for #Filter
func (wrapper *ChainWrapper) Filter(cb Predicate) *ChainWrapper {
	wrapper.Mid = Filter(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Select is a chaining wrapper for #Select
func (wrapper *ChainWrapper) Select(cb Predicate) *ChainWrapper {
	return wrapper.Filter(cb)
}

// Reject is a chaining wrapper for #Reject
func (wrapper *ChainWrapper) Reject(cb Predicate) *ChainWrapper {
	wrapper.Mid = Reject(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Reduce is a chaining wrapper for #Reduce
func (wrapper *ChainWrapper) Reduce(cb Collector, initial Object) *ChainWrapper {
	wrapper.Res = Reduce(wrapper.Mid, cb, initial)
	wrapper.Mid = nil
	return wrapper
}

// Inject is a chaining wrapper for #Inject
func (wrapper *ChainWrapper) Inject(cb Collector, initial Object) *ChainWrapper {
	return wrapper.Reduce(cb, initial)
}

// FoldL is a chaining wrapper for #FoldL
func (wrapper *ChainWrapper) FoldL(cb Collector, initial Object) *ChainWrapper {
	return wrapper.Reduce(cb, initial)
}

// ReduceRight is a chaining wrapper for #ReduceRight
func (wrapper *ChainWrapper) ReduceRight(cb Collector, initial Object) *ChainWrapper {
	wrapper.Res = ReduceRight(wrapper.Mid, cb, initial)
	wrapper.Mid = nil
	return wrapper
}

// FoldR is a chaining wrapper for #FoldR
func (wrapper *ChainWrapper) FoldR(cb Collector, initial Object) *ChainWrapper {
	return wrapper.ReduceRight(cb, initial)
}

// Min is a chaining wrapper for #Min
func (wrapper *ChainWrapper) Min(cb Comparator) *ChainWrapper {
	wrapper.Res = Min(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Max is a chaining wrapper for #Max
func (wrapper *ChainWrapper) Max(cb Comparator) *ChainWrapper {
	wrapper.Res = Max(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Find is a chaining wrapper for #Find
func (wrapper *ChainWrapper) Find(cb Predicate) *ChainWrapper {
	wrapper.Res = Find(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Detect is a chaining wrapper for #Detect
func (wrapper *ChainWrapper) Detect(cb Predicate) *ChainWrapper {
	return wrapper.Find(cb)
}

// FindLast is a chaining wrapper for #FindLast
func (wrapper *ChainWrapper) FindLast(cb Predicate) *ChainWrapper {
	wrapper.Res = FindLast(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// FindIndex is a chaining wrapper for #FindIndex
func (wrapper *ChainWrapper) FindIndex(cb Predicate) *ChainWrapper {
	wrapper.Res = FindIndex(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// FindLastIndex is a chaining wrapper for #FindLastIndex
func (wrapper *ChainWrapper) FindLastIndex(cb Predicate) *ChainWrapper {
	wrapper.Res = FindLastIndex(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Some is a chaining wrapper for #Some
func (wrapper *ChainWrapper) Some(cb Predicate) *ChainWrapper {
	wrapper.Res = Some(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Any is a chaining wrapper for #Any
func (wrapper *ChainWrapper) Any(cb Predicate) *ChainWrapper {
	return wrapper.Some(cb)
}

// IndexOf is a chaining wrapper for #IndexOf
func (wrapper *ChainWrapper) IndexOf(target Object, isSorted bool, cb Comparator) *ChainWrapper {
	wrapper.Res = IndexOf(wrapper.Mid, target, isSorted, cb)
	wrapper.Mid = nil
	return wrapper
}

// LastIndexOf is a chaining wrapper for #LastIndexOf
func (wrapper *ChainWrapper) LastIndexOf(target Object, cb Comparator) *ChainWrapper {
	wrapper.Res = LastIndexOf(wrapper.Mid, target, cb)
	wrapper.Mid = nil
	return wrapper
}

// Contains is a chaining wrapper for #Contains
func (wrapper *ChainWrapper) Contains(target Object, isSorted bool, cb Comparator) *ChainWrapper {
	wrapper.Res = Contains(wrapper.Mid, target, isSorted, cb)
	wrapper.Mid = nil
	return wrapper
}

// Includes is a chaining wrapper for #Includes
func (wrapper *ChainWrapper) Includes(target Object, isSorted bool, cb Comparator) *ChainWrapper {
	return wrapper.Contains(target, isSorted, cb)
}

// Every is a chaining wrapper for #Every
func (wrapper *ChainWrapper) Every(cb Predicate) *ChainWrapper {
	wrapper.Res = Every(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// All is a chaining wrapper for #All
func (wrapper *ChainWrapper) All(cb Predicate) *ChainWrapper {
	return wrapper.Every(cb)
}

// Uniq is a chaining wrapper for #Uniq
func (wrapper *ChainWrapper) Uniq(cb Comparator) *ChainWrapper {
	wrapper.Mid = Uniq(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Unique is a chaining wrapper for #Unique
func (wrapper *ChainWrapper) Unique(cb Comparator) *ChainWrapper {
	return wrapper.Uniq(cb)
}

// Differenc is a chaining wrapper for #Differenc
func (wrapper *ChainWrapper) Difference(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Mid = Difference(wrapper.Mid, other, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Without is a chaining wrapper for #Without
func (wrapper *ChainWrapper) Without(nonGrata Object, cb Comparator) *ChainWrapper {
	wrapper.Mid = Without(wrapper.Mid, nonGrata, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Intersection is a chaining wrapper for #Intersection
func (wrapper *ChainWrapper) Intersection(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Mid = Intersection(wrapper.Mid, other, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Union is a chaining wrapper for #Union
func (wrapper *ChainWrapper) Union(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Mid = Union(wrapper.Mid, other, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// SortBy is a chaining wrapper for #SortBy
func (wrapper *ChainWrapper) SortBy(cb Comparator) *ChainWrapper {
	wrapper.Mid = SortBy(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// CountBy is a chaining wrapper for #CountBy
func (wrapper *ChainWrapper) CountBy(cb Callback) *ChainWrapper {
	wrapper.Res = CountBy(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// GroupBy is a chaining wrapper for #GroupBy
func (wrapper *ChainWrapper) GroupBy(cb Callback) *ChainWrapper {
	wrapper.Res = GroupBy(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Remove is a chaining wrapper for #Remove
func (wrapper *ChainWrapper) Remove(pos int) *ChainWrapper {
	wrapper.Mid = Remove(wrapper.Mid, pos)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Insert is a chaining wrapper for #Insert
func (wrapper *ChainWrapper) Insert(tg Object, pos int) *ChainWrapper {
	wrapper.Mid = Insert(wrapper.Mid, tg, pos)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Concat is a chaining wrapper for #Concat
func (wrapper *ChainWrapper) Concat(next Seq) *ChainWrapper {
	wrapper.Mid = Concat(wrapper.Mid, next)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Shuffle is a chaining wrapper for #Shuffle
func (wrapper *ChainWrapper) Shuffle() *ChainWrapper {
	wrapper.Mid = ShuffledCopy(wrapper.Mid)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Reverse is a chaining wrapper for #Reverse
func (wrapper *ChainWrapper) Reverse() *ChainWrapper {
	wrapper.Mid = ReversedCopy(wrapper.Mid)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// EqualsStrict is a chaining wrapper for #EqualsStrict
func (wrapper *ChainWrapper) EqualsStrict(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Res = EqualsStrict(wrapper.Mid, other, cb)
	wrapper.Mid = nil
	return wrapper
}

// EqualsNotStrict is a chaining wrapper for #EqualsNotStrict
func (wrapper *ChainWrapper) EqualsNotStrict(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Res = EqualsNotStrict(wrapper.Mid, other, cb)
	wrapper.Mid = nil
	return wrapper
}

// Value returns result of calculations, you've done through chaining calls
func (wrapper *ChainWrapper) Value() Object {
	return wrapper.Res
}
<commit_msg>Typo fix<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2016 Alexey Derbyshev
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package ugo

// ChainWrapper is the special struct,
// containing resulting and middleware data
type ChainWrapper struct {
	Mid Seq    // Mid is for middleware calculations
	Res Object // Res if for resulting data
}

// Each is a chaining wrapper for #Each
func (wrapper *ChainWrapper) Each(cb Action) *ChainWrapper {
	Each(wrapper.Mid, cb)
	return wrapper
}

// ForEach is a chaining wrapper for #ForEach
func (wrapper *ChainWrapper) ForEach(cb Action) *ChainWrapper {
	return wrapper.Each(cb)
}

// Map is a chaining wrapper for #Map
func (wrapper *ChainWrapper) Map(cb Callback) *ChainWrapper {
	wrapper.Mid = Map(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Collect is a chaining wrapper for #Collect
func (wrapper *ChainWrapper) Collect(cb Callback) *ChainWrapper {
	return wrapper.Map(cb)
}

// Filter is a chaining wrapper for #Filter
func (wrapper *ChainWrapper) Filter(cb Predicate) *ChainWrapper {
	wrapper.Mid = Filter(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Select is a chaining wrapper for #Select
func (wrapper *ChainWrapper) Select(cb Predicate) *ChainWrapper {
	return wrapper.Filter(cb)
}

// Reject is a chaining wrapper for #Reject
func (wrapper *ChainWrapper) Reject(cb Predicate) *ChainWrapper {
	wrapper.Mid = Reject(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Reduce is a chaining wrapper for #Reduce
func (wrapper *ChainWrapper) Reduce(cb Collector, initial Object) *ChainWrapper {
	wrapper.Res = Reduce(wrapper.Mid, cb, initial)
	wrapper.Mid = nil
	return wrapper
}

// Inject is a chaining wrapper for #Inject
func (wrapper *ChainWrapper) Inject(cb Collector, initial Object) *ChainWrapper {
	return wrapper.Reduce(cb, initial)
}

// FoldL is a chaining wrapper for #FoldL
func (wrapper *ChainWrapper) FoldL(cb Collector, initial Object) *ChainWrapper {
	return wrapper.Reduce(cb, initial)
}

// ReduceRight is a chaining wrapper for #ReduceRight
func (wrapper *ChainWrapper) ReduceRight(cb Collector, initial Object) *ChainWrapper {
	wrapper.Res = ReduceRight(wrapper.Mid, cb, initial)
	wrapper.Mid = nil
	return wrapper
}

// FoldR is a chaining wrapper for #FoldR
func (wrapper *ChainWrapper) FoldR(cb Collector, initial Object) *ChainWrapper {
	return wrapper.ReduceRight(cb, initial)
}

// Min is a chaining wrapper for #Min
func (wrapper *ChainWrapper) Min(cb Comparator) *ChainWrapper {
	wrapper.Res = Min(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Max is a chaining wrapper for #Max
func (wrapper *ChainWrapper) Max(cb Comparator) *ChainWrapper {
	wrapper.Res = Max(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Find is a chaining wrapper for #Find
func (wrapper *ChainWrapper) Find(cb Predicate) *ChainWrapper {
	wrapper.Res = Find(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Detect is a chaining wrapper for #Detect
func (wrapper *ChainWrapper) Detect(cb Predicate) *ChainWrapper {
	return wrapper.Find(cb)
}

// FindLast is a chaining wrapper for #FindLast
func (wrapper *ChainWrapper) FindLast(cb Predicate) *ChainWrapper {
	wrapper.Res = FindLast(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// FindIndex is a chaining wrapper for #FindIndex
func (wrapper *ChainWrapper) FindIndex(cb Predicate) *ChainWrapper {
	wrapper.Res = FindIndex(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// FindLastIndex is a chaining wrapper for #FindLastIndex
func (wrapper *ChainWrapper) FindLastIndex(cb Predicate) *ChainWrapper {
	wrapper.Res = FindLastIndex(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Some is a chaining wrapper for #Some
func (wrapper *ChainWrapper) Some(cb Predicate) *ChainWrapper {
	wrapper.Res = Some(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Any is a chaining wrapper for #Any
func (wrapper *ChainWrapper) Any(cb Predicate) *ChainWrapper {
	return wrapper.Some(cb)
}

// IndexOf is a chaining wrapper for #IndexOf
func (wrapper *ChainWrapper) IndexOf(target Object, isSorted bool, cb Comparator) *ChainWrapper {
	wrapper.Res = IndexOf(wrapper.Mid, target, isSorted, cb)
	wrapper.Mid = nil
	return wrapper
}

// LastIndexOf is a chaining wrapper for #LastIndexOf
func (wrapper *ChainWrapper) LastIndexOf(target Object, cb Comparator) *ChainWrapper {
	wrapper.Res = LastIndexOf(wrapper.Mid, target, cb)
	wrapper.Mid = nil
	return wrapper
}

// Contains is a chaining wrapper for #Contains
func (wrapper *ChainWrapper) Contains(target Object, isSorted bool, cb Comparator) *ChainWrapper {
	wrapper.Res = Contains(wrapper.Mid, target, isSorted, cb)
	wrapper.Mid = nil
	return wrapper
}

// Includes is a chaining wrapper for #Includes
func (wrapper *ChainWrapper) Includes(target Object, isSorted bool, cb Comparator) *ChainWrapper {
	return wrapper.Contains(target, isSorted, cb)
}

// Every is a chaining wrapper for #Every
func (wrapper *ChainWrapper) Every(cb Predicate) *ChainWrapper {
	wrapper.Res = Every(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// All is a chaining wrapper for #All
func (wrapper *ChainWrapper) All(cb Predicate) *ChainWrapper {
	return wrapper.Every(cb)
}

// Uniq is a chaining wrapper for #Uniq
func (wrapper *ChainWrapper) Uniq(cb Comparator) *ChainWrapper {
	wrapper.Mid = Uniq(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Unique is a chaining wrapper for #Unique
func (wrapper *ChainWrapper) Unique(cb Comparator) *ChainWrapper {
	return wrapper.Uniq(cb)
}

// Difference is a chaining wrapper for #Difference
func (wrapper *ChainWrapper) Difference(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Mid = Difference(wrapper.Mid, other, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Without is a chaining wrapper for #Without
func (wrapper *ChainWrapper) Without(nonGrata Object, cb Comparator) *ChainWrapper {
	wrapper.Mid = Without(wrapper.Mid, nonGrata, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Intersection is a chaining wrapper for #Intersection
func (wrapper *ChainWrapper) Intersection(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Mid = Intersection(wrapper.Mid, other, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Union is a chaining wrapper for #Union
func (wrapper *ChainWrapper) Union(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Mid = Union(wrapper.Mid, other, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// SortBy is a chaining wrapper for #SortBy
func (wrapper *ChainWrapper) SortBy(cb Comparator) *ChainWrapper {
	wrapper.Mid = SortBy(wrapper.Mid, cb)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// CountBy is a chaining wrapper for #CountBy
func (wrapper *ChainWrapper) CountBy(cb Callback) *ChainWrapper {
	wrapper.Res = CountBy(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// GroupBy is a chaining wrapper for #GroupBy
func (wrapper *ChainWrapper) GroupBy(cb Callback) *ChainWrapper {
	wrapper.Res = GroupBy(wrapper.Mid, cb)
	wrapper.Mid = nil
	return wrapper
}

// Remove is a chaining wrapper for #Remove
func (wrapper *ChainWrapper) Remove(pos int) *ChainWrapper {
	wrapper.Mid = Remove(wrapper.Mid, pos)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Insert is a chaining wrapper for #Insert
func (wrapper *ChainWrapper) Insert(tg Object, pos int) *ChainWrapper {
	wrapper.Mid = Insert(wrapper.Mid, tg, pos)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Concat is a chaining wrapper for #Concat
func (wrapper *ChainWrapper) Concat(next Seq) *ChainWrapper {
	wrapper.Mid = Concat(wrapper.Mid, next)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Shuffle is a chaining wrapper for #Shuffle
func (wrapper *ChainWrapper) Shuffle() *ChainWrapper {
	wrapper.Mid = ShuffledCopy(wrapper.Mid)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// Reverse is a chaining wrapper for #Reverse
func (wrapper *ChainWrapper) Reverse() *ChainWrapper {
	wrapper.Mid = ReversedCopy(wrapper.Mid)
	wrapper.Res = wrapper.Mid
	return wrapper
}

// EqualsStrict is a chaining wrapper for #EqualsStrict
func (wrapper *ChainWrapper) EqualsStrict(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Res = EqualsStrict(wrapper.Mid, other, cb)
	wrapper.Mid = nil
	return wrapper
}

// EqualsNotStrict is a chaining wrapper for #EqualsNotStrict
func (wrapper *ChainWrapper) EqualsNotStrict(other Seq, cb Comparator) *ChainWrapper {
	wrapper.Res = EqualsNotStrict(wrapper.Mid, other, cb)
	wrapper.Mid = nil
	return wrapper
}

// Value returns result of calculations, you've done through chaining calls
func (wrapper *ChainWrapper) Value() Object {
	return wrapper.Res
}
<|endoftext|>


<commit_before>// Copyright 2018 Frédéric Guillot. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package cli // import "miniflux.app/cli"

import (
	"fmt"
	"runtime"

	"miniflux.app/version"
)

func info() {
	fmt.Println("Version:", version.Version)
	fmt.Println("Build Date:", version.BuildDate)
	fmt.Println("Go Version:", runtime.Version())
}
<commit_msg>Add compiler, Arch, and OS to info command<commit_after>// Copyright 2018 Frédéric Guillot. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package cli // import "miniflux.app/cli"

import (
	"fmt"
	"runtime"

	"miniflux.app/version"
)

func info() {
	fmt.Println("Version:", version.Version)
	fmt.Println("Build Date:", version.BuildDate)
	fmt.Println("Go Version:", runtime.Version())
	fmt.Println("Compiler:", runtime.Compiler)
	fmt.Println("Arch:", runtime.GOARCH)
	fmt.Println("OS:", runtime.GOOS)
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2016 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"time"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	"github.com/unknwon/cae/zip"
	"github.com/unknwon/com"
	"github.com/urfave/cli"
)

// CmdDump represents the available dump sub-command.
var CmdDump = cli.Command{
	Name:  "dump",
	Usage: "Dump Gitea files and database",
	Description: `Dump compresses all related files and database into zip file.
It can be used for backup and capture Gitea server image to send to maintainer`,
	Action: runDump,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "file, f",
			Value: fmt.Sprintf("gitea-dump-%d.zip", time.Now().Unix()),
			Usage: "Name of the dump file which will be created.",
		},
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "Show process details",
		},
		cli.StringFlag{
			Name:  "tempdir, t",
			Value: os.TempDir(),
			Usage: "Temporary dir path",
		},
		cli.StringFlag{
			Name:  "database, d",
			Usage: "Specify the database SQL syntax",
		},
		cli.BoolFlag{
			Name:  "skip-repository, R",
			Usage: "Skip the repository dumping",
		},
	},
}

func fatal(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
	log.Fatal(format, args...)
}

func runDump(ctx *cli.Context) error {
	setting.NewContext()
	setting.NewServices() // cannot access session settings otherwise

	err := models.SetEngine()
	if err != nil {
		return err
	}

	tmpDir := ctx.String("tempdir")
	if _, err := os.Stat(tmpDir); os.IsNotExist(err) {
		fatal("Path does not exist: %s", tmpDir)
	}
	tmpWorkDir, err := ioutil.TempDir(tmpDir, "gitea-dump-")
	if err != nil {
		fatal("Failed to create tmp work directory: %v", err)
	}
	log.Info("Creating tmp work dir: %s", tmpWorkDir)

	// work-around #1103
	if os.Getenv("TMPDIR") == "" {
		os.Setenv("TMPDIR", tmpWorkDir)
	}

	dbDump := path.Join(tmpWorkDir, "gitea-db.sql")

	fileName := ctx.String("file")
	log.Info("Packing dump files...")
	z, err := zip.Create(fileName)
	if err != nil {
		fatal("Failed to create %s: %v", fileName, err)
	}

	zip.Verbose = ctx.Bool("verbose")

	if ctx.IsSet("skip-repository") {
		log.Info("Skip dumping local repositories")
	} else {
		log.Info("Dumping local repositories...%s", setting.RepoRootPath)
		reposDump := path.Join(tmpWorkDir, "gitea-repo.zip")
		if err := zip.PackTo(setting.RepoRootPath, reposDump, true); err != nil {
			fatal("Failed to dump local repositories: %v", err)
		}
		if err := z.AddFile("gitea-repo.zip", reposDump); err != nil {
			fatal("Failed to include gitea-repo.zip: %v", err)
		}
	}

	targetDBType := ctx.String("database")
	if len(targetDBType) > 0 && targetDBType != setting.Database.Type {
		log.Info("Dumping database %s => %s...", setting.Database.Type, targetDBType)
	} else {
		log.Info("Dumping database...")
	}

	if err := models.DumpDatabase(dbDump, targetDBType); err != nil {
		fatal("Failed to dump database: %v", err)
	}

	if err := z.AddFile("gitea-db.sql", dbDump); err != nil {
		fatal("Failed to include gitea-db.sql: %v", err)
	}

	if len(setting.CustomConf) > 0 {
		log.Info("Adding custom configuration file from %s", setting.CustomConf)
		if err := z.AddFile("app.ini", setting.CustomConf); err != nil {
			fatal("Failed to include specified app.ini: %v", err)
		}
	}

	customDir, err := os.Stat(setting.CustomPath)
	if err == nil && customDir.IsDir() {
		if err := z.AddDir("custom", setting.CustomPath); err != nil {
			fatal("Failed to include custom: %v", err)
		}
	} else {
		log.Info("Custom dir %s doesn't exist, skipped", setting.CustomPath)
	}

	if com.IsExist(setting.AppDataPath) {
		log.Info("Packing data directory...%s", setting.AppDataPath)

		var sessionAbsPath string
		if setting.SessionConfig.Provider == "file" {
			sessionAbsPath = setting.SessionConfig.ProviderConfig
		}
		if err := zipAddDirectoryExclude(z, "data", setting.AppDataPath, sessionAbsPath); err != nil {
			fatal("Failed to include data directory: %v", err)
		}
	}

	if err := z.AddDir("log", setting.LogRootPath); err != nil {
		fatal("Failed to include log: %v", err)
	}

	if err = z.Close(); err != nil {
		_ = os.Remove(fileName)
		fatal("Failed to save %s: %v", fileName, err)
	}

	if err := os.Chmod(fileName, 0600); err != nil {
		log.Info("Can't change file access permissions mask to 0600: %v", err)
	}

	log.Info("Removing tmp work dir: %s", tmpWorkDir)

	if err := os.RemoveAll(tmpWorkDir); err != nil {
		fatal("Failed to remove %s: %v", tmpWorkDir, err)
	}
	log.Info("Finish dumping in file %s", fileName)

	return nil
}

// zipAddDirectoryExclude zips absPath to specified zipPath inside z excluding excludeAbsPath
func zipAddDirectoryExclude(zip *zip.ZipArchive, zipPath, absPath string, excludeAbsPath string) error {
	absPath, err := filepath.Abs(absPath)
	if err != nil {
		return err
	}
	dir, err := os.Open(absPath)
	if err != nil {
		return err
	}
	defer dir.Close()

	zip.AddEmptyDir(zipPath)

	files, err := dir.Readdir(0)
	if err != nil {
		return err
	}
	for _, file := range files {
		currentAbsPath := path.Join(absPath, file.Name())
		currentZipPath := path.Join(zipPath, file.Name())
		if file.IsDir() {
			if currentAbsPath != excludeAbsPath {
				if err = zipAddDirectoryExclude(zip, currentZipPath, currentAbsPath, excludeAbsPath); err != nil {
					return err
				}
			}

		} else {
			if err = zip.AddFile(currentZipPath, currentAbsPath); err != nil {
				return err
			}
		}
	}
	return nil
}
<commit_msg>fix dump non-exist log directory (#9818)<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2016 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"time"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	"github.com/unknwon/cae/zip"
	"github.com/unknwon/com"
	"github.com/urfave/cli"
)

// CmdDump represents the available dump sub-command.
var CmdDump = cli.Command{
	Name:  "dump",
	Usage: "Dump Gitea files and database",
	Description: `Dump compresses all related files and database into zip file.
It can be used for backup and capture Gitea server image to send to maintainer`,
	Action: runDump,
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "file, f",
			Value: fmt.Sprintf("gitea-dump-%d.zip", time.Now().Unix()),
			Usage: "Name of the dump file which will be created.",
		},
		cli.BoolFlag{
			Name:  "verbose, V",
			Usage: "Show process details",
		},
		cli.StringFlag{
			Name:  "tempdir, t",
			Value: os.TempDir(),
			Usage: "Temporary dir path",
		},
		cli.StringFlag{
			Name:  "database, d",
			Usage: "Specify the database SQL syntax",
		},
		cli.BoolFlag{
			Name:  "skip-repository, R",
			Usage: "Skip the repository dumping",
		},
	},
}

func fatal(format string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", args...)
	log.Fatal(format, args...)
}

func runDump(ctx *cli.Context) error {
	setting.NewContext()
	setting.NewServices() // cannot access session settings otherwise

	err := models.SetEngine()
	if err != nil {
		return err
	}

	tmpDir := ctx.String("tempdir")
	if _, err := os.Stat(tmpDir); os.IsNotExist(err) {
		fatal("Path does not exist: %s", tmpDir)
	}
	tmpWorkDir, err := ioutil.TempDir(tmpDir, "gitea-dump-")
	if err != nil {
		fatal("Failed to create tmp work directory: %v", err)
	}
	log.Info("Creating tmp work dir: %s", tmpWorkDir)

	// work-around #1103
	if os.Getenv("TMPDIR") == "" {
		os.Setenv("TMPDIR", tmpWorkDir)
	}

	dbDump := path.Join(tmpWorkDir, "gitea-db.sql")

	fileName := ctx.String("file")
	log.Info("Packing dump files...")
	z, err := zip.Create(fileName)
	if err != nil {
		fatal("Failed to create %s: %v", fileName, err)
	}

	zip.Verbose = ctx.Bool("verbose")

	if ctx.IsSet("skip-repository") {
		log.Info("Skip dumping local repositories")
	} else {
		log.Info("Dumping local repositories...%s", setting.RepoRootPath)
		reposDump := path.Join(tmpWorkDir, "gitea-repo.zip")
		if err := zip.PackTo(setting.RepoRootPath, reposDump, true); err != nil {
			fatal("Failed to dump local repositories: %v", err)
		}
		if err := z.AddFile("gitea-repo.zip", reposDump); err != nil {
			fatal("Failed to include gitea-repo.zip: %v", err)
		}
	}

	targetDBType := ctx.String("database")
	if len(targetDBType) > 0 && targetDBType != setting.Database.Type {
		log.Info("Dumping database %s => %s...", setting.Database.Type, targetDBType)
	} else {
		log.Info("Dumping database...")
	}

	if err := models.DumpDatabase(dbDump, targetDBType); err != nil {
		fatal("Failed to dump database: %v", err)
	}

	if err := z.AddFile("gitea-db.sql", dbDump); err != nil {
		fatal("Failed to include gitea-db.sql: %v", err)
	}

	if len(setting.CustomConf) > 0 {
		log.Info("Adding custom configuration file from %s", setting.CustomConf)
		if err := z.AddFile("app.ini", setting.CustomConf); err != nil {
			fatal("Failed to include specified app.ini: %v", err)
		}
	}

	customDir, err := os.Stat(setting.CustomPath)
	if err == nil && customDir.IsDir() {
		if err := z.AddDir("custom", setting.CustomPath); err != nil {
			fatal("Failed to include custom: %v", err)
		}
	} else {
		log.Info("Custom dir %s doesn't exist, skipped", setting.CustomPath)
	}

	if com.IsExist(setting.AppDataPath) {
		log.Info("Packing data directory...%s", setting.AppDataPath)

		var sessionAbsPath string
		if setting.SessionConfig.Provider == "file" {
			sessionAbsPath = setting.SessionConfig.ProviderConfig
		}
		if err := zipAddDirectoryExclude(z, "data", setting.AppDataPath, sessionAbsPath); err != nil {
			fatal("Failed to include data directory: %v", err)
		}
	}

	if com.IsExist(setting.LogRootPath) {
		if err := z.AddDir("log", setting.LogRootPath); err != nil {
			fatal("Failed to include log: %v", err)
		}
	}

	if err = z.Close(); err != nil {
		_ = os.Remove(fileName)
		fatal("Failed to save %s: %v", fileName, err)
	}

	if err := os.Chmod(fileName, 0600); err != nil {
		log.Info("Can't change file access permissions mask to 0600: %v", err)
	}

	log.Info("Removing tmp work dir: %s", tmpWorkDir)

	if err := os.RemoveAll(tmpWorkDir); err != nil {
		fatal("Failed to remove %s: %v", tmpWorkDir, err)
	}
	log.Info("Finish dumping in file %s", fileName)

	return nil
}

// zipAddDirectoryExclude zips absPath to specified zipPath inside z excluding excludeAbsPath
func zipAddDirectoryExclude(zip *zip.ZipArchive, zipPath, absPath string, excludeAbsPath string) error {
	absPath, err := filepath.Abs(absPath)
	if err != nil {
		return err
	}
	dir, err := os.Open(absPath)
	if err != nil {
		return err
	}
	defer dir.Close()

	zip.AddEmptyDir(zipPath)

	files, err := dir.Readdir(0)
	if err != nil {
		return err
	}
	for _, file := range files {
		currentAbsPath := path.Join(absPath, file.Name())
		currentZipPath := path.Join(zipPath, file.Name())
		if file.IsDir() {
			if currentAbsPath != excludeAbsPath {
				if err = zipAddDirectoryExclude(zip, currentZipPath, currentAbsPath, excludeAbsPath); err != nil {
					return err
				}
			}

		} else {
			if err = zip.AddFile(currentZipPath, currentAbsPath); err != nil {
				return err
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
	"io/ioutil"
	"github.com/pkg/errors"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"log"
	"github.com/glassechidna/lastkeypair/common"
	"github.com/aws/aws-sdk-go/aws"
	"os"
)

var hostCmd = &cobra.Command{
	Use:   "host",
	Short: "A brief description of your command",
	Long: `A longer description that spans multiple lines and likely contains examples
and usage of using your command. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.`,
	Run: func(cmd *cobra.Command, args []string) {
		hostKeyPath, _ := cmd.PersistentFlags().GetString("host-key-path")
		signedHostKeyPath, _ := cmd.PersistentFlags().GetString("signed-host-key-path")
		caPubkeyPath, _ := cmd.PersistentFlags().GetString("cert-authority-path")
		sshdConfigPath, _ := cmd.PersistentFlags().GetString("sshd-config-path")
		authorizedPrincipalsPath, _ := cmd.PersistentFlags().GetString("authorized-principals-path")
		functionName, _ := cmd.PersistentFlags().GetString("lambda-name")
		kmsKeyId, _ := cmd.PersistentFlags().GetString("kms-key")
		funcIdentity, _ := cmd.PersistentFlags().GetString("func-identity")
		principals, _ := cmd.PersistentFlags().GetStringSlice("principal")

		err := doit(hostKeyPath, signedHostKeyPath, caPubkeyPath, sshdConfigPath, authorizedPrincipalsPath, functionName, kmsKeyId, funcIdentity, principals)
		if err != nil {
			log.Panicf("err: %s\n", err.Error())
		}
	},
}

func hostSession() (*session.Session, error) {
	sessOpts := session.Options{
		SharedConfigState: session.SharedConfigEnable,
		AssumeRoleTokenProvider: stscreds.StdinTokenProvider,
	}

	sess, err := session.NewSessionWithOptions(sessOpts)
	if err != nil {
		return nil, errors.Wrap(err, "creating aws session")
	}

	client := ec2metadata.New(sess)
	if client.Available() {
		region, err := client.Region()
		if err != nil {
			return nil, errors.Wrap(err, "getting region from ec2 metadata")
		}
		sess = sess.Copy(aws.NewConfig().WithRegion(region))
	}

	return sess, nil
}

func doit(hostKeyPath, signedHostKeyPath, caPubkeyPath, sshdConfigPath, authorizedPrincipalsPath, functionName, kmsKeyId, funcIdentity string, principals []string) error {
	hostKeyBytes, err := ioutil.ReadFile(hostKeyPath)
	if err != nil {
		return errors.Wrap(err, "reading ssh host key")
	}
	hostKey := string(hostKeyBytes)

	sess, err := hostSession()
	client := ec2metadata.New(sess)

	ident, err := common.CallerIdentityUser(sess)
	instanceArn, err := getInstanceArn(client)
	if err != nil {
		return errors.Wrap(err, "fetching instance arn from metadata service")
	}

	principals = append(principals, *instanceArn)
	token, err := hostCertToken(sess, *ident, kmsKeyId, funcIdentity, *instanceArn, principals)

	caPubkey, err := client.GetMetadata("public-keys/0/openssh-key")
	if err != nil {
		return errors.Wrap(err, "fetching ssh CA key")
	}

	response := common.HostCertRespJson{}
	err = common.RequestSignedPayload(sess, functionName, common.HostCertReqJson{
		EventType: "HostCertReq",
		Token: *token,
		PublicKey: hostKey,
	}, &response)
	if err != nil {
		return errors.Wrap(err, "requesting signed host key")
	}

	err = ioutil.WriteFile(signedHostKeyPath, []byte(response.SignedHostPublicKey), 0600)
	if err != nil {
		return errors.Wrap(err, "writing signed host key to filesystem")
	}

	err = ioutil.WriteFile(caPubkeyPath, []byte(caPubkey), 0600)
	if err != nil {
		return errors.Wrap(err, "writing ca pubkey to filesystem")
	}

	authorizedPrincipalsBytes := []byte(fmt.Sprintf("%s\n", *instanceArn))

	err = ioutil.WriteFile(authorizedPrincipalsPath, authorizedPrincipalsBytes, 0444)
	if err != nil {
		return errors.Wrap(err, "writing authorized principals to filesystem")
	}

	err = appendToFile(sshdConfigPath, fmt.Sprintf(`
HostCertificate %s
TrustedUserCAKeys %s
AuthorizedPrincipalsFile %s
`, signedHostKeyPath, caPubkeyPath, authorizedPrincipalsPath))
	if err != nil {
		return errors.Wrap(err, "appending to sshd config")
	}

	return nil
}

func getInstanceArn(client *ec2metadata.EC2Metadata) (*string, error) {
	region, err := client.Region()
	if err != nil {
		return nil, errors.Wrap(err, "getting region")
	}

	ident, err := client.GetInstanceIdentityDocument()
	if err != nil {
		return nil, errors.Wrap(err, "getting identity doc for account id and instance id")
	}

	ret := fmt.Sprintf("arn:aws:ec2:%s:%s:instance/%s", region, ident.AccountID, ident.InstanceID)
	return &ret, nil

}

func hostCertToken(sess *session.Session, ident common.StsIdentity, kmsKeyId, funcIdentity, instanceArn string, principals []string) (*common.Token, error) {
	params := common.TokenParams{
		FromId:          ident.UserId,
		FromAccount:     ident.AccountId,
		To:              funcIdentity,
		Type:            "AssumedRole",
		HostInstanceArn: instanceArn,
		Principals: principals,
	}

	ret := common.CreateToken(sess, params, kmsKeyId)
	return &ret, nil
}

func appendToFile(path, text string) error {
	f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = f.WriteString(text)
	if err != nil {
		return err
	}
	return nil
}

func init() {
	RootCmd.AddCommand(hostCmd)

	hostCmd.PersistentFlags().String("host-key-path", "/etc/ssh/ssh_host_rsa_key.pub", "")
	hostCmd.PersistentFlags().String("signed-host-key-path", "/etc/ssh/ssh_host_rsa_key-cert.pub", "")
	hostCmd.PersistentFlags().String("cert-authority-path", "/etc/ssh/cert_authority.pub", "")
	hostCmd.PersistentFlags().String("authorized-principals-path", "/etc/ssh/authorized_principals", "")
	hostCmd.PersistentFlags().String("sshd-config-path", "/etc/ssh/sshd_config", "")
	hostCmd.PersistentFlags().String("lambda-name", "LastKeypair", "")
	hostCmd.PersistentFlags().String("func-identity", "LastKeypair", "")
	hostCmd.PersistentFlags().StringSliceP("principal", "p", []string{""}, "Additional principals to request from CA")
	hostCmd.PersistentFlags().String("kms-key", "alias/LastKeypair", "ID, ARN or alias of KMS key for auth to CA")
}
<commit_msg>Host cert additional principals should be included in authorized_principals (#22)<commit_after>package cmd

import (
	"fmt"

	"github.com/spf13/cobra"
	"io/ioutil"
	"github.com/pkg/errors"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"log"
	"github.com/glassechidna/lastkeypair/common"
	"github.com/aws/aws-sdk-go/aws"
	"os"
	"strings"
)

var hostCmd = &cobra.Command{
	Use:   "host",
	Short: "A brief description of your command",
	Long: `A longer description that spans multiple lines and likely contains examples
and usage of using your command. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.`,
	Run: func(cmd *cobra.Command, args []string) {
		hostKeyPath, _ := cmd.PersistentFlags().GetString("host-key-path")
		signedHostKeyPath, _ := cmd.PersistentFlags().GetString("signed-host-key-path")
		caPubkeyPath, _ := cmd.PersistentFlags().GetString("cert-authority-path")
		sshdConfigPath, _ := cmd.PersistentFlags().GetString("sshd-config-path")
		authorizedPrincipalsPath, _ := cmd.PersistentFlags().GetString("authorized-principals-path")
		functionName, _ := cmd.PersistentFlags().GetString("lambda-name")
		kmsKeyId, _ := cmd.PersistentFlags().GetString("kms-key")
		funcIdentity, _ := cmd.PersistentFlags().GetString("func-identity")
		principals, _ := cmd.PersistentFlags().GetStringSlice("principal")

		err := doit(hostKeyPath, signedHostKeyPath, caPubkeyPath, sshdConfigPath, authorizedPrincipalsPath, functionName, kmsKeyId, funcIdentity, principals)
		if err != nil {
			log.Panicf("err: %s\n", err.Error())
		}
	},
}

func hostSession() (*session.Session, error) {
	sessOpts := session.Options{
		SharedConfigState: session.SharedConfigEnable,
		AssumeRoleTokenProvider: stscreds.StdinTokenProvider,
	}

	sess, err := session.NewSessionWithOptions(sessOpts)
	if err != nil {
		return nil, errors.Wrap(err, "creating aws session")
	}

	client := ec2metadata.New(sess)
	if client.Available() {
		region, err := client.Region()
		if err != nil {
			return nil, errors.Wrap(err, "getting region from ec2 metadata")
		}
		sess = sess.Copy(aws.NewConfig().WithRegion(region))
	}

	return sess, nil
}

func doit(hostKeyPath, signedHostKeyPath, caPubkeyPath, sshdConfigPath, authorizedPrincipalsPath, functionName, kmsKeyId, funcIdentity string, principals []string) error {
	hostKeyBytes, err := ioutil.ReadFile(hostKeyPath)
	if err != nil {
		return errors.Wrap(err, "reading ssh host key")
	}
	hostKey := string(hostKeyBytes)

	sess, err := hostSession()
	client := ec2metadata.New(sess)

	ident, err := common.CallerIdentityUser(sess)
	instanceArn, err := getInstanceArn(client)
	if err != nil {
		return errors.Wrap(err, "fetching instance arn from metadata service")
	}

	principals = append(principals, *instanceArn)
	token, err := hostCertToken(sess, *ident, kmsKeyId, funcIdentity, *instanceArn, principals)

	caPubkey, err := client.GetMetadata("public-keys/0/openssh-key")
	if err != nil {
		return errors.Wrap(err, "fetching ssh CA key")
	}

	response := common.HostCertRespJson{}
	err = common.RequestSignedPayload(sess, functionName, common.HostCertReqJson{
		EventType: "HostCertReq",
		Token: *token,
		PublicKey: hostKey,
	}, &response)
	if err != nil {
		return errors.Wrap(err, "requesting signed host key")
	}

	err = ioutil.WriteFile(signedHostKeyPath, []byte(response.SignedHostPublicKey), 0600)
	if err != nil {
		return errors.Wrap(err, "writing signed host key to filesystem")
	}

	err = ioutil.WriteFile(caPubkeyPath, []byte(caPubkey), 0600)
	if err != nil {
		return errors.Wrap(err, "writing ca pubkey to filesystem")
	}

	authorizedPrincipalsBytes := []byte(fmt.Sprintf("%s\n", strings.Join(principals, "\n")))

	err = ioutil.WriteFile(authorizedPrincipalsPath, authorizedPrincipalsBytes, 0444)
	if err != nil {
		return errors.Wrap(err, "writing authorized principals to filesystem")
	}

	err = appendToFile(sshdConfigPath, fmt.Sprintf(`
HostCertificate %s
TrustedUserCAKeys %s
AuthorizedPrincipalsFile %s
`, signedHostKeyPath, caPubkeyPath, authorizedPrincipalsPath))
	if err != nil {
		return errors.Wrap(err, "appending to sshd config")
	}

	return nil
}

func getInstanceArn(client *ec2metadata.EC2Metadata) (*string, error) {
	region, err := client.Region()
	if err != nil {
		return nil, errors.Wrap(err, "getting region")
	}

	ident, err := client.GetInstanceIdentityDocument()
	if err != nil {
		return nil, errors.Wrap(err, "getting identity doc for account id and instance id")
	}

	ret := fmt.Sprintf("arn:aws:ec2:%s:%s:instance/%s", region, ident.AccountID, ident.InstanceID)
	return &ret, nil

}

func hostCertToken(sess *session.Session, ident common.StsIdentity, kmsKeyId, funcIdentity, instanceArn string, principals []string) (*common.Token, error) {
	params := common.TokenParams{
		FromId:          ident.UserId,
		FromAccount:     ident.AccountId,
		To:              funcIdentity,
		Type:            "AssumedRole",
		HostInstanceArn: instanceArn,
		Principals: principals,
	}

	ret := common.CreateToken(sess, params, kmsKeyId)
	return &ret, nil
}

func appendToFile(path, text string) error {
	f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = f.WriteString(text)
	if err != nil {
		return err
	}
	return nil
}

func init() {
	RootCmd.AddCommand(hostCmd)

	hostCmd.PersistentFlags().String("host-key-path", "/etc/ssh/ssh_host_rsa_key.pub", "")
	hostCmd.PersistentFlags().String("signed-host-key-path", "/etc/ssh/ssh_host_rsa_key-cert.pub", "")
	hostCmd.PersistentFlags().String("cert-authority-path", "/etc/ssh/cert_authority.pub", "")
	hostCmd.PersistentFlags().String("authorized-principals-path", "/etc/ssh/authorized_principals", "")
	hostCmd.PersistentFlags().String("sshd-config-path", "/etc/ssh/sshd_config", "")
	hostCmd.PersistentFlags().String("lambda-name", "LastKeypair", "")
	hostCmd.PersistentFlags().String("func-identity", "LastKeypair", "")
	hostCmd.PersistentFlags().StringSliceP("principal", "p", []string{""}, "Additional principals to request from CA")
	hostCmd.PersistentFlags().String("kms-key", "alias/LastKeypair", "ID, ARN or alias of KMS key for auth to CA")
}
<|endoftext|>