<commit_before>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/remote_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"io"
	"regexp"
	"strings"
)

func init() {
	Commands = append(Commands, &commandRemoteConfigure{})
}

type commandRemoteConfigure struct {
}

func (c *commandRemoteConfigure) Name() string {
	return "remote.configure"
}

func (c *commandRemoteConfigure) Help() string {
	return `remote storage configuration

	# see the current configurations
	remote.configure

	# set or update a configuration
	remote.configure -name=cloud1 -type=s3 -s3.access_key=xxx -s3.secret_key=yyy
	remote.configure -name=cloud2 -type=gcs -gcs.appCredentialsFile=~/service-account-file.json
	remote.configure -name=cloud3 -type=azure -azure.account_name=xxx -azure.account_key=yyy

	# delete one configuration
	remote.configure -delete -name=cloud1

`
}

var (
	isAlpha = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9]*$`).MatchString
)

func (c *commandRemoteConfigure) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	conf := &remote_pb.RemoteConf{}

	remoteConfigureCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	isDelete := remoteConfigureCommand.Bool("delete", false, "delete one remote storage by its name")

	remoteConfigureCommand.StringVar(&conf.Name, "name", "", "a short name to identify the remote storage")
	remoteConfigureCommand.StringVar(&conf.Type, "type", "s3", "[s3|gcs|azure|b2|aliyun|tencent|baidu|wasabi] storage type")

	remoteConfigureCommand.StringVar(&conf.S3AccessKey, "s3.access_key", "", "s3 access key")
	remoteConfigureCommand.StringVar(&conf.S3SecretKey, "s3.secret_key", "", "s3 secret key")
	remoteConfigureCommand.StringVar(&conf.S3Region, "s3.region", "us-east-2", "s3 region")
	remoteConfigureCommand.StringVar(&conf.S3Endpoint, "s3.endpoint", "", "endpoint for s3-compatible local object store")
	remoteConfigureCommand.StringVar(&conf.S3StorageClass, "s3.storage_class", "", "s3 storage class")
	remoteConfigureCommand.BoolVar(&conf.S3ForcePathStyle, "s3.force_path_style", true, "s3 force path style")

	remoteConfigureCommand.StringVar(&conf.GcsGoogleApplicationCredentials, "gcs.appCredentialsFile", "", "google cloud storage credentials file, default to use env GOOGLE_APPLICATION_CREDENTIALS")

	remoteConfigureCommand.StringVar(&conf.AzureAccountName, "azure.account_name", "", "azure account name, default to use env AZURE_STORAGE_ACCOUNT")
	remoteConfigureCommand.StringVar(&conf.AzureAccountKey, "azure.account_key", "", "azure account name, default to use env AZURE_STORAGE_ACCESS_KEY")

	remoteConfigureCommand.StringVar(&conf.BackblazeKeyId, "b2.key_id", "", "backblaze keyID")
	remoteConfigureCommand.StringVar(&conf.BackblazeApplicationKey, "b2.application_key", "", "backblaze applicationKey. Note that your Master Application Key will not work with the S3 Compatible API. You must create a new key that is eligible for use. For more information: https://help.backblaze.com/hc/en-us/articles/360047425453")
	remoteConfigureCommand.StringVar(&conf.BackblazeEndpoint, "b2.endpoint", "", "backblaze endpoint")

	remoteConfigureCommand.StringVar(&conf.AliyunAccessKey, "aliyun.access_key", "", "Aliyun access key, default to use env ALICLOUD_ACCESS_KEY_ID")
	remoteConfigureCommand.StringVar(&conf.AliyunSecretKey, "aliyun.secret_key", "", "Aliyun secret key, default to use env ALICLOUD_ACCESS_KEY_SECRET")
	remoteConfigureCommand.StringVar(&conf.AliyunEndpoint, "aliyun.endpoint", "", "Aliyun endpoint")
	remoteConfigureCommand.StringVar(&conf.AliyunRegion, "aliyun.region", "", "Aliyun region")

	remoteConfigureCommand.StringVar(&conf.TencentSecretId, "tencent.secret_id", "", "Tencent Secret Id, default to use env COS_SECRETID")
	remoteConfigureCommand.StringVar(&conf.TencentSecretKey, "tencent.secret_key", "", "Tencent secret key, default to use env COS_SECRETKEY")
	remoteConfigureCommand.StringVar(&conf.TencentEndpoint, "tencent.endpoint", "", "Tencent endpoint")
	remoteConfigureCommand.StringVar(&conf.TencentRegion, "tencent.region", "", "Tencent region")

	remoteConfigureCommand.StringVar(&conf.BaiduAccessKey, "baidu.access_key", "", "Baidu access key, default to use env BDCLOUD_ACCESS_KEY")
	remoteConfigureCommand.StringVar(&conf.BaiduSecretKey, "baidu.secret_key", "", "Baidu secret key, default to use env BDCLOUD_SECRET_KEY")
	remoteConfigureCommand.StringVar(&conf.BaiduEndpoint, "baidu.endpoint", "", "Baidu endpoint")
	remoteConfigureCommand.StringVar(&conf.BaiduRegion, "baidu.region", "", "Baidu region")

	remoteConfigureCommand.StringVar(&conf.WasabiAccessKey, "wasabi.access_key", "", "Wasabi access key")
	remoteConfigureCommand.StringVar(&conf.WasabiSecretKey, "wasabi.secret_key", "", "Wasabi secret key")
	remoteConfigureCommand.StringVar(&conf.WasabiEndpoint, "wasabi.endpoint", "", "Wasabi endpoint, see https://wasabi.com/wp-content/themes/wasabi/docs/API_Guide/index.html#t=topics%2Fapidiff-intro.htm")
	remoteConfigureCommand.StringVar(&conf.WasabiRegion, "wasabi.region", "", "Wasabi region")

	if err = remoteConfigureCommand.Parse(args); err != nil {
		return nil
	}

	if conf.Name == "" {
		return c.listExistingRemoteStorages(commandEnv, writer)
	}

	if !isAlpha(conf.Name) {
		return fmt.Errorf("only letters and numbers allowed in name: %v", conf.Name)
	}

	if *isDelete {
		return c.deleteRemoteStorage(commandEnv, writer, conf.Name)
	}

	return c.saveRemoteStorage(commandEnv, writer, conf)

}

func (c *commandRemoteConfigure) listExistingRemoteStorages(commandEnv *CommandEnv, writer io.Writer) error {

	return filer_pb.ReadDirAllEntries(commandEnv, util.FullPath(filer.DirectoryEtcRemote), "", func(entry *filer_pb.Entry, isLast bool) error {
		if len(entry.Content) == 0 {
			fmt.Fprintf(writer, "skipping %s\n", entry.Name)
			return nil
		}
		if !strings.HasSuffix(entry.Name, filer.REMOTE_STORAGE_CONF_SUFFIX) {
			return nil
		}
		conf := &remote_pb.RemoteConf{}

		if err := proto.Unmarshal(entry.Content, conf); err != nil {
			return fmt.Errorf("unmarshal %s/%s: %v", filer.DirectoryEtcRemote, entry.Name, err)
		}

		conf.S3SecretKey = strings.Repeat("*", len(conf.S3SecretKey))

		m := jsonpb.Marshaler{
			EmitDefaults: false,
			Indent:       "  ",
		}

		err := m.Marshal(writer, conf)
		fmt.Fprintln(writer)

		return err
	})

}

func (c *commandRemoteConfigure) deleteRemoteStorage(commandEnv *CommandEnv, writer io.Writer, storageName string) error {

	return commandEnv.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.DeleteEntryRequest{
			Directory:            filer.DirectoryEtcRemote,
			Name:                 storageName + filer.REMOTE_STORAGE_CONF_SUFFIX,
			IgnoreRecursiveError: false,
			IsDeleteData:         true,
			IsRecursive:          true,
			IsFromOtherCluster:   false,
			Signatures:           nil,
		}
		_, err := client.DeleteEntry(context.Background(), request)

		if err == nil {
			fmt.Fprintf(writer, "removed: %s\n", storageName)
		}

		return err

	})

}

func (c *commandRemoteConfigure) saveRemoteStorage(commandEnv *CommandEnv, writer io.Writer, conf *remote_pb.RemoteConf) error {

	data, err := proto.Marshal(conf)
	if err != nil {
		return err
	}

	if err = commandEnv.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {
		return filer.SaveInsideFiler(client, filer.DirectoryEtcRemote, conf.Name+filer.REMOTE_STORAGE_CONF_SUFFIX, data)
	}); err != nil && err != filer_pb.ErrNotFound {
		return err
	}

	return nil

}
<commit_msg>add aliyun example<commit_after>package shell

import (
	"context"
	"flag"
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/filer"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/remote_pb"
	"github.com/chrislusf/seaweedfs/weed/util"
	"github.com/golang/protobuf/jsonpb"
	"github.com/golang/protobuf/proto"
	"io"
	"regexp"
	"strings"
)

func init() {
	Commands = append(Commands, &commandRemoteConfigure{})
}

type commandRemoteConfigure struct {
}

func (c *commandRemoteConfigure) Name() string {
	return "remote.configure"
}

func (c *commandRemoteConfigure) Help() string {
	return `remote storage configuration

	# see the current configurations
	remote.configure

	# set or update a configuration
	remote.configure -name=cloud1 -type=s3 -s3.access_key=xxx -s3.secret_key=yyy
	remote.configure -name=cloud2 -type=gcs -gcs.appCredentialsFile=~/service-account-file.json
	remote.configure -name=cloud3 -type=azure -azure.account_name=xxx -azure.account_key=yyy
	remote.configure -name=cloud4 -type=aliyun -aliyun.access_key=xxx -aliyun.secret_key=yyy -aliyun.endpoint=oss-cn-shenzhen.aliyuncs.com -aliyun.region=cn-sehnzhen

	# delete one configuration
	remote.configure -delete -name=cloud1

`
}

var (
	isAlpha = regexp.MustCompile(`^[A-Za-z][A-Za-z0-9]*$`).MatchString
)

func (c *commandRemoteConfigure) Do(args []string, commandEnv *CommandEnv, writer io.Writer) (err error) {

	conf := &remote_pb.RemoteConf{}

	remoteConfigureCommand := flag.NewFlagSet(c.Name(), flag.ContinueOnError)
	isDelete := remoteConfigureCommand.Bool("delete", false, "delete one remote storage by its name")

	remoteConfigureCommand.StringVar(&conf.Name, "name", "", "a short name to identify the remote storage")
	remoteConfigureCommand.StringVar(&conf.Type, "type", "s3", "[s3|gcs|azure|b2|aliyun|tencent|baidu|wasabi] storage type")

	remoteConfigureCommand.StringVar(&conf.S3AccessKey, "s3.access_key", "", "s3 access key")
	remoteConfigureCommand.StringVar(&conf.S3SecretKey, "s3.secret_key", "", "s3 secret key")
	remoteConfigureCommand.StringVar(&conf.S3Region, "s3.region", "us-east-2", "s3 region")
	remoteConfigureCommand.StringVar(&conf.S3Endpoint, "s3.endpoint", "", "endpoint for s3-compatible local object store")
	remoteConfigureCommand.StringVar(&conf.S3StorageClass, "s3.storage_class", "", "s3 storage class")
	remoteConfigureCommand.BoolVar(&conf.S3ForcePathStyle, "s3.force_path_style", true, "s3 force path style")

	remoteConfigureCommand.StringVar(&conf.GcsGoogleApplicationCredentials, "gcs.appCredentialsFile", "", "google cloud storage credentials file, default to use env GOOGLE_APPLICATION_CREDENTIALS")

	remoteConfigureCommand.StringVar(&conf.AzureAccountName, "azure.account_name", "", "azure account name, default to use env AZURE_STORAGE_ACCOUNT")
	remoteConfigureCommand.StringVar(&conf.AzureAccountKey, "azure.account_key", "", "azure account name, default to use env AZURE_STORAGE_ACCESS_KEY")

	remoteConfigureCommand.StringVar(&conf.BackblazeKeyId, "b2.key_id", "", "backblaze keyID")
	remoteConfigureCommand.StringVar(&conf.BackblazeApplicationKey, "b2.application_key", "", "backblaze applicationKey. Note that your Master Application Key will not work with the S3 Compatible API. You must create a new key that is eligible for use. For more information: https://help.backblaze.com/hc/en-us/articles/360047425453")
	remoteConfigureCommand.StringVar(&conf.BackblazeEndpoint, "b2.endpoint", "", "backblaze endpoint")

	remoteConfigureCommand.StringVar(&conf.AliyunAccessKey, "aliyun.access_key", "", "Aliyun access key, default to use env ALICLOUD_ACCESS_KEY_ID")
	remoteConfigureCommand.StringVar(&conf.AliyunSecretKey, "aliyun.secret_key", "", "Aliyun secret key, default to use env ALICLOUD_ACCESS_KEY_SECRET")
	remoteConfigureCommand.StringVar(&conf.AliyunEndpoint, "aliyun.endpoint", "", "Aliyun endpoint")
	remoteConfigureCommand.StringVar(&conf.AliyunRegion, "aliyun.region", "", "Aliyun region")

	remoteConfigureCommand.StringVar(&conf.TencentSecretId, "tencent.secret_id", "", "Tencent Secret Id, default to use env COS_SECRETID")
	remoteConfigureCommand.StringVar(&conf.TencentSecretKey, "tencent.secret_key", "", "Tencent secret key, default to use env COS_SECRETKEY")
	remoteConfigureCommand.StringVar(&conf.TencentEndpoint, "tencent.endpoint", "", "Tencent endpoint")
	remoteConfigureCommand.StringVar(&conf.TencentRegion, "tencent.region", "", "Tencent region")

	remoteConfigureCommand.StringVar(&conf.BaiduAccessKey, "baidu.access_key", "", "Baidu access key, default to use env BDCLOUD_ACCESS_KEY")
	remoteConfigureCommand.StringVar(&conf.BaiduSecretKey, "baidu.secret_key", "", "Baidu secret key, default to use env BDCLOUD_SECRET_KEY")
	remoteConfigureCommand.StringVar(&conf.BaiduEndpoint, "baidu.endpoint", "", "Baidu endpoint")
	remoteConfigureCommand.StringVar(&conf.BaiduRegion, "baidu.region", "", "Baidu region")

	remoteConfigureCommand.StringVar(&conf.WasabiAccessKey, "wasabi.access_key", "", "Wasabi access key")
	remoteConfigureCommand.StringVar(&conf.WasabiSecretKey, "wasabi.secret_key", "", "Wasabi secret key")
	remoteConfigureCommand.StringVar(&conf.WasabiEndpoint, "wasabi.endpoint", "", "Wasabi endpoint, see https://wasabi.com/wp-content/themes/wasabi/docs/API_Guide/index.html#t=topics%2Fapidiff-intro.htm")
	remoteConfigureCommand.StringVar(&conf.WasabiRegion, "wasabi.region", "", "Wasabi region")

	if err = remoteConfigureCommand.Parse(args); err != nil {
		return nil
	}

	if conf.Name == "" {
		return c.listExistingRemoteStorages(commandEnv, writer)
	}

	if !isAlpha(conf.Name) {
		return fmt.Errorf("only letters and numbers allowed in name: %v", conf.Name)
	}

	if *isDelete {
		return c.deleteRemoteStorage(commandEnv, writer, conf.Name)
	}

	return c.saveRemoteStorage(commandEnv, writer, conf)

}

func (c *commandRemoteConfigure) listExistingRemoteStorages(commandEnv *CommandEnv, writer io.Writer) error {

	return filer_pb.ReadDirAllEntries(commandEnv, util.FullPath(filer.DirectoryEtcRemote), "", func(entry *filer_pb.Entry, isLast bool) error {
		if len(entry.Content) == 0 {
			fmt.Fprintf(writer, "skipping %s\n", entry.Name)
			return nil
		}
		if !strings.HasSuffix(entry.Name, filer.REMOTE_STORAGE_CONF_SUFFIX) {
			return nil
		}
		conf := &remote_pb.RemoteConf{}

		if err := proto.Unmarshal(entry.Content, conf); err != nil {
			return fmt.Errorf("unmarshal %s/%s: %v", filer.DirectoryEtcRemote, entry.Name, err)
		}

		conf.S3SecretKey = strings.Repeat("*", len(conf.S3SecretKey))

		m := jsonpb.Marshaler{
			EmitDefaults: false,
			Indent:       "  ",
		}

		err := m.Marshal(writer, conf)
		fmt.Fprintln(writer)

		return err
	})

}

func (c *commandRemoteConfigure) deleteRemoteStorage(commandEnv *CommandEnv, writer io.Writer, storageName string) error {

	return commandEnv.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {

		request := &filer_pb.DeleteEntryRequest{
			Directory:            filer.DirectoryEtcRemote,
			Name:                 storageName + filer.REMOTE_STORAGE_CONF_SUFFIX,
			IgnoreRecursiveError: false,
			IsDeleteData:         true,
			IsRecursive:          true,
			IsFromOtherCluster:   false,
			Signatures:           nil,
		}
		_, err := client.DeleteEntry(context.Background(), request)

		if err == nil {
			fmt.Fprintf(writer, "removed: %s\n", storageName)
		}

		return err

	})

}

func (c *commandRemoteConfigure) saveRemoteStorage(commandEnv *CommandEnv, writer io.Writer, conf *remote_pb.RemoteConf) error {

	data, err := proto.Marshal(conf)
	if err != nil {
		return err
	}

	if err = commandEnv.WithFilerClient(func(client filer_pb.SeaweedFilerClient) error {
		return filer.SaveInsideFiler(client, filer.DirectoryEtcRemote, conf.Name+filer.REMOTE_STORAGE_CONF_SUFFIX, data)
	}); err != nil && err != filer_pb.ErrNotFound {
		return err
	}

	return nil

}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"strings"

	"github.com/Cepave/open-falcon-backend/g"
	"github.com/spf13/cobra"
)

var Check = &cobra.Command{
	Use:   "check [Module ...]",
	Short: "Check the status of Open-Falcon modules",
	Long: `
Check if the specified Open-Falcon modules are running.

Modules:

  ` + "all " + strings.Join(g.AllModulesInOrder, " "),
	RunE: check,
}

func check(c *cobra.Command, args []string) error {
	if len(args) == 0 {
		return c.Usage()
	}
	for _, moduleName := range args {
		if !g.HasModule(moduleName) {
			return fmt.Errorf("%s doesn't exist\n", moduleName)
		}

		if g.IsRunning(moduleName) {
			fmt.Print("[", g.ModuleApps[moduleName], "] ", g.Pid(moduleName), "\n")
		} else {
			fmt.Print("[", g.ModuleApps[moduleName], "] down\n")
		}
	}
	return nil
}
<commit_msg>add module 'all' for cmd/check<commit_after>package cmd

import (
	"fmt"
	"strings"

	"github.com/Cepave/open-falcon-backend/g"
	"github.com/spf13/cobra"
)

var Check = &cobra.Command{
	Use:   "check [Module ...]",
	Short: "Check the status of Open-Falcon modules",
	Long: `
Check if the specified Open-Falcon modules are running.

Modules:

  ` + "all " + strings.Join(g.AllModulesInOrder, " "),
	RunE: check,
}

func check(c *cobra.Command, args []string) error {
	if len(args) == 0 {
		return c.Usage()
	}

	if len(args) == 1 && args[0] == "all" {
		for moduleName := range g.Modules {
			if g.IsRunning(moduleName) {
				fmt.Printf("%20s %10s %15s \n", g.ModuleApps[moduleName], "UP", g.Pid(moduleName))
			} else {
				fmt.Printf("%20s %10s %15s \n", g.ModuleApps[moduleName], "DOWN", "-")
			}
		}
	} else {
		for _, moduleName := range args {
			if !g.HasModule(moduleName) {
				return fmt.Errorf("%s doesn't exist\n", moduleName)
			}

			if g.IsRunning(moduleName) {
				fmt.Print("[", g.ModuleApps[moduleName], "] ", g.Pid(moduleName), "\n")
			} else {
				fmt.Print("[", g.ModuleApps[moduleName], "] down\n")
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"fmt"
	"github.com/smira/aptly/deb"
	"github.com/smira/commander"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
)

func aptlyGraph(cmd *commander.Command, args []string) error {
	var err error

	if len(args) != 0 {
		cmd.Usage()
		return commander.ErrCommandError
	}

	fmt.Printf("Generating graph...\n")
	graph, err := deb.BuildGraph(context.CollectionFactory())
	if err != nil {
		return err
	}

	buf := bytes.NewBufferString(graph.String())

	tempfile, err := ioutil.TempFile("", "aptly-graph")
	if err != nil {
		return err
	}
	tempfile.Close()
	os.Remove(tempfile.Name())

	format := context.Flags().Lookup("format").Value.String()

	tempfilename := tempfile.Name() + "." + format

	command := exec.Command("dot", "-T"+format, "-o"+tempfilename)
	command.Stderr = os.Stderr

	stdin, err := command.StdinPipe()
	if err != nil {
		return err
	}

	err = command.Start()
	if err != nil {
		return fmt.Errorf("unable to execute dot: %s (is graphviz package installed?)", err)
	}

	_, err = io.Copy(stdin, buf)
	if err != nil {
		return err
	}

	err = stdin.Close()
	if err != nil {
		return err
	}

	err = command.Wait()
	if err != nil {
		return err
	}

	fmt.Printf("Rendered to %s file: %s, trying to open it...\n", format, tempfilename)

	_ = exec.Command("open", tempfilename).Run()

	return err
}

func makeCmdGraph() *commander.Command {
	cmd := &commander.Command{
		Run:       aptlyGraph,
		UsageLine: "graph",
		Short:     "render graph of relationships",
		Long: `
Command graph displays relationship between mirrors, local repositories,
snapshots and published repositories using graphviz package to render
graph as an image.

Example:

  $ aptly graph
`,
	}

	cmd.Flag.String("format", "png", "render graph to specified format (png, svg, pdf, etc.)")

	return cmd
}
<commit_msg>aptly graph specific output filename<commit_after>package cmd

import (
	"bytes"
	"fmt"
	"github.com/smira/aptly/deb"
	"github.com/smira/commander"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
)

func aptlyGraph(cmd *commander.Command, args []string) error {
	var err error

	if len(args) != 0 {
		cmd.Usage()
		return commander.ErrCommandError
	}

	fmt.Printf("Generating graph...\n")
	graph, err := deb.BuildGraph(context.CollectionFactory())
	if err != nil {
		return err
	}

	buf := bytes.NewBufferString(graph.String())

	tempfile, err := ioutil.TempFile("", "aptly-graph")
	if err != nil {
		return err
	}
	tempfile.Close()
	os.Remove(tempfile.Name())

	format := context.Flags().Lookup("format").Value.String()
	output := context.Flags().Lookup("output").Value.String()

	tempfilename := tempfile.Name() + "." + format
	if output != "" {
		fmt.Printf("set file name option as  %s .  will output as %s.%s .\n", output, output, format)
		tempfilename = output + "." + format
	}

	command := exec.Command("dot", "-T"+format, "-o"+tempfilename)
	command.Stderr = os.Stderr

	stdin, err := command.StdinPipe()
	if err != nil {
		return err
	}

	err = command.Start()
	if err != nil {
		return fmt.Errorf("unable to execute dot: %s (is graphviz package installed?)", err)
	}

	_, err = io.Copy(stdin, buf)
	if err != nil {
		return err
	}

	err = stdin.Close()
	if err != nil {
		return err
	}

	err = command.Wait()
	if err != nil {
		return err
	}

	fmt.Printf("Rendered to %s file: %s, trying to open it...\n", format, tempfilename)

	_ = exec.Command("open", tempfilename).Run()

	return err
}

func makeCmdGraph() *commander.Command {
	cmd := &commander.Command{
		Run:       aptlyGraph,
		UsageLine: "graph",
		Short:     "render graph of relationships",
		Long: `
Command graph displays relationship between mirrors, local repositories,
snapshots and published repositories using graphviz package to render
graph as an image.

Example:

  $ aptly graph
`,
	}

	cmd.Flag.String("format", "png", "render graph to specified format (png, svg, pdf, etc.)")
	cmd.Flag.String("output", "", "specically output filename , default is output /tmp/aptlyxxxx.png")

	return cmd
}
<|endoftext|>


<commit_before>package cmd

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

	"github.com/Shopify/sarama"
	"github.com/bsm/sarama-cluster"
	"github.com/gocql/gocql"
	"github.com/golang/protobuf/proto"
)

const (
	maxFlushMsgs       = 5
	maxFlushFreqSecs   = 1
	connectTimeoutSecs = 10
	keepaliveSecs      = 5
)

type kafkaconfig struct {
	brokers      []string
	topic        string
	manager      EventBusManager
	maxOpenSends uint
}

func setupKafka() {
	kafkaConfig.brokers = strings.Split(kafkaBrokerStr, ",")
	if len(kafkaConfig.brokers) < 1 {
		log.Fatalf("At least one Kafka broker is required")
	}
	if kafkaConfig.topic == "" {
		log.Fatalf("Kafka topic is required")
	}
	kp, err := NewKafkaManager(kafkaConfig.brokers, kafkaConfig.topic, kafkaConfig.maxOpenSends, logger)
	if err != nil {
		log.Fatalf("Error creating Kafka producer: %v", err)
	}
	kafkaConfig.manager = kp
}

// EventBusProducer describes an object capable of publishing events somewhere
type EventBusProducer interface {
	PublishEvent(*BuildEvent) error
}

// EventBusConsumer describes an object cabable of subscribing to events somewhere
type EventBusConsumer interface {
	SubscribeToTopic(chan<- *BuildEvent, <-chan struct{}, gocql.UUID) error
}

// EventBusManager describes an object that can publish and subscribe to events somewhere
type EventBusManager interface {
	EventBusProducer
	EventBusConsumer
}

// KafkaManager handles sending event messages to the configured Kafka topic
type KafkaManager struct {
	ap           sarama.AsyncProducer
	topic        string
	brokers      []string
	consumerConf *cluster.Config
	logger       *log.Logger
}

// NewKafkaManager returns a new Kafka manager object
func NewKafkaManager(brokers []string, topic string, maxsends uint, logger *log.Logger) (*KafkaManager, error) {
	pconf := sarama.NewConfig()

	pconf.Net.MaxOpenRequests = int(maxsends)
	pconf.Net.DialTimeout = connTimeoutSecs * time.Second
	pconf.Net.ReadTimeout = connTimeoutSecs * time.Second
	pconf.Net.WriteTimeout = connTimeoutSecs * time.Second
	pconf.Net.KeepAlive = keepaliveSecs * time.Second

	pconf.Producer.Return.Errors = true
	pconf.Producer.Flush.Messages = maxFlushMsgs
	pconf.Producer.Flush.Frequency = maxFlushFreqSecs * time.Second

	asyncp, err := sarama.NewAsyncProducer(brokers, pconf)
	if err != nil {
		return nil, err
	}

	cconf := cluster.NewConfig()
	cconf.Net = pconf.Net
	cconf.Consumer.Return.Errors = true

	kp := &KafkaManager{
		ap:           asyncp,
		topic:        topic,
		brokers:      kafkaConfig.brokers,
		consumerConf: cconf,
		logger:       logger,
	}
	go kp.handlePErrors()
	return kp, nil
}

func (kp *KafkaManager) handlePErrors() {
	var kerr *sarama.ProducerError
	for {
		kerr = <-kp.ap.Errors()
		log.Printf("Kafka producer error: %v", kerr)
	}
}

// PublishEvent publishes a build event to the configured Kafka topic
func (kp *KafkaManager) PublishEvent(event *BuildEvent) error {
	id, err := gocql.ParseUUID(event.BuildId)
	if err != nil {
		return err
	}
	val, err := proto.Marshal(event)
	if err != nil {
		return fmt.Errorf("error marshaling protobuf: %v", err)
	}
	pmsg := &sarama.ProducerMessage{
		Topic: kp.topic,
		Key:   sarama.ByteEncoder(id.Bytes()), // Key is build ID to preserve event order (all events of a build go to the same partition)
		Value: sarama.ByteEncoder(val),
	}
	select { // don't block if Kafka is unavailable for some reason
	case kp.ap.Input() <- pmsg:
		return nil
	default:
		return fmt.Errorf("could not publish Kafka message: channel full")
	}
}

// SubscribeToTopic listens to the configured topic, filters by build_id and writes
// the resulting messages to output. When the subscribed build is finished
// output is closed. done is a signal from the caller to abort the stream subscription
func (kp *KafkaManager) SubscribeToTopic(output chan<- *BuildEvent, done <-chan struct{}, buildID gocql.UUID) error {
	// random group ID for each connection
	groupid, err := gocql.RandomUUID()
	if err != nil {
		return err
	}
	con, err := cluster.NewConsumer(kp.brokers, groupid.String(), []string{kp.topic}, kp.consumerConf)
	if err != nil {
		return err
	}
	handleConsumerErrors := func() {
		var err error
		for {
			err = <-con.Errors()
			if err == nil { // chan closed
				return
			}
			kp.logger.Printf("Kafka consumer error: %v", err)
		}
	}
	go handleConsumerErrors()
	go func() {
		defer close(output)
		defer con.Close()
		var err error
		var msg *sarama.ConsumerMessage
		var event *BuildEvent
		input := con.Messages()
		for {
			select {
			case <-done:
				kp.logger.Printf("SubscribeToTopic: aborting")
				return
			default:
				break
			}
			msg = <-input
			if msg == nil {
				return
			}
			if bytes.Equal(msg.Key, []byte(buildID[:])) {
				event = &BuildEvent{}
				err = proto.Unmarshal(msg.Value, event)
				if err != nil {
					kp.logger.Printf("%v: error unmarshaling event from Kafka stream: %v", buildID.String(), err)
					continue
				}
				output <- event
				if event.BuildFinished || event.EventError.IsError {
					return
				}
			}
		}
	}()
	return nil
}
<commit_msg>kafka: specify version<commit_after>package cmd

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

	"github.com/Shopify/sarama"
	"github.com/bsm/sarama-cluster"
	"github.com/gocql/gocql"
	"github.com/golang/protobuf/proto"
)

const (
	maxFlushMsgs       = 5
	maxFlushFreqSecs   = 1
	connectTimeoutSecs = 10
	keepaliveSecs      = 5
)

type kafkaconfig struct {
	brokers      []string
	topic        string
	manager      EventBusManager
	maxOpenSends uint
}

func setupKafka() {
	kafkaConfig.brokers = strings.Split(kafkaBrokerStr, ",")
	if len(kafkaConfig.brokers) < 1 {
		log.Fatalf("At least one Kafka broker is required")
	}
	if kafkaConfig.topic == "" {
		log.Fatalf("Kafka topic is required")
	}
	kp, err := NewKafkaManager(kafkaConfig.brokers, kafkaConfig.topic, kafkaConfig.maxOpenSends, logger)
	if err != nil {
		log.Fatalf("Error creating Kafka producer: %v", err)
	}
	kafkaConfig.manager = kp
}

// EventBusProducer describes an object capable of publishing events somewhere
type EventBusProducer interface {
	PublishEvent(*BuildEvent) error
}

// EventBusConsumer describes an object cabable of subscribing to events somewhere
type EventBusConsumer interface {
	SubscribeToTopic(chan<- *BuildEvent, <-chan struct{}, gocql.UUID) error
}

// EventBusManager describes an object that can publish and subscribe to events somewhere
type EventBusManager interface {
	EventBusProducer
	EventBusConsumer
}

// KafkaManager handles sending event messages to the configured Kafka topic
type KafkaManager struct {
	ap           sarama.AsyncProducer
	topic        string
	brokers      []string
	consumerConf *cluster.Config
	logger       *log.Logger
}

// NewKafkaManager returns a new Kafka manager object
func NewKafkaManager(brokers []string, topic string, maxsends uint, logger *log.Logger) (*KafkaManager, error) {
	pconf := sarama.NewConfig()
	pconf.Version = sarama.V0_10_0_0

	pconf.Net.MaxOpenRequests = int(maxsends)
	pconf.Net.DialTimeout = connTimeoutSecs * time.Second
	pconf.Net.ReadTimeout = connTimeoutSecs * time.Second
	pconf.Net.WriteTimeout = connTimeoutSecs * time.Second
	pconf.Net.KeepAlive = keepaliveSecs * time.Second

	pconf.Producer.Return.Errors = true
	pconf.Producer.Flush.Messages = maxFlushMsgs
	pconf.Producer.Flush.Frequency = maxFlushFreqSecs * time.Second

	asyncp, err := sarama.NewAsyncProducer(brokers, pconf)
	if err != nil {
		return nil, err
	}

	cconf := cluster.NewConfig()
	cconf.Net = pconf.Net
	cconf.Consumer.Return.Errors = true

	kp := &KafkaManager{
		ap:           asyncp,
		topic:        topic,
		brokers:      kafkaConfig.brokers,
		consumerConf: cconf,
		logger:       logger,
	}
	go kp.handlePErrors()
	return kp, nil
}

func (kp *KafkaManager) handlePErrors() {
	var kerr *sarama.ProducerError
	for {
		kerr = <-kp.ap.Errors()
		log.Printf("Kafka producer error: %v", kerr)
	}
}

// PublishEvent publishes a build event to the configured Kafka topic
func (kp *KafkaManager) PublishEvent(event *BuildEvent) error {
	id, err := gocql.ParseUUID(event.BuildId)
	if err != nil {
		return err
	}
	val, err := proto.Marshal(event)
	if err != nil {
		return fmt.Errorf("error marshaling protobuf: %v", err)
	}
	pmsg := &sarama.ProducerMessage{
		Topic: kp.topic,
		Key:   sarama.ByteEncoder(id.Bytes()), // Key is build ID to preserve event order (all events of a build go to the same partition)
		Value: sarama.ByteEncoder(val),
	}
	select { // don't block if Kafka is unavailable for some reason
	case kp.ap.Input() <- pmsg:
		return nil
	default:
		return fmt.Errorf("could not publish Kafka message: channel full")
	}
}

// SubscribeToTopic listens to the configured topic, filters by build_id and writes
// the resulting messages to output. When the subscribed build is finished
// output is closed. done is a signal from the caller to abort the stream subscription
func (kp *KafkaManager) SubscribeToTopic(output chan<- *BuildEvent, done <-chan struct{}, buildID gocql.UUID) error {
	// random group ID for each connection
	groupid, err := gocql.RandomUUID()
	if err != nil {
		return err
	}
	con, err := cluster.NewConsumer(kp.brokers, groupid.String(), []string{kp.topic}, kp.consumerConf)
	if err != nil {
		return err
	}
	handleConsumerErrors := func() {
		var err error
		for {
			err = <-con.Errors()
			if err == nil { // chan closed
				return
			}
			kp.logger.Printf("Kafka consumer error: %v", err)
		}
	}
	go handleConsumerErrors()
	go func() {
		defer close(output)
		defer con.Close()
		var err error
		var msg *sarama.ConsumerMessage
		var event *BuildEvent
		input := con.Messages()
		for {
			select {
			case <-done:
				kp.logger.Printf("SubscribeToTopic: aborting")
				return
			default:
				break
			}
			msg = <-input
			if msg == nil {
				return
			}
			if bytes.Equal(msg.Key, []byte(buildID[:])) {
				event = &BuildEvent{}
				err = proto.Unmarshal(msg.Value, event)
				if err != nil {
					kp.logger.Printf("%v: error unmarshaling event from Kafka stream: %v", buildID.String(), err)
					continue
				}
				output <- event
				if event.BuildFinished || event.EventError.IsError {
					return
				}
			}
		}
	}()
	return nil
}
<|endoftext|>


<commit_before>// Copyright © 2017-2018 Mikael Berthe <mikael@lilotux.net>
//
// Licensed under the MIT license.
// Please see the LICENSE file is this directory.

package cmd

import (
	"os"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

var mediaOpts struct {
	filePath string
}

// mediaCmd represents the media command
var mediaCmd = &cobra.Command{
	Use:     "media --file FILENAME",
	Aliases: []string{"upload"},
	Short:   "Upload a media attachment",
	//Long: `TBW...`,
	RunE: mediaRunE,
}

func init() {
	RootCmd.AddCommand(mediaCmd)

	mediaCmd.Flags().StringVar(&mediaOpts.filePath, "file", "", "Path of the media file")
	mediaCmd.MarkFlagRequired("file")
}

func mediaRunE(cmd *cobra.Command, args []string) error {
	opt := mediaOpts

	if opt.filePath == "" {
		return errors.New("no media file name provided")
	}

	if err := madonInit(true); err != nil {
		return err
	}

	attachment, err := gClient.UploadMedia(opt.filePath)
	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}

	p, err := getPrinter()
	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}
	return p.printObj(attachment)
}

// uploadFile uploads a media file and returns the attachment ID
func uploadFile(filePath string) (int64, error) {
	attachment, err := gClient.UploadMedia(filePath)
	if err != nil {
		return 0, err
	}
	if attachment == nil {
		return 0, nil
	}
	return attachment.ID, nil
}
<commit_msg>media: add --description and --focus flags for attachments<commit_after>// Copyright © 2017-2018 Mikael Berthe <mikael@lilotux.net>
//
// Licensed under the MIT license.
// Please see the LICENSE file is this directory.

package cmd

import (
	"os"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

var mediaOpts struct {
	filePath    string
	description string
	focus       string
}

// mediaCmd represents the media command
var mediaCmd = &cobra.Command{
	Use:     "media --file FILENAME",
	Aliases: []string{"upload"},
	Short:   "Upload a media attachment",
	//Long: `TBW...`,
	RunE: mediaRunE,
}

func init() {
	RootCmd.AddCommand(mediaCmd)

	mediaCmd.Flags().StringVar(&mediaOpts.filePath, "file", "", "Path of the media file")
	mediaCmd.MarkFlagRequired("file")

	mediaCmd.Flags().StringVar(&mediaOpts.description, "description", "", "Plain text description")
	mediaCmd.Flags().StringVar(&mediaOpts.focus, "focus", "", "Focal point")
}

func mediaRunE(cmd *cobra.Command, args []string) error {
	opt := mediaOpts

	if opt.filePath == "" {
		return errors.New("no media file name provided")
	}

	if err := madonInit(true); err != nil {
		return err
	}

	attachment, err := gClient.UploadMedia(opt.filePath, opt.description, opt.focus)
	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}

	p, err := getPrinter()
	if err != nil {
		errPrint("Error: %s", err.Error())
		os.Exit(1)
	}
	return p.printObj(attachment)
}

// uploadFile uploads a media file and returns the attachment ID
func uploadFile(filePath string) (int64, error) {
	attachment, err := gClient.UploadMedia(filePath, "", "")
	if err != nil {
		return 0, err
	}
	if attachment == nil {
		return 0, nil
	}
	return attachment.ID, nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"mime"
	"net"
	"net/http"
	"strings"

	"github.com/philips/grpc-gateway-example/pkg/ui/data/swagger"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"github.com/philips/go-bindata-assetfs"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	pb "github.com/philips/grpc-gateway-example/echopb"
)

// serveCmd represents the serve command
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Launches the example webserver on https://localhost:10000",
	Run: func(cmd *cobra.Command, args []string) {
		serve()
	},
}

func init() {
	RootCmd.AddCommand(serveCmd)
}

type myService struct{}

func (m *myService) Echo(c context.Context, s *pb.EchoMessage) (*pb.EchoMessage, error) {
	fmt.Printf("rpc request Echo(%q)\n", s.Value)
	return s, nil
}

func newServer() *myService {
	return new(myService)
}

// grpcHandlerFunc returns an http.Handler that delegates to grpcServer on incoming gRPC
// connections or otherHandler otherwise. Copied from cockroachdb.
func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// TODO(tamird): point to merged gRPC code rather than a PR.
		// This is a partial recreation of gRPC's internal checks https://github.com/grpc/grpc-go/pull/514/files#diff-95e9a25b738459a2d3030e1e6fa2a718R61
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcServer.ServeHTTP(w, r)
		} else {
			otherHandler.ServeHTTP(w, r)
		}
	})
}

func serveSwagger(mux *http.ServeMux) {
	mime.AddExtensionType(".svg", "image/svg+xml")

	// Expose files in third_party/swagger-ui/ on <host>/swagger-ui
	fileServer := http.FileServer(&assetfs.AssetFS{
		Asset:    swagger.Asset,
		AssetDir: swagger.AssetDir,
		Prefix:   "third_party/swagger-ui",
	})
	prefix := "/swagger-ui/"
	mux.Handle(prefix, http.StripPrefix(prefix, fileServer))
}

func serve() {
	opts := []grpc.ServerOption{
		grpc.Creds(credentials.NewClientTLSFromCert(demoCertPool, "localhost:10000"))}

	grpcServer := grpc.NewServer(opts...)
	pb.RegisterEchoServiceServer(grpcServer, newServer())
	ctx := context.Background()

	dcreds := credentials.NewTLS(&tls.Config{
		ServerName: demoAddr,
		RootCAs:    demoCertPool,
	})
	dopts := []grpc.DialOption{grpc.WithTransportCredentials(dcreds)}

	mux := http.NewServeMux()
	mux.HandleFunc("/swagger.json", func(w http.ResponseWriter, req *http.Request) {
		io.Copy(w, strings.NewReader(pb.Swagger))
	})

	gwmux := runtime.NewServeMux()
	err := pb.RegisterEchoServiceHandlerFromEndpoint(ctx, gwmux, demoAddr, dopts)
	if err != nil {
		fmt.Printf("serve: %v\n", err)
		return
	}

	mux.Handle("/", gwmux)
	serveSwagger(mux)

	conn, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		panic(err)
	}

	srv := &http.Server{
		Addr:    demoAddr,
		Handler: grpcHandlerFunc(grpcServer, mux),
		TLSConfig: &tls.Config{
			Certificates: []tls.Certificate{*demoKeyPair},
		},
	}

	fmt.Printf("grpc on port: %d\n", port)
	err = srv.Serve(tls.NewListener(conn, srv.TLSConfig))

	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}

	return
}
<commit_msg>cmd: serve: fix go 1.7 regression<commit_after>package cmd

import (
	"crypto/tls"
	"fmt"
	"io"
	"log"
	"mime"
	"net"
	"net/http"
	"strings"

	"github.com/philips/grpc-gateway-example/pkg/ui/data/swagger"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"
	"github.com/philips/go-bindata-assetfs"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	pb "github.com/philips/grpc-gateway-example/echopb"
)

// serveCmd represents the serve command
var serveCmd = &cobra.Command{
	Use:   "serve",
	Short: "Launches the example webserver on https://localhost:10000",
	Run: func(cmd *cobra.Command, args []string) {
		serve()
	},
}

func init() {
	RootCmd.AddCommand(serveCmd)
}

type myService struct{}

func (m *myService) Echo(c context.Context, s *pb.EchoMessage) (*pb.EchoMessage, error) {
	fmt.Printf("rpc request Echo(%q)\n", s.Value)
	return s, nil
}

func newServer() *myService {
	return new(myService)
}

// grpcHandlerFunc returns an http.Handler that delegates to grpcServer on incoming gRPC
// connections or otherHandler otherwise. Copied from cockroachdb.
func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// TODO(tamird): point to merged gRPC code rather than a PR.
		// This is a partial recreation of gRPC's internal checks https://github.com/grpc/grpc-go/pull/514/files#diff-95e9a25b738459a2d3030e1e6fa2a718R61
		if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
			grpcServer.ServeHTTP(w, r)
		} else {
			otherHandler.ServeHTTP(w, r)
		}
	})
}

func serveSwagger(mux *http.ServeMux) {
	mime.AddExtensionType(".svg", "image/svg+xml")

	// Expose files in third_party/swagger-ui/ on <host>/swagger-ui
	fileServer := http.FileServer(&assetfs.AssetFS{
		Asset:    swagger.Asset,
		AssetDir: swagger.AssetDir,
		Prefix:   "third_party/swagger-ui",
	})
	prefix := "/swagger-ui/"
	mux.Handle(prefix, http.StripPrefix(prefix, fileServer))
}

func serve() {
	opts := []grpc.ServerOption{
		grpc.Creds(credentials.NewClientTLSFromCert(demoCertPool, "localhost:10000"))}

	grpcServer := grpc.NewServer(opts...)
	pb.RegisterEchoServiceServer(grpcServer, newServer())
	ctx := context.Background()

	dcreds := credentials.NewTLS(&tls.Config{
		ServerName: demoAddr,
		RootCAs:    demoCertPool,
	})
	dopts := []grpc.DialOption{grpc.WithTransportCredentials(dcreds)}

	mux := http.NewServeMux()
	mux.HandleFunc("/swagger.json", func(w http.ResponseWriter, req *http.Request) {
		io.Copy(w, strings.NewReader(pb.Swagger))
	})

	gwmux := runtime.NewServeMux()
	err := pb.RegisterEchoServiceHandlerFromEndpoint(ctx, gwmux, demoAddr, dopts)
	if err != nil {
		fmt.Printf("serve: %v\n", err)
		return
	}

	mux.Handle("/", gwmux)
	serveSwagger(mux)

	conn, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		panic(err)
	}

	srv := &http.Server{
		Addr:    demoAddr,
		Handler: grpcHandlerFunc(grpcServer, mux),
		TLSConfig: &tls.Config{
			Certificates: []tls.Certificate{*demoKeyPair},
			NextProtos:   []string{"h2"},
		},
	}

	fmt.Printf("grpc on port: %d\n", port)
	err = srv.Serve(tls.NewListener(conn, srv.TLSConfig))

	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"math"
	"net"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/mux"
	_ "github.com/mattn/go-sqlite3"
	uuid "github.com/satori/go.uuid"
)

const (
	actionAllow  = "allow"
	actionBlock  = "block"
	actionIgnore = "ignore"

	saneTime = "2016-01-02 15:04:05 MST"
)

var (
	templates = flag.String("templates", ".", "Template dir")
	staticDir = flag.String("static", ".", "Static dir")
	addr      = flag.String("addr", ":8080", "Address to listen to.")
	squidLog  = flag.String("squidlog", "", "Path to squid log.")
	dbFile    = flag.String("db", "", "sqlite database.")

	db *sql.DB
)

type aclID string
type acl struct {
	ACLID   aclID
	Comment string
}
type ruleID string
type rule struct {
	RuleID  ruleID
	Type    string
	Value   string
	Action  string
	Comment string
}

func host2domain(h string) string {
	tld := 1
	for _, d := range []string{
		".pp.se",
		".co.uk",
		".gov.uk",
	} {
		if strings.HasSuffix(h, d) {
			tld = 2
			break
		}
	}
	s := strings.Split(h, ".")
	if len(s) <= tld {
		return h
	}
	return strings.Join(s[len(s)-tld-1:], ".")
}

func rootHandler(r *http.Request) (template.HTML, error) {
	tmpl := template.Must(template.ParseFiles(path.Join(*templates, "main.html")))
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, nil); err != nil {
		return "", fmt.Errorf("template execute fail: %v", err)
	}
	return template.HTML(buf.String()), nil
}

func openDB() {
	var err error
	db, err = sql.Open("sqlite3", *dbFile)
	if err != nil {
		log.Fatalf("Failed to open database %q: %v", *dbFile, err)
	}
	if _, err := db.Exec("PRAGMA foreign_keys = ON"); err != nil {
		log.Fatalf("Failed to turn on foreign keys")
	}
}

func allowHandler(w http.ResponseWriter, r *http.Request) {
	typ := r.FormValue("type")
	value := r.FormValue("value")
	if typ == "" || value == "" {
		http.Error(w, "Missing parameters", http.StatusBadRequest)
		return
	}
	// TODO: look up 'misc'.
	aclID := "db4c935f-fc5b-4868-ab32-c80459159c3e"
	if err := func() error {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		id := uuid.NewV4()
		log.Printf("Adding rule %q", id)
		if _, err := tx.Exec(`INSERT INTO rules(rule_id, action, type, value) VALUES(?,?,?,?)`, id, actionAllow, typ, value); err != nil {
			return err
		}
		if _, err := tx.Exec(`INSERT INTO aclrules(acl_id, rule_id) VALUES(?, ?)`, aclID, id); err != nil {
			return err
		}
		if err := tx.Commit(); err != nil {
			return err
		}
		return nil
	}(); err != nil {
		log.Printf("Database trouble: %v", err)
		http.Error(w, "DB problems", http.StatusInternalServerError)
	}
}

func reverse(s []string) []string {
	l := len(s)
	o := make([]string, l, l)
	for i, j := 0, l-1; i < j; i, j = i+1, j-1 {
		o[i], o[j] = s[j], s[i]
	}
	return o
}

func errWrap(f func(*http.Request) (template.HTML, error)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		tmpl := template.Must(template.ParseFiles(path.Join(*templates, "page.html")))
		h, err := f(r)
		if err != nil {
			log.Printf("Error in HTTP handler: %v", err)
			http.Error(w, "Internal error", http.StatusInternalServerError)
			return
		}
		if err := tmpl.Execute(w, struct{ Content template.HTML }{Content: h}); err != nil {
			log.Printf("Error in main handler: %v", err)
			http.Error(w, "Internal error", http.StatusInternalServerError)
			return
		}
	}
}

func aclHandler(r *http.Request) (template.HTML, error) {
	current := aclID(mux.Vars(r)["aclID"])

	data := struct {
		ACLs []acl

		Current acl
		Rules   []rule
	}{}
	{
		rows, err := db.Query(`SELECT acl_id, comment FROM acls ORDER BY comment`)
		if err != nil {
			return "", err
		}
		defer rows.Close()

		for rows.Next() {
			var s string
			var c sql.NullString
			if err := rows.Scan(&s, &c); err != nil {
				return "", err
			}
			e := acl{
				ACLID:   aclID(s),
				Comment: c.String,
			}
			if current == e.ACLID {
				data.Current = e
			}
			data.ACLs = append(data.ACLs, e)
		}
	}

	if len(current) > 0 {
		r, err := loadACL(current)
		if err != nil {
			return "", err
		}
		data.Rules = r
	}

	tmpl := template.Must(template.New("acl.html").Funcs(template.FuncMap{
		"aclIDEQ": func(a, b aclID) bool { return a == b },
	}).ParseFiles(path.Join(*templates, "acl.html")))
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, &data); err != nil {
		return "", fmt.Errorf("template execute fail: %v", err)
	}
	return template.HTML(buf.String()), nil
}

func loadACL(id aclID) ([]rule, error) {
	rows, err := db.Query(`
SELECT rules.rule_id, rules.type, rules.value, rules.action, rules.comment
FROM aclrules
JOIN rules ON aclrules.rule_id=rules.rule_id
WHERE aclrules.acl_id=?
ORDER BY rules.comment`, string(id))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var rules []rule
	for rows.Next() {
		var e rule
		var s string
		var c sql.NullString
		if err := rows.Scan(&s, &e.Type, &e.Value, &e.Action, &c); err != nil {
			return nil, err
		}
		e.RuleID = ruleID(s)
		e.Comment = c.String
		rules = append(rules, e)
	}
	return rules, nil
}

func tailLogHandler(w http.ResponseWriter, r *http.Request) {
	b, err := ioutil.ReadFile(*squidLog)
	if err != nil {
		log.Printf("Failed to read squid log: %v", err)
		return
	}
	lines := reverse(strings.Split(string(b), "\n"))
	const n = 30
	if len(lines) > n {
		lines = lines[:n]
	}
	type logEntry struct {
		Time   string
		Client string
		Method string
		Domain string
		Host   string
		Path   string
		URL    string
	}
	entries := []logEntry{}
	//                        time        ms    client     DENIED    size   method  URL           HIER    type
	re := regexp.MustCompile(`([0-9.]+)\s+\d+\s+([^\s]+)\s+([^\s]+)\s+\d+\s+(\w+)\s+([^\s]+)\s+-\s[^\s]+\s([^\s]+)`)
	for _, l := range lines {
		if len(l) == 0 {
			continue
		}
		s := re.FindStringSubmatch(l)
		if len(s) == 0 {
			log.Printf("Bad log line: %q", l)
			continue
		}
		var host, p string
		u := s[5]
		if ur, err := url.Parse(u); strings.Contains(u, "/") && err == nil && ur.Scheme != "" {
			host = ur.Host
			p = ur.Path
		} else {
			host, _, err = net.SplitHostPort(u)
			if err != nil {
				host = u
			}
		}

		ts, err := strconv.ParseFloat(s[1], 64)
		if err != nil {
			log.Printf("Failed to parse epoch time %q: %v", s[1], err)
		}

		entries = append(entries, logEntry{
			Time:   time.Unix(int64(ts), int64(1e9*(ts-math.Trunc(ts)))).UTC().Format(saneTime),
			Client: s[2],
			Method: s[4],
			Domain: "." + host2domain(host),
			Host:   host,
			Path:   p,
			URL:    u,
		})
	}
	b, err = json.Marshal(entries)
	if err != nil {
		panic(err)
	}
	if _, err := w.Write(b); err != nil {
		log.Printf("Failed writing tail stuff: %v", err)
	}
}

func main() {
	flag.Parse()
	if flag.NArg() > 0 {
		log.Fatalf("Extra args on cmdline: %q", flag.Args())
	}
	openDB()
	log.Printf("Running...")
	r := mux.NewRouter()
	r.HandleFunc("/", errWrap(rootHandler)).Methods("GET", "HEAD")
	r.HandleFunc("/acl/", errWrap(aclHandler)).Methods("GET", "HEAD")
	r.HandleFunc("/acl/{aclID}", errWrap(aclHandler)).Methods("GET", "HEAD")
	r.HandleFunc("/ajax/allow", allowHandler).Methods("POST")
	r.HandleFunc("/ajax/tail-log", tailLogHandler).Methods("GET")

	fs := http.FileServer(http.Dir(*staticDir))
	http.Handle("/static/", http.StripPrefix("/static/", fs))
	http.Handle("/", r)

	log.Fatal(http.ListenAndServe(*addr, nil))
}
<commit_msg>Check more DB errors<commit_after>package main

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io/ioutil"
	"log"
	"math"
	"net"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/gorilla/mux"
	_ "github.com/mattn/go-sqlite3"
	uuid "github.com/satori/go.uuid"
)

const (
	actionAllow  = "allow"
	actionBlock  = "block"
	actionIgnore = "ignore"

	saneTime = "2016-01-02 15:04:05 MST"
)

var (
	templates = flag.String("templates", ".", "Template dir")
	staticDir = flag.String("static", ".", "Static dir")
	addr      = flag.String("addr", ":8080", "Address to listen to.")
	squidLog  = flag.String("squidlog", "", "Path to squid log.")
	dbFile    = flag.String("db", "", "sqlite database.")

	db *sql.DB
)

type aclID string
type acl struct {
	ACLID   aclID
	Comment string
}
type ruleID string
type rule struct {
	RuleID  ruleID
	Type    string
	Value   string
	Action  string
	Comment string
}

func host2domain(h string) string {
	tld := 1
	for _, d := range []string{
		".pp.se",
		".co.uk",
		".gov.uk",
	} {
		if strings.HasSuffix(h, d) {
			tld = 2
			break
		}
	}
	s := strings.Split(h, ".")
	if len(s) <= tld {
		return h
	}
	return strings.Join(s[len(s)-tld-1:], ".")
}

func rootHandler(r *http.Request) (template.HTML, error) {
	tmpl := template.Must(template.ParseFiles(path.Join(*templates, "main.html")))
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, nil); err != nil {
		return "", fmt.Errorf("template execute fail: %v", err)
	}
	return template.HTML(buf.String()), nil
}

func openDB() {
	var err error
	db, err = sql.Open("sqlite3", *dbFile)
	if err != nil {
		log.Fatalf("Failed to open database %q: %v", *dbFile, err)
	}
	if _, err := db.Exec("PRAGMA foreign_keys = ON"); err != nil {
		log.Fatalf("Failed to turn on foreign keys")
	}
}

func allowHandler(w http.ResponseWriter, r *http.Request) {
	typ := r.FormValue("type")
	value := r.FormValue("value")
	if typ == "" || value == "" {
		http.Error(w, "Missing parameters", http.StatusBadRequest)
		return
	}
	// TODO: look up 'misc'.
	aclID := "db4c935f-fc5b-4868-ab32-c80459159c3e"
	if err := func() error {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		id := uuid.NewV4()
		log.Printf("Adding rule %q", id)
		if _, err := tx.Exec(`INSERT INTO rules(rule_id, action, type, value) VALUES(?,?,?,?)`, id, actionAllow, typ, value); err != nil {
			return err
		}
		if _, err := tx.Exec(`INSERT INTO aclrules(acl_id, rule_id) VALUES(?, ?)`, aclID, id); err != nil {
			return err
		}
		if err := tx.Commit(); err != nil {
			return err
		}
		return nil
	}(); err != nil {
		log.Printf("Database trouble: %v", err)
		http.Error(w, "DB problems", http.StatusInternalServerError)
	}
}

func reverse(s []string) []string {
	l := len(s)
	o := make([]string, l, l)
	for i, j := 0, l-1; i < j; i, j = i+1, j-1 {
		o[i], o[j] = s[j], s[i]
	}
	return o
}

func errWrap(f func(*http.Request) (template.HTML, error)) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		tmpl := template.Must(template.ParseFiles(path.Join(*templates, "page.html")))
		h, err := f(r)
		if err != nil {
			log.Printf("Error in HTTP handler: %v", err)
			http.Error(w, "Internal error", http.StatusInternalServerError)
			return
		}
		if err := tmpl.Execute(w, struct{ Content template.HTML }{Content: h}); err != nil {
			log.Printf("Error in main handler: %v", err)
			http.Error(w, "Internal error", http.StatusInternalServerError)
			return
		}
	}
}

func aclHandler(r *http.Request) (template.HTML, error) {
	current := aclID(mux.Vars(r)["aclID"])

	data := struct {
		ACLs []acl

		Current acl
		Rules   []rule
	}{}
	{
		rows, err := db.Query(`SELECT acl_id, comment FROM acls ORDER BY comment`)
		if err != nil {
			return "", err
		}
		defer rows.Close()

		for rows.Next() {
			var s string
			var c sql.NullString
			if err := rows.Scan(&s, &c); err != nil {
				return "", err
			}
			e := acl{
				ACLID:   aclID(s),
				Comment: c.String,
			}
			if current == e.ACLID {
				data.Current = e
			}
			data.ACLs = append(data.ACLs, e)
		}
		if err := rows.Err(); err != nil {
			return "", err
		}
	}

	if len(current) > 0 {
		r, err := loadACL(current)
		if err != nil {
			return "", err
		}
		data.Rules = r
	}

	tmpl := template.Must(template.New("acl.html").Funcs(template.FuncMap{
		"aclIDEQ": func(a, b aclID) bool { return a == b },
	}).ParseFiles(path.Join(*templates, "acl.html")))
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, &data); err != nil {
		return "", fmt.Errorf("template execute fail: %v", err)
	}
	return template.HTML(buf.String()), nil
}

func loadACL(id aclID) ([]rule, error) {
	rows, err := db.Query(`
SELECT rules.rule_id, rules.type, rules.value, rules.action, rules.comment
FROM aclrules
JOIN rules ON aclrules.rule_id=rules.rule_id
WHERE aclrules.acl_id=?
ORDER BY rules.comment`, string(id))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var rules []rule
	for rows.Next() {
		var e rule
		var s string
		var c sql.NullString
		if err := rows.Scan(&s, &e.Type, &e.Value, &e.Action, &c); err != nil {
			return nil, err
		}
		e.RuleID = ruleID(s)
		e.Comment = c.String
		rules = append(rules, e)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	return rules, nil
}

func tailLogHandler(w http.ResponseWriter, r *http.Request) {
	b, err := ioutil.ReadFile(*squidLog)
	if err != nil {
		log.Printf("Failed to read squid log: %v", err)
		return
	}
	lines := reverse(strings.Split(string(b), "\n"))
	const n = 30
	if len(lines) > n {
		lines = lines[:n]
	}
	type logEntry struct {
		Time   string
		Client string
		Method string
		Domain string
		Host   string
		Path   string
		URL    string
	}
	entries := []logEntry{}
	//                        time        ms    client     DENIED    size   method  URL           HIER    type
	re := regexp.MustCompile(`([0-9.]+)\s+\d+\s+([^\s]+)\s+([^\s]+)\s+\d+\s+(\w+)\s+([^\s]+)\s+-\s[^\s]+\s([^\s]+)`)
	for _, l := range lines {
		if len(l) == 0 {
			continue
		}
		s := re.FindStringSubmatch(l)
		if len(s) == 0 {
			log.Printf("Bad log line: %q", l)
			continue
		}
		var host, p string
		u := s[5]
		if ur, err := url.Parse(u); strings.Contains(u, "/") && err == nil && ur.Scheme != "" {
			host = ur.Host
			p = ur.Path
		} else {
			host, _, err = net.SplitHostPort(u)
			if err != nil {
				host = u
			}
		}

		ts, err := strconv.ParseFloat(s[1], 64)
		if err != nil {
			log.Printf("Failed to parse epoch time %q: %v", s[1], err)
		}

		entries = append(entries, logEntry{
			Time:   time.Unix(int64(ts), int64(1e9*(ts-math.Trunc(ts)))).UTC().Format(saneTime),
			Client: s[2],
			Method: s[4],
			Domain: "." + host2domain(host),
			Host:   host,
			Path:   p,
			URL:    u,
		})
	}
	b, err = json.Marshal(entries)
	if err != nil {
		panic(err)
	}
	if _, err := w.Write(b); err != nil {
		log.Printf("Failed writing tail stuff: %v", err)
	}
}

func main() {
	flag.Parse()
	if flag.NArg() > 0 {
		log.Fatalf("Extra args on cmdline: %q", flag.Args())
	}
	openDB()
	log.Printf("Running...")
	r := mux.NewRouter()
	r.HandleFunc("/", errWrap(rootHandler)).Methods("GET", "HEAD")
	r.HandleFunc("/acl/", errWrap(aclHandler)).Methods("GET", "HEAD")
	r.HandleFunc("/acl/{aclID}", errWrap(aclHandler)).Methods("GET", "HEAD")
	r.HandleFunc("/ajax/allow", allowHandler).Methods("POST")
	r.HandleFunc("/ajax/tail-log", tailLogHandler).Methods("GET")

	fs := http.FileServer(http.Dir(*staticDir))
	http.Handle("/static/", http.StripPrefix("/static/", fs))
	http.Handle("/", r)

	log.Fatal(http.ListenAndServe(*addr, nil))
}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	flag_ "flag"
	"fmt"
	"github.com/cossacklabs/acra/keystore"
	"github.com/cossacklabs/acra/utils"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
)

var (
	config     = flag_.String("config", "", "path to config")
	dumpconfig = flag_.Bool("dumpconfig", false, "dump config")
)

const (
	LOG_DEBUG = iota
	LOG_VERBOSE
	LOG_DISCARD
)

func init() {
	// override default usage message by ours
	flag_.CommandLine.Usage = PrintDefaults
}

func ValidateClientId(clientId string) {
	if !keystore.ValidateId([]byte(clientId)) {
		fmt.Println("Error: invalid client id, user only digits, letters and '_', '-', ' ' characters")
		os.Exit(1)
	}
}

func isZeroValue(flag *flag_.Flag, value string) bool {
	/* took from flag/flag.go */

	// Build a zero value of the flag's Value type, and see if the
	// result of calling its String method equals the value passed in.
	// This works unless the Value type is itself an interface type.
	typ := reflect.TypeOf(flag.Value)
	var z reflect.Value
	if typ.Kind() == reflect.Ptr {
		z = reflect.New(typ.Elem())
	} else {
		z = reflect.Zero(typ)
	}
	if value == z.Interface().(flag_.Value).String() {
		return true
	}

	switch value {
	case "false":
		return true
	case "":
		return true
	case "0":
		return true
	}
	return false
}

func PrintDefaults() {
	/* took from flag/flag.go and overrided arg display format (-/--) */
	flag_.CommandLine.VisitAll(func(flag *flag_.Flag) {
		var s string
		if len(flag.Name) > 2 {
			s = fmt.Sprintf("  --%s", flag.Name) // Two spaces before -; see next two comments.
		} else {
			s = fmt.Sprintf("  -%s", flag.Name) // Two spaces before -; see next two comments.
		}
		// Boolean flags of one ASCII letter are so common we
		// treat them specially, putting their usage on the same line.
		if len(s) <= 4 {
			// space, space, '-', 'x'.
			s += "\t"
		} else {
			// Four spaces before the tab triggers good alignment
			// for both 4- and 8-space tab stops.
			s += "\n    \t"
		}
		s += flag.Usage
		if !isZeroValue(flag, flag.DefValue) {
			getter, ok := flag.Value.(flag_.Getter)
			if !ok {
				return
			}

			if _, ok := getter.Get().(string); ok {
				// put quotes on the value
				s += fmt.Sprintf(" (default %q)", flag.DefValue)
			} else {
				s += fmt.Sprintf(" (default %v)", flag.DefValue)
			}
		}
		fmt.Fprint(os.Stderr, s, "\n")
	})
}

func GenerateYaml(output io.Writer) {
	flag_.CommandLine.VisitAll(func(flag *flag_.Flag) {
		s := fmt.Sprintf("# %v\n%v: %v\n", flag.Usage, flag.Name, flag.DefValue)
		fmt.Fprint(output, s, "\n")
	})
}

func Parse(configPath string) error {
	/*load from yaml config and cli. if dumpconfig option pass than generate config and exit*/

	// first parse using bultin flag
	err := flag_.CommandLine.Parse(os.Args[1:])
	if err != nil {
		return err
	}

	if *config != "" {
		configPath = *config
	}
	var args []string
	// parse yaml and add params that wasn't passed from cli
	if configPath != "" {
		configPath, err := utils.AbsPath(configPath)
		if err != nil {
			return err
		}
		exists, err := utils.FileExists(configPath)
		if err != nil {
			return err
		}
		if exists {
			data, err := ioutil.ReadFile(configPath)
			if err != nil {
				return err
			}
			yamlConfig := map[string]interface{}{}
			err = yaml.Unmarshal([]byte(data), &yamlConfig)
			if err != nil {
				return err
			}
			setArgs := make(map[string]bool)
			flag_.Visit(func(flag *flag_.Flag) {
				setArgs[flag.Name] = true
			})
			// generate args list for flag.Parse as it was from cli args
			args = make([]string, 0)
			flag_.VisitAll(func(flag *flag_.Flag) {
				// generate only args that wasn't set from cli
				if _, alreadySet := setArgs[flag.Name]; !alreadySet {
					if value, yamlOk := yamlConfig[flag.Name]; yamlOk {
						if value != nil {
							args = append(args, fmt.Sprintf("--%v=%v", flag.Name, value))
						}
					}
				}
			})
		}
	}
	// set options from config that wasn't set by cli
	err = flag_.CommandLine.Parse(args)
	if err != nil {
		return err
	}
	if *dumpconfig {
		var absPath string
		if *config == "" {
			absPath, err = utils.AbsPath(configPath)
			if err != nil {
				return err
			}
		} else {
			absPath, err = utils.AbsPath(*config)
			if err != nil {
				return err
			}
		}

		dirPath := filepath.Dir(absPath)
		err = os.MkdirAll(dirPath, 0744)
		if err != nil {
			return err
		}

		file, err := os.Create(absPath)
		if err != nil {
			return err
		}
		defer file.Close()

		GenerateYaml(file)
		os.Exit(0)
	}
	return nil
}

var DEBUG_PREFIX = []byte("Debug: ")

type NotDebugWriter struct {
	writer io.Writer
}

func NewNotDebugWriter(writer io.Writer) *NotDebugWriter {
	return &NotDebugWriter{writer: writer}
}
func (wr *NotDebugWriter) Write(p []byte) (int, error) {
	if bytes.Contains(p, DEBUG_PREFIX) {
		return 0, nil
	}
	return wr.writer.Write(p)
}

func SetLogLevel(level int) {
	if level == LOG_DEBUG {
		log.SetOutput(os.Stdout)
	} else if level == LOG_VERBOSE {
		log.SetOutput(NewNotDebugWriter(os.Stdout))
	} else if level == LOG_DISCARD {
		log.SetOutput(ioutil.Discard)
	} else {
		panic(fmt.Sprintf("Incorrect log level - %v", level))
	}
}
<commit_msg>extend error message<commit_after>package cmd

import (
	"bytes"
	flag_ "flag"
	"fmt"
	"github.com/cossacklabs/acra/keystore"
	"github.com/cossacklabs/acra/utils"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"reflect"
)

var (
	config     = flag_.String("config", "", "path to config")
	dumpconfig = flag_.Bool("dumpconfig", false, "dump config")
)

const (
	LOG_DEBUG = iota
	LOG_VERBOSE
	LOG_DISCARD
)

func init() {
	// override default usage message by ours
	flag_.CommandLine.Usage = PrintDefaults
}

func ValidateClientId(clientId string) {
	if !keystore.ValidateId([]byte(clientId)) {
		fmt.Printf("Error: invalid client id,  %d <= len(client id) <= %d, only digits, letters and '_', '-', ' ' characters\n",
			keystore.MIN_CLIENT_ID_LENGTH, keystore.MAX_CLIENT_ID_LENGTH)
		os.Exit(1)
	}
}

func isZeroValue(flag *flag_.Flag, value string) bool {
	/* took from flag/flag.go */

	// Build a zero value of the flag's Value type, and see if the
	// result of calling its String method equals the value passed in.
	// This works unless the Value type is itself an interface type.
	typ := reflect.TypeOf(flag.Value)
	var z reflect.Value
	if typ.Kind() == reflect.Ptr {
		z = reflect.New(typ.Elem())
	} else {
		z = reflect.Zero(typ)
	}
	if value == z.Interface().(flag_.Value).String() {
		return true
	}

	switch value {
	case "false":
		return true
	case "":
		return true
	case "0":
		return true
	}
	return false
}

func PrintDefaults() {
	/* took from flag/flag.go and overrided arg display format (-/--) */
	flag_.CommandLine.VisitAll(func(flag *flag_.Flag) {
		var s string
		if len(flag.Name) > 2 {
			s = fmt.Sprintf("  --%s", flag.Name) // Two spaces before -; see next two comments.
		} else {
			s = fmt.Sprintf("  -%s", flag.Name) // Two spaces before -; see next two comments.
		}
		// Boolean flags of one ASCII letter are so common we
		// treat them specially, putting their usage on the same line.
		if len(s) <= 4 {
			// space, space, '-', 'x'.
			s += "\t"
		} else {
			// Four spaces before the tab triggers good alignment
			// for both 4- and 8-space tab stops.
			s += "\n    \t"
		}
		s += flag.Usage
		if !isZeroValue(flag, flag.DefValue) {
			getter, ok := flag.Value.(flag_.Getter)
			if !ok {
				return
			}

			if _, ok := getter.Get().(string); ok {
				// put quotes on the value
				s += fmt.Sprintf(" (default %q)", flag.DefValue)
			} else {
				s += fmt.Sprintf(" (default %v)", flag.DefValue)
			}
		}
		fmt.Fprint(os.Stderr, s, "\n")
	})
}

func GenerateYaml(output io.Writer) {
	flag_.CommandLine.VisitAll(func(flag *flag_.Flag) {
		s := fmt.Sprintf("# %v\n%v: %v\n", flag.Usage, flag.Name, flag.DefValue)
		fmt.Fprint(output, s, "\n")
	})
}

func Parse(configPath string) error {
	/*load from yaml config and cli. if dumpconfig option pass than generate config and exit*/

	// first parse using bultin flag
	err := flag_.CommandLine.Parse(os.Args[1:])
	if err != nil {
		return err
	}

	if *config != "" {
		configPath = *config
	}
	var args []string
	// parse yaml and add params that wasn't passed from cli
	if configPath != "" {
		configPath, err := utils.AbsPath(configPath)
		if err != nil {
			return err
		}
		exists, err := utils.FileExists(configPath)
		if err != nil {
			return err
		}
		if exists {
			data, err := ioutil.ReadFile(configPath)
			if err != nil {
				return err
			}
			yamlConfig := map[string]interface{}{}
			err = yaml.Unmarshal([]byte(data), &yamlConfig)
			if err != nil {
				return err
			}
			setArgs := make(map[string]bool)
			flag_.Visit(func(flag *flag_.Flag) {
				setArgs[flag.Name] = true
			})
			// generate args list for flag.Parse as it was from cli args
			args = make([]string, 0)
			flag_.VisitAll(func(flag *flag_.Flag) {
				// generate only args that wasn't set from cli
				if _, alreadySet := setArgs[flag.Name]; !alreadySet {
					if value, yamlOk := yamlConfig[flag.Name]; yamlOk {
						if value != nil {
							args = append(args, fmt.Sprintf("--%v=%v", flag.Name, value))
						}
					}
				}
			})
		}
	}
	// set options from config that wasn't set by cli
	err = flag_.CommandLine.Parse(args)
	if err != nil {
		return err
	}
	if *dumpconfig {
		var absPath string
		if *config == "" {
			absPath, err = utils.AbsPath(configPath)
			if err != nil {
				return err
			}
		} else {
			absPath, err = utils.AbsPath(*config)
			if err != nil {
				return err
			}
		}

		dirPath := filepath.Dir(absPath)
		err = os.MkdirAll(dirPath, 0744)
		if err != nil {
			return err
		}

		file, err := os.Create(absPath)
		if err != nil {
			return err
		}
		defer file.Close()

		GenerateYaml(file)
		os.Exit(0)
	}
	return nil
}

var DEBUG_PREFIX = []byte("Debug: ")

type NotDebugWriter struct {
	writer io.Writer
}

func NewNotDebugWriter(writer io.Writer) *NotDebugWriter {
	return &NotDebugWriter{writer: writer}
}
func (wr *NotDebugWriter) Write(p []byte) (int, error) {
	if bytes.Contains(p, DEBUG_PREFIX) {
		return 0, nil
	}
	return wr.writer.Write(p)
}

func SetLogLevel(level int) {
	if level == LOG_DEBUG {
		log.SetOutput(os.Stdout)
	} else if level == LOG_VERBOSE {
		log.SetOutput(NewNotDebugWriter(os.Stdout))
	} else if level == LOG_DISCARD {
		log.SetOutput(ioutil.Discard)
	} else {
		panic(fmt.Sprintf("Incorrect log level - %v", level))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"runtime"
	"time"

	"github.com/BurntSushi/xgb"
	"github.com/BurntSushi/xgb/randr"
	"github.com/BurntSushi/xgb/xproto"
)

type CmdWatch struct{}

func init() {
	_, err := parser.AddCommand("watch",
		"watch for changes",
		"The watch command listens for changes and configures the outputs accordingly",
		&CmdWatch{})
	if err != nil {
		panic(err)
	}
}

type Event struct {
	Event xgb.Event
	Error error
}

const eventSendTimeout = 500 * time.Millisecond

func subscribeXEvents(ch chan<- Event, done <-chan struct{}) {
	X, err := xgb.NewConn()
	if err != nil {
		ch <- Event{Error: err}
		return
	}

	defer X.Close()
	if err = randr.Init(X); err != nil {
		ch <- Event{Error: err}
		return
	}

	root := xproto.Setup(X).DefaultScreen(X).Root

	eventMask := randr.NotifyMaskScreenChange |
		randr.NotifyMaskCrtcChange |
		randr.NotifyMaskOutputChange |
		randr.NotifyMaskOutputProperty

	err = randr.SelectInputChecked(X, root, uint16(eventMask)).Check()
	if err != nil {
		ch <- Event{Error: err}
		return
	}

	for {
		ev, err := X.WaitForEvent()
		select {
		case ch <- Event{Event: ev, Error: err}:
		case <-time.After(eventSendTimeout):
			continue
		case <-done:
			return
		}
		if err != nil {
			log.Fatal(err)
		}
	}
}

func (cmd CmdWatch) Execute(args []string) error {
	globalOpts.ReadConfigfile()

	done := make(chan struct{})
	defer close(done)

	ch := make(chan Event)
	go subscribeXEvents(ch, done)

	V("grobi %s, compiled with %v on %v\n", version, runtime.Version(), runtime.GOOS)
	V("successfully subscribed to X RANDR change events\n")

	var tickerCh <-chan time.Time
	if globalOpts.PollInterval > 0 {
		tickerCh = time.NewTicker(time.Duration(globalOpts.PollInterval) * time.Second).C
	}

	var backoffCh <-chan time.Time
	var disablePoll bool
	var eventReceived bool

	var lastRule Rule
	var lastOutputs Outputs
	for {
		if !disablePoll {
			V("--------------------- next iteration")
			var outputs Outputs
			var err error

			if eventReceived {
				outputs, err = DetectOutputs()
				eventReceived = false
			} else {
				outputs, err = GetOutputs()
			}

			if err != nil {
				return err
			}

			V("got outputs: %v", outputs)
			V("       last: %v", lastOutputs)

			// disable outputs which have a changed display
			var off Outputs
			for _, o := range outputs {
				for _, last := range lastOutputs {
					if o.Name != last.Name {
						continue
					}

					if last.Active() && !o.Active() {
						V("  output %v: monitor not active any more, disabling", o.Name)
						off = append(off, o)
						continue
					}

					if o.MonitorId != last.MonitorId {
						V("  output %v: monitor has changed, disabling", o.Name)
						off = append(off, o)
						continue
					}
				}
			}

			if len(off) > 0 {
				V("disable %d outputs", len(off))

				cmd, err := DisableOutputs(off)
				if err != nil {
					return err
				}

				// forget the last rule set, something has changed for sure
				lastRule = Rule{}

				err = RunCommand(cmd)
				if err != nil {
					fmt.Fprintf(os.Stderr, "error disabling outputs: %v", err)
				}

				// refresh outputs again
				outputs, err = GetOutputs()
				if err != nil {
					return err
				}

				V("new outputs after disable: %v", outputs)
			}

			rule, err := MatchRules(globalOpts.cfg.Rules, outputs)
			if err != nil {
				return err
			}

			if rule.Name != lastRule.Name {
				V("outputs: %v", outputs)
				V("new rule found: %v", rule.Name)

				err = ApplyRule(outputs, rule)
				if err != nil {
					return err
				}

				lastRule = rule

				if globalOpts.Pause > 0 {
					V("disable polling for %d seconds\n", globalOpts.Pause)
					disablePoll = true
					backoffCh = time.After(time.Duration(globalOpts.Pause) * time.Second)
				}
			}

			lastOutputs = outputs
		}

		select {
		case ev := <-ch:
			V("new RANDR change event received\n")
			if ev.Error != nil {
				return ev.Error
			}

			eventReceived = true
		case <-tickerCh:
		case <-backoffCh:
			V("reenable polling\n")
			backoffCh = nil
			disablePoll = false
		}
	}
}
<commit_msg>Remove extensive logging<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"runtime"
	"time"

	"github.com/BurntSushi/xgb"
	"github.com/BurntSushi/xgb/randr"
	"github.com/BurntSushi/xgb/xproto"
)

type CmdWatch struct{}

func init() {
	_, err := parser.AddCommand("watch",
		"watch for changes",
		"The watch command listens for changes and configures the outputs accordingly",
		&CmdWatch{})
	if err != nil {
		panic(err)
	}
}

type Event struct {
	Event xgb.Event
	Error error
}

const eventSendTimeout = 500 * time.Millisecond

func subscribeXEvents(ch chan<- Event, done <-chan struct{}) {
	X, err := xgb.NewConn()
	if err != nil {
		ch <- Event{Error: err}
		return
	}

	defer X.Close()
	if err = randr.Init(X); err != nil {
		ch <- Event{Error: err}
		return
	}

	root := xproto.Setup(X).DefaultScreen(X).Root

	eventMask := randr.NotifyMaskScreenChange |
		randr.NotifyMaskCrtcChange |
		randr.NotifyMaskOutputChange |
		randr.NotifyMaskOutputProperty

	err = randr.SelectInputChecked(X, root, uint16(eventMask)).Check()
	if err != nil {
		ch <- Event{Error: err}
		return
	}

	for {
		ev, err := X.WaitForEvent()
		select {
		case ch <- Event{Event: ev, Error: err}:
		case <-time.After(eventSendTimeout):
			continue
		case <-done:
			return
		}
		if err != nil {
			log.Fatal(err)
		}
	}
}

func (cmd CmdWatch) Execute(args []string) error {
	globalOpts.ReadConfigfile()

	done := make(chan struct{})
	defer close(done)

	ch := make(chan Event)
	go subscribeXEvents(ch, done)

	V("grobi %s, compiled with %v on %v\n", version, runtime.Version(), runtime.GOOS)
	V("successfully subscribed to X RANDR change events\n")

	var tickerCh <-chan time.Time
	if globalOpts.PollInterval > 0 {
		tickerCh = time.NewTicker(time.Duration(globalOpts.PollInterval) * time.Second).C
	}

	var backoffCh <-chan time.Time
	var disablePoll bool
	var eventReceived bool

	var lastRule Rule
	var lastOutputs Outputs
	for {
		if !disablePoll {
			var outputs Outputs
			var err error

			if eventReceived {
				outputs, err = DetectOutputs()
				eventReceived = false
			} else {
				outputs, err = GetOutputs()
			}

			if err != nil {
				return err
			}

			// disable outputs which have a changed display
			var off Outputs
			for _, o := range outputs {
				for _, last := range lastOutputs {
					if o.Name != last.Name {
						continue
					}

					if last.Active() && !o.Active() {
						V("  output %v: monitor not active any more, disabling", o.Name)
						off = append(off, o)
						continue
					}

					if o.MonitorId != last.MonitorId {
						V("  output %v: monitor has changed, disabling", o.Name)
						off = append(off, o)
						continue
					}
				}
			}

			if len(off) > 0 {
				V("disable %d outputs", len(off))

				cmd, err := DisableOutputs(off)
				if err != nil {
					return err
				}

				// forget the last rule set, something has changed for sure
				lastRule = Rule{}

				err = RunCommand(cmd)
				if err != nil {
					fmt.Fprintf(os.Stderr, "error disabling outputs: %v", err)
				}

				// refresh outputs again
				outputs, err = GetOutputs()
				if err != nil {
					return err
				}

				V("new outputs after disable: %v", outputs)
			}

			rule, err := MatchRules(globalOpts.cfg.Rules, outputs)
			if err != nil {
				return err
			}

			if rule.Name != lastRule.Name {
				V("outputs: %v", outputs)
				V("new rule found: %v", rule.Name)

				err = ApplyRule(outputs, rule)
				if err != nil {
					return err
				}

				lastRule = rule

				if globalOpts.Pause > 0 {
					V("disable polling for %d seconds\n", globalOpts.Pause)
					disablePoll = true
					backoffCh = time.After(time.Duration(globalOpts.Pause) * time.Second)
				}
			}

			lastOutputs = outputs
		}

		select {
		case ev := <-ch:
			V("new RANDR change event received\n")
			if ev.Error != nil {
				return ev.Error
			}

			eventReceived = true
		case <-tickerCh:
		case <-backoffCh:
			V("reenable polling\n")
			backoffCh = nil
			disablePoll = false
		}
	}
}
<|endoftext|>


<commit_before>package dbus

import (
	"bytes"
	"encoding/binary"
	"errors"
	"net"
	"os"
	"reflect"
	"strings"
	"sync"
)

const defaultSystemBusAddress = "unix:path=/var/run/dbus/system_bus_socket"

// Connection represents a connection to a message bus (usually, the system or
// session bus).
//
// Multiple goroutines may invoke methods on a connection simultaneously.
type Connection struct {
	transport       net.Conn
	uuid            string
	names           []string
	namesLck        sync.RWMutex
	serial          chan uint32
	replies         map[uint32]chan *Reply
	repliesLck      sync.RWMutex
	handlers        map[ObjectPath]*expObject
	handlersLck     sync.RWMutex
	out             chan *Message
	signals         chan Signal
	signalsLck      sync.Mutex
	eavesdropped    chan *Message
	eavesdroppedLck sync.Mutex
	busObj          *Object
}

// ConnectSessionBus connects to the session message bus and returns the
// connection or any error that occured.
func ConnectSessionBus() (*Connection, error) {
	address := os.Getenv("DBUS_SESSION_BUS_ADDRESS")
	if address != "" {
		return NewConnection(address)
	}
	return nil, errors.New("couldn't determine address of the session bus")
}

// ConnectSystemBus connects to the system message bus and returns the
// connection or any error that occured.
func ConnectSystemBus() (*Connection, error) {
	address := os.Getenv("DBUS_SYSTEM_BUS_ADDRESS")
	if address != "" {
		return NewConnection(address)
	}
	return NewConnection(defaultSystemBusAddress)
}

// NewConnection establishes a new connection to the message bus specified by
// address.
func NewConnection(address string) (*Connection, error) {
	var err error
	conn := new(Connection)
	if strings.HasPrefix(address, "unix") {
		abstract := getKey(address, "abstract")
		path := getKey(address, "path")
		switch {
		case abstract == "" && path == "":
			return nil, errors.New("bad address: neither path nor abstract set")
		case abstract != "" && path == "":
			conn.transport, err = net.Dial("unix", "@"+abstract)
			if err != nil {
				return nil, err
			}
		case abstract == "" && path != "":
			conn.transport, err = net.Dial("unix", path)
			if err != nil {
				return nil, err
			}
		case abstract != "" && path != "":
			return nil, errors.New("bad address: both path and abstract set")
		}
	} else {
		return nil, errors.New("bad address: invalid or unsupported transport")
	}
	if err = conn.auth(); err != nil {
		conn.transport.Close()
		return nil, err
	}
	conn.replies = make(map[uint32]chan *Reply)
	conn.out = make(chan *Message, 10)
	conn.handlers = make(map[ObjectPath]*expObject)
	conn.serial = make(chan uint32)
	conn.busObj = conn.Object("org.freedesktop.DBus", "/org/freedesktop/DBus")
	go conn.inWorker()
	go conn.outWorker()
	go conn.genSerials()
	if err = conn.hello(); err != nil {
		conn.transport.Close()
		return nil, err
	}
	return conn, nil
}

// BusObject returns the message bus object (i.e.
// conn.Object("org.freedesktop.DBus", "/org/freedesktop/DBus")).
func (conn *Connection) BusObject() *Object {
	return conn.busObj
}

// Close closes the underlying transport of the connection and stops all
// related goroutines.
func (conn *Connection) Close() error {
	close(conn.out)
	conn.signalsLck.Lock()
	if conn.signals != nil {
		close(conn.signals)
	}
	conn.signalsLck.Unlock()
	conn.eavesdroppedLck.Lock()
	if conn.eavesdropped != nil {
		close(conn.eavesdropped)
	}
	conn.eavesdroppedLck.Unlock()
	return conn.transport.Close()
}

// Eavesdrop changes the channel to which all messages are sent whose
// destination field is not one of the known names of this connection and which
// are not signals. The caller has to make sure that c is sufficiently buffered;
// if a message arrives when a write to c is not possible, the message is
// discarded.
//
// The channel can be reset by passing nil.
//
// If the connection is closed by the server or a call to Close, the channel is
// also closed.
func (conn *Connection) Eavesdrop(c chan *Message) {
	conn.eavesdroppedLck.Lock()
	conn.eavesdropped = c
	conn.eavesdroppedLck.Unlock()
}

func (conn *Connection) hello() error {
	var s string
	err := conn.busObj.Call("org.freedesktop.DBus.Hello", 0).Store(&s)
	if err != nil {
		return err
	}
	conn.namesLck.Lock()
	conn.names = make([]string, 1)
	conn.names[0] = s
	conn.namesLck.Unlock()
	return nil
}

func (conn *Connection) inWorker() {
	for {
		msg, err := conn.readMessage()
		if err == nil {
			dest, _ := msg.Headers[FieldDestination].value.(string)
			found := false
			conn.namesLck.RLock()
			if len(conn.names) == 0 {
				found = true
			}
			for _, v := range conn.names {
				if dest == v {
					found = true
					break
				}
			}
			conn.namesLck.RUnlock()
			conn.eavesdroppedLck.Lock()
			if !found && (msg.Type != TypeSignal || conn.eavesdropped != nil) {
				select {
				case conn.eavesdropped <- msg:
				default:
				}
				conn.eavesdroppedLck.Unlock()
				continue
			}
			conn.eavesdroppedLck.Unlock()
			switch msg.Type {
			case TypeMethodReply, TypeError:
				var reply *Reply
				serial := msg.Headers[FieldReplySerial].value.(uint32)
				if reply == nil {
					if msg.Type == TypeError {
						name, _ := msg.Headers[FieldErrorName].value.(string)
						reply = &Reply{nil, Error{name, msg.Body}}
					} else {
						reply = &Reply{msg.Body, nil}
					}
				}
				conn.repliesLck.Lock()
				if c, ok := conn.replies[serial]; ok {
					c <- reply
					delete(conn.replies, serial)
				}
				conn.repliesLck.Unlock()
			case TypeSignal:
				var signal Signal
				signal.Name = msg.Headers[FieldMember].value.(string)
				signal.Body = msg.Body
				// don't block trying to send a signal
				conn.signalsLck.Lock()
				select {
				case conn.signals <- signal:
				default:
				}
				conn.signalsLck.Unlock()
			case TypeMethodCall:
				go conn.handleCall(msg)
			}
		} else if _, ok := err.(InvalidMessageError); !ok {
			conn.Close()
			conn.repliesLck.RLock()
			for _, v := range conn.replies {
				v <- &Reply{nil, err}
			}
			conn.repliesLck.RUnlock()
			return
		}
		// invalid messages are ignored
	}
}

func (conn *Connection) outWorker() {
	for msg := range conn.out {
		err := msg.EncodeTo(conn.transport)
		conn.repliesLck.RLock()
		if err != nil && conn.replies[msg.Serial] != nil {
			conn.replies[msg.Serial] <- &Reply{nil, err}
		}
		conn.repliesLck.RUnlock()
	}
}

func (conn *Connection) readMessage() (*Message, error) {
	// read the first 16 bytes, from which we can figure out the length of the
	// rest of the message
	var header [16]byte
	if _, err := conn.transport.Read(header[:]); err != nil {
		return nil, err
	}
	var order binary.ByteOrder
	switch header[0] {
	case 'l':
		order = binary.LittleEndian
	case 'B':
		order = binary.BigEndian
	default:
		return nil, InvalidMessageError("invalid byte order")
	}
	// header[4:8] -> length of message body, header[12:16] -> length of header
	// fields (without alignment)
	var blen, hlen uint32
	binary.Read(bytes.NewBuffer(header[4:8]), order, &blen)
	binary.Read(bytes.NewBuffer(header[12:16]), order, &hlen)
	if hlen%8 != 0 {
		hlen += 8 - (hlen % 8)
	}
	rest := make([]byte, int(blen+hlen))
	if _, err := conn.transport.Read(rest); err != nil {
		return nil, err
	}
	all := make([]byte, 16+len(rest))
	copy(all, header[:])
	copy(all[16:], rest)
	return DecodeMessage(bytes.NewBuffer(all))
}

func (conn *Connection) sendError(e Error, dest string, serial uint32) {
	msg := new(Message)
	msg.Order = binary.LittleEndian
	msg.Type = TypeError
	msg.Serial = <-conn.serial
	msg.Headers = make(map[HeaderField]Variant)
	msg.Headers[FieldDestination] = MakeVariant(dest)
	msg.Headers[FieldErrorName] = MakeVariant(e.Name)
	msg.Headers[FieldReplySerial] = MakeVariant(serial)
	msg.Body = e.Body
	if len(e.Body) > 0 {
		msg.Headers[FieldSignature] = MakeVariant(GetSignature(e.Body...))
	}
	conn.out <- msg
}

func (conn *Connection) sendReply(dest string, serial uint32, values ...interface{}) {
	msg := new(Message)
	msg.Order = binary.LittleEndian
	msg.Type = TypeMethodReply
	msg.Serial = <-conn.serial
	msg.Headers = make(map[HeaderField]Variant)
	msg.Headers[FieldDestination] = MakeVariant(dest)
	msg.Headers[FieldReplySerial] = MakeVariant(serial)
	msg.Body = values
	if len(values) > 0 {
		msg.Headers[FieldSignature] = MakeVariant(GetSignature(values...))
	}
	conn.out <- msg
}

func (conn *Connection) genSerials() {
	s := uint32(1)
	for {
		conn.serial <- s
		// let's hope that nobody sends 2^32-1 messages at once
		s++
		if s == 0 {
			s++
		}
	}
}

// Object returns the object identified by the given destination name and path.
func (conn *Connection) Object(dest string, path ObjectPath) *Object {
	if !path.IsValid() {
		panic("invalid DBus path")
	}
	return &Object{conn, dest, path}
}

// Send the given message to the message bus. You usually don't need to use
// this; use the higher-level equivalents (Call, Emit and Export) instead.
// The returned cookie is nil if msg isn't a message call or if NoReplyExpected
// is set.
//
// The serial member is set to a unique serial before sending.
func (conn *Connection) Send(msg *Message) Cookie {
	if err := msg.IsValid(); err != nil {
		c := make(chan *Reply, 1)
		c <- &Reply{nil, err}
		return Cookie(c)
	}
	msg.Serial = <-conn.serial
	if msg.Type == TypeMethodCall && msg.Flags&FlagNoReplyExpected == 0 {
		conn.repliesLck.Lock()
		c := make(chan *Reply, 1)
		conn.replies[msg.Serial] = c
		conn.repliesLck.Unlock()
		conn.out <- msg
		return Cookie(c)
	}
	conn.out <- msg
	return nil
}

// Signal sets the channel to which all received signal messages are forwarded.
// The caller has to make sure that c is sufficiently buffered; if a message
// arrives when a write to c is not possible, it is discarded.
//
// The channel can be reset by passing nil.
//
// This channel is "overwritten" by Eavesdrop; i.e., if there currently is a
// channel for eavesdropped messages, this channel receives all signals, and the
// channel passed to Signal will not receive any signals.
//
// If the connection is closed by the server or a call to Close, the channel is
// also closed.
func (conn *Connection) Signal(c chan Signal) {
	conn.signalsLck.Lock()
	conn.signals = c
	conn.signalsLck.Unlock()
}

// Error represents a DBus message of type Error.
type Error struct {
	Name string
	Body []interface{}
}

func (e Error) Error() string {
	if len(e.Body) > 1 {
		s, ok := e.Body[0].(string)
		if ok {
			return s
		}
	}
	return e.Name
}

// Signal represents a DBus message of type Signal.
type Signal struct {
	Name string
	Body []interface{}
}

func getKey(s, key string) string {
	i := strings.IndexRune(s, ':')
	if i == -1 {
		return ""
	}
	s = s[i+1:]
	i = strings.Index(s, key)
	if i == -1 {
		return ""
	}
	if i+len(key)+1 >= len(s) || s[i+len(key)] != '=' {
		return ""
	}
	j := strings.Index(s, ",")
	if j == -1 {
		j = len(s)
	}
	return s[i+len(key)+1 : j]
}

func dereferenceAll(vs []interface{}) []interface{} {
	for i := range vs {
		v := reflect.ValueOf(vs[i])
		v = v.Elem()
		vs[i] = v.Interface()
	}
	return vs
}
<commit_msg>Never re-use currently used serials<commit_after>package dbus

import (
	"bytes"
	"encoding/binary"
	"errors"
	"net"
	"os"
	"reflect"
	"strings"
	"sync"
)

const defaultSystemBusAddress = "unix:path=/var/run/dbus/system_bus_socket"

// Connection represents a connection to a message bus (usually, the system or
// session bus).
//
// Multiple goroutines may invoke methods on a connection simultaneously.
type Connection struct {
	transport       net.Conn
	uuid            string
	names           []string
	namesLck        sync.RWMutex
	serial          chan uint32
	serialUsed      chan uint32
	replies         map[uint32]chan *Reply
	repliesLck      sync.RWMutex
	handlers        map[ObjectPath]*expObject
	handlersLck     sync.RWMutex
	out             chan *Message
	signals         chan Signal
	signalsLck      sync.Mutex
	eavesdropped    chan *Message
	eavesdroppedLck sync.Mutex
	busObj          *Object
}

// ConnectSessionBus connects to the session message bus and returns the
// connection or any error that occured.
func ConnectSessionBus() (*Connection, error) {
	address := os.Getenv("DBUS_SESSION_BUS_ADDRESS")
	if address != "" {
		return NewConnection(address)
	}
	return nil, errors.New("couldn't determine address of the session bus")
}

// ConnectSystemBus connects to the system message bus and returns the
// connection or any error that occured.
func ConnectSystemBus() (*Connection, error) {
	address := os.Getenv("DBUS_SYSTEM_BUS_ADDRESS")
	if address != "" {
		return NewConnection(address)
	}
	return NewConnection(defaultSystemBusAddress)
}

// NewConnection establishes a new connection to the message bus specified by
// address.
func NewConnection(address string) (*Connection, error) {
	var err error
	conn := new(Connection)
	if strings.HasPrefix(address, "unix") {
		abstract := getKey(address, "abstract")
		path := getKey(address, "path")
		switch {
		case abstract == "" && path == "":
			return nil, errors.New("bad address: neither path nor abstract set")
		case abstract != "" && path == "":
			conn.transport, err = net.Dial("unix", "@"+abstract)
			if err != nil {
				return nil, err
			}
		case abstract == "" && path != "":
			conn.transport, err = net.Dial("unix", path)
			if err != nil {
				return nil, err
			}
		case abstract != "" && path != "":
			return nil, errors.New("bad address: both path and abstract set")
		}
	} else {
		return nil, errors.New("bad address: invalid or unsupported transport")
	}
	if err = conn.auth(); err != nil {
		conn.transport.Close()
		return nil, err
	}
	conn.replies = make(map[uint32]chan *Reply)
	conn.out = make(chan *Message, 10)
	conn.handlers = make(map[ObjectPath]*expObject)
	conn.serial = make(chan uint32)
	conn.serialUsed = make(chan uint32)
	conn.busObj = conn.Object("org.freedesktop.DBus", "/org/freedesktop/DBus")
	go conn.inWorker()
	go conn.outWorker()
	go conn.serials()
	if err = conn.hello(); err != nil {
		conn.transport.Close()
		return nil, err
	}
	return conn, nil
}

// BusObject returns the message bus object (i.e.
// conn.Object("org.freedesktop.DBus", "/org/freedesktop/DBus")).
func (conn *Connection) BusObject() *Object {
	return conn.busObj
}

// Close closes the underlying transport of the connection and stops all
// related goroutines.
func (conn *Connection) Close() error {
	close(conn.out)
	conn.signalsLck.Lock()
	if conn.signals != nil {
		close(conn.signals)
	}
	conn.signalsLck.Unlock()
	conn.eavesdroppedLck.Lock()
	if conn.eavesdropped != nil {
		close(conn.eavesdropped)
	}
	conn.eavesdroppedLck.Unlock()
	return conn.transport.Close()
}

// Eavesdrop changes the channel to which all messages are sent whose
// destination field is not one of the known names of this connection and which
// are not signals. The caller has to make sure that c is sufficiently buffered;
// if a message arrives when a write to c is not possible, the message is
// discarded.
//
// The channel can be reset by passing nil.
//
// If the connection is closed by the server or a call to Close, the channel is
// also closed.
func (conn *Connection) Eavesdrop(c chan *Message) {
	conn.eavesdroppedLck.Lock()
	conn.eavesdropped = c
	conn.eavesdroppedLck.Unlock()
}

func (conn *Connection) hello() error {
	var s string
	err := conn.busObj.Call("org.freedesktop.DBus.Hello", 0).Store(&s)
	if err != nil {
		return err
	}
	conn.namesLck.Lock()
	conn.names = make([]string, 1)
	conn.names[0] = s
	conn.namesLck.Unlock()
	return nil
}

func (conn *Connection) inWorker() {
	for {
		msg, err := conn.readMessage()
		if err == nil {
			dest, _ := msg.Headers[FieldDestination].value.(string)
			found := false
			conn.namesLck.RLock()
			if len(conn.names) == 0 {
				found = true
			}
			for _, v := range conn.names {
				if dest == v {
					found = true
					break
				}
			}
			conn.namesLck.RUnlock()
			conn.eavesdroppedLck.Lock()
			if !found && (msg.Type != TypeSignal || conn.eavesdropped != nil) {
				select {
				case conn.eavesdropped <- msg:
				default:
				}
				conn.eavesdroppedLck.Unlock()
				continue
			}
			conn.eavesdroppedLck.Unlock()
			switch msg.Type {
			case TypeMethodReply, TypeError:
				var reply *Reply
				serial := msg.Headers[FieldReplySerial].value.(uint32)
				if reply == nil {
					if msg.Type == TypeError {
						name, _ := msg.Headers[FieldErrorName].value.(string)
						reply = &Reply{nil, Error{name, msg.Body}}
					} else {
						reply = &Reply{msg.Body, nil}
					}
				}
				conn.repliesLck.Lock()
				if c, ok := conn.replies[serial]; ok {
					c <- reply
					conn.serialUsed <- serial
					delete(conn.replies, serial)
				}
				conn.repliesLck.Unlock()
			case TypeSignal:
				var signal Signal
				signal.Name = msg.Headers[FieldMember].value.(string)
				signal.Body = msg.Body
				// don't block trying to send a signal
				conn.signalsLck.Lock()
				select {
				case conn.signals <- signal:
				default:
				}
				conn.signalsLck.Unlock()
			case TypeMethodCall:
				go conn.handleCall(msg)
			}
		} else if _, ok := err.(InvalidMessageError); !ok {
			conn.Close()
			conn.repliesLck.RLock()
			for _, v := range conn.replies {
				v <- &Reply{nil, err}
			}
			conn.repliesLck.RUnlock()
			return
		}
		// invalid messages are ignored
	}
}

func (conn *Connection) outWorker() {
	for msg := range conn.out {
		err := msg.EncodeTo(conn.transport)
		conn.repliesLck.RLock()
		if err != nil {
			if conn.replies[msg.Serial] != nil {
				conn.replies[msg.Serial] <- &Reply{nil, err}
			}
			conn.serialUsed <- msg.Serial
		} else if msg.Type != TypeMethodCall {
			conn.serialUsed <- msg.Serial
		}
		conn.repliesLck.RUnlock()
	}
}

func (conn *Connection) readMessage() (*Message, error) {
	// read the first 16 bytes, from which we can figure out the length of the
	// rest of the message
	var header [16]byte
	if _, err := conn.transport.Read(header[:]); err != nil {
		return nil, err
	}
	var order binary.ByteOrder
	switch header[0] {
	case 'l':
		order = binary.LittleEndian
	case 'B':
		order = binary.BigEndian
	default:
		return nil, InvalidMessageError("invalid byte order")
	}
	// header[4:8] -> length of message body, header[12:16] -> length of header
	// fields (without alignment)
	var blen, hlen uint32
	binary.Read(bytes.NewBuffer(header[4:8]), order, &blen)
	binary.Read(bytes.NewBuffer(header[12:16]), order, &hlen)
	if hlen%8 != 0 {
		hlen += 8 - (hlen % 8)
	}
	rest := make([]byte, int(blen+hlen))
	if _, err := conn.transport.Read(rest); err != nil {
		return nil, err
	}
	all := make([]byte, 16+len(rest))
	copy(all, header[:])
	copy(all[16:], rest)
	return DecodeMessage(bytes.NewBuffer(all))
}

func (conn *Connection) sendError(e Error, dest string, serial uint32) {
	msg := new(Message)
	msg.Order = binary.LittleEndian
	msg.Type = TypeError
	msg.Serial = <-conn.serial
	msg.Headers = make(map[HeaderField]Variant)
	msg.Headers[FieldDestination] = MakeVariant(dest)
	msg.Headers[FieldErrorName] = MakeVariant(e.Name)
	msg.Headers[FieldReplySerial] = MakeVariant(serial)
	msg.Body = e.Body
	if len(e.Body) > 0 {
		msg.Headers[FieldSignature] = MakeVariant(GetSignature(e.Body...))
	}
	conn.out <- msg
}

func (conn *Connection) sendReply(dest string, serial uint32, values ...interface{}) {
	msg := new(Message)
	msg.Order = binary.LittleEndian
	msg.Type = TypeMethodReply
	msg.Serial = <-conn.serial
	msg.Headers = make(map[HeaderField]Variant)
	msg.Headers[FieldDestination] = MakeVariant(dest)
	msg.Headers[FieldReplySerial] = MakeVariant(serial)
	msg.Body = values
	if len(values) > 0 {
		msg.Headers[FieldSignature] = MakeVariant(GetSignature(values...))
	}
	conn.out <- msg
}

func (conn *Connection) serials() {
	s := uint32(1)
	used := make(map[uint32]bool)
	used[0] = true // ensure that 0 is never used
	for {
		select {
		case conn.serial <- s:
			used[s] = true
			s++
			for used[s] {
				s++
			}
		case n := <-conn.serialUsed:
			delete(used, n)
		}
	}
}

// Object returns the object identified by the given destination name and path.
func (conn *Connection) Object(dest string, path ObjectPath) *Object {
	if !path.IsValid() {
		panic("invalid DBus path")
	}
	return &Object{conn, dest, path}
}

// Send the given message to the message bus. You usually don't need to use
// this; use the higher-level equivalents (Call, Emit and Export) instead.
// The returned cookie is nil if msg isn't a message call or if NoReplyExpected
// is set.
//
// The serial member is set to a unique serial before sending.
func (conn *Connection) Send(msg *Message) Cookie {
	if err := msg.IsValid(); err != nil {
		c := make(chan *Reply, 1)
		c <- &Reply{nil, err}
		return Cookie(c)
	}
	msg.Serial = <-conn.serial
	if msg.Type == TypeMethodCall && msg.Flags&FlagNoReplyExpected == 0 {
		conn.repliesLck.Lock()
		c := make(chan *Reply, 1)
		conn.replies[msg.Serial] = c
		conn.repliesLck.Unlock()
		conn.out <- msg
		return Cookie(c)
	}
	conn.out <- msg
	return nil
}

// Signal sets the channel to which all received signal messages are forwarded.
// The caller has to make sure that c is sufficiently buffered; if a message
// arrives when a write to c is not possible, it is discarded.
//
// The channel can be reset by passing nil.
//
// This channel is "overwritten" by Eavesdrop; i.e., if there currently is a
// channel for eavesdropped messages, this channel receives all signals, and the
// channel passed to Signal will not receive any signals.
//
// If the connection is closed by the server or a call to Close, the channel is
// also closed.
func (conn *Connection) Signal(c chan Signal) {
	conn.signalsLck.Lock()
	conn.signals = c
	conn.signalsLck.Unlock()
}

// Error represents a DBus message of type Error.
type Error struct {
	Name string
	Body []interface{}
}

func (e Error) Error() string {
	if len(e.Body) > 1 {
		s, ok := e.Body[0].(string)
		if ok {
			return s
		}
	}
	return e.Name
}

// Signal represents a DBus message of type Signal.
type Signal struct {
	Name string
	Body []interface{}
}

func getKey(s, key string) string {
	i := strings.IndexRune(s, ':')
	if i == -1 {
		return ""
	}
	s = s[i+1:]
	i = strings.Index(s, key)
	if i == -1 {
		return ""
	}
	if i+len(key)+1 >= len(s) || s[i+len(key)] != '=' {
		return ""
	}
	j := strings.Index(s, ",")
	if j == -1 {
		j = len(s)
	}
	return s[i+len(key)+1 : j]
}

func dereferenceAll(vs []interface{}) []interface{} {
	for i := range vs {
		v := reflect.ValueOf(vs[i])
		v = v.Elem()
		vs[i] = v.Interface()
	}
	return vs
}
<|endoftext|>


<commit_before>package docker

type DeviceSet interface {
	AddDevice(hash, baseHash string) error
	SetInitialized(hash string) error
	DeactivateDevice(hash string) error
	RemoveDevice(hash string) error
	MountDevice(hash, path string) error
	HasDevice(hash string) bool
	HasInitializedDevice(hash string) bool
}
<commit_msg>Add DeviceSetWrapper<commit_after>package docker

type DeviceSet interface {
	AddDevice(hash, baseHash string) error
	SetInitialized(hash string) error
	DeactivateDevice(hash string) error
	RemoveDevice(hash string) error
	MountDevice(hash, path string) error
	HasDevice(hash string) bool
	HasInitializedDevice(hash string) bool
}

type DeviceSetWrapper struct {
	wrapped DeviceSet
	prefix string
}

func (wrapper *DeviceSetWrapper) wrap(hash string) string {
	if hash != "" {
		hash = wrapper.prefix + "-" + hash
	}
	return hash
}

func (wrapper *DeviceSetWrapper) AddDevice(hash, baseHash string) error {
	return wrapper.wrapped.AddDevice(wrapper.wrap(hash), wrapper.wrap(baseHash))
}

func (wrapper *DeviceSetWrapper) SetInitialized(hash string) error {
	return wrapper.wrapped.SetInitialized(wrapper.wrap(hash))
}

func (wrapper *DeviceSetWrapper) DeactivateDevice(hash string) error {
	return wrapper.wrapped.DeactivateDevice(wrapper.wrap(hash))
}

func (wrapper *DeviceSetWrapper) RemoveDevice(hash string) error {
	return wrapper.wrapped.RemoveDevice(wrapper.wrap(hash))
}

func (wrapper *DeviceSetWrapper) MountDevice(hash, path string) error {
	return wrapper.wrapped.MountDevice(wrapper.wrap(hash), path)
}

func (wrapper *DeviceSetWrapper) HasDevice(hash string) bool {
	return wrapper.wrapped.HasDevice(wrapper.wrap(hash))
}

func (wrapper *DeviceSetWrapper) HasInitializedDevice(hash string) bool {
	return wrapper.wrapped.HasInitializedDevice(wrapper.wrap(hash))
}

func NewDeviceSetWrapper(wrapped DeviceSet, prefix string) DeviceSet {
	wrapper := &DeviceSetWrapper{
		wrapped: wrapped,
		prefix: prefix,
	}
	return wrapper
}
<|endoftext|>