<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"
)

const (
	HTTPS_TEMPLATE = `` +
		`  DNS Lookup   TCP Connection   SSL Handshake   Server Processing   Content Transfer` +
		`[   {a0000}  |     {a0001}    |    {a0002}    |      {a0003}      |      {a0004}     ]` +
		`             |                |               |                   |                  |` +
		`    namelookup:{b0000}        |               |                   |                  |` +
		`                        connect:{b0001}       |                   |                  |` +
		`                                    pretransfer:{b0002}           |                  |` +
		`                                                      starttransfer:{b0003}          |` +
		`                                                                                total:{b0004}`

	HTTP_TEMPLATE = `` +
		`   DNS Lookup   TCP Connection   Server Processing   Content Transfer` + "\n" +
		`[ %s  |     %s  |        %s  |       %s  ]` + "\n" +
		`             |                |                   |                  |` + "\n" +
		`    namelookup:%s      |                   |                  |` + "\n" +
		`                        connect:%s         |                  |` + "\n" +
		`                                      starttransfer:%s        |` + "\n" +
		`                                                                 total:%s` + "\n"
)

const ISATTY = true // TODO(dfc) make this respect the actual pty state

func makeColor(code int) func(string) string {
	if !ISATTY {
		return func(s string) string { return s }
	}
	return func(s string) string {
		return fmt.Sprintf("\x1b[%dm%s\x1b[0m", code, s)
	}
}

var (
	red       = makeColor(31)
	green     = makeColor(32)
	yellow    = makeColor(33)
	blue      = makeColor(34)
	magenta   = makeColor(35)
	cyan      = makeColor(36)
	bold      = makeColor(1)
	underline = makeColor(4)

	grayscale = func(code int) func(string) string {
		if !ISATTY {
			return func(s string) string { return s }
		}
		return func(s string) string {
			return fmt.Sprintf("\x1b[;38;5;%dm%s\x1b[0m", code+232, s)
		}
	}
)

func main() {
	flag.Parse()

	args := flag.Args()
	if len(args) != 1 {
		log.Fatalf("usage: %s URL", os.Args[0])
	}

	url, err := url.Parse(args[0])
	if err != nil {
		log.Fatalf("could not parse url %q: %v", args[0], err)
	}

	scheme := url.Scheme
	hostport := url.Host
	host, port := func() (string, int) {
		switch scheme {
		case "https":
			return hostport, 443
		case "http":
			return hostport, 80
		default:
			log.Fatalf("unsupported url scheme %q", scheme)
			return "", 0 // not reached
		}
	}()

	t0 := time.Now() // before dns resolution
	raddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		log.Fatalf("unable to resolve host: %v", err)
	}

	t1 := time.Now() // after dns resolution, before connect
	conn, err := net.DialTCP("tcp", nil, raddr)
	if err != nil {
		log.Fatalf("unable to connect to host %vv %v", raddr, err)
	}

	t2 := time.Now() // after connect, before request
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		log.Fatalf("unable to create request: %v", err)
	}

	if err := req.Write(conn); err != nil {
		log.Fatalf("failed to write request: %v", err)
	}

	t3 := time.Now() // after request, before read response
	resp, err := http.ReadResponse(bufio.NewReader(conn), req)
	if err != nil {
		log.Fatalf("failed to read response: %v", err)
	}

	t4 := time.Now() // after read request, before read body
	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
		log.Fatalf("failed to read response body: %v", err)
	}

	t5 := time.Now() // after read body
	resp.Body.Close()

	_ = t3

	// print status line and headers
	fmt.Printf("\n%s%s%s\n", green("HTTP"), grayscale(14)("/"), cyan(fmt.Sprintf("%d.%d %s", resp.ProtoMajor, resp.ProtoMinor, resp.Status)))

	for k, v := range resp.Header {
		fmt.Println(grayscale(14)(k+":"), cyan(strings.Join(v, ",")))
	}

	fmt.Println("\nBody discarded\n")

	fmta := func(d time.Duration) string {
		return cyan(fmt.Sprintf("%7dms", int(d/time.Millisecond)))
	}

	fmtb := func(d time.Duration) string {
		return cyan(fmt.Sprintf("%7dms", int(d/time.Millisecond)))
	}

	switch scheme {
	case "https":
	case "http":
		fmt.Printf(HTTP_TEMPLATE,
			fmta(t1.Sub(t0)), // dns lookup
			fmta(t2.Sub(t1)), // tcp connection
			fmta(t4.Sub(t2)), // server processing
			fmta(t5.Sub(t4)), // content transfer
			fmtb(t1.Sub(t0)), // namelookup
			fmtb(t2.Sub(t0)), // connect
			fmtb(t4.Sub(t0)), // starttransfer
			fmtb(t5.Sub(t0)), // total
		)

	}
}
<commit_msg>more colorization<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"
)

const (
	HTTPS_TEMPLATE = `` +
		`  DNS Lookup   TCP Connection   SSL Handshake   Server Processing   Content Transfer` +
		`[   {a0000}  |     {a0001}    |    {a0002}    |      {a0003}      |      {a0004}     ]` +
		`             |                |               |                   |                  |` +
		`    namelookup:{b0000}        |               |                   |                  |` +
		`                        connect:{b0001}       |                   |                  |` +
		`                                    pretransfer:{b0002}           |                  |` +
		`                                                      starttransfer:{b0003}          |` +
		`                                                                                total:{b0004}`

	HTTP_TEMPLATE = `` +
		`   DNS Lookup   TCP Connection   Server Processing   Content Transfer` + "\n" +
		`[ %s  |     %s  |        %s  |       %s  ]` + "\n" +
		`             |                |                   |                  |` + "\n" +
		`    namelookup:%s      |                   |                  |` + "\n" +
		`                        connect:%s         |                  |` + "\n" +
		`                                      starttransfer:%s        |` + "\n" +
		`                                                                 total:%s` + "\n"
)

const ISATTY = true // TODO(dfc) make this respect the actual pty state

func makeColor(code int) func(string) string {
	if !ISATTY {
		return func(s string) string { return s }
	}
	return func(s string) string {
		return fmt.Sprintf("\x1b[%dm%s\x1b[0m", code, s)
	}
}

var (
	red       = makeColor(31)
	green     = makeColor(32)
	yellow    = makeColor(33)
	blue      = makeColor(34)
	magenta   = makeColor(35)
	cyan      = makeColor(36)
	bold      = makeColor(1)
	underline = makeColor(4)

	grayscale = func(code int) func(string) string {
		if !ISATTY {
			return func(s string) string { return s }
		}
		return func(s string) string {
			return fmt.Sprintf("\x1b[;38;5;%dm%s\x1b[0m", code+232, s)
		}
	}
)

func main() {
	flag.Parse()

	args := flag.Args()
	if len(args) != 1 {
		log.Fatalf("usage: %s URL", os.Args[0])
	}

	url, err := url.Parse(args[0])
	if err != nil {
		log.Fatalf("could not parse url %q: %v", args[0], err)
	}

	scheme := url.Scheme
	hostport := url.Host
	host, port := func() (string, int) {
		switch scheme {
		case "https":
			return hostport, 443
		case "http":
			return hostport, 80
		default:
			log.Fatalf("unsupported url scheme %q", scheme)
			return "", 0 // not reached
		}
	}()

	t0 := time.Now() // before dns resolution
	raddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%d", host, port))
	if err != nil {
		log.Fatalf("unable to resolve host: %v", err)
	}

	t1 := time.Now() // after dns resolution, before connect
	conn, err := net.DialTCP("tcp", nil, raddr)
	if err != nil {
		log.Fatalf("unable to connect to host %vv %v", raddr, err)
	}

	t2 := time.Now() // after connect, before request
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		log.Fatalf("unable to create request: %v", err)
	}

	if err := req.Write(conn); err != nil {
		log.Fatalf("failed to write request: %v", err)
	}

	t3 := time.Now() // after request, before read response
	resp, err := http.ReadResponse(bufio.NewReader(conn), req)
	if err != nil {
		log.Fatalf("failed to read response: %v", err)
	}

	t4 := time.Now() // after read request, before read body
	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
		log.Fatalf("failed to read response body: %v", err)
	}

	t5 := time.Now() // after read body
	resp.Body.Close()

	_ = t3

	// print status line and headers
	fmt.Printf("\n%s%s%s\n", green("HTTP"), grayscale(14)("/"), cyan(fmt.Sprintf("%d.%d %s", resp.ProtoMajor, resp.ProtoMinor, resp.Status)))

	for k, v := range resp.Header {
		fmt.Println(grayscale(14)(k+":"), cyan(strings.Join(v, ",")))
	}

	fmt.Println("\nBody discarded\n")

	fmta := func(d time.Duration) string {
		return cyan(fmt.Sprintf("%7dms", int(d/time.Millisecond)))
	}

	fmtb := func(d time.Duration) string {
		return cyan(fmt.Sprintf("%7sms", strconv.Itoa(int(d/time.Millisecond)))) // TODO(dfc) left justify, %7_s doesn't seem to work.
	}

	colorize := func(s string) string {
		v := strings.Split(s, "\n")
		v[0] = grayscale(16)(v[0])
		return strings.Join(v, "\n")
	}

	switch scheme {
	case "https":
		// TODO(dfc) handle HTTPS
	case "http":
		fmt.Printf(colorize(HTTP_TEMPLATE),
			fmta(t1.Sub(t0)), // dns lookup
			fmta(t2.Sub(t1)), // tcp connection
			fmta(t4.Sub(t2)), // server processing
			fmta(t5.Sub(t4)), // content transfer
			fmtb(t1.Sub(t0)), // namelookup
			fmtb(t2.Sub(t0)), // connect
			fmtb(t4.Sub(t0)), // starttransfer
			fmtb(t5.Sub(t0)), // total
		)

	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/stianeikeland/go-rpio"
	"net/http"
	"os"
	"time"
)

const Version = "4.3.0"

var options struct {
	http            string
	pinNumber       int
	statusPinNumber int
	sleepTimeout    int
	cert            string
	key             string
	version         bool
}

var sharedSecret = os.Getenv("GARAGE_SECRET")

func logHandler(event string) {
	fmt.Fprintln(os.Stdout, event, "-", time.Now())
}

func verifySignature(signedText []byte, signature []byte) bool {
	mac := hmac.New(sha512.New, []byte(sharedSecret))
	mac.Write(signedText)
	expectedMAC := []byte(hex.EncodeToString(mac.Sum(nil)))
	return hmac.Equal(signature, expectedMAC)
}

func verifyTime(timestamp int64) (int64, error) {
	timeSinceRequest := time.Now().Unix() - timestamp
	if timeSinceRequest > 10 {
		return -1, errors.New("Timestamp is too far in the past")
	}
	return timestamp, nil
}

func toggleSwitch(pinNumber int, sleepTimeout int) (err error) {
	err = rpio.Open()
	if err != nil {
		return err
	}
	pin := rpio.Pin(pinNumber)
	pin.Output()

	pin.Low()
	rpio.Close()

	snooze := time.Duration(sleepTimeout) * time.Millisecond
	time.Sleep(snooze)

	err = rpio.Open()
	if err != nil {
		return err
	}
	pin.High()
	rpio.Close()

	return nil
}

type ClientRequest struct {
	Timestamp int64 `json:"timestamp"`
}

func CheckStatus(pinNumber int) (state string, err error) {
	err = rpio.Open()
	if err != nil {
		return
	}
	defer rpio.Close()

	pin := rpio.Pin(pinNumber)

	status := "open"
	if pin.Read() == 0 {
		status = "closed"
	}

	return status, err
}

func DoorStatus(w http.ResponseWriter, r *http.Request) {
	var jsonResp struct {
		Text string `json:"door_status"`
	}

	status, err := CheckStatus(options.statusPinNumber)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		w.WriteHeader(422)
		jsonResp.Text = fmt.Sprintf("%s", err)
	}

	jsonResp.Text = status
	message, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	w.Write(message)
}

func GetVersion(w http.ResponseWriter, r *http.Request) {
	logHandler("Version")
	var jsonResp struct {
		Text string `json:"version"`
	}
	jsonResp.Text = Version
	message, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	w.Write(message)
}

func Relay(w http.ResponseWriter, r *http.Request) {
	header := r.Header.Get("signature")
	signature, err := base64.URLEncoding.DecodeString(header)
	if err != nil {
		panic(err)
	}

	buf := new(bytes.Buffer)
	buf.ReadFrom(r.Body)
	body := buf.Bytes()

	var jsonResp struct {
		Text string `json:"status"`
	}
	jsonResp.Text = "signal received"

	verified := verifySignature(body, signature)
	if verified {
		// Verify time
		var clientRequest ClientRequest
		err := json.Unmarshal([]byte(body), &clientRequest)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}

		_, err = verifyTime(clientRequest.Timestamp)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}

		// Toggle switch
		logHandler("TOGGLE DOOR")
		err = toggleSwitch(options.pinNumber, options.sleepTimeout)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			w.WriteHeader(422)
			jsonResp.Text = fmt.Sprintf("%s", err)
		}
	} else {
		w.WriteHeader(401)
		logHandler(fmt.Sprintf("Invalid signature: %s", signature))
		jsonResp.Text = fmt.Sprintf("%s", "Invalid signature")
	}

	message, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	w.Write(message)
}

func main() {
	if os.Args[1] == "update" {
		CheckForUpdates()
		os.Exit(0)
	}

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage:  %s [options]\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.IntVar(&options.pinNumber, "pin", 25, "GPIO pin of relay")
	flag.IntVar(&options.statusPinNumber, "status-pin", 10, "GPIO pin of reed switch")
	flag.IntVar(&options.sleepTimeout, "sleep", 100, "Time in milliseconds to keep switch closed")
	flag.StringVar(&options.http, "http", "", "HTTP listen address (e.g. 127.0.0.1:8225)")
	flag.StringVar(&options.cert, "cert", "", "SSL certificate path (e.g. /ssl/example.com.cert)")
	flag.StringVar(&options.key, "key", "", "SSL certificate key (e.g. /ssl/example.com.key)")
	flag.BoolVar(&options.version, "version", false, "print version and exit")
	flag.Parse()

	if options.version {
		fmt.Printf("garage-server v%v\n", Version)
		os.Exit(0)
	}

	if sharedSecret == "" {
		println("You did not set GARAGE_SECRET env var")
		os.Exit(1)
	}

	serveAddress := "127.0.0.1:8225"
	if options.http != "" {
		serveAddress = options.http
	}

	http.HandleFunc("/", Relay)
	http.HandleFunc("/status", DoorStatus)
	http.HandleFunc("/version", GetVersion)

	fmt.Fprintln(os.Stderr, "=> Booting Garage Server ", Version)
	fmt.Fprintln(os.Stderr, "=> Run `garage-server -h` for more startup options")
	fmt.Fprintln(os.Stderr, "=> Ctrl-C to shutdown server")

	var err error
	if options.key != "" && options.cert != "" {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("* Listening on https://%s", serveAddress))
		err = http.ListenAndServeTLS(serveAddress, options.cert, options.key, nil)
	} else {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("* Listening on http://%s", serveAddress))
		err = http.ListenAndServe(serveAddress, nil)
	}

	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}
<commit_msg>Bump to version 4.4.0<commit_after>package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"github.com/stianeikeland/go-rpio"
	"net/http"
	"os"
	"time"
)

const Version = "4.4.0"

var options struct {
	http            string
	pinNumber       int
	statusPinNumber int
	sleepTimeout    int
	cert            string
	key             string
	version         bool
}

var sharedSecret = os.Getenv("GARAGE_SECRET")

func logHandler(event string) {
	fmt.Fprintln(os.Stdout, event, "-", time.Now())
}

func verifySignature(signedText []byte, signature []byte) bool {
	mac := hmac.New(sha512.New, []byte(sharedSecret))
	mac.Write(signedText)
	expectedMAC := []byte(hex.EncodeToString(mac.Sum(nil)))
	return hmac.Equal(signature, expectedMAC)
}

func verifyTime(timestamp int64) (int64, error) {
	timeSinceRequest := time.Now().Unix() - timestamp
	if timeSinceRequest > 10 {
		return -1, errors.New("Timestamp is too far in the past")
	}
	return timestamp, nil
}

func toggleSwitch(pinNumber int, sleepTimeout int) (err error) {
	err = rpio.Open()
	if err != nil {
		return err
	}
	pin := rpio.Pin(pinNumber)
	pin.Output()

	pin.Low()
	rpio.Close()

	snooze := time.Duration(sleepTimeout) * time.Millisecond
	time.Sleep(snooze)

	err = rpio.Open()
	if err != nil {
		return err
	}
	pin.High()
	rpio.Close()

	return nil
}

type ClientRequest struct {
	Timestamp int64 `json:"timestamp"`
}

func CheckStatus(pinNumber int) (state string, err error) {
	err = rpio.Open()
	if err != nil {
		return
	}
	defer rpio.Close()

	pin := rpio.Pin(pinNumber)

	status := "open"
	if pin.Read() == 0 {
		status = "closed"
	}

	return status, err
}

func DoorStatus(w http.ResponseWriter, r *http.Request) {
	var jsonResp struct {
		Text string `json:"door_status"`
	}

	status, err := CheckStatus(options.statusPinNumber)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		w.WriteHeader(422)
		jsonResp.Text = fmt.Sprintf("%s", err)
	}

	jsonResp.Text = status
	message, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	w.Write(message)
}

func GetVersion(w http.ResponseWriter, r *http.Request) {
	logHandler("Version")
	var jsonResp struct {
		Text string `json:"version"`
	}
	jsonResp.Text = Version
	message, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	w.Write(message)
}

func Relay(w http.ResponseWriter, r *http.Request) {
	header := r.Header.Get("signature")
	signature, err := base64.URLEncoding.DecodeString(header)
	if err != nil {
		panic(err)
	}

	buf := new(bytes.Buffer)
	buf.ReadFrom(r.Body)
	body := buf.Bytes()

	var jsonResp struct {
		Text string `json:"status"`
	}
	jsonResp.Text = "signal received"

	verified := verifySignature(body, signature)
	if verified {
		// Verify time
		var clientRequest ClientRequest
		err := json.Unmarshal([]byte(body), &clientRequest)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}

		_, err = verifyTime(clientRequest.Timestamp)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			return
		}

		// Toggle switch
		logHandler("TOGGLE DOOR")
		err = toggleSwitch(options.pinNumber, options.sleepTimeout)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			w.WriteHeader(422)
			jsonResp.Text = fmt.Sprintf("%s", err)
		}
	} else {
		w.WriteHeader(401)
		logHandler(fmt.Sprintf("Invalid signature: %s", signature))
		jsonResp.Text = fmt.Sprintf("%s", "Invalid signature")
	}

	message, err := json.Marshal(jsonResp)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
	}
	w.Write(message)
}

func main() {
	if os.Args[1] == "update" {
		CheckForUpdates()
		os.Exit(0)
	}

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage:  %s [options]\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.IntVar(&options.pinNumber, "pin", 25, "GPIO pin of relay")
	flag.IntVar(&options.statusPinNumber, "status-pin", 10, "GPIO pin of reed switch")
	flag.IntVar(&options.sleepTimeout, "sleep", 100, "Time in milliseconds to keep switch closed")
	flag.StringVar(&options.http, "http", "", "HTTP listen address (e.g. 127.0.0.1:8225)")
	flag.StringVar(&options.cert, "cert", "", "SSL certificate path (e.g. /ssl/example.com.cert)")
	flag.StringVar(&options.key, "key", "", "SSL certificate key (e.g. /ssl/example.com.key)")
	flag.BoolVar(&options.version, "version", false, "print version and exit")
	flag.Parse()

	if options.version {
		fmt.Printf("garage-server v%v\n", Version)
		os.Exit(0)
	}

	if sharedSecret == "" {
		println("You did not set GARAGE_SECRET env var")
		os.Exit(1)
	}

	serveAddress := "127.0.0.1:8225"
	if options.http != "" {
		serveAddress = options.http
	}

	http.HandleFunc("/", Relay)
	http.HandleFunc("/status", DoorStatus)
	http.HandleFunc("/version", GetVersion)

	fmt.Fprintln(os.Stderr, "=> Booting Garage Server ", Version)
	fmt.Fprintln(os.Stderr, "=> Run `garage-server -h` for more startup options")
	fmt.Fprintln(os.Stderr, "=> Ctrl-C to shutdown server")

	var err error
	if options.key != "" && options.cert != "" {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("* Listening on https://%s", serveAddress))
		err = http.ListenAndServeTLS(serveAddress, options.cert, options.key, nil)
	} else {
		fmt.Fprintln(os.Stderr, fmt.Sprintf("* Listening on http://%s", serveAddress))
		err = http.ListenAndServe(serveAddress, nil)
	}

	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/fatih/color"
	"github.com/lucagrulla/cw/cloudwatch"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	timeFormat = "2006-01-02T15:04:05"
	version    = "2.1.2"

	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()
	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.").HintAction(groupsCompletion).Required().String()

	tailCommand     = kp.Command("tail", "Tail a log group.")
	follow          = tailCommand.Flag("follow", "Don't stop when the end of stream is reached, but rather wait for additional data to be appended.").Short('f').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()
	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()
	logGroupName    = tailCommand.Arg("group", "The log group name.").Required().HintAction(groupsCompletion).String()
	logStreamName   = tailCommand.Arg("stream", "The log stream name. Use \\* for tail all the group streams.").Default("*").HintAction(streamsCompletion).String()
	startTime       = tailCommand.Arg("start", "The UTC start time. Passed as either date/time or human-friendly format. The human-friendly format accepts the number of hours and minutes prior to the present. Denote hours with 'h' and minutes with 'm' i.e. 80m, 4h30m. If time is passed (format: hh[:mm]) it is expanded to today at the given time. Full available date/time format: 2017-02-27[T09:00[:00]].").Default(time.Now().UTC().Add(-30 * time.Second).Format(timeFormat)).String()
	endTime         = tailCommand.Arg("end", "The UTC start time. Passed as either date/time or human-friendly format. The human-friendly format accepts the number of hours and minutes prior to the present. Denote hours with 'h' and minutes with 'm' i.e. 80m, 4h30m. If time is passed (format: hh[:mm]) it is expanded to today at the given time. Full available date/time format: 2017-02-27[T09:00[:00]].").String()
	local           = tailCommand.Flag("local", "Treat date and time in Local zone.").Short('l').Default("false").Bool()
)

func groupsCompletion() []string {
	var groups []string
	kingpin.MustParse(kp.Parse(os.Args[1:]))

	for msg := range cloudwatch.New(awsProfile, awsRegion, debug).LsGroups() {
		groups = append(groups, *msg)
	}
	return groups
}

func streamsCompletion() []string {
	var streams []string
	kingpin.MustParse(kp.Parse(os.Args[1:]))

	for msg := range cloudwatch.New(awsProfile, awsRegion, debug).LsStreams(logGroupName, nil) {
		streams = append(streams, *msg)
	}
	return streams
}

func timestampToTime(timeStamp *string) time.Time {
	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
	} 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
	} 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
	} 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)
	} 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)
	} else if regexp.MustCompile(`^\d{1,}h$|^\d{1,}m$|^\d{1,}h\d{1,}m$`).MatchString(*timeStamp) {
		d, _ := time.ParseDuration(*timeStamp)

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

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

func fetchLatestVersion() chan string {
	latestVersionChannel := make(chan string, 1)
	go func() {
		r, _ := http.Get("https://github.com/lucagrulla/cw/releases/latest")

		finalURL := r.Request.URL.String()
		tokens := strings.Split(finalURL, "/")
		latestVersionChannel <- tokens[len(tokens)-1]
	}()
	return latestVersionChannel
}

func newVersionMsg(currentVersion string, latestVersionChannel chan string) {
	latestVersion := <-latestVersionChannel
	if latestVersion != fmt.Sprintf("v%s", currentVersion) {
		fmt.Println("")
		fmt.Println("")
		if *noColor {
			msg := fmt.Sprintf("%s - %s -> %s", "A new version of cw is available!", currentVersion, latestVersion)
			fmt.Println(msg)
		} else {
			msg := fmt.Sprintf("%s - %s -> %s", color.GreenString("A new version of cw is available!"), color.YellowString(currentVersion), color.GreenString(latestVersion))
			fmt.Println(msg)
		}
	}
}

func versionCheckOnSigterm(versionMsgFunc func()) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	<-c
	versionMsgFunc()
	os.Exit(0)
}

func main() {
	kp.Version(version).Author("Luca Grulla")

	versionMsgFunc := func() { newVersionMsg(version, fetchLatestVersion()) }
	defer versionMsgFunc()
	go versionCheckOnSigterm(versionMsgFunc)

	cmd := kingpin.MustParse(kp.Parse(os.Args[1:]))
	c := cloudwatch.New(awsProfile, awsRegion, debug)
	switch cmd {
	case "ls groups":

		for msg := range c.LsGroups() {
			fmt.Println(*msg)
		}
	case "ls streams":
		for msg := range c.LsStreams(lsLogGroupName, nil) {
			fmt.Println(*msg)
		}
	case "tail":
		st := timestampToTime(startTime)
		var et time.Time
		if *endTime != "" {
			et = timestampToTime(endTime)
		}
		for event := range c.Tail(logGroupName, logStreamName, follow, &st, &et, grep, grepv) {
			msg := *event.Message
			if *printEventID {
				if *noColor {
					msg = fmt.Sprintf("%s - %s", *event.EventId, msg)
				} else {
					msg = fmt.Sprintf("%s - %s", color.YellowString(*event.EventId), msg)
				}
			}
			if *printStreamName {
				if *noColor {
					msg = fmt.Sprintf("%s - %s", *event.LogStreamName, msg)
				} else {
					msg = fmt.Sprintf("%s - %s", color.BlueString(*event.LogStreamName), msg)
				}
			}
			if *printTimestamp {
				eventTimestamp := *event.Timestamp / 1000
				ts := time.Unix(eventTimestamp, 0).Format(timeFormat)
				if *noColor {
					msg = fmt.Sprintf("%s - %s", ts, msg)
				} else {
					msg = fmt.Sprintf("%s - %s", color.GreenString(ts), msg)
				}
			}
			fmt.Println(msg)
		}
	}
}
<commit_msg>Fix error versioning<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/fatih/color"
	"github.com/lucagrulla/cw/cloudwatch"
	"gopkg.in/alecthomas/kingpin.v2"
)

var (
	timeFormat = "2006-01-02T15:04:05"
	version    = "2.1.2"

	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()
	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.").HintAction(groupsCompletion).Required().String()

	tailCommand     = kp.Command("tail", "Tail a log group.")
	follow          = tailCommand.Flag("follow", "Don't stop when the end of stream is reached, but rather wait for additional data to be appended.").Short('f').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()
	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()
	logGroupName    = tailCommand.Arg("group", "The log group name.").Required().HintAction(groupsCompletion).String()
	logStreamName   = tailCommand.Arg("stream", "The log stream name. Use \\* for tail all the group streams.").Default("*").HintAction(streamsCompletion).String()
	startTime       = tailCommand.Arg("start", "The UTC start time. Passed as either date/time or human-friendly format. The human-friendly format accepts the number of hours and minutes prior to the present. Denote hours with 'h' and minutes with 'm' i.e. 80m, 4h30m. If time is passed (format: hh[:mm]) it is expanded to today at the given time. Full available date/time format: 2017-02-27[T09:00[:00]].").Default(time.Now().UTC().Add(-30 * time.Second).Format(timeFormat)).String()
	endTime         = tailCommand.Arg("end", "The UTC start time. Passed as either date/time or human-friendly format. The human-friendly format accepts the number of hours and minutes prior to the present. Denote hours with 'h' and minutes with 'm' i.e. 80m, 4h30m. If time is passed (format: hh[:mm]) it is expanded to today at the given time. Full available date/time format: 2017-02-27[T09:00[:00]].").String()
	local           = tailCommand.Flag("local", "Treat date and time in Local zone.").Short('l').Default("false").Bool()
)

func groupsCompletion() []string {
	var groups []string
	kingpin.MustParse(kp.Parse(os.Args[1:]))

	for msg := range cloudwatch.New(awsProfile, awsRegion, debug).LsGroups() {
		groups = append(groups, *msg)
	}
	return groups
}

func streamsCompletion() []string {
	var streams []string
	kingpin.MustParse(kp.Parse(os.Args[1:]))

	for msg := range cloudwatch.New(awsProfile, awsRegion, debug).LsStreams(logGroupName, nil) {
		streams = append(streams, *msg)
	}
	return streams
}

func timestampToTime(timeStamp *string) time.Time {
	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
	} 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
	} 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
	} 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)
	} 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)
	} else if regexp.MustCompile(`^\d{1,}h$|^\d{1,}m$|^\d{1,}h\d{1,}m$`).MatchString(*timeStamp) {
		d, _ := time.ParseDuration(*timeStamp)

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

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

func fetchLatestVersion() chan string {
	latestVersionChannel := make(chan string, 1)
	go func() {
		r, e := http.Get("https://github.com/lucagrulla/cw/releases/latest")

		if e != nil {
			close(latestVersionChannel)
		} else {
			finalURL := r.Request.URL.String()
			tokens := strings.Split(finalURL, "/")
			latestVersionChannel <- tokens[len(tokens)-1]
		}
	}()
	return latestVersionChannel
}

func newVersionMsg(currentVersion string, latestVersionChannel chan string) {
	latestVersion, ok := <-latestVersionChannel
	//if the channel is closed it means we failed to fetch the latest version. Ignore the version message.
	if !ok {
		if latestVersion != fmt.Sprintf("v%s", currentVersion) {
			fmt.Println("")
			fmt.Println("")
			if *noColor {
				msg := fmt.Sprintf("%s - %s -> %s", "A new version of cw is available!", currentVersion, latestVersion)
				fmt.Println(msg)
			} else {
				msg := fmt.Sprintf("%s - %s -> %s", color.GreenString("A new version of cw is available!"), color.YellowString(currentVersion), color.GreenString(latestVersion))
				fmt.Println(msg)
			}
		}
	}
}

func versionCheckOnSigterm() {
	//only way to avoid print of the signal: interrupt message
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	<-c
	os.Exit(0)
}

func main() {
	kp.Version(version).Author("Luca Grulla")

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

	cmd := kingpin.MustParse(kp.Parse(os.Args[1:]))
	c := cloudwatch.New(awsProfile, awsRegion, debug)
	switch cmd {
	case "ls groups":

		for msg := range c.LsGroups() {
			fmt.Println(*msg)
		}
	case "ls streams":
		for msg := range c.LsStreams(lsLogGroupName, nil) {
			fmt.Println(*msg)
		}
	case "tail":
		st := timestampToTime(startTime)
		var et time.Time
		if *endTime != "" {
			et = timestampToTime(endTime)
		}
		for event := range c.Tail(logGroupName, logStreamName, follow, &st, &et, grep, grepv) {
			msg := *event.Message
			if *printEventID {
				if *noColor {
					msg = fmt.Sprintf("%s - %s", *event.EventId, msg)
				} else {
					msg = fmt.Sprintf("%s - %s", color.YellowString(*event.EventId), msg)
				}
			}
			if *printStreamName {
				if *noColor {
					msg = fmt.Sprintf("%s - %s", *event.LogStreamName, msg)
				} else {
					msg = fmt.Sprintf("%s - %s", color.BlueString(*event.LogStreamName), msg)
				}
			}
			if *printTimestamp {
				eventTimestamp := *event.Timestamp / 1000
				ts := time.Unix(eventTimestamp, 0).Format(timeFormat)
				if *noColor {
					msg = fmt.Sprintf("%s - %s", ts, msg)
				} else {
					msg = fmt.Sprintf("%s - %s", color.GreenString(ts), msg)
				}
			}
			fmt.Println(msg)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"flag"
	"io"
	"log"
	"mime"
	"mime/multipart"
	"net/mail"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/alexcesaro/quotedprintable"
	"github.com/mxk/go-imap/imap"
)

var wd, _ = os.Getwd()
var conf = flag.String("conf", wd+"/goatee.cfg", "Path to config file.")
type Config struct {
	Server      string
	User        string
	Password    string
	Destination string
}

type Goatee struct {
	config Config
	client *imap.Client
}

func (g *Goatee) Connect() {
	var err error
	log.Print("Connecting to server..\n")
	g.client, err = imap.DialTLS(g.config.Server, &tls.Config{})

	if err != nil {
		log.Fatalf("Connection to server failed: %s", err)
	}

	if g.client.State() == imap.Login {
		log.Print("Logging in..\n")
		g.client.Login(g.config.User, g.config.Password)
	}

	log.Print("Opening INBOX..\n")
	g.client.Select("INBOX", false)
}

func (g *Goatee) DecodeSubject(msg *mail.Message) string {
	s, _, err := quotedprintable.DecodeHeader(msg.Header.Get("Subject"))

	if err != nil {
		return msg.Header.Get("Subject")
	} else {
		return s
	}
}

func (g *Goatee) ExtractAttachment(r io.Reader, params map[string]string) {
	mr := multipart.NewReader(r, params["boundary"])
	for {
		p, err := mr.NextPart()
		if err == io.EOF {
			break
		} else if err != nil {
			log.Fatalf("Error parsing part: %s", err)
		}

		ct, params, _ := mime.ParseMediaType(p.Header.Get("Content-Type"))

		if strings.HasPrefix(ct, "multipart/mixed") {
			log.Printf("Extracting attachments from %s", ct)
			g.ExtractAttachment(p, params)
		} else if strings.HasPrefix(ct, "application/pdf") {
			path := filepath.Join(wd, g.config.Destination,
				p.FileName())
			dst, err := os.Create(path)
			if err != nil {
				log.Fatalf("Failed to create file: %s", err)
			}
			r := base64.NewDecoder(base64.StdEncoding, p)
			_, err = io.Copy(dst, r)
			if err != nil {
				log.Fatalf("Failed to store attachment: %s", err)
			}
		}
	}
}

func (g *Goatee) FetchMails() {
	log.Print("Fetching unread UIDs..\n")
	cmd, err := g.client.UIDSearch("1:* NOT SEEN")
	cmd.Result(imap.OK)

	if err != nil {
		log.Fatalf("UIDSearch failed: %s", err)
	}

	uids := cmd.Data[0].SearchResults()
	if len(uids) == 0 {
		log.Fatal("No unread messages found.")
	}

	log.Print("Fetching mail bodies..\n")
	set, _ := imap.NewSeqSet("")
	set.AddNum(uids...)
	cmd, err = g.client.Fetch(set, "UID", "FLAGS", "BODY[]")

	if err != nil {
		log.Fatalf("Fetch failed: %s", err)
	}

	for cmd.InProgress() {
		g.client.Recv(10 * time.Second)

		for _, rsp := range cmd.Data {
			body := imap.AsBytes(rsp.MessageInfo().Attrs["BODY[]"])

			if msg, _ := mail.ReadMessage(bytes.NewReader(body)); msg != nil {
				log.Printf("|-- %v", g.DecodeSubject(msg))
				mediaType, params, _ := mime.ParseMediaType(
					msg.Header.Get("Content-Type"))
				if strings.HasPrefix(mediaType, "multipart/") {
					log.Printf("Extracting attachments from %s", mediaType)
					g.ExtractAttachment(msg.Body, params)
				}
			}
		}
		cmd.Data = nil
	}

	if rsp, err := cmd.Result(imap.OK); err != nil {
		if err == imap.ErrAborted {
			log.Fatal("Fetch command aborted")
		} else {
			log.Fatalf("Fetch error:", rsp.Info)
		}
	}

	log.Print("Marking messages seen..\n")
	cmd, err = g.client.UIDStore(set, "+FLAGS.SILENT",
		imap.NewFlagSet(`\Seen`))

	if rsp, err := cmd.Result(imap.OK); err != nil {
		log.Fatalf("UIDStore error:", rsp.Info)
	}

	cmd.Data = nil
}

func (g *Goatee) ReadConfig(path string) {
	if _, err := os.Stat(path); err != nil {
		log.Fatalf("File doesn't exist: %v", err)
	}

	if _, err := toml.DecodeFile(path, &g.config); err != nil {
		log.Fatalf("Error opening config file: %s", err)
	}
}

func main() {
	flag.Parse()
	g := Goatee{}
	g.Connect()
	defer g.client.Logout(30 * time.Second)
	g.FetchMails()
	g.ReadConfig(*conf)
}
<commit_msg>Explicit logfile<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"flag"
	"io"
	"log"
	"mime"
	"mime/multipart"
	"net/mail"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/alexcesaro/quotedprintable"
	"github.com/mxk/go-imap/imap"
)

var wd, _ = os.Getwd()
var conf = flag.String("conf", wd+"/goatee.cfg", "Path to config file.")
var logfile = flag.String("log", wd+"/goatee.log", "Path to log file.")
type Config struct {
	Server      string
	User        string
	Password    string
	Destination string
}

type Goatee struct {
	config Config
	client *imap.Client
}

func (g *Goatee) Connect() {
	var err error
	log.Print("Connecting to server..\n")
	g.client, err = imap.DialTLS(g.config.Server, &tls.Config{})

	if err != nil {
		log.Fatalf("Connection to server failed: %s", err)
	}

	if g.client.State() == imap.Login {
		log.Print("Logging in..\n")
		g.client.Login(g.config.User, g.config.Password)
	}

	log.Print("Opening INBOX..\n")
	g.client.Select("INBOX", false)
}

func (g *Goatee) DecodeSubject(msg *mail.Message) string {
	s, _, err := quotedprintable.DecodeHeader(msg.Header.Get("Subject"))

	if err != nil {
		return msg.Header.Get("Subject")
	} else {
		return s
	}
}

func (g *Goatee) ExtractAttachment(r io.Reader, params map[string]string) {
	mr := multipart.NewReader(r, params["boundary"])
	for {
		p, err := mr.NextPart()
		if err == io.EOF {
			break
		} else if err != nil {
			log.Fatalf("Error parsing part: %s", err)
		}

		ct, params, _ := mime.ParseMediaType(p.Header.Get("Content-Type"))

		if strings.HasPrefix(ct, "multipart/mixed") {
			log.Printf("Extracting attachments from %s", ct)
			g.ExtractAttachment(p, params)
		} else if strings.HasPrefix(ct, "application/pdf") {
			path := filepath.Join(wd, g.config.Destination,
				p.FileName())
			dst, err := os.Create(path)
			if err != nil {
				log.Fatalf("Failed to create file: %s", err)
			}
			r := base64.NewDecoder(base64.StdEncoding, p)
			_, err = io.Copy(dst, r)
			if err != nil {
				log.Fatalf("Failed to store attachment: %s", err)
			}
		}
	}
}

func (g *Goatee) FetchMails() {
	log.Print("Fetching unread UIDs..\n")
	cmd, err := g.client.UIDSearch("1:* NOT SEEN")
	cmd.Result(imap.OK)

	if err != nil {
		log.Fatalf("UIDSearch failed: %s", err)
	}

	uids := cmd.Data[0].SearchResults()
	if len(uids) == 0 {
		log.Fatal("No unread messages found.")
	}

	log.Print("Fetching mail bodies..\n")
	set, _ := imap.NewSeqSet("")
	set.AddNum(uids...)
	cmd, err = g.client.Fetch(set, "UID", "FLAGS", "BODY[]")

	if err != nil {
		log.Fatalf("Fetch failed: %s", err)
	}

	for cmd.InProgress() {
		g.client.Recv(10 * time.Second)

		for _, rsp := range cmd.Data {
			body := imap.AsBytes(rsp.MessageInfo().Attrs["BODY[]"])

			if msg, _ := mail.ReadMessage(bytes.NewReader(body)); msg != nil {
				log.Printf("|-- %v", g.DecodeSubject(msg))
				mediaType, params, _ := mime.ParseMediaType(
					msg.Header.Get("Content-Type"))
				if strings.HasPrefix(mediaType, "multipart/") {
					log.Printf("Extracting attachments from %s", mediaType)
					g.ExtractAttachment(msg.Body, params)
				}
			}
		}
		cmd.Data = nil
	}

	if rsp, err := cmd.Result(imap.OK); err != nil {
		if err == imap.ErrAborted {
			log.Fatal("Fetch command aborted")
		} else {
			log.Fatalf("Fetch error:", rsp.Info)
		}
	}

	log.Print("Marking messages seen..\n")
	cmd, err = g.client.UIDStore(set, "+FLAGS.SILENT",
		imap.NewFlagSet(`\Seen`))

	if rsp, err := cmd.Result(imap.OK); err != nil {
		log.Fatalf("UIDStore error:", rsp.Info)
	}

	cmd.Data = nil
}

func (g *Goatee) OpenLog(path string) {
	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		log.Fatalf("Error opening logfile: %v", err)
	}
	log.SetOutput(io.MultiWriter(os.Stderr, f))
}

func (g *Goatee) ReadConfig(path string) {
	if _, err := os.Stat(path); err != nil {
		log.Fatalf("File doesn't exist: %v", err)
	}

	if _, err := toml.DecodeFile(path, &g.config); err != nil {
		log.Fatalf("Error opening config file: %s", err)
	}
}

func main() {
	flag.Parse()
	g := Goatee{}
	g.Connect()
	defer g.client.Logout(30 * time.Second)
	g.FetchMails()
	g.ReadConfig(*conf)
	g.OpenLog(*logfile)
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"github.com/andrew-d/go-termutil"
)

func main() {
	if termutil.Isatty(os.Stdin.Fd()) {
		receive()
	} else {
		publish()
	}
}
<commit_msg>flags to force publishing or receiving<commit_after>package main

import (
	"flag"
	"os"

	"github.com/andrew-d/go-termutil"
)

var (
	pubOnly  = flag.Bool("p", false, "only publish, regardless of TTY stdin")
	recvOnly = flag.Bool("r", false, "only receive, regardless of TTY stdin")
)

func main() {
	flag.Parse()

	switch {
	case *recvOnly:
		receive()
	case *pubOnly:
		publish()
	case termutil.Isatty(os.Stdin.Fd()):
		receive()
	default:
		publish()
	}
}
<|endoftext|>


<commit_before>package main

import (
  "github.com/pilu/traffic"
)

var router *traffic.Router

func init() {
  router = traffic.New()
  router.Get("/", RootHandler)
}

func main() {
  router.Run()
}
<commit_msg>Set hostname and port<commit_after>package main

import (
  "github.com/pilu/traffic"
)

var router *traffic.Router

func init() {
	traffic.SetHost("0.0.0.0")
	traffic.SetPort(3000)
	router = traffic.New()
	router.Get("/", RootHandler)
}

func main() {
  router.Run()
}
<|endoftext|>


<commit_before>// Copyright 2015 CloudAwan LLC
//
// 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 (
	_ "github.com/cloudawan/cloudone_analysis/execute"
	"github.com/cloudawan/cloudone_analysis/restapi"
	"github.com/cloudawan/cloudone_analysis/utility/logger"
)

var log = logger.GetLogManager().GetLogger("default")

func main() {
	restapi.StartRestAPIServer()
}
<commit_msg>Add close for execute<commit_after>// Copyright 2015 CloudAwan LLC
//
// 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 (
	"github.com/cloudawan/cloudone_analysis/execute"
	"github.com/cloudawan/cloudone_analysis/restapi"
	"github.com/cloudawan/cloudone_analysis/utility/logger"
)

var log = logger.GetLogManager().GetLogger("default")

func main() {
	restapi.StartRestAPIServer()
	execute.Close()
}
<|endoftext|>


<commit_before>package badger

import (
	"image"

	"bytes"

	"path/filepath"

	"math"

	"fmt"

	"io"

	"io/ioutil"

	"github.com/CirrusMD/badger/internal"
	"github.com/bmatcuk/doublestar"
	"github.com/nfnt/resize"
	"golang.org/x/image/font/inconsolata"
	gg "gopkg.in/fogleman/gg.v1"
)

const Version = "0.3.0"

type Options struct {
	// Semantic version, example: 3.0.1
	MarketingVersion string
	BuildNumber      string

	// Flag to use an alpha overlay image (default is beta)
	Alpha bool

	// Flag to use a dark alpha/beta overlay image (default is light)
	Dark bool

	// Glob pattern (default is ./**/*.appiconset)
	Glob string

	Logger io.Writer
}

func (o Options) isBeta() bool {
	return !o.Alpha
}

// Badge applies badge to all images found with the glob pattern given a set of Options
func Badge(opts Options) error {
	opts = validOptions(opts)
	images, err := findImages(opts.Glob)
	if err != nil {
		return err
	}

	for _, imgPath := range images {
		fmt.Fprintf(opts.Logger, "Badging %s...\n", imgPath)
		img, err := gg.LoadImage(imgPath)
		if err != nil {
			return fmt.Errorf("could not open file: %v", err)
		}
		parent := gg.NewContextForImage(img)

		if opts.MarketingVersion != "" {
			drawTopText(parent, opts.MarketingVersion, "#555555", 0)
		}
		if opts.BuildNumber != "" {
			drawTopText(parent, opts.BuildNumber, "f48041", parent.Width()/2)
		}
		if err := overlayBadgeImage(opts, parent); err != nil {
			return err
		}
		err = gg.SavePNG(imgPath, parent.Image())
		if err != nil {
			return fmt.Errorf("could not save png: %v", err)
		}
	}
	return nil
}

func validOptions(opts Options) Options {
	if opts.Glob == "" {
		opts.Glob = "./**/*.appiconset"
	}
	if opts.Logger == nil {
		opts.Logger = ioutil.Discard
	}
	return opts
}

func findImages(glob string) ([]string, error) {
	pattern := filepath.Join(glob, "*.png")
	pattern = filepath.Clean(pattern)
	images, err := doublestar.Glob(pattern)
	if err != nil {
		return images, fmt.Errorf("unable to parse glob pattern: %s, error: %v", pattern, err)
	}
	if len(images) == 0 {
		return images, fmt.Errorf("unable to find images for pattern: %s", pattern)
	}
	return images, nil
}

func drawTopText(parent *gg.Context, text string, hexColor string, x int) {
	temp := gg.NewContext(50, 100)
	temp.SetFontFace(inconsolata.Bold8x16)
	sw, sh := temp.MeasureString(text)

	vc := gg.NewContext(int(math.Ceil(sw)), int(math.Ceil(sh)))
	vc.SetHexColor(hexColor)
	vc.Clear()

	temp.SetFontFace(inconsolata.Bold8x16)
	vc.SetRGB(1, 1, 1)
	vc.DrawStringAnchored(text, sw/2, sh/2, 0.5, 0.5)

	img := vc.Image()
	w, h := versionDimensions(parent)
	img = resize.Resize(uint(w), uint(h), img, resize.NearestNeighbor)

	parent.DrawImage(img, x, 0)
}

func versionDimensions(dc *gg.Context) (int, int) {
	width := dc.Image().Bounds().Max.X / 2
	height := int(float64(dc.Image().Bounds().Max.Y)*0.2) - 1
	return width, height
}

func overlayBadgeImage(opts Options, parent *gg.Context) error {
	badge, err := findBadgeImage(opts)
	if err != nil {
		return err
	}
	badge = resize.Resize(uint(parent.Width()), uint(parent.Height()), badge, resize.NearestNeighbor)
	parent.DrawImage(badge, 0, 0)
	return nil
}

func findBadgeImage(opts Options) (image.Image, error) {
	imgName := ""
	if opts.Alpha && opts.Dark {
		imgName = "alpha_badge_dark.png"
	} else if opts.Alpha {
		imgName = "alpha_badge_light.png"
	} else if opts.isBeta() && opts.Dark {
		imgName = "beta_badge_dark.png"
	} else if opts.isBeta() {
		imgName = "beta_badge_light.png"
	}
	if imgName == "" {
		return nil, nil
	}
	raw, err := internal.Asset("assets/" + imgName)
	if err != nil {
		return nil, fmt.Errorf("could not load overlay image: %v", err)
	}

	img, _, err := image.Decode(bytes.NewReader(raw))
	if err != nil {
		return nil, fmt.Errorf("unable to decode overlay image: %v", err)
	}

	return img, nil
}
<commit_msg>bump version<commit_after>package badger

import (
	"image"

	"bytes"

	"path/filepath"

	"math"

	"fmt"

	"io"

	"io/ioutil"

	"github.com/CirrusMD/badger/internal"
	"github.com/bmatcuk/doublestar"
	"github.com/nfnt/resize"
	"golang.org/x/image/font/inconsolata"
	gg "gopkg.in/fogleman/gg.v1"
)

const Version = "0.3.2"

type Options struct {
	// Semantic version, example: 3.0.1
	MarketingVersion string
	BuildNumber      string

	// Flag to use an alpha overlay image (default is beta)
	Alpha bool

	// Flag to use a dark alpha/beta overlay image (default is light)
	Dark bool

	// Glob pattern (default is ./**/*.appiconset)
	Glob string

	Logger io.Writer
}

func (o Options) isBeta() bool {
	return !o.Alpha
}

// Badge applies badge to all images found with the glob pattern given a set of Options
func Badge(opts Options) error {
	opts = validOptions(opts)
	images, err := findImages(opts.Glob)
	if err != nil {
		return err
	}

	for _, imgPath := range images {
		fmt.Fprintf(opts.Logger, "Badging %s...\n", imgPath)
		img, err := gg.LoadImage(imgPath)
		if err != nil {
			return fmt.Errorf("could not open file: %v", err)
		}
		parent := gg.NewContextForImage(img)

		if opts.MarketingVersion != "" {
			drawTopText(parent, opts.MarketingVersion, "#555555", 0)
		}
		if opts.BuildNumber != "" {
			drawTopText(parent, opts.BuildNumber, "f48041", parent.Width()/2)
		}
		if err := overlayBadgeImage(opts, parent); err != nil {
			return err
		}
		err = gg.SavePNG(imgPath, parent.Image())
		if err != nil {
			return fmt.Errorf("could not save png: %v", err)
		}
	}
	return nil
}

func validOptions(opts Options) Options {
	if opts.Glob == "" {
		opts.Glob = "./**/*.appiconset"
	}
	if opts.Logger == nil {
		opts.Logger = ioutil.Discard
	}
	return opts
}

func findImages(glob string) ([]string, error) {
	pattern := filepath.Join(glob, "*.png")
	pattern = filepath.Clean(pattern)
	images, err := doublestar.Glob(pattern)
	if err != nil {
		return images, fmt.Errorf("unable to parse glob pattern: %s, error: %v", pattern, err)
	}
	if len(images) == 0 {
		return images, fmt.Errorf("unable to find images for pattern: %s", pattern)
	}
	return images, nil
}

func drawTopText(parent *gg.Context, text string, hexColor string, x int) {
	temp := gg.NewContext(50, 100)
	temp.SetFontFace(inconsolata.Bold8x16)
	sw, sh := temp.MeasureString(text)

	vc := gg.NewContext(int(math.Ceil(sw)), int(math.Ceil(sh)))
	vc.SetHexColor(hexColor)
	vc.Clear()

	temp.SetFontFace(inconsolata.Bold8x16)
	vc.SetRGB(1, 1, 1)
	vc.DrawStringAnchored(text, sw/2, sh/2, 0.5, 0.5)

	img := vc.Image()
	w, h := versionDimensions(parent)
	img = resize.Resize(uint(w), uint(h), img, resize.NearestNeighbor)

	parent.DrawImage(img, x, 0)
}

func versionDimensions(dc *gg.Context) (int, int) {
	width := dc.Image().Bounds().Max.X / 2
	height := int(float64(dc.Image().Bounds().Max.Y)*0.2) - 1
	return width, height
}

func overlayBadgeImage(opts Options, parent *gg.Context) error {
	badge, err := findBadgeImage(opts)
	if err != nil {
		return err
	}
	badge = resize.Resize(uint(parent.Width()), uint(parent.Height()), badge, resize.NearestNeighbor)
	parent.DrawImage(badge, 0, 0)
	return nil
}

func findBadgeImage(opts Options) (image.Image, error) {
	imgName := ""
	if opts.Alpha && opts.Dark {
		imgName = "alpha_badge_dark.png"
	} else if opts.Alpha {
		imgName = "alpha_badge_light.png"
	} else if opts.isBeta() && opts.Dark {
		imgName = "beta_badge_dark.png"
	} else if opts.isBeta() {
		imgName = "beta_badge_light.png"
	}
	if imgName == "" {
		return nil, nil
	}
	raw, err := internal.Asset("assets/" + imgName)
	if err != nil {
		return nil, fmt.Errorf("could not load overlay image: %v", err)
	}

	img, _, err := image.Decode(bytes.NewReader(raw))
	if err != nil {
		return nil, fmt.Errorf("unable to decode overlay image: %v", err)
	}

	return img, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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.

// WebBorer is a directory-enumeration tool written in Go.
package main

import (
	"github.com/Matir/webborer/client"
	"github.com/Matir/webborer/filter"
	"github.com/Matir/webborer/logging"
	"github.com/Matir/webborer/results"
	ss "github.com/Matir/webborer/settings"
	"github.com/Matir/webborer/util"
	"github.com/Matir/webborer/wordlist"
	"github.com/Matir/webborer/worker"
	"github.com/Matir/webborer/workqueue"
	"runtime"
)

// This is the main runner for webborer.
// TODO: separate the actual scanning from all of the setup steps
func main() {
	util.EnableStackTraces()

	// Load scan settings
	settings, err := ss.GetScanSettings()
	if err != nil {
		logging.Logf(logging.LogFatal, err.Error())
		return
	}
	logging.ResetLog(settings.LogfilePath, settings.LogLevel)
	logging.Logf(logging.LogInfo, "Flags: %s", settings)

	// Enable CPU profiling
	var cpuProfStop func()
	if settings.DebugCPUProf {
		cpuProfStop = util.EnableCPUProfiling()
	}

	// Set number of threads
	logging.Logf(logging.LogDebug, "Setting GOMAXPROCS to %d.", settings.Threads)
	runtime.GOMAXPROCS(settings.Threads)

	// Load wordlist
	var words []string
	words, err = wordlist.LoadWordlist(settings.WordlistPath)
	if err != nil {
		logging.Logf(logging.LogFatal, "Unable to load wordlist: %s", err.Error())
		return
	}

	// Build an HTTP Client Factory
	logging.Logf(logging.LogDebug, "Creating Client Factory...")
	clientFactory, err := client.NewProxyClientFactory(settings.Proxies, settings.Timeout, settings.UserAgent)
	if err != nil {
		logging.Logf(logging.LogFatal, "Unable to build client factory: %s", err.Error())
		return
	}
	clientFactory.SetUsernamePassword(settings.HTTPUsername, settings.HTTPPassword)

	// Starting point
	scope, err := settings.GetScopes()
	if err != nil {
		logging.Logf(logging.LogFatal, err.Error())
		return
	}

	// Setup the main workqueue
	logging.Logf(logging.LogDebug, "Starting work queue...")
	queue := workqueue.NewWorkQueue(settings.QueueSize, scope, settings.AllowHTTPSUpgrade)
	queue.RunInBackground()

	logging.Logf(logging.LogDebug, "Creating expander and filter...")
	expander := filter.Expander{Wordlist: &words, Adder: queue.GetAddCount()}
	expander.ProcessWordlist()
	filter := filter.NewWorkFilter(settings, queue.GetDoneFunc())

	// Check robots mode
	if settings.RobotsMode == ss.ObeyRobots {
		filter.AddRobotsFilter(scope, clientFactory)
	}

	work := filter.RunFilter(expander.Expand(queue.GetWorkChan()))

	logging.Logf(logging.LogDebug, "Creating results manager...")
	rchan := make(chan results.Result, settings.QueueSize)
	resultsManager, err := results.GetResultsManager(settings)
	if err != nil {
		logging.Logf(logging.LogFatal, "Unable to start results manager: %s", err.Error())
		return
	}

	logging.Logf(logging.LogDebug, "Starting %d workers...", settings.Workers)
	worker.StartWorkers(settings, clientFactory, work, queue.GetAddFunc(), queue.GetDoneFunc(), rchan)

	logging.Logf(logging.LogDebug, "Starting results manager...")
	resultsManager.Run(rchan)

	// Kick things off with the seed URL
	logging.Logf(logging.LogDebug, "Adding starting URLs: %v", scope)
	queue.AddURLs(scope...)

	// Potentially seed from robots
	if settings.RobotsMode == ss.SeedRobots {
		queue.SeedFromRobots(scope, clientFactory)
	}

	// Wait for work to be done
	logging.Logf(logging.LogDebug, "Main goroutine waiting for work...")
	queue.WaitPipe()
	logging.Logf(logging.LogDebug, "Work done.")

	// Cleanup
	queue.InputFinished()
	close(rchan)

	resultsManager.Wait()
	if cpuProfStop != nil {
		cpuProfStop()
	}
	logging.Logf(logging.LogDebug, "Done!")
}
<commit_msg>Minor updates to main.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// 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.

// WebBorer is a directory-enumeration tool written in Go.
package main

import (
	"github.com/Matir/webborer/client"
	"github.com/Matir/webborer/filter"
	"github.com/Matir/webborer/logging"
	"github.com/Matir/webborer/results"
	ss "github.com/Matir/webborer/settings"
	"github.com/Matir/webborer/util"
	"github.com/Matir/webborer/wordlist"
	"github.com/Matir/webborer/worker"
	"github.com/Matir/webborer/workqueue"
	"runtime"
)

// Load settings from flags
func loadSettings() (*ss.ScanSettings, error) {
	// Load scan settings
	settings, err := ss.GetScanSettings()
	if err != nil {
		logging.Logf(logging.LogFatal, err.Error())
		return nil, err
	}
	logging.ResetLog(settings.LogfilePath, settings.LogLevel)
	logging.Logf(logging.LogInfo, "Flags: %s", settings)
	return settings, nil
}

// This is the main runner for webborer.
// TODO: separate the actual scanning from all of the setup steps
func main() {
	util.EnableStackTraces()

	settings, err := loadSettings()
	if err != nil {
		return
	}

	// Enable CPU profiling
	var cpuProfStop func()
	if settings.DebugCPUProf {
		cpuProfStop = util.EnableCPUProfiling()
	}

	// Set number of threads
	logging.Logf(logging.LogDebug, "Setting GOMAXPROCS to %d.", settings.Threads)
	runtime.GOMAXPROCS(settings.Threads)

	// Load wordlist
	var words []string
	words, err = wordlist.LoadWordlist(settings.WordlistPath)
	if err != nil {
		logging.Logf(logging.LogFatal, "Unable to load wordlist: %s", err.Error())
		return
	}

	// Build an HTTP Client Factory
	logging.Logf(logging.LogDebug, "Creating Client Factory...")
	clientFactory, err := client.NewProxyClientFactory(settings.Proxies, settings.Timeout, settings.UserAgent)
	if err != nil {
		logging.Logf(logging.LogFatal, "Unable to build client factory: %s", err.Error())
		return
	}
	clientFactory.SetUsernamePassword(settings.HTTPUsername, settings.HTTPPassword)

	// Starting point
	scope, err := settings.GetScopes()
	if err != nil {
		logging.Logf(logging.LogFatal, err.Error())
		return
	}

	// Setup the main workqueue
	logging.Logf(logging.LogDebug, "Starting work queue...")
	queue := workqueue.NewWorkQueue(settings.QueueSize, scope, settings.AllowHTTPSUpgrade)
	queue.RunInBackground()

	logging.Logf(logging.LogDebug, "Creating expander and filter...")
	expander := filter.Expander{Wordlist: &words, Adder: queue.GetAddCount()}
	expander.ProcessWordlist()
	filter := filter.NewWorkFilter(settings, queue.GetDoneFunc())

	// Check robots mode
	if settings.RobotsMode == ss.ObeyRobots {
		filter.AddRobotsFilter(scope, clientFactory)
	}

	// filter paths after expansion
	logging.Debugf("Starting expansion and filtering...")
	work := filter.RunFilter(expander.Expand(queue.GetWorkChan()))

	logging.Logf(logging.LogDebug, "Creating results manager...")
	rchan := make(chan results.Result, settings.QueueSize)
	resultsManager, err := results.GetResultsManager(settings)
	if err != nil {
		logging.Logf(logging.LogFatal, "Unable to start results manager: %s", err.Error())
		return
	}

	logging.Logf(logging.LogDebug, "Starting %d workers...", settings.Workers)
	worker.StartWorkers(settings, clientFactory, work, queue.GetAddFunc(), queue.GetDoneFunc(), rchan)

	logging.Logf(logging.LogDebug, "Starting results manager...")
	resultsManager.Run(rchan)

	// Kick things off with the seed URL
	logging.Logf(logging.LogDebug, "Adding starting URLs: %v", scope)
	queue.AddURLs(scope...)

	// Potentially seed from robots
	if settings.RobotsMode == ss.SeedRobots {
		queue.SeedFromRobots(scope, clientFactory)
	}

	// Wait for work to be done
	logging.Logf(logging.LogDebug, "Main goroutine waiting for work...")
	queue.WaitPipe()
	logging.Logf(logging.LogDebug, "Work done.")

	// Cleanup
	queue.InputFinished()
	close(rchan)

	resultsManager.Wait()
	if cpuProfStop != nil {
		cpuProfStop()
	}
	logging.Logf(logging.LogDebug, "Done!")
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strconv"
	"strings"
	"sync"
)

const (
	CmdDisplayMax = 32
	// max length of /proc/$PID/comm
	CommMax = 16
	// from ps_mem - average error due to truncation in the kernel
	// pss calculations
	PssAdjust    = .5
	pageSize     = 4096
	mapDetailLen = len("Size:                  4 kB")

	usage = `Usage: %s [OPTION...]
Simple, accurate RAM and swap reporting.

Options:
`
)

var (
	memProfile string
	cpuProfile string

	tyPss          = []byte("Pss:")
	tySwap         = []byte("Swap:")
	tyPrivateClean = []byte("Private_Clean:")
	tyPrivateDirty = []byte("Private_Dirty:")
)

// store info about a command (group of processes), similar to how
// ps_mem works.
type CmdMemInfo struct {
	PIDs    []int
	Name    string
	Pss     float64
	Shared  float64
	Swapped float64
}

type MapInfo struct {
	Inode uint64
	Name  string
}

// mapLine is a line from /proc/$PID/maps, or one of the same header
// lines from smaps.
func NewMapInfo(mapLine []byte) MapInfo {
	var mi MapInfo
	var err error
	pieces := splitSpaces(mapLine)
	if len(pieces) == 6 {
		mi.Name = string(pieces[5])
	}
	mi.Inode, err = ParseUint(pieces[4], 10, 64)
	if err != nil {
		panic(fmt.Sprintf("NewMapInfo: Atoi(%s): %s (%s)",
			string(pieces[4]), err, string(mapLine)))
	}
	return mi
}

func (mi MapInfo) IsAnon() bool {
	return mi.Inode == 0
}

// isDigit returns true if the rune d represents an ascii digit
// between 0 and 9, inclusive.
func isDigit(d uint8) bool {
	return d >= '0' && d <= '9'
}

// pidList returns a list of the process-IDs of every currently
// running process on the local system.
func pidList() ([]int, error) {
	procLs, err := ioutil.ReadDir("/proc")
	if err != nil {
		return nil, fmt.Errorf("ReadDir(/proc): %s", err)
	}

	pids := make([]int, 0, len(procLs))
	for _, pInfo := range procLs {
		if !isDigit(pInfo.Name()[0]) || !pInfo.IsDir() {
			continue
		}
		pidInt, err := strconv.Atoi(pInfo.Name())
		if err != nil {
			return nil, fmt.Errorf("Atoi(%s): %s", pInfo.Name(), err)
		}
		pids = append(pids, pidInt)
	}
	return pids, nil
}

// procName gets the process name for a worker.  It first checks the
// value of /proc/$PID/cmdline.  If setproctitle(3) has been called,
// it will use this.  Otherwise it uses the value of
// path.Base(/proc/$PID/exe), which has info on whether the executable
// has changed since the process was exec'ed.
func procName(pid int) (string, error) {
	p, err := os.Readlink(fmt.Sprintf("/proc/%d/exe", pid))
	// this would return an error if the PID doesn't
	// exist, or if the PID refers to a kernel thread.
	if err != nil {
		return "", nil
	}
	// cmdline is the null separated list of command line
	// arguments for the process, unless setproctitle(3)
	// has been called, in which case it is the
	argsB, err := ioutil.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid))
	if err != nil {
		return "", fmt.Errorf("ReadFile(%s): %s", fmt.Sprintf("/proc/%d/cmdline", pid), err)
	}
	args := strings.Split(string(argsB), "\000")
	n := args[0]

	nTrunc := n
	if len(n) > CommMax {
		nTrunc = n[:CommMax]
	}
	if strings.HasPrefix(p, nTrunc) {
		n = path.Base(p)
	}
	return n, nil
}

// splitSpaces returns a slice of byte slices which are the space
// delimited words from the original byte slice.  Unlike
// strings.Split($X, " "), runs of multiple spaces in a row are
// discarded.  NOTE WELL: this only checks for spaces (' '), other
// unicode whitespace isn't supported.
func splitSpaces(b []byte) [][]byte {
	// most lines in smaps have the form "Swap: 4 kB", so
	// preallocate the slice's array appropriately.
	res := make([][]byte, 0, 3)
	start := 0
	for i := 0; i < len(b)-1; i++ {
		if b[i] == ' ' {
			start = i + 1
		} else if b[i+1] == ' ' {
			res = append(res, b[start:i+1])
			start = i + 1
		}
	}
	if start != len(b) && b[start] != ' ' {
		res = append(res, b[start:])
	}
	return res
}

// procMem returns the amount of Pss, shared, and swapped out memory
// used.  The swapped out amount refers to anonymous pages only.
func procMem(pid int) (pss, shared, swap float64, err error) {
	fPath := fmt.Sprintf("/proc/%d/smaps", pid)
	f, err := os.Open(fPath)
	if err != nil {
		err = fmt.Errorf("os.Open(%s): %s", fPath, err)
		return
	}
	var priv float64
	//var curr MapInfo
	r := bufio.NewReaderSize(f, pageSize)
	for {
		var l []byte
		var isPrefix bool
		l, isPrefix, err = r.ReadLine()
		// this should never happen, so take the easy way out.
		if isPrefix {
			err = fmt.Errorf("ReadLine(%s): isPrefix", fPath)
		}
		if err != nil {
			// if we've got EOF, then we're simply done
			// processing smaps.
			if err == io.EOF {
				err = nil
				break
			}
			// otherwise error out
			err = fmt.Errorf("ReadLine(%s): %s", fPath, err)
			return
		}

		if len(l) != mapDetailLen {
			//curr = NewMapInfo(l)
			continue
		}
		pieces := splitSpaces(l)
		ty := pieces[0]
		var v uint64
		if bytes.Equal(ty, tyPss) {
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			pss += float64(v) + PssAdjust
		} else if bytes.Equal(ty, tyPrivateClean) || bytes.Equal(ty, tyPrivateDirty) {
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			priv += float64(v)
		} else if bytes.Equal(ty, tySwap) {
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			swap += float64(v)
		}
	}
	shared = pss - priv
	return
}

func worker(pidRequest chan int, wg *sync.WaitGroup, result chan *CmdMemInfo) {
	for pid := range pidRequest {
		var err error
		cmi := new(CmdMemInfo)

		cmi.PIDs = []int{pid}
		cmi.Name, err = procName(pid)
		if err != nil {
			log.Printf("procName(%d): %s", pid, err)
			wg.Done()
			continue
		} else if cmi.Name == "" {
			// XXX: This happens with kernel
			// threads. maybe warn? idk.
			wg.Done()
			continue
		}

		cmi.Pss, cmi.Shared, cmi.Swapped, err = procMem(pid)
		if err != nil {
			log.Printf("procMem(%d): %s", pid, err)
			wg.Done()
			continue
		}

		result <- cmi
		wg.Done()
	}
}

type byTotal []*CmdMemInfo

func (c byTotal) Len() int           { return len(c) }
func (c byTotal) Less(i, j int) bool { return c[i].Pss < c[j].Pss }
func (c byTotal) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		flag.PrintDefaults()
	}

	flag.StringVar(&memProfile, "memprofile", "",
		"write memory profile to this file")
	flag.StringVar(&cpuProfile, "cpuprofile", "",
		"write cpu profile to this file")

	flag.Parse()
}

// startProfiling enables memory and/or CPU profiling if the
// appropriate command line flags have been set.
func startProfiling() {
	var err error
	// if we've passed in filenames to dump profiling data too,
	// start collecting profiling data.
	if memProfile != "" {
		runtime.MemProfileRate = 1
	}
	if cpuProfile != "" {
		var f *os.File
		if f, err = os.Create(cpuProfile); err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
	}
}

func stopProfiling() {
	if memProfile != "" {
		runtime.GC()
		f, err := os.Create(memProfile)
		if err != nil {
			log.Println(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
	}
	if cpuProfile != "" {
		pprof.StopCPUProfile()
		cpuProfile = ""
	}
}

func main() {
	startProfiling()
	defer stopProfiling()

	// need to be root to read map info for other user's
	// processes.
	if os.Geteuid() != 0 {
		fmt.Printf("%s requires root privileges. (try 'sudo `which %s`)\n",
			os.Args[0], os.Args[0])
		return
	}

	pids, err := pidList()
	if err != nil {
		log.Printf("pidList: %s", err)
		return
	}

	var wg sync.WaitGroup
	work := make(chan int, len(pids))
	result := make(chan *CmdMemInfo, len(pids))

	// give us as much parallelism as possible
	nCPU := runtime.NumCPU()
	runtime.GOMAXPROCS(nCPU)
	for i := 0; i < nCPU; i++ {
		go worker(work, &wg, result)
	}

	wg.Add(len(pids))
	for _, pid := range pids {
		work <- pid
	}
	wg.Wait()

	// aggregate similar processes by command name.
	cmdMap := map[string]*CmdMemInfo{}
loop:
	for {
		// this only works correctly because we a channel
		// where the buffer size >= the number of potential
		// results.
		select {
		case c := <-result:
			n := c.Name
			if _, ok := cmdMap[n]; !ok {
				cmdMap[n] = c
				continue
			}
			cmdMap[n].PIDs = append(cmdMap[n].PIDs, c.PIDs...)
			cmdMap[n].Pss += c.Pss
			cmdMap[n].Shared += c.Shared
			cmdMap[n].Swapped += c.Swapped
		default:
			break loop
		}
	}

	// extract map values to a slice so we can sort them
	cmds := make([]*CmdMemInfo, 0, len(cmdMap))
	for _, c := range cmdMap {
		cmds = append(cmds, c)
	}
	sort.Sort(byTotal(cmds))

	// keep track of total RAM and swap usage
	var totPss, totSwap float64

	fmt.Printf("%10s%10s%10s\t%s\n", "MB RAM", "SHARED", "SWAPPED", "PROCESS (COUNT)")
	for _, c := range cmds {
		n := c.Name
		if len(n) > CmdDisplayMax {
			if n[0] == '[' {
				n = n[:strings.IndexRune(n, ']')+1]
			} else {
				n = n[:CmdDisplayMax]
			}
		}
		s := ""
		if c.Swapped > 0 {
			swap := float64(c.Swapped) / 1024.
			totSwap += swap
			s = fmt.Sprintf("%10.1f", swap)
		}
		pss := float64(c.Pss) / 1024.
		fmt.Printf("%10.1f%10.1f%10s\t%s (%d)\n", pss, float64(c.Shared)/1024., s, n, len(c.PIDs))
		totPss += pss
	}
	fmt.Printf("#%9.1f%20.1f\tTOTAL USED BY PROCESSES\n", totPss, totSwap)
}
<commit_msg>doc: explain start/stopProfiling<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strconv"
	"strings"
	"sync"
)

const (
	CmdDisplayMax = 32
	// max length of /proc/$PID/comm
	CommMax = 16
	// from ps_mem - average error due to truncation in the kernel
	// pss calculations
	PssAdjust    = .5
	pageSize     = 4096
	mapDetailLen = len("Size:                  4 kB")

	usage = `Usage: %s [OPTION...]
Simple, accurate RAM and swap reporting.

Options:
`
)

var (
	memProfile string
	cpuProfile string

	tyPss          = []byte("Pss:")
	tySwap         = []byte("Swap:")
	tyPrivateClean = []byte("Private_Clean:")
	tyPrivateDirty = []byte("Private_Dirty:")
)

// store info about a command (group of processes), similar to how
// ps_mem works.
type CmdMemInfo struct {
	PIDs    []int
	Name    string
	Pss     float64
	Shared  float64
	Swapped float64
}

type MapInfo struct {
	Inode uint64
	Name  string
}

// mapLine is a line from /proc/$PID/maps, or one of the same header
// lines from smaps.
func NewMapInfo(mapLine []byte) MapInfo {
	var mi MapInfo
	var err error
	pieces := splitSpaces(mapLine)
	if len(pieces) == 6 {
		mi.Name = string(pieces[5])
	}
	mi.Inode, err = ParseUint(pieces[4], 10, 64)
	if err != nil {
		panic(fmt.Sprintf("NewMapInfo: Atoi(%s): %s (%s)",
			string(pieces[4]), err, string(mapLine)))
	}
	return mi
}

func (mi MapInfo) IsAnon() bool {
	return mi.Inode == 0
}

// isDigit returns true if the rune d represents an ascii digit
// between 0 and 9, inclusive.
func isDigit(d uint8) bool {
	return d >= '0' && d <= '9'
}

// pidList returns a list of the process-IDs of every currently
// running process on the local system.
func pidList() ([]int, error) {
	procLs, err := ioutil.ReadDir("/proc")
	if err != nil {
		return nil, fmt.Errorf("ReadDir(/proc): %s", err)
	}

	pids := make([]int, 0, len(procLs))
	for _, pInfo := range procLs {
		if !isDigit(pInfo.Name()[0]) || !pInfo.IsDir() {
			continue
		}
		pidInt, err := strconv.Atoi(pInfo.Name())
		if err != nil {
			return nil, fmt.Errorf("Atoi(%s): %s", pInfo.Name(), err)
		}
		pids = append(pids, pidInt)
	}
	return pids, nil
}

// procName gets the process name for a worker.  It first checks the
// value of /proc/$PID/cmdline.  If setproctitle(3) has been called,
// it will use this.  Otherwise it uses the value of
// path.Base(/proc/$PID/exe), which has info on whether the executable
// has changed since the process was exec'ed.
func procName(pid int) (string, error) {
	p, err := os.Readlink(fmt.Sprintf("/proc/%d/exe", pid))
	// this would return an error if the PID doesn't
	// exist, or if the PID refers to a kernel thread.
	if err != nil {
		return "", nil
	}
	// cmdline is the null separated list of command line
	// arguments for the process, unless setproctitle(3)
	// has been called, in which case it is the
	argsB, err := ioutil.ReadFile(fmt.Sprintf("/proc/%d/cmdline", pid))
	if err != nil {
		return "", fmt.Errorf("ReadFile(%s): %s", fmt.Sprintf("/proc/%d/cmdline", pid), err)
	}
	args := strings.Split(string(argsB), "\000")
	n := args[0]

	nTrunc := n
	if len(n) > CommMax {
		nTrunc = n[:CommMax]
	}
	if strings.HasPrefix(p, nTrunc) {
		n = path.Base(p)
	}
	return n, nil
}

// splitSpaces returns a slice of byte slices which are the space
// delimited words from the original byte slice.  Unlike
// strings.Split($X, " "), runs of multiple spaces in a row are
// discarded.  NOTE WELL: this only checks for spaces (' '), other
// unicode whitespace isn't supported.
func splitSpaces(b []byte) [][]byte {
	// most lines in smaps have the form "Swap: 4 kB", so
	// preallocate the slice's array appropriately.
	res := make([][]byte, 0, 3)
	start := 0
	for i := 0; i < len(b)-1; i++ {
		if b[i] == ' ' {
			start = i + 1
		} else if b[i+1] == ' ' {
			res = append(res, b[start:i+1])
			start = i + 1
		}
	}
	if start != len(b) && b[start] != ' ' {
		res = append(res, b[start:])
	}
	return res
}

// procMem returns the amount of Pss, shared, and swapped out memory
// used.  The swapped out amount refers to anonymous pages only.
func procMem(pid int) (pss, shared, swap float64, err error) {
	fPath := fmt.Sprintf("/proc/%d/smaps", pid)
	f, err := os.Open(fPath)
	if err != nil {
		err = fmt.Errorf("os.Open(%s): %s", fPath, err)
		return
	}
	var priv float64
	//var curr MapInfo
	r := bufio.NewReaderSize(f, pageSize)
	for {
		var l []byte
		var isPrefix bool
		l, isPrefix, err = r.ReadLine()
		// this should never happen, so take the easy way out.
		if isPrefix {
			err = fmt.Errorf("ReadLine(%s): isPrefix", fPath)
		}
		if err != nil {
			// if we've got EOF, then we're simply done
			// processing smaps.
			if err == io.EOF {
				err = nil
				break
			}
			// otherwise error out
			err = fmt.Errorf("ReadLine(%s): %s", fPath, err)
			return
		}

		if len(l) != mapDetailLen {
			//curr = NewMapInfo(l)
			continue
		}
		pieces := splitSpaces(l)
		ty := pieces[0]
		var v uint64
		if bytes.Equal(ty, tyPss) {
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			pss += float64(v) + PssAdjust
		} else if bytes.Equal(ty, tyPrivateClean) || bytes.Equal(ty, tyPrivateDirty) {
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			priv += float64(v)
		} else if bytes.Equal(ty, tySwap) {
			v, err = ParseUint(pieces[1], 10, 64)
			if err != nil {
				err = fmt.Errorf("Atoi(%s): %s", string(pieces[1]), err)
				return
			}
			swap += float64(v)
		}
	}
	shared = pss - priv
	return
}

func worker(pidRequest chan int, wg *sync.WaitGroup, result chan *CmdMemInfo) {
	for pid := range pidRequest {
		var err error
		cmi := new(CmdMemInfo)

		cmi.PIDs = []int{pid}
		cmi.Name, err = procName(pid)
		if err != nil {
			log.Printf("procName(%d): %s", pid, err)
			wg.Done()
			continue
		} else if cmi.Name == "" {
			// XXX: This happens with kernel
			// threads. maybe warn? idk.
			wg.Done()
			continue
		}

		cmi.Pss, cmi.Shared, cmi.Swapped, err = procMem(pid)
		if err != nil {
			log.Printf("procMem(%d): %s", pid, err)
			wg.Done()
			continue
		}

		result <- cmi
		wg.Done()
	}
}

type byTotal []*CmdMemInfo

func (c byTotal) Len() int           { return len(c) }
func (c byTotal) Less(i, j int) bool { return c[i].Pss < c[j].Pss }
func (c byTotal) Swap(i, j int)      { c[i], c[j] = c[j], c[i] }

func init() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage, os.Args[0])
		flag.PrintDefaults()
	}

	flag.StringVar(&memProfile, "memprofile", "",
		"write memory profile to this file")
	flag.StringVar(&cpuProfile, "cpuprofile", "",
		"write cpu profile to this file")

	flag.Parse()
}

// startProfiling enables memory and/or CPU profiling if the
// appropriate command line flags have been set.
func startProfiling() {
	var err error
	// if we've passed in filenames to dump profiling data too,
	// start collecting profiling data.
	if memProfile != "" {
		runtime.MemProfileRate = 1
	}
	if cpuProfile != "" {
		var f *os.File
		if f, err = os.Create(cpuProfile); err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
	}
}

func stopProfiling() {
	if memProfile != "" {
		runtime.GC()
		f, err := os.Create(memProfile)
		if err != nil {
			log.Println(err)
		}
		pprof.WriteHeapProfile(f)
		f.Close()
	}
	if cpuProfile != "" {
		pprof.StopCPUProfile()
		cpuProfile = ""
	}
}

func main() {
	// if -memprof or -cpuprof haven't been set on the command
	// line, these are nops
	startProfiling()
	defer stopProfiling()

	// need to be root to read map info for other user's
	// processes.
	if os.Geteuid() != 0 {
		fmt.Printf("%s requires root privileges. (try 'sudo `which %s`)\n",
			os.Args[0], os.Args[0])
		return
	}

	pids, err := pidList()
	if err != nil {
		log.Printf("pidList: %s", err)
		return
	}

	var wg sync.WaitGroup
	work := make(chan int, len(pids))
	result := make(chan *CmdMemInfo, len(pids))

	// give us as much parallelism as possible
	nCPU := runtime.NumCPU()
	runtime.GOMAXPROCS(nCPU)
	for i := 0; i < nCPU; i++ {
		go worker(work, &wg, result)
	}

	wg.Add(len(pids))
	for _, pid := range pids {
		work <- pid
	}
	wg.Wait()

	// aggregate similar processes by command name.
	cmdMap := map[string]*CmdMemInfo{}
loop:
	for {
		// this only works correctly because we a channel
		// where the buffer size >= the number of potential
		// results.
		select {
		case c := <-result:
			n := c.Name
			if _, ok := cmdMap[n]; !ok {
				cmdMap[n] = c
				continue
			}
			cmdMap[n].PIDs = append(cmdMap[n].PIDs, c.PIDs...)
			cmdMap[n].Pss += c.Pss
			cmdMap[n].Shared += c.Shared
			cmdMap[n].Swapped += c.Swapped
		default:
			break loop
		}
	}

	// extract map values to a slice so we can sort them
	cmds := make([]*CmdMemInfo, 0, len(cmdMap))
	for _, c := range cmdMap {
		cmds = append(cmds, c)
	}
	sort.Sort(byTotal(cmds))

	// keep track of total RAM and swap usage
	var totPss, totSwap float64

	fmt.Printf("%10s%10s%10s\t%s\n", "MB RAM", "SHARED", "SWAPPED", "PROCESS (COUNT)")
	for _, c := range cmds {
		n := c.Name
		if len(n) > CmdDisplayMax {
			if n[0] == '[' {
				n = n[:strings.IndexRune(n, ']')+1]
			} else {
				n = n[:CmdDisplayMax]
			}
		}
		s := ""
		if c.Swapped > 0 {
			swap := float64(c.Swapped) / 1024.
			totSwap += swap
			s = fmt.Sprintf("%10.1f", swap)
		}
		pss := float64(c.Pss) / 1024.
		fmt.Printf("%10.1f%10.1f%10s\t%s (%d)\n", pss, float64(c.Shared)/1024., s, n, len(c.PIDs))
		totPss += pss
	}
	fmt.Printf("#%9.1f%20.1f\tTOTAL USED BY PROCESSES\n", totPss, totSwap)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"
	"runtime"
	"sync"
	"time"

	"github.com/couchbaselabs/gateload/api"
	"github.com/couchbaselabs/gateload/workload"
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	// start up an http server, just to serve up expvars
	go http.ListenAndServe(":9876", nil)

	var config workload.Config
	workload.ReadConfig(&config)

	admin := api.SyncGatewayClient{}
	admin.Init(config.Hostname, config.Database, config.Port, config.AdminPort, config.LogRequests)
	if !admin.Valid() {
		log.Fatalf("unable to connect to sync_gateway, check the hostname and database")
	}

	pendingUsers := make(chan *workload.User)
	users := make([]*workload.User, config.NumPullers+config.NumPushers)

	// start a routine to place pending users into array
	go func() {
		for pendingUser := range pendingUsers {

			// users = append(users, pendingUser)
			users[pendingUser.SeqId-config.UserOffset] = pendingUser
		}
	}()

	rampUpDelay := config.RampUpIntervalMs / (config.NumPullers + config.NumPushers)
	// use a fixed number of workers to create the users/sessions
	userIterator := workload.UserIterator(
		config.NumPullers,
		config.NumPushers,
		config.UserOffset,
		config.ChannelActiveUsers,
		config.ChannelConcurrentUsers,
		config.MinUserOffTimeMs,
		config.MaxUserOffTimeMs,
		rampUpDelay,
		config.RunTimeMs,
	)
	adminWg := sync.WaitGroup{}
	worker := func() {
		defer adminWg.Done()
		for user := range userIterator {
			createSession(&admin, user, config)
			pendingUsers <- user
		}
	}

	for i := 0; i < 16; i++ {
		adminWg.Add(1)
		go worker()
	}

	// wait for all the workers to finish
	adminWg.Wait()
	// close the pending users channel to free that routine
	close(pendingUsers)

	numChannels := (config.NumPullers + config.NumPushers) / config.ChannelActiveUsers
	channelRampUpDelayMs := time.Duration(config.RampUpIntervalMs/numChannels) * time.Millisecond
	wg := sync.WaitGroup{}
	channel := ""
	for _, user := range users {
		nextChannel := user.Channel
		if channel != nextChannel {
			if channel != "" {
				time.Sleep(channelRampUpDelayMs)
			}
			channel = nextChannel
		}
		wg := sync.WaitGroup{}
		go runUser(user, config, &wg)
		wg.Add(1)
	}

	if config.RunTimeMs > 0 {
		time.Sleep(time.Duration(config.RunTimeMs-config.RampUpIntervalMs) * time.Millisecond)
		log.Println("Shutting down clients")
	} else {
		wg.Wait()
	}
}

func createSession(admin *api.SyncGatewayClient, user *workload.User, config workload.Config) {
	userMeta := api.UserAuth{
		Name:          user.Name,
		Password:      "password",
		AdminChannels: []string{user.Channel},
	}
	admin.AddUser(user.Name, userMeta)

	session := api.Session{Name: user.Name, TTL: 2592000} // 1 month
	user.Cookie = admin.CreateSession(user.Name, session)
}

func runUser(user *workload.User, config workload.Config, wg *sync.WaitGroup) {
	c := api.SyncGatewayClient{}
	c.Init(
		config.Hostname,
		config.Database,
		config.Port,
		config.AdminPort,
		config.LogRequests,
	)
	c.AddCookie(&user.Cookie)

	log.Printf("Starting new %s (%s)", user.Type, user.Name)
	if user.Type == "pusher" {
		go workload.RunNewPusher(
			user.Schedule,
			user.Name,
			&c,
			user.Channel,
			config.DocSize,
			config.SendAttachment,
			config.DocSizeDistribution,
			user.SeqId,
			config.SleepTimeMs,
			wg,
		)
	} else {
		go workload.RunNewPuller(
			user.Schedule,
			&c,
			user.Channel,
			user.Name,
			config.FeedType,
			wg,
		)
	}

}
<commit_msg>Reformat some code<commit_after>package main

import (
	"log"
	"net/http"
	"runtime"
	"sync"
	"time"

	"github.com/couchbaselabs/gateload/api"
	"github.com/couchbaselabs/gateload/workload"
)

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	// start up an http server, just to serve up expvars
	go http.ListenAndServe(":9876", nil)

	var config workload.Config
	workload.ReadConfig(&config)

	admin := api.SyncGatewayClient{}
	admin.Init(
		config.Hostname,
		config.Database,
		config.Port,
		config.AdminPort,
		config.LogRequests,
	)
	if !admin.Valid() {
		log.Fatalf("unable to connect to sync_gateway, check the hostname and database")
	}

	pendingUsers := make(chan *workload.User)
	users := make([]*workload.User, config.NumPullers+config.NumPushers)

	// start a routine to place pending users into array
	go func() {
		for pendingUser := range pendingUsers {

			// users = append(users, pendingUser)
			users[pendingUser.SeqId-config.UserOffset] = pendingUser
		}
	}()

	rampUpDelay := config.RampUpIntervalMs / (config.NumPullers + config.NumPushers)

	// use a fixed number of workers to create the users/sessions
	userIterator := workload.UserIterator(
		config.NumPullers,
		config.NumPushers,
		config.UserOffset,
		config.ChannelActiveUsers,
		config.ChannelConcurrentUsers,
		config.MinUserOffTimeMs,
		config.MaxUserOffTimeMs,
		rampUpDelay,
		config.RunTimeMs,
	)
	adminWg := sync.WaitGroup{}
	worker := func() {
		defer adminWg.Done()
		for user := range userIterator {
			createSession(&admin, user, config)
			pendingUsers <- user
		}
	}

	for i := 0; i < 16; i++ {
		adminWg.Add(1)
		go worker()
	}

	// wait for all the workers to finish
	adminWg.Wait()
	// close the pending users channel to free that routine
	close(pendingUsers)

	numChannels := (config.NumPullers + config.NumPushers) / config.ChannelActiveUsers
	channelRampUpDelayMs := time.Duration(config.RampUpIntervalMs/numChannels) * time.Millisecond

	wg := sync.WaitGroup{}
	channel := ""
	for _, user := range users {
		nextChannel := user.Channel
		if channel != nextChannel {
			if channel != "" {
				time.Sleep(channelRampUpDelayMs)
			}
			channel = nextChannel
		}
		wg := sync.WaitGroup{}
		go runUser(user, config, &wg)
		wg.Add(1)
	}

	if config.RunTimeMs > 0 {
		time.Sleep(time.Duration(config.RunTimeMs-config.RampUpIntervalMs) * time.Millisecond)
		log.Println("Shutting down clients")
	} else {
		wg.Wait()
	}
}

func createSession(admin *api.SyncGatewayClient, user *workload.User, config workload.Config) {
	userMeta := api.UserAuth{
		Name:          user.Name,
		Password:      "password",
		AdminChannels: []string{user.Channel},
	}
	admin.AddUser(user.Name, userMeta)

	session := api.Session{Name: user.Name, TTL: 2592000} // 1 month
	user.Cookie = admin.CreateSession(user.Name, session)
}

func runUser(user *workload.User, config workload.Config, wg *sync.WaitGroup) {
	c := api.SyncGatewayClient{}
	c.Init(
		config.Hostname,
		config.Database,
		config.Port,
		config.AdminPort,
		config.LogRequests,
	)
	c.AddCookie(&user.Cookie)

	log.Printf("Starting new %s (%s)", user.Type, user.Name)
	if user.Type == "pusher" {
		go workload.RunNewPusher(
			user.Schedule,
			user.Name,
			&c,
			user.Channel,
			config.DocSize,
			config.SendAttachment,
			config.DocSizeDistribution,
			user.SeqId,
			config.SleepTimeMs,
			wg,
		)
	} else {
		go workload.RunNewPuller(
			user.Schedule,
			&c,
			user.Channel,
			user.Name,
			config.FeedType,
			wg,
		)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"strings"
)

var pattern = flag.String("pattern", "", "Search pattern")
var chars = flag.Int("chars", 20, "Characters around pattern")

// foundPos stores position of occurence and line number
type foundPos struct {
	runePos int
	lineNo  int
}

// runeIndex is rune version of strings.Index
func runeIndex(rs, rsep []rune) int {
	if len(rs) < len(rsep) {
		return -1
	}
	for i, _ := range rs {
		matched := true
		p := i
		for j, _ := range rsep {
			if rs[p] != rsep[j] {
				matched = false
				break
			}
			p++
		}
		if matched {
			return i
		}
	}

	return -1
}

// eachRunePos returns position of sep in s
func eachRunePos(s, sep string) []int {
	index := 0
	result := make([]int, 0)
	rs := []rune(s)
	rsep := []rune(sep)
	pos := runeIndex(rs[index:], rsep)
	for pos != -1 {
		result = append(result, index+pos)
		index += pos + 1
		pos = runeIndex(rs[index:], rsep)
	}
	return result
}

func main() {
	flag.Parse()
	if *pattern == "" {
		if flag.NArg() == 0 {
			flag.Usage()
		}
		os.Exit(1)
	}

	runes := make([]rune, 0) //text converted to runes
	lastLen := 0
	foundPosList := make([]foundPos, 0)

	scanner := bufio.NewScanner(os.Stdin)
	for lineNo := 1; scanner.Scan(); lineNo++ {
		line := strings.TrimSpace(scanner.Text()) + " "
		for _, r := range []rune(line) {
			runes = append(runes, r)
		}

		for _, pos := range eachRunePos(line, *pattern) {
			foundPosList = append(foundPosList, foundPos{lastLen + pos, lineNo})
		}

		lastLen = len(runes)
	}
	if err := scanner.Err(); err != nil {
		fmt.Errorf(err.Error())
		os.Exit(1)
	}

	runePattern := []rune(*pattern)

	for _, foundPos := range foundPosList {
		pos := foundPos.runePos
		fmt.Printf("%4d\t%s\n", foundPos.lineNo, string(runes[pos-*chars:pos+len(runePattern)+*chars]))
	}
}
<commit_msg>fixed crash when pattern appears from head/tail of text<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"
	"strings"
)

var pattern = flag.String("pattern", "", "Search pattern")
var chars = flag.Int("chars", 20, "Characters around pattern")

// foundPos stores position of occurence and line number
type foundPos struct {
	runePos int
	lineNo  int
}

// runeIndex is rune version of strings.Index
func runeIndex(rs, rsep []rune) int {
	if len(rs) < len(rsep) {
		return -1
	}
	for i, _ := range rs {
		matched := true
		p := i
		for j, _ := range rsep {
			if rs[p] != rsep[j] {
				matched = false
				break
			}
			p++
		}
		if matched {
			return i
		}
	}

	return -1
}

// eachRunePos returns position of sep in s
func eachRunePos(s, sep string) []int {
	index := 0
	result := make([]int, 0)
	rs := []rune(s)
	rsep := []rune(sep)
	pos := runeIndex(rs[index:], rsep)
	for pos != -1 {
		result = append(result, index+pos)
		index += pos + 1
		pos = runeIndex(rs[index:], rsep)
	}
	return result
}

// substring with format(centering)
func formatRune(runes []rune, from, to int) string {
	centeringSpace := ""
	if from < 0 {
		centeringSpace = strings.Repeat(" ", 0-from)
		from = 0
	}
	if to > len(runes) {
		to = len(runes)
	}

	return centeringSpace + string(runes[from:to])
}

func main() {
	flag.Parse()
	if *pattern == "" {
		if flag.NArg() == 0 {
			flag.Usage()
		}
		os.Exit(1)
	}

	runes := make([]rune, 0) //text converted to runes
	lastLen := 0
	foundPosList := make([]foundPos, 0)

	scanner := bufio.NewScanner(os.Stdin)
	for lineNo := 1; scanner.Scan(); lineNo++ {
		line := strings.TrimSpace(scanner.Text()) + " "
		for _, r := range []rune(line) {
			runes = append(runes, r)
		}

		for _, pos := range eachRunePos(line, *pattern) {
			foundPosList = append(foundPosList, foundPos{lastLen + pos, lineNo})
		}

		lastLen = len(runes)
	}
	if err := scanner.Err(); err != nil {
		fmt.Errorf(err.Error())
		os.Exit(1)
	}

	runePattern := []rune(*pattern)

	for _, foundPos := range foundPosList {
		from := foundPos.runePos - *chars
		to := foundPos.runePos + len(runePattern) + *chars
		str := formatRune(runes, from, to)

		fmt.Printf("%4d\t%s\n", foundPos.lineNo, str)
	}
}
<|endoftext|>