weave-scope/prog/main.go

package main

import (
	"compress/gzip"
	"flag"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/weaveworks/scope/app"
	"github.com/weaveworks/scope/common/xfer"
	"github.com/weaveworks/scope/probe/kubernetes"
	"github.com/weaveworks/weave/common"
)

var (
	// set at build time
	version = "dev"
	// tokens to be elided when logging
	serviceTokenFlag       = "service-token"
	probeTokenFlag         = "probe.token"
	kubernetesPasswordFlag = "probe.kubernetes.password"
	kubernetesTokenFlag    = "probe.kubernetes.token"
	sensitiveFlags         = []string{
		serviceTokenFlag,
		probeTokenFlag,
		kubernetesPasswordFlag,
		kubernetesTokenFlag,
	}
)

type prefixFormatter struct {
	prefix []byte
	next   log.Formatter
}

func (f *prefixFormatter) Format(entry *log.Entry) ([]byte, error) {
	formatted, err := f.next.Format(entry)
	if err != nil {
		return formatted, err
	}
	return append(f.prefix, formatted...), nil
}

func setLogFormatter(prefix string) {
	if !strings.HasSuffix(prefix, " ") {
		prefix += " "
	}
	f := prefixFormatter{
		prefix: []byte(prefix),
		// reuse weave's log format
		next: common.Log.Formatter,
	}
	log.SetFormatter(&f)
}

func setLogLevel(levelname string) {
	level, err := log.ParseLevel(levelname)
	if err != nil {
		log.Fatal(err)
	}
	log.SetLevel(level)
}

type flags struct {
	probe probeFlags
	app   appFlags
}

type probeFlags struct {
	token           string
	httpListen      string
	publishInterval time.Duration
	spyInterval     time.Duration
	pluginsRoot     string
	insecure        bool
	logPrefix       string
	logLevel        string
	resolver        string
	noApp           bool
	noControls      bool

	useConntrack bool // Use conntrack for endpoint topo
	spyProcs     bool // Associate endpoints with processes (must be root)
	procEnabled  bool // Produce process topology & process nodes in endpoint
	procRoot     string

	dockerEnabled  bool
	dockerInterval time.Duration
	dockerBridge   string

	kubernetesEnabled bool
	kubernetesConfig  kubernetes.ClientConfig

	weaveEnabled  bool
	weaveAddr     string
	weaveHostname string
}

type appFlags struct {
	window    time.Duration
	listen    string
	logLevel  string
	logPrefix string
	logHTTP   bool

	weaveEnabled   bool
	weaveAddr      string
	weaveHostname  string
	containerName  string
	dockerEndpoint string

	collectorURL              string
	s3URL                     string
	controlRouterURL          string
	pipeRouterURL             string
	natsHostname              string
	memcachedHostname         string
	memcachedTimeout          time.Duration
	memcachedService          string
	memcachedExpiration       time.Duration
	memcachedCompressionLevel int
	userIDHeader              string

	blockProfileRate int

	awsCreateTables bool
	consulInf       string
}

func logCensoredArgs() {
	var prettyPrintedArgs string
	// We show the flags followed by the args. This may change the original
	// ordering. However the flag parser doesn't keep positioning
	// information, not allowing for a more accurate reconstruction.
	flag.Visit(func(f *flag.Flag) {
		value := f.Value.String()
		// omit sensitive information
		for _, sensitiveFlag := range sensitiveFlags {
			if f.Name == sensitiveFlag {
				value = "<elided>"
				break
			}
		}
		prettyPrintedArgs += fmt.Sprintf(" --%s=%s", f.Name, value)
	})
	for _, arg := range flag.Args() {
		prettyPrintedArgs += " " + arg
	}
	log.Infof("command line args:%s", prettyPrintedArgs)
}

func main() {
	var (
		flags         = flags{}
		mode          string
		debug         bool
		weaveEnabled  bool
		weaveHostname string
		dryRun        bool
	)

	// Flags that apply to both probe and app
	flag.StringVar(&mode, "mode", "help", "For internal use.")
	flag.BoolVar(&debug, "debug", false, "Force debug logging.")
	flag.BoolVar(&dryRun, "dry-run", false, "Don't start scope, just parse the arguments.  For internal use only.")
	flag.BoolVar(&weaveEnabled, "weave", true, "Enable Weave Net integrations.")
	flag.StringVar(&weaveHostname, "weave.hostname", "", "Hostname to advertise/lookup in WeaveDNS")

	// We need to know how to parse them, but they are mainly interpreted by the entrypoint script.
	// They are also here so they are included in usage, and the probe uses them to decide if to
	// publish to localhost.
	noApp := flag.Bool("no-app", false, "Don't run the app.")
	probeOnly := flag.Bool("app-only", false, "Only run the app")
	flag.Bool("probe-only", false, "Only run the probe.")
	flag.Bool("no-probe", false, "Don't run the probe.")

	// Probe flags
	flag.StringVar(&flags.probe.token, serviceTokenFlag, "", "Token to use to authenticate with cloud.weave.works")
	flag.StringVar(&flags.probe.token, probeTokenFlag, "", "Token to use to authenticate with cloud.weave.works")
	flag.StringVar(&flags.probe.httpListen, "probe.http.listen", "", "listen address for HTTP profiling and instrumentation server")
	flag.DurationVar(&flags.probe.publishInterval, "probe.publish.interval", 3*time.Second, "publish (output) interval")
	flag.DurationVar(&flags.probe.spyInterval, "probe.spy.interval", time.Second, "spy (scan) interval")
	flag.StringVar(&flags.probe.pluginsRoot, "probe.plugins.root", "/var/run/scope/plugins", "Root directory to search for plugins")
	flag.BoolVar(&flags.probe.noControls, "probe.no-controls", false, "Disable controls (e.g. start/stop containers, terminals, logs ...)")

	flag.BoolVar(&flags.probe.insecure, "probe.insecure", false, "(SSL) explicitly allow \"insecure\" SSL connections and transfers")
	flag.StringVar(&flags.probe.resolver, "probe.resolver", "", "IP address & port of resolver to use.  Default is to use system resolver.")
	flag.StringVar(&flags.probe.logPrefix, "probe.log.prefix", "<probe>", "prefix for each log line")
	flag.StringVar(&flags.probe.logLevel, "probe.log.level", "info", "logging threshold level: debug|info|warn|error|fatal|panic")

	// Proc & endpoint
	flag.BoolVar(&flags.probe.useConntrack, "probe.conntrack", true, "also use conntrack to track connections")
	flag.BoolVar(&flags.probe.spyProcs, "probe.proc.spy", true, "associate endpoints with processes (needs root)")
	flag.StringVar(&flags.probe.procRoot, "probe.proc.root", "/proc", "location of the proc filesystem")
	flag.BoolVar(&flags.probe.procEnabled, "probe.processes", true, "produce process topology & include procspied connections")

	// Docker
	flag.BoolVar(&flags.probe.dockerEnabled, "probe.docker", false, "collect Docker-related attributes for processes")
	flag.DurationVar(&flags.probe.dockerInterval, "probe.docker.interval", 10*time.Second, "how often to update Docker attributes")
	flag.StringVar(&flags.probe.dockerBridge, "probe.docker.bridge", "docker0", "the docker bridge name")

	// K8s
	flag.BoolVar(&flags.probe.kubernetesEnabled, "probe.kubernetes", false, "collect kubernetes-related attributes for containers, should only be enabled on the master node")
	flag.DurationVar(&flags.probe.kubernetesConfig.Interval, "probe.kubernetes.interval", 10*time.Second, "how often to do a full resync of the kubernetes data")
	flag.StringVar(&flags.probe.kubernetesConfig.Server, "probe.kubernetes.api", "", "The address and port of the Kubernetes API server (deprecated in favor of equivalent probe.kubernetes.server)")
	flag.StringVar(&flags.probe.kubernetesConfig.CertificateAuthority, "probe.kubernetes.certificate-authority", "", "Path to a cert. file for the certificate authority")
	flag.StringVar(&flags.probe.kubernetesConfig.ClientCertificate, "probe.kubernetes.client-certificate", "", "Path to a client certificate file for TLS")
	flag.StringVar(&flags.probe.kubernetesConfig.ClientKey, "probe.kubernetes.client-key", "", "Path to a client key file for TLS")
	flag.StringVar(&flags.probe.kubernetesConfig.Cluster, "probe.kubernetes.cluster", "", "The name of the kubeconfig cluster to use")
	flag.StringVar(&flags.probe.kubernetesConfig.Context, "probe.kubernetes.context", "", "The name of the kubeconfig context to use")
	flag.BoolVar(&flags.probe.kubernetesConfig.Insecure, "probe.kubernetes.insecure-skip-tls-verify", false, "If true, the server's certificate will not be checked for validity. This will make your HTTPS connections insecure")
	flag.StringVar(&flags.probe.kubernetesConfig.Kubeconfig, "probe.kubernetes.kubeconfig", "", "Path to the kubeconfig file to use")
	flag.StringVar(&flags.probe.kubernetesConfig.Password, kubernetesPasswordFlag, "", "Password for basic authentication to the API server")
	flag.StringVar(&flags.probe.kubernetesConfig.Server, "probe.kubernetes.server", "", "The address and port of the Kubernetes API server")
	flag.StringVar(&flags.probe.kubernetesConfig.Token, kubernetesTokenFlag, "", "Bearer token for authentication to the API server")
	flag.StringVar(&flags.probe.kubernetesConfig.User, "probe.kubernetes.user", "", "The name of the kubeconfig user to use")
	flag.StringVar(&flags.probe.kubernetesConfig.Username, "probe.kubernetes.username", "", "Username for basic authentication to the API server")

	// Weave
	flag.StringVar(&flags.probe.weaveAddr, "probe.weave.addr", "127.0.0.1:6784", "IP address & port of the Weave router")
	flag.StringVar(&flags.probe.weaveHostname, "probe.weave.hostname", app.DefaultHostname, "Hostname to lookup in WeaveDNS")

	// App flags
	flag.DurationVar(&flags.app.window, "app.window", 15*time.Second, "window")
	flag.StringVar(&flags.app.listen, "app.http.address", ":"+strconv.Itoa(xfer.AppPort), "webserver listen address")
	flag.StringVar(&flags.app.logLevel, "app.log.level", "info", "logging threshold level: debug|info|warn|error|fatal|panic")
	flag.StringVar(&flags.app.logPrefix, "app.log.prefix", "<app>", "prefix for each log line")
	flag.BoolVar(&flags.app.logHTTP, "app.log.http", false, "Log individual HTTP requests")

	flag.StringVar(&flags.app.weaveAddr, "app.weave.addr", app.DefaultWeaveURL, "Address on which to contact WeaveDNS")
	flag.StringVar(&flags.app.weaveHostname, "app.weave.hostname", app.DefaultHostname, "Hostname to advertise in WeaveDNS")
	flag.StringVar(&flags.app.containerName, "app.container.name", app.DefaultContainerName, "Name of this container (to lookup container ID)")
	flag.StringVar(&flags.app.dockerEndpoint, "app.docker", app.DefaultDockerEndpoint, "Location of docker endpoint (to lookup container ID)")

	flag.StringVar(&flags.app.collectorURL, "app.collector", "local", "Collector to use (local, dynamodb, or file)")
	flag.StringVar(&flags.app.s3URL, "app.collector.s3", "local", "S3 URL to use (when collector is dynamodb)")
	flag.StringVar(&flags.app.controlRouterURL, "app.control.router", "local", "Control router to use (local or sqs)")
	flag.StringVar(&flags.app.pipeRouterURL, "app.pipe.router", "local", "Pipe router to use (local)")
	flag.StringVar(&flags.app.natsHostname, "app.nats", "", "Hostname for NATS service to use for shortcut reports.  If empty, shortcut reporting will be disabled.")
	flag.StringVar(&flags.app.memcachedHostname, "app.memcached.hostname", "", "Hostname for memcached service to use when caching reports.  If empty, no memcached will be used.")
	flag.DurationVar(&flags.app.memcachedTimeout, "app.memcached.timeout", 100*time.Millisecond, "Maximum time to wait before giving up on memcached requests.")
	flag.DurationVar(&flags.app.memcachedExpiration, "app.memcached.expiration", 2*15*time.Second, "How long reports stay in the memcache.")
	flag.StringVar(&flags.app.memcachedService, "app.memcached.service", "memcached", "SRV service used to discover memcache servers.")
	flag.IntVar(&flags.app.memcachedCompressionLevel, "app.memcached.compression", gzip.DefaultCompression, "How much to compress reports stored in memcached.")
	flag.StringVar(&flags.app.userIDHeader, "app.userid.header", "", "HTTP header to use as userid")

	flag.IntVar(&flags.app.blockProfileRate, "app.block.profile.rate", 0, "If more than 0, enable block profiling. The profiler aims to sample an average of one blocking event per rate nanoseconds spent blocked.")

	flag.BoolVar(&flags.app.awsCreateTables, "app.aws.create.tables", false, "Create the tables in DynamoDB")
	flag.StringVar(&flags.app.consulInf, "app.consul.inf", "", "The interface who's address I should advertise myself under in consul")

	flag.Parse()

	// Deal with common args
	if debug {
		flags.probe.logLevel = "debug"
		flags.app.logLevel = "debug"
	}
	if weaveHostname != "" {
		flags.probe.weaveHostname = weaveHostname
		flags.app.weaveHostname = weaveHostname
	}
	flags.probe.weaveEnabled = weaveEnabled
	flags.app.weaveEnabled = weaveEnabled
	flags.probe.noApp = *noApp || *probeOnly

	// Special case for #1191, check listen address is well formed
	_, _, err := net.SplitHostPort(flags.app.listen)
	if err != nil {
		log.Errorf("Invalid value for -app.http.address: %v", err)
	}
	if flags.probe.httpListen != "" {
		_, _, err := net.SplitHostPort(flags.probe.httpListen)
		if err != nil {
			log.Errorf("Invalid value for -app.http.address: %v", err)
		}
	}

	if dryRun {
		return
	}

	switch mode {
	case "app":
		appMain(flags.app)
	case "probe":
		probeMain(flags.probe)
	case "version":
		fmt.Println("Weave Scope version", version)
	case "help":
		flag.PrintDefaults()
	default:
		fmt.Printf("command '%s' not recognices", mode)
		os.Exit(1)
	}
}