<commit_before>// Copyright 2015-2016 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"runtime/pprof"
	"runtime/trace"
	"sort"
	"strconv"
	"strings"
	"time"
	"errors"
	_ "github.com/trivago/gollum/consumer"
	"github.com/trivago/gollum/core"
	_ "github.com/trivago/gollum/filter"
	_ "github.com/trivago/gollum/format"
	_ "github.com/trivago/gollum/producer"
	_ "github.com/trivago/gollum/router"
	"github.com/trivago/tgo"
	"github.com/trivago/tgo/tlog"
	"github.com/trivago/tgo/tstrings"
	"github.com/trivago/tgo/thealthcheck"
	"net"
)

func main() {
	tlog.SetCacheWriter()
	parseFlags()
	tlog.SetVerbosity(tlog.Verbosity(*flagLoglevel))

	if *flagVersion {
		printVersion()
		return // ### return, version only ###
	}

	if *flagModules {
		printModules()
		return // ### return, modules only ###
	}

	if *flagHelp || *flagConfigFile == "" {
		printFlags()
		return // ### return, nothing to do ###
	}

	// Read and test config

	configFile := *flagTestConfigFile
	if configFile == "" {
		configFile = *flagConfigFile
	}

	config, err := core.ReadConfig(*flagConfigFile)
	if err != nil {
		fmt.Printf("Config: %s\n", err.Error())
		return // ### return, config error ###
	}

	errors := config.Validate()
	for _, err := range errors {
		fmt.Print(err.Error())
	}

	if *flagTestConfigFile != "" {
		if len(errors) == 0 {
			fmt.Print("Config check passed.")
		} else {
			fmt.Print("Config check FAILED.")
		}
		return // ### return, only test config ###
	}

	// Configure runtime

	if *flagPidFile != "" {
		ioutil.WriteFile(*flagPidFile, []byte(strconv.Itoa(os.Getpid())), 0644)
	}

	if *flagNumCPU == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		runtime.GOMAXPROCS(*flagNumCPU)
	}

	// Metrics server start

	setStaticMetrics()
	if *flagMetricsAddress != "" {
		server := tgo.NewMetricServer()
		address, err := parseAddress(*flagMetricsAddress)
		if err != nil {
			fmt.Printf("%s", err)
			return
		}
		go server.Start(address)
		defer server.Stop()
	}

	// Health Check endpoint

	if *flagHealthCheck != "" {
		address, err := parseAddress(*flagHealthCheck)
		if err != nil {
			fmt.Printf("%s", err)
			return
		}
		thealthcheck.Configure(address)

		go thealthcheck.Start()
		defer thealthcheck.Stop()

		// Add a static "ping" endpoint
		thealthcheck.AddEndpoint("/_PING_", func()(code int, body string){
			return thealthcheck.StatusOK, "PONG"
		})
	}

	// Profiling flags

	if *flagCPUProfile != "" {
		if file, err := os.Create(*flagCPUProfile); err != nil {
			panic(err)
		} else {
			defer file.Close()
			if err := pprof.StartCPUProfile(file); err != nil {
				panic(err)
			}
			defer pprof.StopCPUProfile()
		}
	}

	if *flagProfile {
		time.AfterFunc(time.Second*3, printProfile)
	}

	if *flagTrace != "" {
		traceFile, err := os.OpenFile(*flagTrace, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			panic(err)
		}
		defer traceFile.Close()

		if err := trace.Start(traceFile); err != nil {
			panic(err)
		}
		defer trace.Stop()
	}

	if *flagMemProfile != "" {
		defer dumpMemoryProfile()
	}

	// Start the coordinator

	coordinater := NewCoordinator()
	coordinater.Configure(config)

	defer coordinater.Shutdown()
	coordinater.StartPlugins()
	coordinater.Run()
}

func parseAddress(address string) (string, error) {
	// net.SplitHostPort() doesn't support plain port number
	if tstrings.IsInt(address) {
		address = ":" + address
	}

	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return address, errors.New(fmt.Sprintf("Incorrect address %q: %s", address, err))
	}
	fmt.Printf("DEBUG: addr: %q, host: %q, port: %q", address, host, port)
	return host + ":" + port, nil
}

func dumpMemoryProfile() {
	if file, err := os.Create(*flagMemProfile); err != nil {
		panic(err)
	} else {
		defer file.Close()
		pprof.WriteHeapProfile(file)
	}
}

func printVersion() {
	fmt.Printf("Gollum: %s\n", GetVersionString())
	fmt.Printf("Version: %d\n", GetVersionNumber())
	fmt.Println(runtime.Version())
}

func printModules() {
	namespaces := []string{"consumer", "producer", "filter", "format", "router", "contrib"}
	allMods := []string{}
	for _, pkg := range namespaces {
		modules := core.TypeRegistry.GetRegistered(pkg)
		for _, typeName := range modules {
			allMods = append(allMods, typeName)
		}
	}

	sort.Strings(allMods)
	lastCategory := ""

	for _, name := range allMods {
		pkgIdx := strings.LastIndex(name, ".")
		category := name[:pkgIdx]

		if category != lastCategory {
			fmt.Printf("\n-- %s\n", category)
		}

		fmt.Println(name)
		lastCategory = category
	}
}

func printProfile() {
	msgSec, err := tgo.Metric.Get(core.MetricMessagesSec)
	if err == nil {
		fmt.Printf("Processed %d msg/sec\n", msgSec)
	}
	time.AfterFunc(time.Second*3, printProfile)
}

func setStaticMetrics() {
	metricVersion := "Version"
	tgo.Metric.New(metricVersion)
	tgo.Metric.InitSystemMetrics()
	tgo.Metric.Set(metricVersion, GetVersionNumber())
}
<commit_msg>Remove debug statement<commit_after>// Copyright 2015-2016 trivago GmbH
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"runtime"
	"runtime/pprof"
	"runtime/trace"
	"sort"
	"strconv"
	"strings"
	"time"
	"errors"
	_ "github.com/trivago/gollum/consumer"
	"github.com/trivago/gollum/core"
	_ "github.com/trivago/gollum/filter"
	_ "github.com/trivago/gollum/format"
	_ "github.com/trivago/gollum/producer"
	_ "github.com/trivago/gollum/router"
	"github.com/trivago/tgo"
	"github.com/trivago/tgo/tlog"
	"github.com/trivago/tgo/tstrings"
	"github.com/trivago/tgo/thealthcheck"
	"net"
)

func main() {
	tlog.SetCacheWriter()
	parseFlags()
	tlog.SetVerbosity(tlog.Verbosity(*flagLoglevel))

	if *flagVersion {
		printVersion()
		return // ### return, version only ###
	}

	if *flagModules {
		printModules()
		return // ### return, modules only ###
	}

	if *flagHelp || *flagConfigFile == "" {
		printFlags()
		return // ### return, nothing to do ###
	}

	// Read and test config

	configFile := *flagTestConfigFile
	if configFile == "" {
		configFile = *flagConfigFile
	}

	config, err := core.ReadConfig(*flagConfigFile)
	if err != nil {
		fmt.Printf("Config: %s\n", err.Error())
		return // ### return, config error ###
	}

	errors := config.Validate()
	for _, err := range errors {
		fmt.Print(err.Error())
	}

	if *flagTestConfigFile != "" {
		if len(errors) == 0 {
			fmt.Print("Config check passed.")
		} else {
			fmt.Print("Config check FAILED.")
		}
		return // ### return, only test config ###
	}

	// Configure runtime

	if *flagPidFile != "" {
		ioutil.WriteFile(*flagPidFile, []byte(strconv.Itoa(os.Getpid())), 0644)
	}

	if *flagNumCPU == 0 {
		runtime.GOMAXPROCS(runtime.NumCPU())
	} else {
		runtime.GOMAXPROCS(*flagNumCPU)
	}

	// Metrics server start

	setStaticMetrics()
	if *flagMetricsAddress != "" {
		server := tgo.NewMetricServer()
		address, err := parseAddress(*flagMetricsAddress)
		if err != nil {
			fmt.Printf("%s", err)
			return
		}
		go server.Start(address)
		defer server.Stop()
	}

	// Health Check endpoint

	if *flagHealthCheck != "" {
		address, err := parseAddress(*flagHealthCheck)
		if err != nil {
			fmt.Printf("%s", err)
			return
		}
		thealthcheck.Configure(address)

		go thealthcheck.Start()
		defer thealthcheck.Stop()

		// Add a static "ping" endpoint
		thealthcheck.AddEndpoint("/_PING_", func()(code int, body string){
			return thealthcheck.StatusOK, "PONG"
		})
	}

	// Profiling flags

	if *flagCPUProfile != "" {
		if file, err := os.Create(*flagCPUProfile); err != nil {
			panic(err)
		} else {
			defer file.Close()
			if err := pprof.StartCPUProfile(file); err != nil {
				panic(err)
			}
			defer pprof.StopCPUProfile()
		}
	}

	if *flagProfile {
		time.AfterFunc(time.Second*3, printProfile)
	}

	if *flagTrace != "" {
		traceFile, err := os.OpenFile(*flagTrace, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
		if err != nil {
			panic(err)
		}
		defer traceFile.Close()

		if err := trace.Start(traceFile); err != nil {
			panic(err)
		}
		defer trace.Stop()
	}

	if *flagMemProfile != "" {
		defer dumpMemoryProfile()
	}

	// Start the coordinator

	coordinater := NewCoordinator()
	coordinater.Configure(config)

	defer coordinater.Shutdown()
	coordinater.StartPlugins()
	coordinater.Run()
}

func parseAddress(address string) (string, error) {
	// net.SplitHostPort() doesn't support plain port number
	if tstrings.IsInt(address) {
		address = ":" + address
	}

	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return address, errors.New(fmt.Sprintf("Incorrect address %q: %s", address, err))
	}

	return host + ":" + port, nil
}

func dumpMemoryProfile() {
	if file, err := os.Create(*flagMemProfile); err != nil {
		panic(err)
	} else {
		defer file.Close()
		pprof.WriteHeapProfile(file)
	}
}

func printVersion() {
	fmt.Printf("Gollum: %s\n", GetVersionString())
	fmt.Printf("Version: %d\n", GetVersionNumber())
	fmt.Println(runtime.Version())
}

func printModules() {
	namespaces := []string{"consumer", "producer", "filter", "format", "router", "contrib"}
	allMods := []string{}
	for _, pkg := range namespaces {
		modules := core.TypeRegistry.GetRegistered(pkg)
		for _, typeName := range modules {
			allMods = append(allMods, typeName)
		}
	}

	sort.Strings(allMods)
	lastCategory := ""

	for _, name := range allMods {
		pkgIdx := strings.LastIndex(name, ".")
		category := name[:pkgIdx]

		if category != lastCategory {
			fmt.Printf("\n-- %s\n", category)
		}

		fmt.Println(name)
		lastCategory = category
	}
}

func printProfile() {
	msgSec, err := tgo.Metric.Get(core.MetricMessagesSec)
	if err == nil {
		fmt.Printf("Processed %d msg/sec\n", msgSec)
	}
	time.AfterFunc(time.Second*3, printProfile)
}

func setStaticMetrics() {
	metricVersion := "Version"
	tgo.Metric.New(metricVersion)
	tgo.Metric.InitSystemMetrics()
	tgo.Metric.Set(metricVersion, GetVersionNumber())
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"log"
	"os"
	"path"

	yaml "gopkg.in/yaml.v2"
)

func main() {

	init := flag.Bool("i", false, "Initialize the swagger-meta.json file with default values")
	sourceDir := flag.String("s", ".", "The root of the source code you want swagger-gen to scan and build a swagger spec from. Defaults to current directory")
	outDir := flag.String("o", ".", "The path to the directory where the generated swagger file will be output to. Defaults to current directory")
	format := flag.String("f", "json", "Output format. json | yaml. Defaults to json")

	flag.Parse()

	if _, dirErr := os.Stat(path.Dir(*sourceDir)); os.IsNotExist(dirErr) {
		log.Fatal(dirErr)
	}

	swaggerMetaPath := path.Join(*sourceDir, "swagger-meta.json")

	// Init command (-i)
	if *init == true {
		swagger := Swagger{}
		swagger.Swagger = "2.0"
		swagger.Info = SwaggerInfo{}
		swagger.Info.Description = "My API Description"
		swagger.Info.TermsOfService = "TOS"
		swagger.Info.Title = "My api title"
		swagger.Info.Version = "0.1.0"
		swagger.Info.Contact = Contact{
			"example@email.com",
		}

		swagger.Info.License = License{
			"Apache 2.0",
			"http://www.apache.org/licenses/LICENSE-2.0.html",
		}

		swagger.Host = "myhost.com"
		swagger.BasePath = "/v1"
		swagger.Tags = []Tag{}
		swagger.Schemes = []string{
			"http",
		}
		toJSON(swagger, swaggerMetaPath)
		log.Printf("Swagger meta file generated at path %s", swaggerMetaPath)
		os.Exit(0)
	}

	jsonBytes, jsonBytesErr := ReadJSONToBytes(swaggerMetaPath)

	if jsonBytesErr != nil {
		log.Fatal(jsonBytesErr)
	}

	swaggerf := Swaggerf{}

	swaggerf.ParseSwaggerConfig(jsonBytes)

	// Build the swagger object
	swaggerf.BuildSwagger(*sourceDir)
	var outFile string
	switch *format {
	case "json":
		outFile = "swagger.json"
		outPath := path.Join(*outDir, outFile)
		toJSON(swaggerf.Swagger, outPath)
	case "yaml":
		outFile = "swagger.yaml"
		outPath := path.Join(*outDir, outFile)
		toYAML(swaggerf.Swagger, outPath)
	default:
		log.Fatal("Invalid output format. Should be `json` or `yaml`")
	}
}

func toYAML(swagger Swagger, outFile string) {
	data, _ := yaml.Marshal(&swagger)
	file, _ := os.OpenFile(
		outFile,
		os.O_WRONLY|os.O_TRUNC|os.O_CREATE,
		0666,
	)
	defer file.Close()
	_, outFileErr := file.Write(data)
	if outFileErr != nil {
		log.Fatal(outFileErr)
	}
}

func toJSON(swagger Swagger, outFile string) {

	// log.Printf("Outfile Extension: %s", path.Ext(outFile))
	// log.Printf("Outfile base: %s", path.Base(outFile))
	// log.Printf("Outfile dir: %s", path.Dir(outFile))

	if _, dirErr := os.Stat(path.Dir(outFile)); os.IsNotExist(dirErr) {
		log.Fatal(dirErr)
	}

	data, err := json.MarshalIndent(swagger, "", "    ")
	if err != nil {
		log.Fatal(err)
	}
	file, _ := os.OpenFile(
		outFile,
		os.O_WRONLY|os.O_TRUNC|os.O_CREATE,
		0666,
	)
	defer file.Close()
	log.Printf("Writing swagger definition to %s", outFile)
	_, outFileErr := file.Write(data)
	if outFileErr != nil {
		log.Fatal(outFileErr)
	}
}
<commit_msg>Added a help flag<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"path"

	yaml "gopkg.in/yaml.v2"
)

func main() {

	help := flag.Bool("h", false, "Help")
	init := flag.Bool("i", false, "Initialize the swagger-meta.json file with default values")
	sourceDir := flag.String("s", ".", "The root of the source code you want swagger-gen to scan and build a swagger spec from. Defaults to current directory")
	outDir := flag.String("o", ".", "The path to the directory where the generated swagger file will be output to. Defaults to current directory")
	format := flag.String("f", "json", "Output format. json | yaml. Defaults to json")

	flag.Parse()

	if *help == true {
		fmt.Println(`
			Swagger Generation Tool
			------------------------------------------------------
			Usage:
				// Initialize a swagger-meta.json file
				swagger-gen -i

				// Generate swagger documentation
				swagger-gen -s path/to/src -o path/to/out -f json

		`)
		return
	}

	if _, dirErr := os.Stat(path.Dir(*sourceDir)); os.IsNotExist(dirErr) {
		log.Fatal(dirErr)
	}

	swaggerMetaPath := path.Join(*sourceDir, "swagger-meta.json")

	// Init command (-i)
	if *init == true {
		swagger := Swagger{}
		swagger.Swagger = "2.0"
		swagger.Info = SwaggerInfo{}
		swagger.Info.Description = "My API Description"
		swagger.Info.TermsOfService = "TOS"
		swagger.Info.Title = "My api title"
		swagger.Info.Version = "0.1.0"
		swagger.Info.Contact = Contact{
			"example@email.com",
		}

		swagger.Info.License = License{
			"Apache 2.0",
			"http://www.apache.org/licenses/LICENSE-2.0.html",
		}

		swagger.Host = "myhost.com"
		swagger.BasePath = "/v1"
		swagger.Tags = []Tag{}
		swagger.Schemes = []string{
			"http",
		}
		toJSON(swagger, swaggerMetaPath)
		log.Printf("Swagger meta file generated at path %s", swaggerMetaPath)
		os.Exit(0)
	}

	jsonBytes, jsonBytesErr := ReadJSONToBytes(swaggerMetaPath)

	if jsonBytesErr != nil {
		log.Fatal(jsonBytesErr)
	}

	swaggerf := Swaggerf{}

	swaggerf.ParseSwaggerConfig(jsonBytes)

	// Build the swagger object
	swaggerf.BuildSwagger(*sourceDir)
	var outFile string
	switch *format {
	case "json":
		outFile = "swagger.json"
		outPath := path.Join(*outDir, outFile)
		toJSON(swaggerf.Swagger, outPath)
	case "yaml":
		outFile = "swagger.yaml"
		outPath := path.Join(*outDir, outFile)
		toYAML(swaggerf.Swagger, outPath)
	default:
		log.Fatal("Invalid output format. Should be `json` or `yaml`")
	}
}

func toYAML(swagger Swagger, outFile string) {
	data, _ := yaml.Marshal(&swagger)
	file, _ := os.OpenFile(
		outFile,
		os.O_WRONLY|os.O_TRUNC|os.O_CREATE,
		0666,
	)
	defer file.Close()
	_, outFileErr := file.Write(data)
	if outFileErr != nil {
		log.Fatal(outFileErr)
	}
}

func toJSON(swagger Swagger, outFile string) {

	// log.Printf("Outfile Extension: %s", path.Ext(outFile))
	// log.Printf("Outfile base: %s", path.Base(outFile))
	// log.Printf("Outfile dir: %s", path.Dir(outFile))

	if _, dirErr := os.Stat(path.Dir(outFile)); os.IsNotExist(dirErr) {
		log.Fatal(dirErr)
	}

	data, err := json.MarshalIndent(swagger, "", "    ")
	if err != nil {
		log.Fatal(err)
	}
	file, _ := os.OpenFile(
		outFile,
		os.O_WRONLY|os.O_TRUNC|os.O_CREATE,
		0666,
	)
	defer file.Close()
	log.Printf("Writing swagger definition to %s", outFile)
	_, outFileErr := file.Write(data)
	if outFileErr != nil {
		log.Fatal(outFileErr)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode"

	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/fatih/color"
	"github.com/lucagrulla/cw/cloudwatch"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

const (
	timeFormat = "2006-01-02T15:04:05"
	version    = "3.3.0"
)

var (
	kp = kingpin.New("cw", "The best way to tail AWS Cloudwatch Logs from your terminal.")

	awsProfile     = kp.Flag("profile", "The target AWS profile. By default cw will use the default profile defined in the .aws/credentials file.").Short('p').String()
	awsRegion      = kp.Flag("region", "The target AWS region. By default cw will use the default region defined in the .aws/credentials file.").Short('r').String()
	awsEndpointURL = kp.Flag("endpoint-url", "The target AWS endpoint url. By default cw will use the default aws endpoints.").Short('u').String()
	noColor        = kp.Flag("no-color", "Disable coloured output.").Short('c').Default("false").Bool()
	debug          = kp.Flag("debug", "Enable debug logging.").Short('d').Default("false").Hidden().Bool()

	lsCommand = kp.Command("ls", "Show an entity.")

	lsGroups = lsCommand.Command("groups", "Show all groups.")

	lsStreams      = lsCommand.Command("streams", "Show all streams in a given log group.")
	lsLogGroupName = lsStreams.Arg("group", "The group name.").Required().String()

	tailCommand        = kp.Command("tail", "Tail log groups/streams.")
	logGroupStreamName = tailCommand.Arg("groupName[:logStreamPrefix]", "The log group and stream name, with group:prefix syntax."+
		"Stream name can be just the prefix. If no stream name is specified all stream names in the given group will be tailed."+
		"Multiple group/stream tuple can be passed. e.g. cw tail group1:prefix1 group2:prefix2 group3:prefix3.").Strings()

	follow          = tailCommand.Flag("follow", "Don't stop when the end of streams is reached, but rather wait for additional data to be appended.").Short('f').Default("false").Bool()
	retry           = tailCommand.Flag("retry", "retry if groups/streams not avaialble").Short('e').Default("false").Bool()
	printTimestamp  = tailCommand.Flag("timestamp", "Print the event timestamp.").Short('t').Default("false").Bool()
	printEventID    = tailCommand.Flag("event-id", "Print the event Id.").Short('i').Default("false").Bool()
	printStreamName = tailCommand.Flag("stream-name", "Print the log stream name this event belongs to.").Short('s').Default("false").Bool()
	printGroupName  = tailCommand.Flag("group-name", "Print the log group name this event belongs to.").Short('n').Default("false").Bool()
	startTime       = tailCommand.Flag("start", "The UTC start time. Passed as either date/time or human-friendly format."+
		" The human-friendly format accepts the number of days, hours and minutes prior to the present. "+
		"Denote days with 'd', hours with 'h' and minutes with 'm' i.e. 80m, 4h30m, 2d4h."+
		" If just time is used (format: hh[:mm]) it is expanded to today at the given time."+
		" Full available date/time format: 2017-02-27[T09[:00[:00]].").
		Short('b').Default(time.Now().UTC().Add(-30 * time.Second).Format(timeFormat)).String()
	endTime = tailCommand.Flag("end", "The UTC end time. Passed as either date/time or human-friendly format. "+
		" The human-friendly format accepts the number of days, hours and minutes prior to the present. "+
		"Denote days with 'd', hours with 'h' and minutes with 'm' i.e. 80m, 4h30m, 2d4h."+
		"If just time is used (format: hh[:mm]) it is expanded to today at the given time. Full available date/time format: 2017-02-27[T09[:00[:00]].").
		Short('e').Default("").String()
	local = tailCommand.Flag("local", "Treat date and time in Local timezone.").Short('l').Default("false").Bool()
	grep  = tailCommand.Flag("grep", "Pattern to filter logs by. See http://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/FilterAndPatternSyntax.html for syntax.").
		Short('g').Default("").String()
	grepv = tailCommand.Flag("grepv", "Equivalent of grep --invert-match. Invert match pattern to filter logs by.").Short('v').Default("").String()
)

func timestampToTime(timeStamp *string) (time.Time, error) {
	var zone *time.Location
	if *local {
		zone = time.Local
	} else {
		zone = time.UTC
	}
	if regexp.MustCompile(`^\d{4}-\d{2}-\d{2}$`).MatchString(*timeStamp) {
		t, _ := time.ParseInLocation("2006-01-02", *timeStamp, zone)
		return t, nil
	} else if regexp.MustCompile(`^\d{4}-\d{2}-\d{2}T\d{2}$`).MatchString(*timeStamp) {
		t, _ := time.ParseInLocation("2006-01-02T15", *timeStamp, zone)
		return t, nil
	} else if regexp.MustCompile(`^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}$`).MatchString(*timeStamp) {
		t, _ := time.ParseInLocation("2006-01-02T15:04", *timeStamp, zone)
		return t, nil
	} else if regexp.MustCompile(`^\d{1,2}$`).MatchString(*timeStamp) {
		y, m, d := time.Now().In(zone).Date()
		t, _ := strconv.Atoi(*timeStamp)
		return time.Date(y, m, d, t, 0, 0, 0, zone), nil
	} else if res := regexp.MustCompile(`^(?P<Hour>\d{1,2}):(?P<Minute>\d{2})$`).FindStringSubmatch(*timeStamp); res != nil {
		y, m, d := time.Now().Date()

		t, _ := strconv.Atoi(res[1])
		mm, _ := strconv.Atoi(res[2])

		return time.Date(y, m, d, t, mm, 0, 0, zone), nil
	} else if res := regexp.MustCompile(
		`^(?:(?P<Day>\d{1,})(?:d))?(?P<HourMinute>(?:\d{1,}h)?(?:\d{1,}m)?)?$`).FindStringSubmatch(
		*timeStamp); res != nil {
		// Unfortunately, ParseDuration does not support day time unit
		days, _ := strconv.Atoi(res[1])
		d, _ := time.ParseDuration(res[2])

		t := time.Now().In(zone).AddDate(0, 0, -days).Add(-d)
		y, m, dd := t.Date()
		return time.Date(y, m, dd, t.Hour(), t.Minute(), 0, 0, zone), nil
	}

	//TODO check even last scenario and if it's not a recognized pattern throw an error
	t, err := time.ParseInLocation("2006-01-02T15:04:05", *timeStamp, zone)
	if err != nil {
		return t, err
	}
	return t, nil
}

type logEvent struct {
	logEvent cloudwatchlogs.FilteredLogEvent
	logGroup string
}

func formatLogMsg(ev logEvent, printTime *bool, printStreamName *bool, printGroupName *bool) string {
	msg := *ev.logEvent.Message
	if *printEventID {
		msg = fmt.Sprintf("%s - %s", color.YellowString(*ev.logEvent.EventId), msg)
	}
	if *printStreamName {
		msg = fmt.Sprintf("%s - %s", color.BlueString(*ev.logEvent.LogStreamName), msg)
	}

	if *printGroupName {
		msg = fmt.Sprintf("%s - %s", color.CyanString(ev.logGroup), msg)
	}

	if *printTime {
		eventTimestamp := *ev.logEvent.Timestamp / 1000
		ts := time.Unix(eventTimestamp, 0).Format(timeFormat)
		msg = fmt.Sprintf("%s - %s", color.GreenString(ts), msg)
	}
	return msg
}

func fromStdin() []string {
	var groups []string
	info, _ := os.Stdin.Stat()
	if info.Mode()&os.ModeNamedPipe != 0 {
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			input := scanner.Text()
			if len(input) > 0 {
				tokens := strings.FieldsFunc(strings.TrimSpace(scanner.Text()), func(c rune) bool {
					return unicode.IsSpace(c)
				})
				groups = append(groups, tokens...)
			}
		}
	}
	return groups
}

func main() {
	log := log.New(ioutil.Discard, "", log.LstdFlags)
	kp.Version(version).Author("Luca Grulla")

	defer newVersionMsg(version, fetchLatestVersion())
	go versionCheckOnSigterm()

	cmd := kingpin.MustParse(kp.Parse(os.Args[1:]))
	if *debug {
		log.SetOutput(os.Stderr)
		log.Println("Debug mode is on.")
	}
	if *noColor {
		color.NoColor = true
	}

	c := cloudwatch.New(awsEndpointURL, awsProfile, awsRegion, log)

	switch cmd {
	case "ls groups":

		for msg := range cloudwatch.LsGroups(c) {
			fmt.Println(*msg)
		}
	case "ls streams":
		for msg := range cloudwatch.LsStreams(c, lsLogGroupName, nil) {
			fmt.Println(*msg)
		}
	case "tail":
		if additionalInput := fromStdin(); additionalInput != nil {
			*logGroupStreamName = append(*logGroupStreamName, additionalInput...)
		}
		if len(*logGroupStreamName) == 0 {
			fmt.Fprintln(os.Stderr, "cw: error: required argument 'groupName[:logStreamPrefix]' not provided, try --help")
			os.Exit(1)
		}

		st, err := timestampToTime(startTime)
		if err != nil {
			fmt.Fprintf(os.Stderr, "can't parse %s as a valid date/time\n", *startTime)
			os.Exit(1)
		}
		var et time.Time
		if *endTime != "" {
			endT, errr := timestampToTime(endTime)
			if errr != nil {
				fmt.Fprintf(os.Stderr, "can't parse %s as a valid date/time\n", *endTime)
				os.Exit(1)
			} else {
				et = endT
			}
		}
		out := make(chan *logEvent)

		var wg sync.WaitGroup

		triggerChannels := make([]chan<- time.Time, len(*logGroupStreamName))

		coordinator := &tailCoordinator{log: log}
		for idx, gs := range *logGroupStreamName {
			trigger := make(chan time.Time, 1)
			go func(groupStream string) {
				tokens := strings.Split(groupStream, ":")
				var prefix string
				group := tokens[0]
				if len(tokens) > 1 && tokens[1] != "*" {
					prefix = tokens[1]
				}
				ch, e := cloudwatch.Tail(c, &group, &prefix, follow, retry, &st, &et, grep, grepv, trigger, log)
				if e != nil {
					fmt.Fprintln(os.Stderr, e.Error())
				}
				for c := range ch {
					out <- &logEvent{logEvent: *c, logGroup: group}
				}
				coordinator.remove(trigger)
				wg.Done()
			}(gs)
			triggerChannels[idx] = trigger
			wg.Add(1)
		}

		coordinator.start(triggerChannels)

		go func() {
			wg.Wait()
			log.Println("closing main channel...")

			close(out)
		}()

		for logEv := range out {
			fmt.Println(formatLogMsg(*logEv, printTimestamp, printStreamName, printGroupName))
		}
	}
}
<commit_msg>remove short version flag for retry<commit_after>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode"

	"github.com/aws/aws-sdk-go/service/cloudwatchlogs"
	"github.com/fatih/color"
	"github.com/lucagrulla/cw/cloudwatch"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
)

const (
	timeFormat = "2006-01-02T15:04:05"
	version    = "3.3.0"
)

var (
	kp = kingpin.New("cw", "The best way to tail AWS Cloudwatch Logs from your terminal.")

	awsProfile     = kp.Flag("profile", "The target AWS profile. By default cw will use the default profile defined in the .aws/credentials file.").Short('p').String()
	awsRegion      = kp.Flag("region", "The target AWS region. By default cw will use the default region defined in the .aws/credentials file.").Short('r').String()
	awsEndpointURL = kp.Flag("endpoint-url", "The target AWS endpoint url. By default cw will use the default aws endpoints.").Short('u').String()
	noColor        = kp.Flag("no-color", "Disable coloured output.").Short('c').Default("false").Bool()
	debug          = kp.Flag("debug", "Enable debug logging.").Short('d').Default("false").Hidden().Bool()

	lsCommand = kp.Command("ls", "Show an entity.")

	lsGroups = lsCommand.Command("groups", "Show all groups.")

	lsStreams      = lsCommand.Command("streams", "Show all streams in a given log group.")
	lsLogGroupName = lsStreams.Arg("group", "The group name.").Required().String()

	tailCommand        = kp.Command("tail", "Tail log groups/streams.")
	logGroupStreamName = tailCommand.Arg("groupName[:logStreamPrefix]", "The log group and stream name, with group:prefix syntax."+
		"Stream name can be just the prefix. If no stream name is specified all stream names in the given group will be tailed."+
		"Multiple group/stream tuple can be passed. e.g. cw tail group1:prefix1 group2:prefix2 group3:prefix3.").Strings()

	follow          = tailCommand.Flag("follow", "Don't stop when the end of streams is reached, but rather wait for additional data to be appended.").Short('f').Default("false").Bool()
	retry           = tailCommand.Flag("retry", "retry if groups/streams not avaialble").Default("false").Bool()
	printTimestamp  = tailCommand.Flag("timestamp", "Print the event timestamp.").Short('t').Default("false").Bool()
	printEventID    = tailCommand.Flag("event-id", "Print the event Id.").Short('i').Default("false").Bool()
	printStreamName = tailCommand.Flag("stream-name", "Print the log stream name this event belongs to.").Short('s').Default("false").Bool()
	printGroupName  = tailCommand.Flag("group-name", "Print the log group name this event belongs to.").Short('n').Default("false").Bool()
	startTime       = tailCommand.Flag("start", "The UTC start time. Passed as either date/time or human-friendly format."+
		" The human-friendly format accepts the number of days, hours and minutes prior to the present. "+
		"Denote days with 'd', hours with 'h' and minutes with 'm' i.e. 80m, 4h30m, 2d4h."+
		" If just time is used (format: hh[:mm]) it is expanded to today at the given time."+
		" Full available date/time format: 2017-02-27[T09[:00[:00]].").
		Short('b').Default(time.Now().UTC().Add(-30 * time.Second).Format(timeFormat)).String()
	endTime = tailCommand.Flag("end", "The UTC end time. Passed as either date/time or human-friendly format. "+
		" The human-friendly format accepts the number of days, hours and minutes prior to the present. "+
		"Denote days with 'd', hours with 'h' and minutes with 'm' i.e. 80m, 4h30m, 2d4h."+
		"If just time is used (format: hh[:mm]) it is expanded to today at the given time. Full available date/time format: 2017-02-27[T09[:00[:00]].").
		Short('e').Default("").String()
	local = tailCommand.Flag("local", "Treat date and time in Local timezone.").Short('l').Default("false").Bool()
	grep  = tailCommand.Flag("grep", "Pattern to filter logs by. See http://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/FilterAndPatternSyntax.html for syntax.").
		Short('g').Default("").String()
	grepv = tailCommand.Flag("grepv", "Equivalent of grep --invert-match. Invert match pattern to filter logs by.").Short('v').Default("").String()
)

func timestampToTime(timeStamp *string) (time.Time, error) {
	var zone *time.Location
	if *local {
		zone = time.Local
	} else {
		zone = time.UTC
	}
	if regexp.MustCompile(`^\d{4}-\d{2}-\d{2}$`).MatchString(*timeStamp) {
		t, _ := time.ParseInLocation("2006-01-02", *timeStamp, zone)
		return t, nil
	} else if regexp.MustCompile(`^\d{4}-\d{2}-\d{2}T\d{2}$`).MatchString(*timeStamp) {
		t, _ := time.ParseInLocation("2006-01-02T15", *timeStamp, zone)
		return t, nil
	} else if regexp.MustCompile(`^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}$`).MatchString(*timeStamp) {
		t, _ := time.ParseInLocation("2006-01-02T15:04", *timeStamp, zone)
		return t, nil
	} else if regexp.MustCompile(`^\d{1,2}$`).MatchString(*timeStamp) {
		y, m, d := time.Now().In(zone).Date()
		t, _ := strconv.Atoi(*timeStamp)
		return time.Date(y, m, d, t, 0, 0, 0, zone), nil
	} else if res := regexp.MustCompile(`^(?P<Hour>\d{1,2}):(?P<Minute>\d{2})$`).FindStringSubmatch(*timeStamp); res != nil {
		y, m, d := time.Now().Date()

		t, _ := strconv.Atoi(res[1])
		mm, _ := strconv.Atoi(res[2])

		return time.Date(y, m, d, t, mm, 0, 0, zone), nil
	} else if res := regexp.MustCompile(
		`^(?:(?P<Day>\d{1,})(?:d))?(?P<HourMinute>(?:\d{1,}h)?(?:\d{1,}m)?)?$`).FindStringSubmatch(
		*timeStamp); res != nil {
		// Unfortunately, ParseDuration does not support day time unit
		days, _ := strconv.Atoi(res[1])
		d, _ := time.ParseDuration(res[2])

		t := time.Now().In(zone).AddDate(0, 0, -days).Add(-d)
		y, m, dd := t.Date()
		return time.Date(y, m, dd, t.Hour(), t.Minute(), 0, 0, zone), nil
	}

	//TODO check even last scenario and if it's not a recognized pattern throw an error
	t, err := time.ParseInLocation("2006-01-02T15:04:05", *timeStamp, zone)
	if err != nil {
		return t, err
	}
	return t, nil
}

type logEvent struct {
	logEvent cloudwatchlogs.FilteredLogEvent
	logGroup string
}

func formatLogMsg(ev logEvent, printTime *bool, printStreamName *bool, printGroupName *bool) string {
	msg := *ev.logEvent.Message
	if *printEventID {
		msg = fmt.Sprintf("%s - %s", color.YellowString(*ev.logEvent.EventId), msg)
	}
	if *printStreamName {
		msg = fmt.Sprintf("%s - %s", color.BlueString(*ev.logEvent.LogStreamName), msg)
	}

	if *printGroupName {
		msg = fmt.Sprintf("%s - %s", color.CyanString(ev.logGroup), msg)
	}

	if *printTime {
		eventTimestamp := *ev.logEvent.Timestamp / 1000
		ts := time.Unix(eventTimestamp, 0).Format(timeFormat)
		msg = fmt.Sprintf("%s - %s", color.GreenString(ts), msg)
	}
	return msg
}

func fromStdin() []string {
	var groups []string
	info, _ := os.Stdin.Stat()
	if info.Mode()&os.ModeNamedPipe != 0 {
		scanner := bufio.NewScanner(os.Stdin)
		for scanner.Scan() {
			input := scanner.Text()
			if len(input) > 0 {
				tokens := strings.FieldsFunc(strings.TrimSpace(scanner.Text()), func(c rune) bool {
					return unicode.IsSpace(c)
				})
				groups = append(groups, tokens...)
			}
		}
	}
	return groups
}

func main() {
	log := log.New(ioutil.Discard, "", log.LstdFlags)
	kp.Version(version).Author("Luca Grulla")

	defer newVersionMsg(version, fetchLatestVersion())
	go versionCheckOnSigterm()

	cmd := kingpin.MustParse(kp.Parse(os.Args[1:]))
	if *debug {
		log.SetOutput(os.Stderr)
		log.Println("Debug mode is on.")
	}
	if *noColor {
		color.NoColor = true
	}

	c := cloudwatch.New(awsEndpointURL, awsProfile, awsRegion, log)

	switch cmd {
	case "ls groups":

		for msg := range cloudwatch.LsGroups(c) {
			fmt.Println(*msg)
		}
	case "ls streams":
		for msg := range cloudwatch.LsStreams(c, lsLogGroupName, nil) {
			fmt.Println(*msg)
		}
	case "tail":
		if additionalInput := fromStdin(); additionalInput != nil {
			*logGroupStreamName = append(*logGroupStreamName, additionalInput...)
		}
		if len(*logGroupStreamName) == 0 {
			fmt.Fprintln(os.Stderr, "cw: error: required argument 'groupName[:logStreamPrefix]' not provided, try --help")
			os.Exit(1)
		}

		st, err := timestampToTime(startTime)
		if err != nil {
			fmt.Fprintf(os.Stderr, "can't parse %s as a valid date/time\n", *startTime)
			os.Exit(1)
		}
		var et time.Time
		if *endTime != "" {
			endT, errr := timestampToTime(endTime)
			if errr != nil {
				fmt.Fprintf(os.Stderr, "can't parse %s as a valid date/time\n", *endTime)
				os.Exit(1)
			} else {
				et = endT
			}
		}
		out := make(chan *logEvent)

		var wg sync.WaitGroup

		triggerChannels := make([]chan<- time.Time, len(*logGroupStreamName))

		coordinator := &tailCoordinator{log: log}
		for idx, gs := range *logGroupStreamName {
			trigger := make(chan time.Time, 1)
			go func(groupStream string) {
				tokens := strings.Split(groupStream, ":")
				var prefix string
				group := tokens[0]
				if len(tokens) > 1 && tokens[1] != "*" {
					prefix = tokens[1]
				}
				ch, e := cloudwatch.Tail(c, &group, &prefix, follow, retry, &st, &et, grep, grepv, trigger, log)
				if e != nil {
					fmt.Fprintln(os.Stderr, e.Error())
				}
				for c := range ch {
					out <- &logEvent{logEvent: *c, logGroup: group}
				}
				coordinator.remove(trigger)
				wg.Done()
			}(gs)
			triggerChannels[idx] = trigger
			wg.Add(1)
		}

		coordinator.start(triggerChannels)

		go func() {
			wg.Wait()
			log.Println("closing main channel...")

			close(out)
		}()

		for logEv := range out {
			fmt.Println(formatLogMsg(*logEv, printTimestamp, printStreamName, printGroupName))
		}
	}
}
<|endoftext|>


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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"os/user"
	"runtime"
	"syscall"
	"time"

	"github.com/bpowers/fuse"
	"github.com/bpowers/fuse/fs"
)

const (
	usage = `Usage: %s [OPTION...] MOUNTPOINT
Slack as a filesystem.

The -token-path option or the SLACKFS_TOKEN_PATH environmental
variable can be used to specify the location of a file containing your
slack API token.  If neither is set, the default token file location
of '%s' is tried.

Options:
`

	noTokenFound = `No token file could be read.  If you haven't generated one yet, go to
https://api.slack.com/web and scroll to the bottom. If 'none' is
listed as the token, then click the green 'Create Token' button next
to it. Copy the token that appears in red into an otherwise empty file
stored at ~/.slack-token (or a location of your choosing, specified in
the environment by SLACKFS_TOKEN_PATH or after the commandline flag
-token-path).

`
)

var defaultTokenPath string // initalized in init() below

func debugOut(msg interface{}) {
	log.Printf("%s", msg)
}

func init() {
	// with more than 1 maxproc, we end up with wild, terrible
	// memory fragmentation. pss goes from ~14 MB to 180 MB.
	runtime.GOMAXPROCS(1)

	home := ""

	if u, err := user.Current(); err == nil {
		home = u.HomeDir
	} else {
		home = os.Getenv("HOME")
	}
	// if we're statically linked + HOME isn't set (running as a
	// service under systemd?), we won't know our home directory.
	// It is least surprising to fail, rather than look in
	// whatever random directory we're in.
	if home != "" {
		defaultTokenPath = fmt.Sprintf("%s/.slack-token", home)
	}
}

func getToken(flagPath string) string {
	path := defaultTokenPath
	if envPath := os.Getenv("SLACKFS_TOKEN_PATH"); envPath != "" {
		path = envPath
	}
	if flagPath != "" {
		path = flagPath
	}
	if path == "" {
		return ""
	}
	buf, err := ioutil.ReadFile(path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error reading token file: %s\n\n", err)
		return ""
	}
	return string(bytes.TrimSpace(buf))
}

var memProfile, cpuProfile string

func main() {
	flag.StringVar(&memProfile, "memprofile", "",
		"write memory profile to this file")
	flag.StringVar(&cpuProfile, "cpuprofile", "",
		"write cpu profile to this file")
	offline := flag.String("offline", "",
		"specified JSON info response file to use offline")
	tokenPath := flag.String("token-path", "", "Slack API token")

	verbose := flag.Bool("v", false, "verbose FUSE logging")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0], defaultTokenPath)
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	token := getToken(*tokenPath)
	if token == "" {
		fmt.Fprintf(os.Stderr, noTokenFound)
		flag.Usage()
		os.Exit(1)
	}

	var debugFn func(msg interface{})
	if *verbose {
		debugFn = debugOut
	}

	mountpoint := flag.Arg(0)

	prof, err := NewProf(memProfile, cpuProfile)
	if err != nil {
		log.Fatal(err)
	}
	// if -memprof or -cpuprof haven't been set on the command
	// line, these are nops
	prof.Start()
	// Set up channel on which to send signal notifications.  We
	// must use a buffered channel or risk missing the signal if
	// we're not ready to receive when the signal is sent.
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGQUIT)
	go func() {
		for range sigChan {
			prof.Stop()
		}
	}()

	var conn *FSConn
	if *offline != "" {
		conn, err = NewOfflineFSConn(*offline)
	} else {
		conn, err = NewFSConn(token)
	}
	if err != nil {
		log.Fatalf("NewFS: %s", err)
	}

	c, err := fuse.Mount(
		mountpoint,
		fuse.FSName("slack"),
		fuse.Subtype("slackfs"),
		fuse.LocalVolume(),
		fuse.VolumeName("Slack"),
	)
	if err != nil {
		log.Fatalf("Mount(%s): %s", mountpoint, err)
	}
	defer c.Close()

	err = fs.Serve(c, conn.super, debugFn)
	if err != nil {
		log.Fatal(err)
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		log.Fatal(err)
	}

	time.Sleep(120 * time.Second)
	log.Printf("done done")
}
<commit_msg>try to mkdir the mountpoint if it doesn't exist<commit_after>// Copyright 2015 Bobby Powers. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"os/user"
	"runtime"
	"syscall"
	"time"

	"github.com/bpowers/fuse"
	"github.com/bpowers/fuse/fs"
)

const (
	usage = `Usage: %s [OPTION...] MOUNTPOINT
Slack as a filesystem.

The -token-path option or the SLACKFS_TOKEN_PATH environmental
variable can be used to specify the location of a file containing your
slack API token.  If neither is set, the default token file location
of '%s' is tried.

Options:
`

	noTokenFound = `No token file could be read.  If you haven't generated one yet, go to
https://api.slack.com/web and scroll to the bottom. If 'none' is
listed as the token, then click the green 'Create Token' button next
to it. Copy the token that appears in red into an otherwise empty file
stored at ~/.slack-token (or a location of your choosing, specified in
the environment by SLACKFS_TOKEN_PATH or after the commandline flag
-token-path).

`
)

var defaultTokenPath string // initalized in init() below

func debugOut(msg interface{}) {
	log.Printf("%s", msg)
}

func init() {
	// with more than 1 maxproc, we end up with wild, terrible
	// memory fragmentation. pss goes from ~14 MB to 180 MB.
	runtime.GOMAXPROCS(1)

	home := ""

	if u, err := user.Current(); err == nil {
		home = u.HomeDir
	} else {
		home = os.Getenv("HOME")
	}
	// if we're statically linked + HOME isn't set (running as a
	// service under systemd?), we won't know our home directory.
	// It is least surprising to fail, rather than look in
	// whatever random directory we're in.
	if home != "" {
		defaultTokenPath = fmt.Sprintf("%s/.slack-token", home)
	}
}

func getToken(flagPath string) string {
	path := defaultTokenPath
	if envPath := os.Getenv("SLACKFS_TOKEN_PATH"); envPath != "" {
		path = envPath
	}
	if flagPath != "" {
		path = flagPath
	}
	if path == "" {
		return ""
	}
	buf, err := ioutil.ReadFile(path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error reading token file: %s\n\n", err)
		return ""
	}
	return string(bytes.TrimSpace(buf))
}

var memProfile, cpuProfile string

func main() {
	flag.StringVar(&memProfile, "memprofile", "",
		"write memory profile to this file")
	flag.StringVar(&cpuProfile, "cpuprofile", "",
		"write cpu profile to this file")
	offline := flag.String("offline", "",
		"specified JSON info response file to use offline")
	tokenPath := flag.String("token-path", "", "Slack API token")

	verbose := flag.Bool("v", false, "verbose FUSE logging")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0], defaultTokenPath)
		flag.PrintDefaults()
	}
	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	token := getToken(*tokenPath)
	if token == "" {
		fmt.Fprintf(os.Stderr, noTokenFound)
		flag.Usage()
		os.Exit(1)
	}

	var debugFn func(msg interface{})
	if *verbose {
		debugFn = debugOut
	}

	mountpoint := flag.Arg(0)

	prof, err := NewProf(memProfile, cpuProfile)
	if err != nil {
		log.Fatal(err)
	}
	// if -memprof or -cpuprof haven't been set on the command
	// line, these are nops
	prof.Start()
	// Set up channel on which to send signal notifications.  We
	// must use a buffered channel or risk missing the signal if
	// we're not ready to receive when the signal is sent.
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGQUIT)
	go func() {
		for range sigChan {
			prof.Stop()
		}
	}()

	if err := os.MkdirAll(mountpoint, 0777); err != nil {
		log.Fatalf("couldn't create mountpoint: %s", err)
	}

	var conn *FSConn
	if *offline != "" {
		conn, err = NewOfflineFSConn(*offline)
	} else {
		conn, err = NewFSConn(token)
	}
	if err != nil {
		log.Fatalf("NewFS: %s", err)
	}

	c, err := fuse.Mount(
		mountpoint,
		fuse.FSName("slack"),
		fuse.Subtype("slackfs"),
		fuse.LocalVolume(),
		fuse.VolumeName("Slack"),
	)
	if err != nil {
		log.Fatalf("Mount(%s): %s", mountpoint, err)
	}
	defer c.Close()

	err = fs.Serve(c, conn.super, debugFn)
	if err != nil {
		log.Fatal(err)
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		log.Fatal(err)
	}

	time.Sleep(120 * time.Second)
	log.Printf("done done")
}
<|endoftext|>


<commit_before>package miniprofiler

import (
	"fmt"
	"io"
	"math"
	"sort"
	"strings"
	"time"
)

var (
	mp        *MiniProfiler
	enabled   bool
	condition func() bool
)

type MiniProfilerData struct {
	description string
	steps       map[string]int64
	lastStep    time.Time
	memos       []string
}

type MiniProfiler struct {
	profiles []*MiniProfilerData
}

func init() {
	mp = new(MiniProfiler)
	mp.profiles = make([]*MiniProfilerData, 0)
	enabled = true
	condition = func() bool { return true }
}

func Enable() {
	enabled = true
}

func Disable() {
	enabled = false
}

func SetCondition(c func() bool) {
	condition = c
}

func Begin(description string) *MiniProfilerData {
	if !enabled {
		return nil
	}
	if !condition() {
		return nil
	}
	return &MiniProfilerData{description, make(map[string]int64, 0), time.Now(), make([]string, 0)}
}

func Flush(writer io.Writer) {
	for _, prof := range mp.profiles {
		outputs := []string{"log:MP"}

		for tag, val := range prof.steps {
			outputs = append(outputs, fmt.Sprintf("%s:%d", tag, val))
		}

		outputs = append(outputs, fmt.Sprintf("description:%s", prof.description))

		memoOutput := []string{}
		for _, m := range prof.memos {
			memoOutput = append(memoOutput, m)
		}
		outputs = append(outputs, fmt.Sprintf("memo:%s", strings.Join(memoOutput, ",")))

		fmt.Fprintln(writer, strings.Join(outputs, "\t"))
	}
	mp.profiles = make([]*MiniProfilerData, 0)
}

type Measure struct {
	Description string
	Count       int
	Total       int
	Mean        int
	Min         int
	Max         int
}

type By func(a, b *Measure) bool

func (by By) Sort(measures []*Measure) {
	ms := &measureSorter{
		measures: measures,
		by:       by,
	}
	sort.Sort(ms)
}

type measureSorter struct {
	measures []*Measure
	by       func(a, b *Measure) bool
}

func (s *measureSorter) Len() int {
	return len(s.measures)
}

func (s *measureSorter) Swap(i, j int) {
	s.measures[i], s.measures[j] = s.measures[j], s.measures[i]
}

func (s *measureSorter) Less(i, j int) bool {
	return s.by(s.measures[i], s.measures[j])
}

type Column struct {
	Name    string
	Summary string
	Sort    By
}

var (
	columns = []*Column{
		&Column{Name: "Count", Summary: "Count", Sort: func(a, b *Measure) bool { return a.Count > b.Count }},
		&Column{Name: "Total", Summary: "Total", Sort: func(a, b *Measure) bool { return a.Total > b.Total }},
		&Column{Name: "Mean", Summary: "Mean", Sort: func(a, b *Measure) bool { return a.Mean > b.Mean }},
		&Column{Name: "Max", Summary: "Maximum(100 Percentile)", Sort: func(a, b *Measure) bool { return a.Max > b.Max }},
	}
)

type ByTime []*Time

type Time struct {
	Url  string
	Time int64
}

func (a ByTime) Len() int           { return len(a) }
func (a ByTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByTime) Less(i, j int) bool { return a[i].Time > a[j].Time }

func getIntegerDigitWidth(i int) int {
	var w int
	switch {
	case i < 1:
		w = 1
	default:
		w = int(math.Log10(float64(i)) + 1)
	}
	return w
}

func showMeasures(writer io.Writer, measures []*Measure) {
	countWidth := 5
	totalWidth := 5
	meanWidth := 5
	maxWidth := 5

	for _, m := range measures {
		var w int
		w = getIntegerDigitWidth(m.Count)
		if countWidth < w {
			countWidth = w
		}
		w = getIntegerDigitWidth(m.Total)
		if totalWidth < w {
			totalWidth = w
		}
		w = getIntegerDigitWidth(m.Mean)
		if meanWidth < w {
			meanWidth = w
		}
		w = getIntegerDigitWidth(m.Max)
		if maxWidth < w {
			maxWidth = w
		}
	}

	var format string
	for _, column := range columns {
		switch column.Name {
		case "Count":
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", countWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", countWidth)
		case "Total":
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", totalWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", totalWidth)
		case "Mean":
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", meanWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", meanWidth)
		default:
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", maxWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", maxWidth)
		}
	}
	fmt.Fprintf(writer, "Description\n")
	format += "%s\n"

	for _, m := range measures {
		fmt.Fprintf(writer, format, m.Count, m.Total, m.Mean, m.Min, m.Max, m.Description)
	}
}

func Analyze(writer io.Writer) {
	var times = make(map[string][]int)
	var totals = make(map[string]int)
	for _, prof := range mp.profiles {
		for tag, val := range prof.steps {
			times[tag] = append(times[tag], int(val))
			totals[tag] += int(val)
		}
	}
	var measures []*Measure
	for tag, times := range times {
		sort.Ints(times)
		count := len(times)
		total := totals[tag]
		m := &Measure{
			Description: tag,
			Count:       count,
			Total:       total,
			Mean:        total / count,
			Min:         times[0],
			Max:         times[count-1],
		}
		measures = append(measures, m)
	}

	for _, column := range columns {
		if column.Sort != nil {
			fmt.Fprintf(writer, "Sort by %s\n", column.Summary)
			By(column.Sort).Sort(measures)
			showMeasures(writer, measures)
			fmt.Fprintln(writer)
		}
	}
}

func (mpd *MiniProfilerData) Step(tag string) {
	if !enabled {
		return
	}
	if mpd == nil {
		return
	}
	now := time.Now()
	thisstep := now.Sub(mpd.lastStep).Nanoseconds()

	mpd.steps[tag] = thisstep
	mpd.lastStep = now
}

func (mpd *MiniProfilerData) Memo(memo string) {
	if !enabled {
		return
	}
	if mpd == nil {
		return
	}
	memo = memoFilter(memo)
	mpd.memos = append(mpd.memos, memo)
}

/*
remove (,|\t|\r|\n)
*/
func memoFilter(memo string) string {
	memo = strings.Replace(memo, ",", "", -1)
	memo = strings.Replace(memo, "\t", "", -1)
	memo = strings.Replace(memo, "\r", "", -1)
	memo = strings.Replace(memo, "\n", "", -1)
	return memo
}

func (mpd *MiniProfilerData) End() {
	if !enabled {
		return
	}
	if mpd == nil {
		return
	}
	mpd.Step("Last Step to End")
	mp.profiles = append(mp.profiles, mpd)
}
<commit_msg>cosmetic<commit_after>package miniprofiler

import (
	"fmt"
	"io"
	"math"
	"sort"
	"strings"
	"time"
)

var (
	mp        *MiniProfiler
	enabled   bool
	condition func() bool
)

type MiniProfilerData struct {
	description string
	steps       map[string]int64
	lastStep    time.Time
	memos       []string
}

type MiniProfiler struct {
	profiles []*MiniProfilerData
}

func init() {
	mp = new(MiniProfiler)
	mp.profiles = make([]*MiniProfilerData, 0)
	enabled = true
	condition = func() bool { return true }
}

func Enable() {
	enabled = true
}

func Disable() {
	enabled = false
}

func SetCondition(c func() bool) {
	condition = c
}

func Begin(description string) *MiniProfilerData {
	if !enabled {
		return nil
	}
	if !condition() {
		return nil
	}
	return &MiniProfilerData{description, make(map[string]int64, 0), time.Now(), make([]string, 0)}
}

func Flush(writer io.Writer) {
	for _, prof := range mp.profiles {
		outputs := []string{"log:MP"}

		for tag, val := range prof.steps {
			outputs = append(outputs, fmt.Sprintf("%s:%d", tag, val))
		}

		outputs = append(outputs, fmt.Sprintf("description:%s", prof.description))

		memoOutput := []string{}
		for _, m := range prof.memos {
			memoOutput = append(memoOutput, m)
		}
		outputs = append(outputs, fmt.Sprintf("memo:%s", strings.Join(memoOutput, ",")))

		fmt.Fprintln(writer, strings.Join(outputs, "\t"))
	}
	mp.profiles = make([]*MiniProfilerData, 0)
}

type Measure struct {
	Description string
	Count       int
	Total       int
	Mean        int
	Min         int
	Max         int
}

type By func(a, b *Measure) bool

func (by By) Sort(measures []*Measure) {
	ms := &measureSorter{
		measures: measures,
		by:       by,
	}
	sort.Sort(ms)
}

type measureSorter struct {
	measures []*Measure
	by       func(a, b *Measure) bool
}

func (s *measureSorter) Len() int {
	return len(s.measures)
}

func (s *measureSorter) Swap(i, j int) {
	s.measures[i], s.measures[j] = s.measures[j], s.measures[i]
}

func (s *measureSorter) Less(i, j int) bool {
	return s.by(s.measures[i], s.measures[j])
}

type Column struct {
	Name    string
	Summary string
	Sort    By
}

var (
	columns = []*Column{
		&Column{Name: "Count", Summary: "Count", Sort: func(a, b *Measure) bool { return a.Count > b.Count }},
		&Column{Name: "Total", Summary: "Total", Sort: func(a, b *Measure) bool { return a.Total > b.Total }},
		&Column{Name: "Mean", Summary: "Mean", Sort: func(a, b *Measure) bool { return a.Mean > b.Mean }},
		&Column{Name: "Min"},
		&Column{Name: "Max", Summary: "Maximum(100 Percentile)", Sort: func(a, b *Measure) bool { return a.Max > b.Max }},
	}
)

type ByTime []*Time

type Time struct {
	Url  string
	Time int64
}

func (a ByTime) Len() int           { return len(a) }
func (a ByTime) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByTime) Less(i, j int) bool { return a[i].Time > a[j].Time }

func getIntegerDigitWidth(i int) int {
	var w int
	switch {
	case i < 1:
		w = 1
	default:
		w = int(math.Log10(float64(i)) + 1)
	}
	return w
}

func showMeasures(writer io.Writer, measures []*Measure) {
	countWidth := 5
	totalWidth := 5
	meanWidth := 5
	maxWidth := 5

	for _, m := range measures {
		var w int
		w = getIntegerDigitWidth(m.Count)
		if countWidth < w {
			countWidth = w
		}
		w = getIntegerDigitWidth(m.Total)
		if totalWidth < w {
			totalWidth = w
		}
		w = getIntegerDigitWidth(m.Mean)
		if meanWidth < w {
			meanWidth = w
		}
		w = getIntegerDigitWidth(m.Max)
		if maxWidth < w {
			maxWidth = w
		}
	}

	var format string
	for _, column := range columns {
		switch column.Name {
		case "Count":
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", countWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", countWidth)
		case "Total":
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", totalWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", totalWidth)
		case "Mean":
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", meanWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", meanWidth)
		default:
			fmt.Fprintf(writer, fmt.Sprintf("%%%ds  ", maxWidth), column.Name)
			format += fmt.Sprintf("%%%dd  ", maxWidth)
		}
	}
	fmt.Fprintln(writer, "Description")
	format += "%s\n"

	for _, m := range measures {
		fmt.Fprintf(writer, format, m.Count, m.Total, m.Mean, m.Min, m.Max, m.Description)
	}
}

func Analyze(writer io.Writer) {
	var times = make(map[string][]int)
	var totals = make(map[string]int)
	for _, prof := range mp.profiles {
		for tag, val := range prof.steps {
			times[tag] = append(times[tag], int(val))
			totals[tag] += int(val)
		}
	}
	var measures []*Measure
	for tag, times := range times {
		sort.Ints(times)
		count := len(times)
		total := totals[tag]
		m := &Measure{
			Description: tag,
			Count:       count,
			Total:       total,
			Mean:        total / count,
			Min:         times[0],
			Max:         times[count-1],
		}
		measures = append(measures, m)
	}

	for _, column := range columns {
		if column.Sort != nil {
			fmt.Fprintf(writer, "Sort by %s\n", column.Summary)
			By(column.Sort).Sort(measures)
			showMeasures(writer, measures)
			fmt.Fprintln(writer)
		}
	}
}

func (mpd *MiniProfilerData) Step(tag string) {
	if !enabled {
		return
	}
	if mpd == nil {
		return
	}
	now := time.Now()
	thisstep := now.Sub(mpd.lastStep).Nanoseconds()

	mpd.steps[tag] = thisstep
	mpd.lastStep = now
}

func (mpd *MiniProfilerData) Memo(memo string) {
	if !enabled {
		return
	}
	if mpd == nil {
		return
	}
	memo = memoFilter(memo)
	mpd.memos = append(mpd.memos, memo)
}

/*
remove (,|\t|\r|\n)
*/
func memoFilter(memo string) string {
	memo = strings.Replace(memo, ",", "", -1)
	memo = strings.Replace(memo, "\t", "", -1)
	memo = strings.Replace(memo, "\r", "", -1)
	memo = strings.Replace(memo, "\n", "", -1)
	return memo
}

func (mpd *MiniProfilerData) End() {
	if !enabled {
		return
	}
	if mpd == nil {
		return
	}
	mpd.Step("Last Step to End")
	mp.profiles = append(mp.profiles, mpd)
}
<|endoftext|>


<commit_before>// Copyright 2015 Prometheus Team
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"
	"github.com/prometheus/common/version"

	"github.com/prometheus/alertmanager/config"
	"github.com/prometheus/alertmanager/notify"
	"github.com/prometheus/alertmanager/provider/boltmem"
	"github.com/prometheus/alertmanager/template"
	"github.com/prometheus/alertmanager/types"
)

var (
	showVersion = flag.Bool("version", false, "Print version information.")

	configFile = flag.String("config.file", "alertmanager.yml", "Alertmanager configuration file name.")
	dataDir    = flag.String("storage.path", "data/", "Base path for data storage.")

	externalURL   = flag.String("web.external-url", "", "The URL under which Alertmanager is externally reachable (for example, if Alertmanager is served via a reverse proxy). Used for generating relative and absolute links back to Alertmanager itself. If the URL has a path portion, it will be used to prefix all HTTP endpoints served by Alertmanager. If omitted, relevant URL components will be derived automatically.")
	listenAddress = flag.String("web.listen-address", ":9093", "Address to listen on for the web interface and API.")
)

var (
	configSuccess = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "alertmanager",
		Name:      "config_last_reload_successful",
		Help:      "Whether the last configuration reload attempt was successful.",
	})
	configSuccessTime = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "alertmanager",
		Name:      "config_last_reload_success_timestamp_seconds",
		Help:      "Timestamp of the last successful configuration reload.",
	})
)

func init() {
	prometheus.MustRegister(configSuccess)
	prometheus.MustRegister(configSuccessTime)
	prometheus.MustRegister(version.NewCollector("alertmanager"))
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("alertmanager"))
		os.Exit(0)
	}

	log.Infoln("Starting alertmanager", version.Info())
	log.Infoln("Build context", version.BuildContext())

	err := os.MkdirAll(*dataDir, 0777)
	if err != nil {
		log.Fatal(err)
	}

	marker := types.NewMarker()

	alerts, err := boltmem.NewAlerts(*dataDir)
	if err != nil {
		log.Fatal(err)
	}
	defer alerts.Close()

	notifies, err := boltmem.NewNotificationInfo(*dataDir)
	if err != nil {
		log.Fatal(err)
	}
	defer notifies.Close()

	silences, err := boltmem.NewSilences(*dataDir, marker)
	if err != nil {
		log.Fatal(err)
	}
	defer silences.Close()

	var (
		inhibitor *Inhibitor
		tmpl      *template.Template
		disp      *Dispatcher
	)
	defer disp.Stop()

	api := NewAPI(alerts, silences, func() AlertOverview {
		return disp.Groups()
	})

	build := func(rcvs []*config.Receiver) notify.Notifier {
		var (
			router  = notify.Router{}
			fanouts = notify.Build(rcvs, tmpl)
		)
		for name, fo := range fanouts {
			for i, n := range fo {
				n = notify.Retry(n)
				n = notify.Log(n, log.With("step", "retry"))
				n = notify.Dedup(notifies, n)
				n = notify.Log(n, log.With("step", "dedup"))

				fo[i] = n
			}
			router[name] = fo
		}
		n := notify.Notifier(router)

		n = notify.Log(n, log.With("step", "route"))
		n = notify.Silence(silences, n, marker)
		n = notify.Log(n, log.With("step", "silence"))
		n = notify.Inhibit(inhibitor, n, marker)
		n = notify.Log(n, log.With("step", "inhibit"))

		return n
	}

	amURL, err := extURL(*externalURL)
	if err != nil {
		log.Fatal(err)
	}

	reload := func() (err error) {
		log.With("file", *configFile).Infof("Loading configuration file")
		defer func() {
			if err != nil {
				log.With("file", *configFile).Errorf("Loading configuration file failed: %s", err)
				configSuccess.Set(0)
			} else {
				configSuccess.Set(1)
				configSuccessTime.Set(float64(time.Now().Unix()))
			}
		}()

		conf, err := config.LoadFile(*configFile)
		if err != nil {
			return err
		}

		api.Update(conf.String(), time.Duration(conf.Global.ResolveTimeout))

		tmpl, err = template.FromGlobs(conf.Templates...)
		if err != nil {
			return err
		}
		tmpl.ExternalURL = amURL

		inhibitor.Stop()
		disp.Stop()

		inhibitor = NewInhibitor(alerts, conf.InhibitRules, marker)
		disp = NewDispatcher(alerts, NewRoute(conf.Route, nil), build(conf.Receivers), marker)

		go disp.Run()
		go inhibitor.Run()

		return nil
	}

	if err := reload(); err != nil {
		os.Exit(1)
	}

	router := route.New()

	webReload := make(chan struct{})
	RegisterWeb(router.WithPrefix(amURL.Path), webReload)
	api.Register(router.WithPrefix(path.Join(amURL.Path, "/api")))

	log.Infoln("Listening on", *listenAddress)
	go listen(router)

	var (
		hup      = make(chan os.Signal)
		hupReady = make(chan bool)
		term     = make(chan os.Signal)
	)
	signal.Notify(hup, syscall.SIGHUP)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-hupReady
		for {
			select {
			case <-hup:
			case <-webReload:
			}
			reload()
		}
	}()

	// Wait for reload or termination signals.
	close(hupReady) // Unblock SIGHUP handler.

	<-term

	log.Infoln("Received SIGTERM, exiting gracefully...")
}

func extURL(s string) (*url.URL, error) {
	if s == "" {
		hostname, err := os.Hostname()
		if err != nil {
			return nil, err
		}
		_, port, err := net.SplitHostPort(*listenAddress)
		if err != nil {
			return nil, err
		}

		s = fmt.Sprintf("http://%s:%s/", hostname, port)
	}

	u, err := url.Parse(s)
	if err != nil {
		return nil, err
	}

	ppref := strings.TrimRight(u.Path, "/")
	if ppref != "" && !strings.HasPrefix(ppref, "/") {
		ppref = "/" + ppref
	}
	u.Path = ppref

	return u, nil
}

func listen(router *route.Router) {
	if err := http.ListenAndServe(*listenAddress, router); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>main: use flags as arguments instead of globals<commit_after>// Copyright 2015 Prometheus Team
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/route"
	"github.com/prometheus/common/version"

	"github.com/prometheus/alertmanager/config"
	"github.com/prometheus/alertmanager/notify"
	"github.com/prometheus/alertmanager/provider/boltmem"
	"github.com/prometheus/alertmanager/template"
	"github.com/prometheus/alertmanager/types"
)

var (
	showVersion = flag.Bool("version", false, "Print version information.")

	configFile = flag.String("config.file", "alertmanager.yml", "Alertmanager configuration file name.")
	dataDir    = flag.String("storage.path", "data/", "Base path for data storage.")

	externalURL   = flag.String("web.external-url", "", "The URL under which Alertmanager is externally reachable (for example, if Alertmanager is served via a reverse proxy). Used for generating relative and absolute links back to Alertmanager itself. If the URL has a path portion, it will be used to prefix all HTTP endpoints served by Alertmanager. If omitted, relevant URL components will be derived automatically.")
	listenAddress = flag.String("web.listen-address", ":9093", "Address to listen on for the web interface and API.")
)

var (
	configSuccess = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "alertmanager",
		Name:      "config_last_reload_successful",
		Help:      "Whether the last configuration reload attempt was successful.",
	})
	configSuccessTime = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "alertmanager",
		Name:      "config_last_reload_success_timestamp_seconds",
		Help:      "Timestamp of the last successful configuration reload.",
	})
)

func init() {
	prometheus.MustRegister(configSuccess)
	prometheus.MustRegister(configSuccessTime)
	prometheus.MustRegister(version.NewCollector("alertmanager"))
}

func main() {
	flag.Parse()

	if *showVersion {
		fmt.Fprintln(os.Stdout, version.Print("alertmanager"))
		os.Exit(0)
	}

	log.Infoln("Starting alertmanager", version.Info())
	log.Infoln("Build context", version.BuildContext())

	err := os.MkdirAll(*dataDir, 0777)
	if err != nil {
		log.Fatal(err)
	}

	marker := types.NewMarker()

	alerts, err := boltmem.NewAlerts(*dataDir)
	if err != nil {
		log.Fatal(err)
	}
	defer alerts.Close()

	notifies, err := boltmem.NewNotificationInfo(*dataDir)
	if err != nil {
		log.Fatal(err)
	}
	defer notifies.Close()

	silences, err := boltmem.NewSilences(*dataDir, marker)
	if err != nil {
		log.Fatal(err)
	}
	defer silences.Close()

	var (
		inhibitor *Inhibitor
		tmpl      *template.Template
		disp      *Dispatcher
	)
	defer disp.Stop()

	api := NewAPI(alerts, silences, func() AlertOverview {
		return disp.Groups()
	})

	build := func(rcvs []*config.Receiver) notify.Notifier {
		var (
			router  = notify.Router{}
			fanouts = notify.Build(rcvs, tmpl)
		)
		for name, fo := range fanouts {
			for i, n := range fo {
				n = notify.Retry(n)
				n = notify.Log(n, log.With("step", "retry"))
				n = notify.Dedup(notifies, n)
				n = notify.Log(n, log.With("step", "dedup"))

				fo[i] = n
			}
			router[name] = fo
		}
		n := notify.Notifier(router)

		n = notify.Log(n, log.With("step", "route"))
		n = notify.Silence(silences, n, marker)
		n = notify.Log(n, log.With("step", "silence"))
		n = notify.Inhibit(inhibitor, n, marker)
		n = notify.Log(n, log.With("step", "inhibit"))

		return n
	}

	amURL, err := extURL(*listenAddress, *externalURL)
	if err != nil {
		log.Fatal(err)
	}

	reload := func() (err error) {
		log.With("file", *configFile).Infof("Loading configuration file")
		defer func() {
			if err != nil {
				log.With("file", *configFile).Errorf("Loading configuration file failed: %s", err)
				configSuccess.Set(0)
			} else {
				configSuccess.Set(1)
				configSuccessTime.Set(float64(time.Now().Unix()))
			}
		}()

		conf, err := config.LoadFile(*configFile)
		if err != nil {
			return err
		}

		api.Update(conf.String(), time.Duration(conf.Global.ResolveTimeout))

		tmpl, err = template.FromGlobs(conf.Templates...)
		if err != nil {
			return err
		}
		tmpl.ExternalURL = amURL

		inhibitor.Stop()
		disp.Stop()

		inhibitor = NewInhibitor(alerts, conf.InhibitRules, marker)
		disp = NewDispatcher(alerts, NewRoute(conf.Route, nil), build(conf.Receivers), marker)

		go disp.Run()
		go inhibitor.Run()

		return nil
	}

	if err := reload(); err != nil {
		os.Exit(1)
	}

	router := route.New()

	webReload := make(chan struct{})
	RegisterWeb(router.WithPrefix(amURL.Path), webReload)
	api.Register(router.WithPrefix(path.Join(amURL.Path, "/api")))

	log.Infoln("Listening on", *listenAddress)
	go listen(*listenAddress, router)

	var (
		hup      = make(chan os.Signal)
		hupReady = make(chan bool)
		term     = make(chan os.Signal)
	)
	signal.Notify(hup, syscall.SIGHUP)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-hupReady
		for {
			select {
			case <-hup:
			case <-webReload:
			}
			reload()
		}
	}()

	// Wait for reload or termination signals.
	close(hupReady) // Unblock SIGHUP handler.

	<-term

	log.Infoln("Received SIGTERM, exiting gracefully...")
}

func extURL(listen, external string) (*url.URL, error) {
	if external == "" {
		hostname, err := os.Hostname()
		if err != nil {
			return nil, err
		}
		_, port, err := net.SplitHostPort(listen)
		if err != nil {
			return nil, err
		}

		external = fmt.Sprintf("http://%s:%s/", hostname, port)
	}

	u, err := url.Parse(external)
	if err != nil {
		return nil, err
	}

	ppref := strings.TrimRight(u.Path, "/")
	if ppref != "" && !strings.HasPrefix(ppref, "/") {
		ppref = "/" + ppref
	}
	u.Path = ppref

	return u, nil
}

func listen(listen string, router *route.Router) {
	if err := http.ListenAndServe(listen, router); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"bufio"
	"os"
	"net"
	"io"
	"./banner"
)

// Command-line flags
var (
	encoding, outputFileName, inputFileName, logFileName, metadataFileName, messageFileName, interfaceName string
	portFlag uint
	inputFile, metadataFile  *os.File
	senders uint
)

// Module configurations
var (
	grabConfig banner.GrabConfig
	outputConfig banner.OutputConfig
)

// Pre-main bind flags to variables
func init() {

	flag.StringVar(&encoding, "encoding", "string", "Encode banner as string|hex|base64")
	flag.StringVar(&outputFileName, "output-file", "-", "Output filename, use - for stdout")
	flag.StringVar(&inputFileName, "input-file", "-", "Input filename, use - for stdin")
	flag.StringVar(&messageFileName, "data", "", "Optional message to send (%s will be replaced with destination IP)")
	flag.StringVar(&metadataFileName, "metadata-file", "-", "File to record banner-grab metadata, use - for stdout")
	flag.StringVar(&logFileName, "log-file", "-", "File to log to, use - for stderr")
	flag.StringVar(&interfaceName, "interface", "", "Network interface to send on")
	flag.UintVar(&portFlag, "port", 80, "Port to grab on")
	flag.IntVar(&grabConfig.Timeout, "timeout", 4, "Set connection timeout in seconds")
	flag.BoolVar(&grabConfig.Tls, "tls", false, "Grab over TLS")
	flag.BoolVar(&grabConfig.Udp, "udp", false, "Grab over UDP")
	flag.UintVar(&senders, "senders", 10, "Number of send coroutines to use")
	flag.Parse()

	// Validate port
	if portFlag > 65535 {
		log.Fatal("Error: Port", portFlag, "out of range")
	}
	grabConfig.Port = uint16(portFlag)

	// Validate timeout
	if grabConfig.Timeout < 0 {
		log.Fatal("Error: Invalid timeout", grabConfig.Timeout)
	}

	// Check UDP
	if grabConfig.Udp {
		log.Print("Warning: UDP is untested")
	}

	// Validate senders
	if senders == 0 {
		log.Fatal("Error: Need at least one sender")
	}

	// Check output type
	if converter, ok := banner.Converters[encoding]; ok {
		outputConfig.Converter = converter
	} else {
		log.Fatal("Error: Invalid encoding ", encoding)
	}

	// Check the network interface
	var err error
	if interfaceName != "" {
		var iface *net.Interface
		if iface, err = net.InterfaceByName(interfaceName); err != nil {
			log.Fatal("Error: Invalid network interface: ", interfaceName)
		}
		var addrs []net.Addr
		if addrs, err = iface.Addrs(); err != nil || len(addrs) == 0 {
			log.Fatal("Error: No addresses for interface ", interfaceName)
		}
		grabConfig.LocalAddr = addrs[0]
	}

	// Open input and output files
	switch inputFileName {
	case "-":
		inputFile = os.Stdin
	default:
		if inputFile, err = os.Open(inputFileName); err != nil {
			log.Fatal(err)
		}
	}

	switch outputFileName {
	case "-":
		outputConfig.OutputFile = os.Stdout
	default:
		if outputConfig.OutputFile, err = os.Create(outputFileName); err != nil {
			log.Fatal(err)
		}
	}

	// Open message file, if applicable
	if messageFileName != "" {
		if messageFile, err := os.Open(messageFileName); err != nil {
			log.Fatal(err)
		} else {
			buf := make([]byte, 1024)
			n, err := messageFile.Read(buf)
			grabConfig.SendMessage = true
			grabConfig.Message = string(buf[0:n])
			if err != nil && err != io.EOF {
				log.Fatal(err)
			}
			messageFile.Close()
		}
	}

	// Open metadata file
	if metadataFileName == "-" {
		metadataFile = os.Stdout
	} else {
		if metadataFile, err = os.Create(metadataFileName); err != nil {
			log.Fatal(err)
		}
	}

	// Open log file, attach to configs
	var logFile *os.File
	if logFileName == "-" {
		logFile = os.Stderr
	} else {
		if logFile, err = os.Create(logFileName); err != nil {
			log.Fatal(err)
		}
	}
	logger := log.New(logFile, "[BANNER-GRAB]", log.LstdFlags)
	outputConfig.ErrorLog = logger
	grabConfig.ErrorLog = logger
}

func ReadInput(addrChan chan net.IP, inputFile *os.File) {
	scanner := bufio.NewScanner(inputFile)
	for scanner.Scan() {
		ipString := scanner.Text()
		ip := net.ParseIP(ipString)
		if ip == nil {
			fmt.Fprintln(os.Stderr, "Invalid IP address: ", ipString)
			continue
		}
		addrChan <- ip
	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "Reading stdin: ", err)
	}
	close(addrChan)
}

func main() {
	addrChan := make(chan net.IP, senders)
	resultChan := make(chan banner.Result, senders)
	summaryChan := make(chan banner.Summary)
	doneChan := make(chan int)

	go banner.WriteOutput(resultChan, summaryChan, &outputConfig)
	for i := uint(0); i < senders; i += 1 {
		go banner.GrabBanner(addrChan, resultChan, doneChan, &grabConfig)
	}
	ReadInput(addrChan, inputFile)

	// Wait for grabbers to finish
	for i := uint(0); i < senders; i += 1 {
		<- doneChan
	}
	close(doneChan)
	close(resultChan)

	if inputFile != os.Stdin {
		inputFile.Close()
	}
	if outputConfig.OutputFile != os.Stdout {
		outputConfig.OutputFile.Close()
	}

	summary := <- summaryChan
	if s, err := banner.SerializeSummary(&summary); err != nil {
		log.Fatal(err)
	} else {
		metadataFile.Write(s)
	}
}

<commit_msg>Add newline after summary<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"bufio"
	"os"
	"net"
	"io"
	"./banner"
)

// Command-line flags
var (
	encoding, outputFileName, inputFileName, logFileName, metadataFileName, messageFileName, interfaceName string
	portFlag uint
	inputFile, metadataFile  *os.File
	senders uint
)

// Module configurations
var (
	grabConfig banner.GrabConfig
	outputConfig banner.OutputConfig
)

// Pre-main bind flags to variables
func init() {

	flag.StringVar(&encoding, "encoding", "string", "Encode banner as string|hex|base64")
	flag.StringVar(&outputFileName, "output-file", "-", "Output filename, use - for stdout")
	flag.StringVar(&inputFileName, "input-file", "-", "Input filename, use - for stdin")
	flag.StringVar(&messageFileName, "data", "", "Optional message to send (%s will be replaced with destination IP)")
	flag.StringVar(&metadataFileName, "metadata-file", "-", "File to record banner-grab metadata, use - for stdout")
	flag.StringVar(&logFileName, "log-file", "-", "File to log to, use - for stderr")
	flag.StringVar(&interfaceName, "interface", "", "Network interface to send on")
	flag.UintVar(&portFlag, "port", 80, "Port to grab on")
	flag.IntVar(&grabConfig.Timeout, "timeout", 4, "Set connection timeout in seconds")
	flag.BoolVar(&grabConfig.Tls, "tls", false, "Grab over TLS")
	flag.BoolVar(&grabConfig.Udp, "udp", false, "Grab over UDP")
	flag.UintVar(&senders, "senders", 10, "Number of send coroutines to use")
	flag.Parse()

	// Validate port
	if portFlag > 65535 {
		log.Fatal("Error: Port", portFlag, "out of range")
	}
	grabConfig.Port = uint16(portFlag)

	// Validate timeout
	if grabConfig.Timeout < 0 {
		log.Fatal("Error: Invalid timeout", grabConfig.Timeout)
	}

	// Check UDP
	if grabConfig.Udp {
		log.Print("Warning: UDP is untested")
	}

	// Validate senders
	if senders == 0 {
		log.Fatal("Error: Need at least one sender")
	}

	// Check output type
	if converter, ok := banner.Converters[encoding]; ok {
		outputConfig.Converter = converter
	} else {
		log.Fatal("Error: Invalid encoding ", encoding)
	}

	// Check the network interface
	var err error
	if interfaceName != "" {
		var iface *net.Interface
		if iface, err = net.InterfaceByName(interfaceName); err != nil {
			log.Fatal("Error: Invalid network interface: ", interfaceName)
		}
		var addrs []net.Addr
		if addrs, err = iface.Addrs(); err != nil || len(addrs) == 0 {
			log.Fatal("Error: No addresses for interface ", interfaceName)
		}
		grabConfig.LocalAddr = addrs[0]
	}

	// Open input and output files
	switch inputFileName {
	case "-":
		inputFile = os.Stdin
	default:
		if inputFile, err = os.Open(inputFileName); err != nil {
			log.Fatal(err)
		}
	}

	switch outputFileName {
	case "-":
		outputConfig.OutputFile = os.Stdout
	default:
		if outputConfig.OutputFile, err = os.Create(outputFileName); err != nil {
			log.Fatal(err)
		}
	}

	// Open message file, if applicable
	if messageFileName != "" {
		if messageFile, err := os.Open(messageFileName); err != nil {
			log.Fatal(err)
		} else {
			buf := make([]byte, 1024)
			n, err := messageFile.Read(buf)
			grabConfig.SendMessage = true
			grabConfig.Message = string(buf[0:n])
			if err != nil && err != io.EOF {
				log.Fatal(err)
			}
			messageFile.Close()
		}
	}

	// Open metadata file
	if metadataFileName == "-" {
		metadataFile = os.Stdout
	} else {
		if metadataFile, err = os.Create(metadataFileName); err != nil {
			log.Fatal(err)
		}
	}

	// Open log file, attach to configs
	var logFile *os.File
	if logFileName == "-" {
		logFile = os.Stderr
	} else {
		if logFile, err = os.Create(logFileName); err != nil {
			log.Fatal(err)
		}
	}
	logger := log.New(logFile, "[BANNER-GRAB]", log.LstdFlags)
	outputConfig.ErrorLog = logger
	grabConfig.ErrorLog = logger
}

func ReadInput(addrChan chan net.IP, inputFile *os.File) {
	scanner := bufio.NewScanner(inputFile)
	for scanner.Scan() {
		ipString := scanner.Text()
		ip := net.ParseIP(ipString)
		if ip == nil {
			fmt.Fprintln(os.Stderr, "Invalid IP address: ", ipString)
			continue
		}
		addrChan <- ip
	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "Reading stdin: ", err)
	}
	close(addrChan)
}

func main() {
	addrChan := make(chan net.IP, senders)
	resultChan := make(chan banner.Result, senders)
	summaryChan := make(chan banner.Summary)
	doneChan := make(chan int)

	go banner.WriteOutput(resultChan, summaryChan, &outputConfig)
	for i := uint(0); i < senders; i += 1 {
		go banner.GrabBanner(addrChan, resultChan, doneChan, &grabConfig)
	}
	ReadInput(addrChan, inputFile)

	// Wait for grabbers to finish
	for i := uint(0); i < senders; i += 1 {
		<- doneChan
	}
	close(doneChan)
	close(resultChan)

	if inputFile != os.Stdin {
		inputFile.Close()
	}
	if outputConfig.OutputFile != os.Stdout {
		outputConfig.OutputFile.Close()
	}

	summary := <- summaryChan
	if s, err := banner.SerializeSummary(&summary); err != nil {
		log.Fatal(err)
	} else {
		metadataFile.Write(s)
		metadataFile.Write([]byte("\n"))
	}
}

<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"go/build"
	"os"
	"path"
	"sort"
	"strings"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/pointer"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/ssautil"
)

var (
	ignoreFlag = flag.String("ignore", "",
		"ignore package path")

	limitFlag = flag.String("limit", "",
		"limit package path")

	focusFlag = flag.String("focus", "",
		"focus package name")

	subFlag = flag.String("sub", "",
		"subgraph by [type, pkg]")

	minlenFlag = flag.Uint("minlen", 2,
		"minlen of edge")

	testFlag = flag.Bool("test", false,
		"Loads test code (*_test.go) for imported packages")
)

func main() {
	flag.Parse()
	if err := doCallgraph(&build.Default, *focusFlag, *limitFlag, *ignoreFlag, *subFlag, *minlenFlag, *testFlag, flag.Args()); err != nil {
		fmt.Fprintf(os.Stderr, "go-callvis: %s\n", err)
		os.Exit(1)
	}
}

func doCallgraph(ctxt *build.Context, focusPkg, limitPath, ignorePath, subgraph string, minlen uint, tests bool, args []string) error {
	conf := loader.Config{Build: &build.Default}

	ignorePaths := []string{}
	for _, p := range strings.Split(ignorePath, ",") {
		if p != "" {
			ignorePaths = append(ignorePaths, p)
		}
	}

	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "missing arguments")
		return nil
	}
	_, err := conf.FromArgs(args, tests)
	if err != nil {
		return err
	}

	iprog, err := conf.Load()
	if err != nil {
		fmt.Print(err) // type error in some package
		return nil
	}

	prog := ssautil.CreateProgram(iprog, 0)
	prog.Build()

	main, err := mainPackage(prog, tests)
	if err != nil {
		return err
	}
	config := &pointer.Config{
		Mains:          []*ssa.Package{main},
		BuildCallGraph: true,
	}
	result, err := pointer.Analyze(config)
	if err != nil {
		return err // internal error in pointer analysis
	}
	result.CallGraph.DeleteSyntheticNodes()

	subType := subgraph == "type"
	subPkg := subgraph == "pkg"

	var edges []string
	edgeMap := make(map[string]struct{})
	callgraph.GraphVisitEdges(result.CallGraph, func(edge *callgraph.Edge) error {
		caller := edge.Caller.Func
		callee := edge.Callee.Func
		if caller.Pkg == nil || callee.Synthetic != "" {
			return nil
		}
		if !(caller.Pkg.Pkg.Name() == focusPkg || callee.Pkg.Pkg.Name() == focusPkg) {
			return nil
		}
		if strings.HasPrefix(caller.Pkg.Pkg.Path(), path.Join(main.Pkg.Path(), "vendor")) ||
			strings.HasPrefix(callee.Pkg.Pkg.Path(), path.Join(main.Pkg.Path(), "vendor")) {
			return nil
		}
		if !(strings.HasPrefix(caller.Pkg.Pkg.Path(), limitPath) &&
			strings.HasPrefix(callee.Pkg.Pkg.Path(), limitPath)) {
			return nil
		}
		for _, p := range ignorePaths {
			if strings.HasPrefix(caller.Pkg.Pkg.Path(), p) ||
				strings.HasPrefix(callee.Pkg.Pkg.Path(), p) {
				return nil
			}
		}

		callerProps := []string{}
		callerSign := caller.Signature
		if caller.Parent() != nil {
			callerSign = caller.Parent().Signature
		}
		callerLabel := fmt.Sprintf("%s\n%s", caller.Pkg.Pkg.Name(), caller.RelString(caller.Pkg.Pkg))
		if caller.Pkg.Pkg.Name() == focusPkg {
			callerProps = append(callerProps, "fillcolor=lightblue")
			callerLabel = fmt.Sprintf("%s", caller.RelString(caller.Pkg.Pkg))
			if subType && callerSign.Recv() != nil {
				callerParts := strings.Split(callerLabel, ".")
				callerLabel = callerParts[len(callerParts)-1]
			}
		} else if subPkg {
			callerLabel = fmt.Sprintf("%s", caller.RelString(caller.Pkg.Pkg))
		}
		callerProps = append(callerProps, fmt.Sprintf("label=%q", callerLabel))
		if caller.Parent() != nil {
			callerProps = append(callerProps, "style=\"dotted,rounded,filled\"")
		} else if caller.Object() != nil && caller.Object().Exported() {
			callerProps = append(callerProps, "style=\"bold,rounded,filled\"")
		}
		callerNode := fmt.Sprintf("%q [%s]", caller, strings.Join(callerProps, " "))
		if subType && caller.Pkg.Pkg.Name() == focusPkg && callerSign.Recv() != nil {
			callerNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				callerSign.Recv().Type(), strings.Split(fmt.Sprint(callerSign.Recv().Type()), ".")[1], callerNode)
		} else if subPkg && caller.Pkg.Pkg.Name() != focusPkg {
			callerNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				caller.Pkg.Pkg.Name(), caller.Pkg.Pkg.Name(), callerNode)
		}

		calleeProps := []string{}
		calleeSign := callee.Signature
		if callee.Parent() != nil {
			calleeSign = callee.Parent().Signature
		}
		calleeLabel := fmt.Sprintf("%s\n%s", callee.Pkg.Pkg.Name(), callee.RelString(callee.Pkg.Pkg))
		if callee.Pkg.Pkg.Name() == focusPkg {
			calleeProps = append(calleeProps, "fillcolor=lightblue")
			calleeLabel = fmt.Sprintf("%s", callee.RelString(callee.Pkg.Pkg))
			if subType && calleeSign.Recv() != nil {
				calleeParts := strings.Split(calleeLabel, ".")
				calleeLabel = calleeParts[len(calleeParts)-1]
			}
		} else if subPkg {
			calleeLabel = fmt.Sprintf("%s", callee.RelString(callee.Pkg.Pkg))
		}
		calleeProps = append(calleeProps, fmt.Sprintf("label=%q", calleeLabel))
		if callee.Parent() != nil {
			calleeProps = append(calleeProps, "style=\"dotted,rounded,filled\"")
		} else if callee.Object() != nil && callee.Object().Exported() {
			calleeProps = append(calleeProps, "style=\"bold,rounded,filled\"")
		}
		calleeNode := fmt.Sprintf("%q [%s]", callee, strings.Join(calleeProps, " "))
		if subType && callee.Pkg.Pkg.Name() == focusPkg && calleeSign.Recv() != nil {
			calleeNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				calleeSign.Recv().Type(), strings.Split(fmt.Sprint(calleeSign.Recv().Type()), ".")[1], calleeNode)
		} else if subPkg && callee.Pkg.Pkg.Name() != focusPkg {
			calleeNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				callee.Pkg.Pkg.Name(), callee.Pkg.Pkg.Name(), calleeNode)
		}

		edgeProps := []string{}
		if edge.Site != nil && edge.Site.Common().StaticCallee() == nil {
			edgeProps = append(edgeProps, "style=dashed")
		}
		switch edge.Site.(type) {
		case *ssa.Go:
			edgeProps = append(edgeProps, "arrowhead=empty")
		case *ssa.Defer:
			edgeProps = append(edgeProps, "arrowhead=odot")
		}
		if callee.Pkg.Pkg.Name() != focusPkg || caller.Pkg.Pkg.Name() != focusPkg {
			edgeProps = append(edgeProps, "color=saddlebrown")
		}
		s := fmt.Sprintf("%s;%s; %q -> %q [%s]",
			callerNode, calleeNode,
			caller, callee, strings.Join(edgeProps, " "))
		if _, ok := edgeMap[s]; !ok {
			edges = append(edges, s)
			edgeMap[s] = struct{}{}
		}

		return nil
	})

	sort.Strings(edges)

	fmt.Printf(`digraph G {
        label="%s";
        labelloc=t;
        bgcolor=aliceblue;
        rankdir=LR;
        fontsize=22;
        fontname="Ubuntu";
        edge [minlen=%d];
        node [shape=box style="rounded,filled" fillcolor=wheat fontname="Ubuntu"];
`, focusPkg, minlen)
	for _, edge := range edges {
		fmt.Println("\t", edge)
	}
	fmt.Println("}")

	//fmt.Println(len(edges), "edges")
	return nil
}

func mainPackage(prog *ssa.Program, tests bool) (*ssa.Package, error) {
	pkgs := prog.AllPackages()
	/*if tests {
		if len(pkgs) > 0 {
			if main := prog.CreateTestMainPackage(pkgs...); main != nil {
				return main, nil
			}
		}
		return nil, fmt.Errorf("no tests")
	}*/
	for _, pkg := range pkgs {
		if pkg.Pkg.Name() == "main" {
			if pkg.Func("main") == nil {
				return nil, fmt.Errorf("no func main() in main package")
			}
			return pkg, nil
		}
	}
	return nil, fmt.Errorf("no main package")
}
<commit_msg>Fixes #2<commit_after>package main

import (
	"flag"
	"fmt"
	"go/build"
	"os"
	"path"
	"sort"
	"strings"

	"golang.org/x/tools/go/callgraph"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/pointer"
	"golang.org/x/tools/go/ssa"
	"golang.org/x/tools/go/ssa/ssautil"
)

var (
	ignoreFlag = flag.String("ignore", "",
		"ignore package path")

	limitFlag = flag.String("limit", "",
		"limit package path")

	focusFlag = flag.String("focus", "",
		"focus package name")

	subFlag = flag.String("sub", "",
		"subgraph by [type, pkg]")

	minlenFlag = flag.Uint("minlen", 2,
		"minlen of edge")

	testFlag = flag.Bool("test", false,
		"Loads test code (*_test.go) for imported packages")
)

func main() {
	flag.Parse()
	if err := doCallgraph(&build.Default, *focusFlag, *limitFlag, *ignoreFlag, *subFlag, *minlenFlag, *testFlag, flag.Args()); err != nil {
		fmt.Fprintf(os.Stderr, "go-callvis: %s\n", err)
		os.Exit(1)
	}
}

func doCallgraph(ctxt *build.Context, focusPkg, limitPath, ignorePath, subgraph string, minlen uint, tests bool, args []string) error {
	conf := loader.Config{Build: &build.Default}

	ignorePaths := []string{}
	for _, p := range strings.Split(ignorePath, ",") {
		if p != "" {
			ignorePaths = append(ignorePaths, p)
		}
	}

	if len(args) == 0 {
		fmt.Fprintln(os.Stderr, "missing arguments")
		return nil
	}
	_, err := conf.FromArgs(args, tests)
	if err != nil {
		return err
	}

	iprog, err := conf.Load()
	if err != nil {
		fmt.Print(err) // type error in some package
		return nil
	}

	prog := ssautil.CreateProgram(iprog, 0)
	prog.Build()

	main, err := mainPackage(prog, tests)
	if err != nil {
		return err
	}
	config := &pointer.Config{
		Mains:          []*ssa.Package{main},
		BuildCallGraph: true,
	}
	result, err := pointer.Analyze(config)
	if err != nil {
		return err // internal error in pointer analysis
	}
	result.CallGraph.DeleteSyntheticNodes()

	subType := subgraph == "type"
	subPkg := subgraph == "pkg"

	var edges []string
	edgeMap := make(map[string]struct{})
	callgraph.GraphVisitEdges(result.CallGraph, func(edge *callgraph.Edge) error {
		caller := edge.Caller.Func
		callee := edge.Callee.Func
		if caller.Pkg == nil || callee.Synthetic != "" {
			return nil
		}
		if !(caller.Pkg.Pkg.Name() == focusPkg || callee.Pkg.Pkg.Name() == focusPkg) {
			return nil
		}
		if strings.HasPrefix(caller.Pkg.Pkg.Path(), path.Join(main.Pkg.Path(), "vendor")) ||
			strings.HasPrefix(callee.Pkg.Pkg.Path(), path.Join(main.Pkg.Path(), "vendor")) {
			return nil
		}
		if !(strings.HasPrefix(caller.Pkg.Pkg.Path(), limitPath) &&
			strings.HasPrefix(callee.Pkg.Pkg.Path(), limitPath)) {
			return nil
		}
		for _, p := range ignorePaths {
			if strings.HasPrefix(caller.Pkg.Pkg.Path(), p) ||
				strings.HasPrefix(callee.Pkg.Pkg.Path(), p) {
				return nil
			}
		}

		callerProps := []string{}
		callerSign := caller.Signature
		if caller.Parent() != nil {
			callerSign = caller.Parent().Signature
		}
		callerLabel := fmt.Sprintf("%s\n%s", caller.Pkg.Pkg.Name(), caller.RelString(caller.Pkg.Pkg))
		if caller.Pkg.Pkg.Name() == focusPkg {
			callerProps = append(callerProps, "fillcolor=lightblue")
			callerLabel = fmt.Sprintf("%s", caller.RelString(caller.Pkg.Pkg))
			if subType && callerSign.Recv() != nil {
				callerParts := strings.Split(callerLabel, ".")
				callerLabel = callerParts[len(callerParts)-1]
			}
		} else if subPkg {
			callerLabel = fmt.Sprintf("%s", caller.RelString(caller.Pkg.Pkg))
		}
		callerProps = append(callerProps, fmt.Sprintf("label=%q", callerLabel))
		if caller.Parent() != nil {
			callerProps = append(callerProps, "style=\"dotted,rounded,filled\"")
		} else if caller.Object() != nil && caller.Object().Exported() {
			callerProps = append(callerProps, "style=\"bold,rounded,filled\"")
		}
		callerNode := fmt.Sprintf("%q [%s]", caller, strings.Join(callerProps, " "))
		if subType && caller.Pkg.Pkg.Name() == focusPkg && callerSign.Recv() != nil {
			parts := strings.Split(fmt.Sprint(callerSign.Recv().Type()), ".")
			clusterLabel := parts[len(parts)-1]
			callerNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				callerSign.Recv().Type(), clusterLabel, callerNode)
		} else if subPkg && caller.Pkg.Pkg.Name() != focusPkg {
			callerNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				caller.Pkg.Pkg.Name(), caller.Pkg.Pkg.Name(), callerNode)
		}

		calleeProps := []string{}
		calleeSign := callee.Signature
		if callee.Parent() != nil {
			calleeSign = callee.Parent().Signature
		}
		calleeLabel := fmt.Sprintf("%s\n%s", callee.Pkg.Pkg.Name(), callee.RelString(callee.Pkg.Pkg))
		if callee.Pkg.Pkg.Name() == focusPkg {
			calleeProps = append(calleeProps, "fillcolor=lightblue")
			calleeLabel = fmt.Sprintf("%s", callee.RelString(callee.Pkg.Pkg))
			if subType && calleeSign.Recv() != nil {
				calleeParts := strings.Split(calleeLabel, ".")
				calleeLabel = calleeParts[len(calleeParts)-1]
			}
		} else if subPkg {
			calleeLabel = fmt.Sprintf("%s", callee.RelString(callee.Pkg.Pkg))
		}
		calleeProps = append(calleeProps, fmt.Sprintf("label=%q", calleeLabel))
		if callee.Parent() != nil {
			calleeProps = append(calleeProps, "style=\"dotted,rounded,filled\"")
		} else if callee.Object() != nil && callee.Object().Exported() {
			calleeProps = append(calleeProps, "style=\"bold,rounded,filled\"")
		}
		calleeNode := fmt.Sprintf("%q [%s]", callee, strings.Join(calleeProps, " "))
		if subType && callee.Pkg.Pkg.Name() == focusPkg && calleeSign.Recv() != nil {
			parts := strings.Split(fmt.Sprint(calleeSign.Recv().Type()), ".")
			clusterLabel := parts[len(parts)-1]
			calleeNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				calleeSign.Recv().Type(), clusterLabel, calleeNode)
		} else if subPkg && callee.Pkg.Pkg.Name() != focusPkg {
			calleeNode = fmt.Sprintf("subgraph \"cluster_%s\" { penwidth=0.5; fontsize=18; label=\"%s\"; style=filled; fillcolor=snow; %s; }",
				callee.Pkg.Pkg.Name(), callee.Pkg.Pkg.Name(), calleeNode)
		}

		edgeProps := []string{}
		if edge.Site != nil && edge.Site.Common().StaticCallee() == nil {
			edgeProps = append(edgeProps, "style=dashed")
		}
		switch edge.Site.(type) {
		case *ssa.Go:
			edgeProps = append(edgeProps, "arrowhead=empty")
		case *ssa.Defer:
			edgeProps = append(edgeProps, "arrowhead=odot")
		}
		if callee.Pkg.Pkg.Name() != focusPkg || caller.Pkg.Pkg.Name() != focusPkg {
			edgeProps = append(edgeProps, "color=saddlebrown")
		}
		s := fmt.Sprintf("%s;%s; %q -> %q [%s]",
			callerNode, calleeNode,
			caller, callee, strings.Join(edgeProps, " "))
		if _, ok := edgeMap[s]; !ok {
			edges = append(edges, s)
			edgeMap[s] = struct{}{}
		}

		return nil
	})

	sort.Strings(edges)

	fmt.Printf(`digraph G {
        label="%s";
        labelloc=t;
        bgcolor=aliceblue;
        rankdir=LR;
        fontsize=22;
        fontname="Ubuntu";
        edge [minlen=%d];
        node [shape=box style="rounded,filled" fillcolor=wheat fontname="Ubuntu"];
`, focusPkg, minlen)
	for _, edge := range edges {
		fmt.Println("\t", edge)
	}
	fmt.Println("}")

	//fmt.Println(len(edges), "edges")
	return nil
}

func mainPackage(prog *ssa.Program, tests bool) (*ssa.Package, error) {
	pkgs := prog.AllPackages()
	/*if tests {
		if len(pkgs) > 0 {
			if main := prog.CreateTestMainPackage(pkgs...); main != nil {
				return main, nil
			}
		}
		return nil, fmt.Errorf("no tests")
	}*/
	for _, pkg := range pkgs {
		if pkg.Pkg.Name() == "main" {
			if pkg.Func("main") == nil {
				return nil, fmt.Errorf("no func main() in main package")
			}
			return pkg, nil
		}
	}
	return nil, fmt.Errorf("no main package")
}
<|endoftext|>