<commit_before>package main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/web/api/v1"
	"github.com/weaveworks/scope/common/middleware"
	"golang.org/x/net/context"

	"github.com/weaveworks/cortex"
	"github.com/weaveworks/cortex/chunk"
	"github.com/weaveworks/cortex/distributor"
	"github.com/weaveworks/cortex/ingester"
	"github.com/weaveworks/cortex/querier"
	"github.com/weaveworks/cortex/ring"
	"github.com/weaveworks/cortex/ui"
	"github.com/weaveworks/cortex/user"
)

const (
	modeDistributor = "distributor"
	modeIngester    = "ingester"

	infName          = "eth0"
	userIDHeaderName = "X-Scope-OrgID"
)

var (
	requestDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace: "cortex",
		Name:      "request_duration_seconds",
		Help:      "Time (in seconds) spent serving HTTP requests.",

		// Cortex latency can be very low (ingesters typically take a few ms to
		// process a request), all the way to very high (queries can take tens of
		// second).  As its important, we use 10 buckets.  Smallest is 128us
		// biggest is 130s.
		Buckets: prometheus.ExponentialBuckets(0.000128, 4, 10),
	}, []string{"method", "route", "status_code", "ws"})
)

func init() {
	prometheus.MustRegister(requestDuration)
}

type cfg struct {
	mode                 string
	listenPort           int
	consulHost           string
	consulPrefix         string
	s3URL                string
	dynamodbURL          string
	dynamodbCreateTables bool
	memcachedHostname    string
	memcachedTimeout     time.Duration
	memcachedExpiration  time.Duration
	memcachedService     string
	remoteTimeout        time.Duration
	numTokens            int
	logSuccess           bool

	ingesterConfig    ingester.Config
	distributorConfig distributor.Config
}

func main() {
	var cfg cfg
	flag.StringVar(&cfg.mode, "mode", modeDistributor, "Mode (distributor, ingester).")
	flag.IntVar(&cfg.listenPort, "web.listen-port", 9094, "HTTP server listen port.")
	flag.StringVar(&cfg.consulHost, "consul.hostname", "localhost:8500", "Hostname and port of Consul.")
	flag.StringVar(&cfg.consulPrefix, "consul.prefix", "collectors/", "Prefix for keys in Consul.")
	flag.StringVar(&cfg.s3URL, "s3.url", "localhost:4569", "S3 endpoint URL.")
	flag.StringVar(&cfg.dynamodbURL, "dynamodb.url", "localhost:8000", "DynamoDB endpoint URL.")
	flag.BoolVar(&cfg.dynamodbCreateTables, "dynamodb.create-tables", false, "Create required DynamoDB tables on startup.")
	flag.StringVar(&cfg.memcachedHostname, "memcached.hostname", "", "Hostname for memcached service to use when caching chunks. If empty, no memcached will be used.")
	flag.DurationVar(&cfg.memcachedTimeout, "memcached.timeout", 100*time.Millisecond, "Maximum time to wait before giving up on memcached requests.")
	flag.DurationVar(&cfg.memcachedExpiration, "memcached.expiration", 0, "How long chunks stay in the memcache.")
	flag.StringVar(&cfg.memcachedService, "memcached.service", "memcached", "SRV service used to discover memcache servers.")
	flag.DurationVar(&cfg.remoteTimeout, "remote.timeout", 5*time.Second, "Timeout for downstream ingesters.")
	flag.DurationVar(&cfg.ingesterConfig.FlushCheckPeriod, "ingester.flush-period", 1*time.Minute, "Period with which to attempt to flush chunks.")
	flag.DurationVar(&cfg.ingesterConfig.RateUpdatePeriod, "ingester.rate-update-period", 15*time.Second, "Period with which to update the per-user ingestion rates.")
	flag.DurationVar(&cfg.ingesterConfig.MaxChunkAge, "ingester.max-chunk-age", 10*time.Minute, "Maximum chunk age before flushing.")
	flag.IntVar(&cfg.numTokens, "ingester.num-tokens", 128, "Number of tokens for each ingester.")
	flag.IntVar(&cfg.distributorConfig.ReplicationFactor, "distributor.replication-factor", 3, "The number of ingesters to write to and read from.")
	flag.IntVar(&cfg.distributorConfig.MinReadSuccesses, "distributor.min-read-successes", 2, "The minimum number of ingesters from which a read must succeed.")
	flag.DurationVar(&cfg.distributorConfig.HeartbeatTimeout, "distributor.heartbeat-timeout", time.Minute, "The heartbeat timeout after which ingesters are skipped for reads/writes.")
	flag.BoolVar(&cfg.logSuccess, "log.success", false, "Log successful requests")
	flag.Parse()

	chunkStore, err := setupChunkStore(cfg)
	if err != nil {
		log.Fatalf("Error initializing chunk store: %v", err)
	}
	resourceWatcher, err := chunk.WatchDynamo(cfg.dynamodbURL, 2*time.Minute)
	if err != nil {
		log.Fatalf("Error initializing DynamoDB watcher: %v", err)
	}
	defer resourceWatcher.Stop()
	prometheus.MustRegister(resourceWatcher)

	consul, err := ring.NewConsulClient(cfg.consulHost)
	if err != nil {
		log.Fatalf("Error initializing Consul client: %v", err)
	}
	consul = ring.PrefixClient(consul, cfg.consulPrefix)
	r := ring.New(consul, cfg.distributorConfig.HeartbeatTimeout)
	defer r.Stop()

	switch cfg.mode {
	case modeDistributor:
		cfg.distributorConfig.Ring = r
		cfg.distributorConfig.ClientFactory = func(address string) (*distributor.IngesterClient, error) {
			return distributor.NewIngesterClient(address, cfg.remoteTimeout)
		}
		setupDistributor(cfg.distributorConfig, chunkStore, cfg.logSuccess)
	case modeIngester:
		cfg.ingesterConfig.Ring = r
		registration, err := ring.RegisterIngester(consul, cfg.listenPort, cfg.numTokens)
		if err != nil {
			// This only happens for errors in configuration & set-up, not for
			// network errors.
			log.Fatalf("Could not register ingester: %v", err)
		}
		ing := setupIngester(chunkStore, cfg.ingesterConfig, cfg.logSuccess)

		// Deferring a func to make ordering obvious
		defer func() {
			registration.ChangeState(ring.Leaving)
			ing.Stop()
			registration.Unregister()
		}()

		prometheus.MustRegister(registration)
	default:
		log.Fatalf("Mode %s not supported!", cfg.mode)
	}

	http.Handle("/metrics", prometheus.Handler())
	go http.ListenAndServe(fmt.Sprintf(":%d", cfg.listenPort), nil)

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	<-term
	log.Warn("Received SIGTERM, exiting gracefully...")
}

func setupChunkStore(cfg cfg) (chunk.Store, error) {
	var chunkCache *chunk.Cache
	if cfg.memcachedHostname != "" {
		chunkCache = &chunk.Cache{
			Memcache: chunk.NewMemcacheClient(chunk.MemcacheConfig{
				Host:           cfg.memcachedHostname,
				Service:        cfg.memcachedService,
				Timeout:        cfg.memcachedTimeout,
				UpdateInterval: 1 * time.Minute,
			}),
			Expiration: cfg.memcachedExpiration,
		}
	}
	chunkStore, err := chunk.NewAWSStore(chunk.StoreConfig{
		S3URL:       cfg.s3URL,
		DynamoDBURL: cfg.dynamodbURL,
		ChunkCache:  chunkCache,
	})
	if err != nil {
		return nil, err
	}
	if cfg.dynamodbCreateTables {
		if err = chunkStore.CreateTables(); err != nil {
			return nil, err
		}
	}
	return chunkStore, err
}

func setupDistributor(
	cfg distributor.Config,
	chunkStore chunk.Store,
	logSuccess bool,
) {
	dist, err := distributor.New(cfg)
	if err != nil {
		log.Fatal(err)
	}
	prometheus.MustRegister(dist)

	prefix := "/api/prom"
	http.Handle(prefix+"/push", instrument(logSuccess, cortex.AppenderHandler(dist, handleDistributorError)))

	// TODO: Move querier to separate binary.
	setupQuerier(dist, chunkStore, prefix, logSuccess)
}

func handleDistributorError(w http.ResponseWriter, err error) {
	switch e := err.(type) {
	case distributor.IngesterError:
		switch {
		case 400 <= e.StatusCode && e.StatusCode < 500:
			log.Warnf("append err: %v", err)
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
	}
	log.Errorf("append err: %v", err)
	http.Error(w, err.Error(), http.StatusInternalServerError)
}

// setupQuerier sets up a complete querying pipeline:
//
// PromQL -> MergeQuerier -> Distributor -> IngesterQuerier -> Ingester
//              |
//              `----------> ChunkQuerier -> DynamoDB/S3
func setupQuerier(
	distributor *distributor.Distributor,
	chunkStore chunk.Store,
	prefix string,
	logSuccess bool,
) {
	queryable := querier.Queryable{
		Q: querier.MergeQuerier{
			Queriers: []querier.Querier{
				distributor,
				&querier.ChunkQuerier{
					Store: chunkStore,
				},
			},
		},
	}

	engine := promql.NewEngine(queryable, nil)

	api := v1.NewAPI(engine, querier.DummyStorage{queryable})
	router := route.New(func(r *http.Request) (context.Context, error) {
		userID := r.Header.Get(userIDHeaderName)
		if r.Method != "OPTIONS" && userID == "" {
			// For now, getting the user ID from basic auth allows for easy testing
			// with Grafana.
			// TODO: Remove basic auth support.
			userID, _, _ = r.BasicAuth()
			if userID == "" {
				return nil, fmt.Errorf("missing user ID")
			}
		}
		return user.WithID(context.Background(), userID), nil
	})
	api.Register(router.WithPrefix(prefix + "/api/v1"))
	http.Handle("/", router)

	http.Handle(prefix+"/user_stats", instrument(logSuccess, cortex.DistributorUserStatsHandler(distributor.UserStats)))

	http.Handle(prefix+"/graph", instrument(logSuccess, ui.GraphHandler()))
	http.Handle(prefix+"/static/", instrument(logSuccess, ui.StaticAssetsHandler(prefix+"/static/")))
}

func setupIngester(
	chunkStore chunk.Store,
	cfg ingester.Config,
	logSuccess bool,
) *ingester.Ingester {
	ingester, err := ingester.New(cfg, chunkStore)
	if err != nil {
		log.Fatal(err)
	}
	prometheus.MustRegister(ingester)

	http.Handle("/push", instrument(logSuccess, cortex.AppenderHandler(ingester, handleIngesterError)))
	http.Handle("/query", instrument(logSuccess, cortex.QueryHandler(ingester)))
	http.Handle("/label_values", instrument(logSuccess, cortex.LabelValuesHandler(ingester)))
	http.Handle("/user_stats", instrument(logSuccess, cortex.IngesterUserStatsHandler(ingester.UserStats)))
	http.Handle("/ready", instrument(logSuccess, cortex.IngesterReadinessHandler(ingester)))
	return ingester
}

func handleIngesterError(w http.ResponseWriter, err error) {
	switch err {
	case ingester.ErrOutOfOrderSample, ingester.ErrDuplicateSampleForTimestamp:
		log.Warnf("append err: %v", err)
		http.Error(w, err.Error(), http.StatusBadRequest)
	default:
		log.Errorf("append err: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// instrument instruments a handler.
func instrument(logSuccess bool, handler http.Handler) http.Handler {
	return middleware.Merge(
		middleware.Log{
			LogSuccess: logSuccess,
		},
		middleware.Instrument{
			Duration: requestDuration,
		},
	).Wrap(handler)
}
<commit_msg>Make polling interval a parameter<commit_after>package main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"
	"github.com/prometheus/prometheus/promql"
	"github.com/prometheus/prometheus/web/api/v1"
	"github.com/weaveworks/scope/common/middleware"
	"golang.org/x/net/context"

	"github.com/weaveworks/cortex"
	"github.com/weaveworks/cortex/chunk"
	"github.com/weaveworks/cortex/distributor"
	"github.com/weaveworks/cortex/ingester"
	"github.com/weaveworks/cortex/querier"
	"github.com/weaveworks/cortex/ring"
	"github.com/weaveworks/cortex/ui"
	"github.com/weaveworks/cortex/user"
)

const (
	modeDistributor = "distributor"
	modeIngester    = "ingester"

	infName          = "eth0"
	userIDHeaderName = "X-Scope-OrgID"
)

var (
	requestDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{
		Namespace: "cortex",
		Name:      "request_duration_seconds",
		Help:      "Time (in seconds) spent serving HTTP requests.",

		// Cortex latency can be very low (ingesters typically take a few ms to
		// process a request), all the way to very high (queries can take tens of
		// second).  As its important, we use 10 buckets.  Smallest is 128us
		// biggest is 130s.
		Buckets: prometheus.ExponentialBuckets(0.000128, 4, 10),
	}, []string{"method", "route", "status_code", "ws"})
)

func init() {
	prometheus.MustRegister(requestDuration)
}

type cfg struct {
	mode                 string
	listenPort           int
	consulHost           string
	consulPrefix         string
	s3URL                string
	dynamodbURL          string
	dynamodbCreateTables bool
	dynamodbPollInterval time.Duration
	memcachedHostname    string
	memcachedTimeout     time.Duration
	memcachedExpiration  time.Duration
	memcachedService     string
	remoteTimeout        time.Duration
	numTokens            int
	logSuccess           bool

	ingesterConfig    ingester.Config
	distributorConfig distributor.Config
}

func main() {
	var cfg cfg
	flag.StringVar(&cfg.mode, "mode", modeDistributor, "Mode (distributor, ingester).")
	flag.IntVar(&cfg.listenPort, "web.listen-port", 9094, "HTTP server listen port.")
	flag.StringVar(&cfg.consulHost, "consul.hostname", "localhost:8500", "Hostname and port of Consul.")
	flag.StringVar(&cfg.consulPrefix, "consul.prefix", "collectors/", "Prefix for keys in Consul.")
	flag.StringVar(&cfg.s3URL, "s3.url", "localhost:4569", "S3 endpoint URL.")
	flag.StringVar(&cfg.dynamodbURL, "dynamodb.url", "localhost:8000", "DynamoDB endpoint URL.")
	flag.BoolVar(&cfg.dynamodbCreateTables, "dynamodb.create-tables", false, "Create required DynamoDB tables on startup.")
	flag.DurationVar(&cfg.dynamodbPollInterval, "dynamodb.poll-interval", 2*time.Minute, "How frequently to poll DynamoDB to learn our capacity.")
	flag.StringVar(&cfg.memcachedHostname, "memcached.hostname", "", "Hostname for memcached service to use when caching chunks. If empty, no memcached will be used.")
	flag.DurationVar(&cfg.memcachedTimeout, "memcached.timeout", 100*time.Millisecond, "Maximum time to wait before giving up on memcached requests.")
	flag.DurationVar(&cfg.memcachedExpiration, "memcached.expiration", 0, "How long chunks stay in the memcache.")
	flag.StringVar(&cfg.memcachedService, "memcached.service", "memcached", "SRV service used to discover memcache servers.")
	flag.DurationVar(&cfg.remoteTimeout, "remote.timeout", 5*time.Second, "Timeout for downstream ingesters.")
	flag.DurationVar(&cfg.ingesterConfig.FlushCheckPeriod, "ingester.flush-period", 1*time.Minute, "Period with which to attempt to flush chunks.")
	flag.DurationVar(&cfg.ingesterConfig.RateUpdatePeriod, "ingester.rate-update-period", 15*time.Second, "Period with which to update the per-user ingestion rates.")
	flag.DurationVar(&cfg.ingesterConfig.MaxChunkAge, "ingester.max-chunk-age", 10*time.Minute, "Maximum chunk age before flushing.")
	flag.IntVar(&cfg.numTokens, "ingester.num-tokens", 128, "Number of tokens for each ingester.")
	flag.IntVar(&cfg.distributorConfig.ReplicationFactor, "distributor.replication-factor", 3, "The number of ingesters to write to and read from.")
	flag.IntVar(&cfg.distributorConfig.MinReadSuccesses, "distributor.min-read-successes", 2, "The minimum number of ingesters from which a read must succeed.")
	flag.DurationVar(&cfg.distributorConfig.HeartbeatTimeout, "distributor.heartbeat-timeout", time.Minute, "The heartbeat timeout after which ingesters are skipped for reads/writes.")
	flag.BoolVar(&cfg.logSuccess, "log.success", false, "Log successful requests")
	flag.Parse()

	chunkStore, err := setupChunkStore(cfg)
	if err != nil {
		log.Fatalf("Error initializing chunk store: %v", err)
	}
	if cfg.dynamodbPollInterval < 1*time.Minute {
		log.Warnf("Polling DynamoDB more than once a minute. Likely to get throttled: %v", cfg.dynamodbPollInterval)
	}
	resourceWatcher, err := chunk.WatchDynamo(cfg.dynamodbURL, cfg.dynamodbPollInterval)
	if err != nil {
		log.Fatalf("Error initializing DynamoDB watcher: %v", err)
	}
	defer resourceWatcher.Stop()
	prometheus.MustRegister(resourceWatcher)

	consul, err := ring.NewConsulClient(cfg.consulHost)
	if err != nil {
		log.Fatalf("Error initializing Consul client: %v", err)
	}
	consul = ring.PrefixClient(consul, cfg.consulPrefix)
	r := ring.New(consul, cfg.distributorConfig.HeartbeatTimeout)
	defer r.Stop()

	switch cfg.mode {
	case modeDistributor:
		cfg.distributorConfig.Ring = r
		cfg.distributorConfig.ClientFactory = func(address string) (*distributor.IngesterClient, error) {
			return distributor.NewIngesterClient(address, cfg.remoteTimeout)
		}
		setupDistributor(cfg.distributorConfig, chunkStore, cfg.logSuccess)
	case modeIngester:
		cfg.ingesterConfig.Ring = r
		registration, err := ring.RegisterIngester(consul, cfg.listenPort, cfg.numTokens)
		if err != nil {
			// This only happens for errors in configuration & set-up, not for
			// network errors.
			log.Fatalf("Could not register ingester: %v", err)
		}
		ing := setupIngester(chunkStore, cfg.ingesterConfig, cfg.logSuccess)

		// Deferring a func to make ordering obvious
		defer func() {
			registration.ChangeState(ring.Leaving)
			ing.Stop()
			registration.Unregister()
		}()

		prometheus.MustRegister(registration)
	default:
		log.Fatalf("Mode %s not supported!", cfg.mode)
	}

	http.Handle("/metrics", prometheus.Handler())
	go http.ListenAndServe(fmt.Sprintf(":%d", cfg.listenPort), nil)

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
	<-term
	log.Warn("Received SIGTERM, exiting gracefully...")
}

func setupChunkStore(cfg cfg) (chunk.Store, error) {
	var chunkCache *chunk.Cache
	if cfg.memcachedHostname != "" {
		chunkCache = &chunk.Cache{
			Memcache: chunk.NewMemcacheClient(chunk.MemcacheConfig{
				Host:           cfg.memcachedHostname,
				Service:        cfg.memcachedService,
				Timeout:        cfg.memcachedTimeout,
				UpdateInterval: 1 * time.Minute,
			}),
			Expiration: cfg.memcachedExpiration,
		}
	}
	chunkStore, err := chunk.NewAWSStore(chunk.StoreConfig{
		S3URL:       cfg.s3URL,
		DynamoDBURL: cfg.dynamodbURL,
		ChunkCache:  chunkCache,
	})
	if err != nil {
		return nil, err
	}
	if cfg.dynamodbCreateTables {
		if err = chunkStore.CreateTables(); err != nil {
			return nil, err
		}
	}
	return chunkStore, err
}

func setupDistributor(
	cfg distributor.Config,
	chunkStore chunk.Store,
	logSuccess bool,
) {
	dist, err := distributor.New(cfg)
	if err != nil {
		log.Fatal(err)
	}
	prometheus.MustRegister(dist)

	prefix := "/api/prom"
	http.Handle(prefix+"/push", instrument(logSuccess, cortex.AppenderHandler(dist, handleDistributorError)))

	// TODO: Move querier to separate binary.
	setupQuerier(dist, chunkStore, prefix, logSuccess)
}

func handleDistributorError(w http.ResponseWriter, err error) {
	switch e := err.(type) {
	case distributor.IngesterError:
		switch {
		case 400 <= e.StatusCode && e.StatusCode < 500:
			log.Warnf("append err: %v", err)
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
	}
	log.Errorf("append err: %v", err)
	http.Error(w, err.Error(), http.StatusInternalServerError)
}

// setupQuerier sets up a complete querying pipeline:
//
// PromQL -> MergeQuerier -> Distributor -> IngesterQuerier -> Ingester
//              |
//              `----------> ChunkQuerier -> DynamoDB/S3
func setupQuerier(
	distributor *distributor.Distributor,
	chunkStore chunk.Store,
	prefix string,
	logSuccess bool,
) {
	queryable := querier.Queryable{
		Q: querier.MergeQuerier{
			Queriers: []querier.Querier{
				distributor,
				&querier.ChunkQuerier{
					Store: chunkStore,
				},
			},
		},
	}

	engine := promql.NewEngine(queryable, nil)

	api := v1.NewAPI(engine, querier.DummyStorage{Queryable: queryable})
	router := route.New(func(r *http.Request) (context.Context, error) {
		userID := r.Header.Get(userIDHeaderName)
		if r.Method != "OPTIONS" && userID == "" {
			// For now, getting the user ID from basic auth allows for easy testing
			// with Grafana.
			// TODO: Remove basic auth support.
			userID, _, _ = r.BasicAuth()
			if userID == "" {
				return nil, fmt.Errorf("missing user ID")
			}
		}
		return user.WithID(context.Background(), userID), nil
	})
	api.Register(router.WithPrefix(prefix + "/api/v1"))
	http.Handle("/", router)

	http.Handle(prefix+"/user_stats", instrument(logSuccess, cortex.DistributorUserStatsHandler(distributor.UserStats)))

	http.Handle(prefix+"/graph", instrument(logSuccess, ui.GraphHandler()))
	http.Handle(prefix+"/static/", instrument(logSuccess, ui.StaticAssetsHandler(prefix+"/static/")))
}

func setupIngester(
	chunkStore chunk.Store,
	cfg ingester.Config,
	logSuccess bool,
) *ingester.Ingester {
	ingester, err := ingester.New(cfg, chunkStore)
	if err != nil {
		log.Fatal(err)
	}
	prometheus.MustRegister(ingester)

	http.Handle("/push", instrument(logSuccess, cortex.AppenderHandler(ingester, handleIngesterError)))
	http.Handle("/query", instrument(logSuccess, cortex.QueryHandler(ingester)))
	http.Handle("/label_values", instrument(logSuccess, cortex.LabelValuesHandler(ingester)))
	http.Handle("/user_stats", instrument(logSuccess, cortex.IngesterUserStatsHandler(ingester.UserStats)))
	http.Handle("/ready", instrument(logSuccess, cortex.IngesterReadinessHandler(ingester)))
	return ingester
}

func handleIngesterError(w http.ResponseWriter, err error) {
	switch err {
	case ingester.ErrOutOfOrderSample, ingester.ErrDuplicateSampleForTimestamp:
		log.Warnf("append err: %v", err)
		http.Error(w, err.Error(), http.StatusBadRequest)
	default:
		log.Errorf("append err: %v", err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// instrument instruments a handler.
func instrument(logSuccess bool, handler http.Handler) http.Handler {
	return middleware.Merge(
		middleware.Log{
			LogSuccess: logSuccess,
		},
		middleware.Instrument{
			Duration: requestDuration,
		},
	).Wrap(handler)
}
<|endoftext|>


<commit_before>package module

import (
	"io"
	"log"
	"net/http"
	"os"
	ospath "path"
	"strconv"
	"strings"

	"../utility/array"
	"../utility/json"
)

type RoutingModule interface {
	ReturnFiles()
	ReturnBinary()
	Close()
}

type moduleConfig struct {
	name     string
	exts     []string
	types    moduleType
	routing  func(r io.ReadCloser, vpath string, w http.ResponseWriter, size int64) RoutingModule
	routing2 func(path, vpath string, w http.ResponseWriter) RoutingModule
}

type moduleType int

const (
	typeModuleImage moduleType = iota
	typeModuleArch
)

var (
	confs []*moduleConfig
)

func ReturnSupportType(w http.ResponseWriter, r *http.Request) {
	data := supportType()
	json.WriteResponse(w, data)
}

func Routing(path string, w http.ResponseWriter) RoutingModule {
	path = strings.Replace(path, "\\", "/", -1)

	// return directory routing module
	if stat, err := os.Stat(path); err == nil && stat.IsDir() {
		return newDirectoryRouting(path, w)
	}

	// split real path :: virtual path
	paths := strings.Split(path, "/")
	path, paths = paths[0], paths[1:]
	path += "/"

	for len(paths) > 0 {
		newPath := path + "/" + paths[0]
		if _, err := os.Stat(newPath); err != nil {
			break
		}

		path = newPath
		paths = paths[1:]
	}
	vpath := strings.Join(paths, "/")

	// invalid path
	if _, err := os.Stat(path); err != nil {
		http.Error(w, "Not Found File", http.StatusInternalServerError)
		return nil
	}

	// dispatch
	conf := dispatch(path, w)
	if conf == nil {
		return nil
	}
	return conf.routing2(path, vpath, w)
}

func install(h *moduleConfig) *moduleConfig {
	confs = append(confs, h)
	return h
}

func dispatch(path string, w http.ResponseWriter) *moduleConfig {
	// not found extension
	ext := ospath.Ext(path)
	if len(ext) == 0 {
		http.Error(w, "Not Found File Type", http.StatusInternalServerError)
		return nil
	}

	// search extension
	for _, conf := range confs {
		if array.IsInclude(ext, conf.exts) {
			return conf
		}
	}

	// not found file type
	http.Error(w, "No Support Type", http.StatusUnsupportedMediaType)
	return nil
}

func returnBinary(w http.ResponseWriter, r io.Reader, mime string, size int64) {
	w.Header().Set("Content-Type", mime)
	w.Header().Set("Content-Length", strconv.FormatInt(size, 10))
	if _, err := io.Copy(w, r); err != nil {
		log.Printf("ERR: %s: %v\n", mime, err)
	}
}

func supportType() map[string][]string {
	files := []string{}
	archs := []string{}

	for _, conf := range confs {
		switch conf.types {
		case typeModuleImage:
			files = append(files, conf.exts...)
		case typeModuleArch:
			archs = append(archs, conf.exts...)
		default:
			// nothing to do
		}
	}

	return map[string][]string{
		"image":   files,
		"archive": archs,
	}
}
<commit_msg>update module interface<commit_after>package module

import (
	"fmt"
	"io"
	"strings"
)

type ModuleType int

const (
	MODULE_IMAGE ModuleType = iota
	MODULE_ARCH
)

const errNotSupport = "%s is not support this module type"

type File struct {
	data []byte
	mime string
	size int64
}

type Module struct {
	exts []string
	Type ModuleType
	// archive type
	funcArchFiles func(r io.Reader) []string
	funcArchRead  func(r io.Reader, vpath string) *File
	// image type
	funcImageRead func(r io.Reader) *File
}

var modules []*Module

func GetSupportModule(path string) *Module {
	ext := getFileExt(path)
	mod, ok := modules[ext]
	if !ok {
		return nil
	}
	return mod
}

func RegisterArchModule(exts []string, funcArchFiles func(r io.Reader) []string, funcArchRead func(r io.Reader, vpath string) *File) {
	h := &Module{
		exts:          exts,
		Type:          MODULE_ARCH,
		funcArchFiles: funcArchFiles,
		funcArchRead:  funcArchRead,
		funcImageRead: dummyFuncImageRead,
	}
	modules = append(modules, h)
}

func RegisterImageModule(exts []string, funcImageRead func(r io.Reader) *File) {
	h := &Module{
		exts:          exts,
		Type:          MODULE_IMAGE,
		funcArchFiles: dummyFuncArchFiles,
		funcArchRead:  dummyFuncArchRead,
		funcImageRead: funcImageRead,
	}
	modules = append(modules, h)
}

func SupportType() map[ModuleType][]string {
	types := map[ModuleType][]string{
		MODULE_IMAGE: make([]string, 0),
		MODULE_ARCH:  make([]string, 0),
	}

	for _, module := range modules {
		v, _ := types[module.Type]
		v = append(v, module.exts...)
		types[module.types] = v
	}
	return types
}

func dummyFuncArchFiles(r io.Reader) []string {
	fmt.Printf(errNotSupport, "Files")
	return nil
}

func dummyFuncArchRead(r io.Reader, vpath string) *File {
	fmt.Printf(errNotSupport, "Read")
	return nil
}

func dummyFuncImageRead(r io.Reader) *File {
	fmt.Printf(errNotSupport, "Read")
	return nil
}

func getFileExt(path string) (ext string) {
	index := strings.LastIndex(path, ".")
	if index < 0 || index >= len(path)-1 {
		return
	}

	ext = path[index+1:]
	return
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"fmt"
	"net"
	"os"
	"time"

	"github.com/golang/glog"
	"github.com/urfave/cli"
	"google.golang.org/grpc"
)

const (
	defaultTimeout = 10 * time.Second
)

func getRuntimeClientConnection(context *cli.Context) (*grpc.ClientConn, error) {
	runtimeEndpoint := context.GlobalString("runtime-endpoint")
	if runtimeEndpoint == "" {
		return nil, fmt.Errorf("--runtime-endpoint is not set")
	}
	conn, err := grpc.Dial(context.GlobalString("runtime-endpoint"), grpc.WithInsecure(), grpc.WithTimeout(context.GlobalDuration("timeout")),
		grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
			return net.DialTimeout("unix", addr, timeout)
		}))
	if err != nil {
		return nil, fmt.Errorf("failed to connect: %v", err)
	}
	return conn, nil
}

func main() {
	app := cli.NewApp()
	app.Name = "crictl"
	app.Usage = "client for CRI"
	app.Version = "0.0.1"

	app.Commands = []cli.Command{
		runtimeVersionCommand,
		runtimePodSandboxCommand,
		runtimeContainerCommand,
	}

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "runtime-endpoint",
			Usage: "CRI server runtime endpoint",
		},
		cli.StringFlag{
			Name:  "image-endpoint",
			Usage: "CRI server image endpoint",
		},
		cli.DurationFlag{
			Name:  "timeout",
			Value: defaultTimeout,
			Usage: "Timeout of connecting to server",
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "Enable debug output",
		},
	}

	if err := app.Run(os.Args); err != nil {
		glog.Fatal(err)
	}
}
<commit_msg>set default runtime endpoint<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"fmt"
	"net"
	"os"
	"time"

	"github.com/golang/glog"
	"github.com/urfave/cli"
	"google.golang.org/grpc"
)

const (
	defaultTimeout = 10 * time.Second
)

func getRuntimeClientConnection(context *cli.Context) (*grpc.ClientConn, error) {
	runtimeEndpoint := context.GlobalString("runtime-endpoint")
	if runtimeEndpoint == "" {
		return nil, fmt.Errorf("--runtime-endpoint is not set")
	}
	conn, err := grpc.Dial(context.GlobalString("runtime-endpoint"), grpc.WithInsecure(), grpc.WithTimeout(context.GlobalDuration("timeout")),
		grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
			return net.DialTimeout("unix", addr, timeout)
		}))
	if err != nil {
		return nil, fmt.Errorf("failed to connect: %v", err)
	}
	return conn, nil
}

func main() {
	app := cli.NewApp()
	app.Name = "crictl"
	app.Usage = "client for CRI"
	app.Version = "0.0.1"

	app.Commands = []cli.Command{
		runtimeVersionCommand,
		runtimePodSandboxCommand,
		runtimeContainerCommand,
	}

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "runtime-endpoint",
			Value: "/var/run/dockershim.sock",
			Usage: "CRI server runtime endpoint",
		},
		cli.StringFlag{
			Name:  "image-endpoint",
			Usage: "CRI server image endpoint",
		},
		cli.DurationFlag{
			Name:  "timeout",
			Value: defaultTimeout,
			Usage: "Timeout of connecting to server",
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "Enable debug output",
		},
	}

	if err := app.Run(os.Args); err != nil {
		glog.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package module

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/dnaeon/gru/resource"
	"github.com/hashicorp/hcl"
	"github.com/hashicorp/hcl/hcl/ast"
)

// ErrMultipleImport error is returned if there are multiple import
// declarations in the same module
var ErrMultipleImport = errors.New("Multiple import declarations found")

// Import type represents a single import declaration from HCL/JSON
type ImportType struct {
	// Module names being imported
	Module []string
}

// Module type represents a collection of resources and module imports
type Module struct {
	// Name of the module
	Name string

	// Resources loaded from the module
	Resources []resource.Resource

	// Module imports
	ModuleImport ImportType
}

// New creates a new empty module
func New(name string) *Module {
	m := &Module{
		Name:      name,
		Resources: make([]resource.Resource, 0),
		ModuleImport: ImportType{
			Module: make([]string, 0),
		},
	}

	return m
}

// Load loads a module from the given HCL or JSON input
func Load(name string, r io.Reader) (*Module, error) {
	m := New(name)

	input, err := ioutil.ReadAll(r)
	if err != nil {
		return m, err
	}

	// Parse configuration
	obj, err := hcl.Parse(string(input))
	if err != nil {
		return m, err
	}

	// Top-level node should be an object list
	root, ok := obj.Node.(*ast.ObjectList)
	if !ok {
		return m, fmt.Errorf("Missing root node in %s", name)
	}

	err = m.hclLoadImport(root)
	if err != nil {
		return m, err
	}

	err = m.hclLoadResources(root)
	if err != nil {
		return m, err
	}

	return m, nil
}

// hclLoadResources loads all resource declarations from the given HCL input
func (m *Module) hclLoadResources(root *ast.ObjectList) error {
	hclResources := root.Filter("resource")
	for _, item := range hclResources.Items {
		position := item.Val.Pos().String()

		// The item is expected to have exactly one key which
		// represents the resource type.
		if len(item.Keys) != 1 {
			e := fmt.Errorf("Invalid resource declaration found in %s:%s", m.Name, position)
			return e
		}

		// Get the resource from registry and create the actual resource
		resourceType := item.Keys[0].Token.Value().(string)
		registryItem, ok := resource.Registry[resourceType]
		if !ok {
			e := fmt.Errorf("Unknown resource type '%s' found in %s:%s", resourceType, m.Name, position)
			return e
		}

		// Create the actual resource by calling it's provider
		r, err := registryItem.Provider(item)
		if err != nil {
			return err
		}

		m.Resources = append(m.Resources, r)
	}

	return nil
}

// hclLoadImport loads all import declarations from the given HCL input
func (m *Module) hclLoadImport(root *ast.ObjectList) error {
	hclImport := root.Filter("import")

	// We expect to have exactly one import declaration per module file
	if len(hclImport.Items) > 1 {
		return fmt.Errorf("Multiple import declarations found in %s", m.Name)
	}

	if len(hclImport.Items) == 0 {
		return nil
	}

	err := hcl.DecodeObject(&m.ModuleImport, hclImport.Items[0])
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>module.Load() loads all known resource types from the resource registry<commit_after>package module

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/dnaeon/gru/resource"
	"github.com/hashicorp/hcl"
	"github.com/hashicorp/hcl/hcl/ast"
)

// ErrMultipleImport error is returned if there are multiple import
// declarations in the same module
var ErrMultipleImport = errors.New("Multiple import declarations found")

// Import type represents a single import declaration from HCL/JSON
type ImportType struct {
	// Module names being imported
	Module []string
}

// Module type represents a collection of resources and module imports
type Module struct {
	// Name of the module
	Name string

	// Resources loaded from the module
	Resources []resource.Resource

	// Module imports
	ModuleImport ImportType
}

// New creates a new empty module
func New(name string) *Module {
	m := &Module{
		Name:      name,
		Resources: make([]resource.Resource, 0),
		ModuleImport: ImportType{
			Module: make([]string, 0),
		},
	}

	return m
}

// Load loads a module from the given HCL or JSON input
func Load(name string, r io.Reader) (*Module, error) {
	m := New(name)

	input, err := ioutil.ReadAll(r)
	if err != nil {
		return m, err
	}

	// Parse configuration
	obj, err := hcl.Parse(string(input))
	if err != nil {
		return m, err
	}

	// Top-level node should be an object list
	root, ok := obj.Node.(*ast.ObjectList)
	if !ok {
		return m, fmt.Errorf("Missing root node in %s", name)
	}

	err = m.hclLoadImport(root)
	if err != nil {
		return m, err
	}

	// Load all known resource types from the given input
	for name := range resource.Registry {
		err = m.hclLoadResources(name, root)
		if err != nil {
			return m, err
		}
	}

	return m, nil
}

// hclLoadResources loads all declarations with the
// given resource type from the provided HCL input
func (m *Module) hclLoadResources(resourceType string, root *ast.ObjectList) error {
	hclResources := root.Filter(resourceType)
	for _, item := range hclResources.Items {
		position := item.Val.Pos().String()

		// The item is expected to exactly one key which
		// represent the resource name
		if len(item.Keys) != 1 {
			e := fmt.Errorf("Invalid resource declaration found in %s:%s", m.Name, position)
			return e
		}

		// Get the resource from registry and create the actual resource
		resourceName := item.Keys[0].Token.Value().(string)
		registryItem, ok := resource.Registry[resourceType]
		if !ok {
			e := fmt.Errorf("Unknown resource type '%s' found in %s:%s", resourceType, m.Name, position)
			return e
		}

		// Create the actual resource by calling it's provider
		r, err := registryItem.Provider(resourceName, item)
		if err != nil {
			return err
		}

		m.Resources = append(m.Resources, r)
	}

	return nil
}

// hclLoadImport loads all import declarations from the given HCL input
func (m *Module) hclLoadImport(root *ast.ObjectList) error {
	hclImport := root.Filter("import")

	// We expect to have exactly one import declaration per module file
	if len(hclImport.Items) > 1 {
		return fmt.Errorf("Multiple import declarations found in %s", m.Name)
	}

	if len(hclImport.Items) == 0 {
		return nil
	}

	err := hcl.DecodeObject(&m.ModuleImport, hclImport.Items[0])
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/goplus/gop/x/gopmod"
)

type goFile struct {
	file string
	code []byte
}

func newGoProj(file string) *gopmod.Project {
	code, err := os.ReadFile(file)
	if err != nil {
		log.Fatalln(err)
	}
	src := &goFile{file: file, code: code}
	return &gopmod.Project{
		Source:        src,
		AutoGenFile:   file,
		FriendlyFname: filepath.Base(file),
	}
}

func (p *goFile) Fingerp() [20]byte { // source code fingerprint
	return sha1.Sum(p.code)
}

func (p *goFile) IsDirty(outFile string, temp bool) bool {
	return true
}

func (p *goFile) GenGo(outFile string) error {
	if p.file != outFile {
		return os.WriteFile(outFile, p.code, 0666)
	}
	return nil
}

func main() {
	if len(os.Args) < 2 {
		fmt.Fprint(os.Stderr, "Usage: gorun <file.go>\n\n")
		return
	}
	goProj := newGoProj(os.Args[1])
	ctx := gopmod.New("")
	cmd := ctx.GoCommand("run", goProj, os.Args[2:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = os.Environ()
	err := cmd.Run()
	if err != nil {
		switch e := err.(type) {
		case *exec.ExitError:
			os.Exit(e.ExitCode())
		default:
			log.Fatalln(err)
		}
	}
}
<commit_msg>gorun<commit_after>package main

import (
	"crypto/sha1"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"

	"github.com/goplus/gop/x/gopmod"
)

type goFile struct {
	file string
	code []byte
}

func newGoProj(file string) *gopmod.Project {
	code, err := os.ReadFile(file)
	if err != nil {
		log.Fatalln(err)
	}
	src := &goFile{file: file, code: code}
	return &gopmod.Project{
		Source:        src,
		AutoGenFile:   file,
		FriendlyFname: filepath.Base(file),
	}
}

func (p *goFile) Fingerp() [20]byte { // source code fingerprint
	return sha1.Sum(p.code)
}

func (p *goFile) IsDirty(outFile string, temp bool) bool {
	return true
}

func (p *goFile) GenGo(outFile string) error {
	if p.file != outFile {
		return os.WriteFile(outFile, p.code, 0666)
	}
	return nil
}

func main() {
	if len(os.Args) < 2 {
		fmt.Fprint(os.Stderr, "Usage: gorun <file.go> [switch ...]\n\n")
		return
	}
	goProj := newGoProj(os.Args[1])
	ctx := gopmod.New("")
	cmd := ctx.GoCommand("run", goProj, os.Args[2:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = os.Environ()
	err := cmd.Run()
	if err != nil {
		switch e := err.(type) {
		case *exec.ExitError:
			os.Exit(e.ExitCode())
		default:
			log.Fatalln(err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"time"

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

var (
	version = "master"
)

const (
	actionVersion = "version"
	actionInit    = "init"
	actionPush    = "push"
	actionReindex = "reindex"
	actionDelete  = "delete"

	defaultTimeout = time.Minute * 5
)

// Action describes plugin action that can be run.
type Action interface {
	Run(context.Context) error
}

func main() {
	log.SetFlags(0)

	if len(os.Args) == 5 && !isAction(os.Args[1]) {
		cmd := proxyCmd{uri: os.Args[4]}
		ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
		defer cancel()
		if err := cmd.Run(ctx); err != nil {
			log.Fatal(err)
		}
		return
	}

	cli := kingpin.New("helm s3", "")
	cli.Command(actionVersion, "Show plugin version.")

	initCmd := cli.Command(actionInit, "Initialize empty repository on AWS S3.")
	initURI := initCmd.Arg("uri", "URI of repository, e.g. s3://awesome-bucket/charts").
		Required().
		String()

	pushCmd := cli.Command(actionPush, "Push chart to the repository.")
	pushChartPath := pushCmd.Arg("chartPath", "Path to a chart, e.g. ./epicservice-0.5.1.tgz").
		Required().
		String()
	pushTargetRepository := pushCmd.Arg("repo", "Target repository to push to").
		Required().
		String()
	pushForce := pushCmd.Flag("force", "Replace the chart if it already exists. This can cause the repository to lose existing chart; use it with care").
		Bool()

	reindexCmd := cli.Command(actionReindex, "Reindex the repository.")
	reindexTargetRepository := reindexCmd.Arg("repo", "Target repository to reindex").
		Required().
		String()

	deleteCmd := cli.Command(actionDelete, "Delete chart from the repository.").Alias("del")
	deleteChartName := deleteCmd.Arg("chartName", "Name of chart to delete").
		Required().
		String()
	deleteChartVersion := deleteCmd.Flag("version", "Version of chart to delete").
		Required().
		String()
	deleteTargetRepository := deleteCmd.Arg("repo", "Target repository to delete from").
		Required().
		String()

	action := kingpin.MustParse(cli.Parse(os.Args[1:]))
	if action == "" {
		cli.Usage(os.Args[1:])
		os.Exit(0)
	}

	var act Action
	switch action {
	case actionVersion:
		fmt.Print(version)
		return

	case actionInit:
		act = initAction{
			uri: *initURI,
		}
		defer fmt.Printf("Initialized empty repository at %s\n", *initURI)

	case actionPush:
		act = pushAction{
			chartPath: *pushChartPath,
			repoName:  *pushTargetRepository,
			force:     *pushForce,
		}

	case actionReindex:
		fmt.Fprint(os.Stderr, "Warning: reindex feature is in beta. If you experience any issues,\nplease provide your feedback here: https://github.com/hypnoglow/helm-s3/issues/22\n\n")
		act = reindexAction{
			repoName: *reindexTargetRepository,
		}
		defer fmt.Printf("Repository %s was successfully reindexed.\n", *reindexTargetRepository)

	case actionDelete:
		act = deleteAction{
			name:     *deleteChartName,
			version:  *deleteChartVersion,
			repoName: *deleteTargetRepository,
		}
	default:
		return
	}

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

	err := act.Run(ctx)
	switch err {
	case nil:
	case ErrChartExists:
		log.Fatalf("The chart already exists in the repository and cannot be overwritten without an explicit intent. If you want to replace existing chart, use --force flag:\n\n\thelm s3 push --force %s %s\n\n", *pushChartPath, *pushTargetRepository)
	default:
		log.Fatal(err)
	}
}

func isAction(name string) bool {
	return name == actionDelete ||
		name == actionInit ||
		name == actionPush ||
		name == actionReindex ||
		name == actionVersion
}
<commit_msg>Add timeout flag<commit_after>package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"time"

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

var (
	version = "master"
)

const (
	actionVersion = "version"
	actionInit    = "init"
	actionPush    = "push"
	actionReindex = "reindex"
	actionDelete  = "delete"

	defaultTimeout       = time.Minute * 5
	defaultTimeoutString = "5m"

	helpFlagTimeout = `Timeout for the whole operation to complete. Defaults to 5 minutes.

If you don't use MFA, it may be reasonable to lower the timeout
for the most commands, for example to 10 seconds.

In opposite, in cases where you want to reindex big repository
(e.g. 10 000 charts), you definitely want to increase the timeout.
`
)

// Action describes plugin action that can be run.
type Action interface {
	Run(context.Context) error
}

func main() {
	log.SetFlags(0)

	if len(os.Args) == 5 && !isAction(os.Args[1]) {
		cmd := proxyCmd{uri: os.Args[4]}
		ctx, cancel := context.WithTimeout(context.Background(), defaultTimeout)
		defer cancel()
		if err := cmd.Run(ctx); err != nil {
			log.Fatal(err)
		}
		return
	}

	cli := kingpin.New("helm s3", "")
	cli.Command(actionVersion, "Show plugin version.")

	timeout := cli.Flag("timeout", helpFlagTimeout).
		Default(defaultTimeoutString).
		Duration()

	initCmd := cli.Command(actionInit, "Initialize empty repository on AWS S3.")
	initURI := initCmd.Arg("uri", "URI of repository, e.g. s3://awesome-bucket/charts").
		Required().
		String()

	pushCmd := cli.Command(actionPush, "Push chart to the repository.")
	pushChartPath := pushCmd.Arg("chartPath", "Path to a chart, e.g. ./epicservice-0.5.1.tgz").
		Required().
		String()
	pushTargetRepository := pushCmd.Arg("repo", "Target repository to push to").
		Required().
		String()
	pushForce := pushCmd.Flag("force", "Replace the chart if it already exists. This can cause the repository to lose existing chart; use it with care").
		Bool()

	reindexCmd := cli.Command(actionReindex, "Reindex the repository.")
	reindexTargetRepository := reindexCmd.Arg("repo", "Target repository to reindex").
		Required().
		String()

	deleteCmd := cli.Command(actionDelete, "Delete chart from the repository.").Alias("del")
	deleteChartName := deleteCmd.Arg("chartName", "Name of chart to delete").
		Required().
		String()
	deleteChartVersion := deleteCmd.Flag("version", "Version of chart to delete").
		Required().
		String()
	deleteTargetRepository := deleteCmd.Arg("repo", "Target repository to delete from").
		Required().
		String()

	action := kingpin.MustParse(cli.Parse(os.Args[1:]))
	if action == "" {
		cli.Usage(os.Args[1:])
		os.Exit(0)
	}

	var act Action
	switch action {
	case actionVersion:
		fmt.Print(version)
		return

	case actionInit:
		act = initAction{
			uri: *initURI,
		}
		defer fmt.Printf("Initialized empty repository at %s\n", *initURI)

	case actionPush:
		act = pushAction{
			chartPath: *pushChartPath,
			repoName:  *pushTargetRepository,
			force:     *pushForce,
		}

	case actionReindex:
		fmt.Fprint(os.Stderr, "Warning: reindex feature is in beta. If you experience any issues,\nplease provide your feedback here: https://github.com/hypnoglow/helm-s3/issues/22\n\n")
		act = reindexAction{
			repoName: *reindexTargetRepository,
		}
		defer fmt.Printf("Repository %s was successfully reindexed.\n", *reindexTargetRepository)

	case actionDelete:
		act = deleteAction{
			name:     *deleteChartName,
			version:  *deleteChartVersion,
			repoName: *deleteTargetRepository,
		}
	default:
		return
	}

	ctx, cancel := context.WithTimeout(context.Background(), *timeout)
	defer cancel()

	err := act.Run(ctx)
	switch err {
	case nil:
	case ErrChartExists:
		log.Fatalf("The chart already exists in the repository and cannot be overwritten without an explicit intent. If you want to replace existing chart, use --force flag:\n\n\thelm s3 push --force %s %s\n\n", *pushChartPath, *pushTargetRepository)
	default:
		log.Fatal(err)
	}
}

func isAction(name string) bool {
	return name == actionDelete ||
		name == actionInit ||
		name == actionPush ||
		name == actionReindex ||
		name == actionVersion
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/frankbraun/kitchensink/textbuffer"
	"github.com/gdamore/tcell"
)

func draw(s tcell.Screen, tb *textbuffer.TextBuffer, w, h int) {
	for y := 0; y < h && y < tb.Lines(); y++ {
		x := 0
		lineLen := tb.LineLenCell(y)
		for x < w && x < lineLen {
			c, cw := tb.GetCell(x, y)
			if x+cw <= lineLen {
				s.SetContent(x, y, c[0], c[1:], tcell.StyleDefault)
			}
			x += cw
		}
	}
}

func least(filename string) error {
	tb, err := textbuffer.NewFile(filename)
	if err != nil {
		return err
	}
	s, err := tcell.NewScreen()
	if err != nil {
		return err
	}
	defer s.Fini()
	if err = s.Init(); err != nil {
		return err
	}
	s.SetStyle(tcell.StyleDefault.
		Foreground(tcell.ColorBlack).
		Background(tcell.ColorWhite))
	s.Clear()
	w, h := s.Size()
	s.Clear()
	draw(s, tb, w, h)
	s.Show()
	for {
		ev := s.PollEvent()
		switch ev.(type) {
		case *tcell.EventKey:
			return nil
		case *tcell.EventResize:
			w, h = s.Size()
			s.Clear()
			draw(s, tb, w, h)
			s.Sync()
		}
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s filename\n", os.Args[0])
	os.Exit(1)
}

func fatal(err error) {
	fmt.Fprintf(os.Stderr, "%s: error: %s\n", os.Args[0], err)
	os.Exit(1)
}

func main() {
	if len(os.Args) != 2 {
		usage()
	}
	if err := least(os.Args[1]); err != nil {
		fatal(err)
	}
}
<commit_msg>more stuff for least<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/frankbraun/kitchensink/textbuffer"
	"github.com/gdamore/tcell"
)

var (
	curX int
	curY int
)

func draw(s tcell.Screen, tb *textbuffer.TextBuffer, w, h int) {
	for y := 0; y < h && y+curY < tb.Lines(); y++ {
		x := 0
		lineLen := tb.LineLenCell(y + curY)
		for x < w && x+curX < lineLen {
			c, cw := tb.GetCell(x+curX, y+curY)
			if x+cw <= lineLen {
				s.SetContent(x, y, c[0], c[1:], tcell.StyleDefault)
			}
			x += cw
		}
	}
}

func redraw(s tcell.Screen, tb *textbuffer.TextBuffer, w, h int) {
	s.Clear()
	draw(s, tb, w, h)
	s.Show()
}

func least(filename string) error {
	tb, err := textbuffer.NewFile(filename)
	if err != nil {
		return err
	}
	maxX := 0
	maxY := tb.Lines()
	for y := 0; y < maxY; y++ {
		if tb.LineLenCell(y) > maxX {
			maxX = tb.LineLenCell(y)
		}
	}
	s, err := tcell.NewScreen()
	if err != nil {
		return err
	}
	defer s.Fini()
	if err = s.Init(); err != nil {
		return err
	}
	s.SetStyle(tcell.StyleDefault.
		Foreground(tcell.ColorBlack).
		Background(tcell.ColorWhite))
	s.Clear()
	w, h := s.Size()
	redraw(s, tb, w, h)
	for {
		ev := s.PollEvent()
		switch ev.(type) {
		case *tcell.EventKey:
			switch ev.(*tcell.EventKey).Rune() {
			case ' ': // page down
				curY += h
				if curY >= maxY {
					curY = maxY - h
				}
				redraw(s, tb, w, h)
			case 'b': // page down
				curY -= h
				if curY < 0 {
					curY = 0
				}
				redraw(s, tb, w, h)
			case 'g': // top
				curY = 0
				redraw(s, tb, w, h)
			case 'G': // bottom
				curY = maxY - h - 1
				if curY < 0 {
					curY = 0
				}
				redraw(s, tb, w, h)
			case 'j': // down
				if curY < maxY {
					curY++
				}
				redraw(s, tb, w, h)
			case 'k': // up
				if curY > 0 {
					curY--
				}
				redraw(s, tb, w, h)
			case 'h': // left
				if curX > 0 {
					curX--
				}
				redraw(s, tb, w, h)
			case 'l': // right
				if curX < maxX {
					curX++
				}
				redraw(s, tb, w, h)
			case 'q': // quit
				return nil
			}
		case *tcell.EventResize:
			w, h = s.Size()
			s.Clear()
			draw(s, tb, w, h)
			s.Sync()
		}
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s filename\n", os.Args[0])
	os.Exit(1)
}

func fatal(err error) {
	fmt.Fprintf(os.Stderr, "%s: error: %s\n", os.Args[0], err)
	os.Exit(1)
}

func main() {
	if len(os.Args) != 2 {
		usage()
	}
	if err := least(os.Args[1]); err != nil {
		fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"html/template"
	"image"
	"image/png"
	"net/http"
	"os"
	"runtime/pprof"
	"sync"
	"time"

	"github.com/maruel/go-lepton/lepton"
	"github.com/maruel/interrupt"
)

type imageRing struct {
	c chan *lepton.LeptonBuffer
}

func makeImageRing() *imageRing {
	return &imageRing{c: make(chan *lepton.LeptonBuffer, 16)}
}

func (i *imageRing) get() *lepton.LeptonBuffer {
	select {
	case b := <-i.c:
		return b
	default:
		return &lepton.LeptonBuffer{}
	}
}

func (i *imageRing) done(b *lepton.LeptonBuffer) {
	if len(i.c) < 8 {
		i.c <- b
	}
}

type doubleBuffer struct {
	lock        sync.Mutex
	frontBuffer *image.Gray
	backBuffer  *image.Gray
	Stats       lepton.Stats
	Min         uint16
	Max         uint16
}

var currentImage doubleBuffer

var rootTmpl = template.Must(template.New("name").Parse(`

	<title>go-lepton</title>

	Still:

	{{.Stats}}

	{{.Min}} - {{.Max}}

	`))

func root(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html")
	currentImage.lock.Lock()
	rootTmpl.Execute(w, currentImage)
	currentImage.lock.Unlock()
}

func still(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "image/png")
	w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate")
	currentImage.lock.Lock()
	defer currentImage.lock.Unlock()
	png.Encode(w, currentImage.frontBuffer)
}

func mainImpl() error {
	cpuprofile := flag.String("cpuprofile", "", "dump CPU profile in file")
	port := flag.Int("port", 8010, "http port to listen on")
	flag.Parse()

	if len(flag.Args()) != 0 {
		return fmt.Errorf("unexpected argument: %s", flag.Args())
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			return err
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	interrupt.HandleCtrlC()

	l, err := lepton.MakeLepton()
	if l != nil {
		defer l.Close()
	}
	if err != nil {
		return err
	}

	c := make(chan *lepton.LeptonBuffer, 16)

	currentImage.frontBuffer = image.NewGray(image.Rect(0, 0, 80, 60))
	currentImage.backBuffer = image.NewGray(image.Rect(0, 0, 80, 60))
	ring := makeImageRing()

	go func() {
		for {
			// Keep this loop busy to not lose sync on SPI.
			b := ring.get()
			l.ReadImg(b)
			c <- b
		}
	}()

	go func() {
		for {
			// Processing is done in a separate loop to not miss a frame.
			img := <-c
			lepton.Scale(currentImage.backBuffer, img)
			ring.done(img)
			currentImage.lock.Lock()
			currentImage.backBuffer, currentImage.frontBuffer = currentImage.frontBuffer, currentImage.backBuffer
			currentImage.Min = img.Min
			currentImage.Max = img.Max
			currentImage.lock.Unlock()
		}
	}()

	http.HandleFunc("/", root)
	http.HandleFunc("/favicon.ico", still)
	http.HandleFunc("/still.png", still)
	fmt.Printf("Listening on %d\n", *port)
	go http.ListenAndServe(fmt.Sprintf(":%d", *port), nil)

	for !interrupt.IsSet() {
		stats := l.Stats()
		currentImage.lock.Lock()
		currentImage.Stats = stats
		currentImage.lock.Unlock()
		fmt.Printf("\r%d frames %d duped %d dummy %d badsync %d broken %d fail", stats.GoodFrames, stats.DuplicateFrames, stats.DummyLines, stats.SyncFailures, stats.BrokenPackets, stats.TransferFails)
		time.Sleep(time.Second)
	}
	fmt.Print("\n")
	return nil
}

func main() {
	if err := mainImpl(); err != nil {
		fmt.Fprintf(os.Stderr, "\ngo-lepton: %s.\n", err)
		os.Exit(1)
	}
}
<commit_msg>Make image larger<commit_after>// Copyright 2015 Marc-Antoine Ruel. All rights reserved.
// Use of this source code is governed under the Apache License, Version 2.0
// that can be found in the LICENSE file.

package main

import (
	"flag"
	"fmt"
	"html/template"
	"image"
	"image/png"
	"net/http"
	"os"
	"runtime/pprof"
	"sync"
	"time"

	"github.com/maruel/go-lepton/lepton"
	"github.com/maruel/interrupt"
)

type imageRing struct {
	c chan *lepton.LeptonBuffer
}

func makeImageRing() *imageRing {
	return &imageRing{c: make(chan *lepton.LeptonBuffer, 16)}
}

func (i *imageRing) get() *lepton.LeptonBuffer {
	select {
	case b := <-i.c:
		return b
	default:
		return &lepton.LeptonBuffer{}
	}
}

func (i *imageRing) done(b *lepton.LeptonBuffer) {
	if len(i.c) < 8 {
		i.c <- b
	}
}

type doubleBuffer struct {
	lock        sync.Mutex
	frontBuffer *image.Gray
	backBuffer  *image.Gray
	Stats       lepton.Stats
	Min         uint16
	Max         uint16
}

var currentImage doubleBuffer

var rootTmpl = template.Must(template.New("name").Parse(`

		<title>go-lepton</title>

	Still:

	{{.Stats}}

	{{.Min}} - {{.Max}}

	`))

func root(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html")
	currentImage.lock.Lock()
	rootTmpl.Execute(w, currentImage)
	currentImage.lock.Unlock()
}

func still(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "image/png")
	w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate")
	currentImage.lock.Lock()
	defer currentImage.lock.Unlock()
	png.Encode(w, currentImage.frontBuffer)
}

func mainImpl() error {
	cpuprofile := flag.String("cpuprofile", "", "dump CPU profile in file")
	port := flag.Int("port", 8010, "http port to listen on")
	flag.Parse()

	if len(flag.Args()) != 0 {
		return fmt.Errorf("unexpected argument: %s", flag.Args())
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			return err
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	interrupt.HandleCtrlC()

	l, err := lepton.MakeLepton()
	if l != nil {
		defer l.Close()
	}
	if err != nil {
		return err
	}

	c := make(chan *lepton.LeptonBuffer, 16)

	currentImage.frontBuffer = image.NewGray(image.Rect(0, 0, 80, 60))
	currentImage.backBuffer = image.NewGray(image.Rect(0, 0, 80, 60))
	ring := makeImageRing()

	go func() {
		for {
			// Keep this loop busy to not lose sync on SPI.
			b := ring.get()
			l.ReadImg(b)
			c <- b
		}
	}()

	go func() {
		for {
			// Processing is done in a separate loop to not miss a frame.
			img := <-c
			lepton.Scale(currentImage.backBuffer, img)
			ring.done(img)
			currentImage.lock.Lock()
			currentImage.backBuffer, currentImage.frontBuffer = currentImage.frontBuffer, currentImage.backBuffer
			currentImage.Min = img.Min
			currentImage.Max = img.Max
			currentImage.lock.Unlock()
		}
	}()

	http.HandleFunc("/", root)
	http.HandleFunc("/favicon.ico", still)
	http.HandleFunc("/still.png", still)
	fmt.Printf("Listening on %d\n", *port)
	go http.ListenAndServe(fmt.Sprintf(":%d", *port), nil)

	for !interrupt.IsSet() {
		stats := l.Stats()
		currentImage.lock.Lock()
		currentImage.Stats = stats
		currentImage.lock.Unlock()
		fmt.Printf("\r%d frames %d duped %d dummy %d badsync %d broken %d fail", stats.GoodFrames, stats.DuplicateFrames, stats.DummyLines, stats.SyncFailures, stats.BrokenPackets, stats.TransferFails)
		time.Sleep(time.Second)
	}
	fmt.Print("\n")
	return nil
}

func main() {
	if err := mainImpl(); err != nil {
		fmt.Fprintf(os.Stderr, "\ngo-lepton: %s.\n", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"strings"

	"github.com/uli-go/xz/lzma"
	"github.com/uli-go/xz/xlog"
)

type packer interface {
	outputPaths(path string) (outputPath, tmpPath string, err error)
	pack(w io.Writer, r io.Reader, preset int) (n int64, err error)
}

const lzmaSuffix = ".lzma"

// parameters converts the lzmago executable flags to lzma parameters.
//
// I cannot use the preset config from the Tukaani project directly,
// because I don't have two algorithm modes and can't support parameters
// like nice_len or depth. So at this point in time I stay with the
// dictionary sizes the default combination of (LC,LP,LB) = (3,0,2).
// The default preset is 6.
// Following list provides exponents of two for the dictionary sizes:
// 18, 20, 21, 22, 22, 23, 23, 24, 25, 26.
func parameters(preset int) lzma.Parameters {
	dictSizeExps := []uint{18, 20, 21, 22, 22, 23, 23, 24, 25, 26}
	dictSize := int64(1) << dictSizeExps[preset]
	p := lzma.Parameters{
		LC:           3,
		LP:           0,
		PB:           2,
		DictSize:     dictSize,
		EOS:          true,
		ExtraBufSize: 16 * 1024,
	}
	return p
}

type lzmaPacker struct{}

func (p lzmaPacker) outputPaths(path string) (out, tmp string, err error) {
	if path == "-" {
		return "-", "-", nil
	}
	if path == "" {
		err = errors.New("path is empty")
		return
	}
	if strings.HasSuffix(path, lzmaSuffix) {
		err = fmt.Errorf("path %s has suffix %s -- ignored",
			path, lzmaSuffix)
		return
	}
	out = path + lzmaSuffix
	tmp = out + ".pack"
	return
}

func (p lzmaPacker) pack(w io.Writer, r io.Reader, preset int) (n int64, err error) {
	if w == nil {
		panic("writer w is nil")
	}
	if r == nil {
		panic("reader r is nil")
	}
	params := parameters(preset)
	bw := bufio.NewWriter(w)
	err = bw.Flush()
	lw, err := lzma.NewWriterParams(bw, params)
	if err != nil {
		return
	}
	err = bw.Flush()
	n, err = io.Copy(lw, r)
	if err != nil {
		return
	}
	err = bw.Flush()
	if err = lw.Close(); err != nil {
		return
	}
	err = bw.Flush()
	return
}

type lzmaUnpacker struct{}

func (u lzmaUnpacker) outputPaths(path string) (out, tmp string, err error) {
	if path == "-" {
		return "-", "-", nil
	}
	if !strings.HasSuffix(path, lzmaSuffix) {
		err = fmt.Errorf("path %s has no suffix %s",
			path, lzmaSuffix)
		return
	}
	base := filepath.Base(path)
	if base == lzmaSuffix {
		err = fmt.Errorf(
			"path %s has only suffix %s as filename",
			path, lzmaSuffix)
		return
	}
	out = path[:len(path)-len(lzmaSuffix)]
	tmp = out + ".unpack"
	return
}

func (u lzmaUnpacker) pack(w io.Writer, r io.Reader, preset int) (n int64, err error) {
	if w == nil {
		panic("writer w is nil")
	}
	if r == nil {
		panic("reader r is nil")
	}
	// pack actually unpacks
	br := bufio.NewReader(r)
	lr, err := lzma.NewReader(br)
	if err != nil {
		return
	}
	n, err = io.Copy(w, lr)
	return
}

func signalHandler(tmpPath string) chan<- struct{} {
	quit := make(chan struct{})
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, os.Interrupt)
	go func() {
		select {
		case <-quit:
			signal.Stop(sigch)
			return
		case <-sigch:
			if tmpPath != "-" {
				os.Remove(tmpPath)
			}
			os.Exit(7)
		}
	}()
	return quit
}

func packFile(pck packer, path, tmpPath string, opts *options) error {
	var err error

	// open reader
	var r *os.File
	if path == "-" {
		r = os.Stdin
	} else {
		fi, err := os.Lstat(path)
		if err != nil {
			return err
		}
		if !fi.Mode().IsRegular() {
			return fmt.Errorf("%s is not a regular file", path)
		}
		r, err = os.Open(path)
		if err != nil {
			return err
		}
		fi, err = r.Stat()
		if err != nil {
			r.Close()
			return err
		}
		if !fi.Mode().IsRegular() {
			r.Close()
			return fmt.Errorf("%s is not a regular file", path)
		}
	}
	defer func() {
		if err != nil {
			r.Close()
		} else {
			err = r.Close()
		}
	}()

	// open writer
	var w *os.File
	if tmpPath == "-" {
		w = os.Stdout
	} else {
		if opts.force {
			os.Remove(tmpPath)
		}
		w, err = os.OpenFile(tmpPath,
			os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0666)
		if err != nil {
			return err
		}
		defer func() {
			if err != nil {
				w.Close()
			} else {
				err = w.Close()
			}
		}()
		fi, err := w.Stat()
		if err != nil {
			return err
		}
		if !fi.Mode().IsRegular() {
			return fmt.Errorf("%s is not a regular file", tmpPath)
		}
	}

	_, err = pck.pack(w, r, opts.preset)
	return err
}

func processFile(path string, opts *options) {
	var pck packer
	if opts.decompress {
		pck = lzmaUnpacker{}
	} else {
		pck = lzmaPacker{}
	}
	outputPath, tmpPath, err := pck.outputPaths(path)
	if err != nil {
		xlog.Warn(err)
		return
	}
	if opts.stdout {
		outputPath, tmpPath = "-", "-"
	}
	if outputPath != "-" {
		_, err = os.Lstat(outputPath)
		if err == nil && !opts.force {
			xlog.Warnf("file %s exists", outputPath)
			return
		}
	}
	defer func(tmp string) {
		if tmp != "-" {
			os.Remove(tmp)
		}
	}(tmpPath)
	quit := signalHandler(tmpPath)
	defer close(quit)

	if err = packFile(pck, path, tmpPath, opts); err != nil {
		xlog.Warn(err)
		return
	}
	if tmpPath != "-" && outputPath != "-" {
		if err = os.Rename(tmpPath, outputPath); err != nil {
			xlog.Warn(err)
			return
		}
	}
	if !opts.keep && !opts.stdout && path != "-" {
		if err = os.Remove(path); err != nil {
			xlog.Warn(err)
			return
		}
	}
}
<commit_msg>lzma.go: removed unneeded calls to Flush<commit_after>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"os/signal"
	"path/filepath"
	"strings"

	"github.com/uli-go/xz/lzma"
	"github.com/uli-go/xz/xlog"
)

type packer interface {
	outputPaths(path string) (outputPath, tmpPath string, err error)
	pack(w io.Writer, r io.Reader, preset int) (n int64, err error)
}

const lzmaSuffix = ".lzma"

// parameters converts the lzmago executable flags to lzma parameters.
//
// I cannot use the preset config from the Tukaani project directly,
// because I don't have two algorithm modes and can't support parameters
// like nice_len or depth. So at this point in time I stay with the
// dictionary sizes the default combination of (LC,LP,LB) = (3,0,2).
// The default preset is 6.
// Following list provides exponents of two for the dictionary sizes:
// 18, 20, 21, 22, 22, 23, 23, 24, 25, 26.
func parameters(preset int) lzma.Parameters {
	dictSizeExps := []uint{18, 20, 21, 22, 22, 23, 23, 24, 25, 26}
	dictSize := int64(1) << dictSizeExps[preset]
	p := lzma.Parameters{
		LC:           3,
		LP:           0,
		PB:           2,
		DictSize:     dictSize,
		EOS:          true,
		ExtraBufSize: 16 * 1024,
	}
	return p
}

type lzmaPacker struct{}

func (p lzmaPacker) outputPaths(path string) (out, tmp string, err error) {
	if path == "-" {
		return "-", "-", nil
	}
	if path == "" {
		err = errors.New("path is empty")
		return
	}
	if strings.HasSuffix(path, lzmaSuffix) {
		err = fmt.Errorf("path %s has suffix %s -- ignored",
			path, lzmaSuffix)
		return
	}
	out = path + lzmaSuffix
	tmp = out + ".pack"
	return
}

func (p lzmaPacker) pack(w io.Writer, r io.Reader, preset int) (n int64, err error) {
	if w == nil {
		panic("writer w is nil")
	}
	if r == nil {
		panic("reader r is nil")
	}
	params := parameters(preset)
	bw := bufio.NewWriter(w)
	lw, err := lzma.NewWriterParams(bw, params)
	if err != nil {
		return
	}
	n, err = io.Copy(lw, r)
	if err != nil {
		return
	}
	if err = lw.Close(); err != nil {
		return
	}
	err = bw.Flush()
	return
}

type lzmaUnpacker struct{}

func (u lzmaUnpacker) outputPaths(path string) (out, tmp string, err error) {
	if path == "-" {
		return "-", "-", nil
	}
	if !strings.HasSuffix(path, lzmaSuffix) {
		err = fmt.Errorf("path %s has no suffix %s",
			path, lzmaSuffix)
		return
	}
	base := filepath.Base(path)
	if base == lzmaSuffix {
		err = fmt.Errorf(
			"path %s has only suffix %s as filename",
			path, lzmaSuffix)
		return
	}
	out = path[:len(path)-len(lzmaSuffix)]
	tmp = out + ".unpack"
	return
}

func (u lzmaUnpacker) pack(w io.Writer, r io.Reader, preset int) (n int64, err error) {
	if w == nil {
		panic("writer w is nil")
	}
	if r == nil {
		panic("reader r is nil")
	}
	// pack actually unpacks
	br := bufio.NewReader(r)
	lr, err := lzma.NewReader(br)
	if err != nil {
		return
	}
	n, err = io.Copy(w, lr)
	return
}

func signalHandler(tmpPath string) chan<- struct{} {
	quit := make(chan struct{})
	sigch := make(chan os.Signal, 1)
	signal.Notify(sigch, os.Interrupt)
	go func() {
		select {
		case <-quit:
			signal.Stop(sigch)
			return
		case <-sigch:
			if tmpPath != "-" {
				os.Remove(tmpPath)
			}
			os.Exit(7)
		}
	}()
	return quit
}

func packFile(pck packer, path, tmpPath string, opts *options) error {
	var err error

	// open reader
	var r *os.File
	if path == "-" {
		r = os.Stdin
	} else {
		fi, err := os.Lstat(path)
		if err != nil {
			return err
		}
		if !fi.Mode().IsRegular() {
			return fmt.Errorf("%s is not a regular file", path)
		}
		r, err = os.Open(path)
		if err != nil {
			return err
		}
		fi, err = r.Stat()
		if err != nil {
			r.Close()
			return err
		}
		if !fi.Mode().IsRegular() {
			r.Close()
			return fmt.Errorf("%s is not a regular file", path)
		}
	}
	defer func() {
		if err != nil {
			r.Close()
		} else {
			err = r.Close()
		}
	}()

	// open writer
	var w *os.File
	if tmpPath == "-" {
		w = os.Stdout
	} else {
		if opts.force {
			os.Remove(tmpPath)
		}
		w, err = os.OpenFile(tmpPath,
			os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0666)
		if err != nil {
			return err
		}
		defer func() {
			if err != nil {
				w.Close()
			} else {
				err = w.Close()
			}
		}()
		fi, err := w.Stat()
		if err != nil {
			return err
		}
		if !fi.Mode().IsRegular() {
			return fmt.Errorf("%s is not a regular file", tmpPath)
		}
	}

	_, err = pck.pack(w, r, opts.preset)
	return err
}

func processFile(path string, opts *options) {
	var pck packer
	if opts.decompress {
		pck = lzmaUnpacker{}
	} else {
		pck = lzmaPacker{}
	}
	outputPath, tmpPath, err := pck.outputPaths(path)
	if err != nil {
		xlog.Warn(err)
		return
	}
	if opts.stdout {
		outputPath, tmpPath = "-", "-"
	}
	if outputPath != "-" {
		_, err = os.Lstat(outputPath)
		if err == nil && !opts.force {
			xlog.Warnf("file %s exists", outputPath)
			return
		}
	}
	defer func(tmp string) {
		if tmp != "-" {
			os.Remove(tmp)
		}
	}(tmpPath)
	quit := signalHandler(tmpPath)
	defer close(quit)

	if err = packFile(pck, path, tmpPath, opts); err != nil {
		xlog.Warn(err)
		return
	}
	if tmpPath != "-" && outputPath != "-" {
		if err = os.Rename(tmpPath, outputPath); err != nil {
			xlog.Warn(err)
			return
		}
	}
	if !opts.keep && !opts.stdout && path != "-" {
		if err = os.Remove(path); err != nil {
			xlog.Warn(err)
			return
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/jinzhu/gorm"
	"github.com/modrzew/malusers"
	"github.com/modrzew/malusers/api"
)

func logger(inner http.HandlerFunc) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		inner(w, r)
		log.Printf(
			"%s\t%s\t%s",
			r.Method,
			r.RequestURI,
			time.Since(start),
		)
	})
}

func addHandlers(router *mux.Router, db *gorm.DB, cache *api.Cache) {
	handlers := &api.Handlers{DB: db, Cache: cache}
	router.Handle("/user/{username}", logger(handlers.GetUserStats)).Methods("GET")
}

func main() {
	config := malusers.ReadConfig()
	db := malusers.OpenDb()
	cache := api.GetCache(db)

	router := mux.NewRouter()
	addHandlers(router, db, cache)
	url := fmt.Sprintf("%s:%d", config.API.Host, config.API.Port)
	http.ListenAndServe(url, router)
}
<commit_msg>Add CORS header<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/jinzhu/gorm"
	"github.com/modrzew/malusers"
	"github.com/modrzew/malusers/api"
)

func logger(inner http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		inner(w, r)
		log.Printf(
			"%s\t%s\t%s",
			r.Method,
			r.RequestURI,
			time.Since(start),
		)
	})
}

func cors(inner http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		inner(w, r)
	})
}

func addHandlers(router *mux.Router, db *gorm.DB, cache *api.Cache) {
	handlers := &api.Handlers{DB: db, Cache: cache}
	router.Handle("/user/{username}", logger(cors(handlers.GetUserStats))).Methods("GET")
}

func main() {
	config := malusers.ReadConfig()
	db := malusers.OpenDb()
	cache := api.GetCache(db)

	router := mux.NewRouter()
	addHandlers(router, db, cache)
	url := fmt.Sprintf("%s:%d", config.API.Host, config.API.Port)
	http.ListenAndServe(url, router)
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"os"
	"path/filepath"
	"sort"
	"strings"

	notaryclient "github.com/docker/notary/client"
	"github.com/docker/notary/cryptoservice"
	"github.com/docker/notary/passphrase"
	"github.com/docker/notary/trustmanager"

	"github.com/docker/notary/tuf/data"
	"github.com/spf13/cobra"
)

func init() {
	cmdKey.AddCommand(cmdKeyList)
	cmdKey.AddCommand(cmdKeyGenerateRootKey)

	cmdKeyExport.Flags().StringVarP(&keysExportGUN, "gun", "g", "", "Globally Unique Name to export keys for")
	cmdKey.AddCommand(cmdKeyExport)
	cmdKey.AddCommand(cmdKeyExportRoot)
	cmdKeyExportRoot.Flags().BoolVarP(&keysExportRootChangePassphrase, "change-passphrase", "p", false, "Set a new passphrase for the key being exported")
	cmdKey.AddCommand(cmdKeyImport)
	cmdKey.AddCommand(cmdKeyImportRoot)
	cmdKey.AddCommand(cmdRotateKey)
}

var cmdKey = &cobra.Command{
	Use:   "key",
	Short: "Operates on keys.",
	Long:  `Operations on private keys.`,
}

var cmdKeyList = &cobra.Command{
	Use:   "list",
	Short: "Lists keys.",
	Long:  "Lists all keys known to notary.",
	Run:   keysList,
}

var cmdRotateKey = &cobra.Command{
	Use:   "rotate [ GUN ]",
	Short: "Rotate all the signing (non-root) keys for the given Globally Unique Name.",
	Long:  "Removes all old signing (non-root) keys for the given Globally Unique Name, and generates new ones.  This only makes local changes - please use then `notary publish` to push the key rotation changes to the remote server.",
	Run:   keysRotate,
}

var cmdKeyGenerateRootKey = &cobra.Command{
	Use:   "generate [ algorithm ]",
	Short: "Generates a new root key with a given algorithm.",
	Long:  "Generates a new root key with a given algorithm. If a hardware smartcard is available, the key will be stored both on hardware and on disk.  Please make sure to back up the key that is written to disk, and to then take the on-disk key offline.",
	Run:   keysGenerateRootKey,
}

var keysExportGUN string

var cmdKeyExport = &cobra.Command{
	Use:   "export [ filename ]",
	Short: "Exports keys to a ZIP file.",
	Long:  "Exports a collection of keys. The keys are reencrypted with a new passphrase. The output is a ZIP file.  If the --gun option is passed, only signing keys and no root keys will be exported.  Does not work on keys that are only in hardware (smartcards).",
	Run:   keysExport,
}

var keysExportRootChangePassphrase bool

var cmdKeyExportRoot = &cobra.Command{
	Use:   "export-root [ keyID ] [ filename ]",
	Short: "Exports given root key to a file.",
	Long:  "Exports a root key, without reencrypting. The output is a PEM file. Does not work on keys that are only in hardware (smartcards).",
	Run:   keysExportRoot,
}

var cmdKeyImport = &cobra.Command{
	Use:   "import [ filename ]",
	Short: "Imports keys from a ZIP file.",
	Long:  "Imports one or more keys from a ZIP file. If a hardware smartcard is available, the root key will be imported into the smartcard but not to disk.",
	Run:   keysImport,
}

var cmdKeyImportRoot = &cobra.Command{
	Use:   "import-root [ filename ]",
	Short: "Imports root key.",
	Long:  "Imports a root key from a PEM file. If a hardware smartcard is available, the root key will be imported into the smartcard but not to disk.",
	Run:   keysImportRoot,
}

func keysList(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		cmd.Usage()
		os.Exit(1)
	}

	parseConfig()

	stores := getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)

	keys := make(map[trustmanager.KeyStore]map[string]string)
	for _, store := range stores {
		keys[store] = store.ListKeys()
	}

	cmd.Println("")
	cmd.Println("# Root keys: ")
	for store, keysMap := range keys {
		for k, v := range keysMap {
			if v == "root" {
				cmd.Println(k, "-", store.Name())
			}
		}
	}

	cmd.Println("")
	cmd.Println("# Signing keys: ")

	// Get a list of all the keys
	for store, keysMap := range keys {
		var sortedKeys []string
		for k := range keysMap {
			sortedKeys = append(sortedKeys, k)
		}

		// Sort the list of all the keys
		sort.Strings(sortedKeys)

		// Print a sorted list of the key/role
		for _, k := range sortedKeys {
			if keysMap[k] != "root" {
				printKey(cmd, k, keysMap[k], store.Name())
			}
		}
	}
}

func keysGenerateRootKey(cmd *cobra.Command, args []string) {
	// We require one or no arguments (since we have a default value), but if the
	// user passes in more than one argument, we error out.
	if len(args) > 1 {
		cmd.Usage()
		fatalf("Please provide only one Algorithm as an argument to generate (rsa, ecdsa)")
	}

	parseConfig()

	// If no param is given to generate, generates an ecdsa key by default
	algorithm := data.ECDSAKey

	// If we were provided an argument lets attempt to use it as an algorithm
	if len(args) > 0 {
		algorithm = args[0]
	}

	allowedCiphers := map[string]bool{
		data.ECDSAKey: true,
		data.RSAKey:   true,
	}

	if !allowedCiphers[strings.ToLower(algorithm)] {
		fatalf("Algorithm not allowed, possible values are: RSA, ECDSA")
	}

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	pubKey, err := cs.Create(data.CanonicalRootRole, algorithm)
	if err != nil {
		fatalf("Failed to create a new root key: %v", err)
	}

	cmd.Printf("Generated new %s root key with keyID: %s\n", algorithm, pubKey.ID())
}

// keysExport exports a collection of keys to a ZIP file
func keysExport(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cmd.Usage()
		fatalf("Must specify output filename for export")
	}

	parseConfig()
	exportFilename := args[0]

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	exportFile, err := os.Create(exportFilename)
	if err != nil {
		fatalf("Error creating output file: %v", err)
	}

	// Must use a different passphrase retriever to avoid caching the
	// unlocking passphrase and reusing that.
	exportRetriever := getRetriever()
	if keysExportGUN != "" {
		err = cs.ExportKeysByGUN(exportFile, keysExportGUN, exportRetriever)
	} else {
		err = cs.ExportAllKeys(exportFile, exportRetriever)
	}

	exportFile.Close()

	if err != nil {
		os.Remove(exportFilename)
		fatalf("Error exporting keys: %v", err)
	}
}

// keysExportRoot exports a root key by ID to a PEM file
func keysExportRoot(cmd *cobra.Command, args []string) {
	if len(args) < 2 {
		cmd.Usage()
		fatalf("Must specify key ID and output filename for export")
	}

	parseConfig()

	keyID := args[0]
	exportFilename := args[1]

	if len(keyID) != idSize {
		fatalf("Please specify a valid root key ID")
	}

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	exportFile, err := os.Create(exportFilename)
	if err != nil {
		fatalf("Error creating output file: %v", err)
	}
	if keysExportRootChangePassphrase {
		// Must use a different passphrase retriever to avoid caching the
		// unlocking passphrase and reusing that.
		exportRetriever := getRetriever()
		err = cs.ExportRootKeyReencrypt(exportFile, keyID, exportRetriever)
	} else {
		err = cs.ExportRootKey(exportFile, keyID)
	}
	exportFile.Close()
	if err != nil {
		os.Remove(exportFilename)
		fatalf("Error exporting root key: %v", err)
	}
}

// keysImport imports keys from a ZIP file
func keysImport(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cmd.Usage()
		fatalf("Must specify input filename for import")
	}

	importFilename := args[0]

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	zipReader, err := zip.OpenReader(importFilename)
	if err != nil {
		fatalf("Opening file for import: %v", err)
	}
	defer zipReader.Close()

	err = cs.ImportKeysZip(zipReader.Reader)

	if err != nil {
		fatalf("Error importing keys: %v", err)
	}
}

// keysImportRoot imports a root key from a PEM file
func keysImportRoot(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		cmd.Usage()
		fatalf("Must specify input filename for import")
	}

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	importFilename := args[0]

	importFile, err := os.Open(importFilename)
	if err != nil {
		fatalf("Opening file for import: %v", err)
	}
	defer importFile.Close()

	err = cs.ImportRootKey(importFile)

	if err != nil {
		fatalf("Error importing root key: %v", err)
	}
}

func printKey(cmd *cobra.Command, keyPath, alias, loc string) {
	keyID := filepath.Base(keyPath)
	gun := filepath.Dir(keyPath)
	cmd.Printf("%s - %s - %s - %s\n", gun, alias, keyID, loc)
}

func keysRotate(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cmd.Usage()
		fatalf("Must specify a GUN and target")
	}
	parseConfig()

	gun := args[0]
	nRepo, err := notaryclient.NewNotaryRepository(mainViper.GetString("trust_dir"), gun, remoteTrustServer, nil, retriever)
	if err != nil {
		fatalf(err.Error())
	}
	if err := nRepo.RotateKeys(); err != nil {
		fatalf(err.Error())
	}
}

func getKeyStores(cmd *cobra.Command, directory string,
	ret passphrase.Retriever, withHardware bool) []trustmanager.KeyStore {

	fileKeyStore, err := trustmanager.NewKeyFileStore(directory, ret)
	if err != nil {
		fatalf("Failed to create private key store in directory: %s", directory)
	}

	ks := []trustmanager.KeyStore{fileKeyStore}

	if withHardware {
		yubiStore, err := getYubiKeyStore(fileKeyStore, ret)
		if err == nil && yubiStore != nil {
			// Note that the order is important, since we want to prioritize
			// the yubikey store
			ks = []trustmanager.KeyStore{yubiStore, fileKeyStore}
		}
	}

	return ks
}
<commit_msg>set withHardware flag to false for export commands. We can never export from hardware Signed-off-by: David Lawrence <david.lawrence@docker.com> (github: endophage)<commit_after>package main

import (
	"archive/zip"
	"os"
	"path/filepath"
	"sort"
	"strings"

	notaryclient "github.com/docker/notary/client"
	"github.com/docker/notary/cryptoservice"
	"github.com/docker/notary/passphrase"
	"github.com/docker/notary/trustmanager"

	"github.com/docker/notary/tuf/data"
	"github.com/spf13/cobra"
)

func init() {
	cmdKey.AddCommand(cmdKeyList)
	cmdKey.AddCommand(cmdKeyGenerateRootKey)

	cmdKeyExport.Flags().StringVarP(&keysExportGUN, "gun", "g", "", "Globally Unique Name to export keys for")
	cmdKey.AddCommand(cmdKeyExport)
	cmdKey.AddCommand(cmdKeyExportRoot)
	cmdKeyExportRoot.Flags().BoolVarP(&keysExportRootChangePassphrase, "change-passphrase", "p", false, "Set a new passphrase for the key being exported")
	cmdKey.AddCommand(cmdKeyImport)
	cmdKey.AddCommand(cmdKeyImportRoot)
	cmdKey.AddCommand(cmdRotateKey)
}

var cmdKey = &cobra.Command{
	Use:   "key",
	Short: "Operates on keys.",
	Long:  `Operations on private keys.`,
}

var cmdKeyList = &cobra.Command{
	Use:   "list",
	Short: "Lists keys.",
	Long:  "Lists all keys known to notary.",
	Run:   keysList,
}

var cmdRotateKey = &cobra.Command{
	Use:   "rotate [ GUN ]",
	Short: "Rotate all the signing (non-root) keys for the given Globally Unique Name.",
	Long:  "Removes all old signing (non-root) keys for the given Globally Unique Name, and generates new ones.  This only makes local changes - please use then `notary publish` to push the key rotation changes to the remote server.",
	Run:   keysRotate,
}

var cmdKeyGenerateRootKey = &cobra.Command{
	Use:   "generate [ algorithm ]",
	Short: "Generates a new root key with a given algorithm.",
	Long:  "Generates a new root key with a given algorithm. If a hardware smartcard is available, the key will be stored both on hardware and on disk.  Please make sure to back up the key that is written to disk, and to then take the on-disk key offline.",
	Run:   keysGenerateRootKey,
}

var keysExportGUN string

var cmdKeyExport = &cobra.Command{
	Use:   "export [ filename ]",
	Short: "Exports keys to a ZIP file.",
	Long:  "Exports a collection of keys. The keys are reencrypted with a new passphrase. The output is a ZIP file.  If the --gun option is passed, only signing keys and no root keys will be exported.  Does not work on keys that are only in hardware (smartcards).",
	Run:   keysExport,
}

var keysExportRootChangePassphrase bool

var cmdKeyExportRoot = &cobra.Command{
	Use:   "export-root [ keyID ] [ filename ]",
	Short: "Exports given root key to a file.",
	Long:  "Exports a root key, without reencrypting. The output is a PEM file. Does not work on keys that are only in hardware (smartcards).",
	Run:   keysExportRoot,
}

var cmdKeyImport = &cobra.Command{
	Use:   "import [ filename ]",
	Short: "Imports keys from a ZIP file.",
	Long:  "Imports one or more keys from a ZIP file. If a hardware smartcard is available, the root key will be imported into the smartcard but not to disk.",
	Run:   keysImport,
}

var cmdKeyImportRoot = &cobra.Command{
	Use:   "import-root [ filename ]",
	Short: "Imports root key.",
	Long:  "Imports a root key from a PEM file. If a hardware smartcard is available, the root key will be imported into the smartcard but not to disk.",
	Run:   keysImportRoot,
}

func keysList(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		cmd.Usage()
		os.Exit(1)
	}

	parseConfig()

	stores := getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)

	keys := make(map[trustmanager.KeyStore]map[string]string)
	for _, store := range stores {
		keys[store] = store.ListKeys()
	}

	cmd.Println("")
	cmd.Println("# Root keys: ")
	for store, keysMap := range keys {
		for k, v := range keysMap {
			if v == "root" {
				cmd.Println(k, "-", store.Name())
			}
		}
	}

	cmd.Println("")
	cmd.Println("# Signing keys: ")

	// Get a list of all the keys
	for store, keysMap := range keys {
		var sortedKeys []string
		for k := range keysMap {
			sortedKeys = append(sortedKeys, k)
		}

		// Sort the list of all the keys
		sort.Strings(sortedKeys)

		// Print a sorted list of the key/role
		for _, k := range sortedKeys {
			if keysMap[k] != "root" {
				printKey(cmd, k, keysMap[k], store.Name())
			}
		}
	}
}

func keysGenerateRootKey(cmd *cobra.Command, args []string) {
	// We require one or no arguments (since we have a default value), but if the
	// user passes in more than one argument, we error out.
	if len(args) > 1 {
		cmd.Usage()
		fatalf("Please provide only one Algorithm as an argument to generate (rsa, ecdsa)")
	}

	parseConfig()

	// If no param is given to generate, generates an ecdsa key by default
	algorithm := data.ECDSAKey

	// If we were provided an argument lets attempt to use it as an algorithm
	if len(args) > 0 {
		algorithm = args[0]
	}

	allowedCiphers := map[string]bool{
		data.ECDSAKey: true,
		data.RSAKey:   true,
	}

	if !allowedCiphers[strings.ToLower(algorithm)] {
		fatalf("Algorithm not allowed, possible values are: RSA, ECDSA")
	}

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	pubKey, err := cs.Create(data.CanonicalRootRole, algorithm)
	if err != nil {
		fatalf("Failed to create a new root key: %v", err)
	}

	cmd.Printf("Generated new %s root key with keyID: %s\n", algorithm, pubKey.ID())
}

// keysExport exports a collection of keys to a ZIP file
func keysExport(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cmd.Usage()
		fatalf("Must specify output filename for export")
	}

	parseConfig()
	exportFilename := args[0]

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, false)...,
	)

	exportFile, err := os.Create(exportFilename)
	if err != nil {
		fatalf("Error creating output file: %v", err)
	}

	// Must use a different passphrase retriever to avoid caching the
	// unlocking passphrase and reusing that.
	exportRetriever := getRetriever()
	if keysExportGUN != "" {
		err = cs.ExportKeysByGUN(exportFile, keysExportGUN, exportRetriever)
	} else {
		err = cs.ExportAllKeys(exportFile, exportRetriever)
	}

	exportFile.Close()

	if err != nil {
		os.Remove(exportFilename)
		fatalf("Error exporting keys: %v", err)
	}
}

// keysExportRoot exports a root key by ID to a PEM file
func keysExportRoot(cmd *cobra.Command, args []string) {
	if len(args) < 2 {
		cmd.Usage()
		fatalf("Must specify key ID and output filename for export")
	}

	parseConfig()

	keyID := args[0]
	exportFilename := args[1]

	if len(keyID) != idSize {
		fatalf("Please specify a valid root key ID")
	}

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, false)...,
	)

	exportFile, err := os.Create(exportFilename)
	if err != nil {
		fatalf("Error creating output file: %v", err)
	}
	if keysExportRootChangePassphrase {
		// Must use a different passphrase retriever to avoid caching the
		// unlocking passphrase and reusing that.
		exportRetriever := getRetriever()
		err = cs.ExportRootKeyReencrypt(exportFile, keyID, exportRetriever)
	} else {
		err = cs.ExportRootKey(exportFile, keyID)
	}
	exportFile.Close()
	if err != nil {
		os.Remove(exportFilename)
		fatalf("Error exporting root key: %v", err)
	}
}

// keysImport imports keys from a ZIP file
func keysImport(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cmd.Usage()
		fatalf("Must specify input filename for import")
	}

	importFilename := args[0]

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	zipReader, err := zip.OpenReader(importFilename)
	if err != nil {
		fatalf("Opening file for import: %v", err)
	}
	defer zipReader.Close()

	err = cs.ImportKeysZip(zipReader.Reader)

	if err != nil {
		fatalf("Error importing keys: %v", err)
	}
}

// keysImportRoot imports a root key from a PEM file
func keysImportRoot(cmd *cobra.Command, args []string) {
	if len(args) != 1 {
		cmd.Usage()
		fatalf("Must specify input filename for import")
	}

	parseConfig()

	cs := cryptoservice.NewCryptoService(
		"",
		getKeyStores(cmd, mainViper.GetString("trust_dir"), retriever, true)...,
	)

	importFilename := args[0]

	importFile, err := os.Open(importFilename)
	if err != nil {
		fatalf("Opening file for import: %v", err)
	}
	defer importFile.Close()

	err = cs.ImportRootKey(importFile)

	if err != nil {
		fatalf("Error importing root key: %v", err)
	}
}

func printKey(cmd *cobra.Command, keyPath, alias, loc string) {
	keyID := filepath.Base(keyPath)
	gun := filepath.Dir(keyPath)
	cmd.Printf("%s - %s - %s - %s\n", gun, alias, keyID, loc)
}

func keysRotate(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cmd.Usage()
		fatalf("Must specify a GUN and target")
	}
	parseConfig()

	gun := args[0]
	nRepo, err := notaryclient.NewNotaryRepository(mainViper.GetString("trust_dir"), gun, remoteTrustServer, nil, retriever)
	if err != nil {
		fatalf(err.Error())
	}
	if err := nRepo.RotateKeys(); err != nil {
		fatalf(err.Error())
	}
}

func getKeyStores(cmd *cobra.Command, directory string,
	ret passphrase.Retriever, withHardware bool) []trustmanager.KeyStore {

	fileKeyStore, err := trustmanager.NewKeyFileStore(directory, ret)
	if err != nil {
		fatalf("Failed to create private key store in directory: %s", directory)
	}

	ks := []trustmanager.KeyStore{fileKeyStore}

	if withHardware {
		yubiStore, err := getYubiKeyStore(fileKeyStore, ret)
		if err == nil && yubiStore != nil {
			// Note that the order is important, since we want to prioritize
			// the yubikey store
			ks = []trustmanager.KeyStore{yubiStore, fileKeyStore}
		}
	}

	return ks
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"strings"

	"github.com/nboughton/go-dice"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/vg"
)

var defaultprec = 1000000

func main() {
	d := flag.String("d", "2d6", "Dice set to test. Can be a single value (2d10) or multiple values delineated by commas (2d4,3d10...)")
	p := flag.String("p", "high", "Set prec (high, medium, low). Higher prec performs more tests and thus takes longer")
	flag.Parse()

	prec := defaultprec
	switch *p {
	case "medium":
		prec = 100000
	case "low":
		prec = 10000
	}

	bag, err := dice.NewBag(strings.Split(*d, ",")...)
	if err != nil {
		log.Fatal(err)
	}

	// New plot
	pl, err := plot.New()
	if err != nil {
		log.Fatal(err)
	}
	pl.Title.Text = "Probability Distribution For " + *d
	pl.X.Label.Text = "Rolled"
	pl.X.Min = float64(bag.Min())
	pl.X.Max = float64(bag.Max())

	pl.Y.Label.Text = "Probability (%)"
	pl.Y.Min = 0
	//pl.Y.Max =

	pl.X.Tick.Marker = customTicks{}
	pl.Y.Tick.Marker = customTicks{}
	pl.Add(plotter.NewGrid())

	// Generate plot data
	l, err := plotter.NewLine(lineData(bag, prec))
	if err != nil {
		log.Fatal(err)
	}
	l.LineStyle.Width = vg.Points(1)

	// Add plot data
	pl.Add(l)

	// Save to png
	if err := pl.Save(15*vg.Centimeter, 15*vg.Centimeter, fmt.Sprintf("%s.png", *d)); err != nil {
		log.Fatal(err)
	}
}

func lineData(bag *dice.Bag, prec int) plotter.XYs {
	var (
		min  = bag.Min()
		max  = bag.Max()
		xLen = max - min + 1
		pts  = make(plotter.XYs, xLen)
		x    = make([]float64, xLen)
	)

	for i := 0; i < prec; i++ {
		t, _ := bag.Roll()
		x[t-min]++
	}

	for i := range pts {
		pts[i].X = float64(i + min)
		pts[i].Y = x[i] / float64(prec) * 100
	}

	return pts
}

type customTicks struct{}

func (customTicks) Ticks(min, max float64) []plot.Tick {
	var tks []plot.Tick

	for i := 0.; i < max; i++ {
		t := plot.Tick{Value: float64(i + 1)}

		switch {
		case max > 20 && max <= 50:
			t.Label = label(i, 2)
		case max > 50 && max <= 100:
			t.Label = label(i, 5)
		case max > 100:
			t.Label = label(i, int(max/4))
		default:
			t.Label = label(i, 1)
		}

		tks = append(tks, t)
	}

	return tks
}

func label(i float64, mod int) string {
	if int(i+1)%mod == 0 || i == 1 {
		return fmt.Sprintf("%d", int(i+1))
	}

	return ""
}
<commit_msg>label and image size tweak<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"strings"

	"github.com/nboughton/go-dice"
	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/vg"
)

var defaultprec = 1000000

func main() {
	d := flag.String("d", "2d6", "Dice set to test. Can be a single value (2d10) or multiple values delineated by commas (2d4,3d10...)")
	p := flag.String("p", "high", "Set prec (high, medium, low). Higher prec performs more tests and thus takes longer")
	flag.Parse()

	prec := defaultprec
	switch *p {
	case "medium":
		prec = 100000
	case "low":
		prec = 10000
	}

	bag, err := dice.NewBag(strings.Split(*d, ",")...)
	if err != nil {
		log.Fatal(err)
	}

	// New plot
	pl, err := plot.New()
	if err != nil {
		log.Fatal(err)
	}
	pl.Title.Text = "Distribution For " + bag.String()
	pl.X.Label.Text = "Rolled"
	pl.X.Min = float64(bag.Min())
	pl.X.Max = float64(bag.Max())

	pl.Y.Label.Text = "Probability (%)"
	pl.Y.Min = 0
	//pl.Y.Max =

	pl.X.Tick.Marker = customTicks{}
	pl.Y.Tick.Marker = customTicks{}
	pl.Add(plotter.NewGrid())

	// Generate plot data
	l, err := plotter.NewLine(lineData(bag, prec))
	if err != nil {
		log.Fatal(err)
	}
	l.LineStyle.Width = vg.Points(1)

	// Add plot data
	pl.Add(l)

	// Save to png
	if err := pl.Save(20*vg.Centimeter, 15*vg.Centimeter, fmt.Sprintf("%s.png", *d)); err != nil {
		log.Fatal(err)
	}
}

func lineData(bag *dice.Bag, prec int) plotter.XYs {
	var (
		min  = bag.Min()
		max  = bag.Max()
		xLen = max - min + 1
		pts  = make(plotter.XYs, xLen)
		x    = make([]float64, xLen)
	)

	for i := 0; i < prec; i++ {
		t, _ := bag.Roll()
		x[t-min]++
	}

	for i := range pts {
		pts[i].X = float64(i + min)
		pts[i].Y = x[i] / float64(prec) * 100
	}

	return pts
}

type customTicks struct{}

func (customTicks) Ticks(min, max float64) []plot.Tick {
	var tks []plot.Tick

	for i := 0.; i < max; i++ {
		t := plot.Tick{Value: float64(i + 1)}

		switch {
		case max > 20 && max <= 50:
			t.Label = label(i, 2)
		case max > 50 && max <= 100:
			t.Label = label(i, 5)
		case max > 100:
			t.Label = label(i, int(max/4))
		default:
			t.Label = label(i, 1)
		}

		tks = append(tks, t)
	}

	return tks
}

func label(i float64, mod int) string {
	if int(i+1)%mod == 0 {
		return fmt.Sprintf("%d", int(i+1))
	}

	return ""
}
<|endoftext|>