<commit_before>// This software is direct fork of https://github.com/knq/chromedp/tree/master/cmd/chromedp-proxy
// with couple of features added
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"

	"errors"

	"github.com/Sirupsen/logrus"
)

var (
	flagListen         = flag.String("l", "localhost:9223", "listen address")
	flagRemote         = flag.String("r", "localhost:9222", "remote address")
	flagEllipsis       = flag.Bool("s", false, "shorten requests and responses")
	flagOnce           = flag.Bool("once", false, "debug single session")
	flagShowRequests   = flag.Bool("i", false, "include request frames as they are sent")
	flagDistributeLogs = flag.Bool("d", false, "write logs file per targetId")
	flagQuiet          = flag.Bool("q", false, "do not show logs on stdout")
)

var protocolTargetID = center("protocol message", 36)

func main() {
	flag.Parse()
	mux := http.NewServeMux()

	simpleReverseProxy := httputil.NewSingleHostReverseProxy(&url.URL{Scheme: "http", Host: *flagRemote})

	mux.Handle("/json", simpleReverseProxy)
	mux.Handle("/", simpleReverseProxy)

	rootLogger, err := createLogger("connection")
	if err != nil {
		panic(fmt.Sprintf("could not create logger: %s", err))
	}

	logger := rootLogger.WithFields(logrus.Fields{
		fieldLevel:       levelConnection,
	})

	mux.HandleFunc("/devtools/page/", func(res http.ResponseWriter, req *http.Request) {

		stream := make(chan *protocolMessage, 1024)
		id := path.Base(req.URL.Path)

		var protocolLogger *logrus.Entry

		if *flagDistributeLogs {
			logger, err := createLogger("inspector-" + id)
			if err != nil {
				panic(fmt.Sprintf("could not create logger: %s", err))
			}

			protocolLogger = logger.WithFields(logrus.Fields{
				fieldLevel:       levelConnection,
				fieldInspectorId: id,
			})

		} else {
			protocolLogger = logger.WithFields(logrus.Fields{
				fieldInspectorId: id,
			})
		}

		go dumpStream(protocolLogger, stream)

		endpoint := "ws://" + *flagRemote + "/devtools/page/" + id

		logger.Infof("---------- connection from %s ----------", req.RemoteAddr)
		logger.Infof("checking protocol versions on: %s", endpoint)

		ver, err := checkVersion()
		if err != nil {
			protocolLogger.Errorf("could not check version: %v", err)
			http.Error(res, "could not check version", 500)
			return
		}

		logger.Infof("protocol version: %s", ver["Protocol-Version"])
		logger.Infof("versions: Chrome(%s), V8(%s), Webkit(%s)", ver["Browser"], ver["V8-Version"], ver["WebKit-Version"])
		logger.Infof("browser user agent: %s", ver["User-Agent"])
		logger.Infof("connecting to %s... ", endpoint)

		// connecting to ws
		out, pres, err := wsDialer.Dial(endpoint, nil)
		if err != nil {
			msg := fmt.Sprintf("could not connect to %s: %v", endpoint, err)
			logger.Error(protocolError(msg))
			http.Error(res, msg, 500)
			return
		}
		defer pres.Body.Close()
		defer out.Close()

		// connect incoming websocket
		logger.Infof("upgrading connection on %s...", req.RemoteAddr)
		in, err := wsUpgrader.Upgrade(res, req, nil)
		if err != nil {
			logger.Errorf("could not upgrade websocket from %s: %v", req.RemoteAddr, err)
			http.Error(res, "could not upgrade websocket connection", 500)
			return
		}
		defer in.Close()

		ctxt, cancel := context.WithCancel(context.Background())
		defer cancel()

		errc := make(chan error, 1)
		go proxyWS(ctxt, stream, in, out, errc)
		go proxyWS(ctxt, stream, out, in, errc)

		<-errc
		logger.Infof("---------- closing %s ----------", req.RemoteAddr)

		if *flagOnce {
			os.Exit(0)
		}
	})

	log.Fatal(http.ListenAndServe(*flagListen, mux))
}

func dumpStream(logger *logrus.Entry, stream chan *protocolMessage) {
	logger.Printf("Legend: %s, %s, %s, %s, %s, %s", protocolColor("protocol informations"),
		eventsColor("received events"),
		requestColor("sent request frames"),
		requestReplyColor("requests params"),
		responseColor("received responses"),
		errorColor("error response."),
	)

	requests := make(map[uint64]*protocolMessage)
	targetRequests := make(map[uint64]*protocolMessage)

	for {
		select {
		case msg := <-stream:
			if msg.InTarget() {

				var targetLogger *logrus.Entry

				if *flagDistributeLogs {
					logger, err := createLogger(fmt.Sprintf("target-%s", msg.Params["targetId"]))
					if err != nil {
						panic(fmt.Sprintf("could not create logger: %v", err))
					}

					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetId: msg.Params["targetId"],
					})

				} else {
					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetId: msg.Params["targetId"],
					})
				}

				if msg.IsRequest() {
					requests[msg.Id] = nil

					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						targetRequests[protocolMessage.Id] = protocolMessage

						if *flagShowRequests {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeRequest,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Error("Could not deserialize message: %+v", err)
					}
				}

				if msg.IsEvent() {
					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						if protocolMessage.IsEvent() {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeEvent,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

						if protocolMessage.IsResponse() {
							var logMessage string
							var logType int
							var logRequest string
							var logMethod string

							if protocolMessage.IsError() {
								logMessage = serialize(protocolMessage.Error)
								logType = typeRequestResponseError
							} else {
								logMessage = serialize(protocolMessage.Result)
								logType = typeRequestResponse
							}

							if request, ok := targetRequests[protocolMessage.Id]; ok && request != nil {
								delete(targetRequests, protocolMessage.Id)
								logRequest = serialize(request.Params)
								logMethod = request.Method

							} else {
								logRequest = errorColor("could not find request with id: %d", protocolMessage.Id)
							}

							targetLogger.WithFields(logrus.Fields{
								fieldType:    logType,
								fieldMethod:  logMethod,
								fieldRequest: logRequest,
							}).Info(logMessage)
						}
					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Error("Could not deserialize message: %+v", err)
					}
				}

			} else {
				protocolLogger := logger.WithFields(logrus.Fields{
					fieldLevel:    levelProtocol,
					fieldTargetId: protocolTargetID,
				})

				if msg.IsRequest() {
					requests[msg.Id] = msg

					if *flagShowRequests {
						protocolLogger.WithFields(logrus.Fields{
							fieldType:   typeRequest,
							fieldMethod: msg.Method,
						}).Info(serialize(msg.Params))
					}
				}

				if msg.IsResponse() {

					var logMessage string
					var logType int
					var logRequest string
					var logMethod string

					if msg.IsError() {
						logMessage = serialize(msg.Error)
						logType = typeRequestResponseError
					} else {
						logMessage = serialize(msg.Result)
						logType = typeRequestResponse
					}

					if request, ok := requests[msg.Id]; ok && request != nil {
						logRequest = serialize(request.Params)
						logMethod = request.Method

						delete(requests, msg.Id)

						protocolLogger.WithFields(logrus.Fields{
							fieldType:    logType,
							fieldMethod:  logMethod,
							fieldRequest: logRequest,
						}).Info(logMessage)
					}
				}

				if msg.IsEvent() {
					protocolLogger.WithFields(logrus.Fields{
						fieldType:   typeEvent,
						fieldMethod: msg.Method,
					}).Info(serialize(msg.Params))
				}
			}
		}
	}
}

func checkVersion() (map[string]string, error) {
	cl := &http.Client{}
	req, err := http.NewRequest("GET", "http://"+*flagRemote+"/json/version", nil)
	if err != nil {
		return nil, err
	}

	res, err := cl.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	var v map[string]string
	if err := json.NewDecoder(res.Body).Decode(&v); err != nil {
		return nil, errors.New("expected json result")
	}

	return v, nil
}
<commit_msg>gofmt -s -w<commit_after>// This software is direct fork of https://github.com/knq/chromedp/tree/master/cmd/chromedp-proxy
// with couple of features added
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"

	"errors"

	"github.com/Sirupsen/logrus"
)

var (
	flagListen         = flag.String("l", "localhost:9223", "listen address")
	flagRemote         = flag.String("r", "localhost:9222", "remote address")
	flagEllipsis       = flag.Bool("s", false, "shorten requests and responses")
	flagOnce           = flag.Bool("once", false, "debug single session")
	flagShowRequests   = flag.Bool("i", false, "include request frames as they are sent")
	flagDistributeLogs = flag.Bool("d", false, "write logs file per targetId")
	flagQuiet          = flag.Bool("q", false, "do not show logs on stdout")
)

var protocolTargetID = center("protocol message", 36)

func main() {
	flag.Parse()
	mux := http.NewServeMux()

	simpleReverseProxy := httputil.NewSingleHostReverseProxy(&url.URL{Scheme: "http", Host: *flagRemote})

	mux.Handle("/json", simpleReverseProxy)
	mux.Handle("/", simpleReverseProxy)

	rootLogger, err := createLogger("connection")
	if err != nil {
		panic(fmt.Sprintf("could not create logger: %s", err))
	}

	logger := rootLogger.WithFields(logrus.Fields{
		fieldLevel: levelConnection,
	})

	mux.HandleFunc("/devtools/page/", func(res http.ResponseWriter, req *http.Request) {

		stream := make(chan *protocolMessage, 1024)
		id := path.Base(req.URL.Path)

		var protocolLogger *logrus.Entry

		if *flagDistributeLogs {
			logger, err := createLogger("inspector-" + id)
			if err != nil {
				panic(fmt.Sprintf("could not create logger: %s", err))
			}

			protocolLogger = logger.WithFields(logrus.Fields{
				fieldLevel:       levelConnection,
				fieldInspectorId: id,
			})

		} else {
			protocolLogger = logger.WithFields(logrus.Fields{
				fieldInspectorId: id,
			})
		}

		go dumpStream(protocolLogger, stream)

		endpoint := "ws://" + *flagRemote + "/devtools/page/" + id

		logger.Infof("---------- connection from %s ----------", req.RemoteAddr)
		logger.Infof("checking protocol versions on: %s", endpoint)

		ver, err := checkVersion()
		if err != nil {
			protocolLogger.Errorf("could not check version: %v", err)
			http.Error(res, "could not check version", 500)
			return
		}

		logger.Infof("protocol version: %s", ver["Protocol-Version"])
		logger.Infof("versions: Chrome(%s), V8(%s), Webkit(%s)", ver["Browser"], ver["V8-Version"], ver["WebKit-Version"])
		logger.Infof("browser user agent: %s", ver["User-Agent"])
		logger.Infof("connecting to %s... ", endpoint)

		// connecting to ws
		out, pres, err := wsDialer.Dial(endpoint, nil)
		if err != nil {
			msg := fmt.Sprintf("could not connect to %s: %v", endpoint, err)
			logger.Error(protocolError(msg))
			http.Error(res, msg, 500)
			return
		}
		defer pres.Body.Close()
		defer out.Close()

		// connect incoming websocket
		logger.Infof("upgrading connection on %s...", req.RemoteAddr)
		in, err := wsUpgrader.Upgrade(res, req, nil)
		if err != nil {
			logger.Errorf("could not upgrade websocket from %s: %v", req.RemoteAddr, err)
			http.Error(res, "could not upgrade websocket connection", 500)
			return
		}
		defer in.Close()

		ctxt, cancel := context.WithCancel(context.Background())
		defer cancel()

		errc := make(chan error, 1)
		go proxyWS(ctxt, stream, in, out, errc)
		go proxyWS(ctxt, stream, out, in, errc)

		<-errc
		logger.Infof("---------- closing %s ----------", req.RemoteAddr)

		if *flagOnce {
			os.Exit(0)
		}
	})

	log.Fatal(http.ListenAndServe(*flagListen, mux))
}

func dumpStream(logger *logrus.Entry, stream chan *protocolMessage) {
	logger.Printf("Legend: %s, %s, %s, %s, %s, %s", protocolColor("protocol informations"),
		eventsColor("received events"),
		requestColor("sent request frames"),
		requestReplyColor("requests params"),
		responseColor("received responses"),
		errorColor("error response."),
	)

	requests := make(map[uint64]*protocolMessage)
	targetRequests := make(map[uint64]*protocolMessage)

	for {
		select {
		case msg := <-stream:
			if msg.InTarget() {

				var targetLogger *logrus.Entry

				if *flagDistributeLogs {
					logger, err := createLogger(fmt.Sprintf("target-%s", msg.Params["targetId"]))
					if err != nil {
						panic(fmt.Sprintf("could not create logger: %v", err))
					}

					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetId: msg.Params["targetId"],
					})

				} else {
					targetLogger = logger.WithFields(logrus.Fields{
						fieldLevel:    levelTarget,
						fieldTargetId: msg.Params["targetId"],
					})
				}

				if msg.IsRequest() {
					requests[msg.Id] = nil

					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						targetRequests[protocolMessage.Id] = protocolMessage

						if *flagShowRequests {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeRequest,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Error("Could not deserialize message: %+v", err)
					}
				}

				if msg.IsEvent() {
					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						if protocolMessage.IsEvent() {
							targetLogger.WithFields(logrus.Fields{
								fieldType:   typeEvent,
								fieldMethod: protocolMessage.Method,
							}).Info(serialize(protocolMessage.Params))
						}

						if protocolMessage.IsResponse() {
							var logMessage string
							var logType int
							var logRequest string
							var logMethod string

							if protocolMessage.IsError() {
								logMessage = serialize(protocolMessage.Error)
								logType = typeRequestResponseError
							} else {
								logMessage = serialize(protocolMessage.Result)
								logType = typeRequestResponse
							}

							if request, ok := targetRequests[protocolMessage.Id]; ok && request != nil {
								delete(targetRequests, protocolMessage.Id)
								logRequest = serialize(request.Params)
								logMethod = request.Method

							} else {
								logRequest = errorColor("could not find request with id: %d", protocolMessage.Id)
							}

							targetLogger.WithFields(logrus.Fields{
								fieldType:    logType,
								fieldMethod:  logMethod,
								fieldRequest: logRequest,
							}).Info(logMessage)
						}
					} else {
						logger.WithFields(logrus.Fields{
							fieldLevel: levelConnection,
						}).Error("Could not deserialize message: %+v", err)
					}
				}

			} else {
				protocolLogger := logger.WithFields(logrus.Fields{
					fieldLevel:    levelProtocol,
					fieldTargetId: protocolTargetID,
				})

				if msg.IsRequest() {
					requests[msg.Id] = msg

					if *flagShowRequests {
						protocolLogger.WithFields(logrus.Fields{
							fieldType:   typeRequest,
							fieldMethod: msg.Method,
						}).Info(serialize(msg.Params))
					}
				}

				if msg.IsResponse() {

					var logMessage string
					var logType int
					var logRequest string
					var logMethod string

					if msg.IsError() {
						logMessage = serialize(msg.Error)
						logType = typeRequestResponseError
					} else {
						logMessage = serialize(msg.Result)
						logType = typeRequestResponse
					}

					if request, ok := requests[msg.Id]; ok && request != nil {
						logRequest = serialize(request.Params)
						logMethod = request.Method

						delete(requests, msg.Id)

						protocolLogger.WithFields(logrus.Fields{
							fieldType:    logType,
							fieldMethod:  logMethod,
							fieldRequest: logRequest,
						}).Info(logMessage)
					}
				}

				if msg.IsEvent() {
					protocolLogger.WithFields(logrus.Fields{
						fieldType:   typeEvent,
						fieldMethod: msg.Method,
					}).Info(serialize(msg.Params))
				}
			}
		}
	}
}

func checkVersion() (map[string]string, error) {
	cl := &http.Client{}
	req, err := http.NewRequest("GET", "http://"+*flagRemote+"/json/version", nil)
	if err != nil {
		return nil, err
	}

	res, err := cl.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	var v map[string]string
	if err := json.NewDecoder(res.Body).Decode(&v); err != nil {
		return nil, errors.New("expected json result")
	}

	return v, nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/binary"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"

	"github.com/calmh/syncthing/discover"
	"github.com/golang/groupcache/lru"
	"github.com/juju/ratelimit"
)

type node struct {
	addresses []address
	updated   time.Time
}

type address struct {
	ip   []byte
	port uint16
}

var (
	nodes     = make(map[string]node)
	lock      sync.Mutex
	queries   = 0
	announces = 0
	answered  = 0
	limited   = 0
	unknowns  = 0
	debug     = false
	limiter   = lru.New(1024)
)

func main() {
	var listen string
	var timestamp bool
	var statsIntv int
	var statsFile string

	flag.StringVar(&listen, "listen", ":22025", "Listen address")
	flag.BoolVar(&debug, "debug", false, "Enable debug output")
	flag.BoolVar(&timestamp, "timestamp", true, "Timestamp the log output")
	flag.IntVar(&statsIntv, "stats-intv", 0, "Statistics output interval (s)")
	flag.StringVar(&statsFile, "stats-file", "/var/log/discosrv.stats", "Statistics file name")
	flag.Parse()

	log.SetOutput(os.Stdout)
	if !timestamp {
		log.SetFlags(0)
	}

	addr, _ := net.ResolveUDPAddr("udp", listen)
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Fatal(err)
	}

	if statsIntv > 0 {
		go logStats(statsFile, statsIntv)
	}

	var buf = make([]byte, 1024)
	for {
		buf = buf[:cap(buf)]
		n, addr, err := conn.ReadFromUDP(buf)

		if limit(addr) {
			// Rate limit in effect for source
			continue
		}

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

		if n < 4 {
			log.Printf("Received short packet (%d bytes)", n)
			continue
		}

		buf = buf[:n]
		magic := binary.BigEndian.Uint32(buf)

		switch magic {
		case discover.AnnouncementMagicV2:
			handleAnnounceV2(addr, buf)

		case discover.QueryMagicV2:
			handleQueryV2(conn, addr, buf)

		default:
			lock.Lock()
			unknowns++
			lock.Unlock()
		}
	}
}

func limit(addr *net.UDPAddr) bool {
	key := addr.IP.String()

	lock.Lock()
	defer lock.Unlock()

	bkt, ok := limiter.Get(key)
	if ok {
		bkt := bkt.(*ratelimit.Bucket)
		if bkt.TakeAvailable(1) != 1 {
			// Rate limit exceeded; ignore packet
			if debug {
				log.Println("Rate limit exceeded for", key)
			}
			limited++
			return true
		}
	} else {
		if debug {
			log.Println("New limiter for", key)
		}
		// One packet per ten seconds average rate, burst ten packets
		limiter.Add(key, ratelimit.NewBucket(10*time.Second, 10))
	}

	return false
}

func handleAnnounceV2(addr *net.UDPAddr, buf []byte) {
	var pkt discover.AnnounceV2
	err := pkt.UnmarshalXDR(buf)
	if err != nil && err != io.EOF {
		log.Println("AnnounceV2 Unmarshal:", err)
		log.Println(hex.Dump(buf))
		return
	}
	if debug {
		log.Printf("<- %v %#v", addr, pkt)
	}

	lock.Lock()
	announces++
	lock.Unlock()

	ip := addr.IP.To4()
	if ip == nil {
		ip = addr.IP.To16()
	}

	var addrs []address
	for _, addr := range pkt.This.Addresses {
		tip := addr.IP
		if len(tip) == 0 {
			tip = ip
		}
		addrs = append(addrs, address{
			ip:   tip,
			port: addr.Port,
		})
	}

	node := node{
		addresses: addrs,
		updated:   time.Now(),
	}

	lock.Lock()
	nodes[pkt.This.ID] = node
	lock.Unlock()
}

func handleQueryV2(conn *net.UDPConn, addr *net.UDPAddr, buf []byte) {
	var pkt discover.QueryV2
	err := pkt.UnmarshalXDR(buf)
	if err != nil {
		log.Println("QueryV2 Unmarshal:", err)
		log.Println(hex.Dump(buf))
		return
	}
	if debug {
		log.Printf("<- %v %#v", addr, pkt)
	}

	lock.Lock()
	node, ok := nodes[pkt.NodeID]
	queries++
	lock.Unlock()

	if ok && len(node.addresses) > 0 {
		ann := discover.AnnounceV2{
			Magic: discover.AnnouncementMagicV2,
			This: discover.Node{
				ID: pkt.NodeID,
			},
		}
		for _, addr := range node.addresses {
			ann.This.Addresses = append(ann.This.Addresses, discover.Address{IP: addr.ip, Port: addr.port})
		}
		if debug {
			log.Printf("-> %v %#v", addr, pkt)
		}

		tb := ann.MarshalXDR()
		_, _, err = conn.WriteMsgUDP(tb, nil, addr)
		if err != nil {
			log.Println("QueryV2 response write:", err)
		}

		lock.Lock()
		answered++
		lock.Unlock()
	}
}

func next(intv int) time.Time {
	d := time.Duration(intv) * time.Second
	t0 := time.Now()
	t1 := t0.Add(d).Truncate(d)
	time.Sleep(t1.Sub(t0))
	return t1
}

func logStats(file string, intv int) {
	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		log.Fatal(err)
	}
	for {
		t := next(intv)

		lock.Lock()

		var deleted = 0
		for id, node := range nodes {
			if time.Since(node.updated) > 60*time.Minute {
				delete(nodes, id)
				deleted++
			}
		}

		fmt.Fprintf(f, "%d Nr:%d Ne:%d Qt:%d Qa:%d A:%d U:%d Lq:%d Lc:%d\n",
			t.Unix(), len(nodes), deleted, queries, answered, announces, unknowns, limited, limiter.Len())
		f.Sync()

		queries = 0
		announces = 0
		answered = 0
		limited = 0
		unknowns = 0

		lock.Unlock()
	}
}
<commit_msg>Add license header<commit_after>// Copyright (C) 2014 Jakob Borg and other contributors. All rights reserved.
// Use of this source code is governed by an MIT-style license that can be
// found in the LICENSE file.

package main

import (
	"encoding/binary"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"

	"github.com/calmh/syncthing/discover"
	"github.com/golang/groupcache/lru"
	"github.com/juju/ratelimit"
)

type node struct {
	addresses []address
	updated   time.Time
}

type address struct {
	ip   []byte
	port uint16
}

var (
	nodes     = make(map[string]node)
	lock      sync.Mutex
	queries   = 0
	announces = 0
	answered  = 0
	limited   = 0
	unknowns  = 0
	debug     = false
	limiter   = lru.New(1024)
)

func main() {
	var listen string
	var timestamp bool
	var statsIntv int
	var statsFile string

	flag.StringVar(&listen, "listen", ":22025", "Listen address")
	flag.BoolVar(&debug, "debug", false, "Enable debug output")
	flag.BoolVar(&timestamp, "timestamp", true, "Timestamp the log output")
	flag.IntVar(&statsIntv, "stats-intv", 0, "Statistics output interval (s)")
	flag.StringVar(&statsFile, "stats-file", "/var/log/discosrv.stats", "Statistics file name")
	flag.Parse()

	log.SetOutput(os.Stdout)
	if !timestamp {
		log.SetFlags(0)
	}

	addr, _ := net.ResolveUDPAddr("udp", listen)
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		log.Fatal(err)
	}

	if statsIntv > 0 {
		go logStats(statsFile, statsIntv)
	}

	var buf = make([]byte, 1024)
	for {
		buf = buf[:cap(buf)]
		n, addr, err := conn.ReadFromUDP(buf)

		if limit(addr) {
			// Rate limit in effect for source
			continue
		}

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

		if n < 4 {
			log.Printf("Received short packet (%d bytes)", n)
			continue
		}

		buf = buf[:n]
		magic := binary.BigEndian.Uint32(buf)

		switch magic {
		case discover.AnnouncementMagicV2:
			handleAnnounceV2(addr, buf)

		case discover.QueryMagicV2:
			handleQueryV2(conn, addr, buf)

		default:
			lock.Lock()
			unknowns++
			lock.Unlock()
		}
	}
}

func limit(addr *net.UDPAddr) bool {
	key := addr.IP.String()

	lock.Lock()
	defer lock.Unlock()

	bkt, ok := limiter.Get(key)
	if ok {
		bkt := bkt.(*ratelimit.Bucket)
		if bkt.TakeAvailable(1) != 1 {
			// Rate limit exceeded; ignore packet
			if debug {
				log.Println("Rate limit exceeded for", key)
			}
			limited++
			return true
		}
	} else {
		if debug {
			log.Println("New limiter for", key)
		}
		// One packet per ten seconds average rate, burst ten packets
		limiter.Add(key, ratelimit.NewBucket(10*time.Second, 10))
	}

	return false
}

func handleAnnounceV2(addr *net.UDPAddr, buf []byte) {
	var pkt discover.AnnounceV2
	err := pkt.UnmarshalXDR(buf)
	if err != nil && err != io.EOF {
		log.Println("AnnounceV2 Unmarshal:", err)
		log.Println(hex.Dump(buf))
		return
	}
	if debug {
		log.Printf("<- %v %#v", addr, pkt)
	}

	lock.Lock()
	announces++
	lock.Unlock()

	ip := addr.IP.To4()
	if ip == nil {
		ip = addr.IP.To16()
	}

	var addrs []address
	for _, addr := range pkt.This.Addresses {
		tip := addr.IP
		if len(tip) == 0 {
			tip = ip
		}
		addrs = append(addrs, address{
			ip:   tip,
			port: addr.Port,
		})
	}

	node := node{
		addresses: addrs,
		updated:   time.Now(),
	}

	lock.Lock()
	nodes[pkt.This.ID] = node
	lock.Unlock()
}

func handleQueryV2(conn *net.UDPConn, addr *net.UDPAddr, buf []byte) {
	var pkt discover.QueryV2
	err := pkt.UnmarshalXDR(buf)
	if err != nil {
		log.Println("QueryV2 Unmarshal:", err)
		log.Println(hex.Dump(buf))
		return
	}
	if debug {
		log.Printf("<- %v %#v", addr, pkt)
	}

	lock.Lock()
	node, ok := nodes[pkt.NodeID]
	queries++
	lock.Unlock()

	if ok && len(node.addresses) > 0 {
		ann := discover.AnnounceV2{
			Magic: discover.AnnouncementMagicV2,
			This: discover.Node{
				ID: pkt.NodeID,
			},
		}
		for _, addr := range node.addresses {
			ann.This.Addresses = append(ann.This.Addresses, discover.Address{IP: addr.ip, Port: addr.port})
		}
		if debug {
			log.Printf("-> %v %#v", addr, pkt)
		}

		tb := ann.MarshalXDR()
		_, _, err = conn.WriteMsgUDP(tb, nil, addr)
		if err != nil {
			log.Println("QueryV2 response write:", err)
		}

		lock.Lock()
		answered++
		lock.Unlock()
	}
}

func next(intv int) time.Time {
	d := time.Duration(intv) * time.Second
	t0 := time.Now()
	t1 := t0.Add(d).Truncate(d)
	time.Sleep(t1.Sub(t0))
	return t1
}

func logStats(file string, intv int) {
	f, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		log.Fatal(err)
	}
	for {
		t := next(intv)

		lock.Lock()

		var deleted = 0
		for id, node := range nodes {
			if time.Since(node.updated) > 60*time.Minute {
				delete(nodes, id)
				deleted++
			}
		}

		fmt.Fprintf(f, "%d Nr:%d Ne:%d Qt:%d Qa:%d A:%d U:%d Lq:%d Lc:%d\n",
			t.Unix(), len(nodes), deleted, queries, answered, announces, unknowns, limited, limiter.Len())
		f.Sync()

		queries = 0
		announces = 0
		answered = 0
		limited = 0
		unknowns = 0

		lock.Unlock()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/newrelic/sidecar/catalog"
	"github.com/newrelic/sidecar/haproxy"
	"gopkg.in/alecthomas/kingpin.v1"
)

const (
	RELOAD_BUFFER = 256
	// A new service usually comes in as three events.
	// By 5 seconds it's usually alive.
	RELOAD_HOLD_DOWN = 5 * time.Second
)

var (
	proxy         *haproxy.HAproxy
	stateLock     sync.Mutex
	reloadChan    chan time.Time
	currentState  *catalog.ServicesState
	updateSuccess bool
)

type CliOpts struct {
	ConfigFile *string
}

type ApiErrors struct {
	Errors []string `json:"errors"`
}

type ApiStatus struct {
	Message     string `json:"message"`
	LastChanged string `json:"last_changed"`
}

func exitWithError(err error, message string) {
	if err != nil {
		log.Fatal("%s: %s", message, err.Error())
	}
}

func parseCommandLine() *CliOpts {
	var opts CliOpts
	opts.ConfigFile = kingpin.Flag("config-file", "The config file to use").Short('f').Default("haproxy-api.toml").String()
	kingpin.Parse()
	return &opts
}

func run(command string) error {
	cmd := exec.Command("/bin/bash", "-c", command)
	err := cmd.Run()
	if err != nil {
		log.Errorf("Error running '%s': %s", command, err.Error())
	}

	return err
}

func healthHandler(response http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	response.Header().Set("Content-Type", "application/json")

	errors := make([]string, 0)

	err := run("test -f " + proxy.PidFile + " && ps aux `cat " + proxy.PidFile + "`")
	if err != nil {
		errors = append(errors, "No HAproxy running!")
	}

	stateLock.Lock()
	defer stateLock.Unlock()

	if updateSuccess == false {
		errors = append(errors, "Last attempted HAproxy config write failed!")
	}

	if errors != nil && len(errors) != 0 {
		message, _ := json.Marshal(ApiErrors{errors})
		response.WriteHeader(http.StatusInternalServerError)
		response.Write(message)
		return
	}

	var lastChanged time.Time
	if currentState != nil {
		lastChanged = currentState.LastChanged
	}

	message, _ := json.Marshal(ApiStatus{Message: "Healthy!", LastChanged: lastChanged.String()})
	response.Write(message)
}

func stateHandler(response http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	response.Header().Set("Content-Type", "application/json")

	if currentState == nil {
		message, _ := json.Marshal(ApiErrors{[]string{"No currently stored state"}})
		response.WriteHeader(http.StatusInternalServerError)
		response.Write(message)
		return
	}

	response.Write(currentState.Encode())
}

func updateHandler(response http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	response.Header().Set("Content-Type", "application/json")

	bytes, err := ioutil.ReadAll(req.Body)
	if err != nil {
		message, _ := json.Marshal(ApiErrors{[]string{err.Error()}})
		response.WriteHeader(http.StatusInternalServerError)
		response.Write(message)
		return
	}

	state, err := catalog.Decode(bytes)
	if err != nil {
		response.WriteHeader(http.StatusInternalServerError)
		return
	}

	updateState(state)
}

func processUpdates() {
	for {
		// Batch up to RELOAD_BUFFER number updates into a
		// single update.
		first := <-reloadChan
		pending := len(reloadChan)

		writeAndReload(currentState)

		// We just flushed the most recent state, dump all the
		// pending items up to that point.
		var reload time.Time
		for i := 0; i < pending; i++ {
			reload = <-reloadChan
		}

		if first.Before(reload) {
			log.Infof("Skipped %d messages between %s and %s", pending, first, reload)
		}

		// Don't notify more frequently than every RELOAD_HOLD_DOWN period. When a
		// deployment rolls across the cluster it can trigger a bunch of groupable
		// updates.
		log.Debug("Holding down...")
		time.Sleep(RELOAD_HOLD_DOWN)
	}
}

func writeAndReload(state *catalog.ServicesState) {
	log.Info("Updating HAproxy")
	err := proxy.WriteAndReload(state)
	updateSuccess = (err == nil)
}

func updateState(state *catalog.ServicesState) {
	stateLock.Lock()
	defer stateLock.Unlock()
	currentState = state
	reloadChan <- time.Now().UTC()
}

func fetchState(url string) error {
	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Get(url)
	if err != nil {
		return err
	}

	bytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	state, err := catalog.Decode(bytes)
	if err != nil {
		return err
	}

	currentState = state
	writeAndReload(state)

	return nil
}

func serveHttp(listenIp string, listenPort int) {
	listenStr := fmt.Sprintf("%s:%d", listenIp, listenPort)

	log.Infof("Starting up on %s", listenStr)
	router := mux.NewRouter()

	router.HandleFunc("/update", updateHandler).Methods("POST")
	router.HandleFunc("/health", healthHandler).Methods("GET")
	router.HandleFunc("/state", stateHandler).Methods("GET")
	http.Handle("/", handlers.LoggingHandler(os.Stdout, router))

	err := http.ListenAndServe(listenStr, nil)
	if err != nil {
		log.Fatalf("Can't start http server: %s", err.Error())
	}
}

func main() {
	opts := parseCommandLine()
	config := parseConfig(*opts.ConfigFile)

	proxy = config.HAproxy

	reloadChan = make(chan time.Time, RELOAD_BUFFER)

	log.Info("Fetching initial state on startup...")
	err := fetchState(config.Sidecar.StateUrl)
	if err != nil {
		log.Errorf("Failed to fetch state from '%s'... continuing in hopes someone will post it", config.Sidecar.StateUrl)
	} else {
		log.Info("Successfully retrieved state")
	}

	go processUpdates()

	serveHttp(config.HAproxyApi.BindIP, config.HAproxyApi.BindPort)

	close(reloadChan)
}
<commit_msg>Lots of comments. lastChanged -> time.Time.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/newrelic/sidecar/catalog"
	"github.com/newrelic/sidecar/haproxy"
	"gopkg.in/alecthomas/kingpin.v1"
)

const (
	RELOAD_BUFFER = 256
	// A new service usually comes in as three events.
	// By 5 seconds it's usually alive.
	RELOAD_HOLD_DOWN = 5 * time.Second
)

var (
	proxy         *haproxy.HAproxy
	stateLock     sync.Mutex
	reloadChan    chan time.Time
	currentState  *catalog.ServicesState
	updateSuccess bool
)

type CliOpts struct {
	ConfigFile *string
}

type ApiErrors struct {
	Errors []string `json:"errors"`
}

type ApiStatus struct {
	Message     string    `json:"message"`
	LastChanged time.Time `json:"last_changed"`
}

func exitWithError(err error, message string) {
	if err != nil {
		log.Fatal("%s: %s", message, err.Error())
	}
}

func parseCommandLine() *CliOpts {
	var opts CliOpts
	opts.ConfigFile = kingpin.Flag("config-file", "The config file to use").Short('f').Default("haproxy-api.toml").String()
	kingpin.Parse()
	return &opts
}

func run(command string) error {
	cmd := exec.Command("/bin/bash", "-c", command)
	err := cmd.Run()
	if err != nil {
		log.Errorf("Error running '%s': %s", command, err.Error())
	}

	return err
}

// The health check endpoint. Tells us if HAproxy is running and has
// been properly configured. Since this is critical infrastructure this
// helps make sure a host is not "down" by havign the proxy down.
func healthHandler(response http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	response.Header().Set("Content-Type", "application/json")

	errors := make([]string, 0)

	// Do we have an HAproxy instance running?
	err := run("test -f " + proxy.PidFile + " && ps aux `cat " + proxy.PidFile + "`")
	if err != nil {
		errors = append(errors, "No HAproxy running!")
	}

	stateLock.Lock()
	defer stateLock.Unlock()

	// We were able to write out the template and reload the last time we tried?
	if updateSuccess == false {
		errors = append(errors, "Last attempted HAproxy config write failed!")
	}

	// Umm, crap, something went wrong.
	if errors != nil && len(errors) != 0 {
		message, _ := json.Marshal(ApiErrors{errors})
		response.WriteHeader(http.StatusInternalServerError)
		response.Write(message)
		return
	}

	var lastChanged time.Time
	if currentState != nil {
		lastChanged = currentState.LastChanged
	}

	message, _ := json.Marshal(ApiStatus{Message: "Healthy!", LastChanged: lastChanged})
	response.Write(message)
}

func stateHandler(response http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	response.Header().Set("Content-Type", "application/json")

	if currentState == nil {
		message, _ := json.Marshal(ApiErrors{[]string{"No currently stored state"}})
		response.WriteHeader(http.StatusInternalServerError)
		response.Write(message)
		return
	}

	response.Write(currentState.Encode())
}

func updateHandler(response http.ResponseWriter, req *http.Request) {
	defer req.Body.Close()
	response.Header().Set("Content-Type", "application/json")

	bytes, err := ioutil.ReadAll(req.Body)
	if err != nil {
		message, _ := json.Marshal(ApiErrors{[]string{err.Error()}})
		response.WriteHeader(http.StatusInternalServerError)
		response.Write(message)
		return
	}

	state, err := catalog.Decode(bytes)
	if err != nil {
		response.WriteHeader(http.StatusInternalServerError)
		return
	}

	updateState(state)
}

// Loop forever, processing updates to the state.
func processUpdates() {
	for {
		// Batch up to RELOAD_BUFFER number updates into a
		// single update.
		first := <-reloadChan
		pending := len(reloadChan)

		writeAndReload(currentState)

		// We just flushed the most recent state, dump all the
		// pending items up to that point.
		var reload time.Time
		for i := 0; i < pending; i++ {
			reload = <-reloadChan
		}

		if first.Before(reload) {
			log.Infof("Skipped %d messages between %s and %s", pending, first, reload)
		}

		// Don't notify more frequently than every RELOAD_HOLD_DOWN period. When a
		// deployment rolls across the cluster it can trigger a bunch of groupable
		// updates.
		log.Debug("Holding down...")
		time.Sleep(RELOAD_HOLD_DOWN)
	}
}

// Write out the HAproxy config and reload the instance
func writeAndReload(state *catalog.ServicesState) {
	log.Info("Updating HAproxy")
	err := proxy.WriteAndReload(state)
	updateSuccess = (err == nil)
}

// Process and update by setting the current state and queueing
// a writeAndReload().
func updateState(state *catalog.ServicesState) {
	stateLock.Lock()
	defer stateLock.Unlock()
	currentState = state
	reloadChan <- time.Now().UTC()
}

// Used to fetch the current state from a Sidecar endpoint, usually
// on startup of this process, when the currentState is empty.
func fetchState(url string) error {
	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Get(url)
	if err != nil {
		return err
	}

	bytes, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	state, err := catalog.Decode(bytes)
	if err != nil {
		return err
	}

	currentState = state
	writeAndReload(state)

	return nil
}

// Start the HTTP server and begin handling requests. This is a
// blocking call.
func serveHttp(listenIp string, listenPort int) {
	listenStr := fmt.Sprintf("%s:%d", listenIp, listenPort)

	log.Infof("Starting up on %s", listenStr)
	router := mux.NewRouter()

	router.HandleFunc("/update", updateHandler).Methods("POST")
	router.HandleFunc("/health", healthHandler).Methods("GET")
	router.HandleFunc("/state", stateHandler).Methods("GET")
	http.Handle("/", handlers.LoggingHandler(os.Stdout, router))

	err := http.ListenAndServe(listenStr, nil)
	if err != nil {
		log.Fatalf("Can't start http server: %s", err.Error())
	}
}

func main() {
	opts := parseCommandLine()
	config := parseConfig(*opts.ConfigFile)

	proxy = config.HAproxy

	reloadChan = make(chan time.Time, RELOAD_BUFFER)

	log.Info("Fetching initial state on startup...")
	err := fetchState(config.Sidecar.StateUrl)
	if err != nil {
		log.Errorf("Failed to fetch state from '%s'... continuing in hopes someone will post it", config.Sidecar.StateUrl)
	} else {
		log.Info("Successfully retrieved state")
	}

	go processUpdates()

	serveHttp(config.HAproxyApi.BindIP, config.HAproxyApi.BindPort)

	close(reloadChan)
}
<|endoftext|>


<commit_before>package main

import (
	"time"

	"github.com/gin-gonic/gin"
)

type Host struct {
	Hostname   string    `json:"hostname"`
	Expiration time.Time `json:"expiration"`
}

type Status struct {
	Hosts  []Host  `json:"hosts"`
	Uptime float64 `json:"uptime"`
}

var (
	hosts []Host
	TTL   = 30
)

func buildExpiration() time.Time {
	expiration := time.Now().UTC()
	expiration = expiration.Add(time.Duration(TTL) * time.Second)

	return expiration
}

func cleanup() {
	for {
		for i, h := range hosts {
			if time.Now().UTC().After(h.Expiration) {
				diff := hosts
				diff = append(diff[:i], diff[i+1:]...)
				hosts = diff
			}
		}
		time.Sleep(time.Duration(TTL) * time.Second)
	}
}

func main() {
	boot := time.Now()

	go cleanup()

	router := gin.Default()

	router.POST("/checkin/", func(c *gin.Context) {
		var host Host

		c.Bind(&host)

		if host.Hostname != "" {
			exists := false

			for i, h := range hosts {
				if h.Hostname == host.Hostname {
					exists = true
					hosts[i].Expiration = buildExpiration()
				}
			}

			if !exists {
				host.Expiration = buildExpiration()
				hosts = append(hosts, host)
			}

			c.JSON(200, "")
		} else {
			var response struct {
				Message string `json:"message"`
			}
			response.Message = "Invalid hostname."
			c.JSON(400, response)
		}
	})

	router.GET("/status/", func(c *gin.Context) {
		status := Status{}
		status.Hosts = hosts
		status.Uptime = time.Now().Sub(boot).Seconds()

		c.JSON(200, status)
	})

	router.Run(":8080")
}
<commit_msg>Defined a TravisNotification struct which represents a webhook payload<commit_after>package main

import (
	"time"

	"github.com/gin-gonic/gin"
)

type Host struct {
	Hostname   string    `json:"hostname"`
	Expiration time.Time `json:"expiration"`
}

type Status struct {
	Hosts  []Host  `json:"hosts"`
	Uptime float64 `json:"uptime"`
}

type TravisNotification struct {
	Payload TravisPayload `json:"payload"`
}

type TravisPayload struct {
	Status     string           `json:"status_message"`
	Commit     string           `json:"commit"`
	Branch     string           `json:"branch"`
	Message    string           `json:"message"`
	Repository TravisRepository `json:"repository"`
}

type TravisRepository struct {
	Name  string `json:"name"`
	Owner string `json:"owner_name"`
}

var (
	hosts []Host
	TTL   = 30
)

func buildExpiration() time.Time {
	expiration := time.Now().UTC()
	expiration = expiration.Add(time.Duration(TTL) * time.Second)

	return expiration
}

func cleanup() {
	for {
		for i, h := range hosts {
			if time.Now().UTC().After(h.Expiration) {
				diff := hosts
				diff = append(diff[:i], diff[i+1:]...)
				hosts = diff
			}
		}
		time.Sleep(time.Duration(TTL) * time.Second)
	}
}

func main() {
	boot := time.Now()

	go cleanup()

	router := gin.Default()

	router.POST("/checkin/", func(c *gin.Context) {
		var host Host

		c.Bind(&host)

		if host.Hostname != "" {
			exists := false

			for i, h := range hosts {
				if h.Hostname == host.Hostname {
					exists = true
					hosts[i].Expiration = buildExpiration()
				}
			}

			if !exists {
				host.Expiration = buildExpiration()
				hosts = append(hosts, host)
			}

			c.JSON(200, "")
		} else {
			var response struct {
				Message string `json:"message"`
			}
			response.Message = "Invalid hostname."
			c.JSON(400, response)
		}
	})

	router.GET("/status/", func(c *gin.Context) {
		status := Status{}
		status.Hosts = hosts
		status.Uptime = time.Now().Sub(boot).Seconds()

		c.JSON(200, status)
	})

	router.Run(":8080")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	log "github.com/Sirupsen/logrus"
	"io/ioutil"
	"strings"
	"os"
	"os/exec"
	"strconv"
	"sync"
	"html/template"
	"github.com/b00lduck/raspberry_soundboard/templates"
	"math/rand"
)

var mutex = &sync.Mutex{}

type Sound struct {
	SoundFile string
	ImageFile string
	HasImage  bool
	Count     int
}

func handler(w http.ResponseWriter, r *http.Request) {

	log.WithField("requestURi", r.RequestURI).Info("HTTP")

	if strings.HasPrefix(r.RequestURI, "/play") {
		handlePlay(w, r)
	} else if strings.HasPrefix(r.RequestURI, "/random") {
		handleRandomPlay(w, r)
	} else if strings.HasPrefix(r.RequestURI, "/images") {
		handleImage(w, r)
	} else {
		handleIndex(w, r)
	}

}

func handleImage(w http.ResponseWriter, r *http.Request) {
	filename := "sounds/" + r.RequestURI[8:]
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		log.Error(filename)
		w.WriteHeader(404)
		return
	}
	image, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Error(err)
		w.WriteHeader(500)
	}
	w.Write(image)
}

func handleRandomPlay(w http.ResponseWriter, r *http.Request) {

	sounds := getSounds()

	filename := "sounds/" + sounds[rand.Intn(len(sounds))].SoundFile

	err := play(filename)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	http.Redirect(w, r, "/", 307)
}

func handlePlay(w http.ResponseWriter, r *http.Request) {

	filename := "sounds/" + r.RequestURI[6:]

	err := play(filename)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	http.Redirect(w, r, "/", 307)

}

func play(filename string) error {
	if strings.HasSuffix(filename, ".mp3") {

		log.WithField("filename", filename).Info("playing sound")
		if _, err := os.Stat(filename); os.IsNotExist(err) {
			log.Error(filename)
			return fmt.Errorf("Not found")

		}
		go func() {
			cmd := exec.Command("omxplayer", "-o", "hdmi", filename)
			err := cmd.Run()
			if err != nil {
				log.Error(err)
			} else {
				incCounter(filename)
			}
		}()
	} else {
		return fmt.Errorf("no .mp3 suffix")
	}
	return nil
}

func incCounter(filename string) {

	log.WithField("filename", filename).Info("Increasing counter")

	mutex.Lock()
	intCount := getCounter(filename)
	log.WithField("count", intCount).Info("Old count")
	intCount++

	ioutil.WriteFile(filename + ".count", []byte(fmt.Sprintf("%d", intCount)), 0644)
	mutex.Unlock()
}

func getCounter(filename string) int {

	countfile := filename + ".count"

	if _, err := os.Stat(countfile); os.IsNotExist(err) {
		return 0
	}

	count, err := ioutil.ReadFile(countfile)
	intCount, err := strconv.Atoi(string(count))
	if err != nil {
		log.Error(err)
	}
	return intCount
}

func handleIndex(w http.ResponseWriter, r *http.Request) {

	sounds := getSounds()

	t := template.New("html")
	t, err := t.Parse(templates.MainTemplate)
	if err != nil {
		log.Error(err)
	}

	err = t.Execute(w, sounds)
	if err != nil {
		log.Error(err)
	}

}

func getSounds() []Sound {
	sounds := make([]Sound, 0)

	dir, err := ioutil.ReadDir("sounds")
	if err != nil {
		log.Error(err)
	}

	for _, v := range dir {
		if !v.IsDir() {
			filename := v.Name()
			if (strings.HasSuffix(filename, ".mp3")) {
				newSound := Sound{
					SoundFile: filename,
					HasImage: true,
					Count: getCounter("sounds/" + filename),
				}
				filenameWithoutExt := filename[:len(filename)-4]
				pngFilename := filenameWithoutExt + ".png"
				if _, err := os.Stat("sounds/" + pngFilename); os.IsNotExist(err) {
					jpgFilename := filenameWithoutExt + ".jpg"
					if _, err := os.Stat("sounds/" + jpgFilename); os.IsNotExist(err) {
						newSound.HasImage = false
					} else {
						newSound.ImageFile = jpgFilename
					}
				} else {
					newSound.ImageFile = pngFilename
				}
				sounds = append(sounds, newSound)
			}
		}
	}

	return sounds
}

func main() {
	http.HandleFunc("/", handler)
	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		log.Fatal(err)
	}
}<commit_msg>sorting<commit_after>package main

import (
	"fmt"
	"net/http"
	log "github.com/Sirupsen/logrus"
	"io/ioutil"
	"strings"
	"os"
	"os/exec"
	"strconv"
	"sync"
	"html/template"
	"github.com/b00lduck/raspberry_soundboard/templates"
	"math/rand"
	"sort"
)

var mutex = &sync.Mutex{}

type Sound struct {
	SoundFile string
	ImageFile string
	HasImage  bool
	Count     int
}

func handler(w http.ResponseWriter, r *http.Request) {

	log.WithField("requestURi", r.RequestURI).Info("HTTP")

	if strings.HasPrefix(r.RequestURI, "/play") {
		handlePlay(w, r)
	} else if strings.HasPrefix(r.RequestURI, "/random") {
		handleRandomPlay(w, r)
	} else if strings.HasPrefix(r.RequestURI, "/images") {
		handleImage(w, r)
	} else {
		handleIndex(w, r)
	}

}

func handleImage(w http.ResponseWriter, r *http.Request) {
	filename := "sounds/" + r.RequestURI[8:]
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		log.Error(filename)
		w.WriteHeader(404)
		return
	}
	image, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Error(err)
		w.WriteHeader(500)
	}
	w.Write(image)
}

func handleRandomPlay(w http.ResponseWriter, r *http.Request) {

	sounds := getSounds()

	filename := "sounds/" + sounds[rand.Intn(len(sounds))].SoundFile

	err := play(filename)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	http.Redirect(w, r, "/", 307)
}

func handlePlay(w http.ResponseWriter, r *http.Request) {

	filename := "sounds/" + r.RequestURI[6:]

	err := play(filename)
	if err != nil {
		w.WriteHeader(500)
		w.Write([]byte(err.Error()))
		return
	}
	http.Redirect(w, r, "/", 307)

}

func play(filename string) error {
	if strings.HasSuffix(filename, ".mp3") {

		log.WithField("filename", filename).Info("playing sound")
		if _, err := os.Stat(filename); os.IsNotExist(err) {
			log.Error(filename)
			return fmt.Errorf("Not found")

		}
		go func() {
			cmd := exec.Command("omxplayer", "-o", "hdmi", filename)
			err := cmd.Run()
			if err != nil {
				log.Error(err)
			} else {
				incCounter(filename)
			}
		}()
	} else {
		return fmt.Errorf("no .mp3 suffix")
	}
	return nil
}

func incCounter(filename string) {

	log.WithField("filename", filename).Info("Increasing counter")

	mutex.Lock()
	intCount := getCounter(filename)
	log.WithField("count", intCount).Info("Old count")
	intCount++

	ioutil.WriteFile(filename + ".count", []byte(fmt.Sprintf("%d", intCount)), 0644)
	mutex.Unlock()
}

func getCounter(filename string) int {

	countfile := filename + ".count"

	if _, err := os.Stat(countfile); os.IsNotExist(err) {
		return 0
	}

	count, err := ioutil.ReadFile(countfile)
	intCount, err := strconv.Atoi(string(count))
	if err != nil {
		log.Error(err)
	}
	return intCount
}

func handleIndex(w http.ResponseWriter, r *http.Request) {

	sounds := getSounds()

	t := template.New("html")
	t, err := t.Parse(templates.MainTemplate)
	if err != nil {
		log.Error(err)
	}

	err = t.Execute(w, sounds)
	if err != nil {
		log.Error(err)
	}

}

type ByNumPlayed []Sound

func (s ByNumPlayed) Len() int {
	return len(s)
}
func (s ByNumPlayed) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s ByNumPlayed) Less(i, j int) bool {
	return s[i].Count < s[j].Count
}

func getSounds() []Sound {
	sounds := make([]Sound, 0)

	dir, err := ioutil.ReadDir("sounds")
	if err != nil {
		log.Error(err)
	}

	for _, v := range dir {
		if !v.IsDir() {
			filename := v.Name()
			if (strings.HasSuffix(filename, ".mp3")) {
				newSound := Sound{
					SoundFile: filename,
					HasImage: true,
					Count: getCounter("sounds/" + filename),
				}
				filenameWithoutExt := filename[:len(filename)-4]
				pngFilename := filenameWithoutExt + ".png"
				if _, err := os.Stat("sounds/" + pngFilename); os.IsNotExist(err) {
					jpgFilename := filenameWithoutExt + ".jpg"
					if _, err := os.Stat("sounds/" + jpgFilename); os.IsNotExist(err) {
						newSound.HasImage = false
					} else {
						newSound.ImageFile = jpgFilename
					}
				} else {
					newSound.ImageFile = pngFilename
				}
				sounds = append(sounds, newSound)
			}
		}
	}

	sort.Sort(ByNumPlayed(sounds))

	return sounds
}

func main() {
	http.HandleFunc("/", handler)
	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		log.Fatal(err)
	}
}<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	yaml "gopkg.in/yaml.v2"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchevents"
)

var Version = "0.2.0"

func main() {
	var (
		apply, dryrun   bool
		file, awsRegion string
	)

	flag.BoolVar(&apply, "apply", false, "apply to CloudWatch Events")
	flag.BoolVar(&dryrun, "dry-run", false, "dry-run")
	flag.StringVar(&file, "file", "config.yml", "file path to setting yaml")
	flag.StringVar(&file, "f", "config.yml", "file path to setting yaml (shorthand)")
	flag.StringVar(&awsRegion, "region", os.Getenv("AWS_REGION"), "aws region")
	flag.Parse()

	sess, errS := session.NewSession(
		&aws.Config{
			Region: aws.String(awsRegion),
		},
	)
	if errS != nil {
		fmt.Printf("Session error %v\n", errS)
		os.Exit(1)
	}

	cweRulesOutput, errR := cloudwatchevents.New(sess).ListRules(nil)
	if errR != nil {
		fmt.Printf("API error %v\n", errR)
		os.Exit(1)
	}

	describedRules := Rules{}
	errY := loadYaml(file, &describedRules)
	if errY != nil {
		fmt.Printf("File error %v\n", errY)
		os.Exit(1)
	}

	describedRules.Rules = AssociateRules(cweRulesOutput.Rules, describedRules.Rules)
	for i, rule := range describedRules.Rules {
		t, _ := fetchActualTargetsByRule(cloudwatchevents.New(sess), rule)
		describedRules.Rules[i].Targets = AssociateTargets(t, describedRules.Rules[i].Targets)
	}
	CheckIsNeedUpdateOrDelete(describedRules.Rules)
	displayWhatWillChange(describedRules.Rules)

	if apply && !dryrun {
		errU := updateCloudWatchEvents(cloudwatchevents.New(sess), describedRules.Rules)
		if errU != nil {
			fmt.Printf("API error %v\n", errU)
			os.Exit(1)
		}
	}
}

func loadYaml(file string, r *Rules) error {

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(buf, &r)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Create sessionfor lambda<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	yaml "gopkg.in/yaml.v2"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchevents"
	"github.com/aws/aws-sdk-go/service/lambda"
)

var Version = "0.2.0"

func main() {
	var (
		apply, dryrun   bool
		file, awsRegion string
	)

	flag.BoolVar(&apply, "apply", false, "apply to CloudWatch Events")
	flag.BoolVar(&dryrun, "dry-run", false, "dry-run")
	flag.StringVar(&file, "file", "config.yml", "file path to setting yaml")
	flag.StringVar(&file, "f", "config.yml", "file path to setting yaml (shorthand)")
	flag.StringVar(&awsRegion, "region", os.Getenv("AWS_REGION"), "aws region")
	flag.Parse()

	sess, errS := session.NewSession(
		&aws.Config{
			Region: aws.String(awsRegion),
		},
	)
	if errS != nil {
		fmt.Printf("Session error %v\n", errS)
		os.Exit(1)
	}

	cweRulesOutput, errR := cloudwatchevents.New(sess).ListRules(nil)
	if errR != nil {
		fmt.Printf("API error %v\n", errR)
		os.Exit(1)
	}

	lambdaClient, errL := lambda.New(sess)
	if errL != nil {
		fmt.Printf("API error %v\n", errL)
		os.Exit(1)
	}

	describedRules := Rules{}
	errY := loadYaml(file, &describedRules)
	if errY != nil {
		fmt.Printf("File error %v\n", errY)
		os.Exit(1)
	}

	describedRules.Rules = AssociateRules(cweRulesOutput.Rules, describedRules.Rules)
	for i, rule := range describedRules.Rules {
		t, _ := fetchActualTargetsByRule(cloudwatchevents.New(sess), rule)
		describedRules.Rules[i].Targets = AssociateTargets(t, describedRules.Rules[i].Targets)
	}
	CheckIsNeedUpdateOrDelete(describedRules.Rules)
	displayWhatWillChange(describedRules.Rules)

	if apply && !dryrun {
		errU := updateCloudWatchEvents(cloudwatchevents.New(sess), describedRules.Rules)
		if errU != nil {
			fmt.Printf("API error %v\n", errU)
			os.Exit(1)
		}
	}
}

func loadYaml(file string, r *Rules) error {

	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return err
	}

	err = yaml.Unmarshal(buf, &r)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Jack Wakefield
//
// 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 (
	"os"

	"github.com/kdar/factorlog"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	locationKey       = "location"
	verboseKey        = "verbose"
	cacheDirectoryKey = "cache-dir"
	cacheLengthKey    = "cache-length"
	serverPortKey     = "server-port"
)

func main() {
	rootCommand := &cobra.Command{
		Use:   "transpac",
		Short: "A transparent proxy which uses proxy auto-config (PAC) files for forwarding",
		Run:   run,
	}

	flags := rootCommand.Flags()

	flags.String(locationKey, "", "Path or URL of the proxy auto-config (PAC) file")
	flags.Bool(verboseKey, false, "Enable verbose logging")
	flags.String(cacheDirectoryKey, "/var/cache/transpac", "The directory where files are cached to")
	flags.Int64(cacheLengthKey, 86400, "The length of time in seconds to cache downloaded files")
	flags.Int(serverPortKey, 8080, "The port the proxy server will listen on")

	viper.SetConfigType("toml")
	viper.SetConfigName("config")
	viper.AddConfigPath("/etc/transpac")
	viper.AddConfigPath("$HOME/.config/transpac")
	viper.BindPFlag(locationKey, flags.Lookup(locationKey))
	viper.BindPFlag(verboseKey, flags.Lookup(verboseKey))
	viper.BindPFlag(cacheDirectoryKey, flags.Lookup(cacheDirectoryKey))
	viper.BindPFlag(cacheLengthKey, flags.Lookup(cacheLengthKey))
	viper.BindPFlag(serverPortKey, flags.Lookup(serverPortKey))
	viper.ReadInConfig()

	rootCommand.Execute()
}

func run(cmd *cobra.Command, args []string) {
	logger.Info(viper.GetBool(verboseKey))

	if viper.GetBool(verboseKey) {
		logger.SetMinMaxSeverity(factorlog.DEBUG, factorlog.PANIC)
	} else {
		logger.SetMinMaxSeverity(factorlog.INFO, factorlog.PANIC)
	}

	// ensure the cache directory exists
	if err := os.MkdirAll(viper.GetString(cacheDirectoryKey), 0755); err != nil {
		logger.Fatalf("Failed to create cache directory (%s)", err)
	}

	location := viper.GetString(locationKey)

	if len(location) == 0 {
		logger.Fatal("You must provide a location for the proxy auto-config file")
	}

	logger.Debugf("Creating proxy auto-config parser for '%s'", location)
	p, err := newParser(location)

	if err != nil {
		logger.Fatalf("Failed to load proxy file (%s)", err)
	}

	logger.Debugf("Creating proxy client")
	c := newProxyClient(p)

	logger.Debugf("Creating proxy server")
	s := newServer(p, c)

	if err := s.listen(); err != nil {
		logger.Fatalf("The proxy server failed to listen (%s)", err)
	}
}
<commit_msg>chore(main): remove accidental log<commit_after>// Copyright 2014 Jack Wakefield
//
// 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 (
	"os"

	"github.com/kdar/factorlog"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const (
	locationKey       = "location"
	verboseKey        = "verbose"
	cacheDirectoryKey = "cache-dir"
	cacheLengthKey    = "cache-length"
	serverPortKey     = "server-port"
)

func main() {
	rootCommand := &cobra.Command{
		Use:   "transpac",
		Short: "A transparent proxy which uses proxy auto-config (PAC) files for forwarding",
		Run:   run,
	}

	flags := rootCommand.Flags()

	flags.String(locationKey, "", "Path or URL of the proxy auto-config (PAC) file")
	flags.Bool(verboseKey, false, "Enable verbose logging")
	flags.String(cacheDirectoryKey, "/var/cache/transpac", "The directory where files are cached to")
	flags.Int64(cacheLengthKey, 86400, "The length of time in seconds to cache downloaded files")
	flags.Int(serverPortKey, 8080, "The port the proxy server will listen on")

	viper.SetConfigType("toml")
	viper.SetConfigName("config")
	viper.AddConfigPath("/etc/transpac")
	viper.AddConfigPath("$HOME/.config/transpac")
	viper.BindPFlag(locationKey, flags.Lookup(locationKey))
	viper.BindPFlag(verboseKey, flags.Lookup(verboseKey))
	viper.BindPFlag(cacheDirectoryKey, flags.Lookup(cacheDirectoryKey))
	viper.BindPFlag(cacheLengthKey, flags.Lookup(cacheLengthKey))
	viper.BindPFlag(serverPortKey, flags.Lookup(serverPortKey))
	viper.ReadInConfig()

	rootCommand.Execute()
}

func run(cmd *cobra.Command, args []string) {
	if viper.GetBool(verboseKey) {
		logger.SetMinMaxSeverity(factorlog.DEBUG, factorlog.PANIC)
	} else {
		logger.SetMinMaxSeverity(factorlog.INFO, factorlog.PANIC)
	}

	// ensure the cache directory exists
	if err := os.MkdirAll(viper.GetString(cacheDirectoryKey), 0755); err != nil {
		logger.Fatalf("Failed to create cache directory (%s)", err)
	}

	location := viper.GetString(locationKey)

	if len(location) == 0 {
		logger.Fatal("You must provide a location for the proxy auto-config file")
	}

	logger.Debugf("Creating proxy auto-config parser for '%s'", location)
	p, err := newParser(location)

	if err != nil {
		logger.Fatalf("Failed to load proxy file (%s)", err)
	}

	logger.Debugf("Creating proxy client")
	c := newProxyClient(p)

	logger.Debugf("Creating proxy server")
	s := newServer(p, c)

	if err := s.listen(); err != nil {
		logger.Fatalf("The proxy server failed to listen (%s)", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/hashicorp/go-plugin"
	"github.com/hashicorp/terraform/helper/logging"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mattn/go-colorable"
	"github.com/mattn/go-shellwords"
	"github.com/mitchellh/cli"
	"github.com/mitchellh/panicwrap"
	"github.com/mitchellh/prefixedio"
)

const (
	// EnvCLI is the environment variable name to set additional CLI args.
	EnvCLI = "TF_CLI_ARGS"
)

func main() {
	// Override global prefix set by go-dynect during init()
	log.SetPrefix("")
	os.Exit(realMain())
}

func realMain() int {
	var wrapConfig panicwrap.WrapConfig

	// don't re-exec terraform as a child process for easier debugging
	if os.Getenv("TF_FORK") == "0" {
		return wrappedMain()
	}

	if !panicwrap.Wrapped(&wrapConfig) {
		// Determine where logs should go in general (requested by the user)
		logWriter, err := logging.LogOutput()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup log output: %s", err)
			return 1
		}

		// We always send logs to a temporary file that we use in case
		// there is a panic. Otherwise, we delete it.
		logTempFile, err := ioutil.TempFile("", "terraform-log")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup logging tempfile: %s", err)
			return 1
		}
		defer os.Remove(logTempFile.Name())
		defer logTempFile.Close()

		// Setup the prefixed readers that send data properly to
		// stdout/stderr.
		doneCh := make(chan struct{})
		outR, outW := io.Pipe()
		go copyOutput(outR, doneCh)

		// Create the configuration for panicwrap and wrap our executable
		wrapConfig.Handler = panicHandler(logTempFile)
		wrapConfig.Writer = io.MultiWriter(logTempFile, logWriter)
		wrapConfig.Stdout = outW
		wrapConfig.IgnoreSignals = ignoreSignals
		wrapConfig.ForwardSignals = forwardSignals
		exitStatus, err := panicwrap.Wrap(&wrapConfig)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't start Terraform: %s", err)
			return 1
		}

		// If >= 0, we're the parent, so just exit
		if exitStatus >= 0 {
			// Close the stdout writer so that our copy process can finish
			outW.Close()

			// Wait for the output copying to finish
			<-doneCh

			return exitStatus
		}

		// We're the child, so just close the tempfile we made in order to
		// save file handles since the tempfile is only used by the parent.
		logTempFile.Close()
	}

	// Call the real main
	return wrappedMain()
}

func wrappedMain() int {
	// We always need to close the DebugInfo before we exit.
	defer terraform.CloseDebugInfo()

	log.SetOutput(os.Stderr)
	log.Printf(
		"[INFO] Terraform version: %s %s %s",
		Version, VersionPrerelease, GitCommit)
	log.Printf("[INFO] CLI args: %#v", os.Args)

	// Load the configuration
	config := BuiltinConfig
	if err := config.Discover(Ui); err != nil {
		Ui.Error(fmt.Sprintf("Error discovering plugins: %s", err))
		return 1
	}

	// Load the configuration file if we have one, that can be used to
	// define extra providers and provisioners.
	clicfgFile, err := cliConfigFile()
	if err != nil {
		Ui.Error(fmt.Sprintf("Error loading CLI configuration: \n\n%s", err))
		return 1
	}

	if clicfgFile != "" {
		usrcfg, err := LoadConfig(clicfgFile)
		if err != nil {
			Ui.Error(fmt.Sprintf("Error loading CLI configuration: \n\n%s", err))
			return 1
		}

		config = *config.Merge(usrcfg)
	}

	// Run checkpoint
	go runCheckpoint(&config)

	// Make sure we clean up any managed plugins at the end of this
	defer plugin.CleanupClients()

	// Get the command line args.
	args := os.Args[1:]

	// Build the CLI so far, we do this so we can query the subcommand.
	cliRunner := &cli.CLI{
		Args:       args,
		Commands:   Commands,
		HelpFunc:   helpFunc,
		HelpWriter: os.Stdout,
	}

	// Prefix the args with any args from the EnvCLI
	args, err = mergeEnvArgs(EnvCLI, cliRunner.Subcommand(), args)
	if err != nil {
		Ui.Error(err.Error())
		return 1
	}

	// Prefix the args with any args from the EnvCLI targeting this command
	suffix := strings.Replace(strings.Replace(
		cliRunner.Subcommand(), "-", "_", -1), " ", "_", -1)
	args, err = mergeEnvArgs(
		fmt.Sprintf("%s_%s", EnvCLI, suffix), cliRunner.Subcommand(), args)
	if err != nil {
		Ui.Error(err.Error())
		return 1
	}

	// We shortcut "--version" and "-v" to just show the version
	for _, arg := range args {
		if arg == "-v" || arg == "-version" || arg == "--version" {
			newArgs := make([]string, len(args)+1)
			newArgs[0] = "version"
			copy(newArgs[1:], args)
			args = newArgs
			break
		}
	}

	// Rebuild the CLI with any modified args.
	log.Printf("[INFO] CLI command args: %#v", args)
	cliRunner = &cli.CLI{
		Args:       args,
		Commands:   Commands,
		HelpFunc:   helpFunc,
		HelpWriter: os.Stdout,
	}

	// Initialize the TFConfig settings for the commands...
	ContextOpts.Providers = config.ProviderFactories()
	ContextOpts.Provisioners = config.ProvisionerFactories()

	exitCode, err := cliRunner.Run()
	if err != nil {
		Ui.Error(fmt.Sprintf("Error executing CLI: %s", err.Error()))
		return 1
	}

	return exitCode
}

func cliConfigFile() (string, error) {
	mustExist := true
	configFilePath := os.Getenv("TERRAFORM_CONFIG")
	if configFilePath == "" {
		var err error
		configFilePath, err = ConfigFile()
		mustExist = false

		if err != nil {
			log.Printf(
				"[ERROR] Error detecting default CLI config file path: %s",
				err)
		}
	}

	log.Printf("[DEBUG] Attempting to open CLI config file: %s", configFilePath)
	f, err := os.Open(configFilePath)
	if err == nil {
		f.Close()
		return configFilePath, nil
	}

	if mustExist || !os.IsNotExist(err) {
		return "", err
	}

	log.Println("[DEBUG] File doesn't exist, but doesn't need to. Ignoring.")
	return "", nil
}

// copyOutput uses output prefixes to determine whether data on stdout
// should go to stdout or stderr. This is due to panicwrap using stderr
// as the log and error channel.
func copyOutput(r io.Reader, doneCh chan<- struct{}) {
	defer close(doneCh)

	pr, err := prefixedio.NewReader(r)
	if err != nil {
		panic(err)
	}

	stderrR, err := pr.Prefix(ErrorPrefix)
	if err != nil {
		panic(err)
	}
	stdoutR, err := pr.Prefix(OutputPrefix)
	if err != nil {
		panic(err)
	}
	defaultR, err := pr.Prefix("")
	if err != nil {
		panic(err)
	}

	var stdout io.Writer = os.Stdout
	var stderr io.Writer = os.Stderr

	if runtime.GOOS == "windows" {
		stdout = colorable.NewColorableStdout()
		stderr = colorable.NewColorableStderr()
	}

	var wg sync.WaitGroup
	wg.Add(3)
	go func() {
		defer wg.Done()
		io.Copy(stderr, stderrR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(stdout, stdoutR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(stdout, defaultR)
	}()

	wg.Wait()
}

func mergeEnvArgs(envName string, cmd string, args []string) ([]string, error) {
	v := os.Getenv(envName)
	if v == "" {
		return args, nil
	}

	log.Printf("[INFO] %s value: %q", envName, v)
	extra, err := shellwords.Parse(v)
	if err != nil {
		return nil, fmt.Errorf(
			"Error parsing extra CLI args from %s: %s",
			envName, err)
	}

	// Find the command to look for in the args. If there is a space,
	// we need to find the last part.
	search := cmd
	if idx := strings.LastIndex(search, " "); idx >= 0 {
		search = cmd[idx+1:]
	}

	// Find the index to place the flags. We put them exactly
	// after the first non-flag arg.
	idx := -1
	for i, v := range args {
		if v == search {
			idx = i
			break
		}
	}

	// idx points to the exact arg that isn't a flag. We increment
	// by one so that all the copying below expects idx to be the
	// insertion point.
	idx++

	// Copy the args
	newArgs := make([]string, len(args)+len(extra))
	copy(newArgs, args[:idx])
	copy(newArgs[idx:], extra)
	copy(newArgs[len(extra)+idx:], args[idx:])
	return newArgs, nil
}
<commit_msg>log the Go runtime version at TF startup<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strings"
	"sync"

	"github.com/hashicorp/go-plugin"
	"github.com/hashicorp/terraform/helper/logging"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mattn/go-colorable"
	"github.com/mattn/go-shellwords"
	"github.com/mitchellh/cli"
	"github.com/mitchellh/panicwrap"
	"github.com/mitchellh/prefixedio"
)

const (
	// EnvCLI is the environment variable name to set additional CLI args.
	EnvCLI = "TF_CLI_ARGS"
)

func main() {
	// Override global prefix set by go-dynect during init()
	log.SetPrefix("")
	os.Exit(realMain())
}

func realMain() int {
	var wrapConfig panicwrap.WrapConfig

	// don't re-exec terraform as a child process for easier debugging
	if os.Getenv("TF_FORK") == "0" {
		return wrappedMain()
	}

	if !panicwrap.Wrapped(&wrapConfig) {
		// Determine where logs should go in general (requested by the user)
		logWriter, err := logging.LogOutput()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup log output: %s", err)
			return 1
		}

		// We always send logs to a temporary file that we use in case
		// there is a panic. Otherwise, we delete it.
		logTempFile, err := ioutil.TempFile("", "terraform-log")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup logging tempfile: %s", err)
			return 1
		}
		defer os.Remove(logTempFile.Name())
		defer logTempFile.Close()

		// Setup the prefixed readers that send data properly to
		// stdout/stderr.
		doneCh := make(chan struct{})
		outR, outW := io.Pipe()
		go copyOutput(outR, doneCh)

		// Create the configuration for panicwrap and wrap our executable
		wrapConfig.Handler = panicHandler(logTempFile)
		wrapConfig.Writer = io.MultiWriter(logTempFile, logWriter)
		wrapConfig.Stdout = outW
		wrapConfig.IgnoreSignals = ignoreSignals
		wrapConfig.ForwardSignals = forwardSignals
		exitStatus, err := panicwrap.Wrap(&wrapConfig)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't start Terraform: %s", err)
			return 1
		}

		// If >= 0, we're the parent, so just exit
		if exitStatus >= 0 {
			// Close the stdout writer so that our copy process can finish
			outW.Close()

			// Wait for the output copying to finish
			<-doneCh

			return exitStatus
		}

		// We're the child, so just close the tempfile we made in order to
		// save file handles since the tempfile is only used by the parent.
		logTempFile.Close()
	}

	// Call the real main
	return wrappedMain()
}

func wrappedMain() int {
	// We always need to close the DebugInfo before we exit.
	defer terraform.CloseDebugInfo()

	log.SetOutput(os.Stderr)
	log.Printf(
		"[INFO] Terraform version: %s %s %s",
		Version, VersionPrerelease, GitCommit)
	log.Printf("[INFO] Go runtime version: %s", runtime.Version())
	log.Printf("[INFO] CLI args: %#v", os.Args)

	// Load the configuration
	config := BuiltinConfig
	if err := config.Discover(Ui); err != nil {
		Ui.Error(fmt.Sprintf("Error discovering plugins: %s", err))
		return 1
	}

	// Load the configuration file if we have one, that can be used to
	// define extra providers and provisioners.
	clicfgFile, err := cliConfigFile()
	if err != nil {
		Ui.Error(fmt.Sprintf("Error loading CLI configuration: \n\n%s", err))
		return 1
	}

	if clicfgFile != "" {
		usrcfg, err := LoadConfig(clicfgFile)
		if err != nil {
			Ui.Error(fmt.Sprintf("Error loading CLI configuration: \n\n%s", err))
			return 1
		}

		config = *config.Merge(usrcfg)
	}

	// Run checkpoint
	go runCheckpoint(&config)

	// Make sure we clean up any managed plugins at the end of this
	defer plugin.CleanupClients()

	// Get the command line args.
	args := os.Args[1:]

	// Build the CLI so far, we do this so we can query the subcommand.
	cliRunner := &cli.CLI{
		Args:       args,
		Commands:   Commands,
		HelpFunc:   helpFunc,
		HelpWriter: os.Stdout,
	}

	// Prefix the args with any args from the EnvCLI
	args, err = mergeEnvArgs(EnvCLI, cliRunner.Subcommand(), args)
	if err != nil {
		Ui.Error(err.Error())
		return 1
	}

	// Prefix the args with any args from the EnvCLI targeting this command
	suffix := strings.Replace(strings.Replace(
		cliRunner.Subcommand(), "-", "_", -1), " ", "_", -1)
	args, err = mergeEnvArgs(
		fmt.Sprintf("%s_%s", EnvCLI, suffix), cliRunner.Subcommand(), args)
	if err != nil {
		Ui.Error(err.Error())
		return 1
	}

	// We shortcut "--version" and "-v" to just show the version
	for _, arg := range args {
		if arg == "-v" || arg == "-version" || arg == "--version" {
			newArgs := make([]string, len(args)+1)
			newArgs[0] = "version"
			copy(newArgs[1:], args)
			args = newArgs
			break
		}
	}

	// Rebuild the CLI with any modified args.
	log.Printf("[INFO] CLI command args: %#v", args)
	cliRunner = &cli.CLI{
		Args:       args,
		Commands:   Commands,
		HelpFunc:   helpFunc,
		HelpWriter: os.Stdout,
	}

	// Initialize the TFConfig settings for the commands...
	ContextOpts.Providers = config.ProviderFactories()
	ContextOpts.Provisioners = config.ProvisionerFactories()

	exitCode, err := cliRunner.Run()
	if err != nil {
		Ui.Error(fmt.Sprintf("Error executing CLI: %s", err.Error()))
		return 1
	}

	return exitCode
}

func cliConfigFile() (string, error) {
	mustExist := true
	configFilePath := os.Getenv("TERRAFORM_CONFIG")
	if configFilePath == "" {
		var err error
		configFilePath, err = ConfigFile()
		mustExist = false

		if err != nil {
			log.Printf(
				"[ERROR] Error detecting default CLI config file path: %s",
				err)
		}
	}

	log.Printf("[DEBUG] Attempting to open CLI config file: %s", configFilePath)
	f, err := os.Open(configFilePath)
	if err == nil {
		f.Close()
		return configFilePath, nil
	}

	if mustExist || !os.IsNotExist(err) {
		return "", err
	}

	log.Println("[DEBUG] File doesn't exist, but doesn't need to. Ignoring.")
	return "", nil
}

// copyOutput uses output prefixes to determine whether data on stdout
// should go to stdout or stderr. This is due to panicwrap using stderr
// as the log and error channel.
func copyOutput(r io.Reader, doneCh chan<- struct{}) {
	defer close(doneCh)

	pr, err := prefixedio.NewReader(r)
	if err != nil {
		panic(err)
	}

	stderrR, err := pr.Prefix(ErrorPrefix)
	if err != nil {
		panic(err)
	}
	stdoutR, err := pr.Prefix(OutputPrefix)
	if err != nil {
		panic(err)
	}
	defaultR, err := pr.Prefix("")
	if err != nil {
		panic(err)
	}

	var stdout io.Writer = os.Stdout
	var stderr io.Writer = os.Stderr

	if runtime.GOOS == "windows" {
		stdout = colorable.NewColorableStdout()
		stderr = colorable.NewColorableStderr()
	}

	var wg sync.WaitGroup
	wg.Add(3)
	go func() {
		defer wg.Done()
		io.Copy(stderr, stderrR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(stdout, stdoutR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(stdout, defaultR)
	}()

	wg.Wait()
}

func mergeEnvArgs(envName string, cmd string, args []string) ([]string, error) {
	v := os.Getenv(envName)
	if v == "" {
		return args, nil
	}

	log.Printf("[INFO] %s value: %q", envName, v)
	extra, err := shellwords.Parse(v)
	if err != nil {
		return nil, fmt.Errorf(
			"Error parsing extra CLI args from %s: %s",
			envName, err)
	}

	// Find the command to look for in the args. If there is a space,
	// we need to find the last part.
	search := cmd
	if idx := strings.LastIndex(search, " "); idx >= 0 {
		search = cmd[idx+1:]
	}

	// Find the index to place the flags. We put them exactly
	// after the first non-flag arg.
	idx := -1
	for i, v := range args {
		if v == search {
			idx = i
			break
		}
	}

	// idx points to the exact arg that isn't a flag. We increment
	// by one so that all the copying below expects idx to be the
	// insertion point.
	idx++

	// Copy the args
	newArgs := make([]string, len(args)+len(extra))
	copy(newArgs, args[:idx])
	copy(newArgs[idx:], extra)
	copy(newArgs[len(extra)+idx:], args[idx:])
	return newArgs, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/evolsnow/robot/conn"
	"github.com/gorilla/websocket"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

var upgrader = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
	return true
}} // use default options for webSocket
var visitor = 0

func main() {
	var configFile string
	var debug bool

	flag.StringVar(&configFile, "c", "config.json", "specify config file")
	flag.BoolVar(&debug, "d", false, "debug mode")
	flag.Parse()
	config, err := ParseConfig(configFile)
	if err != nil {
		log.Fatal("a vailid json config file must exist")
	}

	//connect to redis
	redisPort := strconv.Itoa(config.RedisPort)
	redisServer := net.JoinHostPort(config.RedisAddress, redisPort)
	if !conn.Ping(redisServer, config.RedisPassword) {
		log.Fatal("connect to redis server failed")
	}
	conn.Pool = conn.NewPool(redisServer, config.RedisPassword, config.RedisDB)

	//create robot and run
	robot := newRobot(config.RobotToken, config.RobotName, config.WebHookUrl)
	robot.bot.Debug = debug
	go robot.run()

	//run server and web samaritan
	srvPort := strconv.Itoa(config.Port)
	http.HandleFunc("/ajax", ajax)
	http.HandleFunc("/websocket", socketHandler)
	http.HandleFunc("/groupTalk", groupTalk)
	log.Fatal(http.ListenAndServeTLS(net.JoinHostPort(config.Server, srvPort), config.Cert, config.CertKey, nil))

}

//3 robot's group talk
func groupTalk(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	visitor++
	tlChan := make(chan string, 5)
	qinChan := make(chan string, 5)
	iceChan := make(chan string, 5)
	result := make(chan string, 10)
	initSentence := "你好"
	tlChan <- tlAI(initSentence)
	go func() {
		for {
			if visitor > 0 {
				msgToTl := <-tlChan
				replyFromTl := tlAI(msgToTl)
				go func() {
					if replyFromTl != "" {
						result <- "samaritan: " + replyFromTl
						qinChan <- replyFromTl
						iceChan <- replyFromTl
					}
				}()
				//c.WriteMessage(websocket.TextMessage, []byte("samaritan: " + replyFromTl))
			}
		}
	}()

	go func() {
		for {
			if visitor > 0 {
				msgToQin := <-qinChan
				replyFromQin := qinAI(msgToQin)
				go func() {
					if replyFromQin != "" {
						result <- "菲菲: " + replyFromQin
						tlChan <- replyFromQin
						iceChan <- replyFromQin
					}
				}()
				//c.WriteMessage(websocket.TextMessage, []byte("菲菲: " + replyFromQin))
			}
		}
	}()

	go func() {
		for {
			if visitor > 0 {
				msgToIce := <-iceChan
				replyFromIce := iceAI(msgToIce)
				go func() {
					if replyFromIce != "" {
						result <- "小冰: " + replyFromIce
						tlChan <- replyFromIce
						qinChan <- replyFromIce
					}
				}()
				//c.WriteMessage(websocket.TextMessage, []byte("小冰: " + replyFromIce))
			}
		}
	}()

	go func() {
		for {
			c.WriteMessage(websocket.TextMessage, []byte(<-result))
		}
	}()
	for {

		_, _, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			visitor--
			break
		}
	}
}

//used for web samaritan robot
func socketHandler(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	for {
		var in []byte
		var ret []string
		mt, in, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		ret = receive(string(in))
		for i := range ret {
			c.WriteMessage(mt, []byte(ret[i]))
			time.Sleep(time.Second)
		}
		c.WriteMessage(mt, []byte(""))
	}
}

//when webSocket unavailable, fallback to ajax long polling
func ajax(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Access-Control-Allow-Origin", "*")
	var messages = make(chan string)
	if r.Method == "GET" {
		w.Write([]byte(<-messages))
	} else {
		body := r.FormValue("text")
		if body != "" {
			go func(string) {
				ret := receive(body)
				for i := range ret {
					messages <- ret[i]
					time.Sleep(time.Second)
				}
				messages <- ""
			}(body)
		}
	}
}

//receive from client
func receive(in string) (ret []string) {
	fmt.Printf("Received: %s\n", in)
	var response string
	var answer = make(chan string)
	sf := func(c rune) bool {
		return c == ',' || c == '，' || c == ';' || c == '。' || c == '.' || c == '？' || c == '?'
	}
	if chinese(in) {
		go func() {
			if ret := iceAI(in); ret != "" {
				answer <- ret
			}
		}()
		go func() {
			if ret := tlAI(in); ret != "" {
				answer <- ret
			}
		}()
		go func() {
			if ret := qinAI(in); ret != "" {
				answer <- strings.Replace(ret, "Jarvis", "samaritan", -1)
			}
		}()
		response = <-answer
		// Separate into fields with func.
		ret = strings.FieldsFunc(response, sf)

	} else {
		response = mitAI(in)
		ret = strings.FieldsFunc(response, sf)
	}
	return
}
<commit_msg>slow down<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/evolsnow/robot/conn"
	"github.com/gorilla/websocket"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
)

var upgrader = websocket.Upgrader{CheckOrigin: func(r *http.Request) bool {
	return true
}} // use default options for webSocket
var visitor = 0

func main() {
	var configFile string
	var debug bool

	flag.StringVar(&configFile, "c", "config.json", "specify config file")
	flag.BoolVar(&debug, "d", false, "debug mode")
	flag.Parse()
	config, err := ParseConfig(configFile)
	if err != nil {
		log.Fatal("a vailid json config file must exist")
	}

	//connect to redis
	redisPort := strconv.Itoa(config.RedisPort)
	redisServer := net.JoinHostPort(config.RedisAddress, redisPort)
	if !conn.Ping(redisServer, config.RedisPassword) {
		log.Fatal("connect to redis server failed")
	}
	conn.Pool = conn.NewPool(redisServer, config.RedisPassword, config.RedisDB)

	//create robot and run
	robot := newRobot(config.RobotToken, config.RobotName, config.WebHookUrl)
	robot.bot.Debug = debug
	go robot.run()

	//run server and web samaritan
	srvPort := strconv.Itoa(config.Port)
	http.HandleFunc("/ajax", ajax)
	http.HandleFunc("/websocket", socketHandler)
	http.HandleFunc("/groupTalk", groupTalk)
	log.Fatal(http.ListenAndServeTLS(net.JoinHostPort(config.Server, srvPort), config.Cert, config.CertKey, nil))

}

//3 robot's group talk
func groupTalk(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	visitor++
	tlChan := make(chan string, 5)
	qinChan := make(chan string, 5)
	iceChan := make(chan string, 5)
	result := make(chan string, 10)
	initSentence := "你好"
	tlChan <- tlAI(initSentence)
	go func() {
		for {
			if visitor > 0 {
				msgToTl := <-tlChan
				replyFromTl := tlAI(msgToTl)
				time.Sleep(time.Second)
				go func() {
					if replyFromTl != "" {
						result <- "samaritan: " + replyFromTl
						qinChan <- replyFromTl
						iceChan <- replyFromTl
					}
				}()
				//c.WriteMessage(websocket.TextMessage, []byte("samaritan: " + replyFromTl))
			}
		}
	}()

	go func() {
		for {
			if visitor > 0 {
				msgToQin := <-qinChan
				replyFromQin := qinAI(msgToQin)
				time.Sleep(time.Second)
				go func() {
					if replyFromQin != "" {
						result <- "菲菲: " + replyFromQin
						tlChan <- replyFromQin
						iceChan <- replyFromQin
					}
				}()
				//c.WriteMessage(websocket.TextMessage, []byte("菲菲: " + replyFromQin))
			}
		}
	}()

	go func() {
		for {
			if visitor > 0 {
				msgToIce := <-iceChan
				replyFromIce := iceAI(msgToIce)
				go func() {
					if replyFromIce != "" {
						result <- "小冰: " + replyFromIce
						tlChan <- replyFromIce
						qinChan <- replyFromIce
					}
				}()
				//c.WriteMessage(websocket.TextMessage, []byte("小冰: " + replyFromIce))
			}
		}
	}()

	go func() {
		for {
			c.WriteMessage(websocket.TextMessage, []byte(<-result))
		}
	}()
	for {

		_, _, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			visitor--
			break
		}
	}
}

//used for web samaritan robot
func socketHandler(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	for {
		var in []byte
		var ret []string
		mt, in, err := c.ReadMessage()
		if err != nil {
			log.Println("read:", err)
			break
		}
		ret = receive(string(in))
		for i := range ret {
			c.WriteMessage(mt, []byte(ret[i]))
			time.Sleep(time.Second)
		}
		c.WriteMessage(mt, []byte(""))
	}
}

//when webSocket unavailable, fallback to ajax long polling
func ajax(w http.ResponseWriter, r *http.Request) {
	w.Header().Add("Access-Control-Allow-Origin", "*")
	var messages = make(chan string)
	if r.Method == "GET" {
		w.Write([]byte(<-messages))
	} else {
		body := r.FormValue("text")
		if body != "" {
			go func(string) {
				ret := receive(body)
				for i := range ret {
					messages <- ret[i]
					time.Sleep(time.Second)
				}
				messages <- ""
			}(body)
		}
	}
}

//receive from client
func receive(in string) (ret []string) {
	fmt.Printf("Received: %s\n", in)
	var response string
	var answer = make(chan string)
	sf := func(c rune) bool {
		return c == ',' || c == '，' || c == ';' || c == '。' || c == '.' || c == '？' || c == '?'
	}
	if chinese(in) {
		go func() {
			if ret := iceAI(in); ret != "" {
				answer <- ret
			}
		}()
		go func() {
			if ret := tlAI(in); ret != "" {
				answer <- ret
			}
		}()
		go func() {
			if ret := qinAI(in); ret != "" {
				answer <- strings.Replace(ret, "Jarvis", "samaritan", -1)
			}
		}()
		response = <-answer
		// Separate into fields with func.
		ret = strings.FieldsFunc(response, sf)

	} else {
		response = mitAI(in)
		ret = strings.FieldsFunc(response, sf)
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"

	"log"

	"github.com/jawher/mow.cli"
)

const (
	Version      = "0.6.0"
	DefaultTheme = "clean"
)

func main() {
	app := cli.App("servemd", "a simple http server for markdown content")
	app.Version("v version", Version)
	app.Spec = "[OPTIONS] [DIR]"

	var (
		// HTTP Options
		host = app.String(cli.StringOpt{
			Name:   "a host",
			Desc:   "Host/IP address to listen on for HTTP",
			Value:  "",
			EnvVar: "HOST",
		})
		port = app.Int(cli.IntOpt{
			Name:   "p port",
			Desc:   "TCP PORT to listen on for HTTP",
			Value:  3000,
			EnvVar: "PORT",
		})
		users = app.Strings(cli.StringsOpt{
			Name:   "u auth",
			Desc:   "Username and password for basic authentication in the form of user:pass",
			EnvVar: "BASIC_AUTH",
		})
		robotsTag = app.String(cli.StringOpt{
			Name:   "r x-robots-tag",
			Desc:   "Sets a X-Robots-Tag header",
			EnvVar: "X_ROBOTS_TAG",
			Value:  "",
		})

		// Content
		dir = app.String(cli.StringArg{
			Name:   "DIR",
			Desc:   "Directory to serve content from",
			Value:  ".",
			EnvVar: "DOCUMENT_ROOT",
		})
		extension = app.String(cli.StringOpt{
			Name:   "e extension",
			Desc:   "Extension used for markdown files",
			Value:  ".md",
			EnvVar: "DOCUMENT_EXTENSION",
		})
		index = app.String(cli.StringOpt{
			Name:   "i index",
			Desc:   "Filename (without extension) to use for directory index",
			Value:  "index",
			EnvVar: "DIRECTORY_INDEX",
		})

		// Theme
		markdownTheme = app.String(cli.StringOpt{
			Name:   "m markdown-theme",
			Desc:   "Theme to use for styling markdown html",
			Value:  DefaultTheme,
			EnvVar: "MARKDOWN_THEME",
		})
		typekitKitID = app.String(cli.StringOpt{
			Name:   "t typekit-kit-id",
			Desc:   "ID of webfont kit to include from typekit",
			Value:  DefaultTheme,
			EnvVar: "TYPEKIT_KIT_ID",
		})
		codeTheme = app.String(cli.StringOpt{
			Name:   "c code-theme",
			Desc:   "Highlight.js theme to use for syntax highlighting",
			Value:  "",
			EnvVar: "CODE_THEME",
		})
	)

	app.Action = func() {
		// Static Asset Handler
		staticAssetHandler := staticAssetServer()
		staticAssetHandlerFunc := func(w http.ResponseWriter, r *http.Request) {
			staticAssetHandler.ServeHTTP(w, r)
		}
		staticAssetHandlerFunc = headerMiddleware(staticAssetHandlerFunc)
		staticAssetHandlerFunc = basicAuthMiddleware(staticAssetHandlerFunc, *users)
		staticAssetHandlerFunc = robotsTagMiddleware(staticAssetHandlerFunc, *robotsTag)
		http.HandleFunc("/assets/", staticAssetHandlerFunc)

		// Setup the markdown theme (may be custom or bundled)
		themePath, themeHandler := theme(*markdownTheme)
		if themeHandler != nil {
			themeHandler = headerMiddleware(themeHandler)
			themeHandler = basicAuthMiddleware(themeHandler, *users)
			themeHandler = robotsTagMiddleware(themeHandler, *robotsTag)
			http.HandleFunc(themePath, themeHandler)
		}

		// Markdown File Handler
		markdownHandlerFunc := markdownHandleFunc(MarkdownHandlerOptions{
			DocRoot:       *dir,
			DocExtension:  *extension,
			DirIndex:      *index,
			MarkdownTheme: themePath,
			TypekitKitID:  *typekitKitID,
			CodeTheme:     *codeTheme,
		})
		markdownHandlerFunc = headerMiddleware(markdownHandlerFunc)
		markdownHandlerFunc = basicAuthMiddleware(markdownHandlerFunc, *users)
		markdownHandlerFunc = robotsTagMiddleware(markdownHandlerFunc, *robotsTag)
		http.HandleFunc("/", markdownHandlerFunc)

		// Start HTTP server
		addr := fmt.Sprintf("%s:%d", *host, *port)
		log.Printf("Starting server on %s", addr)
		err := http.ListenAndServe(addr, nil)
		if err != nil {
			log.Printf("Error starting server: %s", err)
			cli.Exit(1)
		}

	}

	app.Run(os.Args)
}
<commit_msg>Bump version to 0.7.0<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"

	"log"

	"github.com/jawher/mow.cli"
)

const (
	Version      = "0.7.0"
	DefaultTheme = "clean"
)

func main() {
	app := cli.App("servemd", "a simple http server for markdown content")
	app.Version("v version", Version)
	app.Spec = "[OPTIONS] [DIR]"

	var (
		// HTTP Options
		host = app.String(cli.StringOpt{
			Name:   "a host",
			Desc:   "Host/IP address to listen on for HTTP",
			Value:  "",
			EnvVar: "HOST",
		})
		port = app.Int(cli.IntOpt{
			Name:   "p port",
			Desc:   "TCP PORT to listen on for HTTP",
			Value:  3000,
			EnvVar: "PORT",
		})
		users = app.Strings(cli.StringsOpt{
			Name:   "u auth",
			Desc:   "Username and password for basic authentication in the form of user:pass",
			EnvVar: "BASIC_AUTH",
		})
		robotsTag = app.String(cli.StringOpt{
			Name:   "r x-robots-tag",
			Desc:   "Sets a X-Robots-Tag header",
			EnvVar: "X_ROBOTS_TAG",
			Value:  "",
		})

		// Content
		dir = app.String(cli.StringArg{
			Name:   "DIR",
			Desc:   "Directory to serve content from",
			Value:  ".",
			EnvVar: "DOCUMENT_ROOT",
		})
		extension = app.String(cli.StringOpt{
			Name:   "e extension",
			Desc:   "Extension used for markdown files",
			Value:  ".md",
			EnvVar: "DOCUMENT_EXTENSION",
		})
		index = app.String(cli.StringOpt{
			Name:   "i index",
			Desc:   "Filename (without extension) to use for directory index",
			Value:  "index",
			EnvVar: "DIRECTORY_INDEX",
		})

		// Theme
		markdownTheme = app.String(cli.StringOpt{
			Name:   "m markdown-theme",
			Desc:   "Theme to use for styling markdown html",
			Value:  DefaultTheme,
			EnvVar: "MARKDOWN_THEME",
		})
		typekitKitID = app.String(cli.StringOpt{
			Name:   "t typekit-kit-id",
			Desc:   "ID of webfont kit to include from typekit",
			Value:  DefaultTheme,
			EnvVar: "TYPEKIT_KIT_ID",
		})
		codeTheme = app.String(cli.StringOpt{
			Name:   "c code-theme",
			Desc:   "Highlight.js theme to use for syntax highlighting",
			Value:  "",
			EnvVar: "CODE_THEME",
		})
	)

	app.Action = func() {
		// Static Asset Handler
		staticAssetHandler := staticAssetServer()
		staticAssetHandlerFunc := func(w http.ResponseWriter, r *http.Request) {
			staticAssetHandler.ServeHTTP(w, r)
		}
		staticAssetHandlerFunc = headerMiddleware(staticAssetHandlerFunc)
		staticAssetHandlerFunc = basicAuthMiddleware(staticAssetHandlerFunc, *users)
		staticAssetHandlerFunc = robotsTagMiddleware(staticAssetHandlerFunc, *robotsTag)
		http.HandleFunc("/assets/", staticAssetHandlerFunc)

		// Setup the markdown theme (may be custom or bundled)
		themePath, themeHandler := theme(*markdownTheme)
		if themeHandler != nil {
			themeHandler = headerMiddleware(themeHandler)
			themeHandler = basicAuthMiddleware(themeHandler, *users)
			themeHandler = robotsTagMiddleware(themeHandler, *robotsTag)
			http.HandleFunc(themePath, themeHandler)
		}

		// Markdown File Handler
		markdownHandlerFunc := markdownHandleFunc(MarkdownHandlerOptions{
			DocRoot:       *dir,
			DocExtension:  *extension,
			DirIndex:      *index,
			MarkdownTheme: themePath,
			TypekitKitID:  *typekitKitID,
			CodeTheme:     *codeTheme,
		})
		markdownHandlerFunc = headerMiddleware(markdownHandlerFunc)
		markdownHandlerFunc = basicAuthMiddleware(markdownHandlerFunc, *users)
		markdownHandlerFunc = robotsTagMiddleware(markdownHandlerFunc, *robotsTag)
		http.HandleFunc("/", markdownHandlerFunc)

		// Start HTTP server
		addr := fmt.Sprintf("%s:%d", *host, *port)
		log.Printf("Starting server on %s", addr)
		err := http.ListenAndServe(addr, nil)
		if err != nil {
			log.Printf("Error starting server: %s", err)
			cli.Exit(1)
		}

	}

	app.Run(os.Args)
}
<|endoftext|>