<commit_before>package docker

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/docker/cli/cli/connhelper"
	"github.com/docker/cli/cli/context"
	"github.com/docker/cli/cli/context/store"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/pkg/errors"
)

// EndpointMeta is a typed wrapper around a context-store generic endpoint describing
// a Docker Engine endpoint, without its tls config
type EndpointMeta = context.EndpointMetaBase

// Endpoint is a typed wrapper around a context-store generic endpoint describing
// a Docker Engine endpoint, with its tls data
type Endpoint struct {
	EndpointMeta
	TLSData     *context.TLSData
	TLSPassword string
}

// WithTLSData loads TLS materials for the endpoint
func WithTLSData(s store.Reader, contextName string, m EndpointMeta) (Endpoint, error) {
	tlsData, err := context.LoadTLSData(s, contextName, DockerEndpoint)
	if err != nil {
		return Endpoint{}, err
	}
	return Endpoint{
		EndpointMeta: m,
		TLSData:      tlsData,
	}, nil
}

// tlsConfig extracts a context docker endpoint TLS config
func (c *Endpoint) tlsConfig() (*tls.Config, error) {
	if c.TLSData == nil && !c.SkipTLSVerify {
		// there is no specific tls config
		return nil, nil
	}
	var tlsOpts []func(*tls.Config)
	if c.TLSData != nil && c.TLSData.CA != nil {
		certPool := x509.NewCertPool()
		if !certPool.AppendCertsFromPEM(c.TLSData.CA) {
			return nil, errors.New("failed to retrieve context tls info: ca.pem seems invalid")
		}
		tlsOpts = append(tlsOpts, func(cfg *tls.Config) {
			cfg.RootCAs = certPool
		})
	}
	if c.TLSData != nil && c.TLSData.Key != nil && c.TLSData.Cert != nil {
		keyBytes := c.TLSData.Key
		pemBlock, _ := pem.Decode(keyBytes)
		if pemBlock == nil {
			return nil, fmt.Errorf("no valid private key found")
		}

		var err error
		if x509.IsEncryptedPEMBlock(pemBlock) {
			keyBytes, err = x509.DecryptPEMBlock(pemBlock, []byte(c.TLSPassword))
			if err != nil {
				return nil, errors.Wrap(err, "private key is encrypted, but could not decrypt it")
			}
			keyBytes = pem.EncodeToMemory(&pem.Block{Type: pemBlock.Type, Bytes: keyBytes})
		}

		x509cert, err := tls.X509KeyPair(c.TLSData.Cert, keyBytes)
		if err != nil {
			return nil, errors.Wrap(err, "failed to retrieve context tls info")
		}
		tlsOpts = append(tlsOpts, func(cfg *tls.Config) {
			cfg.Certificates = []tls.Certificate{x509cert}
		})
	}
	if c.SkipTLSVerify {
		tlsOpts = append(tlsOpts, func(cfg *tls.Config) {
			cfg.InsecureSkipVerify = true
		})
	}
	return tlsconfig.ClientDefault(tlsOpts...), nil
}

// ClientOpts returns a slice of Client options to configure an API client with this endpoint
func (c *Endpoint) ClientOpts() ([]client.Opt, error) {
	var result []client.Opt
	if c.Host != "" {
		helper, err := connhelper.GetConnectionHelper(c.Host)
		if err != nil {
			return nil, err
		}
		if helper == nil {
			tlsConfig, err := c.tlsConfig()
			if err != nil {
				return nil, err
			}
			result = append(result,
				client.WithHost(c.Host),
				withHTTPClient(tlsConfig),
			)

		} else {
			httpClient := &http.Client{
				// No tls
				// No proxy
				Transport: &http.Transport{
					DialContext: helper.Dialer,
				},
			}
			result = append(result,
				client.WithHTTPClient(httpClient),
				client.WithHost(helper.Host),
				client.WithDialContext(helper.Dialer),
			)
		}
		result = append(result, client.WithTimeout(10*time.Second))
	}

	version := os.Getenv("DOCKER_API_VERSION")
	if version != "" {
		result = append(result, client.WithVersion(version))
	}
	return result, nil
}

func withHTTPClient(tlsConfig *tls.Config) func(*client.Client) error {
	return func(c *client.Client) error {
		if tlsConfig == nil {
			// Use the default HTTPClient
			return nil
		}

		httpClient := &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: tlsConfig,
				DialContext: (&net.Dialer{
					KeepAlive: 30 * time.Second,
					Timeout:   30 * time.Second,
				}).DialContext,
			},
			CheckRedirect: client.CheckRedirect,
		}
		return client.WithHTTPClient(httpClient)(c)
	}
}

// EndpointFromContext parses a context docker endpoint metadata into a typed EndpointMeta structure
func EndpointFromContext(metadata store.Metadata) (EndpointMeta, error) {
	ep, ok := metadata.Endpoints[DockerEndpoint]
	if !ok {
		return EndpointMeta{}, errors.New("cannot find docker endpoint in context")
	}
	typed, ok := ep.(EndpointMeta)
	if !ok {
		return EndpointMeta{}, errors.Errorf("endpoint %q is not of type EndpointMeta", DockerEndpoint)
	}
	return typed, nil
}
<commit_msg>context: ClientOpts() now includes WithAPIVersionNegotiation if version is missing<commit_after>package docker

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/docker/cli/cli/connhelper"
	"github.com/docker/cli/cli/context"
	"github.com/docker/cli/cli/context/store"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/pkg/errors"
)

// EndpointMeta is a typed wrapper around a context-store generic endpoint describing
// a Docker Engine endpoint, without its tls config
type EndpointMeta = context.EndpointMetaBase

// Endpoint is a typed wrapper around a context-store generic endpoint describing
// a Docker Engine endpoint, with its tls data
type Endpoint struct {
	EndpointMeta
	TLSData     *context.TLSData
	TLSPassword string
}

// WithTLSData loads TLS materials for the endpoint
func WithTLSData(s store.Reader, contextName string, m EndpointMeta) (Endpoint, error) {
	tlsData, err := context.LoadTLSData(s, contextName, DockerEndpoint)
	if err != nil {
		return Endpoint{}, err
	}
	return Endpoint{
		EndpointMeta: m,
		TLSData:      tlsData,
	}, nil
}

// tlsConfig extracts a context docker endpoint TLS config
func (c *Endpoint) tlsConfig() (*tls.Config, error) {
	if c.TLSData == nil && !c.SkipTLSVerify {
		// there is no specific tls config
		return nil, nil
	}
	var tlsOpts []func(*tls.Config)
	if c.TLSData != nil && c.TLSData.CA != nil {
		certPool := x509.NewCertPool()
		if !certPool.AppendCertsFromPEM(c.TLSData.CA) {
			return nil, errors.New("failed to retrieve context tls info: ca.pem seems invalid")
		}
		tlsOpts = append(tlsOpts, func(cfg *tls.Config) {
			cfg.RootCAs = certPool
		})
	}
	if c.TLSData != nil && c.TLSData.Key != nil && c.TLSData.Cert != nil {
		keyBytes := c.TLSData.Key
		pemBlock, _ := pem.Decode(keyBytes)
		if pemBlock == nil {
			return nil, fmt.Errorf("no valid private key found")
		}

		var err error
		if x509.IsEncryptedPEMBlock(pemBlock) {
			keyBytes, err = x509.DecryptPEMBlock(pemBlock, []byte(c.TLSPassword))
			if err != nil {
				return nil, errors.Wrap(err, "private key is encrypted, but could not decrypt it")
			}
			keyBytes = pem.EncodeToMemory(&pem.Block{Type: pemBlock.Type, Bytes: keyBytes})
		}

		x509cert, err := tls.X509KeyPair(c.TLSData.Cert, keyBytes)
		if err != nil {
			return nil, errors.Wrap(err, "failed to retrieve context tls info")
		}
		tlsOpts = append(tlsOpts, func(cfg *tls.Config) {
			cfg.Certificates = []tls.Certificate{x509cert}
		})
	}
	if c.SkipTLSVerify {
		tlsOpts = append(tlsOpts, func(cfg *tls.Config) {
			cfg.InsecureSkipVerify = true
		})
	}
	return tlsconfig.ClientDefault(tlsOpts...), nil
}

// ClientOpts returns a slice of Client options to configure an API client with this endpoint
func (c *Endpoint) ClientOpts() ([]client.Opt, error) {
	var result []client.Opt
	if c.Host != "" {
		helper, err := connhelper.GetConnectionHelper(c.Host)
		if err != nil {
			return nil, err
		}
		if helper == nil {
			tlsConfig, err := c.tlsConfig()
			if err != nil {
				return nil, err
			}
			result = append(result,
				client.WithHost(c.Host),
				withHTTPClient(tlsConfig),
			)

		} else {
			httpClient := &http.Client{
				// No tls
				// No proxy
				Transport: &http.Transport{
					DialContext: helper.Dialer,
				},
			}
			result = append(result,
				client.WithHTTPClient(httpClient),
				client.WithHost(helper.Host),
				client.WithDialContext(helper.Dialer),
			)
		}
		result = append(result, client.WithTimeout(10*time.Second))
	}

	version := os.Getenv("DOCKER_API_VERSION")
	if version != "" {
		result = append(result, client.WithVersion(version))
	} else {
		result = append(result, client.WithAPIVersionNegotiation())
	}
	return result, nil
}

func withHTTPClient(tlsConfig *tls.Config) func(*client.Client) error {
	return func(c *client.Client) error {
		if tlsConfig == nil {
			// Use the default HTTPClient
			return nil
		}

		httpClient := &http.Client{
			Transport: &http.Transport{
				TLSClientConfig: tlsConfig,
				DialContext: (&net.Dialer{
					KeepAlive: 30 * time.Second,
					Timeout:   30 * time.Second,
				}).DialContext,
			},
			CheckRedirect: client.CheckRedirect,
		}
		return client.WithHTTPClient(httpClient)(c)
	}
}

// EndpointFromContext parses a context docker endpoint metadata into a typed EndpointMeta structure
func EndpointFromContext(metadata store.Metadata) (EndpointMeta, error) {
	ep, ok := metadata.Endpoints[DockerEndpoint]
	if !ok {
		return EndpointMeta{}, errors.New("cannot find docker endpoint in context")
	}
	typed, ok := ep.(EndpointMeta)
	if !ok {
		return EndpointMeta{}, errors.Errorf("endpoint %q is not of type EndpointMeta", DockerEndpoint)
	}
	return typed, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/mesosphere/dcos-commons/cli"
	"gopkg.in/alecthomas/kingpin.v2"
	"log"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
)

func main() {
	modName, err := cli.GetModuleName()
	if err != nil {
		log.Fatalf(err.Error())
	}
	// Prettify the user-visible service name:
	serviceName := strings.Title(modName) // eg 'Cassandra'
	if modName == "dse" {
		serviceName = "DSE"
	}

	app, err := cli.NewApp(
		"0.1.0",
		"Mesosphere",
		fmt.Sprintf("Deploy and manage %s clusters", serviceName))
	if err != nil {
		log.Fatalf(err.Error())
	}

	// Omit standard "connection" section. Cassandra isn't on the standard paths yet, and omit
	// standard "config" and "state" sections since Cassandra isn't installing
	// ConfigResource/StateResource yet.
	// Once these are fixed, this block can all be replaced with a call to "HandleCommonArgs"
	cli.HandleCommonFlags(app, modName, fmt.Sprintf("%s DC/OS CLI Module", serviceName))
	//cli.HandleConfigSection(app)
	//cli.HandleConnectionSection(app)
	cli.HandlePlanSection(app)
	//cli.HandleStateSection(app)

	handleSeedsCommand(app)
	handleCustomConnectionCommand(app, serviceName)

	handleNodeSection(app, serviceName)
	handleBackupRestoreSections(app, serviceName)
	handleCleanupRepairSections(app)

	// Omit modname:
	kingpin.MustParse(app.Parse(os.Args[2:]))
}

func handleSeedsCommand(app *kingpin.Application) {
	app.Command("seeds", "Retrieve seed node information").Action(
		func(c *kingpin.ParseContext) error {
			cli.PrintJSON(cli.HTTPGet("v1/seeds"))
			return nil
		})
}

type ConnectionHandler struct {
	showAddress bool
	showDns bool
}
func (cmd *ConnectionHandler) runBase(c *kingpin.ParseContext) error {
	if cmd.showAddress {
		cli.PrintJSON(cli.HTTPGet("v1/connection/address"))
	} else if cmd.showDns {
		cli.PrintJSON(cli.HTTPGet("v1/connection/dns"))
	} else {
		cli.PrintJSON(cli.HTTPGet("v1/connection"))
	}
	return nil
}
func handleCustomConnectionCommand(app *kingpin.Application, serviceName string) {
	cmd := &ConnectionHandler{}

	// Have three explicit commands, to ensure that each possibility shows up in --help:
	connection := app.Command("connection", fmt.Sprintf("Provides %s connection information", serviceName)).Action(cmd.runBase)
	connection.Flag("address", fmt.Sprintf("Provide addresses of the %s nodes", serviceName)).BoolVar(&cmd.showAddress)
	connection.Flag("dns", fmt.Sprintf("Provide dns names of the %s nodes", serviceName)).BoolVar(&cmd.showDns)
}

// Supports either '1234' OR 'node-1234', converting both to 1234.
type NodeValue struct {
	val int
}
func (v *NodeValue) Set(value string) error {
	// First check for an integer value of the form '123'. If this succeeds, print a deprecation warning.
	// TODO remove this block in favor of the following regexp to remove support for this style.
	var err error
	v.val, err = strconv.Atoi(value)
	if err == nil {
		log.Printf("WARNING: Task names of the form '%s' are deprecated. Instead use 'node-%s'.", value, value)
		return nil
	}

	// Next check for a string of the form 'node-123'.
	re := regexp.MustCompile("^(node-)?([0-9]+)$")
	vals := re.FindStringSubmatch(value)
	if len(vals) == 0 {
		return fmt.Errorf("Provided node value could not be parsed. Expected either 'node-NUM' or 'NUM' where NUM is an integer: %s", value)
	}
	v.val, err = strconv.Atoi(vals[2])
	if err != nil {
		return fmt.Errorf("Internal error: Can't convert %s to an int value: %s", vals[2], err)
	}
	return nil
}
func (v *NodeValue) String() string {
	return ""
}

type NodeHandler struct {
	nodeId NodeValue
}
func (cmd *NodeHandler) runDescribe(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/info", cmd.nodeId.val)))
	return nil
}
func (cmd *NodeHandler) runList(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet("v1/nodes/list"))
	return nil
}
func (cmd *NodeHandler) runReplace(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId.val))
	cli.PrintJSON(cli.HTTPPutQuery("v1/nodes/replace", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runRestart(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId.val))
	cli.PrintJSON(cli.HTTPPutQuery("v1/nodes/restart", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runStatus(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/status", cmd.nodeId.val)))
	return nil
}
func handleNodeSection(app *kingpin.Application, serviceName string) {
	cmd := &NodeHandler{}
	connection := app.Command("node", fmt.Sprintf("Manage %s nodes", serviceName))

	describe := connection.Command("describe", "Describes a single node").Action(cmd.runDescribe)
	describe.Arg("task_name", "The task name to describe, e.g. 'node-3'").SetValue(&cmd.nodeId)

	connection.Command("list", "Lists all nodes").Action(cmd.runList)

	replace := connection.Command("replace", "Replaces a single node job, moving it to a different agent").Action(cmd.runReplace)
	replace.Arg("task_name", "The task name to replace, e.g. 'node-3'").SetValue(&cmd.nodeId)

	restart := connection.Command("restart", "Restarts a single node job, keeping it on the same agent").Action(cmd.runRestart)
	restart.Arg("task_name", "The task name to restart, e.g. 'node-3'").SetValue(&cmd.nodeId)

	status := connection.Command("status", "Gets the status of a single node").Action(cmd.runStatus)
	status.Arg("task_name", "The task name to check, e.g. 'node-3'").SetValue(&cmd.nodeId)
}

// Reuse same struct for both 'backup start' and 'restore start' (same args)
type BackupRestoreHandler struct {
	backupName string
	externalLocation string
	s3AccessKey string
	s3SecretKey string
	azureAccount string
	azureKey string
}
func (cmd *BackupRestoreHandler) getArgs() map[string]interface{} {
	return map[string]interface{} {
		"backup_name": cmd.backupName,
		"external_location": cmd.externalLocation,
		"s3_access_key": cmd.s3AccessKey,
		"s3_secret_key": cmd.s3SecretKey,
		"azure_account": cmd.azureAccount,
		"azure_key": cmd.azureKey,
	}
}
func (cmd *BackupRestoreHandler) runBackup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/backup/start", string(payload))
	return nil
}
func (cmd *BackupRestoreHandler) runBackupStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/backup/stop")
	return nil
}
func (cmd *BackupRestoreHandler) runRestore(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/restore/start", string(payload))
	return nil
}
func (cmd *BackupRestoreHandler) runRestoreStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/restore/stop")
	return nil
}
func handleBackupRestoreSections(app *kingpin.Application, serviceName string) {
	cmd := &BackupRestoreHandler{}
	planCmd := &cli.PlanHandler{}

	backup := app.Command("backup", fmt.Sprintf("Backup %s cluster data", serviceName))
	backupStart := backup.Command(
		"start",
		"Perform cluster backup via snapshot mechanism").Action(cmd.runBackup)
	backupStart.Flag("backup_name", "Name of the snapshot").StringVar(&cmd.backupName)
	backupStart.Flag("external_location", "External location where the snapshot should be stored").StringVar(&cmd.externalLocation)
	backupStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	backupStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	backupStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	backupStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	backup.Command(
		"stop",
		"Stops a currently running backup").Action(cmd.runBackupStop)
	// same as 'plan show':
	backup.Command(
		"status",
		"Displays the status of the backup").Action(planCmd.RunShow)

	restore := app.Command("restore", fmt.Sprintf("Restore %s cluster from backup", serviceName))
	restoreStart := restore.Command(
		"start",
		"Restores cluster to a previous snapshot").Action(cmd.runRestore)
	restoreStart.Flag("backup_name", "Name of the snapshot to restore").StringVar(&cmd.backupName)
	restoreStart.Flag("external_location", "External location where the snapshot is stored").StringVar(&cmd.externalLocation)
	restoreStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	restoreStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	restoreStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	restoreStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	restore.Command(
		"stop",
		"Stops a currently running restore").Action(cmd.runRestoreStop)
	// same as 'plan show':
	restore.Command(
		"status",
		"Displays the status of the restore").Action(planCmd.RunShow)
}

// Reuse same struct for both 'cleanup start' and 'repair start' (same args)
type CleanupRepairHandler struct {
	nodes string
	keySpaces string
	columnFamilies string
}
func (cmd *CleanupRepairHandler) getArgs() map[string]interface{} {
	nodesList := []string{}
	nodesSplit := strings.Split(cmd.nodes, ",")
	if sliceContains(nodesSplit, "*") {
		nodesList = append(nodesList, "*")
	} else {
		for _, node := range nodesSplit {
			nodesList = append(nodesList, fmt.Sprintf("node-%s", strings.TrimSpace(node)))
		}
	}

	dict := map[string]interface{} {"nodes": nodesList}

	if len(cmd.keySpaces) != 0 {
		dict["key_spaces"] = cmd.keySpaces
	}
	if len(cmd.keySpaces) != 0 {
		dict["column_families"] = cmd.columnFamilies
	}
	return dict
}
func (cmd *CleanupRepairHandler) runCleanup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/cleanup/start", string(payload))
	return nil
}
func (cmd *CleanupRepairHandler) runCleanupStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/cleanup/stop")
	return nil
}
func (cmd *CleanupRepairHandler) runRepair(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/repair/start", string(payload))
	return nil
}
func (cmd *CleanupRepairHandler) runRepairStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/repair/stop")
	return nil
}
func handleCleanupRepairSections(app *kingpin.Application) {
	cmd := &CleanupRepairHandler{}

	cleanup := app.Command("cleanup", "Clean up old token mappings")
	cleanupStart := cleanup.Command(
		"start",
		"Perform cluster cleanup of deleted or moved keys").Action(cmd.runCleanup)
	cleanupStart.Flag("nodes", "A list of the nodes to cleanup or * for all.").Default("*").StringVar(&cmd.nodes)
	cleanupStart.Flag("key_spaces", "The key spaces to cleanup or empty for all.").StringVar(&cmd.keySpaces)
	cleanupStart.Flag("column_families", "The column families to cleanup.").StringVar(&cmd.columnFamilies)
	cleanup.Command(
		"stop",
		"Stops a currently running cleanup").Action(cmd.runCleanupStop)

	repair := app.Command("repair", "Perform primary range repair")
	repairStart := repair.Command(
		"start",
		"Perform primary range anti-entropy repair").Action(cmd.runRepair)
	repairStart.Flag("nodes", "A list of the nodes to repair or * for all.").Default("*").StringVar(&cmd.nodes)
	repairStart.Flag("key_spaces", "The key spaces to repair or empty for all.").StringVar(&cmd.keySpaces)
	repairStart.Flag("column_families", "The column families to repair.").StringVar(&cmd.columnFamilies)
	repair.Command(
		"stop",
		"Stops a currently running repair").Action(cmd.runRepairStop)
}
func sliceContains(stringSlice []string, searchString string) bool {
	for _, value := range stringSlice {
		if value == searchString {
			return true
		}
	}
	return false
}
<commit_msg>Remove comment about HandleCommonArgs<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/mesosphere/dcos-commons/cli"
	"gopkg.in/alecthomas/kingpin.v2"
	"log"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
)

func main() {
	modName, err := cli.GetModuleName()
	if err != nil {
		log.Fatalf(err.Error())
	}
	// Prettify the user-visible service name:
	serviceName := strings.Title(modName) // eg 'Cassandra'
	if modName == "dse" {
		serviceName = "DSE"
	}

	app, err := cli.NewApp(
		"0.1.0",
		"Mesosphere",
		fmt.Sprintf("Deploy and manage %s clusters", serviceName))
	if err != nil {
		log.Fatalf(err.Error())
	}

	// Omit standard "connection" section. Cassandra isn't on the standard paths yet, and omit
	// standard "config" and "state" sections since Cassandra isn't installing
	// ConfigResource/StateResource yet.
	cli.HandleCommonFlags(app, modName, fmt.Sprintf("%s DC/OS CLI Module", serviceName))
	//cli.HandleConfigSection(app)
	//cli.HandleConnectionSection(app)
	cli.HandlePlanSection(app)
	//cli.HandleStateSection(app)

	handleSeedsCommand(app)
	handleCustomConnectionCommand(app, serviceName)

	handleNodeSection(app, serviceName)
	handleBackupRestoreSections(app, serviceName)
	handleCleanupRepairSections(app)

	// Omit modname:
	kingpin.MustParse(app.Parse(os.Args[2:]))
}

func handleSeedsCommand(app *kingpin.Application) {
	app.Command("seeds", "Retrieve seed node information").Action(
		func(c *kingpin.ParseContext) error {
			cli.PrintJSON(cli.HTTPGet("v1/seeds"))
			return nil
		})
}

type ConnectionHandler struct {
	showAddress bool
	showDns bool
}
func (cmd *ConnectionHandler) runBase(c *kingpin.ParseContext) error {
	if cmd.showAddress {
		cli.PrintJSON(cli.HTTPGet("v1/connection/address"))
	} else if cmd.showDns {
		cli.PrintJSON(cli.HTTPGet("v1/connection/dns"))
	} else {
		cli.PrintJSON(cli.HTTPGet("v1/connection"))
	}
	return nil
}
func handleCustomConnectionCommand(app *kingpin.Application, serviceName string) {
	cmd := &ConnectionHandler{}

	// Have three explicit commands, to ensure that each possibility shows up in --help:
	connection := app.Command("connection", fmt.Sprintf("Provides %s connection information", serviceName)).Action(cmd.runBase)
	connection.Flag("address", fmt.Sprintf("Provide addresses of the %s nodes", serviceName)).BoolVar(&cmd.showAddress)
	connection.Flag("dns", fmt.Sprintf("Provide dns names of the %s nodes", serviceName)).BoolVar(&cmd.showDns)
}

// Supports either '1234' OR 'node-1234', converting both to 1234.
type NodeValue struct {
	val int
}
func (v *NodeValue) Set(value string) error {
	// First check for an integer value of the form '123'. If this succeeds, print a deprecation warning.
	// TODO remove this block in favor of the following regexp to remove support for this style.
	var err error
	v.val, err = strconv.Atoi(value)
	if err == nil {
		log.Printf("WARNING: Task names of the form '%s' are deprecated. Instead use 'node-%s'.", value, value)
		return nil
	}

	// Next check for a string of the form 'node-123'.
	re := regexp.MustCompile("^(node-)?([0-9]+)$")
	vals := re.FindStringSubmatch(value)
	if len(vals) == 0 {
		return fmt.Errorf("Provided node value could not be parsed. Expected either 'node-NUM' or 'NUM' where NUM is an integer: %s", value)
	}
	v.val, err = strconv.Atoi(vals[2])
	if err != nil {
		return fmt.Errorf("Internal error: Can't convert %s to an int value: %s", vals[2], err)
	}
	return nil
}
func (v *NodeValue) String() string {
	return ""
}

type NodeHandler struct {
	nodeId NodeValue
}
func (cmd *NodeHandler) runDescribe(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/info", cmd.nodeId.val)))
	return nil
}
func (cmd *NodeHandler) runList(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet("v1/nodes/list"))
	return nil
}
func (cmd *NodeHandler) runReplace(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId.val))
	cli.PrintJSON(cli.HTTPPutQuery("v1/nodes/replace", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runRestart(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId.val))
	cli.PrintJSON(cli.HTTPPutQuery("v1/nodes/restart", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runStatus(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/status", cmd.nodeId.val)))
	return nil
}
func handleNodeSection(app *kingpin.Application, serviceName string) {
	cmd := &NodeHandler{}
	connection := app.Command("node", fmt.Sprintf("Manage %s nodes", serviceName))

	describe := connection.Command("describe", "Describes a single node").Action(cmd.runDescribe)
	describe.Arg("task_name", "The task name to describe, e.g. 'node-3'").SetValue(&cmd.nodeId)

	connection.Command("list", "Lists all nodes").Action(cmd.runList)

	replace := connection.Command("replace", "Replaces a single node job, moving it to a different agent").Action(cmd.runReplace)
	replace.Arg("task_name", "The task name to replace, e.g. 'node-3'").SetValue(&cmd.nodeId)

	restart := connection.Command("restart", "Restarts a single node job, keeping it on the same agent").Action(cmd.runRestart)
	restart.Arg("task_name", "The task name to restart, e.g. 'node-3'").SetValue(&cmd.nodeId)

	status := connection.Command("status", "Gets the status of a single node").Action(cmd.runStatus)
	status.Arg("task_name", "The task name to check, e.g. 'node-3'").SetValue(&cmd.nodeId)
}

// Reuse same struct for both 'backup start' and 'restore start' (same args)
type BackupRestoreHandler struct {
	backupName string
	externalLocation string
	s3AccessKey string
	s3SecretKey string
	azureAccount string
	azureKey string
}
func (cmd *BackupRestoreHandler) getArgs() map[string]interface{} {
	return map[string]interface{} {
		"backup_name": cmd.backupName,
		"external_location": cmd.externalLocation,
		"s3_access_key": cmd.s3AccessKey,
		"s3_secret_key": cmd.s3SecretKey,
		"azure_account": cmd.azureAccount,
		"azure_key": cmd.azureKey,
	}
}
func (cmd *BackupRestoreHandler) runBackup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/backup/start", string(payload))
	return nil
}
func (cmd *BackupRestoreHandler) runBackupStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/backup/stop")
	return nil
}
func (cmd *BackupRestoreHandler) runRestore(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/restore/start", string(payload))
	return nil
}
func (cmd *BackupRestoreHandler) runRestoreStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/restore/stop")
	return nil
}
func handleBackupRestoreSections(app *kingpin.Application, serviceName string) {
	cmd := &BackupRestoreHandler{}
	planCmd := &cli.PlanHandler{}

	backup := app.Command("backup", fmt.Sprintf("Backup %s cluster data", serviceName))
	backupStart := backup.Command(
		"start",
		"Perform cluster backup via snapshot mechanism").Action(cmd.runBackup)
	backupStart.Flag("backup_name", "Name of the snapshot").StringVar(&cmd.backupName)
	backupStart.Flag("external_location", "External location where the snapshot should be stored").StringVar(&cmd.externalLocation)
	backupStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	backupStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	backupStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	backupStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	backup.Command(
		"stop",
		"Stops a currently running backup").Action(cmd.runBackupStop)
	// same as 'plan show':
	backup.Command(
		"status",
		"Displays the status of the backup").Action(planCmd.RunShow)

	restore := app.Command("restore", fmt.Sprintf("Restore %s cluster from backup", serviceName))
	restoreStart := restore.Command(
		"start",
		"Restores cluster to a previous snapshot").Action(cmd.runRestore)
	restoreStart.Flag("backup_name", "Name of the snapshot to restore").StringVar(&cmd.backupName)
	restoreStart.Flag("external_location", "External location where the snapshot is stored").StringVar(&cmd.externalLocation)
	restoreStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	restoreStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	restoreStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	restoreStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	restore.Command(
		"stop",
		"Stops a currently running restore").Action(cmd.runRestoreStop)
	// same as 'plan show':
	restore.Command(
		"status",
		"Displays the status of the restore").Action(planCmd.RunShow)
}

// Reuse same struct for both 'cleanup start' and 'repair start' (same args)
type CleanupRepairHandler struct {
	nodes string
	keySpaces string
	columnFamilies string
}
func (cmd *CleanupRepairHandler) getArgs() map[string]interface{} {
	nodesList := []string{}
	nodesSplit := strings.Split(cmd.nodes, ",")
	if sliceContains(nodesSplit, "*") {
		nodesList = append(nodesList, "*")
	} else {
		for _, node := range nodesSplit {
			nodesList = append(nodesList, fmt.Sprintf("node-%s", strings.TrimSpace(node)))
		}
	}

	dict := map[string]interface{} {"nodes": nodesList}

	if len(cmd.keySpaces) != 0 {
		dict["key_spaces"] = cmd.keySpaces
	}
	if len(cmd.keySpaces) != 0 {
		dict["column_families"] = cmd.columnFamilies
	}
	return dict
}
func (cmd *CleanupRepairHandler) runCleanup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/cleanup/start", string(payload))
	return nil
}
func (cmd *CleanupRepairHandler) runCleanupStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/cleanup/stop")
	return nil
}
func (cmd *CleanupRepairHandler) runRepair(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.HTTPPutJSON("v1/repair/start", string(payload))
	return nil
}
func (cmd *CleanupRepairHandler) runRepairStop(c *kingpin.ParseContext) error {
	cli.HTTPPut("v1/repair/stop")
	return nil
}
func handleCleanupRepairSections(app *kingpin.Application) {
	cmd := &CleanupRepairHandler{}

	cleanup := app.Command("cleanup", "Clean up old token mappings")
	cleanupStart := cleanup.Command(
		"start",
		"Perform cluster cleanup of deleted or moved keys").Action(cmd.runCleanup)
	cleanupStart.Flag("nodes", "A list of the nodes to cleanup or * for all.").Default("*").StringVar(&cmd.nodes)
	cleanupStart.Flag("key_spaces", "The key spaces to cleanup or empty for all.").StringVar(&cmd.keySpaces)
	cleanupStart.Flag("column_families", "The column families to cleanup.").StringVar(&cmd.columnFamilies)
	cleanup.Command(
		"stop",
		"Stops a currently running cleanup").Action(cmd.runCleanupStop)

	repair := app.Command("repair", "Perform primary range repair")
	repairStart := repair.Command(
		"start",
		"Perform primary range anti-entropy repair").Action(cmd.runRepair)
	repairStart.Flag("nodes", "A list of the nodes to repair or * for all.").Default("*").StringVar(&cmd.nodes)
	repairStart.Flag("key_spaces", "The key spaces to repair or empty for all.").StringVar(&cmd.keySpaces)
	repairStart.Flag("column_families", "The column families to repair.").StringVar(&cmd.columnFamilies)
	repair.Command(
		"stop",
		"Stops a currently running repair").Action(cmd.runRepairStop)
}
func sliceContains(stringSlice []string, searchString string) bool {
	for _, value := range stringSlice {
		if value == searchString {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package request

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http/httputil"

	httpsignatures "github.com/99designs/httpsignatures-go"
	dd_opentracing "github.com/DataDog/dd-trace-go/opentracing"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/remind101/pkg/httpx"
)

type Handlers struct {
	Build            HandlerList
	Sign             HandlerList
	Send             HandlerList
	ValidateResponse HandlerList
	Decode           HandlerList
	DecodeError      HandlerList
	Complete         HandlerList
}

func DefaultHandlers() Handlers {
	return Handlers{
		Build:            NewHandlerList(JSONBuilder),
		Sign:             NewHandlerList(),
		Send:             NewHandlerList(WithTracing(BaseSender)),
		ValidateResponse: NewHandlerList(),
		Decode:           NewHandlerList(JSONDecoder),
		DecodeError:      NewHandlerList(),
		Complete:         NewHandlerList(),
	}
}

func (h Handlers) Copy() Handlers {
	return Handlers{
		Build:            h.Build.copy(),
		Sign:             h.Sign.copy(),
		Send:             h.Send.copy(),
		ValidateResponse: h.ValidateResponse.copy(),
		Decode:           h.Decode.copy(),
		DecodeError:      h.DecodeError.copy(),
		Complete:         h.Complete.copy(),
	}
}

type HandlerList struct {
	list []Handler
}

func NewHandlerList(hh ...Handler) HandlerList {
	return HandlerList{
		list: append([]Handler{}, hh...),
	}
}

func (hl *HandlerList) Run(r *Request) {
	for _, h := range hl.list {
		h.Fn(r)
	}
}

func (hl *HandlerList) Append(h Handler) {
	hl.list = append(hl.list, h)
}

func (hl *HandlerList) Prepend(h Handler) {
	hl.list = append([]Handler{h}, hl.list...)
}

func (hl *HandlerList) copy() HandlerList {
	n := HandlerList{}
	if len(hl.list) == 0 {
		return n
	}

	n.list = append(make([]Handler, 0, len(hl.list)), hl.list...)
	return n
}

type Handler struct {
	Name string
	Fn   func(*Request)
}

// BaseSender sends a request using the http.Client.
var BaseSender = Handler{
	Name: "BaseSender",
	Fn: func(r *Request) {
		r.HTTPResponse, r.Error = r.HTTPClient.Do(r.HTTPRequest)
	},
}

var JSONBuilder = Handler{
	Name: "JSONBuilder",
	Fn: func(r *Request) {
		r.HTTPRequest.Header.Set("Content-Type", "application/json")
		r.HTTPRequest.Header.Set("Accept", "application/json")

		if r.HTTPRequest.Method != "GET" && r.Params != nil {
			raw, err := json.Marshal(r.Params)
			if err != nil {
				r.Error = err
				return
			}
			r.HTTPRequest.Body = ioutil.NopCloser(bytes.NewReader(raw))
		}
	},
}

// JSONDecoder decodes a response as JSON.
var JSONDecoder = Handler{
	Name: "JSONDecoder",
	Fn: func(r *Request) {
		if r.HTTPResponse == nil {
			return
		}
		if r.HTTPResponse.Body != nil {
			defer r.HTTPResponse.Body.Close()
		}
		if r.Data == nil {
			_, r.Error = io.Copy(ioutil.Discard, r.HTTPResponse.Body)
			return
		}
		r.Error = json.NewDecoder(r.HTTPResponse.Body).Decode(r.Data)
	},
}

// RequestSigner signs requests.
func RequestSigner(id, key string) Handler {
	return Handler{
		Name: "RequestSigner",
		Fn: func(r *Request) {
			r.Error = httpsignatures.DefaultSha256Signer.SignRequest(id, key, r.HTTPRequest)
		},
	}
}

// BasicAuther sets basic auth on a request.
func BasicAuther(username, password string) Handler {
	return Handler{
		Name: "BasicAuther",
		Fn: func(r *Request) {
			r.HTTPRequest.SetBasicAuth(username, password)
		},
	}
}

// HeadersFromContext adds headers with values from the request context.
// This is useful for forwarding headers to upstream services.
func HeadersFromContext(headers ...string) Handler {
	return Handler{
		Name: "HeadersFromContext",
		Fn: func(r *Request) {
			for _, header := range headers {
				r.HTTPRequest.Header.Add(header, httpx.Header(r.HTTPRequest.Context(), header))
			}
		},
	}
}

// RequestLogger dumps the entire request to stdout.
var RequestLogger = Handler{
	Name: "RequestLogger",
	Fn: func(r *Request) {
		b, err := httputil.DumpRequestOut(r.HTTPRequest, true)
		if err != nil {
			fmt.Printf("error dumping request: %s\n", err.Error())
			return
		}
		fmt.Println(string(b))
	},
}

// ResponseLogger dumps the entire response to stdout.
var ResponseLogger = Handler{
	Name: "ResponseLogger",
	Fn: func(r *Request) {
		b, err := httputil.DumpResponse(r.HTTPResponse, true)
		if err != nil {
			fmt.Printf("error dumping response: %s\n", err.Error())
			return
		}
		fmt.Println(string(b))
	},
}

// WithTracing returns a Send Handler that wraps another Send Handler in a trace
// span.
func WithTracing(h Handler) Handler {
	return Handler{
		Name: "TracedSender",
		Fn: func(r *Request) {
			span, ctx := opentracing.StartSpanFromContext(r.HTTPRequest.Context(), "client.request")
			defer span.Finish()
			r.HTTPRequest = r.HTTPRequest.WithContext(ctx)

			span.SetTag("http.method", r.HTTPRequest.Method)
			span.SetTag("http.url", r.HTTPRequest.URL.Hostname()+r.HTTPRequest.URL.EscapedPath())

			h.Fn(r)

			if r.HTTPResponse != nil {
				span.SetTag("http.status_code", r.HTTPResponse.StatusCode)
			}

			if r.Error != nil {
				span.SetTag(dd_opentracing.Error, r.Error)
			}
		},
	}
}
<commit_msg>Inject opentracing headers<commit_after>package request

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http/httputil"

	httpsignatures "github.com/99designs/httpsignatures-go"
	dd_opentracing "github.com/DataDog/dd-trace-go/opentracing"
	opentracing "github.com/opentracing/opentracing-go"
	"github.com/remind101/pkg/httpx"
)

type Handlers struct {
	Build            HandlerList
	Sign             HandlerList
	Send             HandlerList
	ValidateResponse HandlerList
	Decode           HandlerList
	DecodeError      HandlerList
	Complete         HandlerList
}

func DefaultHandlers() Handlers {
	return Handlers{
		Build:            NewHandlerList(JSONBuilder),
		Sign:             NewHandlerList(),
		Send:             NewHandlerList(WithTracing(BaseSender)),
		ValidateResponse: NewHandlerList(),
		Decode:           NewHandlerList(JSONDecoder),
		DecodeError:      NewHandlerList(),
		Complete:         NewHandlerList(),
	}
}

func (h Handlers) Copy() Handlers {
	return Handlers{
		Build:            h.Build.copy(),
		Sign:             h.Sign.copy(),
		Send:             h.Send.copy(),
		ValidateResponse: h.ValidateResponse.copy(),
		Decode:           h.Decode.copy(),
		DecodeError:      h.DecodeError.copy(),
		Complete:         h.Complete.copy(),
	}
}

type HandlerList struct {
	list []Handler
}

func NewHandlerList(hh ...Handler) HandlerList {
	return HandlerList{
		list: append([]Handler{}, hh...),
	}
}

func (hl *HandlerList) Run(r *Request) {
	for _, h := range hl.list {
		h.Fn(r)
	}
}

func (hl *HandlerList) Append(h Handler) {
	hl.list = append(hl.list, h)
}

func (hl *HandlerList) Prepend(h Handler) {
	hl.list = append([]Handler{h}, hl.list...)
}

func (hl *HandlerList) copy() HandlerList {
	n := HandlerList{}
	if len(hl.list) == 0 {
		return n
	}

	n.list = append(make([]Handler, 0, len(hl.list)), hl.list...)
	return n
}

type Handler struct {
	Name string
	Fn   func(*Request)
}

// BaseSender sends a request using the http.Client.
var BaseSender = Handler{
	Name: "BaseSender",
	Fn: func(r *Request) {
		r.HTTPResponse, r.Error = r.HTTPClient.Do(r.HTTPRequest)
	},
}

var JSONBuilder = Handler{
	Name: "JSONBuilder",
	Fn: func(r *Request) {
		r.HTTPRequest.Header.Set("Content-Type", "application/json")
		r.HTTPRequest.Header.Set("Accept", "application/json")

		if r.HTTPRequest.Method != "GET" && r.Params != nil {
			raw, err := json.Marshal(r.Params)
			if err != nil {
				r.Error = err
				return
			}
			r.HTTPRequest.Body = ioutil.NopCloser(bytes.NewReader(raw))
		}
	},
}

// JSONDecoder decodes a response as JSON.
var JSONDecoder = Handler{
	Name: "JSONDecoder",
	Fn: func(r *Request) {
		if r.HTTPResponse == nil {
			return
		}
		if r.HTTPResponse.Body != nil {
			defer r.HTTPResponse.Body.Close()
		}
		if r.Data == nil {
			_, r.Error = io.Copy(ioutil.Discard, r.HTTPResponse.Body)
			return
		}
		r.Error = json.NewDecoder(r.HTTPResponse.Body).Decode(r.Data)
	},
}

// RequestSigner signs requests.
func RequestSigner(id, key string) Handler {
	return Handler{
		Name: "RequestSigner",
		Fn: func(r *Request) {
			r.Error = httpsignatures.DefaultSha256Signer.SignRequest(id, key, r.HTTPRequest)
		},
	}
}

// BasicAuther sets basic auth on a request.
func BasicAuther(username, password string) Handler {
	return Handler{
		Name: "BasicAuther",
		Fn: func(r *Request) {
			r.HTTPRequest.SetBasicAuth(username, password)
		},
	}
}

// HeadersFromContext adds headers with values from the request context.
// This is useful for forwarding headers to upstream services.
func HeadersFromContext(headers ...string) Handler {
	return Handler{
		Name: "HeadersFromContext",
		Fn: func(r *Request) {
			for _, header := range headers {
				r.HTTPRequest.Header.Add(header, httpx.Header(r.HTTPRequest.Context(), header))
			}
		},
	}
}

// RequestLogger dumps the entire request to stdout.
var RequestLogger = Handler{
	Name: "RequestLogger",
	Fn: func(r *Request) {
		b, err := httputil.DumpRequestOut(r.HTTPRequest, true)
		if err != nil {
			fmt.Printf("error dumping request: %s\n", err.Error())
			return
		}
		fmt.Println(string(b))
	},
}

// ResponseLogger dumps the entire response to stdout.
var ResponseLogger = Handler{
	Name: "ResponseLogger",
	Fn: func(r *Request) {
		b, err := httputil.DumpResponse(r.HTTPResponse, true)
		if err != nil {
			fmt.Printf("error dumping response: %s\n", err.Error())
			return
		}
		fmt.Println(string(b))
	},
}

// WithTracing returns a Send Handler that wraps another Send Handler in a trace
// span.
func WithTracing(h Handler) Handler {
	return Handler{
		Name: "TracedSender",
		Fn: func(r *Request) {
			span, ctx := opentracing.StartSpanFromContext(r.HTTPRequest.Context(), "client.request")
			opentracing.GlobalTracer().Inject(
				span.Context(),
				opentracing.HTTPHeaders,
				opentracing.HTTPHeadersCarrier(r.HTTPRequest.Header),
			)
			defer span.Finish()
			r.HTTPRequest = r.HTTPRequest.WithContext(ctx)

			span.SetTag("http.method", r.HTTPRequest.Method)
			span.SetTag("http.url", r.HTTPRequest.URL.Hostname()+r.HTTPRequest.URL.EscapedPath())

			h.Fn(r)

			if r.HTTPResponse != nil {
				span.SetTag("http.status_code", r.HTTPResponse.StatusCode)
			}

			if r.Error != nil {
				span.SetTag(dd_opentracing.Error, r.Error)
			}
		},
	}
}
<|endoftext|>


<commit_before>package grim

// Copyright 2015 MediaMath <http://www.mediamath.com>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

import (
	"encoding/json"
	"fmt"

	"github.com/awslabs/aws-sdk-go/aws"
	"github.com/awslabs/aws-sdk-go/aws/awserr"
	"github.com/awslabs/aws-sdk-go/service/sqs"
)

var policyFormat = `{
	"Version": "2008-10-17",
	"Id": "grim-policy",
	"Statement": [
	  {
			"Sid": "1",
			"Effect": "Allow",
			"Principal": {
				"AWS": "*"
			},
			"Action": "SQS:*",
			"Resource": "%v",
			"Condition": {
				"ArnEquals": {
					"aws:SourceArn": %v
				}
			}
		}
	]
}`

type sqsQueue struct {
	URL string
	ARN string
}

func prepareSQSQueue(key, secret, region, queue string) (*sqsQueue, error) {
	config := getConfig(key, secret, region)

	queueURL, err := getQueueURLByName(config, queue)
	if err != nil {
		queueURL, err = createQueue(config, queue)
	}

	if err != nil {
		return nil, err
	}

	queueARN, err := getARNForQueueURL(config, queueURL)
	if err != nil {
		return nil, err
	}

	return &sqsQueue{queueURL, queueARN}, nil
}

func getNextMessage(key, secret, region, queueURL string) (string, error) {
	config := getConfig(key, secret, region)

	message, err := getMessage(config, queueURL)
	if err != nil {
		return "", err
	} else if message == nil || message.ReceiptHandle == nil {
		return "", nil
	}

	err = deleteMessage(config, queueURL, *message.ReceiptHandle)
	if err != nil {
		return "", err
	}

	if message.Body == nil {
		return "", nil
	}

	return *message.Body, nil
}

func setPolicy(key, secret, region, queueARN, queueURL string, topicARNs []string) error {
	svc := sqs.New(getConfig(key, secret, region))

	bs, err := json.Marshal(topicARNs)
	if err != nil {
		return fmt.Errorf("error while creating policy for SQS queue: %v", err)
	}

	policy := fmt.Sprintf(policyFormat, queueARN, string(bs))

	params := &sqs.SetQueueAttributesInput{
		Attributes: &map[string]*string{
			"Policy": aws.String(policy),
		},
		QueueURL: aws.String(queueURL),
	}

	_, err = svc.SetQueueAttributes(params)
	if awserr, ok := err.(awserr.Error); ok {
		return fmt.Errorf("aws error while setting policy for SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return fmt.Errorf("error while setting policy for SQS queue: %v", err)
	}

	return nil
}

func getQueueURLByName(config *aws.Config, queue string) (string, error) {
	svc := sqs.New(config)

	params := &sqs.GetQueueURLInput{
		QueueName: aws.String(queue),
	}

	resp, err := svc.GetQueueURL(params)
	if awserr, ok := err.(awserr.Error); ok {
		return "", fmt.Errorf("aws error while getting URL for SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return "", fmt.Errorf("error while getting URL for SQS queue: %v", err)
	} else if resp == nil || resp.QueueURL == nil {
		return "", nil
	}

	return *resp.QueueURL, nil
}

func getARNForQueueURL(config *aws.Config, queueURL string) (string, error) {
	svc := sqs.New(config)

	arnKey := "QueueArn"

	params := &sqs.GetQueueAttributesInput{
		QueueURL: aws.String(string(queueURL)),
		AttributeNames: []*string{
			aws.String(arnKey),
		},
	}

	resp, err := svc.GetQueueAttributes(params)
	if awserr, ok := err.(awserr.Error); ok {
		return "", fmt.Errorf("aws error while getting ARN for SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return "", fmt.Errorf("error while getting ARN for SQS queue: %v", err)
	} else if resp == nil || resp.Attributes == nil {
		return "", nil
	}

	atts := *resp.Attributes

	arnPtr, ok := atts[arnKey]
	if !ok || arnPtr == nil {
		return "", nil
	}

	return *arnPtr, nil
}

func createQueue(config *aws.Config, queue string) (string, error) {
	svc := sqs.New(config)

	params := &sqs.CreateQueueInput{
		QueueName: aws.String(queue),
		Attributes: &map[string]*string{
			"ReceiveMessageWaitTimeSeconds": aws.String("5"),
		},
	}

	resp, err := svc.CreateQueue(params)
	if awserr, ok := err.(awserr.Error); ok {
		return "", fmt.Errorf("aws error while creating SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return "", fmt.Errorf("error while creating SQS queue: %v", err)
	} else if resp == nil || resp.QueueURL == nil {
		return "", nil
	}

	return *resp.QueueURL, nil
}

func getMessage(config *aws.Config, queueURL string) (*sqs.Message, error) {
	svc := sqs.New(config)

	params := &sqs.ReceiveMessageInput{
		QueueURL:            aws.String(queueURL),
		MaxNumberOfMessages: aws.Long(1),
	}

	resp, err := svc.ReceiveMessage(params)
	if awserr, ok := err.(awserr.Error); ok {
		return nil, fmt.Errorf("aws error while receiving message from SQS: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return nil, fmt.Errorf("error while receiving message from SQS: %v", err)
	} else if resp == nil || len(resp.Messages) == 0 {
		return nil, nil
	}

	return resp.Messages[0], nil
}

func deleteMessage(config *aws.Config, queueURL string, receiptHandle string) error {
	svc := sqs.New(config)

	params := &sqs.DeleteMessageInput{
		QueueURL:      aws.String(queueURL),
		ReceiptHandle: aws.String(receiptHandle),
	}

	_, err := svc.DeleteMessage(params)
	if awserr, ok := err.(awserr.Error); ok {
		return fmt.Errorf("aws error while deleting message from SQS: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return fmt.Errorf("error while deleting message from SQS: %v", err)
	}

	return nil
}
<commit_msg>fixing aws-go-sdk breaking changes.<commit_after>package grim

// Copyright 2015 MediaMath <http://www.mediamath.com>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

import (
	"encoding/json"
	"fmt"

	"github.com/awslabs/aws-sdk-go/aws"
	"github.com/awslabs/aws-sdk-go/aws/awserr"
	"github.com/awslabs/aws-sdk-go/service/sqs"
)

var policyFormat = `{
	"Version": "2008-10-17",
	"Id": "grim-policy",
	"Statement": [
	  {
			"Sid": "1",
			"Effect": "Allow",
			"Principal": {
				"AWS": "*"
			},
			"Action": "SQS:*",
			"Resource": "%v",
			"Condition": {
				"ArnEquals": {
					"aws:SourceArn": %v
				}
			}
		}
	]
}`

type sqsQueue struct {
	URL string
	ARN string
}

func prepareSQSQueue(key, secret, region, queue string) (*sqsQueue, error) {
	config := getConfig(key, secret, region)

	queueURL, err := getQueueURLByName(config, queue)
	if err != nil {
		queueURL, err = createQueue(config, queue)
	}

	if err != nil {
		return nil, err
	}

	queueARN, err := getARNForQueueURL(config, queueURL)
	if err != nil {
		return nil, err
	}

	return &sqsQueue{queueURL, queueARN}, nil
}

func getNextMessage(key, secret, region, queueURL string) (string, error) {
	config := getConfig(key, secret, region)

	message, err := getMessage(config, queueURL)
	if err != nil {
		return "", err
	} else if message == nil || message.ReceiptHandle == nil {
		return "", nil
	}

	err = deleteMessage(config, queueURL, *message.ReceiptHandle)
	if err != nil {
		return "", err
	}

	if message.Body == nil {
		return "", nil
	}

	return *message.Body, nil
}

func setPolicy(key, secret, region, queueARN, queueURL string, topicARNs []string) error {
	svc := sqs.New(getConfig(key, secret, region))

	bs, err := json.Marshal(topicARNs)
	if err != nil {
		return fmt.Errorf("error while creating policy for SQS queue: %v", err)
	}

	policy := fmt.Sprintf(policyFormat, queueARN, string(bs))

	params := &sqs.SetQueueAttributesInput{
		Attributes: &map[string]*string{
			"Policy": aws.String(policy),
		},
		QueueURL: aws.String(queueURL),
	}

	_, err = svc.SetQueueAttributes(params)
	if awserr, ok := err.(awserr.Error); ok {
		return fmt.Errorf("aws error while setting policy for SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return fmt.Errorf("error while setting policy for SQS queue: %v", err)
	}

	return nil
}

func getQueueURLByName(config *aws.Config, queue string) (string, error) {
	svc := sqs.New(config)

	params := sqs.GetQueueURLInput{
		QueueName: aws.String(queue),
	}

	resp, err := svc.GetQueueURL(params)
	if awserr, ok := err.(awserr.Error); ok {
		return "", fmt.Errorf("aws error while getting URL for SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return "", fmt.Errorf("error while getting URL for SQS queue: %v", err)
	} else if resp == nil || resp.QueueURL == nil {
		return "", nil
	}

	return *resp.QueueURL, nil
}

func getARNForQueueURL(config *aws.Config, queueURL string) (string, error) {
	svc := sqs.New(config)

	arnKey := "QueueArn"

	params := &sqs.GetQueueAttributesInput{
		QueueURL: aws.String(string(queueURL)),
		AttributeNames: []*string{
			aws.String(arnKey),
		},
	}

	resp, err := svc.GetQueueAttributes(params)
	if awserr, ok := err.(awserr.Error); ok {
		return "", fmt.Errorf("aws error while getting ARN for SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return "", fmt.Errorf("error while getting ARN for SQS queue: %v", err)
	} else if resp == nil || resp.Attributes == nil {
		return "", nil
	}

	atts := resp.Attributes

	arnPtr, ok := atts[arnKey]
	if !ok || arnPtr == nil {
		return "", nil
	}

	return *arnPtr, nil
}

func createQueue(config *aws.Config, queue string) (string, error) {
	svc := sqs.New(config)

	params := &sqs.CreateQueueInput{
		QueueName: aws.String(queue),
		Attributes: map[string]*string{
			"ReceiveMessageWaitTimeSeconds": aws.String("5"),
		},
	}

	resp, err := svc.CreateQueue(params)
	if awserr, ok := err.(awserr.Error); ok {
		return "", fmt.Errorf("aws error while creating SQS queue: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return "", fmt.Errorf("error while creating SQS queue: %v", err)
	} else if resp == nil || resp.QueueURL == nil {
		return "", nil
	}

	return *resp.QueueURL, nil
}

func getMessage(config *aws.Config, queueURL string) (*sqs.Message, error) {
	svc := sqs.New(config)

	params := &sqs.ReceiveMessageInput{
		QueueURL:            aws.String(queueURL),
		MaxNumberOfMessages: aws.Long(1),
	}

	resp, err := svc.ReceiveMessage(params)
	if awserr, ok := err.(awserr.Error); ok {
		return nil, fmt.Errorf("aws error while receiving message from SQS: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return nil, fmt.Errorf("error while receiving message from SQS: %v", err)
	} else if resp == nil || len(resp.Messages) == 0 {
		return nil, nil
	}

	return resp.Messages[0], nil
}

func deleteMessage(config *aws.Config, queueURL string, receiptHandle string) error {
	svc := sqs.New(config)

	params := &sqs.DeleteMessageInput{
		QueueURL:      aws.String(queueURL),
		ReceiptHandle: aws.String(receiptHandle),
	}

	_, err := svc.DeleteMessage(params)
	if awserr, ok := err.(awserr.Error); ok {
		return fmt.Errorf("aws error while deleting message from SQS: %v %v", awserr.Code(), awserr.Message())
	} else if err != nil {
		return fmt.Errorf("error while deleting message from SQS: %v", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium 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 java

import (
	"bytes"
	"fmt"
	"log"
	"strings"

	"v.io/x/ref/lib/vdl/compile"
)

const arrayTmpl = header + `
// Source: {{.SourceFile}}

package {{.Package}};

{{ .Doc }}
@io.v.v23.vdl.GeneratedFromVdl(name = "{{.VdlTypeName}}")
@io.v.v23.vdl.ArrayLength({{.Length}})
{{ .AccessModifier }} class {{.Name}} extends io.v.v23.vdl.VdlArray<{{.ElemType}}> {
    private static final long serialVersionUID = 1L;

    /**
     * Vdl type for {@link {{.Name}}}.
     */
    public static final io.v.v23.vdl.VdlType VDL_TYPE =
            io.v.v23.vdl.Types.getVdlTypeFromReflect({{.Name}}.class);

    /**
     * Creates a new instance of {@link {{.Name}}} with the given underlying array.
     *
     * @param arr underlying array
     */
    public {{.Name}}({{.ElemType}}[] arr) {
        super(VDL_TYPE, arr);
    }

    /**
     * Creates a new zero-value instance of {@link {{.Name}}}.
     */
    public {{.Name}}() {
        this({{.ZeroValue}});
    }

    {{ if .ElemIsPrimitive }}
    /**
     * Creates a new instance of {@link {{.Name}}} with the given underlying array.
     *
     * @param arr underlying array
     */
    public {{.Name}}({{ .ElemPrimitiveType }}[] arr) {
        super(VDL_TYPE, convert(arr));
    }

    private static {{ .ElemType }}[] convert({{ .ElemPrimitiveType }}[] arr) {
        {{ .ElemType }}[] ret = new {{ .ElemType }}[arr.length];
        for (int i = 0; i < arr.length; ++i) {
            ret[i] = arr[i];
        }
        return ret;
    }
    {{ end }}
}
`

// genJavaArrayFile generates the Java class file for the provided named array type.
func genJavaArrayFile(tdef *compile.TypeDef, env *compile.Env) JavaFileInfo {
	name, access := javaTypeName(tdef, env)
	elemType := javaType(tdef.Type.Elem(), true, env)
	elems := strings.TrimSuffix(strings.Repeat(javaZeroValue(tdef.Type.Elem(), env)+", ", tdef.Type.Len()), ", ")
	zeroValue := fmt.Sprintf("new %s[] {%s}", elemType, elems)
	data := struct {
		AccessModifier    string
		Doc               string
		ElemType          string
		ElemIsPrimitive   bool
		ElemPrimitiveType string
		FileDoc           string
		Length            int
		Name              string
		Package           string
		SourceFile        string
		VdlTypeName       string
		VdlTypeString     string
		ZeroValue         string
	}{
		AccessModifier:    access,
		Doc:               javaDoc(tdef.Doc, tdef.DocSuffix),
		ElemType:          elemType,
		ElemIsPrimitive:   !isClass(tdef.Type.Elem(), env),
		ElemPrimitiveType: javaType(tdef.Type.Elem(), false, env),
		FileDoc:           tdef.File.Package.FileDoc,
		Length:            tdef.Type.Len(),
		Name:              name,
		Package:           javaPath(javaGenPkgPath(tdef.File.Package.GenPath)),
		SourceFile:        tdef.File.BaseName,
		VdlTypeName:       tdef.Type.Name(),
		VdlTypeString:     tdef.Type.String(),
		ZeroValue:         zeroValue,
	}
	var buf bytes.Buffer
	err := parseTmpl("array", arrayTmpl).Execute(&buf, data)
	if err != nil {
		log.Fatalf("vdl: couldn't execute array template: %v", err)
	}
	return JavaFileInfo{
		Name: name + ".java",
		Data: buf.Bytes(),
	}
}
<commit_msg>v.io/x/ref/lib/vdl/codegen/java: accessor for primitive arrays<commit_after>// Copyright 2015 The Vanadium 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 java

import (
	"bytes"
	"fmt"
	"log"
	"strings"

	"v.io/x/ref/lib/vdl/compile"
)

const arrayTmpl = header + `
// Source: {{.SourceFile}}

package {{.Package}};

{{ .Doc }}
@io.v.v23.vdl.GeneratedFromVdl(name = "{{.VdlTypeName}}")
@io.v.v23.vdl.ArrayLength({{.Length}})
{{ .AccessModifier }} class {{.Name}} extends io.v.v23.vdl.VdlArray<{{.ElemType}}> {
    private static final long serialVersionUID = 1L;

    /**
     * Vdl type for {@link {{.Name}}}.
     */
    public static final io.v.v23.vdl.VdlType VDL_TYPE =
            io.v.v23.vdl.Types.getVdlTypeFromReflect({{.Name}}.class);

    /**
     * Creates a new instance of {@link {{.Name}}} with the given underlying array.
     *
     * @param arr underlying array
     */
    public {{.Name}}({{.ElemType}}[] arr) {
        super(VDL_TYPE, arr);
    }

    /**
     * Creates a new zero-value instance of {@link {{.Name}}}.
     */
    public {{.Name}}() {
        this({{.ZeroValue}});
    }

    {{ if .ElemIsPrimitive }}
    /**
     * Creates a new instance of {@link {{.Name}}} with the given underlying array.
     *
     * @param arr underlying array
     */
    public {{.Name}}({{ .ElemPrimitiveType }}[] arr) {
        super(VDL_TYPE, convert(arr));
    }

	/**
	 * Converts the array into its primitive (array) type.
	 */
	public {{ .ElemPrimitiveType }}[] toPrimitiveArray() {
		{{ .ElemPrimitiveType }}[] ret = new {{ .ElemPrimitiveType }}[size()];
		for (int i = 0; i < size(); ++i) {
			ret[i] = get(i);
		}
		return ret;
	}

    private static {{ .ElemType }}[] convert({{ .ElemPrimitiveType }}[] arr) {
        {{ .ElemType }}[] ret = new {{ .ElemType }}[arr.length];
        for (int i = 0; i < arr.length; ++i) {
            ret[i] = arr[i];
        }
        return ret;
    }
    {{ end }}
}
`

// genJavaArrayFile generates the Java class file for the provided named array type.
func genJavaArrayFile(tdef *compile.TypeDef, env *compile.Env) JavaFileInfo {
	name, access := javaTypeName(tdef, env)
	elemType := javaType(tdef.Type.Elem(), true, env)
	elems := strings.TrimSuffix(strings.Repeat(javaZeroValue(tdef.Type.Elem(), env)+", ", tdef.Type.Len()), ", ")
	zeroValue := fmt.Sprintf("new %s[] {%s}", elemType, elems)
	data := struct {
		AccessModifier    string
		Doc               string
		ElemType          string
		ElemIsPrimitive   bool
		ElemPrimitiveType string
		FileDoc           string
		Length            int
		Name              string
		Package           string
		SourceFile        string
		VdlTypeName       string
		VdlTypeString     string
		ZeroValue         string
	}{
		AccessModifier:    access,
		Doc:               javaDoc(tdef.Doc, tdef.DocSuffix),
		ElemType:          elemType,
		ElemIsPrimitive:   !isClass(tdef.Type.Elem(), env),
		ElemPrimitiveType: javaType(tdef.Type.Elem(), false, env),
		FileDoc:           tdef.File.Package.FileDoc,
		Length:            tdef.Type.Len(),
		Name:              name,
		Package:           javaPath(javaGenPkgPath(tdef.File.Package.GenPath)),
		SourceFile:        tdef.File.BaseName,
		VdlTypeName:       tdef.Type.Name(),
		VdlTypeString:     tdef.Type.String(),
		ZeroValue:         zeroValue,
	}
	var buf bytes.Buffer
	err := parseTmpl("array", arrayTmpl).Execute(&buf, data)
	if err != nil {
		log.Fatalf("vdl: couldn't execute array template: %v", err)
	}
	return JavaFileInfo{
		Name: name + ".java",
		Data: buf.Bytes(),
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 tsuru 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 dockercommon

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/fsouza/go-dockerclient"
	"github.com/fsouza/go-dockerclient/testing"
	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/app/image"
	"gopkg.in/check.v1"
)

func (s *S) TestPrepareImageForDeploy(c *check.C) {
	srv, err := testing.NewServer("0.0.0.0:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer srv.Stop()
	a := &app.App{Name: "myapp"}
	cli, err := docker.NewClient(srv.URL())
	c.Assert(err, check.IsNil)
	baseImgName := "baseImg"
	err = cli.PullImage(docker.PullImageOptions{Repository: baseImgName}, docker.AuthConfiguration{})
	c.Assert(err, check.IsNil)
	buf := bytes.Buffer{}
	args := PrepareImageArgs{
		Client:      cli,
		App:         a,
		ProcfileRaw: "web: myapp run",
		ImageId:     baseImgName,
		Out:         &buf,
	}
	newImg, err := PrepareImageForDeploy(args)
	c.Assert(err, check.IsNil)
	c.Assert(newImg, check.Equals, "my.registry/tsuru/app-myapp:v1")
	c.Assert(buf.String(), check.Equals, `---- Inspecting image "baseImg" ----
  ---> Process "web" found with commands: ["myapp run"]
---- Pushing image "my.registry/tsuru/app-myapp:v1" to tsuru ----
Pushing...
Pushed
`)
	imd, err := image.GetImageCustomData(newImg)
	c.Assert(err, check.IsNil)
	c.Assert(imd, check.DeepEquals, image.ImageMetadata{
		Name:            "my.registry/tsuru/app-myapp:v1",
		Processes:       map[string][]string{"web": {"myapp run"}},
		CustomData:      map[string]interface{}{},
		LegacyProcesses: map[string]string{},
	})
}

func (s *S) TestPrepareImageForDeployNoProcfile(c *check.C) {
	srv, err := testing.NewServer("0.0.0.0:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer srv.Stop()
	a := &app.App{Name: "myapp"}
	cli, err := docker.NewClient(srv.URL())
	c.Assert(err, check.IsNil)
	baseImgName := "baseImg"
	err = cli.PullImage(docker.PullImageOptions{Repository: baseImgName}, docker.AuthConfiguration{})
	c.Assert(err, check.IsNil)
	srv.CustomHandler(fmt.Sprintf("/images/%s/json", baseImgName), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		response := docker.Image{
			Config: &docker.Config{
				Entrypoint:   []string{"/bin/sh"},
				Cmd:          []string{"python", "test file.py"},
				ExposedPorts: map[docker.Port]struct{}{"3000/tcp": {}},
			},
		}
		j, _ := json.Marshal(response)
		w.Write(j)
	}))
	buf := bytes.Buffer{}
	args := PrepareImageArgs{
		Client:      cli,
		App:         a,
		ProcfileRaw: "",
		ImageId:     baseImgName,
		Out:         &buf,
	}
	newImg, err := PrepareImageForDeploy(args)
	c.Assert(err, check.IsNil)
	c.Assert(newImg, check.Equals, "my.registry/tsuru/app-myapp:v1")
	c.Assert(buf.String(), check.Equals, `---- Inspecting image "baseImg" ----
  ---> Procfile not found, using entrypoint and cmd
  ---> Process "web" found with commands: ["/bin/sh" "python" "test file.py"]
---- Pushing image "my.registry/tsuru/app-myapp:v1" to tsuru ----
Pushing...
Pushed
`)
	imd, err := image.GetImageCustomData(newImg)
	c.Assert(err, check.IsNil)
	c.Assert(imd, check.DeepEquals, image.ImageMetadata{
		Name:            "my.registry/tsuru/app-myapp:v1",
		Processes:       map[string][]string{"web": {"/bin/sh", "python", "test file.py"}},
		CustomData:      map[string]interface{}{},
		LegacyProcesses: map[string]string{},
		ExposedPort:     "3000/tcp",
	})
}

func (s *S) TestUploadToContainer(c *check.C) {
	srv, err := testing.NewServer("0.0.0.0:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer srv.Stop()
	cli, err := docker.NewClient(srv.URL())
	c.Assert(err, check.IsNil)
	baseImgName := "baseImg"
	err = cli.PullImage(docker.PullImageOptions{Repository: baseImgName}, docker.AuthConfiguration{})
	c.Assert(err, check.IsNil)
	cont, err := cli.CreateContainer(docker.CreateContainerOptions{
		Config: &docker.Config{
			Image: baseImgName,
		},
		HostConfig: &docker.HostConfig{},
	})
	c.Assert(err, check.IsNil)
	err = cli.StartContainer(cont.ID, nil)
	c.Assert(err, check.IsNil)
	buf := strings.NewReader("my data")
	tarFile := AddDeployTarFile(buf, int64(buf.Len()), "archive.tar.gz")
	imgId, path, err := UploadToContainer(cli, cont.ID, tarFile)
	c.Assert(err, check.IsNil)
	c.Assert(imgId, check.Matches, "^img-.{32}$")
	c.Assert(path, check.Equals, "file:///home/application/archive.tar.gz")
}

func (s *S) TestWaitDocker(c *check.C) {
	server, err := testing.NewServer("127.0.0.1:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer server.Stop()
	client, err := docker.NewClient(server.URL())
	c.Assert(err, check.IsNil)
	err = WaitDocker(client)
	c.Assert(err, check.IsNil)
	server.CustomHandler("/_ping", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}))
	err = WaitDocker(client)
	c.Assert(err, check.NotNil)
	config.Set("docker:api-timeout", 1)
	defer config.Unset("docker:api-timeout")
	client, err = docker.NewClient("http://169.254.169.254:2375/")
	c.Assert(err, check.IsNil)
	err = WaitDocker(client)
	c.Assert(err, check.NotNil)
	expectedMsg := `Docker API at "http://169.254.169.254:2375/" didn't respond after 1 seconds`
	c.Assert(err.Error(), check.Equals, expectedMsg)
}
<commit_msg>provision/dockercommon: listen only on localhost in tests<commit_after>// Copyright 2016 tsuru 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 dockercommon

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	"github.com/fsouza/go-dockerclient"
	"github.com/fsouza/go-dockerclient/testing"
	"github.com/tsuru/config"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/app/image"
	"gopkg.in/check.v1"
)

func (s *S) TestPrepareImageForDeploy(c *check.C) {
	srv, err := testing.NewServer("127.0.0.1:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer srv.Stop()
	a := &app.App{Name: "myapp"}
	cli, err := docker.NewClient(srv.URL())
	c.Assert(err, check.IsNil)
	baseImgName := "baseImg"
	err = cli.PullImage(docker.PullImageOptions{Repository: baseImgName}, docker.AuthConfiguration{})
	c.Assert(err, check.IsNil)
	buf := bytes.Buffer{}
	args := PrepareImageArgs{
		Client:      cli,
		App:         a,
		ProcfileRaw: "web: myapp run",
		ImageId:     baseImgName,
		Out:         &buf,
	}
	newImg, err := PrepareImageForDeploy(args)
	c.Assert(err, check.IsNil)
	c.Assert(newImg, check.Equals, "my.registry/tsuru/app-myapp:v1")
	c.Assert(buf.String(), check.Equals, `---- Inspecting image "baseImg" ----
  ---> Process "web" found with commands: ["myapp run"]
---- Pushing image "my.registry/tsuru/app-myapp:v1" to tsuru ----
Pushing...
Pushed
`)
	imd, err := image.GetImageCustomData(newImg)
	c.Assert(err, check.IsNil)
	c.Assert(imd, check.DeepEquals, image.ImageMetadata{
		Name:            "my.registry/tsuru/app-myapp:v1",
		Processes:       map[string][]string{"web": {"myapp run"}},
		CustomData:      map[string]interface{}{},
		LegacyProcesses: map[string]string{},
	})
}

func (s *S) TestPrepareImageForDeployNoProcfile(c *check.C) {
	srv, err := testing.NewServer("127.0.0.1:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer srv.Stop()
	a := &app.App{Name: "myapp"}
	cli, err := docker.NewClient(srv.URL())
	c.Assert(err, check.IsNil)
	baseImgName := "baseImg"
	err = cli.PullImage(docker.PullImageOptions{Repository: baseImgName}, docker.AuthConfiguration{})
	c.Assert(err, check.IsNil)
	srv.CustomHandler(fmt.Sprintf("/images/%s/json", baseImgName), http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		response := docker.Image{
			Config: &docker.Config{
				Entrypoint:   []string{"/bin/sh"},
				Cmd:          []string{"python", "test file.py"},
				ExposedPorts: map[docker.Port]struct{}{"3000/tcp": {}},
			},
		}
		j, _ := json.Marshal(response)
		w.Write(j)
	}))
	buf := bytes.Buffer{}
	args := PrepareImageArgs{
		Client:      cli,
		App:         a,
		ProcfileRaw: "",
		ImageId:     baseImgName,
		Out:         &buf,
	}
	newImg, err := PrepareImageForDeploy(args)
	c.Assert(err, check.IsNil)
	c.Assert(newImg, check.Equals, "my.registry/tsuru/app-myapp:v1")
	c.Assert(buf.String(), check.Equals, `---- Inspecting image "baseImg" ----
  ---> Procfile not found, using entrypoint and cmd
  ---> Process "web" found with commands: ["/bin/sh" "python" "test file.py"]
---- Pushing image "my.registry/tsuru/app-myapp:v1" to tsuru ----
Pushing...
Pushed
`)
	imd, err := image.GetImageCustomData(newImg)
	c.Assert(err, check.IsNil)
	c.Assert(imd, check.DeepEquals, image.ImageMetadata{
		Name:            "my.registry/tsuru/app-myapp:v1",
		Processes:       map[string][]string{"web": {"/bin/sh", "python", "test file.py"}},
		CustomData:      map[string]interface{}{},
		LegacyProcesses: map[string]string{},
		ExposedPort:     "3000/tcp",
	})
}

func (s *S) TestUploadToContainer(c *check.C) {
	srv, err := testing.NewServer("127.0.0.1:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer srv.Stop()
	cli, err := docker.NewClient(srv.URL())
	c.Assert(err, check.IsNil)
	baseImgName := "baseImg"
	err = cli.PullImage(docker.PullImageOptions{Repository: baseImgName}, docker.AuthConfiguration{})
	c.Assert(err, check.IsNil)
	cont, err := cli.CreateContainer(docker.CreateContainerOptions{
		Config: &docker.Config{
			Image: baseImgName,
		},
		HostConfig: &docker.HostConfig{},
	})
	c.Assert(err, check.IsNil)
	err = cli.StartContainer(cont.ID, nil)
	c.Assert(err, check.IsNil)
	buf := strings.NewReader("my data")
	tarFile := AddDeployTarFile(buf, int64(buf.Len()), "archive.tar.gz")
	imgId, path, err := UploadToContainer(cli, cont.ID, tarFile)
	c.Assert(err, check.IsNil)
	c.Assert(imgId, check.Matches, "^img-.{32}$")
	c.Assert(path, check.Equals, "file:///home/application/archive.tar.gz")
}

func (s *S) TestWaitDocker(c *check.C) {
	server, err := testing.NewServer("127.0.0.1:0", nil, nil)
	c.Assert(err, check.IsNil)
	defer server.Stop()
	client, err := docker.NewClient(server.URL())
	c.Assert(err, check.IsNil)
	err = WaitDocker(client)
	c.Assert(err, check.IsNil)
	server.CustomHandler("/_ping", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}))
	err = WaitDocker(client)
	c.Assert(err, check.NotNil)
	config.Set("docker:api-timeout", 1)
	defer config.Unset("docker:api-timeout")
	client, err = docker.NewClient("http://169.254.169.254:2375/")
	c.Assert(err, check.IsNil)
	err = WaitDocker(client)
	c.Assert(err, check.NotNil)
	expectedMsg := `Docker API at "http://169.254.169.254:2375/" didn't respond after 1 seconds`
	c.Assert(err.Error(), check.Equals, expectedMsg)
}
<|endoftext|>


<commit_before>package shell

import (
	"bytes"
	"io"
	"testing"
)

func TestUnixReader_impl(t *testing.T) {
	var raw interface{}
	raw = new(UnixReader)
	if _, ok := raw.(io.Reader); !ok {
		t.Fatal("should be reader")
	}
}

func TestUnixReader(t *testing.T) {
	input := "one\r\ntwo\nthree\r\n"
	expected := "one\ntwo\nthree\n"

	r := &UnixReader{
		Reader: bytes.NewReader([]byte(input)),
	}

	result := new(bytes.Buffer)
	if _, err := io.Copy(result, r); err != nil {
		t.Fatalf("err: %s", err)
	}

	if result.String() != expected {
		t.Fatalf("bad: %#v", result.String())
	}
}
<commit_msg>provisioner/shell: add another UnixReader test for sanity<commit_after>package shell

import (
	"bytes"
	"io"
	"testing"
)

func TestUnixReader_impl(t *testing.T) {
	var raw interface{}
	raw = new(UnixReader)
	if _, ok := raw.(io.Reader); !ok {
		t.Fatal("should be reader")
	}
}

func TestUnixReader(t *testing.T) {
	input := "one\r\ntwo\nthree\r\n"
	expected := "one\ntwo\nthree\n"

	r := &UnixReader{
		Reader: bytes.NewReader([]byte(input)),
	}

	result := new(bytes.Buffer)
	if _, err := io.Copy(result, r); err != nil {
		t.Fatalf("err: %s", err)
	}

	if result.String() != expected {
		t.Fatalf("bad: %#v", result.String())
	}
}

func TestUnixReader_unixOnly(t *testing.T) {
	input := "one\ntwo\nthree\n"
	expected := "one\ntwo\nthree\n"

	r := &UnixReader{
		Reader: bytes.NewReader([]byte(input)),
	}

	result := new(bytes.Buffer)
	if _, err := io.Copy(result, r); err != nil {
		t.Fatalf("err: %s", err)
	}

	if result.String() != expected {
		t.Fatalf("bad: %#v", result.String())
	}
}
<|endoftext|>


<commit_before>package verify

import (
	"fmt"
	"os"
	u "os/user"

	"crypto/rand"
	"encoding/base64"

	"io/ioutil"

	"log"
	fp "path/filepath"
)

type AuthTokenStorage struct {
	Path string
}

func NewAuthTokenStorage() (ts *AuthTokenStorage, err error) {
	//TODO: os/user is not supported for cross compiling, find diff way of doing this
	user, err := u.Current()
	if err != nil {
		fmt.Println("error getting current user: ", err)
		return
	}

	return &AuthTokenStorage{
		Path: fp.Join(user.HomeDir, "retailops"),
	}, nil
}

func (ats *AuthTokenStorage) CreateDirectoryIfMissing() (err error) {
	_, err = os.Open(ats.Path)
	if err != nil {
		if pathErr, ok := err.(*os.PathError); ok && pathErr.Err.Error() == "no such file or directory" {
			fmt.Println("Creating directory...")
			return ats.doFolderCreate()
		}
		return
	}

	return
}

func (ats *AuthTokenStorage) OverwriteIntegrationAuthKey(token string) (err error) {
	tokenPath := fp.Join(ats.Path, "integration_auth_token")
	var file *os.File
	if _, err = os.Stat(tokenPath); os.IsNotExist(err) {
		file, err = os.Create(tokenPath)
		if err != nil {
			return
		}
	} else {
		file, err = os.OpenFile(tokenPath, os.O_WRONLY|os.O_TRUNC, 0)
		if err != nil {
			return
		}
	}

	n, err := file.WriteString(token)
	if err != nil {
		return
	}

	if n != len(token) {
		log.Fatal("huh")
	}

	return
}

func (ats *AuthTokenStorage) GenerateIntegrationAuthToken() (err error) {
	tokenPath := fp.Join(ats.Path, "integration_auth_token")
	fmt.Println("Token path: ", tokenPath)

	token, err := randomToken()
	if err != nil {
		return
	}

	file, err := os.Create(tokenPath)
	if err != nil {
		fmt.Println("Could not create folder: ", err)
		return
	}

	_, err = file.WriteString(token)
	if err != nil {
		fmt.Println("Could not create file: ", err)
	}

	return
}

func (ats *AuthTokenStorage) ReadToken() (token string, err error) {
	tokenPath := fp.Join(ats.Path, "integration_auth_token")
	file, err := os.Open(tokenPath)
	if err != nil {
		return
	}

	buf, err := ioutil.ReadAll(file)
	if err != nil {
		return
	}

	token = string(buf)
	return
}

func randomToken() (token string, err error) {
	randBytes := make([]byte, 24, 24)
	read, err := rand.Read(randBytes)
	if err != nil {
		err = fmt.Errorf("could not generate all rand bytes needed. only read %d of 18", read)
		return
	}
	token = base64.StdEncoding.EncodeToString(randBytes)

	return
}

func (ats *AuthTokenStorage) doFolderCreate() (err error) {
	fmt.Println("Creating folder: ", ats.Path)
	err = os.Mkdir(ats.Path, 0700)
	return
}
<commit_msg>handle missing dir across operating systems<commit_after>package verify

import (
	"fmt"
	"os"
	u "os/user"

	"crypto/rand"
	"encoding/base64"

	"io/ioutil"

	"log"
	fp "path/filepath"
)

type AuthTokenStorage struct {
	Path string
}

func NewAuthTokenStorage() (ts *AuthTokenStorage, err error) {
	//TODO: os/user is not supported for cross compiling, find diff way of doing this
	user, err := u.Current()
	if err != nil {
		fmt.Println("error getting current user: ", err)
		return
	}

	return &AuthTokenStorage{
		Path: fp.Join(user.HomeDir, "retailops"),
	}, nil
}

func (ats *AuthTokenStorage) CreateDirectoryIfMissing() (err error) {
	_, err = os.Open(ats.Path)
	if err != nil {
		if pathErr, ok := err.(*os.PathError); ok && os.IsNotExist(pathErr) {
			fmt.Println("Creating directory...")
			return ats.doFolderCreate()
		}
		return
	}

	return
}

func (ats *AuthTokenStorage) OverwriteIntegrationAuthKey(token string) (err error) {
	tokenPath := fp.Join(ats.Path, "integration_auth_token")
	var file *os.File
	if _, err = os.Stat(tokenPath); os.IsNotExist(err) {
		file, err = os.Create(tokenPath)
		if err != nil {
			return
		}
	} else {
		file, err = os.OpenFile(tokenPath, os.O_WRONLY|os.O_TRUNC, 0)
		if err != nil {
			return
		}
	}

	n, err := file.WriteString(token)
	if err != nil {
		return
	}

	if n != len(token) {
		log.Fatal("huh")
	}

	return
}

func (ats *AuthTokenStorage) GenerateIntegrationAuthToken() (err error) {
	tokenPath := fp.Join(ats.Path, "integration_auth_token")
	fmt.Println("Token path: ", tokenPath)

	token, err := randomToken()
	if err != nil {
		return
	}

	file, err := os.Create(tokenPath)
	if err != nil {
		fmt.Println("Could not create folder: ", err)
		return
	}

	_, err = file.WriteString(token)
	if err != nil {
		fmt.Println("Could not create file: ", err)
	}

	return
}

func (ats *AuthTokenStorage) ReadToken() (token string, err error) {
	tokenPath := fp.Join(ats.Path, "integration_auth_token")
	file, err := os.Open(tokenPath)
	if err != nil {
		return
	}

	buf, err := ioutil.ReadAll(file)
	if err != nil {
		return
	}

	token = string(buf)
	return
}

func randomToken() (token string, err error) {
	randBytes := make([]byte, 24, 24)
	read, err := rand.Read(randBytes)
	if err != nil {
		err = fmt.Errorf("could not generate all rand bytes needed. only read %d of 18", read)
		return
	}
	token = base64.StdEncoding.EncodeToString(randBytes)

	return
}

func (ats *AuthTokenStorage) doFolderCreate() (err error) {
	fmt.Println("Creating folder: ", ats.Path)
	err = os.Mkdir(ats.Path, 0700)
	return
}
<|endoftext|>


<commit_before>package main

import (
	"../cmd"
	"../framework"
	"fmt"
	"github.com/bwmarrin/discordgo"
	"strings"
)

// const (
// 	PREFIX = "!"
// )

var (
	conf       *framework.Config
	CmdHandler *framework.CommandHandler
	Sessions   *framework.SessionManager
	youtube    *framework.Youtube
	botId      string
	PREFIX     string
)

func init() {
	conf = framework.LoadConfig("config.json")
	PREFIX = conf.Prefix

}

func main() {
	CmdHandler = framework.NewCommandHandler()
	registerCommands()
	Sessions = framework.NewSessionManager()
	youtube = &framework.Youtube{Conf: conf}
	discord, err := discordgo.New(conf.BotToken)
	if err != nil {
		fmt.Println("Error creating discord session,", err)
		return
	}
	if conf.UseSharding {
		discord.ShardID = conf.ShardId
		discord.ShardCount = conf.ShardCount
	}
	usr, err := discord.User("@me")
	if err != nil {
		fmt.Println("Error obtaining account details,", err)
		return
	}
	botId = usr.ID
	discord.AddHandler(commandHandler)
	discord.AddHandler(func(discord *discordgo.Session, ready *discordgo.Ready) {
		discord.UpdateStatus(0, "A baby AI")
		guilds := discord.State.Guilds
		fmt.Println("Ready with", len(guilds), "guilds.")
	})
	err = discord.Open()
	if err != nil {
		fmt.Println("Error opening connection,", err)
		return
	}
	fmt.Println("Started")
	<-make(chan struct{})
}

func commandHandler(discord *discordgo.Session, message *discordgo.MessageCreate) {
	user := message.Author
	if user.ID == botId || user.Bot {
		return
	}
	content := message.Content
	if len(content) <= len(PREFIX) {
		return
	}
	if content[:len(PREFIX)] != PREFIX {
		return
	}
	content = content[len(PREFIX):]
	if len(content) < 1 {
		return
	}
	args := strings.Fields(content)
	name := strings.ToLower(args[0])
	command, found := CmdHandler.Get(name)
	if !found {
		return
	}
	channel, err := discord.State.Channel(message.ChannelID)
	if err != nil {
		fmt.Println("Error getting channel,", err)
		return
	}
	guild, err := discord.State.Guild(channel.GuildID)
	if err != nil {
		fmt.Println("Error getting guild,", err)
		return
	}
	ctx := framework.NewContext(discord, guild, channel, user, message, conf, CmdHandler, Sessions, youtube)
	ctx.Args = args[1:]
	c := *command
	c(*ctx)
}

func registerCommands() {
	CmdHandler.Register("help", cmd.HelpCommand)
	CmdHandler.Register("admin", cmd.AdminCommand)
	CmdHandler.Register("join", cmd.JoinCommand)
	CmdHandler.Register("leave", cmd.LeaveCommand)
	CmdHandler.Register("play", cmd.PlayCommand)
	CmdHandler.Register("stop", cmd.StopCommand)
	CmdHandler.Register("info", cmd.InfoCommand)
	CmdHandler.Register("add", cmd.AddCommand)
	CmdHandler.Register("skip", cmd.SkipCommand)
	CmdHandler.Register("queue", cmd.QueueCommand)
	CmdHandler.Register("eval", cmd.EvalCommand)
	CmdHandler.Register("debug", cmd.DebugCommand)
	CmdHandler.Register("clear", cmd.ClearCommand)
	CmdHandler.Register("current", cmd.CurrentCommand)
	CmdHandler.Register("youtube", cmd.YoutubeCommand)
    CmdHandler.Register("shuffle", cmd.ShuffleCommand)
    CmdHandler.Register("pausequeue", cmd.PauseCommand)
    CmdHandler.Register("pick", cmd.PickCommand)
}
<commit_msg>cleanup<commit_after>package main

import (
	"../cmd"
	"../framework"
	"fmt"
	"github.com/bwmarrin/discordgo"
	"strings"
)

var (
	conf       *framework.Config
	CmdHandler *framework.CommandHandler
	Sessions   *framework.SessionManager
	youtube    *framework.Youtube
	botId      string
	PREFIX     string
)

func init() {
	conf = framework.LoadConfig("config.json")
	PREFIX = conf.Prefix

}

func main() {
	CmdHandler = framework.NewCommandHandler()
	registerCommands()
	Sessions = framework.NewSessionManager()
	youtube = &framework.Youtube{Conf: conf}
	discord, err := discordgo.New(conf.BotToken)
	if err != nil {
		fmt.Println("Error creating discord session,", err)
		return
	}
	if conf.UseSharding {
		discord.ShardID = conf.ShardId
		discord.ShardCount = conf.ShardCount
	}
	usr, err := discord.User("@me")
	if err != nil {
		fmt.Println("Error obtaining account details,", err)
		return
	}
	botId = usr.ID
	discord.AddHandler(commandHandler)
	discord.AddHandler(func(discord *discordgo.Session, ready *discordgo.Ready) {
		discord.UpdateStatus(0, "A baby AI")
		guilds := discord.State.Guilds
		fmt.Println("Ready with", len(guilds), "guilds.")
	})
	err = discord.Open()
	if err != nil {
		fmt.Println("Error opening connection,", err)
		return
	}
	fmt.Println("Started")
	<-make(chan struct{})
}

func commandHandler(discord *discordgo.Session, message *discordgo.MessageCreate) {
	user := message.Author
	if user.ID == botId || user.Bot {
		return
	}
	content := message.Content
	if len(content) <= len(PREFIX) {
		return
	}
	if content[:len(PREFIX)] != PREFIX {
		return
	}
	content = content[len(PREFIX):]
	if len(content) < 1 {
		return
	}
	args := strings.Fields(content)
	name := strings.ToLower(args[0])
	command, found := CmdHandler.Get(name)
	if !found {
		return
	}
	channel, err := discord.State.Channel(message.ChannelID)
	if err != nil {
		fmt.Println("Error getting channel,", err)
		return
	}
	guild, err := discord.State.Guild(channel.GuildID)
	if err != nil {
		fmt.Println("Error getting guild,", err)
		return
	}
	ctx := framework.NewContext(discord, guild, channel, user, message, conf, CmdHandler, Sessions, youtube)
	ctx.Args = args[1:]
	c := *command
	c(*ctx)
}

func registerCommands() {
	CmdHandler.Register("help", cmd.HelpCommand)
	CmdHandler.Register("admin", cmd.AdminCommand)
	CmdHandler.Register("join", cmd.JoinCommand)
	CmdHandler.Register("leave", cmd.LeaveCommand)
	CmdHandler.Register("play", cmd.PlayCommand)
	CmdHandler.Register("stop", cmd.StopCommand)
	CmdHandler.Register("info", cmd.InfoCommand)
	CmdHandler.Register("add", cmd.AddCommand)
	CmdHandler.Register("skip", cmd.SkipCommand)
	CmdHandler.Register("queue", cmd.QueueCommand)
	CmdHandler.Register("eval", cmd.EvalCommand)
	CmdHandler.Register("debug", cmd.DebugCommand)
	CmdHandler.Register("clear", cmd.ClearCommand)
	CmdHandler.Register("current", cmd.CurrentCommand)
	CmdHandler.Register("youtube", cmd.YoutubeCommand)
    CmdHandler.Register("shuffle", cmd.ShuffleCommand)
    CmdHandler.Register("pausequeue", cmd.PauseCommand)
    CmdHandler.Register("pick", cmd.PickCommand)
}
<|endoftext|>


<commit_before>package logformatter

import (
	"encoding/json"
	"fmt"

	"github.com/Sirupsen/logrus"
	"os"
	"time"
)

// LogstashFormatter generates json in logstash format.
// Logstash site: http://logstash.net/
type LogstashFormatter struct {

	//Type of the fields
	Type string // if not empty use for logstash type field.

	//Env is the environment on which the application is running
	Env string

	//ServiceName will be by default guble
	ServiceName string

	//ApplicationType will be  by default "service".Other values could be "service", "system", "appserver", "webserver"
	ApplicationType string

	//LogType will be by default  application.Other possible values "access", "error", "application", "system"
	LogType string

	// TimestampFormat sets the format used for timestamps.
	TimestampFormat string
}

// Format the logrus entry to a byte slice, or return an error.
func (f *LogstashFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	fields := make(logrus.Fields)
	for k, v := range entry.Data {
		fields[k] = v
	}

	timeStampFormat := f.TimestampFormat
	if timeStampFormat == "" {
		timeStampFormat = time.RFC3339
	}

	fields["@timestamp"] = entry.Time.Format(timeStampFormat)

	if f.ServiceName != "" {
		fields["service"] = f.ServiceName
	} else {
		fields["service"] = "guble"
	}

	if f.ApplicationType != "" {
		fields["application_type"] = f.ServiceName
	} else {
		fields["application_type"] = "service"
	}

	if f.LogType != "" {
		fields["log_type"] = f.LogType
	} else {
		fields["log_type"] = "application"
	}

	fields["environment"] = f.Env

	// set message field
	v, ok := entry.Data["message"]
	if ok {
		fields["fields.message"] = v
	}
	fields["message"] = entry.Message

	// set level field
	v, ok = entry.Data["level"]
	if ok {
		fields["fields.level"] = v
	}
	fields["loglevel"] = entry.Level.String()

	//set host field
	hostname, err := os.Hostname()
	if err == nil {
		fields["host"] = hostname
	} else {
		fields["host"] = ""
	}

	// set type field
	if f.Type != "" {
		v, ok = entry.Data["type"]
		if ok {
			fields["fields.type"] = v
		}
		fields["type"] = f.Type
	}

	serialized, err := json.Marshal(fields)
	if err != nil {
		return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
	}
	return append(serialized, '\n'), nil
}
<commit_msg>improving logstash formatter: error in fields, message -> msg, prefixing host if field is already existing, consts for defaults, minor refactoring<commit_after>package logformatter

import (
	"encoding/json"
	"fmt"

	"github.com/Sirupsen/logrus"
	"os"
	"time"
)

const (
	defaultServiceName     = "guble"
	defaultLogType         = "application"
	defaultApplicationType = "service"
)

// LogstashFormatter generates json in logstash format.
// Logstash site: http://logstash.net/
type LogstashFormatter struct {

	//Type of the fields
	Type string // if not empty use for logstash type field.

	//Env is the environment on which the application is running
	Env string

	//ServiceName will be by default guble
	ServiceName string

	//ApplicationType will be  by default "service". Other values could be "service", "system", "appserver", "webserver"
	ApplicationType string

	//LogType will be by default  application. Other possible values "access", "error", "application", "system"
	LogType string

	//TimestampFormat sets the format used for timestamps.
	TimestampFormat string
}

// Format the logrus entry to a byte slice, or return an error.
func (f *LogstashFormatter) Format(entry *logrus.Entry) ([]byte, error) {
	fields := make(logrus.Fields)

	for k, v := range entry.Data {
		switch v := v.(type) {
		case error:
			// Otherwise errors are ignored by `encoding/json`
			// https://github.com/Sirupsen/logrus/issues/137
			// https://github.com/sirupsen/logrus/issues/377
			fields[k] = v.Error()
		default:
			fields[k] = v
		}
	}

	fields["environment"] = f.Env

	timeStampFormat := f.TimestampFormat
	if timeStampFormat == "" {
		timeStampFormat = time.RFC3339
	}

	fields["@timestamp"] = entry.Time.Format(timeStampFormat)

	if f.ServiceName != "" {
		fields["service"] = f.ServiceName
	} else {
		fields["service"] = defaultServiceName
	}

	if f.ApplicationType != "" {
		fields["application_type"] = f.ServiceName
	} else {
		fields["application_type"] = defaultApplicationType
	}

	if f.LogType != "" {
		fields["log_type"] = f.LogType
	} else {
		fields["log_type"] = defaultLogType
	}

	// set message field, prefixing fields clashes
	if v, ok := entry.Data["msg"]; ok {
		fields["fields.msg"] = v
	}
	fields["msg"] = entry.Message

	// set level field, prefixing fields clashes
	if v, ok := entry.Data["loglevel"]; ok {
		fields["fields.loglevel"] = v
	}
	fields["loglevel"] = entry.Level.String()

	//set host field, prefixing fields clashes
	if v, ok := entry.Data["host"]; ok {
		fields["fields.host"] = v
	}

	if hostname, err := os.Hostname(); err == nil {
		fields["host"] = hostname
	}

	// set type field
	if f.Type != "" {
		if v, ok := entry.Data["type"]; ok {
			fields["fields.type"] = v
		}
		fields["type"] = f.Type
	}

	serialized, err := json.Marshal(fields)
	if err != nil {
		return nil, fmt.Errorf("Failed to marshal fields to JSON, %v", err)
	}
	return append(serialized, '\n'), nil
}
<|endoftext|>