<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/dustin/go-humanize"
	"github.com/imdario/mergo"
	"gopkg.in/alecthomas/kingpin.v2"
	yamlv2 "gopkg.in/yaml.v2"

	"ci.guzzler.io/guzzler/corcel/cmd"
	"ci.guzzler.io/guzzler/corcel/config"
	"ci.guzzler.io/guzzler/corcel/core"
	"ci.guzzler.io/guzzler/corcel/errormanager"
	"ci.guzzler.io/guzzler/corcel/infrastructure/http"
	"ci.guzzler.io/guzzler/corcel/infrastructure/inproc"
	"ci.guzzler.io/guzzler/corcel/logger"
	"ci.guzzler.io/guzzler/corcel/serialisation/yaml"
	"ci.guzzler.io/guzzler/corcel/statistics"
)

var (
	applicationVersion = "v0.1.4-alpha"
)

func check(err error) {
	if err != nil {
		errormanager.Log(err)
	}
}

//GenerateExecutionOutput ...
func GenerateExecutionOutput(file string, output statistics.AggregatorSnapShot) {
	outputPath, err := filepath.Abs(file)
	check(err)
	yamlOutput, err := yamlv2.Marshal(&output)
	check(err)
	err = ioutil.WriteFile(outputPath, yamlOutput, 0644)
	check(err)
}

//AddExecutionToHistory ...
func AddExecutionToHistory(file string, output statistics.AggregatorSnapShot) {

	var summary statistics.AggregatorSnapShot

	outputPath, err := filepath.Abs(file)
	check(err)

	if _, err = os.Stat(outputPath); os.IsNotExist(err) {
		summary = *statistics.NewAggregatorSnapShot()
	} else {
		data, dataErr := ioutil.ReadFile(outputPath)
		if dataErr != nil {
			panic(dataErr)
		}
		yamlErr := yamlv2.Unmarshal(data, &summary)
		if yamlErr != nil {
			panic(yamlErr)
		}
	}
	summary.Update(output)

	yamlOutput, err := yamlv2.Marshal(&summary)
	check(err)
	err = ioutil.WriteFile(outputPath, yamlOutput, 0644)
	check(err)
}

// ServerCommand ...
type ServerCommand struct {
	Port     int
	registry Registry
}

func (cmd *ServerCommand) run(c *kingpin.ParseContext) error {
	// have access to cmd.registry
	//Start HTTP Server
	// construct HTTP Host
	// Start HTTP Host from cmd options
	return nil
}

func configureServerCommand(app *kingpin.Application, registry *Registry) {
	c := &ServerCommand{
		registry: registry,
	}
	server := app.Command("server", "Start HTTP server").Action(c.run)
	server.Flag("port", "Port").Default("54332").IntVar(&c.Port)
}

func main() {
	args := os.Args[1:]

	logger.Initialise()

	//define registry

	configuration := &config.Configuration{}

	kingpin.UsageTemplate(kingpin.CompactUsageTemplate).Version(applicationVersion).Author("Andrew Rea").Author("James Allen")
	kingpin.CommandLine.Help = "An example implementation of curl."
	app := kingpin.New("corcel", "")
	app.HelpFlag.Short('h')
	app.UsageTemplate(kingpin.LongHelpTemplate)

	configureServerCommand(app, registry)

	run := app.Command("run", "Execute performance test thing")
	run.Arg("file", "Corcel file contains URLs or an ExecutionPlan (see the --plan argument)").Required().StringVar(&configuration.FilePath)
	run.Flag("summary", "Output summary to STDOUT").BoolVar(&configuration.Summary)
	run.Flag("duration", "The duration of the run e.g. 10s 10m 10h etc... valid values are  ms, s, m, h").Default("0s").DurationVar(&configuration.Duration)
	run.Flag("wait-time", "Time to wait between each execution").Default("0s").DurationVar(&configuration.WaitTime)
	run.Flag("workers", "The number of workers to execute the requests").IntVar(&configuration.Workers)
	run.Flag("random", "Select the url at random for each execution").BoolVar(&configuration.Random)
	run.Flag("plan", "Indicate that the corcel file is an ExecutionPlan").BoolVar(&configuration.Plan)
	run.Flag("verbose", "verbosity").Short('v').Action(config.Counter).Bool()
	run.Flag("progress", "Progress reporter").EnumVar(&configuration.Progress, "bar", "logo", "none")

	var err error
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
	case run.FullCommand():
		configuration, err = config.CmdConfig(configuration)
		if err != nil {
			app.Usage(args)
			os.Exit(1)
		}
		//log.SetLevel(logLevel)

		pwd, err := config.PwdConfig()
		if err != nil {
			app.Usage(args)
			os.Exit(1)
		}
		usr, err := config.UserDirConfig()
		if err != nil {
			app.Usage(args)
			os.Exit(1)
		}

		defaults := config.DefaultConfig()
		eachConfig := []interface{}{configuration, pwd, usr, &defaults}
		for _, item := range eachConfig {
			if err := mergo.Merge(configuration, item); err != nil {
				app.Usage(args)
				os.Exit(1)
			}
		}
		config.SetLogLevel(configuration, eachConfig)
		//log.WithFields(log.Fields{"config": config}).Info("Configuration")
		if err != nil {
			app.Usage(args)
			os.Exit(1)
		}

		logger.ConfigureLogging(configuration)

		//TODO: This is not as efficient as it could be for example:
		//Ideally we would only add the HTTP result processor IF an HTTP Action was used
		//Currently every result processor needs to be added.
		//TODO add a ScanForActions .ScanForAssertions .ScanForProcessors .ScanForExtractors .ScanForContexts
		registry := core.CreateRegistry().
			AddActionParser(inproc.YamlDummyActionParser{}).
			AddActionParser(http.YamlHTTPRequestParser{}).
			AddAssertionParser(yaml.ExactAssertionParser{}).
			AddAssertionParser(yaml.NotEqualAssertionParser{}).
			AddAssertionParser(yaml.EmptyAssertionParser{}).
			AddAssertionParser(yaml.NotEmptyAssertionParser{}).
			AddAssertionParser(yaml.GreaterThanAssertionParser{}).
			AddAssertionParser(yaml.GreaterThanOrEqualAssertionParser{}).
			AddAssertionParser(yaml.LessThanAssertionParser{}).
			AddAssertionParser(yaml.LessThanOrEqualAssertionParser{}).
			AddResultProcessor(http.NewHTTPExecutionResultProcessor()).
			AddResultProcessor(inproc.NewGeneralExecutionResultProcessor())

		_, err = filepath.Abs(configuration.FilePath)
		check(err)

		host := cmd.NewConsoleHost(configuration, registry)
		id, _ := host.Control.Start(configuration) //will this block?
		output := host.Control.Stop(id)

		//TODO these should probably be pushed behind the host.Control.Stop afterall the host is a cmd host
		GenerateExecutionOutput("./output.yml", output)

		AddExecutionToHistory("./history.yml", output)

		if configuration.Summary {
			OutputSummary(output)
		}
	}

}

//OutputSummary ...
func OutputSummary(snapshot statistics.AggregatorSnapShot) {
	summary := statistics.CreateSummary(snapshot)

	top(os.Stdout)
	line(os.Stdout, "Running Time", summary.RunningTime)
	line(os.Stdout, "Throughput", fmt.Sprintf("%-.0f req/s", summary.Throughput))
	line(os.Stdout, "Total Requests", fmt.Sprintf("%-.0f", summary.TotalRequests))
	line(os.Stdout, "Number of Errors", fmt.Sprintf("%-.0f", summary.TotalErrors))
	line(os.Stdout, "Availability", fmt.Sprintf("%-.4f%%", summary.Availability))
	line(os.Stdout, "Bytes Sent", fmt.Sprintf("%v", humanize.Bytes(uint64(summary.TotalBytesSent))))
	line(os.Stdout, "Bytes Received", fmt.Sprintf("%v", humanize.Bytes(uint64(summary.TotalBytesReceived))))
	line(os.Stdout, "Mean Response Time", fmt.Sprintf("%.4f ms", summary.MeanResponseTime))
	line(os.Stdout, "Min Response Time", fmt.Sprintf("%.4f ms", summary.MinResponseTime))
	line(os.Stdout, "Max Response Time", fmt.Sprintf("%.4f ms", summary.MaxResponseTime))
	tail(os.Stdout)
}

func top(writer io.Writer) {
	fmt.Fprintln(writer, "╔═══════════════════════════════════════════════════════════════════╗")
	fmt.Fprintln(writer, "║                           Summary                                 ║")
	fmt.Fprintln(writer, "╠═══════════════════════════════════════════════════════════════════╣")
}

func tail(writer io.Writer) {
	fmt.Fprintln(writer, "╚═══════════════════════════════════════════════════════════════════╝")
}

func line(writer io.Writer, label string, value string) {
	fmt.Fprintf(writer, "║ %20s: %-43s ║\n", label, value)
}
<commit_msg>Moving behaviour into the RunCommand.run function<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/dustin/go-humanize"
	"github.com/imdario/mergo"
	"gopkg.in/alecthomas/kingpin.v2"
	yamlv2 "gopkg.in/yaml.v2"

	"ci.guzzler.io/guzzler/corcel/cmd"
	"ci.guzzler.io/guzzler/corcel/config"
	"ci.guzzler.io/guzzler/corcel/core"
	"ci.guzzler.io/guzzler/corcel/errormanager"
	"ci.guzzler.io/guzzler/corcel/infrastructure/http"
	"ci.guzzler.io/guzzler/corcel/infrastructure/inproc"
	"ci.guzzler.io/guzzler/corcel/logger"
	"ci.guzzler.io/guzzler/corcel/serialisation/yaml"
	"ci.guzzler.io/guzzler/corcel/statistics"
)

var (
	applicationVersion = "v0.1.4-alpha"
)

func check(err error) {
	if err != nil {
		errormanager.Log(err)
	}
}

//GenerateExecutionOutput ...
func GenerateExecutionOutput(file string, output statistics.AggregatorSnapShot) {
	outputPath, err := filepath.Abs(file)
	check(err)
	yamlOutput, err := yamlv2.Marshal(&output)
	check(err)
	err = ioutil.WriteFile(outputPath, yamlOutput, 0644)
	check(err)
}

//AddExecutionToHistory ...
func AddExecutionToHistory(file string, output statistics.AggregatorSnapShot) {

	var summary statistics.AggregatorSnapShot

	outputPath, err := filepath.Abs(file)
	check(err)

	if _, err = os.Stat(outputPath); os.IsNotExist(err) {
		summary = *statistics.NewAggregatorSnapShot()
	} else {
		data, dataErr := ioutil.ReadFile(outputPath)
		if dataErr != nil {
			panic(dataErr)
		}
		yamlErr := yamlv2.Unmarshal(data, &summary)
		if yamlErr != nil {
			panic(yamlErr)
		}
	}
	summary.Update(output)

	yamlOutput, err := yamlv2.Marshal(&summary)
	check(err)
	err = ioutil.WriteFile(outputPath, yamlOutput, 0644)
	check(err)
}

// ServerCommand ...
type ServerCommand struct {
	Port     int
	registry *core.Registry
}

func (instance *ServerCommand) run(c *kingpin.ParseContext) error {
	// have access to c.registry
	//Start HTTP Server
	// construct HTTP Host
	// Start HTTP Host from cmd options
	fmt.Printf("Would now be starting the HTTP server on %v\n", instance.Port)
	return nil
}

func configureServerCommand(app *kingpin.Application, registry *core.Registry) {
	c := &ServerCommand{
		registry: registry,
	}
	server := app.Command("server", "Start HTTP server").Action(c.run)
	server.Flag("port", "Port").Default("54332").IntVar(&c.Port)
}

// RunCommand ...
type RunCommand struct {
	Config   *config.Configuration
	registry *core.Registry
}

func (instance *RunCommand) run(c *kingpin.ParseContext) error {
	configuration, err := config.CmdConfig(instance.Config)
	if err != nil {
		return err
	}
	//log.SetLevel(logLevel)

	pwd, err := config.PwdConfig()
	if err != nil {
		return err
	}
	usr, err := config.UserDirConfig()
	if err != nil {
		return err
	}

	defaults := config.DefaultConfig()
	eachConfig := []interface{}{configuration, pwd, usr, &defaults}
	for _, item := range eachConfig {
		if err := mergo.Merge(configuration, item); err != nil {
			return err
		}
	}
	config.SetLogLevel(configuration, eachConfig)
	//log.WithFields(log.Fields{"config": config}).Info("Configuration")
	logger.ConfigureLogging(configuration)

	_, err = filepath.Abs(configuration.FilePath)
	//TODO Why is this a check rather than return err?
	check(err)

	host := cmd.NewConsoleHost(configuration, *instance.registry)
	id, _ := host.Control.Start(configuration) //will this block?
	output := host.Control.Stop(id)

	//TODO these should probably be pushed behind the host.Control.Stop afterall the host is a cmd host
	GenerateExecutionOutput("./output.yml", output)

	AddExecutionToHistory("./history.yml", output)

	if configuration.Summary {
		OutputSummary(output)
	}
	return nil
}

func configureRunCommand(app *kingpin.Application, registry *core.Registry) {
	configuration := &config.Configuration{}
	c := &RunCommand{
		Config:   configuration,
		registry: registry,
	}
	run := app.Command("run", "Execute performance test thing").Action(c.run)
	run.Arg("file", "Corcel file contains URLs or an ExecutionPlan (see the --plan argument)").Required().StringVar(&configuration.FilePath)
	run.Flag("summary", "Output summary to STDOUT").BoolVar(&configuration.Summary)
	run.Flag("duration", "The duration of the run e.g. 10s 10m 10h etc... valid values are  ms, s, m, h").Default("0s").DurationVar(&configuration.Duration)
	run.Flag("wait-time", "Time to wait between each execution").Default("0s").DurationVar(&configuration.WaitTime)
	run.Flag("workers", "The number of workers to execute the requests").IntVar(&configuration.Workers)
	run.Flag("random", "Select the url at random for each execution").BoolVar(&configuration.Random)
	run.Flag("plan", "Indicate that the corcel file is an ExecutionPlan").BoolVar(&configuration.Plan)
	run.Flag("verbose", "verbosity").Short('v').Action(config.Counter).Bool()
	run.Flag("progress", "Progress reporter").EnumVar(&configuration.Progress, "bar", "logo", "none")
}

func main() {
	logger.Initialise()

	//TODO: This is not as efficient as it could be for example:
	//Ideally we would only add the HTTP result processor IF an HTTP Action was used
	//Currently every result processor needs to be added.
	//TODO add a ScanForActions .ScanForAssertions .ScanForProcessors .ScanForExtractors .ScanForContexts
	registry := core.CreateRegistry().
		AddActionParser(inproc.YamlDummyActionParser{}).
		AddActionParser(http.YamlHTTPRequestParser{}).
		AddAssertionParser(yaml.ExactAssertionParser{}).
		AddAssertionParser(yaml.NotEqualAssertionParser{}).
		AddAssertionParser(yaml.EmptyAssertionParser{}).
		AddAssertionParser(yaml.NotEmptyAssertionParser{}).
		AddAssertionParser(yaml.GreaterThanAssertionParser{}).
		AddAssertionParser(yaml.GreaterThanOrEqualAssertionParser{}).
		AddAssertionParser(yaml.LessThanAssertionParser{}).
		AddAssertionParser(yaml.LessThanOrEqualAssertionParser{}).
		AddResultProcessor(http.NewHTTPExecutionResultProcessor()).
		AddResultProcessor(inproc.NewGeneralExecutionResultProcessor())

	kingpin.UsageTemplate(kingpin.CompactUsageTemplate).Version(applicationVersion).Author("Andrew Rea").Author("James Allen")
	kingpin.CommandLine.Help = "An example implementation of curl."
	app := kingpin.New("corcel", "")
	app.HelpFlag.Short('h')
	app.UsageTemplate(kingpin.LongHelpTemplate)

	configureRunCommand(app, &registry)
	configureServerCommand(app, &registry)

	kingpin.MustParse(app.Parse(os.Args[1:]))
}

//OutputSummary ...
func OutputSummary(snapshot statistics.AggregatorSnapShot) {
	summary := statistics.CreateSummary(snapshot)

	top(os.Stdout)
	line(os.Stdout, "Running Time", summary.RunningTime)
	line(os.Stdout, "Throughput", fmt.Sprintf("%-.0f req/s", summary.Throughput))
	line(os.Stdout, "Total Requests", fmt.Sprintf("%-.0f", summary.TotalRequests))
	line(os.Stdout, "Number of Errors", fmt.Sprintf("%-.0f", summary.TotalErrors))
	line(os.Stdout, "Availability", fmt.Sprintf("%-.4f%%", summary.Availability))
	line(os.Stdout, "Bytes Sent", fmt.Sprintf("%v", humanize.Bytes(uint64(summary.TotalBytesSent))))
	line(os.Stdout, "Bytes Received", fmt.Sprintf("%v", humanize.Bytes(uint64(summary.TotalBytesReceived))))
	line(os.Stdout, "Mean Response Time", fmt.Sprintf("%.4f ms", summary.MeanResponseTime))
	line(os.Stdout, "Min Response Time", fmt.Sprintf("%.4f ms", summary.MinResponseTime))
	line(os.Stdout, "Max Response Time", fmt.Sprintf("%.4f ms", summary.MaxResponseTime))
	tail(os.Stdout)
}

func top(writer io.Writer) {
	fmt.Fprintln(writer, "╔═══════════════════════════════════════════════════════════════════╗")
	fmt.Fprintln(writer, "║                           Summary                                 ║")
	fmt.Fprintln(writer, "╠═══════════════════════════════════════════════════════════════════╣")
}

func tail(writer io.Writer) {
	fmt.Fprintln(writer, "╚═══════════════════════════════════════════════════════════════════╝")
}

func line(writer io.Writer, label string, value string) {
	fmt.Fprintf(writer, "║ %20s: %-43s ║\n", label, value)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/boltdb/bolt"
	homedir "github.com/mitchellh/go-homedir"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
	versionNumber = "0.7"

	verbLevel int

	db     *bolt.DB
	dbFile string

	request  Request
	response Response
)

func init() {
	dir, err := homedir.Dir()
	if err != nil {
		panic(err)
	}
	dbFile = fmt.Sprintf("%s/%s", dir, ".rest.db")

	kingpin.Version(versionNumber)
	kingpin.Flag("verbose", "Verbose mode").Short('v').CounterVar(&verbLevel)
	kingpin.UsageTemplate(usageTemplate)
	log.SetFlags(0)

	addAliases()
}

func main() {
	command := kingpin.Parse()

	var err error
	db, err = bolt.Open(dbFile, 0600, nil)
	if err != nil {
		log.Println(err)
		os.Exit(1)
	}
	defer db.Close()
	if luaState != nil {
		defer luaState.Close()
	}

	switch command {
	case "version":
		displayVersion()
	case "service init":
		if err := initService(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service remove":
		if err := removeService(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service list":
		if err := listServices(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service set":
		if err := setValue(); err != nil {
			log.Println(err)
			os.Exit(1)
		}

	case "service unset":
		if err := unsetValue(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service use":
		if err := useService(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service config":
		if err := displayConfig(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service alias":
		if err := addAlias(); err != nil {
			log.Println(err)
			os.Exit(1)
		}

	case "get", "post", "put", "delete", "patch", "options", "head":
		Do(command)

	default:
		Perform(command)
	}
}

// Do perform the request, display the response, and exit.
func Do(command string) {
	request.Method = command
	request.verbose = verbLevel

	resp, err := request.Perform()
	if err != nil {
		log.Println("error making request:", err)
		os.Exit(1)
	}

	response.verbose = verbLevel
	if err := response.Load(resp, request.Settings); err != nil {
		log.Println("error displaying result:", err)
		os.Exit(1)
	}

	fmt.Println(response)

	os.Exit(response.ExitCode())
}

func useService() error {
	err := db.Update(func(tx *bolt.Tx) error {
		serviceBucket := tx.Bucket([]byte("services"))
		if serviceBucket == nil {
			return ErrInitDB
		}

		// Check that the service exists
		if b := serviceBucket.Bucket([]byte(request.Service)); b == nil {
			return ErrNoService{Name: request.Service}
		}

		info := tx.Bucket([]byte("info"))
		if info == nil {
			// If we get here then the db is malformed, examine careully
			// how it happened.
			return ErrNoInfoBucket
		}

		if err := info.Put([]byte("current"), []byte(request.Service)); err != nil {
			return err
		}

		return nil
	})

	return err
}

func displayVersion() {
	var dbVersion string
	noDB := errors.New("no db")
	db.View(func(tx *bolt.Tx) error {
		info := tx.Bucket([]byte("info"))
		if info == nil {
			return noDB
		}

		v := info.Get([]byte("version"))
		if v == nil {
			return noDB
		}

		dbVersion = string(v)
		return nil
	})

	if dbVersion == "" || versionNumber == dbVersion {
		fmt.Println(versionNumber)
	} else {
		fmt.Printf("%s, db: %s, version don't match, the database will have to be updated", versionNumber, dbVersion)
	}

}

func usedFlag(b *bool) func(*kingpin.ParseContext) error {
	return func(*kingpin.ParseContext) error {
		*b = true
		return nil
	}
}

func paramReplacer(parameters map[string]string) *strings.Replacer {
	rep := make([]string, 0, len(parameters))
	for key, value := range parameters {
		rep = append(rep, ":"+key)
		rep = append(rep, value)
	}
	return strings.NewReplacer(rep...)
}
<commit_msg>Fix not showing response on error<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/boltdb/bolt"
	homedir "github.com/mitchellh/go-homedir"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

var (
	versionNumber = "0.7"

	verbLevel int

	db     *bolt.DB
	dbFile string

	request  Request
	response Response
)

func init() {
	dir, err := homedir.Dir()
	if err != nil {
		panic(err)
	}
	dbFile = fmt.Sprintf("%s/%s", dir, ".rest.db")

	kingpin.Version(versionNumber)
	kingpin.Flag("verbose", "Verbose mode").Short('v').CounterVar(&verbLevel)
	kingpin.UsageTemplate(usageTemplate)
	log.SetFlags(0)

	addAliases()
}

func main() {
	command := kingpin.Parse()

	var err error
	db, err = bolt.Open(dbFile, 0600, nil)
	if err != nil {
		log.Println(err)
		os.Exit(1)
	}
	defer db.Close()
	if luaState != nil {
		defer luaState.Close()
	}

	switch command {
	case "version":
		displayVersion()
	case "service init":
		if err := initService(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service remove":
		if err := removeService(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service list":
		if err := listServices(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service set":
		if err := setValue(); err != nil {
			log.Println(err)
			os.Exit(1)
		}

	case "service unset":
		if err := unsetValue(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service use":
		if err := useService(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service config":
		if err := displayConfig(); err != nil {
			log.Println(err)
			os.Exit(1)
		}
	case "service alias":
		if err := addAlias(); err != nil {
			log.Println(err)
			os.Exit(1)
		}

	case "get", "post", "put", "delete", "patch", "options", "head":
		Do(command)

	default:
		Perform(command)
	}
}

// Do perform the request, display the response, and exit.
func Do(command string) {
	request.Method = command
	request.verbose = verbLevel

	resp, err := request.Perform()
	if err != nil {
		log.Println("error making request:", err)
		if resp != nil {
			log.Println(resp.Body)
			resp.Body.Close()
		}
		os.Exit(1)
	}

	response.verbose = verbLevel
	if err := response.Load(resp, request.Settings); err != nil {
		log.Println("error displaying result:", err)
		os.Exit(1)
	}

	fmt.Println(response)

	os.Exit(response.ExitCode())
}

func useService() error {
	err := db.Update(func(tx *bolt.Tx) error {
		serviceBucket := tx.Bucket([]byte("services"))
		if serviceBucket == nil {
			return ErrInitDB
		}

		// Check that the service exists
		if b := serviceBucket.Bucket([]byte(request.Service)); b == nil {
			return ErrNoService{Name: request.Service}
		}

		info := tx.Bucket([]byte("info"))
		if info == nil {
			// If we get here then the db is malformed, examine careully
			// how it happened.
			return ErrNoInfoBucket
		}

		if err := info.Put([]byte("current"), []byte(request.Service)); err != nil {
			return err
		}

		return nil
	})

	return err
}

func displayVersion() {
	var dbVersion string
	noDB := errors.New("no db")
	db.View(func(tx *bolt.Tx) error {
		info := tx.Bucket([]byte("info"))
		if info == nil {
			return noDB
		}

		v := info.Get([]byte("version"))
		if v == nil {
			return noDB
		}

		dbVersion = string(v)
		return nil
	})

	if dbVersion == "" || versionNumber == dbVersion {
		fmt.Println(versionNumber)
	} else {
		fmt.Printf("%s, db: %s, version don't match, the database will have to be updated", versionNumber, dbVersion)
	}

}

func usedFlag(b *bool) func(*kingpin.ParseContext) error {
	return func(*kingpin.ParseContext) error {
		*b = true
		return nil
	}
}

func paramReplacer(parameters map[string]string) *strings.Replacer {
	rep := make([]string, 0, len(parameters))
	for key, value := range parameters {
		rep = append(rep, ":"+key)
		rep = append(rep, value)
	}
	return strings.NewReplacer(rep...)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
)

type Schema struct {
	Swagger             string      `json:"swagger"`
	Info                interface{} `json:"info"`
	Host                string      `json:"host,omitempty"`
	BasePath            string      `json:"basePath,omitempty"`
	Schemes             []string    `json:"schemes,omitempty"`
	Consumes            []string    `json:"consumes,omitempty"`
	Produces            []string    `json:"produces,omitempty"`
	Paths               interface{} `json:"paths"`
	Definitions         interface{} `json:"definitions,omitempty"`
	Parameters          interface{} `json:"parameters,omitempty"`
	Responses           interface{} `json:"responses,omitempty"`
	SecurityDefinitions interface{} `json:"securityDefinitions,omitempty"`
	Security            interface{} `json:"security,omitempty"`
	Tags                interface{} `json:"tags,omitempty"`
	ExternalDocs        interface{} `json:"externalDocs,omitempty"`
}

func main() {
	schema, err := NewSchema("http://petstore.swagger.io/v2/swagger.json")
	if err != nil {
		return
	}
	fmt.Printf("%#v\r\n", schema)
}

func NewSchema(url string) (*Schema, error) {
	schema, err := load(url)
	if err != nil {
		return nil, err
	}
	return schema, nil
}

func load(url string) (*Schema, error) {
	var schema Schema
	response, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	jsonString, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}
	json.Unmarshal(jsonString, &schema)
	if err != nil {
		return nil, err
	}
	return &schema, nil
}
<commit_msg>Remove all model definitions from main.go<commit_after>package main

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

	"github.com/rightlag/go-swagger-object-model"
)

func main() {
	schema, err := NewSchema("http://petstore.swagger.io/v2/swagger.json")
	if err != nil {
		return
	}
	// fmt.Printf("%#v\r\n", schema.Paths["/pet"].Put.Responses["400"])
	fmt.Printf("%#v\r\n", schema)
}

func NewSchema(url string) (*models.Schema, error) {
	schema, err := load(url)
	if err != nil {
		return nil, err
	}
	return schema, nil
}

func load(url string) (*models.Schema, error) {
	var schema models.Schema
	response, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	jsonString, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}
	json.Unmarshal(jsonString, &schema)
	if err != nil {
		return nil, err
	}
	return &schema, nil
}
<|endoftext|>


<commit_before>package tokenizer

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	testContent = `#!/usr/bin/ruby

#!/usr/bin/env node

aaa

#!/usr/bin/env A=B foo=bar awk -f

#!python

func Tokenize(content []byte) []string {
	splitted := bytes.Fields(content)
	tokens := /* make([]string, 0, len(splitted))
	no comment -- comment
	for _, tokenByte := range splitted {
		token64 := base64.StdEncoding.EncodeToString(tokenByte)
		tokens = append(tokens, token64)
		notcatchasanumber3.5
	}*/
othercode
	/* testing multiple

		multiline comments*/

ppp no comment # comment

"literal1"

abb (tokenByte, 0xAF02) | ,3.2L

'literal2' notcatchasanumber3.5

	5 += number * anotherNumber
	if isTrue && isToo {
		0b00001000 >> 1
	}

	return tokens

oneBool = 3 <= 2
varBool = 3<=2>

#ifndef
#i'm not a comment if the single line comment symbol is not followed by a white

  PyErr_SetString(PyExc_RuntimeError, "Relative import is not supported for Python <=2.4.");

<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-strict.dtd">

        <title id="hola" class="">This is a XHTML sample file</title>

            Just a simple XHTML test page.

`
)

var (
	tokensFromTestContent = []string{"SHEBANG#!ruby", "SHEBANG#!node", "SHEBANG#!awk", "<!DOCTYPE>", "PUBLIC", "W3C", "DTD", "XHTML", "1", "0",
		"Strict", "EN", "http", "www", "w3", "org", "TR", "xhtml1", "DTD", "xhtml1", "strict", "dtd", "", "", "<title>", "class=",
		"</title>", "", "", "", "
", "",
		"", "
", "", "", "(", "[", "]", ")", "[", "]", "{", "(", ")", "(", ")", "{", "}", "(", ")", ";", "{", ";",
		"}", "]", "]", "#", "/usr/bin/ruby", "#", "/usr/bin/env", "node", "aaa", "#", "/usr/bin/env", "A", "B", "foo", "bar", "awk", "f", "#",
		"python", "func", "Tokenize", "content", "byte", "string", "splitted", "bytes.Fields", "content", "tokens", "othercode", "ppp", "no",
		"comment", "abb", "tokenByte", "notcatchasanumber", "number", "*", "anotherNumber", "if", "isTrue", "isToo", "b", "return", "tokens",
		"oneBool", "varBool", "#ifndef", "#i", "m", "not", "a", "comment", "if", "the", "single", "line", "comment", "symbol", "is", "not",
		"followed", "by", "a", "white", "PyErr_SetString", "PyExc_RuntimeError", "html", "PUBLIC", "xmlns", "id", "class", "This", "is", "a",
		"XHTML", "sample", "file", "type", "#example", "background", "color", "yellow", "id", "Just", "a", "simple", "XHTML", "test", "page.",
		"-", "|", "+", "&&", "<", "<", "-", "!", "!", "!", "=", "=", "!", ":", "=", ":", "=", ",", ",", "=", ">", ">", "=", "=", "=", "=", ">",
		"'", ",", ">", "=", ">", "=", "=", ">", "=", ">", ":", ">", "=", ">"}

	Tests = []struct {
		name     string
		content  []byte
		expected []string
	}{
		{name: "content", content: []byte(testContent), expected: tokensFromTestContent},
	}
)

func TestTokenize(t *testing.T) {
	for _, test := range Tests {
		t.Run(test.name, func(t *testing.T) {
			before := string(test.content)
			tokens := Tokenize(test.content)
			after := string(test.content)
			require.Equal(t, before, after, "the input slice was modified")
			require.Equal(t, len(test.expected), len(tokens), fmt.Sprintf("token' slice length = %v, want %v", len(test.expected), len(tokens)))
			for i, expectedToken := range test.expected {
				assert.Equal(t, expectedToken, tokens[i], fmt.Sprintf("token = %v, want %v", tokens[i], expectedToken))
			}
		})
	}
}

func BenchmarkTokenizer_BaselineCopy(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i  ByteLimit {
				test.content = test.content[:ByteLimit]
			}
			_ = append([]byte(nil), test.content...)
		}
	}
}

func BenchmarkTokenizer(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		for _, test := range Tests {
			Tokenize(test.content)
		}
	}
}
<commit_msg>test: don't export tokenizer fixtures<commit_after>package tokenizer

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	testContent = `#!/usr/bin/ruby

#!/usr/bin/env node

aaa

#!/usr/bin/env A=B foo=bar awk -f

#!python

func Tokenize(content []byte) []string {
	splitted := bytes.Fields(content)
	tokens := /* make([]string, 0, len(splitted))
	no comment -- comment
	for _, tokenByte := range splitted {
		token64 := base64.StdEncoding.EncodeToString(tokenByte)
		tokens = append(tokens, token64)
		notcatchasanumber3.5
	}*/
othercode
	/* testing multiple

		multiline comments*/

ppp no comment # comment

"literal1"

abb (tokenByte, 0xAF02) | ,3.2L

'literal2' notcatchasanumber3.5

	5 += number * anotherNumber
	if isTrue && isToo {
		0b00001000 >> 1
	}

	return tokens

oneBool = 3 <= 2
varBool = 3<=2>

#ifndef
#i'm not a comment if the single line comment symbol is not followed by a white

  PyErr_SetString(PyExc_RuntimeError, "Relative import is not supported for Python <=2.4.");

<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN" "http://www.w3.org/TR/xhtml1/DTD/xhtml1-strict.dtd">

        <title id="hola" class="">This is a XHTML sample file</title>

            Just a simple XHTML test page.

`
)

var (
	tokensFromTestContent = []string{"SHEBANG#!ruby", "SHEBANG#!node", "SHEBANG#!awk", "<!DOCTYPE>", "PUBLIC", "W3C", "DTD", "XHTML", "1", "0",
		"Strict", "EN", "http", "www", "w3", "org", "TR", "xhtml1", "DTD", "xhtml1", "strict", "dtd", "", "", "<title>", "class=",
		"</title>", "", "", "", "
", "",
		"", "
", "", "", "(", "[", "]", ")", "[", "]", "{", "(", ")", "(", ")", "{", "}", "(", ")", ";", "{", ";",
		"}", "]", "]", "#", "/usr/bin/ruby", "#", "/usr/bin/env", "node", "aaa", "#", "/usr/bin/env", "A", "B", "foo", "bar", "awk", "f", "#",
		"python", "func", "Tokenize", "content", "byte", "string", "splitted", "bytes.Fields", "content", "tokens", "othercode", "ppp", "no",
		"comment", "abb", "tokenByte", "notcatchasanumber", "number", "*", "anotherNumber", "if", "isTrue", "isToo", "b", "return", "tokens",
		"oneBool", "varBool", "#ifndef", "#i", "m", "not", "a", "comment", "if", "the", "single", "line", "comment", "symbol", "is", "not",
		"followed", "by", "a", "white", "PyErr_SetString", "PyExc_RuntimeError", "html", "PUBLIC", "xmlns", "id", "class", "This", "is", "a",
		"XHTML", "sample", "file", "type", "#example", "background", "color", "yellow", "id", "Just", "a", "simple", "XHTML", "test", "page.",
		"-", "|", "+", "&&", "<", "<", "-", "!", "!", "!", "=", "=", "!", ":", "=", ":", "=", ",", ",", "=", ">", ">", "=", "=", "=", "=", ">",
		"'", ",", ">", "=", ">", "=", "=", ">", "=", ">", ":", ">", "=", ">"}

	tests = []struct {
		name     string
		content  []byte
		expected []string
	}{
		{name: "content", content: []byte(testContent), expected: tokensFromTestContent},
	}
)

func TestTokenize(t *testing.T) {
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			before := string(test.content)
			tokens := Tokenize(test.content)
			after := string(test.content)
			require.Equal(t, before, after, "the input slice was modified")
			require.Equal(t, len(test.expected), len(tokens), fmt.Sprintf("token' slice length = %v, want %v", len(test.expected), len(tokens)))
			for i, expectedToken := range test.expected {
				assert.Equal(t, expectedToken, tokens[i], fmt.Sprintf("token = %v, want %v", tokens[i], expectedToken))
			}
		})
	}
}

func BenchmarkTokenizer_BaselineCopy(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i  ByteLimit {
				test.content = test.content[:ByteLimit]
			}
			_ = append([]byte(nil), test.content...)
		}
	}
}

func BenchmarkTokenizer(b *testing.B) {
	b.ReportAllocs()
	for i := 0; i < b.N; i++ {
		for _, test := range tests {
			Tokenize(test.content)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/amzfans/appstarter/app"
	"github.com/amzfans/appstarter/domain"
	"log"
	"os"
	"os/signal"
)

const (
	stdout_socket_file = "stdout.sock"
	stderr_socket_file = "stderr.sock"
)

func main() {
	cmdAndArgs := os.Args[1:]
	if len(cmdAndArgs) == 0 {
		log.Fatal("Cannot get the runnable commands.")
	}

	cmd := cmdAndArgs[0]
	var args []string
	if len(cmdAndArgs) > 1 {
		args = cmdAndArgs[1:]
	}

	dsServer := domain.NewServer(stdout_socket_file, stderr_socket_file)
	err := dsServer.Start()
	if err != nil {
		panic(fmt.Sprintf("Cannot start the domain socket server as %s.", err.Error()))
	}
	defer dsServer.Stop()

	application := app.NewApp(dsServer, cmd, args...)
	err = application.Start()
	if err != nil {
		panic(fmt.Sprintf("Cannot start the application as %s.", err.Error()))
	}

	killChan := make(chan os.Signal)
	signal.Notify(killChan, os.Kill, os.Interrupt)
	for {
		select {
		case <-dsServer.NeedStop:
			shutdown(dsServer, application, nil)
		case <-application.NeedStop:
			shutdown(dsServer, application, nil)
		case osSig := <-killChan:
			shutdown(dsServer, application, osSig)
		}
	}

}

func shutdown(dsServer *domain.DomainSocketServer, application *app.Application, osSig os.Signal) {
	log.Println("Shutdown the domain socket server and kill the underline application.")
	dsServer.Stop()
	if !application.Cmd.ProcessState.Exited() {
		if osSig != nil {
			application.Cmd.Process.Signal(osSig)
		} else {
			application.Cmd.Process.Kill()
		}
	}
	os.Exit(0)
}
<commit_msg>fix the ProcessState is nil issue.<commit_after>package main

import (
	"fmt"
	"github.com/amzfans/appstarter/app"
	"github.com/amzfans/appstarter/domain"
	"log"
	"os"
	"os/signal"
)

const (
	stdout_socket_file = "stdout.sock"
	stderr_socket_file = "stderr.sock"
)

func main() {
	cmdAndArgs := os.Args[1:]
	if len(cmdAndArgs) == 0 {
		log.Fatal("Cannot get the runnable commands.")
	}

	cmd := cmdAndArgs[0]
	var args []string
	if len(cmdAndArgs) > 1 {
		args = cmdAndArgs[1:]
	}

	dsServer := domain.NewServer(stdout_socket_file, stderr_socket_file)
	err := dsServer.Start()
	if err != nil {
		panic(fmt.Sprintf("Cannot start the domain socket server as %s.", err.Error()))
	}
	defer dsServer.Stop()

	application := app.NewApp(dsServer, cmd, args...)
	err = application.Start()
	if err != nil {
		panic(fmt.Sprintf("Cannot start the application as %s.", err.Error()))
	}

	killChan := make(chan os.Signal)
	signal.Notify(killChan, os.Kill, os.Interrupt)
	for {
		select {
		case <-dsServer.NeedStop:
			shutdown(dsServer, application, nil)
		case <-application.NeedStop:
			shutdown(dsServer, application, nil)
		case osSig := <-killChan:
			shutdown(dsServer, application, osSig)
		}
	}

}

func shutdown(dsServer *domain.DomainSocketServer, application *app.Application, osSig os.Signal) {
	log.Println("Shutdown the domain socket server and kill the underline application.")
	dsServer.Stop()
	if application.Cmd.ProcessState != nil && !application.Cmd.ProcessState.Exited() {
		if osSig != nil {
			application.Cmd.Process.Signal(osSig)
		} else {
			application.Cmd.Process.Kill()
		}
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/vbatts/power_info/linux"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"time"
)

var (
	quiet bool
)

func main() {
	flag.Parse()
	linux.SetQuiet(quiet)

	powers, err := filepath.Glob(linux.PowerSupplyPath + "*")
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	for _, power := range powers {
		info := linux.Info{
			Key:     filepath.Base(power),
			Time:    time.Now().UnixNano(),
			Values:  map[string]string{},
			Load:    linux.GetLoadAvg(),
			Version: linux.GetVersion(),
		}

		files, err := filepath.Glob(power + "/*")
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
			os.Exit(2)
		}
		for _, file := range files {
			basename := filepath.Base(file)
			if linux.IsFile(file) {
				fh, err := os.Open(file)
				if err != nil {
					if !quiet {
						fmt.Fprintf(os.Stderr, "WARN: %s\n", err)
					}
					continue
				}
				b, err := ioutil.ReadAll(fh)
				if err != nil {
					if !quiet {
						fmt.Fprintf(os.Stderr, "WARN: %s\n", err)
					}
					continue
				}
				info.Values[basename] = strings.TrimRight(bytes.NewBuffer(b).String(), " \n")
			} else {
				if !quiet {
					fmt.Fprintf(os.Stderr, "WARN: [%s] does not appear to be a file\n", file)
				}
			}
		}
		buf, err := json.Marshal(info)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
			os.Exit(2)
		}
		fmt.Printf("%s\n", buf)
	}
}

func init() {
	flag.BoolVar(&quiet, "quiet", false, "less output")
}
<commit_msg>reusing the StringFromFile method<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/vbatts/power_info/linux"
	"os"
	"path/filepath"
	"time"
)

var (
	quiet bool
)

func main() {
	flag.Parse()
	linux.SetQuiet(quiet)

	powers, err := filepath.Glob(linux.PowerSupplyPath + "*")
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	for _, power := range powers {
		info := linux.Info{
			Key:     filepath.Base(power),
			Time:    time.Now().UnixNano(),
			Values:  map[string]string{},
			Load:    linux.GetLoadAvg(),
			Version: linux.GetVersion(),
		}

		files, err := filepath.Glob(power + "/*")
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
			os.Exit(2)
		}
		for _, file := range files {
			basename := filepath.Base(file)
			if linux.IsFile(file) {
				str, err := linux.StringFromFile(file)
				if err != nil {
					fmt.Fprintf(os.Stderr, "ERROR: reading file [%s]: %s\n", file, err)
				} else {
					info.Values[basename] = str
				}
			} else {
				if !quiet {
					fmt.Fprintf(os.Stderr, "WARN: [%s] does not appear to be a file\n", file)
				}
			}
		}
		buf, err := json.Marshal(info)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
			os.Exit(2)
		}
		fmt.Printf("%s\n", buf)
	}
}

func init() {
	flag.BoolVar(&quiet, "quiet", false, "less output")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	clog "github.com/cenkalti/log"
	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/console"
	"github.com/cenkalti/rain/internal/jsonutil"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/rainrpc"
	"github.com/cenkalti/rain/torrent"
	"github.com/cenkalti/rain/torrent/resumer/boltdbresumer"
	"github.com/cenkalti/rain/torrent/storage/filestorage"
	"github.com/hokaccha/go-prettyjson"
	"github.com/jroimartin/gocui"
	"github.com/mitchellh/go-homedir"
	"github.com/urfave/cli"
	"gopkg.in/yaml.v2"
)

var (
	app = cli.NewApp()
	clt *rainrpc.Client
	log = logger.New("rain")
)

func main() {
	app.Version = clientversion.Version
	app.Usage = "BitTorrent client"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "config, c",
			Usage: "read config from `FILE`",
		},
		cli.StringFlag{
			Name:  "cpuprofile",
			Usage: "write cpu profile to `FILE`",
		},
		cli.StringFlag{
			Name:  "pprof",
			Usage: "run pprof server on `ADDR`",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "enable debug log",
		},
		cli.StringFlag{
			Name:  "logfile",
			Usage: "write log to `FILE`",
		},
	}
	app.Before = handleBeforeCommand
	app.After = handleAfterCommand
	app.Commands = []cli.Command{
		{
			Name:      "download",
			Usage:     "download torrent or magnet",
			ArgsUsage: "[torrent path or magnet link]",
			Action:    handleDownload,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "dest",
					Usage: "save files under `DIR`",
					Value: ".",
				},
				cli.IntFlag{
					Name:  "port",
					Usage: "peer listen port",
				},
				cli.BoolFlag{
					Name:  "seed",
					Usage: "continue seeding after download finishes",
				},
			},
		},
		{
			Name:  "server",
			Usage: "run rpc server and torrent client",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "config",
					Usage: "path to the config file",
					Value: "~/.rain/config.yaml",
				},
			},
			Action: handleServer,
		},
		{
			Name:  "client",
			Usage: "send rpc request to server",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "url",
					Usage: "URL of RPC server",
					Value: "localhost:" + strconv.Itoa(rainrpc.DefaultServerConfig.Port),
				},
			},
			Before: handleBeforeClient,
			Subcommands: []cli.Command{
				{
					Name:   "list",
					Usage:  "list torrents",
					Action: handleList,
				},
				{
					Name:   "add",
					Usage:  "add torrent or magnet",
					Action: handleAdd,
				},
				{
					Name:   "remove",
					Usage:  "remove torrent",
					Action: handleRemove,
				},
				{
					Name:   "stats",
					Usage:  "get stats of torrent",
					Action: handleStats,
				},
				{
					Name:   "start",
					Usage:  "start torrent",
					Action: handleStart,
				},
				{
					Name:   "stop",
					Usage:  "stop",
					Action: handleStop,
				},
				{
					Name:   "console",
					Usage:  "show client console",
					Action: handleConsole,
				},
			},
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func handleBeforeCommand(c *cli.Context) error {
	cpuprofile := c.GlobalString("cpuprofile")
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
	}
	pprofFlag := c.GlobalString("pprof")
	if pprofFlag != "" {
		go func() {
			log.Notice(http.ListenAndServe(pprofFlag, nil))
		}()
	}
	logFile := c.GlobalString("logfile")
	if logFile != "" {
		f, err := os.Create(logFile)
		if err != nil {
			log.Fatal("could not create log file: ", err)
		}
		logger.SetHandler(clog.NewFileHandler(f))
	}
	if c.GlobalBool("debug") {
		logger.SetLevel(clog.DEBUG)
	}
	return nil
}

func handleAfterCommand(c *cli.Context) error {
	if c.GlobalString("cpuprofile") != "" {
		pprof.StopCPUProfile()
	}
	return nil
}

func handleDownload(c *cli.Context) error {
	path := c.Args().Get(0)
	if path == "" {
		return errors.New("first argument must be a torrent file or magnet link")
	}

	sto, err := filestorage.New(c.String("dest"))
	if err != nil {
		return err
	}
	var t *torrent.Torrent
	if strings.HasPrefix(path, "magnet:") {
		t, err = torrent.NewMagnet(path, c.Int("port"), sto, torrent.DefaultConfig)
	} else {
		f, err2 := os.Open(path) // nolint: gosec
		if err2 != nil {
			return err
		}
		t, err = torrent.New(f, c.Int("port"), sto, torrent.DefaultConfig)
		_ = f.Close()
	}
	if err != nil {
		return err
	}
	defer t.Close()

	res, err := boltdbresumer.Open(t.Name()+"."+t.InfoHash()+".resume", []byte(t.InfoHash()))
	if err != nil {
		return err
	}
	defer func() { _ = res.Close() }()

	err = t.SetResume(res)
	if err != nil {
		return err
	}

	g, err := gocui.NewGui(gocui.OutputNormal)
	if err != nil {
		return err
	}
	defer g.Close()

	g.SetManagerFunc(func(g *gocui.Gui) error {
		x, y := g.Size()
		if v, err2 := g.SetView("stats", 0, 0, x-1, y-1); err != nil {
			if err2 != gocui.ErrUnknownView {
				return err2
			}
			v.Wrap = true
		}
		return nil
	})

	_ = g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, func(g *gocui.Gui, v *gocui.View) error {
		t.Stop()
		return nil
	})

	go func() {
		err = g.MainLoop()
		if err == gocui.ErrQuit {
			err = nil
		}
		if err != nil {
			log.Fatal(err)
		}
	}()

	go func() {
		for range time.Tick(100 * time.Millisecond) {
			s := t.Stats()
			b, err2 := jsonutil.MarshalCompactPretty(s)
			if err2 != nil {
				log.Fatal(err2)
			}
			g.Update(func(g *gocui.Gui) error {
				v, err2 := g.View("stats")
				if err2 != nil {
					return err2
				}
				v.Clear()
				_, _ = fmt.Fprintln(v, string(b))
				return nil
			})
		}
	}()

	t.Start()
	completeC := t.NotifyComplete()
	errC := t.NotifyError()
	for {
		select {
		case <-completeC:
			completeC = nil
			if !c.Bool("seed") {
				t.Stop()
				continue
			}
		case err = <-errC:
			if err != nil {
				err = res.Close()
			}
			return err
		}
	}
}

func handleServer(c *cli.Context) error {
	cfg := rainrpc.DefaultServerConfig

	configPath := c.String("config")
	if configPath != "" {
		cp, err := homedir.Expand(configPath)
		if err != nil {
			return err
		}
		b, err := ioutil.ReadFile(cp) // nolint: gosec
		if os.IsNotExist(err) {
			log.Noticef("config file not found at %q, using default config", cp)
		} else if err != nil {
			return err
		} else {
			err = yaml.Unmarshal(b, &cfg)
			if err != nil {
				return err
			}
			log.Infoln("config loaded from:", cp)
		}
	}

	srv, err := rainrpc.NewServer(cfg)
	if err != nil {
		return err
	}
	return srv.ListenAndServe()
}

func handleBeforeClient(c *cli.Context) error {
	var err error
	clt, err = rainrpc.NewClient(c.String("url"))
	return err
}

func handleList(c *cli.Context) error {
	resp, err := clt.ListTorrents()
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleAdd(c *cli.Context) error {
	var b []byte
	var marshalErr error
	arg := c.Args().Get(0)
	if strings.HasPrefix(arg, "magnet:") {
		resp, err := clt.AddMagnet(arg)
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	} else {
		f, err := os.Open(arg) // nolint: gosec
		if err != nil {
			return err
		}
		resp, err := clt.AddTorrent(f)
		_ = f.Close()
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	}
	if marshalErr != nil {
		return marshalErr
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleRemove(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.RemoveTorrent(id)
	return err
}

func handleStats(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentStats(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleStart(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StartTorrent(id)
	return err
}

func handleStop(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StopTorrent(id)
	return err
}

func handleConsole(c *cli.Context) error {
	con := console.New(clt)
	return con.Run()
}
<commit_msg>log parsed config value<commit_after>package main

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	clog "github.com/cenkalti/log"
	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/console"
	"github.com/cenkalti/rain/internal/jsonutil"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/rainrpc"
	"github.com/cenkalti/rain/torrent"
	"github.com/cenkalti/rain/torrent/resumer/boltdbresumer"
	"github.com/cenkalti/rain/torrent/storage/filestorage"
	"github.com/hokaccha/go-prettyjson"
	"github.com/jroimartin/gocui"
	"github.com/mitchellh/go-homedir"
	"github.com/urfave/cli"
	"gopkg.in/yaml.v2"
)

var (
	app = cli.NewApp()
	clt *rainrpc.Client
	log = logger.New("rain")
)

func main() {
	app.Version = clientversion.Version
	app.Usage = "BitTorrent client"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "config, c",
			Usage: "read config from `FILE`",
		},
		cli.StringFlag{
			Name:  "cpuprofile",
			Usage: "write cpu profile to `FILE`",
		},
		cli.StringFlag{
			Name:  "pprof",
			Usage: "run pprof server on `ADDR`",
		},
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "enable debug log",
		},
		cli.StringFlag{
			Name:  "logfile",
			Usage: "write log to `FILE`",
		},
	}
	app.Before = handleBeforeCommand
	app.After = handleAfterCommand
	app.Commands = []cli.Command{
		{
			Name:      "download",
			Usage:     "download torrent or magnet",
			ArgsUsage: "[torrent path or magnet link]",
			Action:    handleDownload,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "dest",
					Usage: "save files under `DIR`",
					Value: ".",
				},
				cli.IntFlag{
					Name:  "port",
					Usage: "peer listen port",
				},
				cli.BoolFlag{
					Name:  "seed",
					Usage: "continue seeding after download finishes",
				},
			},
		},
		{
			Name:  "server",
			Usage: "run rpc server and torrent client",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "config",
					Usage: "path to the config file",
					Value: "~/.rain/config.yaml",
				},
			},
			Action: handleServer,
		},
		{
			Name:  "client",
			Usage: "send rpc request to server",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "url",
					Usage: "URL of RPC server",
					Value: "localhost:" + strconv.Itoa(rainrpc.DefaultServerConfig.Port),
				},
			},
			Before: handleBeforeClient,
			Subcommands: []cli.Command{
				{
					Name:   "list",
					Usage:  "list torrents",
					Action: handleList,
				},
				{
					Name:   "add",
					Usage:  "add torrent or magnet",
					Action: handleAdd,
				},
				{
					Name:   "remove",
					Usage:  "remove torrent",
					Action: handleRemove,
				},
				{
					Name:   "stats",
					Usage:  "get stats of torrent",
					Action: handleStats,
				},
				{
					Name:   "start",
					Usage:  "start torrent",
					Action: handleStart,
				},
				{
					Name:   "stop",
					Usage:  "stop",
					Action: handleStop,
				},
				{
					Name:   "console",
					Usage:  "show client console",
					Action: handleConsole,
				},
			},
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(err)
	}
}

func handleBeforeCommand(c *cli.Context) error {
	cpuprofile := c.GlobalString("cpuprofile")
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatal("could not create CPU profile: ", err)
		}
		if err := pprof.StartCPUProfile(f); err != nil {
			log.Fatal("could not start CPU profile: ", err)
		}
	}
	pprofFlag := c.GlobalString("pprof")
	if pprofFlag != "" {
		go func() {
			log.Notice(http.ListenAndServe(pprofFlag, nil))
		}()
	}
	logFile := c.GlobalString("logfile")
	if logFile != "" {
		f, err := os.Create(logFile)
		if err != nil {
			log.Fatal("could not create log file: ", err)
		}
		logger.SetHandler(clog.NewFileHandler(f))
	}
	if c.GlobalBool("debug") {
		logger.SetLevel(clog.DEBUG)
	}
	return nil
}

func handleAfterCommand(c *cli.Context) error {
	if c.GlobalString("cpuprofile") != "" {
		pprof.StopCPUProfile()
	}
	return nil
}

func handleDownload(c *cli.Context) error {
	path := c.Args().Get(0)
	if path == "" {
		return errors.New("first argument must be a torrent file or magnet link")
	}

	sto, err := filestorage.New(c.String("dest"))
	if err != nil {
		return err
	}
	var t *torrent.Torrent
	if strings.HasPrefix(path, "magnet:") {
		t, err = torrent.NewMagnet(path, c.Int("port"), sto, torrent.DefaultConfig)
	} else {
		f, err2 := os.Open(path) // nolint: gosec
		if err2 != nil {
			return err
		}
		t, err = torrent.New(f, c.Int("port"), sto, torrent.DefaultConfig)
		_ = f.Close()
	}
	if err != nil {
		return err
	}
	defer t.Close()

	res, err := boltdbresumer.Open(t.Name()+"."+t.InfoHash()+".resume", []byte(t.InfoHash()))
	if err != nil {
		return err
	}
	defer func() { _ = res.Close() }()

	err = t.SetResume(res)
	if err != nil {
		return err
	}

	g, err := gocui.NewGui(gocui.OutputNormal)
	if err != nil {
		return err
	}
	defer g.Close()

	g.SetManagerFunc(func(g *gocui.Gui) error {
		x, y := g.Size()
		if v, err2 := g.SetView("stats", 0, 0, x-1, y-1); err != nil {
			if err2 != gocui.ErrUnknownView {
				return err2
			}
			v.Wrap = true
		}
		return nil
	})

	_ = g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, func(g *gocui.Gui, v *gocui.View) error {
		t.Stop()
		return nil
	})

	go func() {
		err = g.MainLoop()
		if err == gocui.ErrQuit {
			err = nil
		}
		if err != nil {
			log.Fatal(err)
		}
	}()

	go func() {
		for range time.Tick(100 * time.Millisecond) {
			s := t.Stats()
			b, err2 := jsonutil.MarshalCompactPretty(s)
			if err2 != nil {
				log.Fatal(err2)
			}
			g.Update(func(g *gocui.Gui) error {
				v, err2 := g.View("stats")
				if err2 != nil {
					return err2
				}
				v.Clear()
				_, _ = fmt.Fprintln(v, string(b))
				return nil
			})
		}
	}()

	t.Start()
	completeC := t.NotifyComplete()
	errC := t.NotifyError()
	for {
		select {
		case <-completeC:
			completeC = nil
			if !c.Bool("seed") {
				t.Stop()
				continue
			}
		case err = <-errC:
			if err != nil {
				err = res.Close()
			}
			return err
		}
	}
}

func handleServer(c *cli.Context) error {
	cfg := rainrpc.DefaultServerConfig

	configPath := c.String("config")
	if configPath != "" {
		cp, err := homedir.Expand(configPath)
		if err != nil {
			return err
		}
		b, err := ioutil.ReadFile(cp) // nolint: gosec
		if os.IsNotExist(err) {
			log.Noticef("config file not found at %q, using default config", cp)
		} else if err != nil {
			return err
		} else {
			err = yaml.Unmarshal(b, &cfg)
			if err != nil {
				return err
			}
			log.Infoln("config loaded from:", cp)
			b, err = yaml.Marshal(&cfg)
			if err != nil {
				return err
			}
			log.Debug("\n" + string(b))
		}
	}

	srv, err := rainrpc.NewServer(cfg)
	if err != nil {
		return err
	}
	return srv.ListenAndServe()
}

func handleBeforeClient(c *cli.Context) error {
	var err error
	clt, err = rainrpc.NewClient(c.String("url"))
	return err
}

func handleList(c *cli.Context) error {
	resp, err := clt.ListTorrents()
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleAdd(c *cli.Context) error {
	var b []byte
	var marshalErr error
	arg := c.Args().Get(0)
	if strings.HasPrefix(arg, "magnet:") {
		resp, err := clt.AddMagnet(arg)
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	} else {
		f, err := os.Open(arg) // nolint: gosec
		if err != nil {
			return err
		}
		resp, err := clt.AddTorrent(f)
		_ = f.Close()
		if err != nil {
			return err
		}
		b, marshalErr = prettyjson.Marshal(resp)
	}
	if marshalErr != nil {
		return marshalErr
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleRemove(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.RemoveTorrent(id)
	return err
}

func handleStats(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	resp, err := clt.GetTorrentStats(id)
	if err != nil {
		return err
	}
	b, err := prettyjson.Marshal(resp)
	if err != nil {
		return err
	}
	_, _ = os.Stdout.Write(b)
	_, _ = os.Stdout.WriteString("\n")
	return nil
}

func handleStart(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StartTorrent(id)
	return err
}

func handleStop(c *cli.Context) error {
	id, err := strconv.ParseUint(c.Args().Get(0), 10, 64)
	if err != nil {
		return err
	}
	_, err = clt.StopTorrent(id)
	return err
}

func handleConsole(c *cli.Context) error {
	con := console.New(clt)
	return con.Run()
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"html"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/dchest/uniuri"
	"github.com/ewhal/pygments"
	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/mux"
)

const (
	ADDRESS  = "http://localhost:9900"
	LENGTH   = 6
	PORT     = ":9900"
	USERNAME = ""
	PASS     = ""
	NAME     = ""
	DATABASE = USERNAME + ":" + PASS + "@/" + NAME + "?charset=utf8"
)

type Response struct {
	ID     string `json:"id"`
	TITLE  string `json:"title"`
	HASH   string `json:"hash"`
	URL    string `json:"url"`
	SIZE   int    `json:"size"`
	DELKEY string `json:"delkey"`
}

type Page struct {
	Title    string
	Body     []byte
	Raw      string
	Home     string
	Download string
	Clone    string
}

func check(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

func generateName() string {
	s := uniuri.NewLen(LENGTH)
	db, err := sql.Open("mysql", DATABASE)
	check(err)

	query, err := db.Query("select id from pastebin")
	for query.Next() {
		var id string
		err := query.Scan(&id)
		if err != nil {

		}
		if id == s {
			generateName()
		}
	}
	db.Close()

	return s

}
func hash(paste string) string {
	hasher := sha1.New()

	hasher.Write([]byte(paste))
	sha := base64.URLEncoding.EncodeToString(hasher.Sum(nil))
	return sha
}

func save(raw string, lang string, title string, expiry string) []string {
	db, err := sql.Open("mysql", DATABASE)
	check(err)

	sha := hash(raw)
	query, err := db.Query("select id, title, hash, data, delkey from pastebin")
	for query.Next() {
		var id, title, hash, paste, delkey string
		err := query.Scan(&id, &title, &hash, &paste, &delkey)
		check(err)
		if hash == sha {
			url := ADDRESS + "/p/" + id
			return []string{id, title, hash, url, paste, delkey}
		}
	}
	id := generateName()
	var url string
	if lang == "" {
		url = ADDRESS + "/p/" + id
	} else {
		url = ADDRESS + "/p/" + id + "/" + lang
	}
	now := time.Now()
	var expiryTime string

	switch expiry {
	case "5 minutes":
		expiryTime = now.Add(time.Minute * 5).Format("2006-01-02 15:04:05")
		break

	case "1 hour":
		expiryTime = now.Add(time.Hour + 1).Format("2006-01-02 15:04:05")
		break

	case "1 day":
		expiryTime = now.Add(time.Hour * 24 * 1).Format("2006-01-02 15:04:05")
		break

	case "1 week":
		expiryTime = now.Add(time.Hour * 24 * 7).Format("2006-01-02 15:04:05")
		break

	case "1 month":
		expiryTime = now.Add(time.Hour * 24 * 30).Format("2006-01-02 15:04:05")
		break

	case "1 year":
		expiryTime = now.Add(time.Hour * 24 * 365).Format("2006-01-02 15:04:05")
		break

	case "forever":
		expiryTime = now.Add(time.Hour * 24 * (365 * 20)).Format("2006-01-02 15:04:05")
		break

	default:
		expiryTime = now.Add(time.Hour * 24 * (365 * 20)).Format("2006-01-02 15:04:05")
		break

	}
	delKey := uniuri.NewLen(40)
	paste := html.EscapeString(raw)

	stmt, err := db.Prepare("INSERT INTO pastebin(id, title, hash, data, delkey, expiry) values(?,?,?,?,?,?)")
	check(err)
	if title == "" {
		_, err = stmt.Exec(id, id, sha, paste, delKey, expiryTime)
		check(err)
	} else {
		_, err = stmt.Exec(id, html.EscapeString(title), sha, paste, delKey, expiryTime)
		check(err)
	}
	db.Close()
	return []string{id, title, sha, url, paste, delKey}
}

func delHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	delkey := vars["delKey"]

	db, err := sql.Open("mysql", DATABASE)
	check(err)

	stmt, err := db.Prepare("delete from pastebin where delkey=?")
	check(err)

	res, err := stmt.Exec(html.EscapeString(delkey))
	check(err)

	_, err = res.RowsAffected()
	if err == sql.ErrNoRows {
		io.WriteString(w, "Error invalid paste")
	} else {
		io.WriteString(w, paste+" deleted")
	}
	db.Close()

}
func saveHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	output := vars["output"]
	switch r.Method {
	case "POST":
		paste := r.FormValue("p")
		lang := r.FormValue("lang")
		title := r.FormValue("title")
		expiry := r.FormValue("expiry")
		if paste == "" {
			http.Error(w, "Empty paste", 500)
			return
		}
		values := save(paste, lang, title, expiry)
		b := &Response{
			ID:     values[0],
			TITLE:  values[1],
			HASH:   values[2],
			URL:    values[3],
			SIZE:   len(values[4]),
			DELKEY: values[5],
		}

		switch output {
		case "json":

			w.Header().Set("Content-Type", "application/json")
			err := json.NewEncoder(w).Encode(b)

			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
		case "xml":
			x, err := xml.MarshalIndent(b, "", "  ")
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			w.Header().Set("Content-Type", "application/xml")
			w.Write(x)

		case "html":
			w.Header().Set("Content-Type", "text/html")
			io.WriteString(w, "
URL: "+b.URL+"
")
			io.WriteString(w, "
Delete Key: "+b.DELKEY+"
")

		case "redirect":
			http.Redirect(w, r, b.URL, 301)

		default:
			w.Header().Set("Content-Type", "plain/text")
			io.WriteString(w, b.URL+"\n")
			io.WriteString(w, "delete key: "+b.DELKEY+"\n")
		}
	}

}

func highlight(s string, lang string) (string, error) {

	highlight, err := pygments.Highlight(html.UnescapeString(s), html.EscapeString(lang), "html", "style=autumn,linenos=True, lineanchors=True,anchorlinenos=True,noclasses=True,", "utf-8")
	if err != nil {
		return "", err
	}
	return highlight, nil

}

func getPaste(paste string, lang string) (string, string) {
	param1 := html.EscapeString(paste)
	db, err := sql.Open("mysql", DATABASE)
	var title, s string
	var expiry string
	err = db.QueryRow("select title, data, expiry from pastebin where id=?", param1).Scan(&title, &s, &expiry)
	check(err)
	if time.Now().Format("2006-01-02 15:04:05") > expiry {
		stmt, err := db.Prepare("delete from pastebin where id=?")
		check(err)
		_, err = stmt.Exec(param1)
		check(err)
		return "Error invalid paste", ""
	}
	db.Close()

	if err == sql.ErrNoRows {
		return "Error invalid paste", ""
	} else {
		if lang == "" {
			return html.UnescapeString(s), html.UnescapeString(title)
		} else {
			high, err := highlight(s, lang)
			check(err)
			return high, html.UnescapeString(title)

		}
	}

}

var templates = template.Must(template.ParseFiles("assets/paste.html", "assets/index.html", "assets/clone.html"))
var syntax, _ = ioutil.ReadFile("assets/syntax.html")

func pasteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	lang := vars["lang"]
	s, title := getPaste(paste, lang)
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	if lang == "" {

		err := templates.ExecuteTemplate(w, "paste.html", p)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}

	} else {
		fmt.Fprintf(w, string(syntax), p.Title, p.Title, s, p.Home, p.Download, p.Raw, p.Clone)

	}
}

func cloneHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, title := getPaste(paste, "")
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	err := templates.ExecuteTemplate(w, "clone.html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

}
func downloadHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")
	w.Header().Set("Content-Disposition", "attachment; filename="+paste)
	w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
	io.WriteString(w, s)

}
func rawHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")
	io.WriteString(w, s)

}

func main() {
	router := mux.NewRouter()
	router.HandleFunc("/p/{pasteId}", pasteHandler)
	router.HandleFunc("/raw/{pasteId}", rawHandler)
	router.HandleFunc("/p/{pasteId}/{lang}", pasteHandler)
	router.HandleFunc("/clone/{pasteId}", cloneHandler)
	router.HandleFunc("/download/{pasteId}", downloadHandler)
	router.HandleFunc("/save", saveHandler)
	router.HandleFunc("/save/{output}", saveHandler)
	router.HandleFunc("/del/{pasteId}/{delKey}", delHandler)
	router.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("assets/"))))
	err := http.ListenAndServe(PORT, router)
	if err != nil {
		log.Fatal(err)
	}

}
<commit_msg>Let sql do the searching<commit_after>package main

import (
	"crypto/sha1"
	"database/sql"
	"encoding/base64"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"html"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/dchest/uniuri"
	"github.com/ewhal/pygments"
	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/mux"
)

const (
	ADDRESS  = "http://localhost:9900"
	LENGTH   = 6
	PORT     = ":9900"
	USERNAME = ""
	PASS     = ""
	NAME     = ""
	DATABASE = USERNAME + ":" + PASS + "@/" + NAME + "?charset=utf8"
)

type Response struct {
	ID     string `json:"id"`
	TITLE  string `json:"title"`
	HASH   string `json:"hash"`
	URL    string `json:"url"`
	SIZE   int    `json:"size"`
	DELKEY string `json:"delkey"`
}

type Page struct {
	Title    string
	Body     []byte
	Raw      string
	Home     string
	Download string
	Clone    string
}

func check(err error) {
	if err != nil {
		fmt.Println(err)
	}
}

func generateName() string {
	s := uniuri.NewLen(LENGTH)
	db, err := sql.Open("mysql", DATABASE)
	check(err)

	query, err := db.Query("select id from pastebin")
	for query.Next() {
		var id string
		err := query.Scan(&id)
		if err != nil {

		}
		if id == s {
			generateName()
		}
	}
	db.Close()

	return s

}
func hash(paste string) string {
	hasher := sha1.New()

	hasher.Write([]byte(paste))
	sha := base64.URLEncoding.EncodeToString(hasher.Sum(nil))
	return sha
}

func save(raw string, lang string, title string, expiry string) []string {
	db, err := sql.Open("mysql", DATABASE)
	check(err)

	sha := hash(raw)
	query, err := db.Query("select id, title, hash, data, delkey from pastebin where hash=?", sha)
	if err != sql.ErrNoRows {
		for query.Next() {
			var id, title, hash, paste, delkey string
			err := query.Scan(&id, &title, &hash, &paste, &delkey)
			check(err)
			url := ADDRESS + "/p/" + id
			return []string{id, title, hash, url, paste, delkey}
		}
	}
	id := generateName()
	var url string
	if lang == "" {
		url = ADDRESS + "/p/" + id
	} else {
		url = ADDRESS + "/p/" + id + "/" + lang
	}
	now := time.Now()
	var expiryTime string

	switch expiry {
	case "5 minutes":
		expiryTime = now.Add(time.Minute * 5).Format("2006-01-02 15:04:05")
		break

	case "1 hour":
		expiryTime = now.Add(time.Hour + 1).Format("2006-01-02 15:04:05")
		break

	case "1 day":
		expiryTime = now.Add(time.Hour * 24 * 1).Format("2006-01-02 15:04:05")
		break

	case "1 week":
		expiryTime = now.Add(time.Hour * 24 * 7).Format("2006-01-02 15:04:05")
		break

	case "1 month":
		expiryTime = now.Add(time.Hour * 24 * 30).Format("2006-01-02 15:04:05")
		break

	case "1 year":
		expiryTime = now.Add(time.Hour * 24 * 365).Format("2006-01-02 15:04:05")
		break

	case "forever":
		expiryTime = now.Add(time.Hour * 24 * (365 * 20)).Format("2006-01-02 15:04:05")
		break

	default:
		expiryTime = now.Add(time.Hour * 24 * (365 * 20)).Format("2006-01-02 15:04:05")
		break

	}
	delKey := uniuri.NewLen(40)
	paste := html.EscapeString(raw)

	stmt, err := db.Prepare("INSERT INTO pastebin(id, title, hash, data, delkey, expiry) values(?,?,?,?,?,?)")
	check(err)
	if title == "" {
		_, err = stmt.Exec(id, id, sha, paste, delKey, expiryTime)
		check(err)
	} else {
		_, err = stmt.Exec(id, html.EscapeString(title), sha, paste, delKey, expiryTime)
		check(err)
	}
	db.Close()
	return []string{id, title, sha, url, paste, delKey}
}

func delHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	delkey := vars["delKey"]

	db, err := sql.Open("mysql", DATABASE)
	check(err)

	stmt, err := db.Prepare("delete from pastebin where delkey=?")
	check(err)

	res, err := stmt.Exec(html.EscapeString(delkey))
	check(err)

	_, err = res.RowsAffected()
	if err == sql.ErrNoRows {
		io.WriteString(w, "Error invalid paste")
	} else {
		io.WriteString(w, paste+" deleted")
	}
	db.Close()

}
func saveHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	output := vars["output"]
	switch r.Method {
	case "POST":
		paste := r.FormValue("p")
		lang := r.FormValue("lang")
		title := r.FormValue("title")
		expiry := r.FormValue("expiry")
		if paste == "" {
			http.Error(w, "Empty paste", 500)
			return
		}
		values := save(paste, lang, title, expiry)
		b := &Response{
			ID:     values[0],
			TITLE:  values[1],
			HASH:   values[2],
			URL:    values[3],
			SIZE:   len(values[4]),
			DELKEY: values[5],
		}

		switch output {
		case "json":

			w.Header().Set("Content-Type", "application/json")
			err := json.NewEncoder(w).Encode(b)

			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
		case "xml":
			x, err := xml.MarshalIndent(b, "", "  ")
			if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			w.Header().Set("Content-Type", "application/xml")
			w.Write(x)

		case "html":
			w.Header().Set("Content-Type", "text/html")
			io.WriteString(w, "
URL: "+b.URL+"
")
			io.WriteString(w, "
Delete Key: "+b.DELKEY+"
")

		case "redirect":
			http.Redirect(w, r, b.URL, 301)

		default:
			w.Header().Set("Content-Type", "plain/text")
			io.WriteString(w, b.URL+"\n")
			io.WriteString(w, "delete key: "+b.DELKEY+"\n")
		}
	}

}

func highlight(s string, lang string) (string, error) {

	highlight, err := pygments.Highlight(html.UnescapeString(s), html.EscapeString(lang), "html", "style=autumn,linenos=True, lineanchors=True,anchorlinenos=True,noclasses=True,", "utf-8")
	if err != nil {
		return "", err
	}
	return highlight, nil

}

func getPaste(paste string, lang string) (string, string) {
	param1 := html.EscapeString(paste)
	db, err := sql.Open("mysql", DATABASE)
	var title, s string
	var expiry string
	err = db.QueryRow("select title, data, expiry from pastebin where id=?", param1).Scan(&title, &s, &expiry)
	check(err)
	if time.Now().Format("2006-01-02 15:04:05") > expiry {
		stmt, err := db.Prepare("delete from pastebin where id=?")
		check(err)
		_, err = stmt.Exec(param1)
		check(err)
		return "Error invalid paste", ""
	}
	db.Close()

	if err == sql.ErrNoRows {
		return "Error invalid paste", ""
	} else {
		if lang == "" {
			return html.UnescapeString(s), html.UnescapeString(title)
		} else {
			high, err := highlight(s, lang)
			check(err)
			return high, html.UnescapeString(title)

		}
	}

}

var templates = template.Must(template.ParseFiles("assets/paste.html", "assets/index.html", "assets/clone.html"))
var syntax, _ = ioutil.ReadFile("assets/syntax.html")

func pasteHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	lang := vars["lang"]
	s, title := getPaste(paste, lang)
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	if lang == "" {

		err := templates.ExecuteTemplate(w, "paste.html", p)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}

	} else {
		fmt.Fprintf(w, string(syntax), p.Title, p.Title, s, p.Home, p.Download, p.Raw, p.Clone)

	}
}

func cloneHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, title := getPaste(paste, "")
	link := ADDRESS + "/raw/" + paste
	download := ADDRESS + "/download/" + paste
	clone := ADDRESS + "/clone/" + paste
	p := &Page{
		Title:    title,
		Body:     []byte(s),
		Raw:      link,
		Home:     ADDRESS,
		Download: download,
		Clone:    clone,
	}
	err := templates.ExecuteTemplate(w, "clone.html", p)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

}
func downloadHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")
	w.Header().Set("Content-Disposition", "attachment; filename="+paste)
	w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
	io.WriteString(w, s)

}
func rawHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	paste := vars["pasteId"]
	s, _ := getPaste(paste, "")
	io.WriteString(w, s)

}

func main() {
	router := mux.NewRouter()
	router.HandleFunc("/p/{pasteId}", pasteHandler)
	router.HandleFunc("/raw/{pasteId}", rawHandler)
	router.HandleFunc("/p/{pasteId}/{lang}", pasteHandler)
	router.HandleFunc("/clone/{pasteId}", cloneHandler)
	router.HandleFunc("/download/{pasteId}", downloadHandler)
	router.HandleFunc("/save", saveHandler)
	router.HandleFunc("/save/{output}", saveHandler)
	router.HandleFunc("/del/{pasteId}/{delKey}", delHandler)
	router.PathPrefix("/").Handler(http.StripPrefix("/", http.FileServer(http.Dir("assets/"))))
	err := http.ListenAndServe(PORT, router)
	if err != nil {
		log.Fatal(err)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aduermael/crypto/ssh"
	"github.com/spf13/cobra"
)

var (
	// path to private key or empty to use default
	sshIdentityFile = ""
	// open a bash session by default, but different option can be used
	shell = "bash"
	// proxy mode (don't start shell session)
	proxyMode = false
	// verbose mode (debug logs)
	verbose = false
)

func main() {

	rootCmd := &cobra.Command{
		Use:   "docker-tunnel [user@]host",
		Short: "Docker-tunnel connects you to a remote Docker host through an SSH tunnel",
		Run: func(cmd *cobra.Command, args []string) {
			if verbose {
				logLevel = logLevelDebug
			}

			if len(args) != 1 {
				cmd.Usage()
				return
			}

			sshClient, err := sshConnect(args[0], sshIdentityFile)
			if err != nil {
				printFatal(err)
			}
			defer sshClient.Close()

			if proxyMode {
				printDebug("proxy mode")

				ln, err := net.Listen("tcp", ":2375")
				if err != nil {
					printFatal(err)
				}
				print("listening on port 2375...")
				for {
					conn, err := ln.Accept()
					if err != nil {
						printFatal(err)
					}
					go handleProxyConnection(conn, sshClient)
				}
			}

			// proxyMode == false
			// open shell session, connected to remote Docker host
			printDebug("shell mode")

			socketPath := tmpSocketPath()
			printDebug("socket path:", socketPath)

			ln, err := net.Listen("unix", socketPath)
			if err != nil {
				printFatal(err)
			}
			defer os.RemoveAll(socketPath)

			// listen in background
			go func() {
				for {
					conn, err := ln.Accept()
					if err != nil {
						printFatal(err)
					}
					printDebug("handle socket connection")
					go handleProxyConnection(conn, sshClient)
				}
			}()

			os.Setenv("PS1", "🐳  $ ")
			os.Setenv("DOCKER_HOST", "unix://"+socketPath)

			sh := exec.Command(shell)
			sh.Stdout = os.Stdout
			sh.Stderr = os.Stderr
			sh.Stdin = os.Stdin

			_ = sh.Run()
		},
	}

	rootCmd.Flags().StringVarP(&sshIdentityFile, "sshid", "i", "", "path to private key")
	rootCmd.Flags().StringVarP(&shell, "shell", "s", "bash", "shell to open session")
	rootCmd.Flags().BoolVarP(&proxyMode, "proxy", "p", false, "proxy mode (don't start shell session)")
	rootCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "verbose mode (debug logs)")

	if err := rootCmd.Execute(); err != nil {
		printFatal(err.Error())
	}
}

func tmpSocketPath() string {
	randBytes := make([]byte, 16)
	rand.Read(randBytes)
	return filepath.Join(os.TempDir(), "docker-"+hex.EncodeToString(randBytes)+".sock")
}

func sshConnect(userAtHost, privateKeyPath string) (*ssh.Client, error) {
	// root user by default
	user := "root"
	host := ""
	userAndHost := strings.SplitN(userAtHost, "@", 2)
	if len(userAndHost) == 1 {
		host = userAndHost[0]
	} else {
		user = userAndHost[0]
		host = userAndHost[1]
	}

	printDebug("user:", user)

	u, err := url.Parse(host)
	if err != nil {
		return nil, fmt.Errorf("ssh connection can't be established: %s", err)
	}
	if u.Scheme == "" {
		u.Scheme = "tcp"
	}

	authMethod, err := authMethodPublicKeys(privateKeyPath)
	if err != nil {
		printFatal(err)
	}
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{authMethod},
	}

	addr := u.Host + u.Path
	parts := strings.Split(addr, ":")
	lastPart := parts[len(parts)-1]
	_, err = strconv.Atoi(lastPart)
	// port is required, used 22 by default
	if err != nil {
		addr += ":22"
	}

	printDebug("address:", u.Scheme+"://"+addr)

	sshClientConn, err := ssh.Dial(u.Scheme, addr, config)
	if err != nil {
		return nil, fmt.Errorf("ssh connection can't be established: %s", err)
	}

	printDebug("ssh connection established")

	return sshClientConn, nil
}

func handleProxyConnection(conn net.Conn, sshClient *ssh.Client) {
	err := forward(conn, sshClient, "unix:///var/run/docker.sock")
	if err != nil {
		printError(os.Stderr, "can't forward connection:", err.Error())
	}
}

func forward(conn net.Conn, sshClient *ssh.Client, remoteAddr string) error {

	// parse OpenSSH version, 6.7 is the minimum required
	reOpenSSH := regexp.MustCompile("OpenSSH_[.0-9]+")
	reOpenSSHVersion := regexp.MustCompile("[.0-9]+")
	match := reOpenSSH.Find(sshClient.ServerVersion())
	openSSHVersionStr := string(reOpenSSHVersion.Find(match))
	openSSHVersion, err := strconv.ParseFloat(openSSHVersionStr, 64)
	if err != nil {
		return errors.New("can't parse server OpenSSH version")
	}
	if openSSHVersion < 6.7 {
		return errors.New("OpenSSH 6.7 minimum required on server side")
	}

	// remote addr
	u, err := url.Parse(remoteAddr)
	if err != nil {
		return fmt.Errorf("can't parse remote address: %s", remoteAddr)
	}

	addr := filepath.Join(u.Host, u.Path)

	sshConn, err := sshClient.Dial(u.Scheme, addr)
	if err != nil {
		return fmt.Errorf("can't connect to %s (from remote)", remoteAddr)
	}

	chan1 := make(chan struct{})
	chan2 := make(chan struct{})
	chan3 := make(chan struct{})
	var o sync.Once
	closeChan2 := func() {
		close(chan2)
	}

	// Copy conn.Reader to sshConn.Writer
	go func() {
		printDebug("copy: read from client conn, write to server conn")
		_, err := io.Copy(sshConn, conn)
		if err != nil {
			printFatal(err)
		}
		close(chan1)

		if c, ok := sshConn.(ssh.Channel); ok {
			_ = c.CloseWrite()
			printDebug("closed sshConn writer")
		}

		for {
			printDebug("can't read from client anymore, trying to write...")
			_, err := conn.Write(make([]byte, 0))
			if err != nil {
				printDebug("can't write, closing both conns")
				o.Do(closeChan2)
				break
			}
			time.Sleep(500 * time.Millisecond)
		}
	}()

	// Copy sshConn.Reader to localConn.Writer
	go func() {
		printDebug("copy: read from server conn, write to client conn")
		_, err := io.Copy(conn, sshConn)
		if err != nil {
			printFatal(err)
		}
		printDebug("can't read from server anymore...")

		if tcpConn, ok := conn.(*net.TCPConn); ok {
			_ = tcpConn.CloseWrite()
			printDebug("closed TCPconn writer")
		} else if unixConn, ok := conn.(*net.UnixConn); ok {
			_ = unixConn.CloseWrite()
			printDebug("closed unixConn writer")
		} else {
			printDebug("can't close conn writer")
		}

		o.Do(closeChan2)
		close(chan3)
	}()

	<-chan1
	<-chan2
	conn.Close()
	sshConn.Close()

	<-chan3

	printDebug("closed socket connection")

	return nil
}
<commit_msg>added error logs<commit_after>package main

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aduermael/crypto/ssh"
	"github.com/spf13/cobra"
)

var (
	// path to private key or empty to use default
	sshIdentityFile = ""
	// open a bash session by default, but different option can be used
	shell = "bash"
	// proxy mode (don't start shell session)
	proxyMode = false
	// verbose mode (debug logs)
	verbose = false
)

func main() {

	rootCmd := &cobra.Command{
		Use:   "docker-tunnel [user@]host",
		Short: "Docker-tunnel connects you to a remote Docker host through an SSH tunnel",
		Run: func(cmd *cobra.Command, args []string) {
			if verbose {
				logLevel = logLevelDebug
			}

			if len(args) != 1 {
				cmd.Usage()
				return
			}

			sshClient, err := sshConnect(args[0], sshIdentityFile)
			if err != nil {
				printFatal(err)
			}
			defer sshClient.Close()

			if proxyMode {
				printDebug("proxy mode")

				ln, err := net.Listen("tcp", ":2375")
				if err != nil {
					printFatal(err)
				}
				print("listening on port 2375...")
				for {
					conn, err := ln.Accept()
					if err != nil {
						printFatal(err)
					}
					go handleProxyConnection(conn, sshClient)
				}
			}

			// proxyMode == false
			// open shell session, connected to remote Docker host
			printDebug("shell mode")

			socketPath := tmpSocketPath()
			printDebug("socket path:", socketPath)

			ln, err := net.Listen("unix", socketPath)
			if err != nil {
				printFatal(err)
			}
			defer os.RemoveAll(socketPath)

			// listen in background
			go func() {
				for {
					conn, err := ln.Accept()
					if err != nil {
						printFatal(err)
					}
					printDebug("handle socket connection")
					go handleProxyConnection(conn, sshClient)
				}
			}()

			os.Setenv("PS1", "🐳  $ ")
			os.Setenv("DOCKER_HOST", "unix://"+socketPath)

			sh := exec.Command(shell)
			sh.Stdout = os.Stdout
			sh.Stderr = os.Stderr
			sh.Stdin = os.Stdin

			_ = sh.Run()
		},
	}

	rootCmd.Flags().StringVarP(&sshIdentityFile, "sshid", "i", "", "path to private key")
	rootCmd.Flags().StringVarP(&shell, "shell", "s", "bash", "shell to open session")
	rootCmd.Flags().BoolVarP(&proxyMode, "proxy", "p", false, "proxy mode (don't start shell session)")
	rootCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "verbose mode (debug logs)")

	if err := rootCmd.Execute(); err != nil {
		printFatal(err.Error())
	}
}

func tmpSocketPath() string {
	randBytes := make([]byte, 16)
	rand.Read(randBytes)
	return filepath.Join(os.TempDir(), "docker-"+hex.EncodeToString(randBytes)+".sock")
}

func sshConnect(userAtHost, privateKeyPath string) (*ssh.Client, error) {
	// root user by default
	user := "root"
	host := ""
	userAndHost := strings.SplitN(userAtHost, "@", 2)
	if len(userAndHost) == 1 {
		host = userAndHost[0]
	} else {
		user = userAndHost[0]
		host = userAndHost[1]
	}

	printDebug("user:", user)

	u, err := url.Parse(host)
	if err != nil {
		return nil, fmt.Errorf("ssh connection can't be established: %s", err)
	}
	if u.Scheme == "" {
		u.Scheme = "tcp"
	}

	authMethod, err := authMethodPublicKeys(privateKeyPath)
	if err != nil {
		printFatal(err)
	}
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{authMethod},
	}

	addr := u.Host + u.Path
	parts := strings.Split(addr, ":")
	lastPart := parts[len(parts)-1]
	_, err = strconv.Atoi(lastPart)
	// port is required, used 22 by default
	if err != nil {
		addr += ":22"
	}

	printDebug("address:", u.Scheme+"://"+addr)

	sshClientConn, err := ssh.Dial(u.Scheme, addr, config)
	if err != nil {
		return nil, fmt.Errorf("ssh connection can't be established: %s", err)
	}

	printDebug("ssh connection established")

	return sshClientConn, nil
}

func handleProxyConnection(conn net.Conn, sshClient *ssh.Client) {
	err := forward(conn, sshClient, "unix:///var/run/docker.sock")
	if err != nil {
		printError(os.Stderr, "can't forward connection:", err.Error())
	}
}

func forward(conn net.Conn, sshClient *ssh.Client, remoteAddr string) error {

	// parse OpenSSH version, 6.7 is the minimum required
	reOpenSSH := regexp.MustCompile("OpenSSH_[.0-9]+")
	reOpenSSHVersion := regexp.MustCompile("[.0-9]+")
	match := reOpenSSH.Find(sshClient.ServerVersion())
	openSSHVersionStr := string(reOpenSSHVersion.Find(match))
	openSSHVersion, err := strconv.ParseFloat(openSSHVersionStr, 64)
	if err != nil {
		return errors.New("can't parse server OpenSSH version")
	}
	if openSSHVersion < 6.7 {
		return errors.New("OpenSSH 6.7 minimum required on server side")
	}

	// remote addr
	u, err := url.Parse(remoteAddr)
	if err != nil {
		return fmt.Errorf("can't parse remote address: %s", remoteAddr)
	}

	addr := filepath.Join(u.Host, u.Path)

	sshConn, err := sshClient.Dial(u.Scheme, addr)
	if err != nil {
		return fmt.Errorf("can't connect to %s (from remote)", remoteAddr)
	}

	chan1 := make(chan struct{})
	chan2 := make(chan struct{})
	chan3 := make(chan struct{})
	var o sync.Once
	closeChan2 := func() {
		close(chan2)
	}

	// Copy conn.Reader to sshConn.Writer
	go func() {
		printDebug("copy: read from client conn, write to server conn")
		_, err := io.Copy(sshConn, conn)
		if err != nil {
			printFatal(err)
		}
		close(chan1)

		if c, ok := sshConn.(ssh.Channel); ok {
			if err := c.CloseWrite(); err != nil {
				printError("can't close sshConn writer")
			} else {
				printDebug("closed sshConn writer")
			}

		}

		for {
			printDebug("can't read from client anymore, trying to write...")
			_, err := conn.Write(make([]byte, 0))
			if err != nil {
				printDebug("can't write, closing both connections")
				o.Do(closeChan2)
				break
			}
			time.Sleep(500 * time.Millisecond)
		}
	}()

	// Copy sshConn.Reader to localConn.Writer
	go func() {
		printDebug("copy: read from server conn, write to client conn")
		_, err := io.Copy(conn, sshConn)
		if err != nil {
			printFatal(err)
		}
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			if err := tcpConn.CloseWrite(); err != nil {
				printError("can't close TCPconn writer")
			} else {
				printDebug("closed TCPconn writer")
			}
		} else if unixConn, ok := conn.(*net.UnixConn); ok {
			if err := unixConn.CloseWrite(); err != nil {
				printError("can't close UnixConn writer")
			} else {
				printDebug("closed UnixConn writer")
			}
		} else {
			printDebug("can't close conn writer")
		}
		o.Do(closeChan2)
		close(chan3)
	}()

	<-chan1
	<-chan2
	conn.Close()
	sshConn.Close()

	<-chan3

	printDebug("closed socket connection")

	return nil
}
<|endoftext|>


<commit_before>package main

import (
    "flag"
    "bytes"
    "os/exec"
    "io/ioutil"
    "syscall"
)

var (
    storeType = flag.String("storeType", "", "zookeeper, consul or etcd.")
    storeConnection = flag.String("storeConnection", "", "Key-value store connection string.")
    rootPath = flag.String("rootPath", "", "Scheduled workflow key-value store root path.")

    workflowPath = flag.String("workflowPath", "/opt/vamp", "Path to workflow files.")

    logo = flag.Bool("logo", true, "Show logo.")
    help = flag.Bool("help", false, "Print usage.")

    logger = createLogger()
)

func printLogo(version string) string {
    return `
██╗   ██╗ █████╗ ███╗   ███╗██████╗
██║   ██║██╔══██╗████╗ ████║██╔══██╗
██║   ██║███████║██╔████╔██║██████╔╝
╚██╗ ██╔╝██╔══██║██║╚██╔╝██║██╔═══╝
 ╚████╔╝ ██║  ██║██║ ╚═╝ ██║██║
  ╚═══╝  ╚═╝  ╚═╝╚═╝     ╚═╝╚═╝
                       workflow agent
                       version ` + version + `
                       by magnetic.io
                                      `
}

func main() {

    flag.Parse()

    if *logo {
        logger.Notice(printLogo("0.9.0"))
    }

    if *help {
        flag.Usage()
        return
    }

    if len(*storeType) == 0 {
        logger.Panic("Key-value store type not speciffed.")
        return
    }

    if len(*rootPath) == 0 {
        logger.Panic("Key-value store root key path not speciffed.")
        return
    }

    if len(*storeConnection) == 0 {
        logger.Panic("Key-value store connection not speciffed.")
        return
    }

    logger.Notice("Starting Vamp Workflow Agent")

    logger.Info("Key-value store type          : %s", *storeType)
    logger.Info("Key-value store connection    : %s", *storeConnection)
    logger.Info("Key-value store root key path : %s", *rootPath)
    logger.Info("Workflow file path            : %s", *workflowPath)

    workflowKey := *rootPath + "/workflow"
    logger.Info("Reading workflow from         : %s", workflowKey)

    content, err := readFromKeyValueStore(workflowKey)

    if err != nil {
        logger.Panic("Can't read the workflow: ", err)
        return
    }

    workflowFile := *workflowPath + "/workflow.js"
    logger.Info("Writing workflow script       : %s", workflowFile)
    err = ioutil.WriteFile(workflowFile, content, 0644)

    if err != nil {
        logger.Panic("Can't write to the workflow script: ", err)
        return
    }

    logger.Info("Executing 'workflow.js' by Node.js.")
    var cmd *exec.Cmd
    cmd = exec.Command("node", workflowFile)

    var stdout, stderr bytes.Buffer
    cmd.Stdout = &stdout
    cmd.Stderr = &stderr

    err = cmd.Run()

    exitStatus := 0

    if err != nil {
        logger.Error("Error during execution: %s, %s", err.Error(), string(stderr.Bytes()[:]))
        if exitError, ok := err.(*exec.ExitError); ok {
            if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
                exitStatus = status.ExitStatus()
            }
        }

    } else {
        logger.Info("Execution standard output: %s", string(stdout.Bytes()[:]))
    }

    logger.Notice("Exit status: %d", exitStatus)
}
<commit_msg>exit with the same exit code as node.js<commit_after>package main

import (
    "flag"
    "bytes"
    "os/exec"
    "io/ioutil"
    "syscall"
    "os"
)

var (
    storeType = flag.String("storeType", "", "zookeeper, consul or etcd.")
    storeConnection = flag.String("storeConnection", "", "Key-value store connection string.")
    rootPath = flag.String("rootPath", "", "Scheduled workflow key-value store root path.")

    workflowPath = flag.String("workflowPath", "/opt/vamp", "Path to workflow files.")

    logo = flag.Bool("logo", true, "Show logo.")
    help = flag.Bool("help", false, "Print usage.")

    logger = createLogger()
)

func printLogo(version string) string {
    return `
██╗   ██╗ █████╗ ███╗   ███╗██████╗
██║   ██║██╔══██╗████╗ ████║██╔══██╗
██║   ██║███████║██╔████╔██║██████╔╝
╚██╗ ██╔╝██╔══██║██║╚██╔╝██║██╔═══╝
 ╚████╔╝ ██║  ██║██║ ╚═╝ ██║██║
  ╚═══╝  ╚═╝  ╚═╝╚═╝     ╚═╝╚═╝
                       workflow agent
                       version ` + version + `
                       by magnetic.io
                                      `
}

func main() {

    flag.Parse()

    if *logo {
        logger.Notice(printLogo("0.9.0"))
    }

    if *help {
        flag.Usage()
        return
    }

    if len(*storeType) == 0 {
        logger.Panic("Key-value store type not speciffed.")
        return
    }

    if len(*rootPath) == 0 {
        logger.Panic("Key-value store root key path not speciffed.")
        return
    }

    if len(*storeConnection) == 0 {
        logger.Panic("Key-value store connection not speciffed.")
        return
    }

    logger.Notice("Starting Vamp Workflow Agent")

    logger.Info("Key-value store type          : %s", *storeType)
    logger.Info("Key-value store connection    : %s", *storeConnection)
    logger.Info("Key-value store root key path : %s", *rootPath)
    logger.Info("Workflow file path            : %s", *workflowPath)

    workflowKey := *rootPath + "/workflow"
    logger.Info("Reading workflow from         : %s", workflowKey)

    content, err := readFromKeyValueStore(workflowKey)

    if err != nil {
        logger.Panic("Can't read the workflow: ", err)
        return
    }

    workflowFile := *workflowPath + "/workflow.js"
    logger.Info("Writing workflow script       : %s", workflowFile)
    err = ioutil.WriteFile(workflowFile, content, 0644)

    if err != nil {
        logger.Panic("Can't write to the workflow script: ", err)
        return
    }

    logger.Info("Executing 'workflow.js' by Node.js.")
    var cmd *exec.Cmd
    cmd = exec.Command("node", workflowFile)

    var stdout, stderr bytes.Buffer
    cmd.Stdout = &stdout
    cmd.Stderr = &stderr

    err = cmd.Run()

    exitStatusCode := 0

    if err != nil {
        logger.Error("Error during execution: %s, %s", err.Error(), string(stderr.Bytes()[:]))
        if exitError, ok := err.(*exec.ExitError); ok {
            if status, ok := exitError.Sys().(syscall.WaitStatus); ok {
                exitStatusCode = status.ExitStatus()
            }
        }

    } else {
        logger.Info("Execution standard output: %s", string(stdout.Bytes()[:]))
    }

    logger.Notice("Exit status code: %d", exitStatusCode)

    os.Exit(exitStatusCode)
}
<|endoftext|>