<commit_before>/*-
 * Copyright (c) 2016, 1&1 Internet SE
 * Copyright (c) 2016, Jörg Pernfuß
 *
 * Use of this source code is governed by a 2-clause BSD license
 * that can be found in the LICENSE file.
 */

package adm

import (
	"fmt"
	"log"
	"strings"
)

// ParseVariadicArguments parses whitespace separated argument lists
// of keyword/value pairs were keywords can be specified multiple
// times, some keywords are required and some only allowed once.
// Sequence of multiple keywords are detected and lead to abort
//
// multKeys => [ "port", "transport" ]
// uniqKeys => [ "team" ]
// reqKeys  => [ "team" ]
// args     => [ "port", "53", "transport", "tcp", "transport",
//               "udp", "team", "GenericOps" ]
//
// result => result["team"] = [ "GenericOps" ]
//           result["port"] = [ "53" ]
//           result["transport"] = [ "tcp", "udp" ]
func ParseVariadicArguments(
	multKeys []string, // keys that may appear multiple times
	uniqKeys []string, // keys that are allowed at most once
	reqKeys []string, // keys that are required at least one
	args []string, // arguments to parse
) map[string][]string {
	// returns a map of slices of string
	result := make(map[string][]string)

	// merge key slices
	keys := append(multKeys, uniqKeys...)

	// helper to skip over next value in args slice
	skip := false

	for pos, val := range args {
		// skip current arg if last argument was a keyword
		if skip {
			skip = false
			continue
		}

		if SliceContainsString(val, keys) {
			// there must be at least one arguments left
			if len(args[pos+1:]) < 1 {
				Abort("Syntax error, incomplete key/value specification (too few items left to parse)")
			}
			// check for back-to-back keyswords
			CheckStringNotAKeyword(args[pos+1], keys)

			// append value of current keyword into result map
			result[val] = append(result[val], args[pos+1])
			skip = true
			continue
		}
		// keywords trigger continue before this
		// values after keywords are skip'ed
		// reaching this is an error
		Abort(fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if we managed to collect all required keywords
	for _, key := range reqKeys {
		// ok is false if slice is nil
		if _, ok := result[key]; !ok {
			Abort(fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specified once
	for _, key := range uniqKeys {
		if sl, ok := result[key]; ok && (len(sl) > 1) {
			Abort(fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

	return result
}

func SliceContainsString(s string, sl []string) bool {
	for _, v := range sl {
		if v == s {
			return true
		}
	}
	return false
}

func CheckStringNotAKeyword(s string, keys []string) {
	if SliceContainsString(s, keys) {
		log.Fatal(`Syntax error, back-to-back keywords`) // XXX
	}
}

// combineStrings takes an arbitray number of strings and combines them
// into one, separated by `.\n`
func combineStrings(s ...string) string {
	var out string
	spacer := ``
	for _, in := range s {
		// ensure a single trailing .
		out = fmt.Sprintf("%s%s", out+spacer, strings.TrimRight(in, `.`)+`.`)
		spacer = "\n"
	}
	return out
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Update adm.ParseVariadicArguments()<commit_after>/*-
 * Copyright (c) 2016, 1&1 Internet SE
 * Copyright (c) 2016, Jörg Pernfuß
 *
 * Use of this source code is governed by a 2-clause BSD license
 * that can be found in the LICENSE file.
 */

package adm

import (
	"fmt"
	"strings"
)

// ParseVariadicArguments parses split up argument lists of
// keyword/value pairs were keywords can be specified multiple
// times, some keywords are required and some only allowed once.
// Sequence of multiple keywords are detected and lead to abort
//
//	multKeys => [ "port", "transport" ]
//	uniqKeys => [ "team" ]
//	reqKeys  => [ "team" ]
//	args     => [ "port", "53", "transport", "tcp", "transport",
//	              "udp", "team", "GenericOps" ]
//
//	result => result["team"] = [ "GenericOps" ]
//	          result["port"] = [ "53" ]
//	          result["transport"] = [ "tcp", "udp" ]
func ParseVariadicArguments(
	result map[string][]string, // provided result map
	multKeys []string, // keys that may appear multiple times
	uniqKeys []string, // keys that are allowed at most once
	reqKeys []string, // keys that are required at least one
	args []string, // arguments to parse
) error {
	// ensure the result is empty
	result = make(map[string][]string)
	// used to hold found errors, so if three keywords are missing they can
	// all be mentioned in one call
	errors := []string{}

	// merge key slices
	keys := append(multKeys, uniqKeys...)

	// helper to skip over next value in args slice
	skip := false

	for pos, val := range args {
		// skip current arg if last argument was a keyword
		if skip {
			skip = false
			continue
		}

		if sliceContainsString(val, keys) {
			// there must be at least one arguments left
			if len(args[pos+1:]) < 1 {
				errors = append(errors,
					`Syntax error, incomplete key/value specification (too few items left to parse)`,
				)
				goto abort
			}
			// check for back-to-back keyswords
			if err := checkStringNotAKeyword(args[pos+1], keys); err != nil {
				errors = append(errors, err.Error())
				goto abort
			}

			// append value of current keyword into result map
			result[val] = append(result[val], args[pos+1])
			skip = true
			continue
		}
		// keywords trigger continue before this
		// values after keywords are skip'ed
		// reaching this is an error
		errors = append(errors, fmt.Sprintf("Syntax error, erroneus argument: %s", val))
	}

	// check if we managed to collect all required keywords
	for _, key := range reqKeys {
		// ok is false if slice is nil
		if _, ok := result[key]; !ok {
			errors = append(errors, fmt.Sprintf("Syntax error, missing keyword: %s", key))
		}
	}

	// check if unique keywords were only specified once
	for _, key := range uniqKeys {
		if sl, ok := result[key]; ok && (len(sl) > 1) {
			errors = append(errors, fmt.Sprintf("Syntax error, keyword must only be provided once: %s", key))
		}
	}

abort:
	if len(errors) > 0 {
		result = nil
		return fmt.Errorf(combineStrings(errors...))
	}

	return nil
}

func sliceContainsString(s string, sl []string) bool {
	for _, v := range sl {
		if v == s {
			return true
		}
	}
	return false
}

func checkStringNotAKeyword(s string, keys []string) error {
	if sliceContainsString(s, keys) {
		return fmt.Errorf("Syntax error, back-to-back keyword: %s", s)
	}
	return nil
}

// combineStrings takes an arbitray number of strings and combines them
// into one, separated by `.\n`
func combineStrings(s ...string) string {
	var out string
	spacer := ``
	for _, in := range s {
		// ensure a single trailing .
		out = fmt.Sprintf("%s%s", out+spacer, strings.TrimRight(in, `.`)+`.`)
		spacer = "\n"
	}
	return out
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package config

import (
	"errors"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	yaml "gopkg.in/yaml.v2"

	"gitlab.com/gitlab-org/gitlab-shell/client"
	"gitlab.com/gitlab-org/gitlab-shell/internal/metrics"

	"gitlab.com/gitlab-org/labkit/log"
)

const (
	configFile            = "config.yml"
	defaultSecretFileName = ".gitlab_shell_secret"
)

type ServerConfig struct {
	Listen                  string   `yaml:"listen,omitempty"`
	ProxyProtocol           bool     `yaml:"proxy_protocol,omitempty"`
	WebListen               string   `yaml:"web_listen,omitempty"`
	ConcurrentSessionsLimit int64    `yaml:"concurrent_sessions_limit,omitempty"`
	GracePeriodSeconds      uint64   `yaml:"grace_period"`
	ReadinessProbe          string   `yaml:"readiness_probe"`
	LivenessProbe           string   `yaml:"liveness_probe"`
	HostKeyFiles            []string `yaml:"host_key_files,omitempty"`
}

type HttpSettingsConfig struct {
	User               string `yaml:"user"`
	Password           string `yaml:"password"`
	ReadTimeoutSeconds uint64 `yaml:"read_timeout"`
	CaFile             string `yaml:"ca_file"`
	CaPath             string `yaml:"ca_path"`
	SelfSignedCert     bool   `yaml:"self_signed_cert"`
}

type Config struct {
	User                  string `yaml:"user,omitempty"`
	RootDir               string
	LogFile               string `yaml:"log_file,omitempty"`
	LogFormat             string `yaml:"log_format,omitempty"`
	LogLevel              string `yaml:"log_level,omitempty"`
	GitlabUrl             string `yaml:"gitlab_url"`
	GitlabRelativeURLRoot string `yaml:"gitlab_relative_url_root"`
	GitlabTracing         string `yaml:"gitlab_tracing"`
	// SecretFilePath is only for parsing. Application code should always use Secret.
	SecretFilePath string             `yaml:"secret_file"`
	Secret         string             `yaml:"secret"`
	SslCertDir     string             `yaml:"ssl_cert_dir"`
	HttpSettings   HttpSettingsConfig `yaml:"http_settings"`
	Server         ServerConfig       `yaml:"sshd"`

	httpClient     *client.HttpClient
	httpClientErr  error
	httpClientOnce sync.Once
}

// The defaults to apply before parsing the config file(s).
var (
	DefaultConfig = Config{
		LogFile:   "gitlab-shell.log",
		LogFormat: "json",
		LogLevel:  "info",
		Server:    DefaultServerConfig,
		User:      "git",
	}

	DefaultServerConfig = ServerConfig{
		Listen:                  "[::]:22",
		WebListen:               "localhost:9122",
		ConcurrentSessionsLimit: 10,
		GracePeriodSeconds:      10,
		ReadinessProbe:          "/start",
		LivenessProbe:           "/health",
		HostKeyFiles: []string{
			"/run/secrets/ssh-hostkeys/ssh_host_rsa_key",
			"/run/secrets/ssh-hostkeys/ssh_host_ecdsa_key",
			"/run/secrets/ssh-hostkeys/ssh_host_ed25519_key",
		},
	}
)

func (sc *ServerConfig) GracePeriod() time.Duration {
	return time.Duration(sc.GracePeriodSeconds) * time.Second
}

func (c *Config) ApplyGlobalState() {
	if c.SslCertDir != "" {
		log.WithFields(log.Fields{"ssl_cert_dir": c.SslCertDir}).Info("SSL_CERT_DIR is configured")

		os.Setenv("SSL_CERT_DIR", c.SslCertDir)
	}
}

func (c *Config) HttpClient() (*client.HttpClient, error) {
	c.httpClientOnce.Do(func() {
		client, err := client.NewHTTPClientWithOpts(
			c.GitlabUrl,
			c.GitlabRelativeURLRoot,
			c.HttpSettings.CaFile,
			c.HttpSettings.CaPath,
			c.HttpSettings.SelfSignedCert,
			c.HttpSettings.ReadTimeoutSeconds,
			nil,
		)
		if err != nil {
			c.httpClientErr = err
			return
		}

		tr := client.Transport
		client.Transport = promhttp.InstrumentRoundTripperDuration(metrics.HttpRequestDuration, tr)

		c.httpClient = client
	})

	return c.httpClient, c.httpClientErr
}

// NewFromDirExternal returns a new config from a given root dir. It also applies defaults appropriate for
// gitlab-shell running in an external SSH server.
func NewFromDirExternal(dir string) (*Config, error) {
	cfg, err := newFromFile(filepath.Join(dir, configFile))
	if err != nil {
		return nil, err
	}

	cfg.ApplyGlobalState()

	return cfg, nil
}

// NewFromDir returns a new config given a root directory. It looks for the config file name in the
// given directory and reads the config from it. It doesn't apply any defaults. New code should prefer
// this over NewFromDirIntegrated and apply the right default via one of the Apply... functions.
func NewFromDir(dir string) (*Config, error) {
	return newFromFile(filepath.Join(dir, configFile))
}

// newFromFile reads a new Config instance from the given file path. It doesn't apply any defaults.
func newFromFile(path string) (*Config, error) {
	cfg := &Config{}
	*cfg = DefaultConfig
	cfg.RootDir = filepath.Dir(path)

	configBytes, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	if err := yaml.Unmarshal(configBytes, cfg); err != nil {
		return nil, err
	}

	if cfg.GitlabUrl != "" {
		// This is only done for historic reasons, don't implement it for new config sources.
		unescapedUrl, err := url.PathUnescape(cfg.GitlabUrl)
		if err != nil {
			return nil, err
		}

		cfg.GitlabUrl = unescapedUrl
	}

	if err := parseSecret(cfg); err != nil {
		return nil, err
	}

	if len(cfg.LogFile) > 0 && cfg.LogFile[0] != '/' && cfg.RootDir != "" {
		cfg.LogFile = filepath.Join(cfg.RootDir, cfg.LogFile)
	}

	return cfg, nil
}

func parseSecret(cfg *Config) error {
	// The secret was parsed from yaml no need to read another file
	if cfg.Secret != "" {
		return nil
	}

	if cfg.SecretFilePath == "" {
		cfg.SecretFilePath = defaultSecretFileName
	}

	if !filepath.IsAbs(cfg.SecretFilePath) {
		cfg.SecretFilePath = path.Join(cfg.RootDir, cfg.SecretFilePath)
	}

	secretFileContent, err := os.ReadFile(cfg.SecretFilePath)
	if err != nil {
		return err
	}
	cfg.Secret = string(secretFileContent)

	return nil
}

// IsSane checks if the given config fulfills the minimum requirements to be able to run.
// Any error returned by this function should be a startup error. On the other hand
// if this function returns nil, this doesn't guarantee the config will work, but it's
// at least worth a try.
func (cfg *Config) IsSane() error {
	if cfg.GitlabUrl == "" {
		return errors.New("gitlab_url is required")
	}
	if cfg.Secret == "" {
		return errors.New("secret or secret_file_path is required")
	}
	return nil
}
<commit_msg>Remove SSL_CERT_DIR logging<commit_after>package config

import (
	"errors"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus/promhttp"
	yaml "gopkg.in/yaml.v2"

	"gitlab.com/gitlab-org/gitlab-shell/client"
	"gitlab.com/gitlab-org/gitlab-shell/internal/metrics"
)

const (
	configFile            = "config.yml"
	defaultSecretFileName = ".gitlab_shell_secret"
)

type ServerConfig struct {
	Listen                  string   `yaml:"listen,omitempty"`
	ProxyProtocol           bool     `yaml:"proxy_protocol,omitempty"`
	WebListen               string   `yaml:"web_listen,omitempty"`
	ConcurrentSessionsLimit int64    `yaml:"concurrent_sessions_limit,omitempty"`
	GracePeriodSeconds      uint64   `yaml:"grace_period"`
	ReadinessProbe          string   `yaml:"readiness_probe"`
	LivenessProbe           string   `yaml:"liveness_probe"`
	HostKeyFiles            []string `yaml:"host_key_files,omitempty"`
}

type HttpSettingsConfig struct {
	User               string `yaml:"user"`
	Password           string `yaml:"password"`
	ReadTimeoutSeconds uint64 `yaml:"read_timeout"`
	CaFile             string `yaml:"ca_file"`
	CaPath             string `yaml:"ca_path"`
	SelfSignedCert     bool   `yaml:"self_signed_cert"`
}

type Config struct {
	User                  string `yaml:"user,omitempty"`
	RootDir               string
	LogFile               string `yaml:"log_file,omitempty"`
	LogFormat             string `yaml:"log_format,omitempty"`
	LogLevel              string `yaml:"log_level,omitempty"`
	GitlabUrl             string `yaml:"gitlab_url"`
	GitlabRelativeURLRoot string `yaml:"gitlab_relative_url_root"`
	GitlabTracing         string `yaml:"gitlab_tracing"`
	// SecretFilePath is only for parsing. Application code should always use Secret.
	SecretFilePath string             `yaml:"secret_file"`
	Secret         string             `yaml:"secret"`
	SslCertDir     string             `yaml:"ssl_cert_dir"`
	HttpSettings   HttpSettingsConfig `yaml:"http_settings"`
	Server         ServerConfig       `yaml:"sshd"`

	httpClient     *client.HttpClient
	httpClientErr  error
	httpClientOnce sync.Once
}

// The defaults to apply before parsing the config file(s).
var (
	DefaultConfig = Config{
		LogFile:   "gitlab-shell.log",
		LogFormat: "json",
		LogLevel:  "info",
		Server:    DefaultServerConfig,
		User:      "git",
	}

	DefaultServerConfig = ServerConfig{
		Listen:                  "[::]:22",
		WebListen:               "localhost:9122",
		ConcurrentSessionsLimit: 10,
		GracePeriodSeconds:      10,
		ReadinessProbe:          "/start",
		LivenessProbe:           "/health",
		HostKeyFiles: []string{
			"/run/secrets/ssh-hostkeys/ssh_host_rsa_key",
			"/run/secrets/ssh-hostkeys/ssh_host_ecdsa_key",
			"/run/secrets/ssh-hostkeys/ssh_host_ed25519_key",
		},
	}
)

func (sc *ServerConfig) GracePeriod() time.Duration {
	return time.Duration(sc.GracePeriodSeconds) * time.Second
}

func (c *Config) ApplyGlobalState() {
	if c.SslCertDir != "" {
		os.Setenv("SSL_CERT_DIR", c.SslCertDir)
	}
}

func (c *Config) HttpClient() (*client.HttpClient, error) {
	c.httpClientOnce.Do(func() {
		client, err := client.NewHTTPClientWithOpts(
			c.GitlabUrl,
			c.GitlabRelativeURLRoot,
			c.HttpSettings.CaFile,
			c.HttpSettings.CaPath,
			c.HttpSettings.SelfSignedCert,
			c.HttpSettings.ReadTimeoutSeconds,
			nil,
		)
		if err != nil {
			c.httpClientErr = err
			return
		}

		tr := client.Transport
		client.Transport = promhttp.InstrumentRoundTripperDuration(metrics.HttpRequestDuration, tr)

		c.httpClient = client
	})

	return c.httpClient, c.httpClientErr
}

// NewFromDirExternal returns a new config from a given root dir. It also applies defaults appropriate for
// gitlab-shell running in an external SSH server.
func NewFromDirExternal(dir string) (*Config, error) {
	cfg, err := newFromFile(filepath.Join(dir, configFile))
	if err != nil {
		return nil, err
	}

	cfg.ApplyGlobalState()

	return cfg, nil
}

// NewFromDir returns a new config given a root directory. It looks for the config file name in the
// given directory and reads the config from it. It doesn't apply any defaults. New code should prefer
// this over NewFromDirIntegrated and apply the right default via one of the Apply... functions.
func NewFromDir(dir string) (*Config, error) {
	return newFromFile(filepath.Join(dir, configFile))
}

// newFromFile reads a new Config instance from the given file path. It doesn't apply any defaults.
func newFromFile(path string) (*Config, error) {
	cfg := &Config{}
	*cfg = DefaultConfig
	cfg.RootDir = filepath.Dir(path)

	configBytes, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	if err := yaml.Unmarshal(configBytes, cfg); err != nil {
		return nil, err
	}

	if cfg.GitlabUrl != "" {
		// This is only done for historic reasons, don't implement it for new config sources.
		unescapedUrl, err := url.PathUnescape(cfg.GitlabUrl)
		if err != nil {
			return nil, err
		}

		cfg.GitlabUrl = unescapedUrl
	}

	if err := parseSecret(cfg); err != nil {
		return nil, err
	}

	if len(cfg.LogFile) > 0 && cfg.LogFile[0] != '/' && cfg.RootDir != "" {
		cfg.LogFile = filepath.Join(cfg.RootDir, cfg.LogFile)
	}

	return cfg, nil
}

func parseSecret(cfg *Config) error {
	// The secret was parsed from yaml no need to read another file
	if cfg.Secret != "" {
		return nil
	}

	if cfg.SecretFilePath == "" {
		cfg.SecretFilePath = defaultSecretFileName
	}

	if !filepath.IsAbs(cfg.SecretFilePath) {
		cfg.SecretFilePath = path.Join(cfg.RootDir, cfg.SecretFilePath)
	}

	secretFileContent, err := os.ReadFile(cfg.SecretFilePath)
	if err != nil {
		return err
	}
	cfg.Secret = string(secretFileContent)

	return nil
}

// IsSane checks if the given config fulfills the minimum requirements to be able to run.
// Any error returned by this function should be a startup error. On the other hand
// if this function returns nil, this doesn't guarantee the config will work, but it's
// at least worth a try.
func (cfg *Config) IsSane() error {
	if cfg.GitlabUrl == "" {
		return errors.New("gitlab_url is required")
	}
	if cfg.Secret == "" {
		return errors.New("secret or secret_file_path is required")
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package frontend

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/google/safehtml"
	"github.com/google/safehtml/uncheckedconversions"
	"golang.org/x/mod/module"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/cookie"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/experiment"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/stdlib"
)

// UnitPage contains data needed to render the unit template.
type UnitPage struct {
	basePage
	// Unit is the unit for this page.
	Unit *internal.UnitMeta

	// Breadcrumb contains data used to render breadcrumb UI elements.
	Breadcrumb breadcrumb

	// Title is the title of the page.
	Title string

	// URLPath is the path suitable for links on the page.
	// See the unitURLPath for details.
	URLPath string

	// CanonicalURLPath is a permanent representation of the URL path for a
	// unit.
	// It uses the resolved module path and version.
	// For example, if the latest version of /my.module/pkg is version v1.5.2,
	// the canonical URL path for that unit would be /my.module@v1.5.2/pkg
	CanonicalURLPath string

	// The version string formatted for display.
	DisplayVersion string

	// LinkVersion is version string suitable for links used to compute
	// latest badges.
	LinkVersion string

	// LatestURL is a url pointing to the latest version of a unit.
	LatestURL string

	// LatestMinorClass is the CSS class that describes the current unit's minor
	// version in relationship to the latest version of the unit.
	LatestMinorClass string

	// Information about the latest major version of the module.
	LatestMajorVersion    string
	LatestMajorVersionURL string

	// PageType is the type of page (pkg, cmd, dir, std, or mod).
	PageType string

	// PageLabels are the labels that will be displayed
	// for a given page.
	PageLabels []string

	// CanShowDetails indicates whether details can be shown or must be
	// hidden due to issues like license restrictions.
	CanShowDetails bool

	// Settings contains settings for the selected tab.
	SelectedTab TabSettings

	// RedirectedFromPath is the path that redirected to the current page.
	// If non-empty, a "redirected from" banner will be displayed
	// (see content/static/html/helpers/_unit_header.tmpl).
	RedirectedFromPath string

	// Details contains data specific to the type of page being rendered.
	Details interface{}
}

// serveUnitPage serves a unit page for a path using the paths,
// modules, documentation, readmes, licenses, and package_imports tables.
func (s *Server) serveUnitPage(ctx context.Context, w http.ResponseWriter, r *http.Request,
	ds internal.DataSource, info *urlPathInfo) (err error) {
	defer derrors.Wrap(&err, "serveUnitPage(ctx, w, r, ds, %v)", info)

	tab := r.FormValue("tab")
	if tab == "" {
		// Default to details tab when there is no tab param.
		tab = tabMain
	}
	// Redirect to clean URL path when tab param is invalid.
	if _, ok := unitTabLookup[tab]; !ok {
		http.Redirect(w, r, r.URL.Path, http.StatusFound)
		return nil
	}

	um, err := ds.GetUnitMeta(ctx, info.fullPath, info.modulePath, info.requestedVersion)
	if err != nil {
		if !errors.Is(err, derrors.NotFound) {
			return err
		}
		return s.servePathNotFoundPage(w, r, ds, info.fullPath, info.modulePath, info.requestedVersion)
	}

	recordVersionTypeMetric(ctx, info.requestedVersion)
	if _, ok := internal.DefaultBranches[info.requestedVersion]; ok {
		// Since path@master is a moving target, we don't want it to be stale.
		// As a result, we enqueue every request of path@master to the frontend
		// task queue, which will initiate a fetch request depending on the
		// last time we tried to fetch this module version.
		//
		// Use a separate context here to prevent the context from being canceled
		// elsewhere before a task is enqueued.
		go func() {
			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
			defer cancel()
			log.Infof(ctx, "serveUnitPage: Scheduling %q@%q to be fetched", info.modulePath, info.requestedVersion)
			if _, err := s.queue.ScheduleFetch(ctx, info.modulePath, info.requestedVersion, "", false); err != nil {
				log.Errorf(ctx, "serveUnitPage(%q): %v", r.URL.Path, err)
			}
		}()
	}

	if !isValidTabForUnit(tab, um) {
		// Redirect to clean URL path when tab param is invalid for the unit
		// type.
		http.Redirect(w, r, r.URL.Path, http.StatusFound)
		return nil
	}

	latestInfo := s.GetLatestInfo(r.Context(), um.Path, um.ModulePath)
	var redirectPath string
	redirectPath, err = cookie.Extract(w, r, cookie.AlternativeModuleFlash)
	if err != nil {
		// Don't fail, but don't display a banner either.
		log.Errorf(ctx, "extracting AlternativeModuleFlash cookie: %v", err)
	}
	tabSettings := unitTabLookup[tab]
	title := pageTitle(um)
	basePage := s.newBasePage(r, title)
	basePage.AllowWideContent = true
	lv := linkVersion(um.Version, um.ModulePath)
	_, majorVersion, _ := module.SplitPathVersion(um.ModulePath)
	_, latestMajorVersion, ok := module.SplitPathVersion(latestInfo.MajorModulePath)
	// Show the banner if there was no error getting the latest major version,
	// and it is different from the major version of the current module path.
	var latestMajorVersionNum string
	if ok && majorVersion != latestMajorVersion && latestMajorVersion != "" {
		latestMajorVersionNum = strings.TrimPrefix(latestMajorVersion, "/")
	}
	page := UnitPage{
		basePage:              basePage,
		Unit:                  um,
		Breadcrumb:            displayBreadcrumb(um, info.requestedVersion),
		Title:                 title,
		SelectedTab:           tabSettings,
		URLPath:               constructUnitURL(um.Path, um.ModulePath, info.requestedVersion),
		CanonicalURLPath:      canonicalURLPath(um),
		DisplayVersion:        displayVersion(um.Version, um.ModulePath),
		LinkVersion:           lv,
		LatestURL:             constructUnitURL(um.Path, um.ModulePath, internal.LatestVersion),
		LatestMinorClass:      latestMinorClass(r.Context(), lv, latestInfo),
		LatestMajorVersion:    latestMajorVersionNum,
		LatestMajorVersionURL: latestInfo.MajorUnitPath,
		PageLabels:            pageLabels(um),
		PageType:              pageType(um),
		RedirectedFromPath:    redirectPath,
	}

	// Use GOOS and GOARCH query parameters to create a build context, which
	// affects the documentation and synopsis. Omitting both results in an empty
	// build context, which will match the first (and preferred) build context.
	// It's also okay to provide just one (e.g. GOOS=windows), which will select
	// the first doc with that value, ignoring the other one.
	bc := internal.BuildContext{GOOS: r.FormValue("GOOS"), GOARCH: r.FormValue("GOARCH")}
	d, err := fetchDetailsForUnit(ctx, r, tab, ds, um, bc)
	if err != nil {
		return err
	}
	page.Details = d
	main, ok := d.(*MainDetails)
	if ok {
		page.MetaDescription = metaDescription(strconv.Itoa(main.ImportedByCount))
	}
	s.servePage(ctx, w, tabSettings.TemplateName, page)
	return nil
}

func latestMinorClass(ctx context.Context, version string, latest internal.LatestInfo) string {
	c := "DetailsHeader-badge"
	switch {
	case latest.MinorVersion == "":
		c += "--unknown"
	case latest.MinorVersion == version && !latest.UnitExistsAtMinor && experiment.IsActive(ctx, internal.ExperimentNotAtLatest):
		c += "--notAtLatest"
	case latest.MinorVersion == version:
		c += "--latest"
	default:
		c += "--goToLatest"
	}
	return c
}

// metaDescription uses a safehtml escape hatch to build HTML used
// to render the <meta name="Description"> for unit pages as a
// workaround for https://github.com/google/safehtml/issues/6.
func metaDescription(synopsis string) safehtml.HTML {
	if synopsis == "" {
		return safehtml.HTML{}
	}
	return safehtml.HTMLConcat(
		uncheckedconversions.HTMLFromStringKnownToSatisfyTypeContract(`<meta name="Description" content="`),
		safehtml.HTMLEscaped(synopsis),
		uncheckedconversions.HTMLFromStringKnownToSatisfyTypeContract(`">`),
	)
}

// isValidTabForUnit reports whether the tab is valid for the given unit.
// It is assumed that tab is a key in unitTabLookup.
func isValidTabForUnit(tab string, um *internal.UnitMeta) bool {
	if tab == tabLicenses && !um.IsRedistributable {
		return false
	}
	if !um.IsPackage() && (tab == tabImports || tab == tabImportedBy) {
		return false
	}
	return true
}

// constructUnitURL returns a URL path that refers to the given unit at the requested
// version. If requestedVersion is "latest", then the resulting path has no
// version; otherwise, it has requestedVersion.
func constructUnitURL(fullPath, modulePath, requestedVersion string) string {
	if requestedVersion == internal.LatestVersion {
		return "/" + fullPath
	}
	v := linkVersion(requestedVersion, modulePath)
	if fullPath == modulePath || modulePath == stdlib.ModulePath {
		return fmt.Sprintf("/%s@%s", fullPath, v)
	}
	return fmt.Sprintf("/%s@%s/%s", modulePath, v, strings.TrimPrefix(fullPath, modulePath+"/"))
}

// canonicalURLPath constructs a URL path to the unit that always includes the
// resolved version.
func canonicalURLPath(um *internal.UnitMeta) string {
	return constructUnitURL(um.Path, um.ModulePath, linkVersion(um.Version, um.ModulePath))
}
<commit_msg>internal/frontend: add JSON endpoint<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package frontend

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/google/safehtml"
	"github.com/google/safehtml/uncheckedconversions"
	"golang.org/x/mod/module"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/cookie"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/experiment"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/stdlib"
)

// UnitPage contains data needed to render the unit template.
type UnitPage struct {
	basePage
	// Unit is the unit for this page.
	Unit *internal.UnitMeta

	// Breadcrumb contains data used to render breadcrumb UI elements.
	Breadcrumb breadcrumb

	// Title is the title of the page.
	Title string

	// URLPath is the path suitable for links on the page.
	// See the unitURLPath for details.
	URLPath string

	// CanonicalURLPath is a permanent representation of the URL path for a
	// unit.
	// It uses the resolved module path and version.
	// For example, if the latest version of /my.module/pkg is version v1.5.2,
	// the canonical URL path for that unit would be /my.module@v1.5.2/pkg
	CanonicalURLPath string

	// The version string formatted for display.
	DisplayVersion string

	// LinkVersion is version string suitable for links used to compute
	// latest badges.
	LinkVersion string

	// LatestURL is a url pointing to the latest version of a unit.
	LatestURL string

	// LatestMinorClass is the CSS class that describes the current unit's minor
	// version in relationship to the latest version of the unit.
	LatestMinorClass string

	// Information about the latest major version of the module.
	LatestMajorVersion    string
	LatestMajorVersionURL string

	// PageType is the type of page (pkg, cmd, dir, std, or mod).
	PageType string

	// PageLabels are the labels that will be displayed
	// for a given page.
	PageLabels []string

	// CanShowDetails indicates whether details can be shown or must be
	// hidden due to issues like license restrictions.
	CanShowDetails bool

	// Settings contains settings for the selected tab.
	SelectedTab TabSettings

	// RedirectedFromPath is the path that redirected to the current page.
	// If non-empty, a "redirected from" banner will be displayed
	// (see content/static/html/helpers/_unit_header.tmpl).
	RedirectedFromPath string

	// Details contains data specific to the type of page being rendered.
	Details interface{}
}

// serveUnitPage serves a unit page for a path using the paths,
// modules, documentation, readmes, licenses, and package_imports tables.
func (s *Server) serveUnitPage(ctx context.Context, w http.ResponseWriter, r *http.Request,
	ds internal.DataSource, info *urlPathInfo) (err error) {
	defer derrors.Wrap(&err, "serveUnitPage(ctx, w, r, ds, %v)", info)

	tab := r.FormValue("tab")
	if tab == "" {
		// Default to details tab when there is no tab param.
		tab = tabMain
	}
	// Redirect to clean URL path when tab param is invalid.
	if _, ok := unitTabLookup[tab]; !ok {
		http.Redirect(w, r, r.URL.Path, http.StatusFound)
		return nil
	}

	um, err := ds.GetUnitMeta(ctx, info.fullPath, info.modulePath, info.requestedVersion)
	if err != nil {
		if !errors.Is(err, derrors.NotFound) {
			return err
		}
		return s.servePathNotFoundPage(w, r, ds, info.fullPath, info.modulePath, info.requestedVersion)
	}

	// Use GOOS and GOARCH query parameters to create a build context, which
	// affects the documentation and synopsis. Omitting both results in an empty
	// build context, which will match the first (and preferred) build context.
	// It's also okay to provide just one (e.g. GOOS=windows), which will select
	// the first doc with that value, ignoring the other one.
	bc := internal.BuildContext{GOOS: r.FormValue("GOOS"), GOARCH: r.FormValue("GOARCH")}
	d, err := fetchDetailsForUnit(ctx, r, tab, ds, um, bc)
	if err != nil {
		return err
	}
	if s.serveStats && r.FormValue("m") == "json" {
		data, err := json.Marshal(d)
		if err != nil {
			return fmt.Errorf("json.Marshal: %v", err)
		}
		if _, err := w.Write(data); err != nil {
			return fmt.Errorf("w.Write: %v", err)
		}
		return nil
	}

	recordVersionTypeMetric(ctx, info.requestedVersion)
	if _, ok := internal.DefaultBranches[info.requestedVersion]; ok {
		// Since path@master is a moving target, we don't want it to be stale.
		// As a result, we enqueue every request of path@master to the frontend
		// task queue, which will initiate a fetch request depending on the
		// last time we tried to fetch this module version.
		//
		// Use a separate context here to prevent the context from being canceled
		// elsewhere before a task is enqueued.
		go func() {
			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute)
			defer cancel()
			log.Infof(ctx, "serveUnitPage: Scheduling %q@%q to be fetched", info.modulePath, info.requestedVersion)
			if _, err := s.queue.ScheduleFetch(ctx, info.modulePath, info.requestedVersion, "", false); err != nil {
				log.Errorf(ctx, "serveUnitPage(%q): %v", r.URL.Path, err)
			}
		}()
	}

	if !isValidTabForUnit(tab, um) {
		// Redirect to clean URL path when tab param is invalid for the unit
		// type.
		http.Redirect(w, r, r.URL.Path, http.StatusFound)
		return nil
	}

	latestInfo := s.GetLatestInfo(r.Context(), um.Path, um.ModulePath)
	var redirectPath string
	redirectPath, err = cookie.Extract(w, r, cookie.AlternativeModuleFlash)
	if err != nil {
		// Don't fail, but don't display a banner either.
		log.Errorf(ctx, "extracting AlternativeModuleFlash cookie: %v", err)
	}
	tabSettings := unitTabLookup[tab]
	title := pageTitle(um)
	basePage := s.newBasePage(r, title)
	basePage.AllowWideContent = true
	lv := linkVersion(um.Version, um.ModulePath)
	_, majorVersion, _ := module.SplitPathVersion(um.ModulePath)
	_, latestMajorVersion, ok := module.SplitPathVersion(latestInfo.MajorModulePath)
	// Show the banner if there was no error getting the latest major version,
	// and it is different from the major version of the current module path.
	var latestMajorVersionNum string
	if ok && majorVersion != latestMajorVersion && latestMajorVersion != "" {
		latestMajorVersionNum = strings.TrimPrefix(latestMajorVersion, "/")
	}
	page := UnitPage{
		basePage:              basePage,
		Unit:                  um,
		Breadcrumb:            displayBreadcrumb(um, info.requestedVersion),
		Title:                 title,
		SelectedTab:           tabSettings,
		URLPath:               constructUnitURL(um.Path, um.ModulePath, info.requestedVersion),
		CanonicalURLPath:      canonicalURLPath(um),
		DisplayVersion:        displayVersion(um.Version, um.ModulePath),
		LinkVersion:           lv,
		LatestURL:             constructUnitURL(um.Path, um.ModulePath, internal.LatestVersion),
		LatestMinorClass:      latestMinorClass(r.Context(), lv, latestInfo),
		LatestMajorVersion:    latestMajorVersionNum,
		LatestMajorVersionURL: latestInfo.MajorUnitPath,
		PageLabels:            pageLabels(um),
		PageType:              pageType(um),
		RedirectedFromPath:    redirectPath,
	}

	page.Details = d
	main, ok := d.(*MainDetails)
	if ok {
		page.MetaDescription = metaDescription(strconv.Itoa(main.ImportedByCount))
	}
	s.servePage(ctx, w, tabSettings.TemplateName, page)
	return nil
}

func latestMinorClass(ctx context.Context, version string, latest internal.LatestInfo) string {
	c := "DetailsHeader-badge"
	switch {
	case latest.MinorVersion == "":
		c += "--unknown"
	case latest.MinorVersion == version && !latest.UnitExistsAtMinor && experiment.IsActive(ctx, internal.ExperimentNotAtLatest):
		c += "--notAtLatest"
	case latest.MinorVersion == version:
		c += "--latest"
	default:
		c += "--goToLatest"
	}
	return c
}

// metaDescription uses a safehtml escape hatch to build HTML used
// to render the <meta name="Description"> for unit pages as a
// workaround for https://github.com/google/safehtml/issues/6.
func metaDescription(synopsis string) safehtml.HTML {
	if synopsis == "" {
		return safehtml.HTML{}
	}
	return safehtml.HTMLConcat(
		uncheckedconversions.HTMLFromStringKnownToSatisfyTypeContract(`<meta name="Description" content="`),
		safehtml.HTMLEscaped(synopsis),
		uncheckedconversions.HTMLFromStringKnownToSatisfyTypeContract(`">`),
	)
}

// isValidTabForUnit reports whether the tab is valid for the given unit.
// It is assumed that tab is a key in unitTabLookup.
func isValidTabForUnit(tab string, um *internal.UnitMeta) bool {
	if tab == tabLicenses && !um.IsRedistributable {
		return false
	}
	if !um.IsPackage() && (tab == tabImports || tab == tabImportedBy) {
		return false
	}
	return true
}

// constructUnitURL returns a URL path that refers to the given unit at the requested
// version. If requestedVersion is "latest", then the resulting path has no
// version; otherwise, it has requestedVersion.
func constructUnitURL(fullPath, modulePath, requestedVersion string) string {
	if requestedVersion == internal.LatestVersion {
		return "/" + fullPath
	}
	v := linkVersion(requestedVersion, modulePath)
	if fullPath == modulePath || modulePath == stdlib.ModulePath {
		return fmt.Sprintf("/%s@%s", fullPath, v)
	}
	return fmt.Sprintf("/%s@%s/%s", modulePath, v, strings.TrimPrefix(fullPath, modulePath+"/"))
}

// canonicalURLPath constructs a URL path to the unit that always includes the
// resolved version.
func canonicalURLPath(um *internal.UnitMeta) string {
	return constructUnitURL(um.Path, um.ModulePath, linkVersion(um.Version, um.ModulePath))
}
<|endoftext|>


<commit_before>// Package grerr (graw-error) defines error values graw can encounter so that
// bots can understand and define scenarios for known error types.
package grerr

import (
	"fmt"
)

var (
	// PermissionDenied usually occurs because a bot provided invalid
	// credentials or tried to access a private subreddit.
	PermissionDenied = fmt.Errorf("reddit returned 403; permission denied")

	// Busy usually occurs when Reddit is under heavy load and does not have
	// anything to do with the running bot.
	Busy = fmt.Errorf("reddit returned 503; it is busy right now")

	// RateLimit means Reddit has received too many requests in the allowed
	// interval for the bot's user agent. This usually means the bot did not
	// correctly define or is using a default user agent, or is running on
	// multiple instances of graw, because graw automatically enforces rule
	// abiding rate limits on all bots.
	RateLimit = fmt.Errorf("reddit returned 429; too many requests")
)
<commit_msg>Remove grerr package.<commit_after><|endoftext|>


<commit_before>package engine

import (
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchevents"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/adapter"
	"github.com/stitchfix/flotilla-os/queue"
	"github.com/stitchfix/flotilla-os/state"
	"strings"
)

//
// ECSExecutionEngine submits runs to ecs
//
type ECSExecutionEngine struct {
	ecsClient  ecsServiceClient
	cwClient   cloudwatchServiceClient
	sqsClient  sqsClient
	adapter    adapter.ECSAdapter
	qm         queue.Manager
	statusQurl string
}

type ecsServiceClient interface {
	RunTask(input *ecs.RunTaskInput) (*ecs.RunTaskOutput, error)
	StopTask(input *ecs.StopTaskInput) (*ecs.StopTaskOutput, error)
	DeregisterTaskDefinition(input *ecs.DeregisterTaskDefinitionInput) (*ecs.DeregisterTaskDefinitionOutput, error)
	RegisterTaskDefinition(input *ecs.RegisterTaskDefinitionInput) (*ecs.RegisterTaskDefinitionOutput, error)
	DescribeContainerInstances(input *ecs.DescribeContainerInstancesInput) (*ecs.DescribeContainerInstancesOutput, error)
}

type cloudwatchServiceClient interface {
	PutRule(input *cloudwatchevents.PutRuleInput) (*cloudwatchevents.PutRuleOutput, error)
	PutTargets(input *cloudwatchevents.PutTargetsInput) (*cloudwatchevents.PutTargetsOutput, error)
	ListRuleNamesByTarget(input *cloudwatchevents.ListRuleNamesByTargetInput) (*cloudwatchevents.ListRuleNamesByTargetOutput, error)
}

type sqsClient interface {
	GetQueueAttributes(input *sqs.GetQueueAttributesInput) (*sqs.GetQueueAttributesOutput, error)
	SetQueueAttributes(input *sqs.SetQueueAttributesInput) (*sqs.SetQueueAttributesOutput, error)
}

type ecsUpdate struct {
	Detail ecs.Task `json:"detail"`
}

//
// Initialize configures the ECSExecutionEngine and initializes internal clients
//
func (ee *ECSExecutionEngine) Initialize(conf config.Config) error {
	if !conf.IsSet("aws_default_region") {
		return errors.Errorf("ECSExecutionEngine needs [aws_default_region] set in config")
	}

	if !conf.IsSet("queue.status") {
		return errors.Errorf("ECSExecutionEngine needs [queue.status] set in config")
	}

	if !conf.IsSet("queue.status_rule") {
		return errors.Errorf("ECSExecutionEngine needs [queue.status_rule] set in config")
	}

	var (
		adpt adapter.ECSAdapter
		err  error
	)

	flotillaMode := conf.GetString("flotilla_mode")

	//
	// When mode is not test, setup and initialize all aws clients
	// - this isn't ideal; is there another way?
	//
	if flotillaMode != "test" {
		sess := session.Must(session.NewSession(&aws.Config{
			Region: aws.String(conf.GetString("aws_default_region"))}))

		ecsClient := ecs.New(sess)
		ec2Client := ec2.New(sess)

		ee.ecsClient = ecsClient
		ee.cwClient = cloudwatchevents.New(sess)
		ee.sqsClient = sqs.New(sess)
		adpt, err = adapter.NewECSAdapter(conf, ecsClient, ec2Client)
		if err != nil {
			return errors.Wrap(err, "problem initializing ECSAdapter")
		}
	}

	ee.adapter = adpt

	if ee.qm == nil {
		return errors.Errorf("no queue.Manager implementation; ECSExecutionEngine needs a queue.Manager")
	}

	//
	// Calling QurlFor creates the status queue if it does not exist
	// - this is necessary for the next step of creating an ecs
	//   task update rule in cloudwatch which routes task updates
	//   to the status queue
	//
	statusQueue := conf.GetString("queue.status")
	ee.statusQurl, err = ee.qm.QurlFor(statusQueue, false)
	if err != nil {
		return errors.Wrapf(err, "problem getting queue url for status queue with name [%s]", statusQueue)
	}

	statusRule := conf.GetString("queue.status_rule")
	return ee.createOrUpdateEventRule(statusRule, statusQueue)
}

func (ee *ECSExecutionEngine) createOrUpdateEventRule(statusRule string, statusQueue string) error {
	createUpdate, err := ee.cwClient.PutRule(&cloudwatchevents.PutRuleInput{
		Description:  aws.String("Routes ecs task status events to flotilla status queues"),
		Name:         &statusRule,
		EventPattern: aws.String(`{"source":["aws.ecs"],"detail-type":["ECS Task State Change"]}`),
	})

	if err != nil {
		return errors.Wrap(err, "problem creating ecs status event routing rule")
	}

	// Route status events to the status queue
	targetArn, err := ee.getTargetArn(ee.statusQurl)
	if err != nil {
		return errors.Wrapf(err, "problem getting target arn for [%s]", ee.statusQurl)
	}

	names, err := ee.cwClient.ListRuleNamesByTarget(&cloudwatchevents.ListRuleNamesByTargetInput{
		TargetArn: &targetArn,
	})
	if err != nil {
		return errors.Wrapf(err, "problem listing rules for target [%s]", targetArn)
	}

	targetExists := len(names.RuleNames) > 0 && *names.RuleNames[0] == statusRule
	if !targetExists {
		res, err := ee.cwClient.PutTargets(&cloudwatchevents.PutTargetsInput{
			Rule: &statusRule,
			Targets: []*cloudwatchevents.Target{
				{
					Arn: &targetArn,
					Id:  &statusQueue,
				},
			},
		})

		if err != nil {
			return errors.Wrapf(
				err, "problem adding [%s] as queue target for status rule [%s]", targetArn, statusRule)
		}

		if *res.FailedEntryCount > 0 {
			failed := res.FailedEntries[0]
			return errors.Errorf("error adding routing rule for ecs status messages [%s]", *failed.ErrorMessage)
		}
	}
	// Finally, add permissions to target queue
	return ee.setTargetPermissions(*createUpdate.RuleArn, targetArn)
}

func (ee *ECSExecutionEngine) getTargetArn(qurl string) (string, error) {
	var arn string
	res, err := ee.sqsClient.GetQueueAttributes(&sqs.GetQueueAttributesInput{
		QueueUrl: &qurl,
		AttributeNames: []*string{
			aws.String("QueueArn"),
		},
	})
	if err != nil {
		return arn, errors.Wrapf(err, "problem getting attribute QueueArn for sqs queue with url [%s]", qurl)
	}
	if res.Attributes["QueueArn"] != nil {
		return *res.Attributes["QueueArn"], nil
	}
	return arn, errors.Errorf("couldn't get queue arn")
}

func (ee *ECSExecutionEngine) setTargetPermissions(sourceArn string, targetArn string) error {
	policyDoc := fmt.Sprintf(`{
		"Version":"2012-10-17",
		"Id":"flotilla-task-status-updates-to-sqs",
		"Statement": [{
			"Sid": "flotilla-task-status-updates-to-sqs-sid",
			"Effect": "Allow",
			"Principal": {
				"AWS": "*"
			},
			"Action": "sqs:SendMessage",
			"Resource": "%s",
			"Condition": {
				"ArnEquals": {
					"aws:SourceArn": "%s"
				}
			}
		}]
	}`, targetArn, sourceArn)

	// Check first
	res, err := ee.sqsClient.GetQueueAttributes(&sqs.GetQueueAttributesInput{
		QueueUrl: &ee.statusQurl,
	})
	if err != nil {
		return errors.Wrapf(err, "problem getting queue attributes for sqs queue [%s]", ee.statusQurl)
	}

	if policy, ok := res.Attributes["Policy"]; ok && *policy == policyDoc {
		return nil
	}

	if _, err := ee.sqsClient.SetQueueAttributes(&sqs.SetQueueAttributesInput{
		Attributes: map[string]*string{
			"Policy": &policyDoc,
		},
		QueueUrl: &ee.statusQurl,
	}); err != nil {
		return errors.Wrapf(
			err, "problem setting permissions allowing [%s] to send events to [%s]", sourceArn, targetArn)
	}
	return nil
}

//
// PollStatus pops status updates from the status queue using the QueueManager
//
func (ee *ECSExecutionEngine) PollStatus() (RunReceipt, error) {
	var (
		receipt RunReceipt
		update  ecsUpdate
		err     error
	)

	rawReceipt, err := ee.qm.ReceiveStatus(ee.statusQurl)
	if err != nil {
		return receipt, errors.Wrapf(err, "problem getting status from [%s]", ee.statusQurl)
	}

	//
	// If we receive an update that is empty, don't try to deserialize it
	//
	if rawReceipt.StatusUpdate != nil {
		err = json.Unmarshal([]byte(*rawReceipt.StatusUpdate), &update)
		if err != nil {
			return receipt, errors.Wrapf(err, "unable to parse status update with json [%s]", *rawReceipt.StatusUpdate)
		}
		adapted := ee.adapter.AdaptTask(update.Detail)
		receipt.Run = &adapted
	}

	receipt.Done = rawReceipt.Done
	return receipt, nil
}

//
// PollRuns receives -at most- one run per queue that is pending execution
//
func (ee *ECSExecutionEngine) PollRuns() ([]RunReceipt, error) {
	queues, err := ee.qm.List()
	if err != nil {
		return nil, errors.Wrap(err, "problem listing queues to poll")
	}

	var runs []RunReceipt
	for _, qurl := range queues {
		//
		// Get new queued Run
		//
		runReceipt, err := ee.qm.ReceiveRun(qurl)

		if err != nil {
			return runs, errors.Wrapf(err, "problem receiving run from queue url [%s]", qurl)
		}

		if runReceipt.Run == nil {
			continue
		}

		runs = append(runs, RunReceipt{runReceipt})
	}
	return runs, nil
}

//
// Enqueue pushes a run onto the queue using the QueueManager
//
func (ee *ECSExecutionEngine) Enqueue(run state.Run) error {
	// Get qurl
	qurl, err := ee.qm.QurlFor(run.ClusterName, true)
	if err != nil {
		return errors.Wrapf(err, "problem getting queue url for [%s]", run.ClusterName)
	}

	// Queue run
	if err = ee.qm.Enqueue(qurl, run); err != nil {
		return errors.Wrapf(err, "problem enqueing run [%s] to queue [%s]", run.RunID, qurl)
	}
	return nil
}

//
// Execute takes a pre-configured run and definition and submits them for execution
// to AWS ECS
//
func (ee *ECSExecutionEngine) Execute(definition state.Definition, run state.Run) (state.Run, bool, error) {
	var executed state.Run
	rti := ee.toRunTaskInput(definition, run)
	result, err := ee.ecsClient.RunTask(&rti)
	if err != nil {
		retryable := false
		if aerr, ok := err.(awserr.Error); ok {
			if aerr.Code() == ecs.ErrCodeInvalidParameterException {
				if strings.Contains(aerr.Message(), "no container instances") {
					retryable = true
				}
			}
		}
		return executed, retryable, errors.Wrapf(err, "problem executing run [%s]", run.RunID)
	}
	if len(result.Failures) != 0 {
		msg := make([]string, len(result.Failures))
		for i, failure := range result.Failures {
			msg[i] = *failure.Reason
		}
		//
		// Retry these, they are very rare;
		// our upfront validation catches the obvious image and cluster resources
		//
		// IMPORTANT - log these messages
		//
		return executed, true, errors.Errorf("ERRORS: %s", strings.Join(msg, "\n"))
	}

	return ee.translateTask(*result.Tasks[0]), false, nil
}

//
// Terminate takes a valid run and stops it
//
func (ee *ECSExecutionEngine) Terminate(run state.Run) error {
	if _, err := ee.ecsClient.StopTask(&ecs.StopTaskInput{
		Cluster: &run.ClusterName,
		Task:    &run.TaskArn,
	}); err != nil {
		return errors.Wrapf(err, "problem stopping run [%s] with task arn [%s]", run.RunID, run.TaskArn)
	}
	return nil
}

//
// Define creates or updates a task definition with ecs
//
func (ee *ECSExecutionEngine) Define(definition state.Definition) (state.Definition, error) {
	rti := ee.adapter.AdaptDefinition(definition)
	result, err := ee.ecsClient.RegisterTaskDefinition(&rti)
	if err != nil {
		return state.Definition{}, errors.Wrapf(
			err, "problem registering definition [%s] with ecs", definition.DefinitionID)
	}

	//
	// We wrap the command of a definition before registering it with
	// ECS. What this means is that the command returned from registration
	// contains only the *wrapped* version. Reversing the wrapping process
	// using string parsing is brittle. Instead, we make the following
	// assumptions:
	//
	// * Definitions are pre-validated using their `IsValid` method meaning
	//   they must have a non-empty user command
	// * Registering a task definition with ECS does not mutate the user command
	// ** The command acknowledged by ECS is -exactly- the wrapped version
	//    of the command contained in the passed in Definition
	// Hence it should be safe to simply attach the passed in definition's
	// Command field to the output.
	//
	defined := ee.adapter.AdaptTaskDef(*result.TaskDefinition)
	defined.Command = definition.Command
	return defined, nil
}

//
// Deregister deregisters the task definition from ecs
//
func (ee *ECSExecutionEngine) Deregister(definition state.Definition) error {
	if _, err := ee.ecsClient.DeregisterTaskDefinition(&ecs.DeregisterTaskDefinitionInput{
		TaskDefinition: &definition.Arn,
	}); err != nil {
		return errors.Wrapf(err, "problem deregistering definition [%s] with ecs", definition.DefinitionID)
	}
	return nil
}

func (ee *ECSExecutionEngine) toRunTaskInput(definition state.Definition, run state.Run) ecs.RunTaskInput {
	return ee.adapter.AdaptRun(definition, run)
}

func (ee *ECSExecutionEngine) translateTask(task ecs.Task) state.Run {
	return ee.adapter.AdaptTask(task)
}
<commit_msg>lower case "No Container Instances" InvalidParameterException when submitting tasks with ecs engine to handle any case (#119)<commit_after>package engine

import (
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchevents"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/adapter"
	"github.com/stitchfix/flotilla-os/queue"
	"github.com/stitchfix/flotilla-os/state"
	"strings"
)

//
// ECSExecutionEngine submits runs to ecs
//
type ECSExecutionEngine struct {
	ecsClient  ecsServiceClient
	cwClient   cloudwatchServiceClient
	sqsClient  sqsClient
	adapter    adapter.ECSAdapter
	qm         queue.Manager
	statusQurl string
}

type ecsServiceClient interface {
	RunTask(input *ecs.RunTaskInput) (*ecs.RunTaskOutput, error)
	StopTask(input *ecs.StopTaskInput) (*ecs.StopTaskOutput, error)
	DeregisterTaskDefinition(input *ecs.DeregisterTaskDefinitionInput) (*ecs.DeregisterTaskDefinitionOutput, error)
	RegisterTaskDefinition(input *ecs.RegisterTaskDefinitionInput) (*ecs.RegisterTaskDefinitionOutput, error)
	DescribeContainerInstances(input *ecs.DescribeContainerInstancesInput) (*ecs.DescribeContainerInstancesOutput, error)
}

type cloudwatchServiceClient interface {
	PutRule(input *cloudwatchevents.PutRuleInput) (*cloudwatchevents.PutRuleOutput, error)
	PutTargets(input *cloudwatchevents.PutTargetsInput) (*cloudwatchevents.PutTargetsOutput, error)
	ListRuleNamesByTarget(input *cloudwatchevents.ListRuleNamesByTargetInput) (*cloudwatchevents.ListRuleNamesByTargetOutput, error)
}

type sqsClient interface {
	GetQueueAttributes(input *sqs.GetQueueAttributesInput) (*sqs.GetQueueAttributesOutput, error)
	SetQueueAttributes(input *sqs.SetQueueAttributesInput) (*sqs.SetQueueAttributesOutput, error)
}

type ecsUpdate struct {
	Detail ecs.Task `json:"detail"`
}

//
// Initialize configures the ECSExecutionEngine and initializes internal clients
//
func (ee *ECSExecutionEngine) Initialize(conf config.Config) error {
	if !conf.IsSet("aws_default_region") {
		return errors.Errorf("ECSExecutionEngine needs [aws_default_region] set in config")
	}

	if !conf.IsSet("queue.status") {
		return errors.Errorf("ECSExecutionEngine needs [queue.status] set in config")
	}

	if !conf.IsSet("queue.status_rule") {
		return errors.Errorf("ECSExecutionEngine needs [queue.status_rule] set in config")
	}

	var (
		adpt adapter.ECSAdapter
		err  error
	)

	flotillaMode := conf.GetString("flotilla_mode")

	//
	// When mode is not test, setup and initialize all aws clients
	// - this isn't ideal; is there another way?
	//
	if flotillaMode != "test" {
		sess := session.Must(session.NewSession(&aws.Config{
			Region: aws.String(conf.GetString("aws_default_region"))}))

		ecsClient := ecs.New(sess)
		ec2Client := ec2.New(sess)

		ee.ecsClient = ecsClient
		ee.cwClient = cloudwatchevents.New(sess)
		ee.sqsClient = sqs.New(sess)
		adpt, err = adapter.NewECSAdapter(conf, ecsClient, ec2Client)
		if err != nil {
			return errors.Wrap(err, "problem initializing ECSAdapter")
		}
	}

	ee.adapter = adpt

	if ee.qm == nil {
		return errors.Errorf("no queue.Manager implementation; ECSExecutionEngine needs a queue.Manager")
	}

	//
	// Calling QurlFor creates the status queue if it does not exist
	// - this is necessary for the next step of creating an ecs
	//   task update rule in cloudwatch which routes task updates
	//   to the status queue
	//
	statusQueue := conf.GetString("queue.status")
	ee.statusQurl, err = ee.qm.QurlFor(statusQueue, false)
	if err != nil {
		return errors.Wrapf(err, "problem getting queue url for status queue with name [%s]", statusQueue)
	}

	statusRule := conf.GetString("queue.status_rule")
	return ee.createOrUpdateEventRule(statusRule, statusQueue)
}

func (ee *ECSExecutionEngine) createOrUpdateEventRule(statusRule string, statusQueue string) error {
	createUpdate, err := ee.cwClient.PutRule(&cloudwatchevents.PutRuleInput{
		Description:  aws.String("Routes ecs task status events to flotilla status queues"),
		Name:         &statusRule,
		EventPattern: aws.String(`{"source":["aws.ecs"],"detail-type":["ECS Task State Change"]}`),
	})

	if err != nil {
		return errors.Wrap(err, "problem creating ecs status event routing rule")
	}

	// Route status events to the status queue
	targetArn, err := ee.getTargetArn(ee.statusQurl)
	if err != nil {
		return errors.Wrapf(err, "problem getting target arn for [%s]", ee.statusQurl)
	}

	names, err := ee.cwClient.ListRuleNamesByTarget(&cloudwatchevents.ListRuleNamesByTargetInput{
		TargetArn: &targetArn,
	})
	if err != nil {
		return errors.Wrapf(err, "problem listing rules for target [%s]", targetArn)
	}

	targetExists := len(names.RuleNames) > 0 && *names.RuleNames[0] == statusRule
	if !targetExists {
		res, err := ee.cwClient.PutTargets(&cloudwatchevents.PutTargetsInput{
			Rule: &statusRule,
			Targets: []*cloudwatchevents.Target{
				{
					Arn: &targetArn,
					Id:  &statusQueue,
				},
			},
		})

		if err != nil {
			return errors.Wrapf(
				err, "problem adding [%s] as queue target for status rule [%s]", targetArn, statusRule)
		}

		if *res.FailedEntryCount > 0 {
			failed := res.FailedEntries[0]
			return errors.Errorf("error adding routing rule for ecs status messages [%s]", *failed.ErrorMessage)
		}
	}
	// Finally, add permissions to target queue
	return ee.setTargetPermissions(*createUpdate.RuleArn, targetArn)
}

func (ee *ECSExecutionEngine) getTargetArn(qurl string) (string, error) {
	var arn string
	res, err := ee.sqsClient.GetQueueAttributes(&sqs.GetQueueAttributesInput{
		QueueUrl: &qurl,
		AttributeNames: []*string{
			aws.String("QueueArn"),
		},
	})
	if err != nil {
		return arn, errors.Wrapf(err, "problem getting attribute QueueArn for sqs queue with url [%s]", qurl)
	}
	if res.Attributes["QueueArn"] != nil {
		return *res.Attributes["QueueArn"], nil
	}
	return arn, errors.Errorf("couldn't get queue arn")
}

func (ee *ECSExecutionEngine) setTargetPermissions(sourceArn string, targetArn string) error {
	policyDoc := fmt.Sprintf(`{
		"Version":"2012-10-17",
		"Id":"flotilla-task-status-updates-to-sqs",
		"Statement": [{
			"Sid": "flotilla-task-status-updates-to-sqs-sid",
			"Effect": "Allow",
			"Principal": {
				"AWS": "*"
			},
			"Action": "sqs:SendMessage",
			"Resource": "%s",
			"Condition": {
				"ArnEquals": {
					"aws:SourceArn": "%s"
				}
			}
		}]
	}`, targetArn, sourceArn)

	// Check first
	res, err := ee.sqsClient.GetQueueAttributes(&sqs.GetQueueAttributesInput{
		QueueUrl: &ee.statusQurl,
	})
	if err != nil {
		return errors.Wrapf(err, "problem getting queue attributes for sqs queue [%s]", ee.statusQurl)
	}

	if policy, ok := res.Attributes["Policy"]; ok && *policy == policyDoc {
		return nil
	}

	if _, err := ee.sqsClient.SetQueueAttributes(&sqs.SetQueueAttributesInput{
		Attributes: map[string]*string{
			"Policy": &policyDoc,
		},
		QueueUrl: &ee.statusQurl,
	}); err != nil {
		return errors.Wrapf(
			err, "problem setting permissions allowing [%s] to send events to [%s]", sourceArn, targetArn)
	}
	return nil
}

//
// PollStatus pops status updates from the status queue using the QueueManager
//
func (ee *ECSExecutionEngine) PollStatus() (RunReceipt, error) {
	var (
		receipt RunReceipt
		update  ecsUpdate
		err     error
	)

	rawReceipt, err := ee.qm.ReceiveStatus(ee.statusQurl)
	if err != nil {
		return receipt, errors.Wrapf(err, "problem getting status from [%s]", ee.statusQurl)
	}

	//
	// If we receive an update that is empty, don't try to deserialize it
	//
	if rawReceipt.StatusUpdate != nil {
		err = json.Unmarshal([]byte(*rawReceipt.StatusUpdate), &update)
		if err != nil {
			return receipt, errors.Wrapf(err, "unable to parse status update with json [%s]", *rawReceipt.StatusUpdate)
		}
		adapted := ee.adapter.AdaptTask(update.Detail)
		receipt.Run = &adapted
	}

	receipt.Done = rawReceipt.Done
	return receipt, nil
}

//
// PollRuns receives -at most- one run per queue that is pending execution
//
func (ee *ECSExecutionEngine) PollRuns() ([]RunReceipt, error) {
	queues, err := ee.qm.List()
	if err != nil {
		return nil, errors.Wrap(err, "problem listing queues to poll")
	}

	var runs []RunReceipt
	for _, qurl := range queues {
		//
		// Get new queued Run
		//
		runReceipt, err := ee.qm.ReceiveRun(qurl)

		if err != nil {
			return runs, errors.Wrapf(err, "problem receiving run from queue url [%s]", qurl)
		}

		if runReceipt.Run == nil {
			continue
		}

		runs = append(runs, RunReceipt{runReceipt})
	}
	return runs, nil
}

//
// Enqueue pushes a run onto the queue using the QueueManager
//
func (ee *ECSExecutionEngine) Enqueue(run state.Run) error {
	// Get qurl
	qurl, err := ee.qm.QurlFor(run.ClusterName, true)
	if err != nil {
		return errors.Wrapf(err, "problem getting queue url for [%s]", run.ClusterName)
	}

	// Queue run
	if err = ee.qm.Enqueue(qurl, run); err != nil {
		return errors.Wrapf(err, "problem enqueing run [%s] to queue [%s]", run.RunID, qurl)
	}
	return nil
}

//
// Execute takes a pre-configured run and definition and submits them for execution
// to AWS ECS
//
func (ee *ECSExecutionEngine) Execute(definition state.Definition, run state.Run) (state.Run, bool, error) {
	var executed state.Run
	rti := ee.toRunTaskInput(definition, run)
	result, err := ee.ecsClient.RunTask(&rti)
	if err != nil {
		retryable := false
		if aerr, ok := err.(awserr.Error); ok {
			if aerr.Code() == ecs.ErrCodeInvalidParameterException {
				if strings.Contains(strings.ToLower(aerr.Message()), "no container instances") {
					retryable = true
				}
			}
		}
		return executed, retryable, errors.Wrapf(err, "problem executing run [%s]", run.RunID)
	}
	if len(result.Failures) != 0 {
		msg := make([]string, len(result.Failures))
		for i, failure := range result.Failures {
			msg[i] = *failure.Reason
		}
		//
		// Retry these, they are very rare;
		// our upfront validation catches the obvious image and cluster resources
		//
		// IMPORTANT - log these messages
		//
		return executed, true, errors.Errorf("ERRORS: %s", strings.Join(msg, "\n"))
	}

	return ee.translateTask(*result.Tasks[0]), false, nil
}

//
// Terminate takes a valid run and stops it
//
func (ee *ECSExecutionEngine) Terminate(run state.Run) error {
	if _, err := ee.ecsClient.StopTask(&ecs.StopTaskInput{
		Cluster: &run.ClusterName,
		Task:    &run.TaskArn,
	}); err != nil {
		return errors.Wrapf(err, "problem stopping run [%s] with task arn [%s]", run.RunID, run.TaskArn)
	}
	return nil
}

//
// Define creates or updates a task definition with ecs
//
func (ee *ECSExecutionEngine) Define(definition state.Definition) (state.Definition, error) {
	rti := ee.adapter.AdaptDefinition(definition)
	result, err := ee.ecsClient.RegisterTaskDefinition(&rti)
	if err != nil {
		return state.Definition{}, errors.Wrapf(
			err, "problem registering definition [%s] with ecs", definition.DefinitionID)
	}

	//
	// We wrap the command of a definition before registering it with
	// ECS. What this means is that the command returned from registration
	// contains only the *wrapped* version. Reversing the wrapping process
	// using string parsing is brittle. Instead, we make the following
	// assumptions:
	//
	// * Definitions are pre-validated using their `IsValid` method meaning
	//   they must have a non-empty user command
	// * Registering a task definition with ECS does not mutate the user command
	// ** The command acknowledged by ECS is -exactly- the wrapped version
	//    of the command contained in the passed in Definition
	// Hence it should be safe to simply attach the passed in definition's
	// Command field to the output.
	//
	defined := ee.adapter.AdaptTaskDef(*result.TaskDefinition)
	defined.Command = definition.Command
	return defined, nil
}

//
// Deregister deregisters the task definition from ecs
//
func (ee *ECSExecutionEngine) Deregister(definition state.Definition) error {
	if _, err := ee.ecsClient.DeregisterTaskDefinition(&ecs.DeregisterTaskDefinitionInput{
		TaskDefinition: &definition.Arn,
	}); err != nil {
		return errors.Wrapf(err, "problem deregistering definition [%s] with ecs", definition.DefinitionID)
	}
	return nil
}

func (ee *ECSExecutionEngine) toRunTaskInput(definition state.Definition, run state.Run) ecs.RunTaskInput {
	return ee.adapter.AdaptRun(definition, run)
}

func (ee *ECSExecutionEngine) translateTask(task ecs.Task) state.Run {
	return ee.adapter.AdaptTask(task)
}
<|endoftext|>


<commit_before>package ssh

const (
	// This is a constant representing a script to install and uninstall public
	// key in remote hosts.
	DefaultPublicKeyInstallScript = `
#!/bin/bash
#
# This is a default script which installs or uninstalls an RSA public key to/from
# authoried_keys file in a typical linux machine.
#
# If the platform differs or if the binaries used in this script are not available
# in target machine, use the 'install_script' parameter with 'roles/' endpoint to
# register a custom script (applicable for Dynamic type only).
#
# Vault server runs this script on the target machine with the following params:
#
# $1:INSTALL_OPTION: "install" or "uninstall"
#
# $2:PUBLIC_KEY_FILE: File name containing public key to be installed. Vault server
# uses UUID as name to avoid collisions with public keys generated for other requests.
#
# $3:AUTH_KEYS_FILE: Absolute path of the authorized_keys file.
# Currently, vault uses /home/<username>/.ssh/authorized_keys as the path.
#
# [Note: This script will be run by Vault using the registered admin username.
# Notice that some commands below are run as 'sudo'. For graceful execution of
# this script there should not be any password prompts. So, disable password
# prompt for the admin username registered with Vault.

set -e

# Storing arguments into variables, to increase readability of the script.
INSTALL_OPTION=$1
PUBLIC_KEY_FILE=$2
AUTH_KEYS_FILE=$3

# Delete the public key file and the temporary file
function cleanup
{
	rm -f "$PUBLIC_KEY_FILE" temp_$PUBLIC_KEY_FILE
}

# 'cleanup' will be called if the script ends or if any command fails.
trap cleanup EXIT

# Return if the option is anything other than 'install' or 'uninstall'.
if [ "$INSTALL_OPTION" != "install" ] && [ "$INSTALL_OPTION" != "uninstall" ]; then
	exit 1
fi

# Create the .ssh directory and authorized_keys file if it does not exist
SSH_DIR=$(dirname $AUTH_KEYS_FILE)
sudo mkdir -p "$SSH_DIR"
sudo touch "$AUTH_KEYS_FILE"

# Remove the key from authorized_keys file if it is already present.
# This step is common for both install and uninstall.  Note that grep's
# return code is ignored, thus if grep fails all keys will be removed
# rather than none and it fails secure
sudo grep -vFf "$PUBLIC_KEY_FILE" "$AUTH_KEYS_FILE" > temp_$PUBLIC_KEY_FILE || true
cat temp_$PUBLIC_KEY_FILE | sudo tee "$AUTH_KEYS_FILE"

# Append the new public key to authorized_keys file
if [ "$INSTALL_OPTION" == "install" ]; then
	cat "$PUBLIC_KEY_FILE" | sudo tee --append "$AUTH_KEYS_FILE"
fi
`
)
<commit_msg>Fix typo and remove trailing whitespace. (#2074)<commit_after>package ssh

const (
	// This is a constant representing a script to install and uninstall public
	// key in remote hosts.
	DefaultPublicKeyInstallScript = `
#!/bin/bash
#
# This is a default script which installs or uninstalls an RSA public key to/from
# authorized_keys file in a typical linux machine.
#
# If the platform differs or if the binaries used in this script are not available
# in target machine, use the 'install_script' parameter with 'roles/' endpoint to
# register a custom script (applicable for Dynamic type only).
#
# Vault server runs this script on the target machine with the following params:
#
# $1:INSTALL_OPTION: "install" or "uninstall"
#
# $2:PUBLIC_KEY_FILE: File name containing public key to be installed. Vault server
# uses UUID as name to avoid collisions with public keys generated for other requests.
#
# $3:AUTH_KEYS_FILE: Absolute path of the authorized_keys file.
# Currently, vault uses /home/<username>/.ssh/authorized_keys as the path.
#
# [Note: This script will be run by Vault using the registered admin username.
# Notice that some commands below are run as 'sudo'. For graceful execution of
# this script there should not be any password prompts. So, disable password
# prompt for the admin username registered with Vault.

set -e

# Storing arguments into variables, to increase readability of the script.
INSTALL_OPTION=$1
PUBLIC_KEY_FILE=$2
AUTH_KEYS_FILE=$3

# Delete the public key file and the temporary file
function cleanup
{
	rm -f "$PUBLIC_KEY_FILE" temp_$PUBLIC_KEY_FILE
}

# 'cleanup' will be called if the script ends or if any command fails.
trap cleanup EXIT

# Return if the option is anything other than 'install' or 'uninstall'.
if [ "$INSTALL_OPTION" != "install" ] && [ "$INSTALL_OPTION" != "uninstall" ]; then
	exit 1
fi

# Create the .ssh directory and authorized_keys file if it does not exist
SSH_DIR=$(dirname $AUTH_KEYS_FILE)
sudo mkdir -p "$SSH_DIR"
sudo touch "$AUTH_KEYS_FILE"

# Remove the key from authorized_keys file if it is already present.
# This step is common for both install and uninstall.  Note that grep's
# return code is ignored, thus if grep fails all keys will be removed
# rather than none and it fails secure
sudo grep -vFf "$PUBLIC_KEY_FILE" "$AUTH_KEYS_FILE" > temp_$PUBLIC_KEY_FILE || true
cat temp_$PUBLIC_KEY_FILE | sudo tee "$AUTH_KEYS_FILE"

# Append the new public key to authorized_keys file
if [ "$INSTALL_OPTION" == "install" ]; then
	cat "$PUBLIC_KEY_FILE" | sudo tee --append "$AUTH_KEYS_FILE"
fi
`
)
<|endoftext|>


<commit_before>//  Copyright (c) 2019 Cisco and/or its affiliates.
//
//  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 commands

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"path"
	"reflect"
	"strings"
	"time"

	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"go.ligato.io/cn-infra/v2/logging"
	"go.ligato.io/cn-infra/v2/servicelabel"

	agentcli "go.ligato.io/vpp-agent/v3/cmd/agentctl/cli"
	"go.ligato.io/vpp-agent/v3/pkg/models"
)

const (
	defaultTimeout = time.Second * 30
	defaultTxOps   = 128
)

func NewImportCommand(cli agentcli.Cli) *cobra.Command {
	var (
		opts    ImportOptions
		timeout uint
	)
	cmd := &cobra.Command{
		Use:   "import FILE",
		Args:  cobra.ExactArgs(1),
		Short: "Import config data from file",
		Long: `Import config data from file into Etcd or via gRPC.
FILE FORMAT
  Contents of the import file must contain single key-value pair per line:

	<key1> <value1>
    <key2> <value2>
    ...
    <keyN> <valueN>

    NOTE: Empty lines and lines starting with '#' are ignored.

  Sample file:
  	config/vpp/v2/interfaces/loop1 {"name":"loop1","type":"SOFTWARE_LOOPBACK"}
  	config/vpp/l2/v2/bridge-domain/bd1 {"name":"bd1"}

KEY FORMAT
    Keys can be defined in two ways:

    - Full  - /vnf-agent/vpp1/config/vpp/v2/interfaces/iface1
    - Short - config/vpp/v2/interfaces/iface1

    When using short keys, import will use configured microservice label (e.g. --service-label flag).`,
		Example: `
  Import data into Etcd:
	{{.CommandPath}} input.txt

  Import data directly into agent via gRPC::
	{{.CommandPath}} --grpc input.txt
`,
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.InputFile = args[0]
			opts.Timeout = time.Duration(timeout)
			return RunImport(cli, opts)
		},
	}
	flags := cmd.Flags()
	flags.UintVar(&opts.TxOps, "txops", defaultTxOps, "Number of ops per transaction")
	flags.DurationVarP(&opts.Timeout, "time", "t", defaultTimeout, "Timeout to wait for server response")
	flags.BoolVar(&opts.ViaGrpc, "grpc", false, "Import config directly to agent via gRPC")
	return cmd
}

type ImportOptions struct {
	InputFile string
	TxOps     uint
	Timeout   time.Duration
	ViaGrpc   bool
}

func RunImport(cli agentcli.Cli, opts ImportOptions) error {
	keyVals, err := parseImportFile(opts.InputFile)
	if err != nil {
		return fmt.Errorf("parsing import data failed: %v", err)
	}
	fmt.Printf("importing %d key-value pairs\n", len(keyVals))

	if opts.ViaGrpc {
		// Set up a connection to the server.
		c, err := cli.Client().GenericClient()
		if err != nil {
			return err
		}
		req := c.ChangeRequest()
		for _, keyVal := range keyVals {
			fmt.Printf(" - %s\n", keyVal.Key)
			req.Update(keyVal.Val)
		}
		fmt.Printf("sending via gRPC\n")

		ctx, cancel := context.WithTimeout(context.Background(), opts.Timeout)
		defer cancel()

		if err := req.Send(ctx); err != nil {
			return fmt.Errorf("send failed: %v", err)
		}
	} else {
		c, err := cli.Client().KVDBClient()
		if err != nil {
			return fmt.Errorf("KVDB error: %v", err)
		}
		db := c.ProtoBroker()
		var txn = db.NewTxn()
		ops := 0
		for i := 0; i < len(keyVals); i++ {
			keyVal := keyVals[i]
			key, err := c.CompleteFullKey(keyVal.Key)
			if err != nil {
				return fmt.Errorf("key processing failed: %v", err)
			}
			fmt.Printf(" - %s\n", key)
			txn.Put(key, keyVal.Val)
			ops++

			if ops == int(opts.TxOps) || i+1 == len(keyVals) {
				fmt.Printf("commiting tx with %d ops\n", ops)
				ctx, cancel := context.WithTimeout(context.Background(), opts.Timeout)
				err = txn.Commit(ctx)
				cancel()
				if err != nil {
					return fmt.Errorf("commit failed: %v", err)
				}
				ops = 0
				txn = db.NewTxn()
			}
		}
	}

	logging.Debug("OK")
	return nil
}

type keyVal struct {
	Key string
	Val proto.Message
}

func parseImportFile(importFile string) (keyVals []keyVal, err error) {
	b, err := ioutil.ReadFile(importFile)
	if err != nil {
		return nil, fmt.Errorf("reading input file failed: %v", err)
	}
	// parse lines
	lines := bytes.Split(b, []byte("\n"))
	for _, l := range lines {
		line := bytes.TrimSpace(l)
		if bytes.HasPrefix(line, []byte("#")) {
			continue
		}
		parts := bytes.SplitN(line, []byte(" "), 2)
		if len(parts) < 2 {
			continue
		}
		key := string(parts[0])
		data := string(parts[1])
		if key == "" || data == "" {
			continue
		}
		logrus.Debugf("parse line: %s %s\n", key, data)

		//key = completeFullKey(key)
		val, err := unmarshalKeyVal(key, data)
		if err != nil {
			return nil, fmt.Errorf("decoding value failed: %v", err)
		}
		logrus.Debugf("KEY: %s - %v\n", key, val)
		keyVals = append(keyVals, keyVal{key, val})
	}
	return
}

func unmarshalKeyVal(fullKey string, data string) (proto.Message, error) {
	key := stripAgentPrefix(fullKey)

	model, err := models.GetModelForKey(key)
	if err != nil {
		return nil, err
	}
	valueType := proto.MessageType(model.ProtoName())
	if valueType == nil {
		return nil, fmt.Errorf("unknown proto message defined for key %s", key)
	}
	value := reflect.New(valueType.Elem()).Interface().(proto.Message)

	if err := jsonpb.UnmarshalString(data, value); err != nil {
		return nil, err
	}
	return value, nil
}

func stripAgentPrefix(key string) string {
	if !strings.HasPrefix(key, servicelabel.GetAllAgentsPrefix()) {
		return key
	}
	keyParts := strings.Split(key, "/")
	if len(keyParts) < 4 || keyParts[0] != "" {
		return path.Join(keyParts[2:]...)
	}
	return path.Join(keyParts[3:]...)
}
<commit_msg>fix: Fix grpc context timeout for agentctl import command (#1718)<commit_after>//  Copyright (c) 2019 Cisco and/or its affiliates.
//
//  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 commands

import (
	"bytes"
	"context"
	"fmt"
	"io/ioutil"
	"path"
	"reflect"
	"strings"
	"time"

	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"go.ligato.io/cn-infra/v2/logging"
	"go.ligato.io/cn-infra/v2/servicelabel"

	agentcli "go.ligato.io/vpp-agent/v3/cmd/agentctl/cli"
	"go.ligato.io/vpp-agent/v3/pkg/models"
)

const (
	defaultTimeout = time.Second * 30
	defaultTxOps   = 128
)

func NewImportCommand(cli agentcli.Cli) *cobra.Command {
	var opts ImportOptions
	cmd := &cobra.Command{
		Use:   "import FILE",
		Args:  cobra.ExactArgs(1),
		Short: "Import config data from file",
		Long: `Import config data from file into Etcd or via gRPC.
FILE FORMAT
  Contents of the import file must contain single key-value pair per line:

	<key1> <value1>
    <key2> <value2>
    ...
    <keyN> <valueN>

    NOTE: Empty lines and lines starting with '#' are ignored.

  Sample file:
  	config/vpp/v2/interfaces/loop1 {"name":"loop1","type":"SOFTWARE_LOOPBACK"}
  	config/vpp/l2/v2/bridge-domain/bd1 {"name":"bd1"}

KEY FORMAT
    Keys can be defined in two ways:

    - Full  - /vnf-agent/vpp1/config/vpp/v2/interfaces/iface1
    - Short - config/vpp/v2/interfaces/iface1

    When using short keys, import will use configured microservice label (e.g. --service-label flag).`,
		Example: `
  Import data into Etcd:
	{{.CommandPath}} input.txt

  Import data directly into agent via gRPC::
	{{.CommandPath}} --grpc input.txt
`,
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.InputFile = args[0]
			return RunImport(cli, opts)
		},
	}
	flags := cmd.Flags()
	flags.UintVar(&opts.TxOps, "txops", defaultTxOps, "Number of ops per transaction")
	flags.DurationVarP(&opts.Timeout, "time", "t", defaultTimeout, "Timeout to wait for server response")
	flags.BoolVar(&opts.ViaGrpc, "grpc", false, "Import config directly to agent via gRPC")
	return cmd
}

type ImportOptions struct {
	InputFile string
	TxOps     uint
	Timeout   time.Duration
	ViaGrpc   bool
}

func RunImport(cli agentcli.Cli, opts ImportOptions) error {
	keyVals, err := parseImportFile(opts.InputFile)
	if err != nil {
		return fmt.Errorf("parsing import data failed: %v", err)
	}
	fmt.Printf("importing %d key-value pairs\n", len(keyVals))

	if opts.ViaGrpc {
		// Set up a connection to the server.
		c, err := cli.Client().GenericClient()
		if err != nil {
			return err
		}
		req := c.ChangeRequest()
		for _, keyVal := range keyVals {
			fmt.Printf(" - %s\n", keyVal.Key)
			req.Update(keyVal.Val)
		}
		fmt.Printf("sending via gRPC\n")

		ctx, cancel := context.WithTimeout(context.Background(), opts.Timeout)
		defer cancel()

		if err := req.Send(ctx); err != nil {
			return fmt.Errorf("send failed: %v", err)
		}
	} else {
		c, err := cli.Client().KVDBClient()
		if err != nil {
			return fmt.Errorf("KVDB error: %v", err)
		}
		db := c.ProtoBroker()
		var txn = db.NewTxn()
		ops := 0
		for i := 0; i < len(keyVals); i++ {
			keyVal := keyVals[i]
			key, err := c.CompleteFullKey(keyVal.Key)
			if err != nil {
				return fmt.Errorf("key processing failed: %v", err)
			}
			fmt.Printf(" - %s\n", key)
			txn.Put(key, keyVal.Val)
			ops++

			if ops == int(opts.TxOps) || i+1 == len(keyVals) {
				fmt.Printf("commiting tx with %d ops\n", ops)
				ctx, cancel := context.WithTimeout(context.Background(), opts.Timeout)
				err = txn.Commit(ctx)
				cancel()
				if err != nil {
					return fmt.Errorf("commit failed: %v", err)
				}
				ops = 0
				txn = db.NewTxn()
			}
		}
	}

	logging.Debug("OK")
	return nil
}

type keyVal struct {
	Key string
	Val proto.Message
}

func parseImportFile(importFile string) (keyVals []keyVal, err error) {
	b, err := ioutil.ReadFile(importFile)
	if err != nil {
		return nil, fmt.Errorf("reading input file failed: %v", err)
	}
	// parse lines
	lines := bytes.Split(b, []byte("\n"))
	for _, l := range lines {
		line := bytes.TrimSpace(l)
		if bytes.HasPrefix(line, []byte("#")) {
			continue
		}
		parts := bytes.SplitN(line, []byte(" "), 2)
		if len(parts) < 2 {
			continue
		}
		key := string(parts[0])
		data := string(parts[1])
		if key == "" || data == "" {
			continue
		}
		logrus.Debugf("parse line: %s %s\n", key, data)

		//key = completeFullKey(key)
		val, err := unmarshalKeyVal(key, data)
		if err != nil {
			return nil, fmt.Errorf("decoding value failed: %v", err)
		}
		logrus.Debugf("KEY: %s - %v\n", key, val)
		keyVals = append(keyVals, keyVal{key, val})
	}
	return
}

func unmarshalKeyVal(fullKey string, data string) (proto.Message, error) {
	key := stripAgentPrefix(fullKey)

	model, err := models.GetModelForKey(key)
	if err != nil {
		return nil, err
	}
	valueType := proto.MessageType(model.ProtoName())
	if valueType == nil {
		return nil, fmt.Errorf("unknown proto message defined for key %s", key)
	}
	value := reflect.New(valueType.Elem()).Interface().(proto.Message)

	if err := jsonpb.UnmarshalString(data, value); err != nil {
		return nil, err
	}
	return value, nil
}

func stripAgentPrefix(key string) string {
	if !strings.HasPrefix(key, servicelabel.GetAllAgentsPrefix()) {
		return key
	}
	keyParts := strings.Split(key, "/")
	if len(keyParts) < 4 || keyParts[0] != "" {
		return path.Join(keyParts[2:]...)
	}
	return path.Join(keyParts[3:]...)
}
<|endoftext|>


<commit_before>package impdev

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"path"
	"sync"

	"github.com/spf13/cobra"
	"github.com/srinandan/apigeecli/cmd/shared"
	types "github.com/srinandan/apigeecli/cmd/types"
)

type Developer struct {
	EMail      string            `json:"email,omitempty"`
	FirstName  string            `json:"firstName,omitempty"`
	LastName   string            `json:"lastName,omitempty"`
	Attributes []types.Attribute `json:"attributes,omitempty"`
	Username   string            `json:"userName,omitempty"`
}

var Cmd = &cobra.Command{
	Use:   "import",
	Short: "Import a file containing App Developers",
	Long:  "Import a file containing App Developers",
	RunE: func(cmd *cobra.Command, args []string) error {
		u, _ := url.Parse(shared.BaseURL)
		u.Path = path.Join(u.Path, shared.RootArgs.Org, "developers")
		return createDevelopers(u.String())
	},
}

var conn int
var file string

func init() {

	Cmd.Flags().StringVarP(&file, "file", "f",
		"", "File containing App Developers")
	Cmd.Flags().IntVarP(&conn, "conn", "c",
		4, "Number of connections")

	_ = Cmd.MarkFlagRequired("file")
}

func createAsyncDeveloper(url string, developer Developer, wg *sync.WaitGroup, errChan chan<- *types.ImportError) {
	defer wg.Done()
	out, err := json.Marshal(developer)
	if err != nil {
		errChan <- &types.ImportError{Err: err}
		return
	}
	_, err = shared.HttpClient(true, url, string(out))
	if err != nil {
		errChan <- &types.ImportError{Err: err}
		return
	}

	errChan <- &types.ImportError{Err: nil}
}

func createDevelopers(url string) error {

	var errChan = make(chan *types.ImportError)
	var wg sync.WaitGroup

	developers, err := readDevelopersFile()
	if err != nil {
		shared.Error.Fatalln("Error reading file: ", err)
		return err
	}

	numDev := len(developers)
	shared.Info.Printf("Found %d products in the file\n", numDev)
	shared.Info.Printf("Create products with %d connections\n", conn)

	if numDev < conn {
		wg.Add(numDev)
		for i := 0; i < numDev; i++ {
			go createAsyncDeveloper(url, developers[i], &wg, errChan)
		}
		go func() {
			wg.Wait()
			close(errChan)
		}()
	} else {
		numOfLoops, remaining := numDev/conn, numDev%conn
		for i := 0; i < numOfLoops; i++ {
			shared.Info.Printf("Create %d batch of products\n", i)
			wg.Add(conn)
			for j := 0; j < conn; j++ {
				go createAsyncDeveloper(url, developers[j+(i*conn)], &wg, errChan)
			}
			go func() {
				wg.Wait()
			}()
		}

		wg.Add(remaining)
		shared.Info.Printf("Create remaining %d products\n", remaining)
		for i := (numDev - remaining); i < numDev; i++ {
			go createAsyncDeveloper(url, developers[i], &wg, errChan)
		}
		go func() {
			wg.Wait()
			close(errChan)
		}()
	}

	//print any errors and return an err
	var errs = false
	for errProd := range errChan {
		if errProd.Err != nil {
			shared.Error.Fatalln(errProd.Err)
			errs = true
		}
	}

	if errs {
		return fmt.Errorf("problem creating one of more products")
	}
	return nil
}

func readDevelopersFile() ([]Developer, error) {

	developers := []Developer{}

	jsonFile, err := os.Open(file)

	if err != nil {
		return developers, err
	}

	defer jsonFile.Close()

	byteValue, err := ioutil.ReadAll(jsonFile)

	if err != nil {
		return developers, err
	}

	err = json.Unmarshal(byteValue, &developers)

	if err != nil {
		return developers, err
	}

	return developers, nil

}
<commit_msg>reimpl batching logic<commit_after>package impdev

import (
	"encoding/json"
	"io/ioutil"
	"net/url"
	"os"
	"path"
	"sync"

	"github.com/spf13/cobra"
	"github.com/srinandan/apigeecli/cmd/shared"
	types "github.com/srinandan/apigeecli/cmd/types"
)

type Developer struct {
	EMail      string            `json:"email,omitempty"`
	FirstName  string            `json:"firstName,omitempty"`
	LastName   string            `json:"lastName,omitempty"`
	Attributes []types.Attribute `json:"attributes,omitempty"`
	Username   string            `json:"userName,omitempty"`
}

var Cmd = &cobra.Command{
	Use:   "import",
	Short: "Import a file containing App Developers",
	Long:  "Import a file containing App Developers",
	RunE: func(cmd *cobra.Command, args []string) error {
		u, _ := url.Parse(shared.BaseURL)
		u.Path = path.Join(u.Path, shared.RootArgs.Org, "developers")
		return createDevelopers(u.String())
	},
}

var conn int
var file string

func init() {

	Cmd.Flags().StringVarP(&file, "file", "f",
		"", "File containing App Developers")
	Cmd.Flags().IntVarP(&conn, "conn", "c",
		4, "Number of connections")

	_ = Cmd.MarkFlagRequired("file")
}

func createAsyncDeveloper(url string, developer Developer, wg *sync.WaitGroup) {
	defer wg.Done()
	out, err := json.Marshal(developer)
	if err != nil {
		shared.Error.Fatalln(err)
		return
	}
	_, err = shared.HttpClient(true, url, string(out))
	if err != nil {
		shared.Error.Fatalln(err)
		return
	}

	shared.Info.Printf("Completed entity: %s", developer.EMail)
}

//batch creates a batch of developers to create
func batch(url string, entities []Developer, pwg *sync.WaitGroup) {

	defer pwg.Done()
	//batch workgroup
	var bwg sync.WaitGroup

	bwg.Add(len(entities))

	for _, entity := range entities {
		go createAsyncDeveloper(url, entity, &bwg)
	}
	bwg.Wait()
}

func createDevelopers(url string) error {

	var pwg sync.WaitGroup

	entities, err := readDevelopersFile()
	if err != nil {
		shared.Error.Fatalln("Error reading file: ", err)
		return err
	}

	numEntities := len(entities)
	shared.Info.Printf("Found %d developers in the file\n", numEntities)
	shared.Info.Printf("Create developers with %d connections\n", conn)

	numOfLoops, remaining := numEntities/conn, numEntities%conn

	//ensure connections aren't greater than products
	if conn > numEntities {
		conn = numEntities
	}

	start := 0

	for i, end := 0, 0; i < numOfLoops; i++ {
		pwg.Add(1)
		end = (i * conn) + conn
		shared.Info.Printf("Creating batch %d of developers\n", (i + 1))
		go batch(url, entities[start:end], &pwg)
		start = end
		pwg.Wait()
	}

	if remaining > 0 {
		pwg.Add(1)
		shared.Info.Printf("Creating remaining %d developers\n", remaining)
		go batch(url, entities[start:numEntities], &pwg)
		pwg.Wait()
	}

	return nil
}

func readDevelopersFile() ([]Developer, error) {

	developers := []Developer{}

	jsonFile, err := os.Open(file)

	if err != nil {
		return developers, err
	}

	defer jsonFile.Close()

	byteValue, err := ioutil.ReadAll(jsonFile)

	if err != nil {
		return developers, err
	}

	err = json.Unmarshal(byteValue, &developers)

	if err != nil {
		return developers, err
	}

	return developers, nil

}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	vaultShowCommand = vaultCommands.Command("show", "Show contents of a vault item")
	vaultShowID      = vaultShowCommand.Arg("id", "ID of the vault item to show").String()
)

func init() {
	vaultShowCommand.Action(showVaultObject)
}

func showVaultObject(context *kingpin.ParseContext) error {
	v, err := openVault()
	if err != nil {
		return err
	}

	b, err := v.Get(*vaultShowID)
	if err != nil {
		return err
	}

	os.Stdout.Write(b)

	return nil
}
<commit_msg>cmdline tool<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"os"

	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	vaultShowCommand   = vaultCommands.Command("show", "Show contents of a vault item")
	vaultShowID        = vaultShowCommand.Arg("id", "ID of the vault item to show").String()
	vaultShowJSON      = vaultShowCommand.Flag("json", "Pretty-print JSON").Short('j').Bool()
	vaultShowNoNewLine = vaultShowCommand.Flag("nonewline", "Do not emit newline").Short('n').Bool()
)

func init() {
	vaultShowCommand.Action(showVaultObject)
}

func showVaultObject(context *kingpin.ParseContext) error {
	v, err := openVault()
	if err != nil {
		return err
	}

	b, err := v.Get(*vaultShowID)
	if err != nil {
		return err
	}

	if *vaultShowJSON {
		var buf bytes.Buffer
		json.Indent(&buf, b, "", "  ")
		buf.WriteTo(os.Stdout)
	} else {
		os.Stdout.Write(b)
	}

	if !*vaultShowNoNewLine {
		os.Stdout.WriteString("\n")
	}

	return nil
}
<|endoftext|>