<commit_before>package command

import (
	"fmt"
	"io"
	"math/rand"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/hashicorp/nomad/api"
)

const (
	// bytesToLines is an estimation of how many bytes are in each log line.
	// This is used to set the offset to read from when a user specifies how
	// many lines to tail from.
	bytesToLines int64 = 120

	// defaultTailLines is the number of lines to tail by default if the value
	// is not overriden.
	defaultTailLines int64 = 10
)

type FSCommand struct {
	Meta
}

func (f *FSCommand) Help() string {
	helpText := `
Usage: nomad fs <alloc-id> <path>

  fs displays either the contents of an allocation directory for the passed allocation,
  or displays the file at the given path. The path is relative to the root of the alloc
  dir and defaults to root if unspecified.

General Options:

  ` + generalOptionsUsage() + `

FS Specific Options:

  -H
    Machine friendly output.

  -verbose
    Show full information.

  -job <job-id>
    Use a random allocation from the specified job ID.

  -stat
    Show file stat information instead of displaying the file, or listing the directory.

  -f
    Causes the output to not stop when the end of the file is reached, but rather to
    wait for additional output.

  -tail
    Show the files contents with offsets relative to the end of the file. If no
    offset is given, -n is defaulted to 10.

  -n
    Sets the tail location in best-efforted number of lines relative to the end
    of the file.

  -c
    Sets the tail location in number of bytes relative to the end of the file.
`
	return strings.TrimSpace(helpText)
}

func (f *FSCommand) Synopsis() string {
	return "Inspect the contents of an allocation directory"
}

func (f *FSCommand) Run(args []string) int {
	var verbose, machine, job, stat, tail, follow bool
	var numLines, numBytes int64

	flags := f.Meta.FlagSet("fs", FlagSetClient)
	flags.Usage = func() { f.Ui.Output(f.Help()) }
	flags.BoolVar(&verbose, "verbose", false, "")
	flags.BoolVar(&machine, "H", false, "")
	flags.BoolVar(&job, "job", false, "")
	flags.BoolVar(&stat, "stat", false, "")
	flags.BoolVar(&follow, "f", false, "")
	flags.BoolVar(&tail, "tail", false, "")
	flags.Int64Var(&numLines, "n", -1, "")
	flags.Int64Var(&numBytes, "c", -1, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()

	if len(args) < 1 {
		if job {
			f.Ui.Error("job ID is required")
		} else {
			f.Ui.Error("allocation ID is required")
		}
		return 1
	}

	if len(args) > 2 {
		f.Ui.Error(f.Help())
		return 1
	}

	path := "/"
	if len(args) == 2 {
		path = args[1]
	}

	client, err := f.Meta.Client()
	if err != nil {
		f.Ui.Error(fmt.Sprintf("Error initializing client: %v", err))
		return 1
	}

	// If -job is specified, use random allocation, otherwise use provided allocation
	allocID := args[0]
	if job {
		allocID, err = getRandomJobAlloc(client, args[0])
		if err != nil {
			f.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err))
			return 1
		}
	}

	// Truncate the id unless full length is requested
	length := shortId
	if verbose {
		length = fullId
	}
	// Query the allocation info
	if len(allocID) == 1 {
		f.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters."))
		return 1
	}
	if len(allocID)%2 == 1 {
		// Identifiers must be of even length, so we strip off the last byte
		// to provide a consistent user experience.
		allocID = allocID[:len(allocID)-1]
	}

	allocs, _, err := client.Allocations().PrefixList(allocID)
	if err != nil {
		f.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err))
		return 1
	}
	if len(allocs) == 0 {
		f.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID))
		return 1
	}
	if len(allocs) > 1 {
		// Format the allocs
		out := make([]string, len(allocs)+1)
		out[0] = "ID|Eval ID|Job ID|Task Group|Desired Status|Client Status"
		for i, alloc := range allocs {
			out[i+1] = fmt.Sprintf("%s|%s|%s|%s|%s|%s",
				limit(alloc.ID, length),
				limit(alloc.EvalID, length),
				alloc.JobID,
				alloc.TaskGroup,
				alloc.DesiredStatus,
				alloc.ClientStatus,
			)
		}
		f.Ui.Output(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", formatList(out)))
		return 0
	}
	// Prefix lookup matched a single allocation
	alloc, _, err := client.Allocations().Info(allocs[0].ID, nil)
	if err != nil {
		f.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err))
		return 1
	}

	// Get file stat info
	file, _, err := client.AllocFS().Stat(alloc, path, nil)
	if err != nil {
		f.Ui.Error(err.Error())
		return 1
	}

	// If we want file stats, print those and exit.
	if stat {
		// Display the file information
		out := make([]string, 2)
		out[0] = "Mode|Size|Modified Time|Name"
		if file != nil {
			fn := file.Name
			if file.IsDir {
				fn = fmt.Sprintf("%s/", fn)
			}
			var size string
			if machine {
				size = fmt.Sprintf("%d", file.Size)
			} else {
				size = humanize.IBytes(uint64(file.Size))
			}
			out[1] = fmt.Sprintf("%s|%s|%s|%s", file.FileMode, size,
				formatTime(file.ModTime), fn)
		}
		f.Ui.Output(formatList(out))
		return 0
	}

	// Determine if the path is a file or a directory.
	if file.IsDir {
		// We have a directory, list it.
		files, _, err := client.AllocFS().List(alloc, path, nil)
		if err != nil {
			f.Ui.Error(fmt.Sprintf("Error listing alloc dir: %s", err))
			return 1
		}
		// Display the file information in a tabular format
		out := make([]string, len(files)+1)
		out[0] = "Mode|Size|Modified Time|Name"
		for i, file := range files {
			fn := file.Name
			if file.IsDir {
				fn = fmt.Sprintf("%s/", fn)
			}
			var size string
			if machine {
				size = fmt.Sprintf("%d", file.Size)
			} else {
				size = humanize.IBytes(uint64(file.Size))
			}
			out[i+1] = fmt.Sprintf("%s|%s|%s|%s",
				file.FileMode,
				size,
				formatTime(file.ModTime),
				fn,
			)
		}
		f.Ui.Output(formatList(out))
		return 0
	}

	// We have a file, output it.
	var r io.ReadCloser
	var readErr error
	if !tail {
		if follow {
			r, readErr = f.followFile(client, alloc, path, api.OriginStart, 0, -1)
		} else {
			r, readErr = client.AllocFS().Cat(alloc, path, nil)
		}

		if readErr != nil {
			readErr = fmt.Errorf("Error reading file: %v", readErr)
		}
	} else {
		// Parse the offset
		var offset int64 = defaultTailLines * bytesToLines

		if nLines, nBytes := numLines != -1, numBytes != -1; nLines && nBytes {
			f.Ui.Error("Both -n and -c are not allowed")
			return 1
		} else if numLines < -1 || numBytes < -1 {
			f.Ui.Error("Invalid size is specified")
			return 1
		} else if nLines {
			offset = numLines * bytesToLines
		} else if nBytes {
			offset = numBytes
		} else {
			numLines = defaultTailLines
		}

		if offset > file.Size {
			offset = file.Size
		}

		if follow {
			r, readErr = f.followFile(client, alloc, path, api.OriginEnd, offset, numLines)
		} else {
			// This offset needs to be relative from the front versus the follow
			// is relative to the end
			offset = file.Size - offset
			r, readErr = client.AllocFS().ReadAt(alloc, path, offset, -1, nil)

			// If numLines is set, wrap the reader
			if numLines != -1 {
				r = NewLineLimitReader(r, int(numLines), int(numLines*bytesToLines), 0)
			}
		}

		if readErr != nil {
			readErr = fmt.Errorf("Error tailing file: %v", readErr)
		}
	}

	defer r.Close()
	if readErr != nil {
		f.Ui.Error(readErr.Error())
		return 1
	}

	io.Copy(os.Stdout, r)
	return 0
}

// followFile outputs the contents of the file to stdout relative to the end of
// the file. If numLines does not equal -1, then tail -n behavior is used.
func (f *FSCommand) followFile(client *api.Client, alloc *api.Allocation,
	path, origin string, offset, numLines int64) (io.ReadCloser, error) {

	cancel := make(chan struct{})
	frames, err := client.AllocFS().Stream(alloc, path, origin, offset, cancel, nil)
	if err != nil {
		return nil, err
	}
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Create a reader
	var r io.ReadCloser
	frameReader := api.NewFrameReader(frames, cancel)
	r = frameReader

	// If numLines is set, wrap the reader
	if numLines != -1 {
		r = NewLineLimitReader(r, int(numLines), int(numLines*bytesToLines), 0)
	}

	go func() {
		<-signalCh

		// End the streaming
		r.Close()
	}()

	return r, nil
}

// Get Random Allocation ID from a known jobID. Prefer to use a running allocation,
// but use a dead allocation if no running allocations are found
func getRandomJobAlloc(client *api.Client, jobID string) (string, error) {
	var runningAllocs []*api.AllocationListStub
	allocs, _, err := client.Jobs().Allocations(jobID, nil)

	// Check that the job actually has allocations
	if len(allocs) == 0 {
		return "", fmt.Errorf("job %q doesn't exist or it has no allocations", jobID)
	}

	for _, v := range allocs {
		if v.ClientStatus == "running" {
			runningAllocs = append(runningAllocs, v)
		}
	}
	// If we don't have any allocations running, use dead allocations
	if len(runningAllocs) < 1 {
		runningAllocs = allocs
	}

	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	allocID := runningAllocs[r.Intn(len(runningAllocs))].ID
	return allocID, err
}
<commit_msg>fs uses timed unblocker like logs command<commit_after>package command

import (
	"fmt"
	"io"
	"math/rand"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/hashicorp/nomad/api"
)

const (
	// bytesToLines is an estimation of how many bytes are in each log line.
	// This is used to set the offset to read from when a user specifies how
	// many lines to tail from.
	bytesToLines int64 = 120

	// defaultTailLines is the number of lines to tail by default if the value
	// is not overriden.
	defaultTailLines int64 = 10
)

type FSCommand struct {
	Meta
}

func (f *FSCommand) Help() string {
	helpText := `
Usage: nomad fs <alloc-id> <path>

  fs displays either the contents of an allocation directory for the passed allocation,
  or displays the file at the given path. The path is relative to the root of the alloc
  dir and defaults to root if unspecified.

General Options:

  ` + generalOptionsUsage() + `

FS Specific Options:

  -H
    Machine friendly output.

  -verbose
    Show full information.

  -job <job-id>
    Use a random allocation from the specified job ID.

  -stat
    Show file stat information instead of displaying the file, or listing the directory.

  -f
    Causes the output to not stop when the end of the file is reached, but rather to
    wait for additional output.

  -tail
    Show the files contents with offsets relative to the end of the file. If no
    offset is given, -n is defaulted to 10.

  -n
    Sets the tail location in best-efforted number of lines relative to the end
    of the file.

  -c
    Sets the tail location in number of bytes relative to the end of the file.
`
	return strings.TrimSpace(helpText)
}

func (f *FSCommand) Synopsis() string {
	return "Inspect the contents of an allocation directory"
}

func (f *FSCommand) Run(args []string) int {
	var verbose, machine, job, stat, tail, follow bool
	var numLines, numBytes int64

	flags := f.Meta.FlagSet("fs", FlagSetClient)
	flags.Usage = func() { f.Ui.Output(f.Help()) }
	flags.BoolVar(&verbose, "verbose", false, "")
	flags.BoolVar(&machine, "H", false, "")
	flags.BoolVar(&job, "job", false, "")
	flags.BoolVar(&stat, "stat", false, "")
	flags.BoolVar(&follow, "f", false, "")
	flags.BoolVar(&tail, "tail", false, "")
	flags.Int64Var(&numLines, "n", -1, "")
	flags.Int64Var(&numBytes, "c", -1, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()

	if len(args) < 1 {
		if job {
			f.Ui.Error("job ID is required")
		} else {
			f.Ui.Error("allocation ID is required")
		}
		return 1
	}

	if len(args) > 2 {
		f.Ui.Error(f.Help())
		return 1
	}

	path := "/"
	if len(args) == 2 {
		path = args[1]
	}

	client, err := f.Meta.Client()
	if err != nil {
		f.Ui.Error(fmt.Sprintf("Error initializing client: %v", err))
		return 1
	}

	// If -job is specified, use random allocation, otherwise use provided allocation
	allocID := args[0]
	if job {
		allocID, err = getRandomJobAlloc(client, args[0])
		if err != nil {
			f.Ui.Error(fmt.Sprintf("Error fetching allocations: %v", err))
			return 1
		}
	}

	// Truncate the id unless full length is requested
	length := shortId
	if verbose {
		length = fullId
	}
	// Query the allocation info
	if len(allocID) == 1 {
		f.Ui.Error(fmt.Sprintf("Alloc ID must contain at least two characters."))
		return 1
	}
	if len(allocID)%2 == 1 {
		// Identifiers must be of even length, so we strip off the last byte
		// to provide a consistent user experience.
		allocID = allocID[:len(allocID)-1]
	}

	allocs, _, err := client.Allocations().PrefixList(allocID)
	if err != nil {
		f.Ui.Error(fmt.Sprintf("Error querying allocation: %v", err))
		return 1
	}
	if len(allocs) == 0 {
		f.Ui.Error(fmt.Sprintf("No allocation(s) with prefix or id %q found", allocID))
		return 1
	}
	if len(allocs) > 1 {
		// Format the allocs
		out := make([]string, len(allocs)+1)
		out[0] = "ID|Eval ID|Job ID|Task Group|Desired Status|Client Status"
		for i, alloc := range allocs {
			out[i+1] = fmt.Sprintf("%s|%s|%s|%s|%s|%s",
				limit(alloc.ID, length),
				limit(alloc.EvalID, length),
				alloc.JobID,
				alloc.TaskGroup,
				alloc.DesiredStatus,
				alloc.ClientStatus,
			)
		}
		f.Ui.Output(fmt.Sprintf("Prefix matched multiple allocations\n\n%s", formatList(out)))
		return 0
	}
	// Prefix lookup matched a single allocation
	alloc, _, err := client.Allocations().Info(allocs[0].ID, nil)
	if err != nil {
		f.Ui.Error(fmt.Sprintf("Error querying allocation: %s", err))
		return 1
	}

	// Get file stat info
	file, _, err := client.AllocFS().Stat(alloc, path, nil)
	if err != nil {
		f.Ui.Error(err.Error())
		return 1
	}

	// If we want file stats, print those and exit.
	if stat {
		// Display the file information
		out := make([]string, 2)
		out[0] = "Mode|Size|Modified Time|Name"
		if file != nil {
			fn := file.Name
			if file.IsDir {
				fn = fmt.Sprintf("%s/", fn)
			}
			var size string
			if machine {
				size = fmt.Sprintf("%d", file.Size)
			} else {
				size = humanize.IBytes(uint64(file.Size))
			}
			out[1] = fmt.Sprintf("%s|%s|%s|%s", file.FileMode, size,
				formatTime(file.ModTime), fn)
		}
		f.Ui.Output(formatList(out))
		return 0
	}

	// Determine if the path is a file or a directory.
	if file.IsDir {
		// We have a directory, list it.
		files, _, err := client.AllocFS().List(alloc, path, nil)
		if err != nil {
			f.Ui.Error(fmt.Sprintf("Error listing alloc dir: %s", err))
			return 1
		}
		// Display the file information in a tabular format
		out := make([]string, len(files)+1)
		out[0] = "Mode|Size|Modified Time|Name"
		for i, file := range files {
			fn := file.Name
			if file.IsDir {
				fn = fmt.Sprintf("%s/", fn)
			}
			var size string
			if machine {
				size = fmt.Sprintf("%d", file.Size)
			} else {
				size = humanize.IBytes(uint64(file.Size))
			}
			out[i+1] = fmt.Sprintf("%s|%s|%s|%s",
				file.FileMode,
				size,
				formatTime(file.ModTime),
				fn,
			)
		}
		f.Ui.Output(formatList(out))
		return 0
	}

	// We have a file, output it.
	var r io.ReadCloser
	var readErr error
	if !tail {
		if follow {
			r, readErr = f.followFile(client, alloc, path, api.OriginStart, 0, -1)
		} else {
			r, readErr = client.AllocFS().Cat(alloc, path, nil)
		}

		if readErr != nil {
			readErr = fmt.Errorf("Error reading file: %v", readErr)
		}
	} else {
		// Parse the offset
		var offset int64 = defaultTailLines * bytesToLines

		if nLines, nBytes := numLines != -1, numBytes != -1; nLines && nBytes {
			f.Ui.Error("Both -n and -c are not allowed")
			return 1
		} else if numLines < -1 || numBytes < -1 {
			f.Ui.Error("Invalid size is specified")
			return 1
		} else if nLines {
			offset = numLines * bytesToLines
		} else if nBytes {
			offset = numBytes
		} else {
			numLines = defaultTailLines
		}

		if offset > file.Size {
			offset = file.Size
		}

		if follow {
			r, readErr = f.followFile(client, alloc, path, api.OriginEnd, offset, numLines)
		} else {
			// This offset needs to be relative from the front versus the follow
			// is relative to the end
			offset = file.Size - offset
			r, readErr = client.AllocFS().ReadAt(alloc, path, offset, -1, nil)

			// If numLines is set, wrap the reader
			if numLines != -1 {
				r = NewLineLimitReader(r, int(numLines), int(numLines*bytesToLines), 1*time.Second)
			}
		}

		if readErr != nil {
			readErr = fmt.Errorf("Error tailing file: %v", readErr)
		}
	}

	defer r.Close()
	if readErr != nil {
		f.Ui.Error(readErr.Error())
		return 1
	}

	io.Copy(os.Stdout, r)
	return 0
}

// followFile outputs the contents of the file to stdout relative to the end of
// the file. If numLines does not equal -1, then tail -n behavior is used.
func (f *FSCommand) followFile(client *api.Client, alloc *api.Allocation,
	path, origin string, offset, numLines int64) (io.ReadCloser, error) {

	cancel := make(chan struct{})
	frames, err := client.AllocFS().Stream(alloc, path, origin, offset, cancel, nil)
	if err != nil {
		return nil, err
	}
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Create a reader
	var r io.ReadCloser
	frameReader := api.NewFrameReader(frames, cancel)
	frameReader.SetUnblockTime(500 * time.Millisecond)
	r = frameReader

	// If numLines is set, wrap the reader
	if numLines != -1 {
		r = NewLineLimitReader(r, int(numLines), int(numLines*bytesToLines), 1*time.Second)
	}

	go func() {
		<-signalCh

		// End the streaming
		r.Close()
	}()

	return r, nil
}

// Get Random Allocation ID from a known jobID. Prefer to use a running allocation,
// but use a dead allocation if no running allocations are found
func getRandomJobAlloc(client *api.Client, jobID string) (string, error) {
	var runningAllocs []*api.AllocationListStub
	allocs, _, err := client.Jobs().Allocations(jobID, nil)

	// Check that the job actually has allocations
	if len(allocs) == 0 {
		return "", fmt.Errorf("job %q doesn't exist or it has no allocations", jobID)
	}

	for _, v := range allocs {
		if v.ClientStatus == "running" {
			runningAllocs = append(runningAllocs, v)
		}
	}
	// If we don't have any allocations running, use dead allocations
	if len(runningAllocs) < 1 {
		runningAllocs = allocs
	}

	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	allocID := runningAllocs[r.Intn(len(runningAllocs))].ID
	return allocID, err
}
<|endoftext|>


<commit_before>package command

import (
	"fmt"
	"net/url"
	"strings"

	"github.com/hashicorp/nomad/api/contexts"
	"github.com/posener/complete"
	"github.com/skratchdot/open-golang/open"
)

var (
	// uiContexts is the contexts the ui can open automatically.
	uiContexts = []contexts.Context{contexts.Jobs, contexts.Allocs, contexts.Nodes}
)

type UiCommand struct {
	Meta
}

func (c *UiCommand) Help() string {
	helpText := `
Usage: nomad ui [options] <identifier>

Open the Nomad Web UI in the default browser. An optional identifier may be
provided, in which case the UI will be opened to view the details for that
object. Supported identifiers are jobs, allocations and nodes.

General Options:

  ` + generalOptionsUsage(usageOptsDefault|usageOptsNoNamespace)

	return strings.TrimSpace(helpText)
}

func (c *UiCommand) AutocompleteFlags() complete.Flags {
	return c.Meta.AutocompleteFlags(FlagSetClient)
}

func (c *UiCommand) AutocompleteArgs() complete.Predictor {
	return complete.PredictFunc(func(a complete.Args) []string {
		client, err := c.Meta.Client()
		if err != nil {
			return nil
		}

		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.All, nil)
		if err != nil {
			return []string{}
		}

		final := make([]string, 0)

		for _, allowed := range uiContexts {
			matches, ok := resp.Matches[allowed]
			if !ok {
				continue
			}
			if len(matches) == 0 {
				continue
			}

			final = append(final, matches...)
		}

		return final
	})
}

func (c *UiCommand) Synopsis() string {
	return "Open the Nomad Web UI"
}

func (c *UiCommand) Name() string { return "ui" }

func (c *UiCommand) Run(args []string) int {
	var authenticate bool

	flags := c.Meta.FlagSet(c.Name(), FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&authenticate, "authenticate", false, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Check that we got no more than one argument
	args = flags.Args()
	if l := len(args); l > 1 {
		c.Ui.Error("This command takes no or one optional argument, [<identifier>]")
		c.Ui.Error(commandErrorText(c))
		return 1
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	url, err := url.Parse(client.Address())
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing Nomad address %q: %s", client.Address(), err))
		return 1
	}

	var ottSecret string
	if authenticate {
		ott, _, err := client.ACLTokens().UpsertOneTimeToken(nil)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Could not get one-time token: %s", err))
			return 1
		}
		ottSecret = ott.OneTimeSecretID
	}

	// We were given an id so look it up
	if len(args) == 1 {
		id := args[0]

		// Query for the context associated with the id
		res, _, err := client.Search().PrefixSearch(id, contexts.All, nil)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error querying search with id: %q", err))
			return 1
		}

		if res.Matches == nil {
			c.Ui.Error(fmt.Sprintf("No matches returned for query: %q", err))
			return 1
		}

		var match contexts.Context
		var fullID string
		matchCount := 0
		for _, ctx := range uiContexts {
			vers, ok := res.Matches[ctx]
			if !ok {
				continue
			}

			if l := len(vers); l == 1 {
				match = ctx
				fullID = vers[0]
				matchCount++
			} else if l > 0 && vers[0] == id {
				// Exact match
				match = ctx
				fullID = vers[0]
				break
			}

			// Only a single result should return, as this is a match against a full id
			if matchCount > 1 || len(vers) > 1 {
				c.logMultiMatchError(id, res.Matches)
				return 1
			}
		}

		switch match {
		case contexts.Nodes:
			url.Path = fmt.Sprintf("ui/clients/%s", fullID)
		case contexts.Allocs:
			url.Path = fmt.Sprintf("ui/allocations/%s", fullID)
		case contexts.Jobs:
			url.Path = fmt.Sprintf("ui/jobs/%s", fullID)
		default:
			c.Ui.Error(fmt.Sprintf("Unable to resolve ID: %q", id))
			return 1
		}
	}

	if authenticate && ottSecret != "" {
		c.Ui.Output(fmt.Sprintf("Opening URL %q with one-time token", url.String()))
		url.RawQuery = fmt.Sprintf("ott=%s", ottSecret)
	} else {
		c.Ui.Output(fmt.Sprintf("Opening URL %q", url.String()))
	}
	if err := open.Start(url.String()); err != nil {
		c.Ui.Error(fmt.Sprintf("Error opening URL: %s", err))
		return 1
	}

	return 0
}

// logMultiMatchError is used to log an error message when multiple matches are
// found. The error message logged displays the matched IDs per context.
func (c *UiCommand) logMultiMatchError(id string, matches map[contexts.Context][]string) {
	c.Ui.Error(fmt.Sprintf("Multiple matches found for id %q", id))
	for _, ctx := range uiContexts {
		vers, ok := matches[ctx]
		if !ok {
			continue
		}
		if len(vers) == 0 {
			continue
		}

		c.Ui.Error(fmt.Sprintf("\n%s:", strings.Title(string(ctx))))
		c.Ui.Error(strings.Join(vers, ", "))
	}
}
<commit_msg>cli: add help for 'ui -authenticate' flag<commit_after>package command

import (
	"fmt"
	"net/url"
	"strings"

	"github.com/hashicorp/nomad/api/contexts"
	"github.com/posener/complete"
	"github.com/skratchdot/open-golang/open"
)

var (
	// uiContexts is the contexts the ui can open automatically.
	uiContexts = []contexts.Context{contexts.Jobs, contexts.Allocs, contexts.Nodes}
)

type UiCommand struct {
	Meta
}

func (c *UiCommand) Help() string {
	helpText := `
Usage: nomad ui [options] <identifier>

Open the Nomad Web UI in the default browser. An optional identifier may be
provided, in which case the UI will be opened to view the details for that
object. Supported identifiers are jobs, allocations and nodes.

General Options:

  ` + generalOptionsUsage(usageOptsDefault|usageOptsNoNamespace) + `

UI Options

  -authenticate: Exchange your Nomad ACL token for a one-time token in the
    web UI, if ACLs are enabled.
`

	return strings.TrimSpace(helpText)
}

func (c *UiCommand) AutocompleteFlags() complete.Flags {
	return c.Meta.AutocompleteFlags(FlagSetClient)
}

func (c *UiCommand) AutocompleteArgs() complete.Predictor {
	return complete.PredictFunc(func(a complete.Args) []string {
		client, err := c.Meta.Client()
		if err != nil {
			return nil
		}

		resp, _, err := client.Search().PrefixSearch(a.Last, contexts.All, nil)
		if err != nil {
			return []string{}
		}

		final := make([]string, 0)

		for _, allowed := range uiContexts {
			matches, ok := resp.Matches[allowed]
			if !ok {
				continue
			}
			if len(matches) == 0 {
				continue
			}

			final = append(final, matches...)
		}

		return final
	})
}

func (c *UiCommand) Synopsis() string {
	return "Open the Nomad Web UI"
}

func (c *UiCommand) Name() string { return "ui" }

func (c *UiCommand) Run(args []string) int {
	var authenticate bool

	flags := c.Meta.FlagSet(c.Name(), FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.BoolVar(&authenticate, "authenticate", false, "")

	if err := flags.Parse(args); err != nil {
		return 1
	}

	// Check that we got no more than one argument
	args = flags.Args()
	if l := len(args); l > 1 {
		c.Ui.Error("This command takes no or one optional argument, [<identifier>]")
		c.Ui.Error(commandErrorText(c))
		return 1
	}

	// Get the HTTP client
	client, err := c.Meta.Client()
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error initializing client: %s", err))
		return 1
	}

	url, err := url.Parse(client.Address())
	if err != nil {
		c.Ui.Error(fmt.Sprintf("Error parsing Nomad address %q: %s", client.Address(), err))
		return 1
	}

	var ottSecret string
	if authenticate {
		ott, _, err := client.ACLTokens().UpsertOneTimeToken(nil)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Could not get one-time token: %s", err))
			return 1
		}
		ottSecret = ott.OneTimeSecretID
	}

	// We were given an id so look it up
	if len(args) == 1 {
		id := args[0]

		// Query for the context associated with the id
		res, _, err := client.Search().PrefixSearch(id, contexts.All, nil)
		if err != nil {
			c.Ui.Error(fmt.Sprintf("Error querying search with id: %q", err))
			return 1
		}

		if res.Matches == nil {
			c.Ui.Error(fmt.Sprintf("No matches returned for query: %q", err))
			return 1
		}

		var match contexts.Context
		var fullID string
		matchCount := 0
		for _, ctx := range uiContexts {
			vers, ok := res.Matches[ctx]
			if !ok {
				continue
			}

			if l := len(vers); l == 1 {
				match = ctx
				fullID = vers[0]
				matchCount++
			} else if l > 0 && vers[0] == id {
				// Exact match
				match = ctx
				fullID = vers[0]
				break
			}

			// Only a single result should return, as this is a match against a full id
			if matchCount > 1 || len(vers) > 1 {
				c.logMultiMatchError(id, res.Matches)
				return 1
			}
		}

		switch match {
		case contexts.Nodes:
			url.Path = fmt.Sprintf("ui/clients/%s", fullID)
		case contexts.Allocs:
			url.Path = fmt.Sprintf("ui/allocations/%s", fullID)
		case contexts.Jobs:
			url.Path = fmt.Sprintf("ui/jobs/%s", fullID)
		default:
			c.Ui.Error(fmt.Sprintf("Unable to resolve ID: %q", id))
			return 1
		}
	}

	if authenticate && ottSecret != "" {
		c.Ui.Output(fmt.Sprintf("Opening URL %q with one-time token", url.String()))
		url.RawQuery = fmt.Sprintf("ott=%s", ottSecret)
	} else {
		c.Ui.Output(fmt.Sprintf("Opening URL %q", url.String()))
	}
	if err := open.Start(url.String()); err != nil {
		c.Ui.Error(fmt.Sprintf("Error opening URL: %s", err))
		return 1
	}

	return 0
}

// logMultiMatchError is used to log an error message when multiple matches are
// found. The error message logged displays the matched IDs per context.
func (c *UiCommand) logMultiMatchError(id string, matches map[contexts.Context][]string) {
	c.Ui.Error(fmt.Sprintf("Multiple matches found for id %q", id))
	for _, ctx := range uiContexts {
		vers, ok := matches[ctx]
		if !ok {
			continue
		}
		if len(vers) == 0 {
			continue
		}

		c.Ui.Error(fmt.Sprintf("\n%s:", strings.Title(string(ctx))))
		c.Ui.Error(strings.Join(vers, ", "))
	}
}
<|endoftext|>


<commit_before>package cli

import (
	"strings"

	"github.com/go-errors/errors"
	"github.com/spf13/cobra"

	"github.com/adrienkohlbecker/ejson-kms/model"
	"github.com/adrienkohlbecker/ejson-kms/utils"
)

const docExport = `
export: Export a secrets file in it's decrypted form.

Each secret in the file will be decrypted and output to standard out.
A number of formats are available:

  * bash:   SECRET='password'
  * bash-export:   export SECRET='password'
  * dotenv: SECRET="password"
  * json:   { "secret": "password" }
  * yaml:   secret: password

Please be careful when exporting your secrets, do not save them to disk!
`
const exampleExport = `
ejson-kms export
ejson-kms export --format=json
ejson-kms export --path=secrets.json --format=dotenv
`

func exportCmd() *cobra.Command {

	cmd := &cobra.Command{
		Use:     "export",
		Short:   "export the decrypted secrets",
		Long:    strings.TrimSpace(docExport),
		Example: strings.TrimSpace(exampleExport),
	}

	var (
		storePath = ".secrets.json"
		format    = "bash"
	)

	cmd.Flags().StringVar(&storePath, "path", storePath, "path of the secrets file")
	cmd.Flags().StringVar(&format, "format", format, "format of the generated output (bash|dotenv|json|yaml|bash-export|bash-ifnotset|bash-ifempty)")

	cmd.RunE = func(_ *cobra.Command, args []string) error {

		err := utils.ValidSecretsPath(storePath)
		if err != nil {
			return errors.WrapPrefix(err, "Invalid path", 0)
		}

		store, err := model.Load(storePath)
		if err != nil {
			return errors.WrapPrefix(err, "Unable to load JSON", 0)
		}

		formatter, err := utils.ValidFormatter(format)
		if err != nil {
			return errors.WrapPrefix(err, "Invalid formatter", 0)
		}

		client, err := kmsDefaultClient()
		if err != nil {
			return errors.WrapPrefix(err, "Unable to initialize AWS client", 0)
		}

		items, err := store.ExportPlaintext(client)
		if err != nil {
			return errors.WrapPrefix(err, "Unable to export items", 0)
		}

		err = formatter(cmd.OutOrStdout(), items)
		if err != nil {
			return errors.WrapPrefix(err, "Unable to export items", 0)
		}

		return nil
	}

	return cmd

}
<commit_msg>update docExport for Long Description<commit_after>package cli

import (
	"strings"

	"github.com/go-errors/errors"
	"github.com/spf13/cobra"

	"github.com/adrienkohlbecker/ejson-kms/model"
	"github.com/adrienkohlbecker/ejson-kms/utils"
)

const docExport = `
export: Export a secrets file in it's decrypted form.

Each secret in the file will be decrypted and output to standard out.
A number of formats are available:

  * bash:          SECRET='password'
  * dotenv:        SECRET="password"
  * json:          { "secret": "password" }
  * yaml:          secret: password
  * bash-export:   export SECRET='password'
  * bash-ifnotset: : ${SECRET='password'}
  * bash-ifempty:  : ${SECRET:='password'}

Please be careful when exporting your secrets, do not save them to disk!
`
const exampleExport = `
ejson-kms export
ejson-kms export --format=json
ejson-kms export --path=secrets.json --format=dotenv
`

func exportCmd() *cobra.Command {

	cmd := &cobra.Command{
		Use:     "export",
		Short:   "export the decrypted secrets",
		Long:    strings.TrimSpace(docExport),
		Example: strings.TrimSpace(exampleExport),
	}

	var (
		storePath = ".secrets.json"
		format    = "bash"
	)

	cmd.Flags().StringVar(&storePath, "path", storePath, "path of the secrets file")
	cmd.Flags().StringVar(&format, "format", format, "format of the generated output (bash|dotenv|json|yaml|bash-export|bash-ifnotset|bash-ifempty)")

	cmd.RunE = func(_ *cobra.Command, args []string) error {

		err := utils.ValidSecretsPath(storePath)
		if err != nil {
			return errors.WrapPrefix(err, "Invalid path", 0)
		}

		store, err := model.Load(storePath)
		if err != nil {
			return errors.WrapPrefix(err, "Unable to load JSON", 0)
		}

		formatter, err := utils.ValidFormatter(format)
		if err != nil {
			return errors.WrapPrefix(err, "Invalid formatter", 0)
		}

		client, err := kmsDefaultClient()
		if err != nil {
			return errors.WrapPrefix(err, "Unable to initialize AWS client", 0)
		}

		items, err := store.ExportPlaintext(client)
		if err != nil {
			return errors.WrapPrefix(err, "Unable to export items", 0)
		}

		err = formatter(cmd.OutOrStdout(), items)
		if err != nil {
			return errors.WrapPrefix(err, "Unable to export items", 0)
		}

		return nil
	}

	return cmd

}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"reflect"
	"strconv"
)

type option struct {
	field    string
	isFlag   bool
	desc     string
	short    string
	long     string
	required bool
	value    string
}

// Reflect the gathered information into the concrete action instance.
func (o *option) reflectTo(value reflect.Value) (e error) {
	if o.value == "" {
		if o.required {
			return fmt.Errorf("option %q is required but not set", o.field)
		}
		return nil
	}

	field := value.FieldByName(o.field)

	switch field.Kind() {
	case reflect.String:
		field.SetString(o.value)
	case reflect.Int:
		i, e := strconv.Atoi(o.value)
		if e != nil {
			return e
		}
		field.SetInt(int64(i))
	case reflect.Bool:
		field.SetBool(o.value == "true")
	default:
		return fmt.Errorf("invalid type %q", field.Type().String())
	}
	return nil
}

func (o *option) description() string {
	desc := "    "
	desc += o.shortDescription(" ")
	desc += fmt.Sprintf("%-*s", 30-len(desc), " ") + o.desc
	return desc
}

func (o *option) shortDescription(sep string) (desc string) {
	if o.short != "" {
		desc += "-" + o.short
	}
	if o.long != "" {
		if o.short != "" {
			desc += sep
		}
		desc += "--" + o.long
	}

	if !o.isFlag {
		desc += " <" + o.field + ">"
	}

	return desc
}

func (a *action) createOption(field reflect.StructField, value reflect.Value, tagMap map[string]string) (e error) {
	if e := validateTagMap(tagMap, "type", "desc", "short", "long", "required", "default"); e != nil {
		return fmt.Errorf("[option:%s] %s", field.Name, e.Error())
	}
	opt := &option{field: field.Name}

	if field.Type.Kind() == reflect.Bool {
		opt.isFlag = true
	}

	opt.short, e = handleShortIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.long, e = handleLongIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.required, e = handleRequired(tagMap)
	if e != nil {
		return e
	}

	if opt.required && opt.isFlag {
		return fmt.Errorf(`field %q is a flag and required, that doesn't make much sense`, field.Name)
	}

	opt.value = handlePresetValue(field, value)

	if opt.value == "" {
		opt.value, e = handleDefault(field, tagMap)
		if e != nil {
			return fmt.Errorf(`wrong value for "default" tag: %s`, e.Error())
		}
	} else { // With a preset value that value is not required any more.
		opt.required = false
	}

	opt.desc = handleDescription(tagMap)

	if opt.short == "" && opt.long == "" {
		return fmt.Errorf("option %q has neither long nor short accessor set", field.Name)
	}

	if opt.short != "" {
		if o, found := a.params[opt.short]; found {
			return fmt.Errorf("short option %q already used (option %q)", opt.short, o.field)
		}
		a.params[opt.short] = opt
	}
	if opt.long != "" {
		if o, found := a.params[opt.long]; found {
			return fmt.Errorf("long option %q already used (option %q)", opt.long, o.field)
		}
		a.params[opt.long] = opt
	}
	a.opts = append(a.opts, opt)
	return nil
}
<commit_msg>show preset values (aka defaults) in the help<commit_after>package cli

import (
	"fmt"
	"reflect"
	"strconv"
)

type option struct {
	field    string
	isFlag   bool
	desc     string
	short    string
	long     string
	required bool
	value    string
}

// Reflect the gathered information into the concrete action instance.
func (o *option) reflectTo(value reflect.Value) (e error) {
	if o.value == "" {
		if o.required {
			return fmt.Errorf("option %q is required but not set", o.field)
		}
		return nil
	}

	field := value.FieldByName(o.field)

	switch field.Kind() {
	case reflect.String:
		field.SetString(o.value)
	case reflect.Int:
		i, e := strconv.Atoi(o.value)
		if e != nil {
			return e
		}
		field.SetInt(int64(i))
	case reflect.Bool:
		field.SetBool(o.value == "true")
	default:
		return fmt.Errorf("invalid type %q", field.Type().String())
	}
	return nil
}

func (o *option) description() string {
	desc := "    "
	desc += o.shortDescription(" ")
	desc += fmt.Sprintf("%-*s", 30-len(desc), " ") + o.desc
	if o.value != "" {
		desc += " (default: " + o.value + ")"
	}
	return desc
}

func (o *option) shortDescription(sep string) (desc string) {
	if o.short != "" {
		desc += "-" + o.short
	}
	if o.long != "" {
		if o.short != "" {
			desc += sep
		}
		desc += "--" + o.long
	}

	if !o.isFlag {
		desc += " <" + o.field + ">"
	}

	return desc
}

func (a *action) createOption(field reflect.StructField, value reflect.Value, tagMap map[string]string) (e error) {
	if e := validateTagMap(tagMap, "type", "desc", "short", "long", "required", "default"); e != nil {
		return fmt.Errorf("[option:%s] %s", field.Name, e.Error())
	}
	opt := &option{field: field.Name}

	if field.Type.Kind() == reflect.Bool {
		opt.isFlag = true
	}

	opt.short, e = handleShortIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.long, e = handleLongIdentifier(tagMap)
	if e != nil {
		return e
	}

	opt.required, e = handleRequired(tagMap)
	if e != nil {
		return e
	}

	if opt.required && opt.isFlag {
		return fmt.Errorf(`field %q is a flag and required, that doesn't make much sense`, field.Name)
	}

	opt.value = handlePresetValue(field, value)

	if opt.value == "" {
		opt.value, e = handleDefault(field, tagMap)
		if e != nil {
			return fmt.Errorf(`wrong value for "default" tag: %s`, e.Error())
		}
	} else { // With a preset value that value is not required any more.
		opt.required = false
	}

	opt.desc = handleDescription(tagMap)

	if opt.short == "" && opt.long == "" {
		return fmt.Errorf("option %q has neither long nor short accessor set", field.Name)
	}

	if opt.short != "" {
		if o, found := a.params[opt.short]; found {
			return fmt.Errorf("short option %q already used (option %q)", opt.short, o.field)
		}
		a.params[opt.short] = opt
	}
	if opt.long != "" {
		if o, found := a.params[opt.long]; found {
			return fmt.Errorf("long option %q already used (option %q)", opt.long, o.field)
		}
		a.params[opt.long] = opt
	}
	a.opts = append(a.opts, opt)
	return nil
}
<|endoftext|>


<commit_before>/* Gorilla Sessions backend for MySQL.

Copyright (c) 2013 Contributors. See the list of contributors in the CONTRIBUTORS file for details.

This software is licensed under a MIT style license available in the LICENSE file.
*/
package mysqlstore

import (
	"database/sql"
	"encoding/gob"
	"errors"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
	"log"
	"net/http"
	"strings"
	"time"
)

type MySQLStore struct {
	db         *sql.DB
	stmtInsert *sql.Stmt
	stmtDelete *sql.Stmt
	stmtUpdate *sql.Stmt
	stmtSelect *sql.Stmt

	Codecs  []securecookie.Codec
	Options *sessions.Options
	table   string
}

type sessionRow struct {
	id         string
	data       string
	createdOn  time.Time
	modifiedOn time.Time
	expiresOn  time.Time
}

func init() {
	gob.Register(time.Time{})
}

func NewMySQLStore(endpoint string, tableName string, path string, maxAge int, keyPairs ...[]byte) (*MySQLStore, error) {
	db, err := sql.Open("mysql", endpoint)
	if err != nil {
		return nil, err
	}

	return NewMySQLStoreFromConnection(db, tableName, path, maxAge, keyPairs...)
}

func NewMySQLStoreFromConnection(db *sql.DB, tableName string, path string, maxAge int, keyPairs ...[]byte) (*MySQLStore, error) {
	tableNameOrig := tableName
	// Make sure table name is enclosed.
	tableName = "`" + strings.Trim(tableName, "`") + "`"

	var count int
	err := db.QueryRow("SELECT COUNT(1) as count FROM information_schema.tables WHERE table_name = ?", tableNameOrig).Scan(&count)
	if err != nil {
		return nil, err
	}
	if count == 0 {
		cTableQ := "CREATE TABLE IF NOT EXISTS " +
			tableName + " (id INT NOT NULL AUTO_INCREMENT, " +
			"session_data LONGBLOB, " +
			"created_on TIMESTAMP DEFAULT 0, " +
			"modified_on TIMESTAMP NOT NULL ON UPDATE CURRENT_TIMESTAMP, " +
			"expires_on TIMESTAMP DEFAULT 0, PRIMARY KEY(`id`)) ENGINE=InnoDB;"
		if _, err := db.Exec(cTableQ); err != nil {
			return nil, err
		}
	}

	insQ := "INSERT INTO " + tableName +
		"(id, session_data, created_on, modified_on, expires_on) VALUES (NULL, ?, ?, ?, ?)"
	stmtInsert, stmtErr := db.Prepare(insQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	delQ := "DELETE FROM " + tableName + " WHERE id = ?"
	stmtDelete, stmtErr := db.Prepare(delQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	updQ := "UPDATE " + tableName + " SET session_data = ?, created_on = ?, expires_on = ? " +
		"WHERE id = ?"
	stmtUpdate, stmtErr := db.Prepare(updQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	selQ := "SELECT id, session_data, created_on, modified_on, expires_on from " +
		tableName + " WHERE id = ?"
	stmtSelect, stmtErr := db.Prepare(selQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	return &MySQLStore{
		db:         db,
		stmtInsert: stmtInsert,
		stmtDelete: stmtDelete,
		stmtUpdate: stmtUpdate,
		stmtSelect: stmtSelect,
		Codecs:     securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   path,
			MaxAge: maxAge,
		},
		table: tableName,
	}, nil
}

func (m *MySQLStore) Close() {
	m.stmtSelect.Close()
	m.stmtUpdate.Close()
	m.stmtDelete.Close()
	m.stmtInsert.Close()
	m.db.Close()
}

func (m *MySQLStore) Get(r *http.Request, name string) (*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(m, name)
}

func (m *MySQLStore) New(r *http.Request, name string) (*sessions.Session, error) {
	session := sessions.NewSession(m, name)
	session.Options = &sessions.Options{
		Path:   m.Options.Path,
		MaxAge: m.Options.MaxAge,
	}
	session.IsNew = true
	var err error
	if cook, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cook.Value, &session.ID, m.Codecs...)
		if err == nil {
			err = m.load(session)
			if err == nil {
				session.IsNew = false
			} else {
				err = nil
			}
		}
	}
	return session, err
}

func (m *MySQLStore) Save(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
	var err error
	if session.ID == "" {
		if err = m.insert(session); err != nil {
			return err
		}
	} else if err = m.save(session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID, m.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded, session.Options))
	return nil
}

func (m *MySQLStore) insert(session *sessions.Session) error {
	var createdOn time.Time
	var modifiedOn time.Time
	var expiresOn time.Time
	crOn := session.Values["created_on"]
	if crOn == nil {
		createdOn = time.Now()
	} else {
		createdOn = crOn.(time.Time)
	}
	modifiedOn = createdOn
	exOn := session.Values["expires_on"]
	if exOn == nil {
		expiresOn = time.Now().Add(time.Second * time.Duration(session.Options.MaxAge))
	} else {
		expiresOn = exOn.(time.Time)
	}
	delete(session.Values, "created_on")
	delete(session.Values, "expires_on")
	delete(session.Values, "modified_on")

	encoded, encErr := securecookie.EncodeMulti(session.Name(), session.Values, m.Codecs...)
	if encErr != nil {
		return encErr
	}
	res, insErr := m.stmtInsert.Exec(encoded, createdOn, modifiedOn, expiresOn)
	if insErr != nil {
		return insErr
	}
	lastInserted, lInsErr := res.LastInsertId()
	if lInsErr != nil {
		return lInsErr
	}
	session.ID = fmt.Sprintf("%d", lastInserted)
	return nil
}

func (m *MySQLStore) Delete(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {

	// Set cookie to expire.
	options := *session.Options
	options.MaxAge = -1
	http.SetCookie(w, sessions.NewCookie(session.Name(), "", &options))
	// Clear session values.
	for k := range session.Values {
		delete(session.Values, k)
	}

	_, delErr := m.stmtDelete.Exec(session.ID)
	if delErr != nil {
		return delErr
	}
	return nil
}

func (m *MySQLStore) save(session *sessions.Session) error {
	if session.IsNew == true {
		return m.insert(session)
	}
	var createdOn time.Time
	var expiresOn time.Time
	crOn := session.Values["created_on"]
	if crOn == nil {
		createdOn = time.Now()
	} else {
		createdOn = crOn.(time.Time)
	}

	exOn := session.Values["expires_on"]
	if exOn == nil {
		expiresOn = time.Now().Add(time.Second * time.Duration(session.Options.MaxAge))
		log.Print("nil")
	} else {
		expiresOn = exOn.(time.Time)
		if expiresOn.Sub(time.Now().Add(time.Second*time.Duration(session.Options.MaxAge))) < 0 {
			expiresOn = time.Now().Add(time.Second * time.Duration(session.Options.MaxAge))
		}
	}

	delete(session.Values, "created_on")
	delete(session.Values, "expires_on")
	delete(session.Values, "modified_on")
	encoded, encErr := securecookie.EncodeMulti(session.Name(), session.Values, m.Codecs...)
	if encErr != nil {
		return encErr
	}
	_, updErr := m.stmtUpdate.Exec(encoded, createdOn, expiresOn, session.ID)
	if updErr != nil {
		return updErr
	}
	return nil
}

func (m *MySQLStore) load(session *sessions.Session) error {
	row := m.stmtSelect.QueryRow(session.ID)
	sess := sessionRow{}
	scanErr := row.Scan(&sess.id, &sess.data, &sess.createdOn, &sess.modifiedOn, &sess.expiresOn)
	if scanErr != nil {
		return scanErr
	}
	if sess.expiresOn.Sub(time.Now()) < 0 {
		log.Printf("Session expired on %s, but it is %s now.", sess.expiresOn, time.Now())
		return errors.New("Session expired")
	}
	err := securecookie.DecodeMulti(session.Name(), sess.data, &session.Values, m.Codecs...)
	if err != nil {
		return err
	}
	session.Values["created_on"] = sess.createdOn
	session.Values["modified_on"] = sess.modifiedOn
	session.Values["expires_on"] = sess.expiresOn
	return nil

}
<commit_msg>Revert "check for table before creation"<commit_after>/* Gorilla Sessions backend for MySQL.

Copyright (c) 2013 Contributors. See the list of contributors in the CONTRIBUTORS file for details.

This software is licensed under a MIT style license available in the LICENSE file.
*/
package mysqlstore

import (
	"database/sql"
	"encoding/gob"
	"errors"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/securecookie"
	"github.com/gorilla/sessions"
	"log"
	"net/http"
	"strings"
	"time"
)

type MySQLStore struct {
	db         *sql.DB
	stmtInsert *sql.Stmt
	stmtDelete *sql.Stmt
	stmtUpdate *sql.Stmt
	stmtSelect *sql.Stmt

	Codecs  []securecookie.Codec
	Options *sessions.Options
	table   string
}

type sessionRow struct {
	id         string
	data       string
	createdOn  time.Time
	modifiedOn time.Time
	expiresOn  time.Time
}

func init() {
	gob.Register(time.Time{})
}

func NewMySQLStore(endpoint string, tableName string, path string, maxAge int, keyPairs ...[]byte) (*MySQLStore, error) {
	db, err := sql.Open("mysql", endpoint)
	if err != nil {
		return nil, err
	}

	return NewMySQLStoreFromConnection(db, tableName, path, maxAge, keyPairs...)
}

func NewMySQLStoreFromConnection(db *sql.DB, tableName string, path string, maxAge int, keyPairs ...[]byte) (*MySQLStore, error) {
	// Make sure table name is enclosed.
	tableName = "`" + strings.Trim(tableName, "`") + "`"

	cTableQ := "CREATE TABLE IF NOT EXISTS " +
		tableName + " (id INT NOT NULL AUTO_INCREMENT, " +
		"session_data LONGBLOB, " +
		"created_on TIMESTAMP DEFAULT 0, " +
		"modified_on TIMESTAMP NOT NULL ON UPDATE CURRENT_TIMESTAMP, " +
		"expires_on TIMESTAMP DEFAULT 0, PRIMARY KEY(`id`)) ENGINE=InnoDB;"
	if _, err := db.Exec(cTableQ); err != nil {
		return nil, err
	}

	insQ := "INSERT INTO " + tableName +
		"(id, session_data, created_on, modified_on, expires_on) VALUES (NULL, ?, ?, ?, ?)"
	stmtInsert, stmtErr := db.Prepare(insQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	delQ := "DELETE FROM " + tableName + " WHERE id = ?"
	stmtDelete, stmtErr := db.Prepare(delQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	updQ := "UPDATE " + tableName + " SET session_data = ?, created_on = ?, expires_on = ? " +
		"WHERE id = ?"
	stmtUpdate, stmtErr := db.Prepare(updQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	selQ := "SELECT id, session_data, created_on, modified_on, expires_on from " +
		tableName + " WHERE id = ?"
	stmtSelect, stmtErr := db.Prepare(selQ)
	if stmtErr != nil {
		return nil, stmtErr
	}

	return &MySQLStore{
		db:         db,
		stmtInsert: stmtInsert,
		stmtDelete: stmtDelete,
		stmtUpdate: stmtUpdate,
		stmtSelect: stmtSelect,
		Codecs:     securecookie.CodecsFromPairs(keyPairs...),
		Options: &sessions.Options{
			Path:   path,
			MaxAge: maxAge,
		},
		table: tableName,
	}, nil
}

func (m *MySQLStore) Close() {
	m.stmtSelect.Close()
	m.stmtUpdate.Close()
	m.stmtDelete.Close()
	m.stmtInsert.Close()
	m.db.Close()
}

func (m *MySQLStore) Get(r *http.Request, name string) (*sessions.Session, error) {
	return sessions.GetRegistry(r).Get(m, name)
}

func (m *MySQLStore) New(r *http.Request, name string) (*sessions.Session, error) {
	session := sessions.NewSession(m, name)
	session.Options = &sessions.Options{
		Path:   m.Options.Path,
		MaxAge: m.Options.MaxAge,
	}
	session.IsNew = true
	var err error
	if cook, errCookie := r.Cookie(name); errCookie == nil {
		err = securecookie.DecodeMulti(name, cook.Value, &session.ID, m.Codecs...)
		if err == nil {
			err = m.load(session)
			if err == nil {
				session.IsNew = false
			} else {
				err = nil
			}
		}
	}
	return session, err
}

func (m *MySQLStore) Save(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {
	var err error
	if session.ID == "" {
		if err = m.insert(session); err != nil {
			return err
		}
	} else if err = m.save(session); err != nil {
		return err
	}
	encoded, err := securecookie.EncodeMulti(session.Name(), session.ID, m.Codecs...)
	if err != nil {
		return err
	}
	http.SetCookie(w, sessions.NewCookie(session.Name(), encoded, session.Options))
	return nil
}

func (m *MySQLStore) insert(session *sessions.Session) error {
	var createdOn time.Time
	var modifiedOn time.Time
	var expiresOn time.Time
	crOn := session.Values["created_on"]
	if crOn == nil {
		createdOn = time.Now()
	} else {
		createdOn = crOn.(time.Time)
	}
	modifiedOn = createdOn
	exOn := session.Values["expires_on"]
	if exOn == nil {
		expiresOn = time.Now().Add(time.Second * time.Duration(session.Options.MaxAge))
	} else {
		expiresOn = exOn.(time.Time)
	}
	delete(session.Values, "created_on")
	delete(session.Values, "expires_on")
	delete(session.Values, "modified_on")

	encoded, encErr := securecookie.EncodeMulti(session.Name(), session.Values, m.Codecs...)
	if encErr != nil {
		return encErr
	}
	res, insErr := m.stmtInsert.Exec(encoded, createdOn, modifiedOn, expiresOn)
	if insErr != nil {
		return insErr
	}
	lastInserted, lInsErr := res.LastInsertId()
	if lInsErr != nil {
		return lInsErr
	}
	session.ID = fmt.Sprintf("%d", lastInserted)
	return nil
}

func (m *MySQLStore) Delete(r *http.Request, w http.ResponseWriter, session *sessions.Session) error {

	// Set cookie to expire.
	options := *session.Options
	options.MaxAge = -1
	http.SetCookie(w, sessions.NewCookie(session.Name(), "", &options))
	// Clear session values.
	for k := range session.Values {
		delete(session.Values, k)
	}

	_, delErr := m.stmtDelete.Exec(session.ID)
	if delErr != nil {
		return delErr
	}
	return nil
}

func (m *MySQLStore) save(session *sessions.Session) error {
	if session.IsNew == true {
		return m.insert(session)
	}
	var createdOn time.Time
	var expiresOn time.Time
	crOn := session.Values["created_on"]
	if crOn == nil {
		createdOn = time.Now()
	} else {
		createdOn = crOn.(time.Time)
	}

	exOn := session.Values["expires_on"]
	if exOn == nil {
		expiresOn = time.Now().Add(time.Second * time.Duration(session.Options.MaxAge))
		log.Print("nil")
	} else {
		expiresOn = exOn.(time.Time)
		if expiresOn.Sub(time.Now().Add(time.Second*time.Duration(session.Options.MaxAge))) < 0 {
			expiresOn = time.Now().Add(time.Second * time.Duration(session.Options.MaxAge))
		}
	}

	delete(session.Values, "created_on")
	delete(session.Values, "expires_on")
	delete(session.Values, "modified_on")
	encoded, encErr := securecookie.EncodeMulti(session.Name(), session.Values, m.Codecs...)
	if encErr != nil {
		return encErr
	}
	_, updErr := m.stmtUpdate.Exec(encoded, createdOn, expiresOn, session.ID)
	if updErr != nil {
		return updErr
	}
	return nil
}

func (m *MySQLStore) load(session *sessions.Session) error {
	row := m.stmtSelect.QueryRow(session.ID)
	sess := sessionRow{}
	scanErr := row.Scan(&sess.id, &sess.data, &sess.createdOn, &sess.modifiedOn, &sess.expiresOn)
	if scanErr != nil {
		return scanErr
	}
	if sess.expiresOn.Sub(time.Now()) < 0 {
		log.Printf("Session expired on %s, but it is %s now.", sess.expiresOn, time.Now())
		return errors.New("Session expired")
	}
	err := securecookie.DecodeMulti(session.Name(), sess.data, &session.Values, m.Codecs...)
	if err != nil {
		return err
	}
	session.Values["created_on"] = sess.createdOn
	session.Values["modified_on"] = sess.modifiedOn
	session.Values["expires_on"] = sess.expiresOn
	return nil

}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/nikepan/go-datastructures/queue"
)

// ClickhouseServer - clickhouse server instance object struct
type ClickhouseServer struct {
	URL         string
	LastRequest time.Time
	Bad         bool
	Client      *http.Client
}

// Clickhouse - main clickhouse sender object
type Clickhouse struct {
	Servers        []*ClickhouseServer
	Queue          *queue.Queue
	mu             sync.Mutex
	DownTimeout    int
	ConnectTimeout int
	Dumper         Dumper
	wg             sync.WaitGroup
}

// ClickhouseRequest - request struct for queue
type ClickhouseRequest struct {
	Params  string
	Query   string
	Content string
	Count   int
}

// ErrServerIsDown - signals about server is down
var ErrServerIsDown = errors.New("server is down")

// ErrNoServers - signals about no working servers
var ErrNoServers = errors.New("No working clickhouse servers")

// NewClickhouse - get clickhouse object
func NewClickhouse(downTimeout int, connectTimeout int) (c *Clickhouse) {
	c = new(Clickhouse)
	c.DownTimeout = downTimeout
	c.ConnectTimeout = connectTimeout
	if c.ConnectTimeout > 0 {
		c.ConnectTimeout = 10
	}
	c.Servers = make([]*ClickhouseServer, 0)
	c.Queue = queue.New(1000)
	go c.Run()
	return c
}

// AddServer - add clickhouse server url
func (c *Clickhouse) AddServer(url string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.Servers = append(c.Servers, &ClickhouseServer{URL: url, Client: &http.Client{
		Timeout: time.Second * time.Duration(c.ConnectTimeout),
	}})
}

// DumpServers - dump servers state to prometheus
func (c *Clickhouse) DumpServers() {
	c.mu.Lock()
	defer c.mu.Unlock()
	good := 0
	bad := 0
	for _, s := range c.Servers {
		if s.Bad {
			bad++
		} else {
			good++
		}
	}
	goodServers.Set(float64(good))
	badServers.Set(float64(bad))
}

// GetNextServer - getting next server for request
func (c *Clickhouse) GetNextServer() (srv *ClickhouseServer) {
	c.mu.Lock()
	defer c.mu.Unlock()
	tnow := time.Now()
	for _, s := range c.Servers {
		if s.Bad {
			if tnow.Sub(s.LastRequest) > time.Second*time.Duration(c.DownTimeout) {
				s.Bad = false
			} else {
				continue
			}
		}
		if srv != nil {
			if srv.LastRequest.Sub(s.LastRequest) > 0 {
				srv = s
			}
		} else {
			srv = s
		}
	}
	if srv != nil {
		srv.LastRequest = time.Now()
	}
	return srv

}

// Send - send request to next server
func (c *Clickhouse) Send(r *ClickhouseRequest) {
	c.wg.Add(1)
	c.Queue.Put(r)
}

// Dump - save query to file
func (c *Clickhouse) Dump(params string, content string, response string, prefix string, status int) error {
	dumpCounter.Inc()
	if c.Dumper != nil {
		c.mu.Lock()
		defer c.mu.Unlock()
		return c.Dumper.Dump(params, content, response, prefix, status)
	}
	return nil
}

// Len - returns queries queue length
func (c *Clickhouse) Len() int64 {
	return c.Queue.Len()
}

// Empty - check if queue is empty
func (c *Clickhouse) Empty() bool {
	return c.Queue.Empty()
}

// Run server
func (c *Clickhouse) Run() {
	var err error
	var datas []interface{}
	for {
		datas, err = c.Queue.Poll(1, time.Second*5)
		if err == nil {
			data := datas[0].(*ClickhouseRequest)
			resp, status, err := c.SendQuery(data)
			if err != nil {
				log.Printf("ERROR: Send (%+v) %+v; response %+v\n", status, err, resp)
				prefix := "1"
				if status >= 400 && status < 502 {
					prefix = "2"
				}
				c.Dump(data.Params, data.Content, resp, prefix, status)
			} else {
				sentCounter.Inc()
			}
			c.DumpServers()
			c.wg.Done()
		}
	}
}

// WaitFlush - wait for flush ends
func (c *Clickhouse) WaitFlush() (err error) {
	c.wg.Wait()
	return nil
}

// SendQuery - sends query to server and return result
func (srv *ClickhouseServer) SendQuery(r *ClickhouseRequest) (response string, status int, err error) {
	if srv.URL != "" {
		url := srv.URL
		if r.Params != "" {
			url += "?" + r.Params
		}
		log.Printf("INFO: send %+v rows to %+v of %+v\n", r.Count, url, r.Query)
		resp, err := srv.Client.Post(url, "", strings.NewReader(r.Content))
		if err != nil {
			srv.Bad = true
			return err.Error(), http.StatusBadGateway, ErrServerIsDown
		}
		buf, _ := ioutil.ReadAll(resp.Body)
		s := string(buf)
		if resp.StatusCode >= 502 {
			srv.Bad = true
			err = ErrServerIsDown
		} else if resp.StatusCode >= 400 {
			err = fmt.Errorf("Wrong server status %+v:\nresponse: %+v\nrequest: %#v", resp.StatusCode, s, r.Content)
		}
		return s, resp.StatusCode, err
	}

	return "", http.StatusOK, err
}

// SendQuery - sends query to server and return result (with server cycle)
func (c *Clickhouse) SendQuery(r *ClickhouseRequest) (response string, status int, err error) {
	for {
		s := c.GetNextServer()
		if s != nil {
			response, status, err = s.SendQuery(r)
			if errors.Is(err, ErrServerIsDown) {
				log.Printf("ERROR: server down (%+v): %+v\n", status, response)
				continue
			}
			return response, status, err
		}
		return response, status, ErrNoServers
	}
}
<commit_msg>remove password from log<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/nikepan/go-datastructures/queue"
)

// ClickhouseServer - clickhouse server instance object struct
type ClickhouseServer struct {
	URL         string
	LastRequest time.Time
	Bad         bool
	Client      *http.Client
}

// Clickhouse - main clickhouse sender object
type Clickhouse struct {
	Servers        []*ClickhouseServer
	Queue          *queue.Queue
	mu             sync.Mutex
	DownTimeout    int
	ConnectTimeout int
	Dumper         Dumper
	wg             sync.WaitGroup
}

// ClickhouseRequest - request struct for queue
type ClickhouseRequest struct {
	Params  string
	Query   string
	Content string
	Count   int
}

// ErrServerIsDown - signals about server is down
var ErrServerIsDown = errors.New("server is down")

// ErrNoServers - signals about no working servers
var ErrNoServers = errors.New("No working clickhouse servers")

// NewClickhouse - get clickhouse object
func NewClickhouse(downTimeout int, connectTimeout int) (c *Clickhouse) {
	c = new(Clickhouse)
	c.DownTimeout = downTimeout
	c.ConnectTimeout = connectTimeout
	if c.ConnectTimeout > 0 {
		c.ConnectTimeout = 10
	}
	c.Servers = make([]*ClickhouseServer, 0)
	c.Queue = queue.New(1000)
	go c.Run()
	return c
}

// AddServer - add clickhouse server url
func (c *Clickhouse) AddServer(url string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.Servers = append(c.Servers, &ClickhouseServer{URL: url, Client: &http.Client{
		Timeout: time.Second * time.Duration(c.ConnectTimeout),
	}})
}

// DumpServers - dump servers state to prometheus
func (c *Clickhouse) DumpServers() {
	c.mu.Lock()
	defer c.mu.Unlock()
	good := 0
	bad := 0
	for _, s := range c.Servers {
		if s.Bad {
			bad++
		} else {
			good++
		}
	}
	goodServers.Set(float64(good))
	badServers.Set(float64(bad))
}

// GetNextServer - getting next server for request
func (c *Clickhouse) GetNextServer() (srv *ClickhouseServer) {
	c.mu.Lock()
	defer c.mu.Unlock()
	tnow := time.Now()
	for _, s := range c.Servers {
		if s.Bad {
			if tnow.Sub(s.LastRequest) > time.Second*time.Duration(c.DownTimeout) {
				s.Bad = false
			} else {
				continue
			}
		}
		if srv != nil {
			if srv.LastRequest.Sub(s.LastRequest) > 0 {
				srv = s
			}
		} else {
			srv = s
		}
	}
	if srv != nil {
		srv.LastRequest = time.Now()
	}
	return srv

}

// Send - send request to next server
func (c *Clickhouse) Send(r *ClickhouseRequest) {
	c.wg.Add(1)
	c.Queue.Put(r)
}

// Dump - save query to file
func (c *Clickhouse) Dump(params string, content string, response string, prefix string, status int) error {
	dumpCounter.Inc()
	if c.Dumper != nil {
		c.mu.Lock()
		defer c.mu.Unlock()
		return c.Dumper.Dump(params, content, response, prefix, status)
	}
	return nil
}

// Len - returns queries queue length
func (c *Clickhouse) Len() int64 {
	return c.Queue.Len()
}

// Empty - check if queue is empty
func (c *Clickhouse) Empty() bool {
	return c.Queue.Empty()
}

// Run server
func (c *Clickhouse) Run() {
	var err error
	var datas []interface{}
	for {
		datas, err = c.Queue.Poll(1, time.Second*5)
		if err == nil {
			data := datas[0].(*ClickhouseRequest)
			resp, status, err := c.SendQuery(data)
			if err != nil {
				log.Printf("ERROR: Send (%+v) %+v; response %+v\n", status, err, resp)
				prefix := "1"
				if status >= 400 && status < 502 {
					prefix = "2"
				}
				c.Dump(data.Params, data.Content, resp, prefix, status)
			} else {
				sentCounter.Inc()
			}
			c.DumpServers()
			c.wg.Done()
		}
	}
}

// WaitFlush - wait for flush ends
func (c *Clickhouse) WaitFlush() (err error) {
	c.wg.Wait()
	return nil
}

// SendQuery - sends query to server and return result
func (srv *ClickhouseServer) SendQuery(r *ClickhouseRequest) (response string, status int, err error) {
	if srv.URL != "" {
		url := srv.URL
		if r.Params != "" {
			url += "?" + r.Params
		}
		log.Printf("INFO: send %+v rows to %+v of %+v\n", r.Count, srv.URL, r.Query)
		resp, err := srv.Client.Post(url, "", strings.NewReader(r.Content))
		if err != nil {
			srv.Bad = true
			return err.Error(), http.StatusBadGateway, ErrServerIsDown
		}
		buf, _ := ioutil.ReadAll(resp.Body)
		s := string(buf)
		if resp.StatusCode >= 502 {
			srv.Bad = true
			err = ErrServerIsDown
		} else if resp.StatusCode >= 400 {
			err = fmt.Errorf("Wrong server status %+v:\nresponse: %+v\nrequest: %#v", resp.StatusCode, s, r.Content)
		}
		return s, resp.StatusCode, err
	}

	return "", http.StatusOK, err
}

// SendQuery - sends query to server and return result (with server cycle)
func (c *Clickhouse) SendQuery(r *ClickhouseRequest) (response string, status int, err error) {
	for {
		s := c.GetNextServer()
		if s != nil {
			response, status, err = s.SendQuery(r)
			if errors.Is(err, ErrServerIsDown) {
				log.Printf("ERROR: server down (%+v): %+v\n", status, response)
				continue
			}
			return response, status, err
		}
		return response, status, ErrNoServers
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 The Cockroach 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. See the AUTHORS file
// for names of contributors.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

/*
Package client provides clients for accessing the various
externally-facing Cockroach database endpoints.

KV Client

The KV client is a fully-featured client of Cockroach's key-value
database. It provides a simple, synchronous interface well-suited to
parallel updates and queries.

The simplest way to use the client is through the Call method. Call
synchronously invokes the method and returns the reply and an
error. The example below shows a get and a put.

  kv := client.NewKV(client.NewHTTPSender("localhost:8080", tlsConfig))

  getResp := &proto.GetResponse{}
  if err := kv.Call(proto.Get, proto.GetArgs(proto.Key("a")), getResp); err != nil {
    log.Fatal(err)
  }
  putResp := &proto.PutResponse{}
  if _, err := kv.Call(proto.Put, proto.PutArgs(proto.Key("b"), getResp.Value.Bytes), putResp) err != nil {
    log.Fatal(err)
  }

The API is synchronous, but accommodates efficient parallel updates
and queries using the Prepare method. An arbitrary number of Prepare
invocations are followed up with a call to Flush. Until the Flush,
requests are buffered locally in anticipation of being sent to
Cockroach as part of a batch. The Flush batches prepared calls and
sends them together. Note however that API calls which are buffered
and sent together are not guaranteed to have atomic semantics. A
transaction must be used to guarantee atomicity. A simple example of
using the API which does two scans in parallel and then sends a
sequence of puts in parallel:

  kv := client.NewKV(client.NewHTTPSender("localhost:8080", tlsConfig))

  acResp, xzResp := &proto.ScanResponse{}, &proto.ScanResponse{}
  kv.Prepare(proto.Scan, proto.ScanArgs(proto.Key("a"), proto.Key("c").Next()), acResp)
  kv.Prepare(proto.Scan, proto.ScanArgs(proto.Key("x"), proto.Key("z").Next()), xzResp)

  // Flush sends both scans in parallel and returns first error or nil.
  if err := kv.Flush(); err != nil {
    log.Fatal(err)
  }

  // Append maximum value from "a"-"c" to all values from "x"-"z".
  max := []byte(nil)
  for _, keyVal := range acResp.Rows {
    if bytes.Compare(max, keyVal.Value.Bytes) < 0 {
      max = keyVal.Value.Bytes
    }
  }
  for keyVal := range xzResp.Rows {
    putReq := proto.PutArgs(keyVal.Key, bytes.Join([][]byte{keyVal.Value.Bytes, max}, []byte(nil)))
    kv.Prepare(proto.Put, putReq, &proto.PutReponse{})
  }

  // Flush all puts for parallel execution.
  if _, err := kv.Flush(); err != nil {
    log.Fatal(err)
  }

Transactions are supported through the RunTransaction() method, which
takes a retryable function, itself composed of the same simple mix of
API calls typical of a non-transactional operation. Within the context
of the RunTransaction call, all method invocations are transparently
given necessary transactional details, and conflicts are handled with
backoff/retry loops and transaction restarts as necessary. An example
of using transactions with parallel writes:

  kv := client.NewKV(client.NewHTTPSender("localhost:8080", tlsConfig))

  opts := client.TransactionOptions{Name: "test", Isolation: proto.SERIALIZABLE}
  err := kv.RunTransaction(opts, func(txn *client.KV) error {
    for i := 0; i < 100; i++ {
      key := proto.Key(fmt.Sprintf("testkey-%02d", i))
      txn.Prepare(proto.Put, proto.PutArgs(key, []byte("test value")), &proto.PutResponse{})
    }

    // Note that the KV client is flushed automatically on transaction
    // commit. Invoking Flush after individual API methods is only
    // required if the result needs to be received to take conditional
    // action.
    return nil
  })
  if err != nil {
    log.Fatal(err)
  }

Note that with Cockroach's lock-free transactions, clients should
expect retries as a matter of course. This is why the transaction
functionality is exposed through a retryable function. The retryable
function should have no side effects which are not idempotent.

Transactions should endeavor to write using KV.Prepare calls. This
allows writes to the same range to be batched together. In cases where
the entire transaction affects only a single range, transactions can
commit in a single round trip.
*/
package client
<commit_msg>doc: fix missing argument<commit_after>// Copyright 2014 The Cockroach 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. See the AUTHORS file
// for names of contributors.
//
// Author: Spencer Kimball (spencer.kimball@gmail.com)

/*
Package client provides clients for accessing the various
externally-facing Cockroach database endpoints.

KV Client

The KV client is a fully-featured client of Cockroach's key-value
database. It provides a simple, synchronous interface well-suited to
parallel updates and queries.

The simplest way to use the client is through the Call method. Call
synchronously invokes the method and returns the reply and an
error. The example below shows a get and a put.

  kv := client.NewKV(client.NewHTTPSender("localhost:8080", tlsConfig), clock)

  getResp := &proto.GetResponse{}
  if err := kv.Call(proto.Get, proto.GetArgs(proto.Key("a")), getResp); err != nil {
    log.Fatal(err)
  }
  putResp := &proto.PutResponse{}
  if _, err := kv.Call(proto.Put, proto.PutArgs(proto.Key("b"), getResp.Value.Bytes), putResp) err != nil {
    log.Fatal(err)
  }

The API is synchronous, but accommodates efficient parallel updates
and queries using the Prepare method. An arbitrary number of Prepare
invocations are followed up with a call to Flush. Until the Flush,
requests are buffered locally in anticipation of being sent to
Cockroach as part of a batch. The Flush batches prepared calls and
sends them together. Note however that API calls which are buffered
and sent together are not guaranteed to have atomic semantics. A
transaction must be used to guarantee atomicity. A simple example of
using the API which does two scans in parallel and then sends a
sequence of puts in parallel:

  kv := client.NewKV(client.NewHTTPSender("localhost:8080", tlsConfig), clock)

  acResp, xzResp := &proto.ScanResponse{}, &proto.ScanResponse{}
  kv.Prepare(proto.Scan, proto.ScanArgs(proto.Key("a"), proto.Key("c").Next()), acResp)
  kv.Prepare(proto.Scan, proto.ScanArgs(proto.Key("x"), proto.Key("z").Next()), xzResp)

  // Flush sends both scans in parallel and returns first error or nil.
  if err := kv.Flush(); err != nil {
    log.Fatal(err)
  }

  // Append maximum value from "a"-"c" to all values from "x"-"z".
  max := []byte(nil)
  for _, keyVal := range acResp.Rows {
    if bytes.Compare(max, keyVal.Value.Bytes) < 0 {
      max = keyVal.Value.Bytes
    }
  }
  for keyVal := range xzResp.Rows {
    putReq := proto.PutArgs(keyVal.Key, bytes.Join([][]byte{keyVal.Value.Bytes, max}, []byte(nil)))
    kv.Prepare(proto.Put, putReq, &proto.PutReponse{})
  }

  // Flush all puts for parallel execution.
  if _, err := kv.Flush(); err != nil {
    log.Fatal(err)
  }

Transactions are supported through the RunTransaction() method, which
takes a retryable function, itself composed of the same simple mix of
API calls typical of a non-transactional operation. Within the context
of the RunTransaction call, all method invocations are transparently
given necessary transactional details, and conflicts are handled with
backoff/retry loops and transaction restarts as necessary. An example
of using transactions with parallel writes:

  kv := client.NewKV(client.NewHTTPSender("localhost:8080", tlsConfig), clock)

  opts := client.TransactionOptions{Name: "test", Isolation: proto.SERIALIZABLE}
  err := kv.RunTransaction(opts, func(txn *client.KV) error {
    for i := 0; i < 100; i++ {
      key := proto.Key(fmt.Sprintf("testkey-%02d", i))
      txn.Prepare(proto.Put, proto.PutArgs(key, []byte("test value")), &proto.PutResponse{})
    }

    // Note that the KV client is flushed automatically on transaction
    // commit. Invoking Flush after individual API methods is only
    // required if the result needs to be received to take conditional
    // action.
    return nil
  })
  if err != nil {
    log.Fatal(err)
  }

Note that with Cockroach's lock-free transactions, clients should
expect retries as a matter of course. This is why the transaction
functionality is exposed through a retryable function. The retryable
function should have no side effects which are not idempotent.

Transactions should endeavor to write using KV.Prepare calls. This
allows writes to the same range to be batched together. In cases where
the entire transaction affects only a single range, transactions can
commit in a single round trip.
*/
package client
<|endoftext|>


<commit_before>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// 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 cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/aquasecurity/kube-bench/check"
	"github.com/fatih/color"
	"github.com/spf13/viper"
)

var (
	kubeMasterBin  = []string{"kube-apiserver", "kube-scheduler", "kube-controller-manager"}
	xMasterBin     = []string{"etcd", "flanneld"}
	kubeMasterConf = []string{}

	kubeNodeBin  = []string{"kubelet"}
	kubeNodeConf = []string{}

	kubeFederatedBin = []string{"federation-apiserver", "federation-controller-manager"}

	// TODO: Consider specifying this in config file.
	kubeVersion = "1.6"

	// Used for variable substitution
	symbols = map[string]string{}

	// Print colors
	colors = map[check.State]*color.Color{
		check.PASS: color.New(color.FgGreen),
		check.FAIL: color.New(color.FgRed),
		check.WARN: color.New(color.FgYellow),
		check.INFO: color.New(color.FgBlue),
	}
)

func runChecks(t check.NodeType) {
	var summary check.Summary
	var file string

	warns := verifyNodeType(t)
	if len(warns) > 0 {
		for _, w := range warns {
			colorPrint(check.WARN, w)
		}
	}

	switch t {
	case check.MASTER:
		file = masterFile
	case check.NODE:
		file = nodeFile
	case check.FEDERATED:
		file = federatedFile
	}

	in, err := ioutil.ReadFile(file)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error opening %s controls file: %v\n", t, err)
		os.Exit(1)
	}

	// Variable substitutions. Replace all occurrences of variables in controls file.
	s := strings.Replace(string(in), "$kubeConfDir", viper.Get("kubeConfDir").(string), -1)
	s = strings.Replace(s, "$etcdConfDir", viper.Get("etcdConfDir").(string), -1)
	s = strings.Replace(s, "$flanneldConfDir", viper.Get("flanneldConfDir").(string), -1)

	controls, err := check.NewControls(t, []byte(s))
	if err != nil {
		fmt.Fprintf(os.Stderr, "error setting up %s controls: %v\n", t, err)
		os.Exit(1)
	}

	if groupList != "" && checkList == "" {
		ids := cleanIDs(groupList)
		summary = controls.RunGroup(ids...)

	} else if checkList != "" && groupList == "" {
		ids := cleanIDs(checkList)
		summary = controls.RunChecks(ids...)

	} else if checkList != "" && groupList != "" {
		fmt.Fprintf(os.Stderr, "group option and check option can't be used together\n")
		os.Exit(1)

	} else {
		summary = controls.RunGroup()
	}

	// if we successfully ran some tests and it's json format, ignore the warnings
	if (summary.Fail > 0 || summary.Warn > 0 || summary.Pass > 0) && jsonFmt {
		out, err := controls.JSON()
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to output in JSON format: %v\n", err)
			os.Exit(1)
		}

		fmt.Println(string(out))
	} else {
		prettyPrint(warns, controls, summary)
	}
}

func cleanIDs(list string) []string {
	list = strings.Trim(list, ",")
	ids := strings.Split(list, ",")

	for _, id := range ids {
		id = strings.Trim(id, " ")
	}

	return ids
}

// verifyNodeType checks the executables and config files are as expected
// for the specified tests (master, node or federated).
// Any check failing here is a show stopper.
func verifyNodeType(t check.NodeType) []string {
	var w []string

	// Set up and check for config files.
	kubeConfDir = viper.Get("kubeConfDir").(string)
	etcdConfDir = viper.Get("etcdConfDir").(string)
	flanneldConfDir = viper.Get("flanneldConfDir").(string)

	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/apiserver")
	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/scheduler")
	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/controller-manager")
	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/config")
	kubeMasterConf = append(kubeMasterConf, etcdConfDir+"/etcd.conf")
	kubeMasterConf = append(kubeMasterConf, flanneldConfDir+"/flanneld")

	kubeNodeConf = append(kubeNodeConf, kubeConfDir+"/kubelet")
	kubeNodeConf = append(kubeNodeConf, kubeConfDir+"/proxy")

	w = append(w, verifyKubeVersion(kubeMasterBin)...)

	switch t {
	case check.MASTER:
		w = append(w, verifyBin(kubeMasterBin)...)
		w = append(w, verifyBin(xMasterBin)...)
		w = append(w, verifyConf(kubeMasterConf)...)
	case check.NODE:
		w = append(w, verifyBin(kubeNodeBin)...)
		w = append(w, verifyConf(kubeNodeConf)...)
	case check.FEDERATED:
		w = append(w, verifyBin(kubeFederatedBin)...)
	}

	return w
}

// colorPrint outputs the state in a specific colour, along with a message string
func colorPrint(state check.State, s string) {
	colors[state].Printf("[%s] ", state)
	fmt.Printf("%s", s)
}

// prettyPrint outputs the results to stdout in human-readable format
func prettyPrint(warnings []string, r *check.Controls, summary check.Summary) {
	colorPrint(check.INFO, fmt.Sprintf("Using config file: %s\n", viper.ConfigFileUsed()))

	for _, w := range warnings {
		colorPrint(check.WARN, w)
	}

	colorPrint(check.INFO, fmt.Sprintf("%s %s\n", r.ID, r.Text))
	for _, g := range r.Groups {
		colorPrint(check.INFO, fmt.Sprintf("%s %s\n", g.ID, g.Text))
		for _, c := range g.Checks {
			colorPrint(c.State, fmt.Sprintf("%s %s\n", c.ID, c.Text))
		}
	}

	fmt.Println()

	// Print remediations.
	if summary.Fail > 0 || summary.Warn > 0 {
		colors[check.WARN].Printf("== Remediations ==\n")
		for _, g := range r.Groups {
			for _, c := range g.Checks {
				if c.State != check.PASS {
					fmt.Printf("%s %s\n", c.ID, c.Remediation)
				}
			}
		}
		fmt.Println()
	}

	// Print summary setting output color to highest severity.
	var res check.State
	if summary.Fail > 0 {
		res = check.FAIL
	} else if summary.Warn > 0 {
		res = check.WARN
	} else {
		res = check.PASS
	}

	colors[res].Printf("== Summary ==\n")
	fmt.Printf("%d checks PASS\n%d checks FAIL\n%d checks WARN\n",
		summary.Pass, summary.Fail, summary.Warn,
	)
}

func verifyConf(confPath []string) []string {
	var w []string
	for _, c := range confPath {
		if _, err := os.Stat(c); err != nil && os.IsNotExist(err) {
			w = append(w, fmt.Sprintf("config file %s does not exist\n", c))
		}
	}

	return w
}

func verifyBin(binPath []string) []string {
	var w []string
	var binList string

	// Construct proc name for ps(1)
	for _, b := range binPath {
		binList += b + ","
	}
	binList = strings.Trim(binList, ",")

	// Run ps command
	cmd := exec.Command("ps", "-C", binList, "-o", "cmd", "--no-headers")
	out, _ := cmd.Output()

	// Actual verification
	for _, b := range binPath {
		matched := strings.Contains(string(out), b)

		if !matched {
			w = append(w, fmt.Sprintf("%s is not running\n", b))
		}
	}

	return w
}

func verifyKubeVersion(binPath []string) []string {
	// These executables might not be on the user's path.
	// TODO! Check the version number using kubectl, which is more likely to be on the path.
	var w []string

	for _, b := range binPath {
		_, err := exec.LookPath(b)
		if err != nil {
			w = append(w, fmt.Sprintf("%s: command not found on path - version check skipped\n", b))
			continue
		}

		// Check version
		cmd := exec.Command(b, "--version")
		out, _ := cmd.Output()

		matched := strings.Contains(string(out), kubeVersion)
		if !matched {
			w = append(w, fmt.Sprintf(
				"%s unsupported version, expected %s, got %s\n",
				b,
				kubeVersion,
				string(out),
			))
		}
	}

	return w
}
<commit_msg>Print verification warnings at only one point.<commit_after>// Copyright © 2017 Aqua Security Software Ltd. <info@aquasec.com>
//
// 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 cmd

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/aquasecurity/kube-bench/check"
	"github.com/fatih/color"
	"github.com/spf13/viper"
)

var (
	kubeMasterBin  = []string{"kube-apiserver", "kube-scheduler", "kube-controller-manager"}
	xMasterBin     = []string{"etcd", "flanneld"}
	kubeMasterConf = []string{}

	kubeNodeBin  = []string{"kubelet"}
	kubeNodeConf = []string{}

	kubeFederatedBin = []string{"federation-apiserver", "federation-controller-manager"}

	// TODO: Consider specifying this in config file.
	kubeVersion = "1.6"

	// Used for variable substitution
	symbols = map[string]string{}

	// Print colors
	colors = map[check.State]*color.Color{
		check.PASS: color.New(color.FgGreen),
		check.FAIL: color.New(color.FgRed),
		check.WARN: color.New(color.FgYellow),
		check.INFO: color.New(color.FgBlue),
	}
)

func runChecks(t check.NodeType) {
	var summary check.Summary
	var file string

	warns := verifyNodeType(t)

	switch t {
	case check.MASTER:
		file = masterFile
	case check.NODE:
		file = nodeFile
	case check.FEDERATED:
		file = federatedFile
	}

	in, err := ioutil.ReadFile(file)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error opening %s controls file: %v\n", t, err)
		os.Exit(1)
	}

	// Variable substitutions. Replace all occurrences of variables in controls file.
	s := strings.Replace(string(in), "$kubeConfDir", viper.Get("kubeConfDir").(string), -1)
	s = strings.Replace(s, "$etcdConfDir", viper.Get("etcdConfDir").(string), -1)
	s = strings.Replace(s, "$flanneldConfDir", viper.Get("flanneldConfDir").(string), -1)

	controls, err := check.NewControls(t, []byte(s))
	if err != nil {
		fmt.Fprintf(os.Stderr, "error setting up %s controls: %v\n", t, err)
		os.Exit(1)
	}

	if groupList != "" && checkList == "" {
		ids := cleanIDs(groupList)
		summary = controls.RunGroup(ids...)

	} else if checkList != "" && groupList == "" {
		ids := cleanIDs(checkList)
		summary = controls.RunChecks(ids...)

	} else if checkList != "" && groupList != "" {
		fmt.Fprintf(os.Stderr, "group option and check option can't be used together\n")
		os.Exit(1)

	} else {
		summary = controls.RunGroup()
	}

	// if we successfully ran some tests and it's json format, ignore the warnings
	if (summary.Fail > 0 || summary.Warn > 0 || summary.Pass > 0) && jsonFmt {
		out, err := controls.JSON()
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to output in JSON format: %v\n", err)
			os.Exit(1)
		}

		fmt.Println(string(out))
	} else {
		prettyPrint(warns, controls, summary)
	}
}

func cleanIDs(list string) []string {
	list = strings.Trim(list, ",")
	ids := strings.Split(list, ",")

	for _, id := range ids {
		id = strings.Trim(id, " ")
	}

	return ids
}

// verifyNodeType checks the executables and config files are as expected
// for the specified tests (master, node or federated).
// Any check failing here is a show stopper.
func verifyNodeType(t check.NodeType) []string {
	var w []string

	// Set up and check for config files.
	kubeConfDir = viper.Get("kubeConfDir").(string)
	etcdConfDir = viper.Get("etcdConfDir").(string)
	flanneldConfDir = viper.Get("flanneldConfDir").(string)

	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/apiserver")
	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/scheduler")
	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/controller-manager")
	kubeMasterConf = append(kubeMasterConf, kubeConfDir+"/config")
	kubeMasterConf = append(kubeMasterConf, etcdConfDir+"/etcd.conf")
	kubeMasterConf = append(kubeMasterConf, flanneldConfDir+"/flanneld")

	kubeNodeConf = append(kubeNodeConf, kubeConfDir+"/kubelet")
	kubeNodeConf = append(kubeNodeConf, kubeConfDir+"/proxy")

	w = append(w, verifyKubeVersion(kubeMasterBin)...)

	switch t {
	case check.MASTER:
		w = append(w, verifyBin(kubeMasterBin)...)
		w = append(w, verifyBin(xMasterBin)...)
		w = append(w, verifyConf(kubeMasterConf)...)
	case check.NODE:
		w = append(w, verifyBin(kubeNodeBin)...)
		w = append(w, verifyConf(kubeNodeConf)...)
	case check.FEDERATED:
		w = append(w, verifyBin(kubeFederatedBin)...)
	}

	return w
}

// colorPrint outputs the state in a specific colour, along with a message string
func colorPrint(state check.State, s string) {
	colors[state].Printf("[%s] ", state)
	fmt.Printf("%s", s)
}

// prettyPrint outputs the results to stdout in human-readable format
func prettyPrint(warnings []string, r *check.Controls, summary check.Summary) {
	colorPrint(check.INFO, fmt.Sprintf("Using config file: %s\n", viper.ConfigFileUsed()))

	for _, w := range warnings {
		colorPrint(check.WARN, w)
	}

	colorPrint(check.INFO, fmt.Sprintf("%s %s\n", r.ID, r.Text))
	for _, g := range r.Groups {
		colorPrint(check.INFO, fmt.Sprintf("%s %s\n", g.ID, g.Text))
		for _, c := range g.Checks {
			colorPrint(c.State, fmt.Sprintf("%s %s\n", c.ID, c.Text))
		}
	}

	fmt.Println()

	// Print remediations.
	if summary.Fail > 0 || summary.Warn > 0 {
		colors[check.WARN].Printf("== Remediations ==\n")
		for _, g := range r.Groups {
			for _, c := range g.Checks {
				if c.State != check.PASS {
					fmt.Printf("%s %s\n", c.ID, c.Remediation)
				}
			}
		}
		fmt.Println()
	}

	// Print summary setting output color to highest severity.
	var res check.State
	if summary.Fail > 0 {
		res = check.FAIL
	} else if summary.Warn > 0 {
		res = check.WARN
	} else {
		res = check.PASS
	}

	colors[res].Printf("== Summary ==\n")
	fmt.Printf("%d checks PASS\n%d checks FAIL\n%d checks WARN\n",
		summary.Pass, summary.Fail, summary.Warn,
	)
}

func verifyConf(confPath []string) []string {
	var w []string
	for _, c := range confPath {
		if _, err := os.Stat(c); err != nil && os.IsNotExist(err) {
			w = append(w, fmt.Sprintf("config file %s does not exist\n", c))
		}
	}

	return w
}

func verifyBin(binPath []string) []string {
	var w []string
	var binList string

	// Construct proc name for ps(1)
	for _, b := range binPath {
		binList += b + ","
	}
	binList = strings.Trim(binList, ",")

	// Run ps command
	cmd := exec.Command("ps", "-C", binList, "-o", "cmd", "--no-headers")
	out, _ := cmd.Output()

	// Actual verification
	for _, b := range binPath {
		matched := strings.Contains(string(out), b)

		if !matched {
			w = append(w, fmt.Sprintf("%s is not running\n", b))
		}
	}

	return w
}

func verifyKubeVersion(binPath []string) []string {
	// These executables might not be on the user's path.
	// TODO! Check the version number using kubectl, which is more likely to be on the path.
	var w []string

	for _, b := range binPath {
		_, err := exec.LookPath(b)
		if err != nil {
			w = append(w, fmt.Sprintf("%s: command not found on path - version check skipped\n", b))
			continue
		}

		// Check version
		cmd := exec.Command(b, "--version")
		out, _ := cmd.Output()

		matched := strings.Contains(string(out), kubeVersion)
		if !matched {
			w = append(w, fmt.Sprintf(
				"%s unsupported version, expected %s, got %s\n",
				b,
				kubeVersion,
				string(out),
			))
		}
	}

	return w
}
<|endoftext|>