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

// The frontend runs a service to serve user-facing traffic.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/integrations/ocsql"
	"github.com/go-redis/redis/v7"
	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/cmd/internal/cmdconfig"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/dcensus"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/frontend"
	"golang.org/x/pkgsite/internal/localdatasource"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/middleware"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/proxydatasource"
	"golang.org/x/pkgsite/internal/queue"
	"golang.org/x/pkgsite/internal/source"
)

var (
	queueName      = config.GetEnv("GO_DISCOVERY_FRONTEND_TASK_QUEUE", "")
	workers        = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	_              = flag.String("static", "content/static", "path to folder containing static files served")
	thirdPartyPath = flag.String("third_party", "third_party", "path to folder containing third-party libraries")
	devMode        = flag.Bool("dev", false, "enable developer mode (reload templates on each page load, serve non-minified JS/CSS, etc.)")
	disableCSP     = flag.Bool("nocsp", false, "disable Content Security Policy")
	proxyURL       = flag.String("proxy_url", "https://proxy.golang.org", "Uses the module proxy referred to by this URL "+
		"for direct proxy mode and frontend fetches")
	directProxy = flag.Bool("direct_proxy", false, "if set to true, uses the module proxy referred to by this URL "+
		"as a direct backend, bypassing the database")
	localPaths         = flag.String("local", "", "run locally, accepts a GOPATH-like collection of local paths for modules to load to memory")
	gopathMode         = flag.Bool("gopath_mode", false, "assume that local modules' paths are relative to GOPATH/src, used only with -local")
	bypassLicenseCheck = flag.Bool("bypass_license_check", false, "display all information, even for non-redistributable paths")
)

func main() {
	flag.Parse()
	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)
	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	log.SetLevel(cfg.LogLevel)

	var (
		dsg        func(context.Context) internal.DataSource
		fetchQueue queue.Queue
	)
	if *bypassLicenseCheck {
		log.Info(ctx, "BYPASSING LICENSE CHECKING: DISPLAYING NON-REDISTRIBUTABLE INFORMATION")
	}

	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.ExperimentGetter")
	expg := cmdconfig.ExperimentGetter(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.ExperimentGetter")

	if *localPaths != "" {
		lds := localdatasource.New()
		paths := filepath.SplitList(*localPaths)
		if *gopathMode {
			for _, path := range paths {
				err := lds.LoadFromGOPATH(ctx, path)
				if err != nil {
					log.Error(ctx, err)
				}
			}
		} else {
			for _, path := range paths {
				err := lds.Load(ctx, path)
				if err != nil {
					log.Error(ctx, err)
				}
			}
		}
		dsg = func(context.Context) internal.DataSource { return lds }
	} else {
		proxyClient, err := proxy.New(*proxyURL)
		if err != nil {
			log.Fatal(ctx, err)
		}

		if *directProxy {
			var pds *proxydatasource.DataSource
			if *bypassLicenseCheck {
				pds = proxydatasource.NewBypassingLicenseCheck(proxyClient)
			} else {
				pds = proxydatasource.New(proxyClient)
			}
			dsg = func(context.Context) internal.DataSource { return pds }
		} else {
			// Wrap the postgres driver with OpenCensus instrumentation.
			ocDriver, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
			if err != nil {
				log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
			}

			log.Infof(ctx, "cmd/frontend: openDB start")
			ddb, err := openDB(ctx, cfg, ocDriver)
			if err != nil {
				log.Fatal(ctx, err)
			}
			log.Infof(ctx, "cmd/frontend: openDB finished")

			var db *postgres.DB
			if *bypassLicenseCheck {
				db = postgres.NewBypassingLicenseCheck(ddb)
			} else {
				db = postgres.New(ddb)
			}
			defer db.Close()
			dsg = func(context.Context) internal.DataSource { return db }
			sourceClient := source.NewClient(config.SourceTimeout)
			// The closure passed to queue.New is only used for testing and local
			// execution, not in production. So it's okay that it doesn't use a
			// per-request connection.
			fetchQueue, err = queue.New(ctx, cfg, queueName, *workers, expg,
				func(ctx context.Context, modulePath, version string) (int, error) {
					return frontend.FetchAndUpdateState(ctx, modulePath, version, proxyClient, sourceClient, db)
				})
			if err != nil {
				log.Fatalf(ctx, "queue.New: %v", err)
			}
		}
	}

	var haClient *redis.Client
	if cfg.RedisHAHost != "" {
		haClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisHAHost + ":" + cfg.RedisHAPort,
		})
	}
	server, err := frontend.NewServer(frontend.ServerConfig{
		DataSourceGetter:     dsg,
		Queue:                fetchQueue,
		CompletionClient:     haClient,
		TaskIDChangeInterval: config.TaskIDChangeIntervalFrontend,
		StaticPath:           template.TrustedSourceFromFlag(flag.Lookup("static").Value),
		ThirdPartyPath:       *thirdPartyPath,
		DevMode:              *devMode,
		AppVersionLabel:      cfg.AppVersionLabel(),
		GoogleTagManagerID:   cfg.GoogleTagManagerID,
		ServeStats:           cfg.ServeStats,
	})
	if err != nil {
		log.Fatalf(ctx, "frontend.NewServer: %v", err)
	}
	router := dcensus.NewRouter(frontend.TagRoute)
	var cacheClient *redis.Client
	if cfg.RedisCacheHost != "" {
		cacheClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisCacheHost + ":" + cfg.RedisCachePort,
		})
	}
	server.Install(router.Handle, cacheClient, cfg.AuthValues)
	views := append(dcensus.ServerViews,
		postgres.SearchLatencyDistribution,
		postgres.SearchResponseCount,
		frontend.FetchLatencyDistribution,
		frontend.FetchResponseCount,
		frontend.PlaygroundShareRequestCount,
		frontend.VersionTypeCount,
		middleware.CacheResultCount,
		middleware.CacheErrorCount,
		middleware.CacheLatency,
		middleware.QuotaResultCount,
	)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8081"), dcensusServer)
	}
	panicHandler, err := server.PanicHandler()
	if err != nil {
		log.Fatal(ctx, err)
	}
	rc := cmdconfig.ReportingClient(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.Experimenter")
	experimenter := cmdconfig.Experimenter(ctx, cfg, expg, rc)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.Experimenter")

	ermw := middleware.Identity()
	if rc != nil {
		ermw = middleware.ErrorReporting(rc.Report)
	}
	mw := middleware.Chain(
		middleware.RequestLog(cmdconfig.Logger(ctx, cfg, "frontend-log")),
		middleware.AcceptRequests(http.MethodGet, http.MethodPost), // accept only GETs and POSTs
		middleware.Quota(cfg.Quota),
		middleware.GodocURL(),                                                                 // potentially redirects so should be early in chain
		middleware.SecureHeaders(!*disableCSP),                                                // must come before any caching for nonces to work
		middleware.LatestVersions(server.GetLatestMinorVersion, server.GetLatestMajorVersion), // must come before caching for version badge to work
		middleware.Panic(panicHandler),
		ermw,
		middleware.Timeout(54*time.Second),
		middleware.Experiment(experimenter),
	)
	addr := cfg.HostAddr("localhost:8080")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, mw(router)))
}

// TODO(https://github.com/golang/go/issues/40097): factor out to reduce
// duplication with cmd/worker/main.go.

// openDB opens a connection to a database with the given driver, using connection info from
// the given config.
// It first tries the main connection info (DBConnInfo), and if that fails, it uses backup
// connection info it if exists (DBSecondaryConnInfo).
func openDB(ctx context.Context, cfg *config.Config, driver string) (_ *database.DB, err error) {
	defer derrors.Wrap(&err, "openDB(ctx, cfg, %q)", driver)
	log.Infof(ctx, "opening database on host %s", cfg.DBHost)
	ddb, err := database.Open(driver, cfg.DBConnInfo(), cfg.InstanceID)
	if err == nil {
		log.Infof(ctx, "connected to primary host: %s", cfg.DBHost)
		return ddb, nil
	}

	ci := cfg.DBSecondaryConnInfo()
	if ci == "" {
		log.Infof(ctx, "no secondary DB host")
		return nil, err
	}
	log.Errorf(ctx, "database.Open for primary host %s failed with %v; trying secondary host %s ",
		cfg.DBHost, err, cfg.DBSecondaryHost)
	db, err := database.Open(driver, ci, cfg.InstanceID)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "connected to secondary host %s", cfg.DBSecondaryHost)
	return db, nil
}
<commit_msg>cmd/frontend: delay loading of local modules<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The frontend runs a service to serve user-facing traffic.
package main

import (
	"context"
	"flag"
	"net/http"
	"os"
	"path/filepath"
	"time"

	"cloud.google.com/go/profiler"
	"contrib.go.opencensus.io/integrations/ocsql"
	"github.com/go-redis/redis/v7"
	"github.com/google/safehtml/template"
	"golang.org/x/pkgsite/cmd/internal/cmdconfig"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/config"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/dcensus"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/frontend"
	"golang.org/x/pkgsite/internal/localdatasource"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/middleware"
	"golang.org/x/pkgsite/internal/postgres"
	"golang.org/x/pkgsite/internal/proxy"
	"golang.org/x/pkgsite/internal/proxydatasource"
	"golang.org/x/pkgsite/internal/queue"
	"golang.org/x/pkgsite/internal/source"
)

var (
	queueName      = config.GetEnv("GO_DISCOVERY_FRONTEND_TASK_QUEUE", "")
	workers        = flag.Int("workers", 10, "number of concurrent requests to the fetch service, when running locally")
	_              = flag.String("static", "content/static", "path to folder containing static files served")
	thirdPartyPath = flag.String("third_party", "third_party", "path to folder containing third-party libraries")
	devMode        = flag.Bool("dev", false, "enable developer mode (reload templates on each page load, serve non-minified JS/CSS, etc.)")
	disableCSP     = flag.Bool("nocsp", false, "disable Content Security Policy")
	proxyURL       = flag.String("proxy_url", "https://proxy.golang.org", "Uses the module proxy referred to by this URL "+
		"for direct proxy mode and frontend fetches")
	directProxy = flag.Bool("direct_proxy", false, "if set to true, uses the module proxy referred to by this URL "+
		"as a direct backend, bypassing the database")
	localPaths         = flag.String("local", "", "run locally, accepts a GOPATH-like collection of local paths for modules to load to memory")
	gopathMode         = flag.Bool("gopath_mode", false, "assume that local modules' paths are relative to GOPATH/src, used only with -local")
	bypassLicenseCheck = flag.Bool("bypass_license_check", false, "display all information, even for non-redistributable paths")
)

func main() {
	flag.Parse()
	ctx := context.Background()
	cfg, err := config.Init(ctx)
	if err != nil {
		log.Fatal(ctx, err)
	}
	cfg.Dump(os.Stderr)
	if cfg.UseProfiler {
		if err := profiler.Start(profiler.Config{}); err != nil {
			log.Fatalf(ctx, "profiler.Start: %v", err)
		}
	}

	log.SetLevel(cfg.LogLevel)

	var (
		dsg        func(context.Context) internal.DataSource
		fetchQueue queue.Queue
	)
	if *bypassLicenseCheck {
		log.Info(ctx, "BYPASSING LICENSE CHECKING: DISPLAYING NON-REDISTRIBUTABLE INFORMATION")
	}

	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.ExperimentGetter")
	expg := cmdconfig.ExperimentGetter(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.ExperimentGetter")

	if *localPaths != "" {
		lds := localdatasource.New()
		dsg = func(context.Context) internal.DataSource { return lds }
	} else {
		proxyClient, err := proxy.New(*proxyURL)
		if err != nil {
			log.Fatal(ctx, err)
		}

		if *directProxy {
			var pds *proxydatasource.DataSource
			if *bypassLicenseCheck {
				pds = proxydatasource.NewBypassingLicenseCheck(proxyClient)
			} else {
				pds = proxydatasource.New(proxyClient)
			}
			dsg = func(context.Context) internal.DataSource { return pds }
		} else {
			// Wrap the postgres driver with OpenCensus instrumentation.
			ocDriver, err := ocsql.Register("postgres", ocsql.WithAllTraceOptions())
			if err != nil {
				log.Fatalf(ctx, "unable to register the ocsql driver: %v\n", err)
			}

			log.Infof(ctx, "cmd/frontend: openDB start")
			ddb, err := openDB(ctx, cfg, ocDriver)
			if err != nil {
				log.Fatal(ctx, err)
			}
			log.Infof(ctx, "cmd/frontend: openDB finished")

			var db *postgres.DB
			if *bypassLicenseCheck {
				db = postgres.NewBypassingLicenseCheck(ddb)
			} else {
				db = postgres.New(ddb)
			}
			defer db.Close()
			dsg = func(context.Context) internal.DataSource { return db }
			sourceClient := source.NewClient(config.SourceTimeout)
			// The closure passed to queue.New is only used for testing and local
			// execution, not in production. So it's okay that it doesn't use a
			// per-request connection.
			fetchQueue, err = queue.New(ctx, cfg, queueName, *workers, expg,
				func(ctx context.Context, modulePath, version string) (int, error) {
					return frontend.FetchAndUpdateState(ctx, modulePath, version, proxyClient, sourceClient, db)
				})
			if err != nil {
				log.Fatalf(ctx, "queue.New: %v", err)
			}
		}
	}

	var haClient *redis.Client
	if cfg.RedisHAHost != "" {
		haClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisHAHost + ":" + cfg.RedisHAPort,
		})
	}
	server, err := frontend.NewServer(frontend.ServerConfig{
		DataSourceGetter:     dsg,
		Queue:                fetchQueue,
		CompletionClient:     haClient,
		TaskIDChangeInterval: config.TaskIDChangeIntervalFrontend,
		StaticPath:           template.TrustedSourceFromFlag(flag.Lookup("static").Value),
		ThirdPartyPath:       *thirdPartyPath,
		DevMode:              *devMode,
		AppVersionLabel:      cfg.AppVersionLabel(),
		GoogleTagManagerID:   cfg.GoogleTagManagerID,
		ServeStats:           cfg.ServeStats,
	})
	if err != nil {
		log.Fatalf(ctx, "frontend.NewServer: %v", err)
	}

	if *localPaths != "" {
		lds, ok := dsg(ctx).(*localdatasource.DataSource)
		if ok {
			load(ctx, lds, *localPaths)
		}
	}

	router := dcensus.NewRouter(frontend.TagRoute)
	var cacheClient *redis.Client
	if cfg.RedisCacheHost != "" {
		cacheClient = redis.NewClient(&redis.Options{
			Addr: cfg.RedisCacheHost + ":" + cfg.RedisCachePort,
		})
	}
	server.Install(router.Handle, cacheClient, cfg.AuthValues)
	views := append(dcensus.ServerViews,
		postgres.SearchLatencyDistribution,
		postgres.SearchResponseCount,
		frontend.FetchLatencyDistribution,
		frontend.FetchResponseCount,
		frontend.PlaygroundShareRequestCount,
		frontend.VersionTypeCount,
		middleware.CacheResultCount,
		middleware.CacheErrorCount,
		middleware.CacheLatency,
		middleware.QuotaResultCount,
	)
	if err := dcensus.Init(cfg, views...); err != nil {
		log.Fatal(ctx, err)
	}
	// We are not currently forwarding any ports on AppEngine, so serving debug
	// information is broken.
	if !cfg.OnAppEngine() {
		dcensusServer, err := dcensus.NewServer()
		if err != nil {
			log.Fatal(ctx, err)
		}
		go http.ListenAndServe(cfg.DebugAddr("localhost:8081"), dcensusServer)
	}
	panicHandler, err := server.PanicHandler()
	if err != nil {
		log.Fatal(ctx, err)
	}
	rc := cmdconfig.ReportingClient(ctx, cfg)
	log.Infof(ctx, "cmd/frontend: initializing cmdconfig.Experimenter")
	experimenter := cmdconfig.Experimenter(ctx, cfg, expg, rc)
	log.Infof(ctx, "cmd/frontend: initialized cmdconfig.Experimenter")

	ermw := middleware.Identity()
	if rc != nil {
		ermw = middleware.ErrorReporting(rc.Report)
	}
	mw := middleware.Chain(
		middleware.RequestLog(cmdconfig.Logger(ctx, cfg, "frontend-log")),
		middleware.AcceptRequests(http.MethodGet, http.MethodPost), // accept only GETs and POSTs
		middleware.Quota(cfg.Quota),
		middleware.GodocURL(),                                                                 // potentially redirects so should be early in chain
		middleware.SecureHeaders(!*disableCSP),                                                // must come before any caching for nonces to work
		middleware.LatestVersions(server.GetLatestMinorVersion, server.GetLatestMajorVersion), // must come before caching for version badge to work
		middleware.Panic(panicHandler),
		ermw,
		middleware.Timeout(54*time.Second),
		middleware.Experiment(experimenter),
	)
	addr := cfg.HostAddr("localhost:8080")
	log.Infof(ctx, "Listening on addr %s", addr)
	log.Fatal(ctx, http.ListenAndServe(addr, mw(router)))
}

// TODO(https://github.com/golang/go/issues/40097): factor out to reduce
// duplication with cmd/worker/main.go.

// openDB opens a connection to a database with the given driver, using connection info from
// the given config.
// It first tries the main connection info (DBConnInfo), and if that fails, it uses backup
// connection info it if exists (DBSecondaryConnInfo).
func openDB(ctx context.Context, cfg *config.Config, driver string) (_ *database.DB, err error) {
	defer derrors.Wrap(&err, "openDB(ctx, cfg, %q)", driver)
	log.Infof(ctx, "opening database on host %s", cfg.DBHost)
	ddb, err := database.Open(driver, cfg.DBConnInfo(), cfg.InstanceID)
	if err == nil {
		log.Infof(ctx, "connected to primary host: %s", cfg.DBHost)
		return ddb, nil
	}

	ci := cfg.DBSecondaryConnInfo()
	if ci == "" {
		log.Infof(ctx, "no secondary DB host")
		return nil, err
	}
	log.Errorf(ctx, "database.Open for primary host %s failed with %v; trying secondary host %s ",
		cfg.DBHost, err, cfg.DBSecondaryHost)
	db, err := database.Open(driver, ci, cfg.InstanceID)
	if err != nil {
		return nil, err
	}
	log.Infof(ctx, "connected to secondary host %s", cfg.DBSecondaryHost)
	return db, nil
}

// load loads local modules from pathList.
func load(ctx context.Context, ds *localdatasource.DataSource, pathList string) {
	paths := filepath.SplitList(pathList)
	for _, path := range paths {
		var err error
		if *gopathMode {
			err = ds.LoadFromGOPATH(ctx, path)
		} else {
			err = ds.Load(ctx, path)
		}
		if err != nil {
			log.Error(ctx, err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sort"
	"strconv"
	"text/tabwriter"

	_ "github.com/mxk/go-sqlite/sqlite3"
	"github.com/rwcarlsen/cyan/query"
)

var (
	help     = flag.Bool("h", false, "print this help message")
	custom   = flag.String("custom", "", "path to custom sql query spec file")
	dbname   = flag.String("db", "", "cyclus sqlite database to query")
	simidstr = flag.String("simid", "", "simulation id in hex (empty string defaults to first sim id in database")
)

var simid []byte

var command string

var db *sql.DB

var cmds = NewCmdSet()

// map[cmdname]sqltext
var customSql = map[string]string{}

func init() {
	cmds.Register("agents", "list all agents in the simulation", doAgents)
	cmds.Register("sims", "list all simulations in the database", doSims)
	cmds.Register("inv", "show inventory of one or more agents at a specific timestep", doInv)
	cmds.Register("created", "show material created by one or more agents between specific timesteps", doCreated)
	cmds.Register("deployseries", "print a time-series of a prototype's total active deployments", doDeploySeries)
	cmds.Register("flow", "Show total transacted material between two groups of agents between specific timesteps", doFlow)
	cmds.Register("invseries", "print a time series of an agent's inventory for specified isotopes", doInvSeries)
	cmds.Register("flowgraph", "print a graphviz dot graph of resource arcs between facilities", doFlowGraph)
	cmds.Register("energy", "print thermal energy (J) generated by the simulation between 2 timesteps", doEnergy)
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	if *help || flag.NArg() < 1 {
		fmt.Println("Usage: metric -db <cyclus-db> [opts] <command> [args...]")
		fmt.Println("Calculates metrics for cyclus simulation data in a sqlite database.")
		flag.PrintDefaults()
		fmt.Println("\nCommands:")
		for i := range cmds.Names {
			fmt.Printf("    %v: %v\n", cmds.Names[i], cmds.Helps[i])
		}
		return
	} else if *dbname == "" {
		log.Fatal("must specify database with db flag")
	}

	if *custom != "" {
		data, err := ioutil.ReadFile(*custom)
		fatalif(err)
		fatalif(json.Unmarshal(data, &customSql))
	}

	var err error
	db, err = sql.Open("sqlite3", *dbname)
	fatalif(err)
	defer db.Close()

	if *simidstr == "" {
		ids, err := query.SimIds(db)
		fatalif(err)
		simid = ids[0]
	} else {
		simid, err = hex.DecodeString(*simidstr)
		fatalif(err)
	}

	cmds.Execute(flag.Args())
}

func doCustom(cmd string, oargs []string) {
	s, ok := customSql[cmd]
	if !ok {
		log.Fatalf("Invalid command %v", cmd)
	}
	args := make([]interface{}, len(oargs))
	for i := range args {
		args[i] = oargs[i]
	}
	rows, err := db.Query(s, args...)
	fatalif(err)

	tw := tabwriter.NewWriter(os.Stdout, 4, 4, 1, ' ', 0)
	cols, err := rows.Columns()
	fatalif(err)
	for _, c := range cols {
		_, err := tw.Write([]byte(c + "\t"))
		fatalif(err)
	}
	_, err = tw.Write([]byte("\n"))
	fatalif(err)

	for rows.Next() {
		vs := make([]interface{}, len(cols))
		vals := make([]string, len(cols))
		for i := range vals {
			vs[i] = &vals[i]
		}
		err := rows.Scan(vs...)
		fatalif(err)

		for _, v := range vals {
			_, err := tw.Write([]byte(v + "\t"))
			fatalif(err)
		}
		_, err = tw.Write([]byte("\n"))
		fatalif(err)
	}
	fatalif(rows.Err())
	fatalif(tw.Flush())
	return
}

func doSims(cmd string, args []string) {
	ids, err := query.SimIds(db)
	fatalif(err)
	for _, id := range ids {
		info, err := query.SimStat(db, id)
		fatalif(err)
		fmt.Println(info)
	}
}

func doAgents(cmd string, args []string) {
	fs := flag.NewFlagSet("agents", flag.ExitOnError)
	proto := fs.String("proto", "", "filter by prototype (default \"\" is all prototypes)")
	fs.Usage = func() {
		log.Print("Usage: agents")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	ags, err := query.AllAgents(db, simid, *proto)
	fatalif(err)
	for _, a := range ags {
		fmt.Println(a)
	}
}

func doInv(cmd string, args []string) {
	fs := flag.NewFlagSet("inv", flag.ExitOnError)
	t := fs.Int("t", -1, "timestep of inventory (-1 = end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: inv [agent-id...]\nZero agents uses all agent inventories")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var agents []int
	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.InvAt(db, simid, *t, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

type Row struct {
	X  int
	Ys []float64
}

type MultiSeries [][]query.XY

func (ms MultiSeries) Rows() []Row {
	rowmap := map[int]Row{}
	xs := []int{}
	for i, s := range ms {
		for _, xy := range s {
			row, ok := rowmap[xy.X]
			if !ok {
				xs = append(xs, xy.X)
				row.Ys = make([]float64, len(ms))
				row.X = xy.X
			}
			row.Ys[i] = xy.Y
			rowmap[xy.X] = row
		}
	}

	sort.Ints(xs)
	rows := make([]Row, 0, len(rowmap))
	for _, x := range xs {
		rows = append(rows, rowmap[x])
	}
	return rows
}

func doInvSeries(cmd string, args []string) {
	fs := flag.NewFlagSet("invseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: invseries <agent-id> <isotope> [isotope...]"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 2 {
		fs.Usage()
		return
	}

	agent, err := strconv.Atoi(fs.Arg(0))
	fatalif(err)

	isos := []int{}
	for _, arg := range fs.Args()[1:] {
		iso, err := strconv.Atoi(arg)
		fatalif(err)
		isos = append(isos, iso)
	}

	ms := MultiSeries{}
	for _, iso := range isos {
		xys, err := query.InvSeries(db, simid, agent, iso)
		ms = append(ms, xys)
		fatalif(err)
	}

	fmt.Printf("# Agent %v inventory in kg\n", agent)
	fmt.Printf("# [Timestep]")
	for _, iso := range isos {
		fmt.Printf(" [%v]", iso)
	}
	fmt.Printf("\n")
	for _, row := range ms.Rows() {
		fmt.Printf("%v", row.X)
		for _, y := range row.Ys {
			fmt.Printf(" %v ", y)
		}
		fmt.Printf("\n")
	}
}

func doFlowGraph(cmd string, args []string) {
	fs := flag.NewFlagSet("flowgraph", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: flowgraph"); fs.PrintDefaults() }
	proto := fs.Bool("proto", false, "aggregate nodes by prototype")
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	arcs, err := query.FlowGraph(db, simid, *t0, *t1, *proto)
	fatalif(err)

	fmt.Println("digraph ResourceFlows {")
	fmt.Println("    overlap = false;")
	fmt.Println("    nodesep=1.0;")
	fmt.Println("    edge [fontsize=9];")
	for _, arc := range arcs {
		fmt.Printf("    \"%v\" -> \"%v\" [label=\"%v\\n(%.3g kg)\"];\n", arc.Src, arc.Dst, arc.Commod, arc.Quantity)
	}
	fmt.Println("}")
}

func doDeploySeries(cmd string, args []string) {
	fs := flag.NewFlagSet("deployseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: deployseries <prototype>"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 1 {
		fs.Usage()
		return
	}

	proto := fs.Arg(0)
	xys, err := query.DeployCumulative(db, simid, proto)
	fatalif(err)

	fmt.Printf("# Prototype %v total active deployments\n", proto)
	fmt.Println("# [Timestep] [Count]")
	for _, xy := range xys {
		fmt.Printf("%v %v\n", xy.X, xy.Y)
	}
}

func doCreated(cmd string, args []string) {
	fs := flag.NewFlagSet("created", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: created [agent-id...]\nZero agents uses all agents"); fs.PrintDefaults() }
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	var agents []int

	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.MatCreated(db, simid, *t0, *t1, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func doFlow(cmd string, args []string) {
	fs := flag.NewFlagSet("flow", flag.ExitOnError)
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: flow <from-agents...> .. <to-agents...>")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var from []int
	var to []int

	if flag.NArg() < 3 {
		fs.Usage()
		return
	}

	before := true
	for _, arg := range args {
		if arg == ".." {
			before = false
			continue
		}

		id, err := strconv.Atoi(arg)
		fatalif(err)
		if before {
			from = append(from, id)
		} else {
			to = append(to, id)
		}
	}
	if len(to) < 1 {
		fs.Usage()
		return
	}

	m, err := query.Flow(db, simid, *t0, *t1, from, to)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func doEnergy(cmd string, args []string) {
	fs := flag.NewFlagSet("energy", flag.ExitOnError)
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: energy")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	e, err := query.EnergyProduced(db, simid, *t0, *t1)
	fatalif(err)
	fmt.Println(e)
}

func fatalif(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

type CmdSet struct {
	funcs map[string]func(string, []string) // map[cmdname]func(cmdname, args)
	Names []string
	Helps []string
}

func NewCmdSet() *CmdSet {
	return &CmdSet{funcs: map[string]func(string, []string){}}
}

func (cs *CmdSet) Register(name, brief string, f func(string, []string)) {
	cs.Names = append(cs.Names, name)
	cs.Helps = append(cs.Helps, brief)
	cs.funcs[name] = f
}

func (cs *CmdSet) Execute(args []string) {
	cmd := args[0]
	f, ok := cs.funcs[cmd]
	if !ok {
		doCustom(cmd, args[1:])
		return
	}
	f(cmd, args[1:])
}
<commit_msg>auto post process on metric cmd if needed<commit_after>package main

import (
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sort"
	"strconv"
	"text/tabwriter"

	_ "github.com/mxk/go-sqlite/sqlite3"
	"github.com/rwcarlsen/cyan/post"
	"github.com/rwcarlsen/cyan/query"
)

var (
	help     = flag.Bool("h", false, "print this help message")
	custom   = flag.String("custom", "", "path to custom sql query spec file")
	dbname   = flag.String("db", "", "cyclus sqlite database to query")
	simidstr = flag.String("simid", "", "simulation id in hex (empty string defaults to first sim id in database")
)

var simid []byte

var command string

var db *sql.DB

var cmds = NewCmdSet()

// map[cmdname]sqltext
var customSql = map[string]string{}

func init() {
	cmds.Register("agents", "list all agents in the simulation", doAgents)
	cmds.Register("sims", "list all simulations in the database", doSims)
	cmds.Register("inv", "show inventory of one or more agents at a specific timestep", doInv)
	cmds.Register("created", "show material created by one or more agents between specific timesteps", doCreated)
	cmds.Register("deployseries", "print a time-series of a prototype's total active deployments", doDeploySeries)
	cmds.Register("flow", "Show total transacted material between two groups of agents between specific timesteps", doFlow)
	cmds.Register("invseries", "print a time series of an agent's inventory for specified isotopes", doInvSeries)
	cmds.Register("flowgraph", "print a graphviz dot graph of resource arcs between facilities", doFlowGraph)
	cmds.Register("energy", "print thermal energy (J) generated by the simulation between 2 timesteps", doEnergy)
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	if *help || flag.NArg() < 1 {
		fmt.Println("Usage: metric -db <cyclus-db> [opts] <command> [args...]")
		fmt.Println("Calculates metrics for cyclus simulation data in a sqlite database.")
		flag.PrintDefaults()
		fmt.Println("\nCommands:")
		for i := range cmds.Names {
			fmt.Printf("    %v: %v\n", cmds.Names[i], cmds.Helps[i])
		}
		return
	} else if *dbname == "" {
		log.Fatal("must specify database with db flag")
	}

	if *custom != "" {
		data, err := ioutil.ReadFile(*custom)
		fatalif(err)
		fatalif(json.Unmarshal(data, &customSql))
	}

	var err error
	db, err = sql.Open("sqlite3", *dbname)
	fatalif(err)
	defer db.Close()

	if *simidstr == "" {
		ids, err := query.SimIds(db)
		fatalif(err)
		simid = ids[0]
	} else {
		simid, err = hex.DecodeString(*simidstr)
		fatalif(err)
	}

	// post process if necessary
	fatalif(post.Prepare(db))
	simids, err := post.GetSimIds(db)
	fatalif(err)
	for _, simid := range simids {
		ctx := post.NewContext(db, simid)
		err := ctx.WalkAll()
		fatalif(err)
	}
	fatalif(post.Finish(db))

	// run command
	cmds.Execute(flag.Args())
}

func doCustom(cmd string, oargs []string) {
	s, ok := customSql[cmd]
	if !ok {
		log.Fatalf("Invalid command %v", cmd)
	}
	args := make([]interface{}, len(oargs))
	for i := range args {
		args[i] = oargs[i]
	}
	rows, err := db.Query(s, args...)
	fatalif(err)

	tw := tabwriter.NewWriter(os.Stdout, 4, 4, 1, ' ', 0)
	cols, err := rows.Columns()
	fatalif(err)
	for _, c := range cols {
		_, err := tw.Write([]byte(c + "\t"))
		fatalif(err)
	}
	_, err = tw.Write([]byte("\n"))
	fatalif(err)

	for rows.Next() {
		vs := make([]interface{}, len(cols))
		vals := make([]string, len(cols))
		for i := range vals {
			vs[i] = &vals[i]
		}
		err := rows.Scan(vs...)
		fatalif(err)

		for _, v := range vals {
			_, err := tw.Write([]byte(v + "\t"))
			fatalif(err)
		}
		_, err = tw.Write([]byte("\n"))
		fatalif(err)
	}
	fatalif(rows.Err())
	fatalif(tw.Flush())
	return
}

func doSims(cmd string, args []string) {
	ids, err := query.SimIds(db)
	fatalif(err)
	for _, id := range ids {
		info, err := query.SimStat(db, id)
		fatalif(err)
		fmt.Println(info)
	}
}

func doAgents(cmd string, args []string) {
	fs := flag.NewFlagSet("agents", flag.ExitOnError)
	proto := fs.String("proto", "", "filter by prototype (default \"\" is all prototypes)")
	fs.Usage = func() {
		log.Print("Usage: agents")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	ags, err := query.AllAgents(db, simid, *proto)
	fatalif(err)
	for _, a := range ags {
		fmt.Println(a)
	}
}

func doInv(cmd string, args []string) {
	fs := flag.NewFlagSet("inv", flag.ExitOnError)
	t := fs.Int("t", -1, "timestep of inventory (-1 = end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: inv [agent-id...]\nZero agents uses all agent inventories")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var agents []int
	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.InvAt(db, simid, *t, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

type Row struct {
	X  int
	Ys []float64
}

type MultiSeries [][]query.XY

func (ms MultiSeries) Rows() []Row {
	rowmap := map[int]Row{}
	xs := []int{}
	for i, s := range ms {
		for _, xy := range s {
			row, ok := rowmap[xy.X]
			if !ok {
				xs = append(xs, xy.X)
				row.Ys = make([]float64, len(ms))
				row.X = xy.X
			}
			row.Ys[i] = xy.Y
			rowmap[xy.X] = row
		}
	}

	sort.Ints(xs)
	rows := make([]Row, 0, len(rowmap))
	for _, x := range xs {
		rows = append(rows, rowmap[x])
	}
	return rows
}

func doInvSeries(cmd string, args []string) {
	fs := flag.NewFlagSet("invseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: invseries <agent-id> <isotope> [isotope...]"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 2 {
		fs.Usage()
		return
	}

	agent, err := strconv.Atoi(fs.Arg(0))
	fatalif(err)

	isos := []int{}
	for _, arg := range fs.Args()[1:] {
		iso, err := strconv.Atoi(arg)
		fatalif(err)
		isos = append(isos, iso)
	}

	ms := MultiSeries{}
	for _, iso := range isos {
		xys, err := query.InvSeries(db, simid, agent, iso)
		ms = append(ms, xys)
		fatalif(err)
	}

	fmt.Printf("# Agent %v inventory in kg\n", agent)
	fmt.Printf("# [Timestep]")
	for _, iso := range isos {
		fmt.Printf(" [%v]", iso)
	}
	fmt.Printf("\n")
	for _, row := range ms.Rows() {
		fmt.Printf("%v", row.X)
		for _, y := range row.Ys {
			fmt.Printf(" %v ", y)
		}
		fmt.Printf("\n")
	}
}

func doFlowGraph(cmd string, args []string) {
	fs := flag.NewFlagSet("flowgraph", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: flowgraph"); fs.PrintDefaults() }
	proto := fs.Bool("proto", false, "aggregate nodes by prototype")
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	arcs, err := query.FlowGraph(db, simid, *t0, *t1, *proto)
	fatalif(err)

	fmt.Println("digraph ResourceFlows {")
	fmt.Println("    overlap = false;")
	fmt.Println("    nodesep=1.0;")
	fmt.Println("    edge [fontsize=9];")
	for _, arc := range arcs {
		fmt.Printf("    \"%v\" -> \"%v\" [label=\"%v\\n(%.3g kg)\"];\n", arc.Src, arc.Dst, arc.Commod, arc.Quantity)
	}
	fmt.Println("}")
}

func doDeploySeries(cmd string, args []string) {
	fs := flag.NewFlagSet("deployseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: deployseries <prototype>"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 1 {
		fs.Usage()
		return
	}

	proto := fs.Arg(0)
	xys, err := query.DeployCumulative(db, simid, proto)
	fatalif(err)

	fmt.Printf("# Prototype %v total active deployments\n", proto)
	fmt.Println("# [Timestep] [Count]")
	for _, xy := range xys {
		fmt.Printf("%v %v\n", xy.X, xy.Y)
	}
}

func doCreated(cmd string, args []string) {
	fs := flag.NewFlagSet("created", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: created [agent-id...]\nZero agents uses all agents"); fs.PrintDefaults() }
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	var agents []int

	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.MatCreated(db, simid, *t0, *t1, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func doFlow(cmd string, args []string) {
	fs := flag.NewFlagSet("flow", flag.ExitOnError)
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: flow <from-agents...> .. <to-agents...>")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var from []int
	var to []int

	if flag.NArg() < 3 {
		fs.Usage()
		return
	}

	before := true
	for _, arg := range args {
		if arg == ".." {
			before = false
			continue
		}

		id, err := strconv.Atoi(arg)
		fatalif(err)
		if before {
			from = append(from, id)
		} else {
			to = append(to, id)
		}
	}
	if len(to) < 1 {
		fs.Usage()
		return
	}

	m, err := query.Flow(db, simid, *t0, *t1, from, to)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func doEnergy(cmd string, args []string) {
	fs := flag.NewFlagSet("energy", flag.ExitOnError)
	t0 := fs.Int("t1", 0, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: energy")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	e, err := query.EnergyProduced(db, simid, *t0, *t1)
	fatalif(err)
	fmt.Println(e)
}

func fatalif(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

type CmdSet struct {
	funcs map[string]func(string, []string) // map[cmdname]func(cmdname, args)
	Names []string
	Helps []string
}

func NewCmdSet() *CmdSet {
	return &CmdSet{funcs: map[string]func(string, []string){}}
}

func (cs *CmdSet) Register(name, brief string, f func(string, []string)) {
	cs.Names = append(cs.Names, name)
	cs.Helps = append(cs.Helps, brief)
	cs.funcs[name] = f
}

func (cs *CmdSet) Execute(args []string) {
	cmd := args[0]
	f, ok := cs.funcs[cmd]
	if !ok {
		doCustom(cmd, args[1:])
		return
	}
	f(cmd, args[1:])
}
<|endoftext|>


<commit_before>package main

import (
	"log"

	lua "github.com/yuin/gopher-lua"
	luar "layeh.com/gopher-luar"

	"github.com/zyedidia/micro/internal/action"
	"github.com/zyedidia/micro/internal/buffer"
	"github.com/zyedidia/micro/internal/config"
	"github.com/zyedidia/micro/internal/display"
	ulua "github.com/zyedidia/micro/internal/lua"
	"github.com/zyedidia/micro/internal/screen"
	"github.com/zyedidia/micro/internal/shell"
	"github.com/zyedidia/micro/internal/util"
)

func init() {
	ulua.L = lua.NewState()
	ulua.L.SetGlobal("import", luar.New(ulua.L, LuaImport))
}

func LuaImport(pkg string) *lua.LTable {
	switch pkg {
	case "micro":
		return luaImportMicro()
	case "micro/shell":
		return luaImportMicroShell()
	case "micro/buffer":
		return luaImportMicroBuffer()
	case "micro/config":
		return luaImportMicroConfig()
	case "micro/util":
		return luaImportMicroUtil()
	default:
		return ulua.Import(pkg)
	}
}

func luaImportMicro() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "TermMessage", luar.New(ulua.L, screen.TermMessage))
	ulua.L.SetField(pkg, "TermError", luar.New(ulua.L, screen.TermError))
	ulua.L.SetField(pkg, "InfoBar", luar.New(ulua.L, action.GetInfoBar))
	ulua.L.SetField(pkg, "Log", luar.New(ulua.L, log.Println))
	ulua.L.SetField(pkg, "SetStatusInfoFn", luar.New(ulua.L, display.SetStatusInfoFnLua))

	return pkg
}

func luaImportMicroConfig() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "MakeCommand", luar.New(ulua.L, action.LuaMakeCommand))
	ulua.L.SetField(pkg, "FileComplete", luar.New(ulua.L, buffer.FileComplete))
	ulua.L.SetField(pkg, "HelpComplete", luar.New(ulua.L, action.HelpComplete))
	ulua.L.SetField(pkg, "OptionComplete", luar.New(ulua.L, action.OptionComplete))
	ulua.L.SetField(pkg, "OptionValueComplete", luar.New(ulua.L, action.OptionValueComplete))
	ulua.L.SetField(pkg, "NoComplete", luar.New(ulua.L, nil))
	ulua.L.SetField(pkg, "TryBindKey", luar.New(ulua.L, action.TryBindKey))
	ulua.L.SetField(pkg, "Reload", luar.New(ulua.L, action.ReloadConfig))
	ulua.L.SetField(pkg, "AddRuntimeFileFromMemory", luar.New(ulua.L, config.PluginAddRuntimeFileFromMemory))
	ulua.L.SetField(pkg, "AddRuntimeFilesFromDirectory", luar.New(ulua.L, config.PluginAddRuntimeFileFromMemory))
	ulua.L.SetField(pkg, "AddRuntimeFile", luar.New(ulua.L, config.PluginAddRuntimeFile))
	ulua.L.SetField(pkg, "ListRuntimeFiles", luar.New(ulua.L, config.PluginListRuntimeFiles))
	ulua.L.SetField(pkg, "ReadRuntimeFile", luar.New(ulua.L, config.PluginReadRuntimeFile))
	ulua.L.SetField(pkg, "RTColorscheme", luar.New(ulua.L, config.RTColorscheme))
	ulua.L.SetField(pkg, "RTSyntax", luar.New(ulua.L, config.RTSyntax))
	ulua.L.SetField(pkg, "RTHelp", luar.New(ulua.L, config.RTHelp))
	ulua.L.SetField(pkg, "RTPlugin", luar.New(ulua.L, config.RTPlugin))

	return pkg
}

func luaImportMicroShell() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "ExecCommand", luar.New(ulua.L, shell.ExecCommand))
	ulua.L.SetField(pkg, "RunCommand", luar.New(ulua.L, shell.RunCommand))
	ulua.L.SetField(pkg, "RunBackgroundShell", luar.New(ulua.L, shell.RunBackgroundShell))
	ulua.L.SetField(pkg, "RunInteractiveShell", luar.New(ulua.L, shell.RunInteractiveShell))
	ulua.L.SetField(pkg, "JobStart", luar.New(ulua.L, shell.JobStart))
	ulua.L.SetField(pkg, "JobSpawn", luar.New(ulua.L, shell.JobSpawn))
	ulua.L.SetField(pkg, "JobStop", luar.New(ulua.L, shell.JobStop))
	ulua.L.SetField(pkg, "JobSend", luar.New(ulua.L, shell.JobSend))

	return pkg
}

func luaImportMicroBuffer() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "NewMessage", luar.New(ulua.L, buffer.NewMessage))
	ulua.L.SetField(pkg, "NewMessageAtLine", luar.New(ulua.L, buffer.NewMessageAtLine))
	ulua.L.SetField(pkg, "MTInfo", luar.New(ulua.L, buffer.MTInfo))
	ulua.L.SetField(pkg, "MTWarning", luar.New(ulua.L, buffer.MTWarning))
	ulua.L.SetField(pkg, "MTError", luar.New(ulua.L, buffer.MTError))
	ulua.L.SetField(pkg, "Loc", luar.New(ulua.L, func(x, y int) buffer.Loc {
		return buffer.Loc{x, y}
	}))

	return pkg
}

func luaImportMicroUtil() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "RuneAt", luar.New(ulua.L, util.LuaRuneAt))
	ulua.L.SetField(pkg, "GetLeadingWhitespace", luar.New(ulua.L, util.LuaGetLeadingWhitespace))
	ulua.L.SetField(pkg, "IsWordChar", luar.New(ulua.L, util.LuaIsWordChar))

	return pkg
}
<commit_msg>Add buftype access for plugins<commit_after>package main

import (
	"log"

	lua "github.com/yuin/gopher-lua"
	luar "layeh.com/gopher-luar"

	"github.com/zyedidia/micro/internal/action"
	"github.com/zyedidia/micro/internal/buffer"
	"github.com/zyedidia/micro/internal/config"
	"github.com/zyedidia/micro/internal/display"
	ulua "github.com/zyedidia/micro/internal/lua"
	"github.com/zyedidia/micro/internal/screen"
	"github.com/zyedidia/micro/internal/shell"
	"github.com/zyedidia/micro/internal/util"
)

func init() {
	ulua.L = lua.NewState()
	ulua.L.SetGlobal("import", luar.New(ulua.L, LuaImport))
}

func LuaImport(pkg string) *lua.LTable {
	switch pkg {
	case "micro":
		return luaImportMicro()
	case "micro/shell":
		return luaImportMicroShell()
	case "micro/buffer":
		return luaImportMicroBuffer()
	case "micro/config":
		return luaImportMicroConfig()
	case "micro/util":
		return luaImportMicroUtil()
	default:
		return ulua.Import(pkg)
	}
}

func luaImportMicro() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "TermMessage", luar.New(ulua.L, screen.TermMessage))
	ulua.L.SetField(pkg, "TermError", luar.New(ulua.L, screen.TermError))
	ulua.L.SetField(pkg, "InfoBar", luar.New(ulua.L, action.GetInfoBar))
	ulua.L.SetField(pkg, "Log", luar.New(ulua.L, log.Println))
	ulua.L.SetField(pkg, "SetStatusInfoFn", luar.New(ulua.L, display.SetStatusInfoFnLua))

	return pkg
}

func luaImportMicroConfig() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "MakeCommand", luar.New(ulua.L, action.LuaMakeCommand))
	ulua.L.SetField(pkg, "FileComplete", luar.New(ulua.L, buffer.FileComplete))
	ulua.L.SetField(pkg, "HelpComplete", luar.New(ulua.L, action.HelpComplete))
	ulua.L.SetField(pkg, "OptionComplete", luar.New(ulua.L, action.OptionComplete))
	ulua.L.SetField(pkg, "OptionValueComplete", luar.New(ulua.L, action.OptionValueComplete))
	ulua.L.SetField(pkg, "NoComplete", luar.New(ulua.L, nil))
	ulua.L.SetField(pkg, "TryBindKey", luar.New(ulua.L, action.TryBindKey))
	ulua.L.SetField(pkg, "Reload", luar.New(ulua.L, action.ReloadConfig))
	ulua.L.SetField(pkg, "AddRuntimeFileFromMemory", luar.New(ulua.L, config.PluginAddRuntimeFileFromMemory))
	ulua.L.SetField(pkg, "AddRuntimeFilesFromDirectory", luar.New(ulua.L, config.PluginAddRuntimeFileFromMemory))
	ulua.L.SetField(pkg, "AddRuntimeFile", luar.New(ulua.L, config.PluginAddRuntimeFile))
	ulua.L.SetField(pkg, "ListRuntimeFiles", luar.New(ulua.L, config.PluginListRuntimeFiles))
	ulua.L.SetField(pkg, "ReadRuntimeFile", luar.New(ulua.L, config.PluginReadRuntimeFile))
	ulua.L.SetField(pkg, "RTColorscheme", luar.New(ulua.L, config.RTColorscheme))
	ulua.L.SetField(pkg, "RTSyntax", luar.New(ulua.L, config.RTSyntax))
	ulua.L.SetField(pkg, "RTHelp", luar.New(ulua.L, config.RTHelp))
	ulua.L.SetField(pkg, "RTPlugin", luar.New(ulua.L, config.RTPlugin))

	return pkg
}

func luaImportMicroShell() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "ExecCommand", luar.New(ulua.L, shell.ExecCommand))
	ulua.L.SetField(pkg, "RunCommand", luar.New(ulua.L, shell.RunCommand))
	ulua.L.SetField(pkg, "RunBackgroundShell", luar.New(ulua.L, shell.RunBackgroundShell))
	ulua.L.SetField(pkg, "RunInteractiveShell", luar.New(ulua.L, shell.RunInteractiveShell))
	ulua.L.SetField(pkg, "JobStart", luar.New(ulua.L, shell.JobStart))
	ulua.L.SetField(pkg, "JobSpawn", luar.New(ulua.L, shell.JobSpawn))
	ulua.L.SetField(pkg, "JobStop", luar.New(ulua.L, shell.JobStop))
	ulua.L.SetField(pkg, "JobSend", luar.New(ulua.L, shell.JobSend))

	return pkg
}

func luaImportMicroBuffer() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "NewMessage", luar.New(ulua.L, buffer.NewMessage))
	ulua.L.SetField(pkg, "NewMessageAtLine", luar.New(ulua.L, buffer.NewMessageAtLine))
	ulua.L.SetField(pkg, "MTInfo", luar.New(ulua.L, buffer.MTInfo))
	ulua.L.SetField(pkg, "MTWarning", luar.New(ulua.L, buffer.MTWarning))
	ulua.L.SetField(pkg, "MTError", luar.New(ulua.L, buffer.MTError))
	ulua.L.SetField(pkg, "Loc", luar.New(ulua.L, func(x, y int) buffer.Loc {
		return buffer.Loc{x, y}
	}))
	ulua.L.SetField(pkg, "BTDefault", luar.New(ulua.L, buffer.BTDefault.Kind))
	ulua.L.SetField(pkg, "BTHelp", luar.New(ulua.L, buffer.BTHelp.Kind))
	ulua.L.SetField(pkg, "BTLog", luar.New(ulua.L, buffer.BTLog.Kind))
	ulua.L.SetField(pkg, "BTScratch", luar.New(ulua.L, buffer.BTScratch.Kind))
	ulua.L.SetField(pkg, "BTRaw", luar.New(ulua.L, buffer.BTRaw.Kind))
	ulua.L.SetField(pkg, "BTInfo", luar.New(ulua.L, buffer.BTInfo.Kind))

	return pkg
}

func luaImportMicroUtil() *lua.LTable {
	pkg := ulua.L.NewTable()

	ulua.L.SetField(pkg, "RuneAt", luar.New(ulua.L, util.LuaRuneAt))
	ulua.L.SetField(pkg, "GetLeadingWhitespace", luar.New(ulua.L, util.LuaGetLeadingWhitespace))
	ulua.L.SetField(pkg, "IsWordChar", luar.New(ulua.L, util.LuaIsWordChar))

	return pkg
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd 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 (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"syscall"

	"github.com/google/crfs/stargz"
	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/logs"
	"github.com/google/go-containerregistry/pkg/name"
	regpkg "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/mutate"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/v1/stream"
	"github.com/google/go-containerregistry/pkg/v1/types"
	"github.com/ktock/remote-snapshotter/cmd/optimize/logger"
	"github.com/ktock/remote-snapshotter/cmd/optimize/sampler"
	"github.com/ktock/remote-snapshotter/cmd/optimize/sorter"
	imgpkg "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"
)

var (
	insecure   = flag.Bool("insecure", false, "allow HTTP connections to the registry which has the prefix \"http://\"")
	noopt      = flag.Bool("noopt", false, "only stargzify and do not optimize layers")
	period     = flag.Int("period", 10, "time period to monitor access log")
	username   = flag.String("user", "", "user name to override image's default config")
	cwd        = flag.String("cwd", "", "working dir to override image's default config")
	cArgs      = flag.String("args", "", "command arguments to override image's default config(in JSON array)")
	entrypoint = flag.String("entrypoint", "", "entrypoint to override image's default config(in JSON array)")
	terminal   = flag.Bool("t", false, "enable terminal for sample container")
	cEnvs      envs
)

type envs struct {
	list []string
}

func (e *envs) String() string {
	return strings.Join(e.list, ", ")
}

func (e *envs) Set(value string) error {
	e.list = append(e.list, value)
	return nil
}

func main() {

	// Set up logs package to get useful messages i.e. progress.
	logs.Warn.SetOutput(os.Stderr)
	logs.Progress.SetOutput(os.Stderr)

	// Parse arguments
	flag.Var(&cEnvs, "env", "environment valulable to add or override to the image's default config")
	flag.Parse()
	if flag.NArg() < 2 {
		fmt.Println("usage: $0 [OPTION]... INPUT_IMAGE OUTPUT_IMAGE")
		os.Exit(1)
	}
	args := flag.Args()
	src, dst, opts, err := parseArgs(args)
	if err != nil {
		log.Fatal(err)
	}

	// Convert and push image
	srcRef, err := parseReference(src)
	if err != nil {
		log.Fatal(err)
	}
	dstRef, err := parseReference(dst)
	if err != nil {
		log.Fatal(err)
	}
	err = convert(srcRef, dstRef, opts...)
	if err != nil {
		log.Fatal(err)
	}
}

func parseArgs(args []string) (src string, dst string, opts []sampler.Option, err error) {
	src = args[0]
	dst = args[1]

	if len(cEnvs.list) > 0 {
		opts = append(opts, sampler.WithEnvs(cEnvs.list))
	}
	if *cArgs != "" {
		var as []string
		err = json.Unmarshal([]byte(*cArgs), &as)
		if err != nil {
			return "", "", nil, errors.Wrapf(err, "invalid option \"args\"")
		}
		opts = append(opts, sampler.WithArgs(as))
	}
	if *entrypoint != "" {
		var es []string
		err = json.Unmarshal([]byte(*entrypoint), &es)
		if err != nil {
			return "", "", nil, errors.Wrapf(err, "invalid option \"entrypoint\"")
		}
		opts = append(opts, sampler.WithEntrypoint(es))
	}
	if *username != "" {
		opts = append(opts, sampler.WithUser(*username))
	}
	if *cwd != "" {
		opts = append(opts, sampler.WithWorkingDir(*cwd))
	}
	if *terminal {
		opts = append(opts, sampler.WithTerminal())
	}

	return
}

func parseReference(path string) (name.Reference, error) {
	var opts []name.Option
	if strings.HasPrefix(path, "http://") {
		path = strings.TrimPrefix(path, "http://")
		if *insecure {
			opts = append(opts, name.Insecure)
		} else {
			return nil, fmt.Errorf("\"-insecure\" option must be specified to connect to %q using HTTP", path)
		}
	}
	ref, err := name.ParseReference(path, opts...)
	if err != nil {
		return nil, err
	}

	return ref, nil
}

func convert(srcRef, dstRef name.Reference, runopts ...sampler.Option) error {
	// Pull source image
	srcImg, err := remote.Image(srcRef, remote.WithAuthFromKeychain(authn.DefaultKeychain))
	if err != nil {
		return err
	}

	// Optimize the image
	// The order of the list is base layer first, top layer last.
	layers, err := srcImg.Layers()
	if err != nil {
		return err
	}
	if !*noopt {
		configData, err := srcImg.RawConfigFile()
		if err != nil {
			return err
		}
		var config imgpkg.Image
		if err := json.Unmarshal(configData, &config); err != nil {
			return fmt.Errorf("failed to parse image config file: %v", err)
		}
		layers, err = optimize(layers, config, runopts...)
		if err != nil {
			return err
		}
	} else {
		for i, l := range layers {
			r, err := l.Uncompressed()
			if err != nil {
				return err
			}
			layers[i] = &layer{r: r}
		}
	}
	srcCfg, err := srcImg.ConfigFile()
	if err != nil {
		return err
	}
	srcCfg.RootFS.DiffIDs = []regpkg.Hash{}
	srcCfg.History = []regpkg.History{}
	img, err := mutate.ConfigFile(empty.Image, srcCfg)
	if err != nil {
		return err
	}
	img, err = mutate.AppendLayers(img, layers...)
	if err != nil {
		return err
	}

	// Push the optimized image.
	if err := remote.Write(dstRef, img, remote.WithAuthFromKeychain(authn.DefaultKeychain)); err != nil {
		return err
	}

	return nil
}

// The order of the "in" list must be base layer first, top layer last.
func optimize(in []regpkg.Layer, config imgpkg.Image, opts ...sampler.Option) (out []regpkg.Layer, err error) {
	root := ""
	mktemp := func() (path string, err error) {
		if path, err = ioutil.TempDir(root, "optimize"); err != nil {
			return "", err
		}
		if err = os.Chmod(path, 0777); err != nil {
			return "", err
		}
		return path, nil
	}
	if root, err = mktemp(); err != nil {
		return nil, err
	}
	defer os.RemoveAll(root)

	// mount layer loggers on temp directories.
	var (
		eg        errgroup.Group
		lowerdirs []string
		addLayer  = make([](func([]regpkg.Layer) ([]regpkg.Layer, error)), len(in))
	)
	for i := range in {
		i := i
		mp, err := mktemp()
		if err != nil {
			return nil, err
		}
		defer syscall.Unmount(mp, syscall.MNT_FORCE)
		lowerdirs = append([]string{mp}, lowerdirs...) // top layer first, base layer last (for overlayfs).
		eg.Go(func() error {
			dgst, err := in[i].Digest()
			if err != nil {
				return err
			}
			r, err := in[i].Uncompressed()
			if err != nil {
				return err
			}
			defer r.Close()
			tarBytes, err := ioutil.ReadAll(r)
			if err != nil {
				return err
			}
			mon := logger.NewOpenReadMonitor()
			if _, err := logger.Mount(mp, bytes.NewReader(tarBytes), mon); err != nil {
				return errors.Wrapf(err, "failed to mount on %s", mp)
			}
			addLayer[i] = func(lowers []regpkg.Layer) ([]regpkg.Layer, error) {
				r, err := sorter.Sort(bytes.NewReader(tarBytes), mon.DumpLog())
				if err != nil {
					return nil, errors.Wrap(err, "failed to sort tar")
				}
				return append(lowers, &layer{r: r}), nil
			}
			fmt.Printf("Unpacked %v\n", dgst)
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		return nil, err
	}

	// prepare FileSystem Bundle
	var (
		bundle   string
		upperdir string
		workdir  string
	)
	if bundle, err = mktemp(); err != nil {
		return nil, err
	}
	if upperdir, err = mktemp(); err != nil {
		return nil, err
	}
	if workdir, err = mktemp(); err != nil {
		return nil, err
	}
	var (
		rootfs = sampler.GetRootfsPathUnder(bundle)
		option = fmt.Sprintf("lowerdir=%s,upperdir=%s,workdir=%s",
			strings.Join(lowerdirs, ":"), upperdir, workdir)
	)
	if err = os.Mkdir(rootfs, 0777); err != nil {
		return nil, err
	}
	if err = syscall.Mount("overlay", rootfs, "overlay", 0, option); err != nil {
		return nil, errors.Wrapf(err, "mount overlayfs on %q with data %q", rootfs, option)
	}
	defer syscall.Unmount(rootfs, syscall.MNT_FORCE)

	// run workload
	if err = sampler.Run(bundle, config, *period, opts...); err != nil {
		return nil, errors.Wrap(err, "failed to run the sampler")
	}

	// get converted layers
	for _, f := range addLayer {
		if out, err = f(out); err != nil {
			return nil, errors.Wrap(err, "failed to get converted layer")
		}
	}
	return
}

type writer func([]byte) (int, error)

func (f writer) Write(b []byte) (int, error) { return f(b) }

type layer struct {
	r    io.Reader
	diff *regpkg.Hash // registered after compression completed
	hash *hash.Hash   // registered after compression completed
	size *int64       // registered after compression completed
}

func (l *layer) Digest() (regpkg.Hash, error) {
	if l.hash == nil {
		return regpkg.Hash{}, stream.ErrNotComputed
	}
	return regpkg.Hash{
		Algorithm: "sha256",
		Hex:       hex.EncodeToString((*l.hash).Sum(nil)),
	}, nil
}

func (l *layer) Size() (int64, error) {
	if l.size == nil {
		return -1, stream.ErrNotComputed
	}
	return *l.size, nil
}

func (l *layer) DiffID() (regpkg.Hash, error) {
	if l.diff == nil {
		return regpkg.Hash{}, stream.ErrNotComputed
	}
	return *l.diff, nil
}

func (l *layer) MediaType() (types.MediaType, error) {
	return types.DockerLayer, nil
}

func (l *layer) Compressed() (io.ReadCloser, error) {
	pr, pw := io.Pipe()
	go func() {
		var (
			diff regpkg.Hash
			h    = sha256.New()
			size int64
			err  error
		)
		w := stargz.NewWriter(io.MultiWriter(pw, writer(func(b []byte) (int, error) {
			n, err := h.Write(b)
			size += int64(n)
			return n, err
		})))
		if err := w.AppendTar(l.r); err != nil {
			pw.CloseWithError(err)
			return
		}
		if err := w.Close(); err != nil {
			pw.CloseWithError(err)
			return
		}
		diff, err = regpkg.NewHash(w.DiffID())
		if err != nil {
			pw.CloseWithError(err)
			return
		}

		// registers all computed information
		l.diff = &diff
		l.hash = &h
		l.size = &size
		pw.Close()
	}()
	return ioutil.NopCloser(pr), nil
}

func (l *layer) Uncompressed() (io.ReadCloser, error) {
	return ioutil.NopCloser(l.r), nil
}
<commit_msg>optimize: $0 -> os.Args[0]<commit_after>/*
   Copyright The containerd 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 (
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"syscall"

	"github.com/google/crfs/stargz"
	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/logs"
	"github.com/google/go-containerregistry/pkg/name"
	regpkg "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/mutate"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/google/go-containerregistry/pkg/v1/stream"
	"github.com/google/go-containerregistry/pkg/v1/types"
	"github.com/ktock/remote-snapshotter/cmd/optimize/logger"
	"github.com/ktock/remote-snapshotter/cmd/optimize/sampler"
	"github.com/ktock/remote-snapshotter/cmd/optimize/sorter"
	imgpkg "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
	"golang.org/x/sync/errgroup"
)

var (
	insecure   = flag.Bool("insecure", false, "allow HTTP connections to the registry which has the prefix \"http://\"")
	noopt      = flag.Bool("noopt", false, "only stargzify and do not optimize layers")
	period     = flag.Int("period", 10, "time period to monitor access log")
	username   = flag.String("user", "", "user name to override image's default config")
	cwd        = flag.String("cwd", "", "working dir to override image's default config")
	cArgs      = flag.String("args", "", "command arguments to override image's default config(in JSON array)")
	entrypoint = flag.String("entrypoint", "", "entrypoint to override image's default config(in JSON array)")
	terminal   = flag.Bool("t", false, "enable terminal for sample container")
	cEnvs      envs
)

type envs struct {
	list []string
}

func (e *envs) String() string {
	return strings.Join(e.list, ", ")
}

func (e *envs) Set(value string) error {
	e.list = append(e.list, value)
	return nil
}

func main() {

	// Set up logs package to get useful messages i.e. progress.
	logs.Warn.SetOutput(os.Stderr)
	logs.Progress.SetOutput(os.Stderr)

	// Parse arguments
	flag.Var(&cEnvs, "env", "environment valulable to add or override to the image's default config")
	flag.Parse()
	if flag.NArg() < 2 {
		fmt.Printf("usage: %s [OPTION]... INPUT_IMAGE OUTPUT_IMAGE\n", os.Args[0])
		os.Exit(1)
	}
	args := flag.Args()
	src, dst, opts, err := parseArgs(args)
	if err != nil {
		log.Fatal(err)
	}

	// Convert and push image
	srcRef, err := parseReference(src)
	if err != nil {
		log.Fatal(err)
	}
	dstRef, err := parseReference(dst)
	if err != nil {
		log.Fatal(err)
	}
	err = convert(srcRef, dstRef, opts...)
	if err != nil {
		log.Fatal(err)
	}
}

func parseArgs(args []string) (src string, dst string, opts []sampler.Option, err error) {
	src = args[0]
	dst = args[1]

	if len(cEnvs.list) > 0 {
		opts = append(opts, sampler.WithEnvs(cEnvs.list))
	}
	if *cArgs != "" {
		var as []string
		err = json.Unmarshal([]byte(*cArgs), &as)
		if err != nil {
			return "", "", nil, errors.Wrapf(err, "invalid option \"args\"")
		}
		opts = append(opts, sampler.WithArgs(as))
	}
	if *entrypoint != "" {
		var es []string
		err = json.Unmarshal([]byte(*entrypoint), &es)
		if err != nil {
			return "", "", nil, errors.Wrapf(err, "invalid option \"entrypoint\"")
		}
		opts = append(opts, sampler.WithEntrypoint(es))
	}
	if *username != "" {
		opts = append(opts, sampler.WithUser(*username))
	}
	if *cwd != "" {
		opts = append(opts, sampler.WithWorkingDir(*cwd))
	}
	if *terminal {
		opts = append(opts, sampler.WithTerminal())
	}

	return
}

func parseReference(path string) (name.Reference, error) {
	var opts []name.Option
	if strings.HasPrefix(path, "http://") {
		path = strings.TrimPrefix(path, "http://")
		if *insecure {
			opts = append(opts, name.Insecure)
		} else {
			return nil, fmt.Errorf("\"-insecure\" option must be specified to connect to %q using HTTP", path)
		}
	}
	ref, err := name.ParseReference(path, opts...)
	if err != nil {
		return nil, err
	}

	return ref, nil
}

func convert(srcRef, dstRef name.Reference, runopts ...sampler.Option) error {
	// Pull source image
	srcImg, err := remote.Image(srcRef, remote.WithAuthFromKeychain(authn.DefaultKeychain))
	if err != nil {
		return err
	}

	// Optimize the image
	// The order of the list is base layer first, top layer last.
	layers, err := srcImg.Layers()
	if err != nil {
		return err
	}
	if !*noopt {
		configData, err := srcImg.RawConfigFile()
		if err != nil {
			return err
		}
		var config imgpkg.Image
		if err := json.Unmarshal(configData, &config); err != nil {
			return fmt.Errorf("failed to parse image config file: %v", err)
		}
		layers, err = optimize(layers, config, runopts...)
		if err != nil {
			return err
		}
	} else {
		for i, l := range layers {
			r, err := l.Uncompressed()
			if err != nil {
				return err
			}
			layers[i] = &layer{r: r}
		}
	}
	srcCfg, err := srcImg.ConfigFile()
	if err != nil {
		return err
	}
	srcCfg.RootFS.DiffIDs = []regpkg.Hash{}
	srcCfg.History = []regpkg.History{}
	img, err := mutate.ConfigFile(empty.Image, srcCfg)
	if err != nil {
		return err
	}
	img, err = mutate.AppendLayers(img, layers...)
	if err != nil {
		return err
	}

	// Push the optimized image.
	if err := remote.Write(dstRef, img, remote.WithAuthFromKeychain(authn.DefaultKeychain)); err != nil {
		return err
	}

	return nil
}

// The order of the "in" list must be base layer first, top layer last.
func optimize(in []regpkg.Layer, config imgpkg.Image, opts ...sampler.Option) (out []regpkg.Layer, err error) {
	root := ""
	mktemp := func() (path string, err error) {
		if path, err = ioutil.TempDir(root, "optimize"); err != nil {
			return "", err
		}
		if err = os.Chmod(path, 0777); err != nil {
			return "", err
		}
		return path, nil
	}
	if root, err = mktemp(); err != nil {
		return nil, err
	}
	defer os.RemoveAll(root)

	// mount layer loggers on temp directories.
	var (
		eg        errgroup.Group
		lowerdirs []string
		addLayer  = make([](func([]regpkg.Layer) ([]regpkg.Layer, error)), len(in))
	)
	for i := range in {
		i := i
		mp, err := mktemp()
		if err != nil {
			return nil, err
		}
		defer syscall.Unmount(mp, syscall.MNT_FORCE)
		lowerdirs = append([]string{mp}, lowerdirs...) // top layer first, base layer last (for overlayfs).
		eg.Go(func() error {
			dgst, err := in[i].Digest()
			if err != nil {
				return err
			}
			r, err := in[i].Uncompressed()
			if err != nil {
				return err
			}
			defer r.Close()
			tarBytes, err := ioutil.ReadAll(r)
			if err != nil {
				return err
			}
			mon := logger.NewOpenReadMonitor()
			if _, err := logger.Mount(mp, bytes.NewReader(tarBytes), mon); err != nil {
				return errors.Wrapf(err, "failed to mount on %s", mp)
			}
			addLayer[i] = func(lowers []regpkg.Layer) ([]regpkg.Layer, error) {
				r, err := sorter.Sort(bytes.NewReader(tarBytes), mon.DumpLog())
				if err != nil {
					return nil, errors.Wrap(err, "failed to sort tar")
				}
				return append(lowers, &layer{r: r}), nil
			}
			fmt.Printf("Unpacked %v\n", dgst)
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		return nil, err
	}

	// prepare FileSystem Bundle
	var (
		bundle   string
		upperdir string
		workdir  string
	)
	if bundle, err = mktemp(); err != nil {
		return nil, err
	}
	if upperdir, err = mktemp(); err != nil {
		return nil, err
	}
	if workdir, err = mktemp(); err != nil {
		return nil, err
	}
	var (
		rootfs = sampler.GetRootfsPathUnder(bundle)
		option = fmt.Sprintf("lowerdir=%s,upperdir=%s,workdir=%s",
			strings.Join(lowerdirs, ":"), upperdir, workdir)
	)
	if err = os.Mkdir(rootfs, 0777); err != nil {
		return nil, err
	}
	if err = syscall.Mount("overlay", rootfs, "overlay", 0, option); err != nil {
		return nil, errors.Wrapf(err, "mount overlayfs on %q with data %q", rootfs, option)
	}
	defer syscall.Unmount(rootfs, syscall.MNT_FORCE)

	// run workload
	if err = sampler.Run(bundle, config, *period, opts...); err != nil {
		return nil, errors.Wrap(err, "failed to run the sampler")
	}

	// get converted layers
	for _, f := range addLayer {
		if out, err = f(out); err != nil {
			return nil, errors.Wrap(err, "failed to get converted layer")
		}
	}
	return
}

type writer func([]byte) (int, error)

func (f writer) Write(b []byte) (int, error) { return f(b) }

type layer struct {
	r    io.Reader
	diff *regpkg.Hash // registered after compression completed
	hash *hash.Hash   // registered after compression completed
	size *int64       // registered after compression completed
}

func (l *layer) Digest() (regpkg.Hash, error) {
	if l.hash == nil {
		return regpkg.Hash{}, stream.ErrNotComputed
	}
	return regpkg.Hash{
		Algorithm: "sha256",
		Hex:       hex.EncodeToString((*l.hash).Sum(nil)),
	}, nil
}

func (l *layer) Size() (int64, error) {
	if l.size == nil {
		return -1, stream.ErrNotComputed
	}
	return *l.size, nil
}

func (l *layer) DiffID() (regpkg.Hash, error) {
	if l.diff == nil {
		return regpkg.Hash{}, stream.ErrNotComputed
	}
	return *l.diff, nil
}

func (l *layer) MediaType() (types.MediaType, error) {
	return types.DockerLayer, nil
}

func (l *layer) Compressed() (io.ReadCloser, error) {
	pr, pw := io.Pipe()
	go func() {
		var (
			diff regpkg.Hash
			h    = sha256.New()
			size int64
			err  error
		)
		w := stargz.NewWriter(io.MultiWriter(pw, writer(func(b []byte) (int, error) {
			n, err := h.Write(b)
			size += int64(n)
			return n, err
		})))
		if err := w.AppendTar(l.r); err != nil {
			pw.CloseWithError(err)
			return
		}
		if err := w.Close(); err != nil {
			pw.CloseWithError(err)
			return
		}
		diff, err = regpkg.NewHash(w.DiffID())
		if err != nil {
			pw.CloseWithError(err)
			return
		}

		// registers all computed information
		l.diff = &diff
		l.hash = &h
		l.size = &size
		pw.Close()
	}()
	return ioutil.NopCloser(pr), nil
}

func (l *layer) Uncompressed() (io.ReadCloser, error) {
	return ioutil.NopCloser(l.r), nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2018, Google, Inc.
//
// 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 pipeline

import (
	"fmt"
	"net/http"

	"github.com/spf13/cobra"
	"github.com/spinnaker/spin/cmd/gateclient"

	"github.com/spinnaker/spin/util"
)

type SaveOptions struct {
	*pipelineOptions
	output       string
	pipelineFile string
}

var (
	savePipelineShort   = "Save the provided pipeline"
	savePipelineLong    = "Save the provided pipeline"
)

func NewSaveCmd(pipelineOptions pipelineOptions) *cobra.Command {
	options := SaveOptions{
		pipelineOptions: &pipelineOptions,
	}
	cmd := &cobra.Command{
		Use:     "save",
		Aliases: []string{},
		Short:   savePipelineShort,
		Long:    savePipelineLong,
		RunE: func(cmd *cobra.Command, args []string) error {
			return savePipeline(cmd, options)
		},
	}

	cmd.PersistentFlags().StringVarP(&options.pipelineFile, "file", "f", "", "path to the pipeline file")

	return cmd
}

func savePipeline(cmd *cobra.Command, options SaveOptions) error {
	gateClient, err := gateclient.NewGateClient(cmd.InheritedFlags())
	if err != nil {
		return err
	}

	pipelineJson, err := util.ParseJsonFromFileOrStdin(options.pipelineFile, false)
	if err != nil {
		return err
	}
	valid := true
	if _, exists := pipelineJson["name"]; !exists {
		util.UI.Error("Required pipeline key 'name' missing...\n")
		valid = false
	}

	if _, exists := pipelineJson["application"]; !exists {
		util.UI.Error("Required pipeline key 'application' missing...\n")
		valid = false
	}

	if template, exists := pipelineJson["template"]; exists && len(template.(map[string]interface{})) > 0 {
		if _, exists := pipelineJson["schema"]; !exists {
			util.UI.Error("Required pipeline key 'schema' missing for templated pipeline...\n")
			valid = false
		}
		pipelineJson["type"] = "templatedPipeline"
	}

	if !valid {
		return fmt.Errorf("Submitted pipeline is invalid: %s\n", pipelineJson)
	}

	resp, err := gateClient.PipelineControllerApi.SavePipelineUsingPOST(gateClient.Context, pipelineJson)

	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("Encountered an error saving pipeline, status code: %d\n", resp.StatusCode)
	}

	util.UI.Info(util.Colorize().Color(fmt.Sprintf("[reset][bold][green]Pipeline save succeeded")))
	return nil
}
<commit_msg>fix(pipeline): Query pipeline id for subsequent saves. (#124)<commit_after>// Copyright (c) 2018, Google, Inc.
//
// 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 pipeline

import (
	"fmt"
	"net/http"

	"github.com/spf13/cobra"
	"github.com/spinnaker/spin/cmd/gateclient"

	"github.com/spinnaker/spin/util"
)

type SaveOptions struct {
	*pipelineOptions
	output       string
	pipelineFile string
}

var (
	savePipelineShort = "Save the provided pipeline"
	savePipelineLong  = "Save the provided pipeline"
)

func NewSaveCmd(pipelineOptions pipelineOptions) *cobra.Command {
	options := SaveOptions{
		pipelineOptions: &pipelineOptions,
	}
	cmd := &cobra.Command{
		Use:     "save",
		Aliases: []string{},
		Short:   savePipelineShort,
		Long:    savePipelineLong,
		RunE: func(cmd *cobra.Command, args []string) error {
			return savePipeline(cmd, options)
		},
	}

	cmd.PersistentFlags().StringVarP(&options.pipelineFile, "file", "f", "", "path to the pipeline file")

	return cmd
}

func savePipeline(cmd *cobra.Command, options SaveOptions) error {
	gateClient, err := gateclient.NewGateClient(cmd.InheritedFlags())
	if err != nil {
		return err
	}

	pipelineJson, err := util.ParseJsonFromFileOrStdin(options.pipelineFile, false)
	if err != nil {
		return err
	}
	valid := true
	if _, exists := pipelineJson["name"]; !exists {
		util.UI.Error("Required pipeline key 'name' missing...\n")
		valid = false
	}

	if _, exists := pipelineJson["application"]; !exists {
		util.UI.Error("Required pipeline key 'application' missing...\n")
		valid = false
	}

	if template, exists := pipelineJson["template"]; exists && len(template.(map[string]interface{})) > 0 {
		if _, exists := pipelineJson["schema"]; !exists {
			util.UI.Error("Required pipeline key 'schema' missing for templated pipeline...\n")
			valid = false
		}
	}

	if !valid {
		return fmt.Errorf("Submitted pipeline is invalid: %s\n", pipelineJson)
	}
	pipelineJson["type"] = "templatedPipeline"
	application := pipelineJson["application"].(string)
	pipelineName := pipelineJson["name"].(string)

	foundPipeline, queryResp, _ := gateClient.ApplicationControllerApi.GetPipelineConfigUsingGET(gateClient.Context, application, pipelineName)

	if queryResp.StatusCode != http.StatusOK {
		return fmt.Errorf("Encountered an error querying pipeline, status code: %d\n", queryResp.StatusCode)
	}

	_, exists := pipelineJson["id"].(string)
	var foundPipelineId string
	if len(foundPipeline) > 0 {
		foundPipelineId = foundPipeline["id"].(string)
	}
	if !exists && foundPipelineId != "" {
		pipelineJson["id"] = foundPipelineId
	}

	saveResp, saveErr := gateClient.PipelineControllerApi.SavePipelineUsingPOST(gateClient.Context, pipelineJson)

	if saveErr != nil {
		fmt.Printf("s err: %v", saveErr)
		return err
	}
	if saveResp.StatusCode != http.StatusOK {
		return fmt.Errorf("Encountered an error saving pipeline, status code: %d\n", saveResp.StatusCode)
	}

	util.UI.Info(util.Colorize().Color(fmt.Sprintf("[reset][bold][green]Pipeline save succeeded")))
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/jmmcatee/cracklord/common"
	"github.com/jmmcatee/cracklord/common/log"
	"github.com/jmmcatee/cracklord/common/queue"
	"github.com/unrolled/secure"
	"github.com/vaughan0/go-ini"
	"net/http"
	"os"
	"strconv"
)

const (
	QUEUED_INIT_FILE = "/etc/cracklord/queued.conf"
	CA_CERT_FILE     = "/etc/cracklord/ssl/cracklord_ca.pem"
	QUEUED_KEY_FILE  = "/etc/cracklord/ssl/queued.key"
	QUEUED_CERT_FILE = "/etc/cracklord/ssl/queued.crt"
)

func main() {
	// Define the flags
	var confPath = flag.String("conf", "", "Configuration file to use")
	var webRoot = flag.String("webroot", "./public", "Location of the web server root")
	var runIP = flag.String("host", "0.0.0.0", "IP to bind to")
	var runPort = flag.String("port", "443", "Port to bind to")
	var caCertPath = flag.String("cacert", "", "CA Certficate to use for validation")
	var KeyPath = flag.String("key", "", "Private key for the resource to use over TLS")
	var CertPath = flag.String("cert", "", "Certicate to use with the resource for TLS")

	// Parse the flags
	flag.Parse()

	// Read the configuration file
	var confFile ini.File
	var confErr error
	if *confPath == "" {
		confFile, confErr = ini.LoadFile(QUEUED_INIT_FILE)
	} else {
		confFile, confErr = ini.LoadFile(*confPath)
	}

	// Build the App Controller
	var server AppController

	// Check for errors
	if confErr != nil {
		println("ERROR: Unable to " + confErr.Error())
		println("See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files.")
		return
	}

	genConf := confFile.Section("General")
	switch common.StripQuotes(genConf["LogLevel"]) {
	case "Debug":
		log.SetLevel(log.DebugLevel)
		log.Warn("Please note the debug level logs may contain sensitive information!")
	case "Info":
		log.SetLevel(log.InfoLevel)
	case "Warn":
		log.SetLevel(log.WarnLevel)
	case "Error":
		log.SetLevel(log.ErrorLevel)
	case "Fatal":
		log.SetLevel(log.FatalLevel)
	case "Panic":
		log.SetLevel(log.PanicLevel)
	default:
		log.SetLevel(log.InfoLevel)
	}

	lf := common.StripQuotes(genConf["LogFile"])
	if lf != "" {
		hook, err := cracklog.NewFileHook(lf)
		if err != nil {
			println("ERROR: Unable to open log file: " + err.Error())
		} else {
			log.AddHook(hook)
		}
	}

	var statefile string
	statefile = common.StripQuotes(genConf["StateFile"])

	var updatetime int
	var resourcetimeout int
	utconf := common.StripQuotes(genConf["UpdateTime"])
	if utconf != "" {
		var err error
		updatetime, err = strconv.Atoi(utconf)
		if err != nil {
			log.WithField("error", err.Error()).Error("Unable to parse update time in config file.")
			updatetime = 30
		}
	} else {
		updatetime = 30
	}
	restimeconf := common.StripQuotes(genConf["ResourceTimeout"])
	if restimeconf != "" {
		var err error
		resourcetimeout, err = strconv.Atoi(restimeconf)
		if err != nil {
			log.WithField("error", err.Error()).Error("Unable to parse resource timeout in config file.")
			resourcetimeout = 5
		}
	} else {
		resourcetimeout = 5
	}

	log.WithFields(log.Fields{
		"ip":   *runIP,
		"port": *runPort,
	}).Info("Starting queue server up.")

	// Get the Authentication configuration
	confAuth := confFile.Section("Authentication")
	if confAuth == nil {
		println("Error: Authentication configuration is required.")
		println("See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files.")
		return
	}

	_, weberr := os.Stat(*webRoot)
	if weberr != nil {
		println("Error: Public web root '" + *webRoot + "' does not exist.")
		return
	}

	// Check for type of authentication and set conf
	switch confAuth["type"] {
	case "INI":
		var i INIAuth

		// Get the users
		umap := map[string]string{}

		au := common.StripQuotes(confAuth["adminuser"])
		if au != "" {
			log.Fatal("An administrative user was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		ap := common.StripQuotes(confAuth["adminpass"])
		if ap != "" {
			log.Fatal("An administrative password was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		su := common.StripQuotes(confAuth["standarduser"])
		if su != "" {
			log.Fatal("An standard user was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		sp := common.StripQuotes(confAuth["standarduser"])
		if sp != "" {
			log.Fatal("An standard password was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		ru := common.StripQuotes(confAuth["readonlyuser"])
		if ru != "" {
			log.Fatal("An read only user was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		rp := common.StripQuotes(confAuth["readonlypass"])
		if rp != "" {
			log.Fatal("An read only password was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		umap[au] = ap
		umap[su] = sp
		umap[ru] = rp

		// Setup group mappings
		gmap := map[string]string{}

		gmap[au] = Administrator
		gmap[su] = StandardUser
		gmap[ru] = ReadOnly

		i.Setup(umap, gmap)

		server.Auth = &i

		log.Info("INI authentication setup complete.")
	case "ActiveDirectory":
		var ad ADAuth

		realm := common.StripQuotes(confAuth["realm"])
		if realm != "" {
			log.Fatal("No Active Directory realm was configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		ad.SetRealm(realm)

		gmap := map[string]string{}
		ro := common.StripQuotes(confAuth["ReadOnlyGroup"])
		if ro != "" {
			log.Fatal("A read only group was not provided. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		st := common.StripQuotes(confAuth["StandardGroup"])
		if st != "" {
			log.Fatal("A group for standard access was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		admin := common.StripQuotes(confAuth["AdminGroup"])
		if admin != "" {
			log.Fatal("A group for read only access was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		gmap[ReadOnly] = ro
		gmap[StandardUser] = st
		gmap[Administrator] = admin

		ad.Setup(gmap)

		server.Auth = &ad
		log.WithFields(log.Fields{
			"readonly": ro,
			"standard": st,
			"admin":    admin,
		}).Info("Active directory authentication configured successfully.")
	}

	// Configure the TokenStore
	server.T = NewTokenStore()

	// Configure the Queue
	server.Q = queue.NewQueue(statefile, updatetime, resourcetimeout)

	// Get the CA
	if *caCertPath == "" {
		// Use default path to get the CA certificate
		*caCertPath = CA_CERT_FILE
	}

	caBytes, err := ioutil.ReadFile(*caCertPath)
	if err != nil {
		println("ERROR: " + err.Error())
	}

	caPool := x509.NewCertPool()
	caPool.AppendCertsFromPEM(caBytes)

	// Get the certificates and private key
	if *CertPath == "" || *KeyPath == "" {
		// The private key and/or the certificate were not given so go with defaults
		*KeyPath = QUEUED_KEY_FILE
		*CertPath = QUEUED_CERT_FILE
	}

	tlscert, err := tls.LoadX509KeyPair(*CertPath, *KeyPath)

	// Setup TLS connection
	tlsconfig := &tls.Config{}
	tlsconfig.Certificates = make([]tls.Certificate, 1)
	tlsconfig.Certificates[0] = tlscert
	tlsconfig.RootCAs = caPool
	tlsconfig.ClientCAs = caPool
	tlsconfig.CipherSuites = []uint16{tls.TLS_RSA_WITH_AES_128_CBC_SHA,
		tls.TLS_RSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
		tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
		tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256}
	tlsconfig.MinVersion = tls.VersionTLS12
	tlsconfig.SessionTicketsDisabled = true

	// Set the TLS policy for the Queue to communicate with the resources
	server.TLS = tlsconfig

	// Add some nice security stuff
	secureMiddleware := secure.New(secure.Options{
		SSLRedirect:             true,
		FrameDeny:               true,
		CustomFrameOptionsValue: "SAMEORIGIN",
		BrowserXssFilter:        true,
		IsDevelopment:           true,
	})

	// Build the Negroni handler
	n := negroni.New(negroni.NewRecovery(),
		cracklog.NewNegroniLogger(),
		negroni.NewStatic(http.Dir(*webRoot)))

	n.Use(negroni.HandlerFunc(secureMiddleware.HandlerFuncWithNext))
	n.UseHandler(server.Router())
	log.Debug("Negroni handler started.")

	listen, err := tls.Listen("tcp", *runIP+":"+*runPort, tlsconfig)
	if err != nil {
		println("ERROR: Unable to bind to '" + *runIP + ":" + *runPort + "':" + err.Error())
		return
	}

	err = http.Serve(listen, n)
	if err != nil {
		log.Fatal("Unable to start up web server: " + err.Error())
	}
}
<commit_msg>Update include statement for @jmmcatee<commit_after>package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/jmmcatee/cracklord/common"
	"github.com/jmmcatee/cracklord/common/log"
	"github.com/jmmcatee/cracklord/common/queue"
	"github.com/unrolled/secure"
	"github.com/vaughan0/go-ini"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"
)

const (
	QUEUED_INIT_FILE = "/etc/cracklord/queued.conf"
	CA_CERT_FILE     = "/etc/cracklord/ssl/cracklord_ca.pem"
	QUEUED_KEY_FILE  = "/etc/cracklord/ssl/queued.key"
	QUEUED_CERT_FILE = "/etc/cracklord/ssl/queued.crt"
)

func main() {
	// Define the flags
	var confPath = flag.String("conf", "", "Configuration file to use")
	var webRoot = flag.String("webroot", "./public", "Location of the web server root")
	var runIP = flag.String("host", "0.0.0.0", "IP to bind to")
	var runPort = flag.String("port", "443", "Port to bind to")
	var caCertPath = flag.String("cacert", "", "CA Certficate to use for validation")
	var KeyPath = flag.String("key", "", "Private key for the resource to use over TLS")
	var CertPath = flag.String("cert", "", "Certicate to use with the resource for TLS")

	// Parse the flags
	flag.Parse()

	// Read the configuration file
	var confFile ini.File
	var confErr error
	if *confPath == "" {
		confFile, confErr = ini.LoadFile(QUEUED_INIT_FILE)
	} else {
		confFile, confErr = ini.LoadFile(*confPath)
	}

	// Build the App Controller
	var server AppController

	// Check for errors
	if confErr != nil {
		println("ERROR: Unable to " + confErr.Error())
		println("See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files.")
		return
	}

	genConf := confFile.Section("General")
	switch common.StripQuotes(genConf["LogLevel"]) {
	case "Debug":
		log.SetLevel(log.DebugLevel)
		log.Warn("Please note the debug level logs may contain sensitive information!")
	case "Info":
		log.SetLevel(log.InfoLevel)
	case "Warn":
		log.SetLevel(log.WarnLevel)
	case "Error":
		log.SetLevel(log.ErrorLevel)
	case "Fatal":
		log.SetLevel(log.FatalLevel)
	case "Panic":
		log.SetLevel(log.PanicLevel)
	default:
		log.SetLevel(log.InfoLevel)
	}

	lf := common.StripQuotes(genConf["LogFile"])
	if lf != "" {
		hook, err := cracklog.NewFileHook(lf)
		if err != nil {
			println("ERROR: Unable to open log file: " + err.Error())
		} else {
			log.AddHook(hook)
		}
	}

	var statefile string
	statefile = common.StripQuotes(genConf["StateFile"])

	var updatetime int
	var resourcetimeout int
	utconf := common.StripQuotes(genConf["UpdateTime"])
	if utconf != "" {
		var err error
		updatetime, err = strconv.Atoi(utconf)
		if err != nil {
			log.WithField("error", err.Error()).Error("Unable to parse update time in config file.")
			updatetime = 30
		}
	} else {
		updatetime = 30
	}
	restimeconf := common.StripQuotes(genConf["ResourceTimeout"])
	if restimeconf != "" {
		var err error
		resourcetimeout, err = strconv.Atoi(restimeconf)
		if err != nil {
			log.WithField("error", err.Error()).Error("Unable to parse resource timeout in config file.")
			resourcetimeout = 5
		}
	} else {
		resourcetimeout = 5
	}

	log.WithFields(log.Fields{
		"ip":   *runIP,
		"port": *runPort,
	}).Info("Starting queue server up.")

	// Get the Authentication configuration
	confAuth := confFile.Section("Authentication")
	if confAuth == nil {
		println("Error: Authentication configuration is required.")
		println("See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files.")
		return
	}

	_, weberr := os.Stat(*webRoot)
	if weberr != nil {
		println("Error: Public web root '" + *webRoot + "' does not exist.")
		return
	}

	// Check for type of authentication and set conf
	switch confAuth["type"] {
	case "INI":
		var i INIAuth

		// Get the users
		umap := map[string]string{}

		au := common.StripQuotes(confAuth["adminuser"])
		if au != "" {
			log.Fatal("An administrative user was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		ap := common.StripQuotes(confAuth["adminpass"])
		if ap != "" {
			log.Fatal("An administrative password was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		su := common.StripQuotes(confAuth["standarduser"])
		if su != "" {
			log.Fatal("An standard user was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		sp := common.StripQuotes(confAuth["standarduser"])
		if sp != "" {
			log.Fatal("An standard password was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		ru := common.StripQuotes(confAuth["readonlyuser"])
		if ru != "" {
			log.Fatal("An read only user was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		rp := common.StripQuotes(confAuth["readonlypass"])
		if rp != "" {
			log.Fatal("An read only password was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		umap[au] = ap
		umap[su] = sp
		umap[ru] = rp

		// Setup group mappings
		gmap := map[string]string{}

		gmap[au] = Administrator
		gmap[su] = StandardUser
		gmap[ru] = ReadOnly

		i.Setup(umap, gmap)

		server.Auth = &i

		log.Info("INI authentication setup complete.")
	case "ActiveDirectory":
		var ad ADAuth

		realm := common.StripQuotes(confAuth["realm"])
		if realm != "" {
			log.Fatal("No Active Directory realm was configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		ad.SetRealm(realm)

		gmap := map[string]string{}
		ro := common.StripQuotes(confAuth["ReadOnlyGroup"])
		if ro != "" {
			log.Fatal("A read only group was not provided. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		st := common.StripQuotes(confAuth["StandardGroup"])
		if st != "" {
			log.Fatal("A group for standard access was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}
		admin := common.StripQuotes(confAuth["AdminGroup"])
		if admin != "" {
			log.Fatal("A group for read only access was not configured. See https://github.com/jmmcatee/cracklord/src/wiki/Configuration-Files#queue-auth")
		}

		gmap[ReadOnly] = ro
		gmap[StandardUser] = st
		gmap[Administrator] = admin

		ad.Setup(gmap)

		server.Auth = &ad
		log.WithFields(log.Fields{
			"readonly": ro,
			"standard": st,
			"admin":    admin,
		}).Info("Active directory authentication configured successfully.")
	}

	// Configure the TokenStore
	server.T = NewTokenStore()

	// Configure the Queue
	server.Q = queue.NewQueue(statefile, updatetime, resourcetimeout)

	// Get the CA
	if *caCertPath == "" {
		// Use default path to get the CA certificate
		*caCertPath = CA_CERT_FILE
	}

	caBytes, err := ioutil.ReadFile(*caCertPath)
	if err != nil {
		println("ERROR: " + err.Error())
	}

	caPool := x509.NewCertPool()
	caPool.AppendCertsFromPEM(caBytes)

	// Get the certificates and private key
	if *CertPath == "" || *KeyPath == "" {
		// The private key and/or the certificate were not given so go with defaults
		*KeyPath = QUEUED_KEY_FILE
		*CertPath = QUEUED_CERT_FILE
	}

	tlscert, err := tls.LoadX509KeyPair(*CertPath, *KeyPath)

	// Setup TLS connection
	tlsconfig := &tls.Config{}
	tlsconfig.Certificates = make([]tls.Certificate, 1)
	tlsconfig.Certificates[0] = tlscert
	tlsconfig.RootCAs = caPool
	tlsconfig.ClientCAs = caPool
	tlsconfig.CipherSuites = []uint16{tls.TLS_RSA_WITH_AES_128_CBC_SHA,
		tls.TLS_RSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA,
		tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA,
		tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
		tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256}
	tlsconfig.MinVersion = tls.VersionTLS12
	tlsconfig.SessionTicketsDisabled = true

	// Set the TLS policy for the Queue to communicate with the resources
	server.TLS = tlsconfig

	// Add some nice security stuff
	secureMiddleware := secure.New(secure.Options{
		SSLRedirect:             true,
		FrameDeny:               true,
		CustomFrameOptionsValue: "SAMEORIGIN",
		BrowserXssFilter:        true,
		IsDevelopment:           true,
	})

	// Build the Negroni handler
	n := negroni.New(negroni.NewRecovery(),
		cracklog.NewNegroniLogger(),
		negroni.NewStatic(http.Dir(*webRoot)))

	n.Use(negroni.HandlerFunc(secureMiddleware.HandlerFuncWithNext))
	n.UseHandler(server.Router())
	log.Debug("Negroni handler started.")

	listen, err := tls.Listen("tcp", *runIP+":"+*runPort, tlsconfig)
	if err != nil {
		println("ERROR: Unable to bind to '" + *runIP + ":" + *runPort + "':" + err.Error())
		return
	}

	err = http.Serve(listen, n)
	if err != nil {
		log.Fatal("Unable to start up web server: " + err.Error())
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright © 2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 *
 * 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.
 *
 * @author		Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @copyright 	2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @license 	Apache-2.0
 */

package server

var banner = `
                ////
            ////////////
        ////////////////////
    ////////////////////////////
///////////////     ////////////////
///////////,            ////////////
///////                     ////////
///                             ////

          &//////////////                                        &&&&&&&&&&&&&&@                   &&&&&&&&&&&&&&&&&&&&&               &&&&&&&         &&&&&&&&
       //////////////////////                                @&&&&&&&&&&&&&&&&&&&&&                &&&&&&&&&&&&&&&&&&&&&&&             &&&&&&&&       @&&&&&&&
     //////////////////////////                            &&&&&&&&&&&&&&&&&&&&&&&&&&              &&&&&&&&&&&&&&&&&&&&&&&&&            &&&&&&&&     @&&&&&&&
   ////////////      ////////////                        &&&&&&&&&&&&      &&&&&&&&&&&&             &&&&&&&&        &&&&&&&&&            &&&&&&&&   &&&&&&&&
  /////////              /////////                      &&&&&&&&&              &&&&&&&&&             &&&&&&&&         &&&&&&&             &&&&&&&& &&&&&&&&
 ////////                  ////////                    &&&&&&&&                  &&&&&&&&             &&&&&&&&@        &&&&&&&              &&&&&&&&&&&&&&
////////                    ////////                  &&&&&&&&                    @&&&&&&&             &&&&&&&&&      &&&&&&&                &&&&&&&&&&&&
///////                      ///////                  &&&&&&&                      &&&&&&&               &&&&&&&&   &&&&&&&&&                 &&&&&&&&&&
///////                      ///////                  &&&&&&&                       &&&&&&&        &&&&&&&&&&&&&&&&&&&&&&&&&                   &&&&&&&&
///////                      ///////                  &&&&&&&                      &&&&&&&&        &&&&&&&&&&&&&&&&&&&&&&&                    &&&&&&&&
///////                      ///////                  &&&&&&&                      &&&&&&&         &&&&&&&&&&&&&&&&&&&&                      &&&&&&&&
////////                    ////////                   &&&&&&&                    &&&&&&&&         &&&&&&&    &&&&&&&&                      &&&&&&&&
 ////////                  ////////                    &&&&&&&&                  &&&&&&&&          &&&&&&&     &&&&&&&&                    &&&&&&&&
  /////////              /////////                      &&&&&&&&&              &&&&&&&&&           &&&&&&&      &&&&&&&&                  &&&&&&&&
   ////////////%    (////////////                         &&&&&&&&&&&&    @&&&&&&&&&&&&            &&&&&&&       &&&&&&&&&               &&&&&&&&
     //////////////////////////                            &&&&&&&&&&&&&&&&&&&&&&&&&&              &&&&&&&         &&&&&&&&             &&&&&&&&
       /////////////////////                                  &&&&&&&&&&&&&&&&&&&&&                &&&&&&&          &&&&&&&&           &&&&&&&&
           //////////////                                        &&&&&&&&&&&&&&                    &&&&&&&           &&&&&&&&         &&&&&&&&

////                            ////
////////                    ////////                  Take security seriously and subscribe to the ORY newsletter. Stay on top of new patches and security
////////////            ////////////                  insights.
////////////////    ////////////////
    ////////////////////////////                      >> Subscribe now: http://eepurl.com/di390P <<
        ////////////////////
            //////////// `
<commit_msg>cmd/server: Shortens long banner message<commit_after>/*
 * Copyright © 2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 *
 * 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.
 *
 * @author		Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @copyright 	2015-2018 Aeneas Rekkas <aeneas+oss@aeneas.io>
 * @license 	Apache-2.0
 */

package server

var banner = `Thank you for using ORY Hydra!

Take security seriously and subscribe to the ORY newsletter. Stay on top of new patches and security insights.

>> Subscribe now: http://eepurl.com/di390P <<    `
<|endoftext|>