<commit_before>package main

import (
	"encoding/json"
	"log"
	"strings"
)

func listFiles(path string, includeMeta bool) (map[string]interface{}, error) {
	viewRes := struct {
		Rows []struct {
			Key   []string
			Value float64
		}
	}{}

	// use the requested path to build our view query parameters
	startKey := []string{}
	if path != "" {
		startKey = strings.Split(path, "/")
	}
	endKey := make([]string, len(startKey)+1, len(startKey)+1)
	copy(endKey, startKey)
	endKey[len(startKey)] = "ZZZZZZ" // FIXME use {} instead
	groupLevel := len(startKey) + 1

	// query the view
	err := couchbase.ViewCustom("cbfs", "file_browse",
		map[string]interface{}{"group_level": groupLevel, "start_key": startKey, "end_key": endKey}, &viewRes)
	if err != nil {
		return nil, err
	}

	// use the view result to build a list of keys
	keys := make([]string, len(viewRes.Rows), len(viewRes.Rows))
	for i, r := range viewRes.Rows {
		keys[i] = strings.Join(r.Key, "/")
	}

	// do a multi-get on the all the keys returned
	bulkResult := couchbase.GetBulk(keys)

	// divide items up into files and directories
	files := map[string]interface{}{}
	dirs := map[string]interface{}{}
	for _, r := range viewRes.Rows {
		key := strings.Join(r.Key, "/")
		res, ok := bulkResult[key]
		if ok == true {
			// this means we have a file
			rv := map[string]interface{}{}
			err := json.Unmarshal(res.Body, &rv)
			if err != nil {
				log.Printf("Error deserializing json, ignoring: %v", err)
			} else {
				if includeMeta {
					files[r.Key[len(r.Key)-1]] = rv
				} else {
					files[r.Key[len(r.Key)-1]] = map[string]interface{}{}
				}
			}
		} else {
			// no record in the multi-get metans this is a directory
			dirs[r.Key[len(r.Key)-1]] = map[string]interface{}{}
		}
	}

	// assemble the final return value
	rv := map[string]interface{}{"dirs": dirs, "files": files}

	return rv, nil
}
<commit_msg>added requested path to the response useful when you're handed one of these objects with no context<commit_after>package main

import (
	"encoding/json"
	"log"
	"strings"
)

func listFiles(path string, includeMeta bool) (map[string]interface{}, error) {
	viewRes := struct {
		Rows []struct {
			Key   []string
			Value float64
		}
	}{}

	// use the requested path to build our view query parameters
	startKey := []string{}
	if path != "" {
		startKey = strings.Split(path, "/")
	}
	endKey := make([]string, len(startKey)+1, len(startKey)+1)
	copy(endKey, startKey)
	endKey[len(startKey)] = "ZZZZZZ" // FIXME use {} instead
	groupLevel := len(startKey) + 1

	// query the view
	err := couchbase.ViewCustom("cbfs", "file_browse",
		map[string]interface{}{"group_level": groupLevel, "start_key": startKey, "end_key": endKey}, &viewRes)
	if err != nil {
		return nil, err
	}

	// use the view result to build a list of keys
	keys := make([]string, len(viewRes.Rows), len(viewRes.Rows))
	for i, r := range viewRes.Rows {
		keys[i] = strings.Join(r.Key, "/")
	}

	// do a multi-get on the all the keys returned
	bulkResult := couchbase.GetBulk(keys)

	// divide items up into files and directories
	files := map[string]interface{}{}
	dirs := map[string]interface{}{}
	for _, r := range viewRes.Rows {
		key := strings.Join(r.Key, "/")
		res, ok := bulkResult[key]
		if ok == true {
			// this means we have a file
			rv := map[string]interface{}{}
			err := json.Unmarshal(res.Body, &rv)
			if err != nil {
				log.Printf("Error deserializing json, ignoring: %v", err)
			} else {
				if includeMeta {
					files[r.Key[len(r.Key)-1]] = rv
				} else {
					files[r.Key[len(r.Key)-1]] = map[string]interface{}{}
				}
			}
		} else {
			// no record in the multi-get metans this is a directory
			dirs[r.Key[len(r.Key)-1]] = map[string]interface{}{}
		}
	}

	// assemble the final return value
	rv := map[string]interface{}{"path": "/" + path, "dirs": dirs, "files": files}

	return rv, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"path/filepath"

	"golang.org/x/tools/go/loader"
)

func loadPaths(wd string, fset *token.FileSet, paths []string) ([]ast.Node, error) {
	var nodes []ast.Node
	ctx := build.Default
	for _, path := range paths {
		pkg, err := ctx.Import(path, wd, 0)
		if err != nil {
			return nil, err
		}
		for _, names := range [...][]string{
			pkg.GoFiles, pkg.CgoFiles, pkg.IgnoredGoFiles,
			pkg.TestGoFiles, pkg.XTestGoFiles,
		} {
			for _, name := range names {
				path := filepath.Join(pkg.Dir, name)
				f, err := parser.ParseFile(fset, path, nil, 0)
				if err != nil {
					return nil, err
				}
				nodes = append(nodes, f)
			}
		}
	}
	return nodes, nil
}

func loadTyped(wd string, fset *token.FileSet, paths []string) (*loader.Program, error) {
	conf := loader.Config{Fset: fset, Cwd: wd}
	if _, err := conf.FromArgs(paths, true); err != nil {
		return nil, err
	}
	return conf.Load()
}
<commit_msg>add loading of separate .go files without types<commit_after>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"path/filepath"
	"strings"

	"golang.org/x/tools/go/loader"
)

func loadPaths(wd string, fset *token.FileSet, paths []string) ([]ast.Node, error) {
	var nodes []ast.Node
	ctx := build.Default
	addFile := func(path string) error {
		f, err := parser.ParseFile(fset, path, nil, 0)
		if err != nil {
			return err
		}
		nodes = append(nodes, f)
		return nil
	}
	for _, path := range paths {
		if strings.HasSuffix(path, ".go") {
			if err := addFile(path); err != nil {
				return nil, err
			}
			continue
		}
		pkg, err := ctx.Import(path, wd, 0)
		if err != nil {
			return nil, err
		}
		for _, names := range [...][]string{
			pkg.GoFiles, pkg.CgoFiles, pkg.IgnoredGoFiles,
			pkg.TestGoFiles, pkg.XTestGoFiles,
		} {
			for _, name := range names {
				path := filepath.Join(pkg.Dir, name)
				if err := addFile(path); err != nil {
					return nil, err
				}
			}
		}
	}
	return nodes, nil
}

func loadTyped(wd string, fset *token.FileSet, paths []string) (*loader.Program, error) {
	conf := loader.Config{Fset: fset, Cwd: wd}
	if _, err := conf.FromArgs(paths, true); err != nil {
		return nil, err
	}
	return conf.Load()
}
<|endoftext|>


<commit_before>package scmd

import (
	"path"
	"runtime"

	"github.com/BurntSushi/toml"
)

type Config struct {
	TOKEN string
}

func LoadToml(confname string) *Config {
	var conf Config
	_, cur, _, _ := runtime.Caller(1)
	confpath := path.Join(path.Dir(cur), confname)

	_, err := toml.DecodeFile(confpath, &conf)
	if err != nil {
		panic(err)
	}

	return &conf
}
<commit_msg>some change<commit_after>package scmd

import (
	"os"
	"path"

	"github.com/BurntSushi/toml"
)

type Config struct {
	TOKEN string
}

func LoadToml(confname string) *Config {
	var conf Config

	cur, _ := os.Getwd()
	confpath := path.Join(cur, confname)

	_, err := toml.DecodeFile(confpath, &conf)
	if err != nil {
		panic(err)
	}

	return &conf
}
<|endoftext|>


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

package main

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"

	"github.com/sdboyer/gps"
)

type lock struct {
	Memo []byte
	P    []gps.LockedProject
}

type rawLock struct {
	Memo string      `json:"memo"`
	P    []lockedDep `json:"projects"`
}

type lockedDep struct {
	Name       string   `json:"name"`
	Version    string   `json:"version,omitempty"`
	Branch     string   `json:"branch,omitempty"`
	Revision   string   `json:"revision"`
	Repository string   `json:"repo,omitempty"`
	Packages   []string `json:"packages"`
}

func readLock(r io.Reader) (*lock, error) {
	rl := rawLock{}
	err := json.NewDecoder(r).Decode(&rl)
	if err != nil {
		return nil, err
	}

	b, err := hex.DecodeString(rl.Memo)
	if err != nil {
		return nil, fmt.Errorf("invalid hash digest in lock's memo field")
	}
	l := &lock{
		Memo: b,
		P:    make([]gps.LockedProject, len(rl.P)),
	}

	for i, ld := range rl.P {
		r := gps.Revision(ld.Revision)

		var v gps.Version
		if ld.Version != "" {
			if ld.Branch != "" {
				return nil, fmt.Errorf("lock file specified both a branch (%s) and version (%s) for %s", ld.Branch, ld.Version, ld.Name)
			}
			v = gps.NewVersion(ld.Version).Is(r)
		} else if ld.Branch != "" {
			v = gps.NewBranch(ld.Branch).Is(r)
		} else if r == "" {
			return nil, fmt.Errorf("lock file has entry for %s, but specifies no version", ld.Name)
		} else {
			v = r
		}

		id := gps.ProjectIdentifier{
			ProjectRoot: gps.ProjectRoot(ld.Name),
			NetworkName: ld.Repository,
		}
		l.P[i] = gps.NewLockedProject(id, v, ld.Packages)
	}

	return l, nil
}

func (l *lock) InputHash() []byte {
	return l.Memo
}

func (l *lock) Projects() []gps.LockedProject {
	return l.P
}
<commit_msg>dep: JSON marshaling for locks<commit_after>// Copyright 2016 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"

	"github.com/sdboyer/gps"
)

type lock struct {
	Memo []byte
	P    []gps.LockedProject
}

type rawLock struct {
	Memo string      `json:"memo"`
	P    []lockedDep `json:"projects"`
}

type lockedDep struct {
	Name       string   `json:"name"`
	Version    string   `json:"version,omitempty"`
	Branch     string   `json:"branch,omitempty"`
	Revision   string   `json:"revision"`
	Repository string   `json:"repo,omitempty"`
	Packages   []string `json:"packages"`
}

func readLock(r io.Reader) (*lock, error) {
	rl := rawLock{}
	err := json.NewDecoder(r).Decode(&rl)
	if err != nil {
		return nil, err
	}

	b, err := hex.DecodeString(rl.Memo)
	if err != nil {
		return nil, fmt.Errorf("invalid hash digest in lock's memo field")
	}
	l := &lock{
		Memo: b,
		P:    make([]gps.LockedProject, len(rl.P)),
	}

	for i, ld := range rl.P {
		r := gps.Revision(ld.Revision)

		var v gps.Version
		if ld.Version != "" {
			if ld.Branch != "" {
				return nil, fmt.Errorf("lock file specified both a branch (%s) and version (%s) for %s", ld.Branch, ld.Version, ld.Name)
			}
			v = gps.NewVersion(ld.Version).Is(r)
		} else if ld.Branch != "" {
			v = gps.NewBranch(ld.Branch).Is(r)
		} else if r == "" {
			return nil, fmt.Errorf("lock file has entry for %s, but specifies no version", ld.Name)
		} else {
			v = r
		}

		id := gps.ProjectIdentifier{
			ProjectRoot: gps.ProjectRoot(ld.Name),
			NetworkName: ld.Repository,
		}
		l.P[i] = gps.NewLockedProject(id, v, ld.Packages)
	}

	return l, nil
}

func (l *lock) InputHash() []byte {
	return l.Memo
}

func (l *lock) Projects() []gps.LockedProject {
	return l.P
}

func (l *lock) MarshalJSON() ([]byte, error) {
	raw := rawLock{
		Memo: hex.EncodeToString(l.Memo),
		P: make([]lockedDep, len(l.P))
	}

	for k, lp := range l.P {
		id := lp.Ident()
		ld := lockedDep{
			Name:       string(id.ProjectRoot()),
			Repository: id.NetworkName,
			Packages:   lp.Packages(),
		}

		switch tv := lp.Version().(type) {
		case gps.UnpairedVersion:
			ld.Version = tv
		case gps.Revision:
			ld.Revision = tv
		case gps.PairedVersion:
			ld.Version = tv.Unpair()
			ld.Revision = tv.Underlying()
		}

		raw.P[k] = ld
	}

	return json.Marshal(raw)
}
<|endoftext|>


<commit_before>package store

import (
	"fmt"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/log"
	"launchpad.net/lpad"
	"strings"
	"time"
)

type PublishBranchError struct {
	URL string
	Err error
}

type PublishBranchErrors []PublishBranchError

func (errs PublishBranchErrors) Error() string {
	return fmt.Sprintf("%d branch(es) failed to be published", len(errs))
}

// PublishCharmsDistro publishes all branch tips found in
// the /charms distribution in Launchpad onto store under
// the "cs:" scheme.
// apiBase specifies the Launchpad base API URL, such
// as lpad.Production or lpad.Staging.
// Errors found while processing one or more branches are
// all returned as a PublishBranchErrors value.
func PublishCharmsDistro(store *Store, apiBase lpad.APIBase) error {
	oauth := &lpad.OAuth{Anonymous: true, Consumer: "juju"}
	root, err := lpad.Login(apiBase, oauth)
	if err != nil {
		return err
	}
	distro, err := root.Distro("charms")
	if err != nil {
		return err
	}
	tips, err := distro.BranchTips(time.Time{})
	if err != nil {
		return err
	}

	var errs PublishBranchErrors
	for _, tip := range tips {
		if !strings.HasSuffix(tip.UniqueName, "/trunk") {
			continue
		}
		burl, curl, err := uniqueNameURLs(tip.UniqueName)
		if err != nil {
			errs = append(errs, PublishBranchError{tip.UniqueName, err})
			log.Printf("error: %v\n", err)
			continue
		}
		log.Printf("----- %s\n", burl)
		if tip.Revision == "" {
			errs = append(errs, PublishBranchError{burl, fmt.Errorf("branch has no revisions")})
			log.Printf("error: branch has no revisions\n")
			continue
		}
		// Charm is published in the personal URL and in any explicitly
		// assigned official series.
		urls := []*charm.URL{curl}
		schema, name := curl.Schema, curl.Name
		for _, series := range tip.OfficialSeries {
			curl = &charm.URL{Schema: schema, Name: name, Series: series, Revision: -1}
			curl.Series = series
			curl.User = ""
			urls = append(urls, curl)
		}

		err = PublishBazaarBranch(store, urls, burl, tip.Revision)
		if err == ErrRedundantUpdate {
			continue
		}
		if err != nil {
			errs = append(errs, PublishBranchError{burl, err})
			log.Printf("error: %v\n", err)
		}
	}
	if errs != nil {
		return errs
	}
	return nil
}

// uniqueNameURLs returns the branch URL and the charm URL for the
// provided Launchpad branch unique name. The unique name must be
// in the form:
//
//     ~<user>/charms/<series>/<charm name>/trunk
//
// For testing purposes, if name has a prefix preceding a string in
// this format, the prefix is stripped out for computing the charm
// URL, and the unique name is returned unchanged as the branch URL.
func uniqueNameURLs(name string) (burl string, curl *charm.URL, err error) {
	u := strings.Split(name, "/")
	if len(u) > 5 {
		u = u[len(u)-5:]
		burl = name
	} else {
		burl = "lp:" + name
	}
	if len(u) < 5 || u[1] != "charms" || u[4] != "trunk" || len(u[0]) < 0 || u[0][0] != '~' {
		return "", nil, fmt.Errorf("unwanted branch name: %s", name)
	}
	curl, err = charm.ParseURL(fmt.Sprintf("cs:%s/%s/%s", u[0], u[2], u[3]))
	if err != nil {
		return "", nil, err
	}
	return burl, curl, nil
}
<commit_msg>fix string length check<commit_after>package store

import (
	"fmt"
	"launchpad.net/juju/go/charm"
	"launchpad.net/juju/go/log"
	"launchpad.net/lpad"
	"strings"
	"time"
)

type PublishBranchError struct {
	URL string
	Err error
}

type PublishBranchErrors []PublishBranchError

func (errs PublishBranchErrors) Error() string {
	return fmt.Sprintf("%d branch(es) failed to be published", len(errs))
}

// PublishCharmsDistro publishes all branch tips found in
// the /charms distribution in Launchpad onto store under
// the "cs:" scheme.
// apiBase specifies the Launchpad base API URL, such
// as lpad.Production or lpad.Staging.
// Errors found while processing one or more branches are
// all returned as a PublishBranchErrors value.
func PublishCharmsDistro(store *Store, apiBase lpad.APIBase) error {
	oauth := &lpad.OAuth{Anonymous: true, Consumer: "juju"}
	root, err := lpad.Login(apiBase, oauth)
	if err != nil {
		return err
	}
	distro, err := root.Distro("charms")
	if err != nil {
		return err
	}
	tips, err := distro.BranchTips(time.Time{})
	if err != nil {
		return err
	}

	var errs PublishBranchErrors
	for _, tip := range tips {
		if !strings.HasSuffix(tip.UniqueName, "/trunk") {
			continue
		}
		burl, curl, err := uniqueNameURLs(tip.UniqueName)
		if err != nil {
			errs = append(errs, PublishBranchError{tip.UniqueName, err})
			log.Printf("error: %v\n", err)
			continue
		}
		log.Printf("----- %s\n", burl)
		if tip.Revision == "" {
			errs = append(errs, PublishBranchError{burl, fmt.Errorf("branch has no revisions")})
			log.Printf("error: branch has no revisions\n")
			continue
		}
		// Charm is published in the personal URL and in any explicitly
		// assigned official series.
		urls := []*charm.URL{curl}
		schema, name := curl.Schema, curl.Name
		for _, series := range tip.OfficialSeries {
			curl = &charm.URL{Schema: schema, Name: name, Series: series, Revision: -1}
			curl.Series = series
			curl.User = ""
			urls = append(urls, curl)
		}

		err = PublishBazaarBranch(store, urls, burl, tip.Revision)
		if err == ErrRedundantUpdate {
			continue
		}
		if err != nil {
			errs = append(errs, PublishBranchError{burl, err})
			log.Printf("error: %v\n", err)
		}
	}
	if errs != nil {
		return errs
	}
	return nil
}

// uniqueNameURLs returns the branch URL and the charm URL for the
// provided Launchpad branch unique name. The unique name must be
// in the form:
//
//     ~<user>/charms/<series>/<charm name>/trunk
//
// For testing purposes, if name has a prefix preceding a string in
// this format, the prefix is stripped out for computing the charm
// URL, and the unique name is returned unchanged as the branch URL.
func uniqueNameURLs(name string) (burl string, curl *charm.URL, err error) {
	u := strings.Split(name, "/")
	if len(u) > 5 {
		u = u[len(u)-5:]
		burl = name
	} else {
		burl = "lp:" + name
	}
	if len(u) < 5 || u[1] != "charms" || u[4] != "trunk" || len(u[0]) == 0 || u[0][0] != '~' {
		return "", nil, fmt.Errorf("unwanted branch name: %s", name)
	}
	curl, err = charm.ParseURL(fmt.Sprintf("cs:%s/%s/%s", u[0], u[2], u[3]))
	if err != nil {
		return "", nil, err
	}
	return burl, curl, nil
}
<|endoftext|>


<commit_before>// Package mail implements a parser for electronic mail messages as specified
// in RFC5322.
//
// We allow both CRLF and LF to be used in the input, possibly mixed.
package mail

import (
	"bytes"
	"encoding/base64"
	"errors"
	"strings"
	"time"
)

var benc = base64.URLEncoding

type Message struct {
	FullHeaders []Header // all headers
	OptHeaders  []Header // unprocessed headers

	MessageId   string
	Id          string
	Date        time.Time
	From        []Address
	Sender      Address
	ReplyTo     []Address
	To          []Address
	Cc          []Address
	Bcc         []Address
	Subject     string
	Comments    []string
	Keywords    []string

	InReply     []string
	References  []string

	Text        string
}

type Header struct {
	Key, Value string
}

func Parse(s []byte) (m Message, e error) {
	r, e := ParseRaw(s)
	if e != nil {
		return
	}
	return Process(r)
}

func Process(r RawMessage) (m Message, e error) {
	m.FullHeaders = []Header{}
	m.OptHeaders = []Header{}
	m.Text = string(r.Body) // TODO mime
	for _, rh := range r.RawHeaders {
		h := Header{string(rh.Key), string(rh.Value)}
		m.FullHeaders = append(m.FullHeaders, h)
		switch string(rh.Key) {
		case `Message-ID`:
			v := bytes.Trim(rh.Value, `<>`)
			m.MessageId = string(v)
			m.Id = benc.EncodeToString(v)
		case `In-Reply-To`:
			ids := strings.Split(string(rh.Value), ` `)
			for _, id := range ids {
				m.InReply = append(m.InReply, strings.Trim(id, `<> `))
			}
		case `References`:
			ids := strings.Split(string(rh.Value), ` `)
			for _, id := range ids {
				m.References = append(m.References, strings.Trim(id, `<> `))
			}
		case `Date`:
			m.Date = ParseDate(string(rh.Value))
		case `From`:
			m.From, e = parseAddressList(rh.Value)
		case `Sender`:
			m.Sender, e = ParseAddress(rh.Value)
		case `Reply-To`:
			m.ReplyTo, e = parseAddressList(rh.Value)
		case `To`:
			m.To, e = parseAddressList(rh.Value)
		case `Cc`:
			m.Cc, e = parseAddressList(rh.Value)
		case `Bcc`:
			m.Bcc, e = parseAddressList(rh.Value)
		case `Subject`:
			m.Subject = string(rh.Value)
		case `Comments`:
			m.Comments = append(m.Comments, string(rh.Value))
		case `Keywords`:
			ks := strings.Split(string(rh.Value), ",")
			for _, k := range ks {
				m.Keywords = append(m.Keywords, strings.TrimSpace(k))
			}
		default:
			m.OptHeaders = append(m.OptHeaders, h)
		}
		if e != nil {
			return
		}
	}
	if m.Sender == nil && len(m.From) > 0 {
		m.Sender = m.From[0]
	}
	return
}

type RawHeader struct {
	Key, Value []byte
}

type RawMessage struct {
	RawHeaders []RawHeader
	Body       []byte
}

func isWSP(b byte) bool {
	return b == ' ' || b == '\t'
}

func ParseRaw(s []byte) (m RawMessage, e error) {
	// parser states
	const (
		READY = iota
		HKEY
		HVWS
		HVAL
	)

	const (
		CR = '\r'
		LF = '\n'
	)
	CRLF := []byte{CR, LF}

	state := READY
	kstart, kend, vstart := 0, 0, 0
	done := false

	m.RawHeaders = []RawHeader{}

	for i := 0; i < len(s); i++ {
		b := s[i]
		switch state {
		case READY:
			if b == CR && i < len(s)-1 && s[i+1] == LF {
				// we are at the beginning of an empty header
				m.Body = s[i+2:]
				done = true
				goto Done
			}
			if b == LF {
				m.Body = s[i+1:]
				done = true
				goto Done
			}
			// otherwise this character is the first in a header
			// key
			kstart = i
			state = HKEY
		case HKEY:
			if b == ':' {
				kend = i
				state = HVWS
			}
		case HVWS:
			if !isWSP(b) {
				vstart = i
				state = HVAL
			}
		case HVAL:
			if b == CR && i < len(s)-2 && s[i+1] == LF && !isWSP(s[i+2]) {
				v := bytes.Replace(s[vstart:i], CRLF, nil, -1)
				hdr := RawHeader{s[kstart:kend], v}
				m.RawHeaders = append(m.RawHeaders, hdr)
				state = READY
				i++
			} else if b == LF && i < len(s)-1 && !isWSP(s[i+1]) {
				v := bytes.Replace(s[vstart:i], CRLF, nil, -1)
				hdr := RawHeader{s[kstart:kend], v}
				m.RawHeaders = append(m.RawHeaders, hdr)
				state = READY
			}
		}
	}
Done:
	if !done {
		e = errors.New("unexpected EOF")
	}
	return
}
<commit_msg>Bugfix<commit_after>// Package mail implements a parser for electronic mail messages as specified
// in RFC5322.
//
// We allow both CRLF and LF to be used in the input, possibly mixed.
package mail

import (
	"bytes"
	"encoding/base64"
	"errors"
	"strings"
	"time"
)

var benc = base64.URLEncoding

type Message struct {
	FullHeaders []Header // all headers
	OptHeaders  []Header // unprocessed headers

	MessageId   string
	Id          string
	Date        time.Time
	From        []Address
	Sender      Address
	ReplyTo     []Address
	To          []Address
	Cc          []Address
	Bcc         []Address
	Subject     string
	Comments    []string
	Keywords    []string

	InReply     []string
	References  []string

	Text        string
}

type Header struct {
	Key, Value string
}

func Parse(s []byte) (m Message, e error) {
	r, e := ParseRaw(s)
	if e != nil {
		return
	}
	return Process(r)
}

func trunc(s string, i int) string {
	if len(s) < i {
		return s
	}
	return s[0:i]
}

func Process(r RawMessage) (m Message, e error) {
	m.FullHeaders = []Header{}
	m.OptHeaders = []Header{}
	m.Text = string(r.Body) // TODO mime
	for _, rh := range r.RawHeaders {
		h := Header{string(rh.Key), string(rh.Value)}
		m.FullHeaders = append(m.FullHeaders, h)
		switch string(rh.Key) {
		case `Message-ID`:
			v := bytes.Trim(rh.Value, `<>`)
			m.MessageId = string(v)
			m.Id = trunc(benc.EncodeToString(v), 64)
		case `In-Reply-To`:
			ids := strings.Split(string(rh.Value), ` `)
			for _, id := range ids {
				m.InReply = append(m.InReply, strings.Trim(id, `<> `))
			}
		case `References`:
			ids := strings.Split(string(rh.Value), ` `)
			for _, id := range ids {
				m.References = append(m.References, strings.Trim(id, `<> `))
			}
		case `Date`:
			m.Date = ParseDate(string(rh.Value))
		case `From`:
			m.From, e = parseAddressList(rh.Value)
		case `Sender`:
			m.Sender, e = ParseAddress(rh.Value)
		case `Reply-To`:
			m.ReplyTo, e = parseAddressList(rh.Value)
		case `To`:
			m.To, e = parseAddressList(rh.Value)
		case `Cc`:
			m.Cc, e = parseAddressList(rh.Value)
		case `Bcc`:
			m.Bcc, e = parseAddressList(rh.Value)
		case `Subject`:
			m.Subject = string(rh.Value)
		case `Comments`:
			m.Comments = append(m.Comments, string(rh.Value))
		case `Keywords`:
			ks := strings.Split(string(rh.Value), ",")
			for _, k := range ks {
				m.Keywords = append(m.Keywords, strings.TrimSpace(k))
			}
		default:
			m.OptHeaders = append(m.OptHeaders, h)
		}
		if e != nil {
			return
		}
	}
	if m.Sender == nil && len(m.From) > 0 {
		m.Sender = m.From[0]
	}
	return
}

type RawHeader struct {
	Key, Value []byte
}

type RawMessage struct {
	RawHeaders []RawHeader
	Body       []byte
}

func isWSP(b byte) bool {
	return b == ' ' || b == '\t'
}

func ParseRaw(s []byte) (m RawMessage, e error) {
	// parser states
	const (
		READY = iota
		HKEY
		HVWS
		HVAL
	)

	const (
		CR = '\r'
		LF = '\n'
	)
	CRLF := []byte{CR, LF}

	state := READY
	kstart, kend, vstart := 0, 0, 0
	done := false

	m.RawHeaders = []RawHeader{}

	for i := 0; i < len(s); i++ {
		b := s[i]
		switch state {
		case READY:
			if b == CR && i < len(s)-1 && s[i+1] == LF {
				// we are at the beginning of an empty header
				m.Body = s[i+2:]
				done = true
				goto Done
			}
			if b == LF {
				m.Body = s[i+1:]
				done = true
				goto Done
			}
			// otherwise this character is the first in a header
			// key
			kstart = i
			state = HKEY
		case HKEY:
			if b == ':' {
				kend = i
				state = HVWS
			}
		case HVWS:
			if !isWSP(b) {
				vstart = i
				state = HVAL
			}
		case HVAL:
			if b == CR && i < len(s)-2 && s[i+1] == LF && !isWSP(s[i+2]) {
				v := bytes.Replace(s[vstart:i], CRLF, nil, -1)
				hdr := RawHeader{s[kstart:kend], v}
				m.RawHeaders = append(m.RawHeaders, hdr)
				state = READY
				i++
			} else if b == LF && i < len(s)-1 && !isWSP(s[i+1]) {
				v := bytes.Replace(s[vstart:i], CRLF, nil, -1)
				hdr := RawHeader{s[kstart:kend], v}
				m.RawHeaders = append(m.RawHeaders, hdr)
				state = READY
			}
		}
	}
Done:
	if !done {
		e = errors.New("unexpected EOF")
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/mdlayher/goat/goat"
)

func main() {
	// Set up command line options
	test := flag.Bool("test", false, "Make goat start, and exit shortly after. Used for testing.")
	flag.Parse()

	// If test mode, trigger quit shortly after startup
	// Used for CI tests, so that we ensure goat starts up and is able to stop gracefully
	if *test {
		go func() {
			os.Setenv("GOAT_TEST", "1")
			fmt.Println(goat.App, ": launched in test mode")
			time.Sleep(2 * time.Second)

			fmt.Println(goat.App, ": exiting test mode")
			os.Exit(0)
		}()
	}

	// Launch manager via goroutine
	killChan := make(chan bool)
	exitChan := make(chan int)
	go goat.Manager(killChan, exitChan)

	// Gracefully handle termination via UNIX signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	signal.Notify(sigChan, syscall.SIGTERM)
	for sig := range sigChan {
		// Trigger manager shutdown
		fmt.Println(goat.App, ": caught signal:", sig)
		killChan <- true
		break
	}

	// Force terminate if signaled twice
	go func(sigChan chan os.Signal) {
		for sig := range sigChan {
			_ = sig
			fmt.Println(goat.App, ": force halting now!")
			os.Exit(1)
		}
	}(sigChan)

	// Exit with specified code from manager
	code := <-exitChan
	fmt.Println(goat.App, ": graceful shutdown complete")
	os.Exit(code)
}
<commit_msg>Test commit for Eliot<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/mdlayher/goat/goat"
)

func main() {
	// Set up command line options
	test := flag.Bool("test", false, "Make goat start, and exit shortly after. Used for testing.")
	flag.Parse()

	// If test mode, trigger quit shortly after startup
	// Used for CI tests, so that we ensure goat starts up and is able to stop gracefully
	if *test {
		go func() {
			// Set environment variable to indicate test mode
			os.Setenv("GOAT_TEST", "1")
			fmt.Println(goat.App, ": launched in test mode")
			time.Sleep(2 * time.Second)

			fmt.Println(goat.App, ": exiting test mode")
			os.Exit(0)
		}()
	}

	// Launch manager via goroutine
	killChan := make(chan bool)
	exitChan := make(chan int)
	go goat.Manager(killChan, exitChan)

	// Gracefully handle termination via UNIX signal
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	signal.Notify(sigChan, syscall.SIGTERM)
	for sig := range sigChan {
		// Trigger manager shutdown
		fmt.Println(goat.App, ": caught signal:", sig)
		killChan <- true
		break
	}

	// Force terminate if signaled twice
	go func(sigChan chan os.Signal) {
		for sig := range sigChan {
			_ = sig
			fmt.Println(goat.App, ": force halting now!")
			os.Exit(1)
		}
	}(sigChan)

	// Exit with specified code from manager
	code := <-exitChan
	fmt.Println(goat.App, ": graceful shutdown complete")
	os.Exit(code)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"io/ioutil"
	"encoding/json"
	"strings"
	"github.com/line/line-bot-sdk-go/linebot"
	"github.com/go-redis/redis"
	"strconv"
	// "time"
)

type device struct {
	// Gps_num float32 `json:"gps_num"`
	// App string `json:"app"`
	// Gps_alt float32 `json:"gps_alt"`
	// Fmt_opt int `json:"fmt_opt"`
	// Device string `json:"device"`
	// S_d2 float32 `json:"s_d2"`
	S_d0 float32 `json:"s_d0"`
	// S_d1 float32 `json:"s_d1"`
	S_h0 float32 `json:"s_h0"`
	SiteName string `json:"SiteName"`
	// Gps_fix float32 `json:"gps_fix"`
	// Ver_app string `json:"ver_app"`
	Gps_lat float32 `json:"gps_lat"`
	S_t0 float32 `json:"s_t0"`
	Timestamp string `json:"timestamp"`
	Gps_lon float32 `json:"gps_lon"`
	// Date string `json:"date"`
	// Tick float32 `json:"tick"`
	Device_id string `json:"device_id"`
	// S_1 float32 `json:"s_1"`
	// S_0 float32 `json:"s_0"`
	// S_3 float32 `json:"s_3"`
	// S_2 float32 `json:"s_2"`
	// Ver_format string `json:"ver_format"`
	// Time string `json:"time"`
}

type airbox struct {
	Source string `json:"source"`
	Feeds []device `json:"feeds"`
	Version string `json:"version"`
	Num_of_records int `json:"num_of_records"`
}

var bot *linebot.Client
var airbox_json airbox
var	client=redis.NewClient(&redis.Options{
		Addr:"hipposerver.ddns.net:6379",
		Password:"",
		DB:0,
	})

func main() {
	url := "https://data.lass-net.org/data/last-all-airbox.json"
	req, _ := http.NewRequest("GET", url, nil)
	res, _ := http.DefaultClient.Do(req)
	defer res.Body.Close()
	body, _ := ioutil.ReadAll(res.Body)
	errs := json.Unmarshal(body, &airbox_json)
	if errs != nil {
		fmt.Println(errs)
	}

	// fmt.Println(airbox_json)
	var err error
	bot, err = linebot.New(os.Getenv("ChannelSecret"), os.Getenv("ChannelAccessToken"))
	log.Println("Bot:", bot, " err:", err)
	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)

	// t:=time.Now()
	// _, min, _:=t.Clock()
	// if min==46{
	// 	pushmessage()
	// }
}
// func pushmessage(){
// 	_,err:=bot.PushMessage("U3617adbdd46283d7e859f36302f4f471", linebot.NewTextMessage("hi!")).Do()
// 	if err!=nil{
// 		panic(err)
// 	}
// }
func callbackHandler(w http.ResponseWriter, r *http.Request) {
	events, err := bot.ParseRequest(r)

	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}

	for _, event := range events {
		if event.Type == linebot.EventTypeMessage {
			switch message := event.Message.(type) {
			case *linebot.TextMessage:
				var txtmessage string
				inText := strings.ToLower(message.Text)
				if strings.Contains(inText,"訂閱"){
					userID:=event.Source.UserID
					// pong, _ := client.Ping().Result()
					// txtmessage=pong
					for i:=0; i<len(airbox_json.Feeds); i++ {
						if strings.Contains(inText,strings.ToLower(airbox_json.Feeds[i].Device_id)) {
							val, err:=client.Get(airbox_json.Feeds[i].Device_id).Result()
							if err!=nil{
								client.Set(airbox_json.Feeds[i].Device_id,userID,0)
								txtmessage="訂閱成功!"
								break
							}
							stringSlice:=strings.Split(val,",")
							if stringInSlice(userID,stringSlice){
								txtmessage="您已訂閱過此ID!"
								break
							} else{
								val=val+","+userID
								client.Set(airbox_json.Feeds[i].Device_id,val,0)
								txtmessage="訂閱成功!"
								break
							}
						}
					}
				} else{
					for i:=0; i<len(airbox_json.Feeds); i++ {
						if strings.Contains(inText,strings.ToLower(airbox_json.Feeds[i].Device_id)) {
							txtmessage="Device_id: "+airbox_json.Feeds[i].Device_id+"\n"
							txtmessage=txtmessage+"Site Name: "+airbox_json.Feeds[i].SiteName+"\n"
							txtmessage=txtmessage+"Location: ("+strconv.FormatFloat(float64(airbox_json.Feeds[i].Gps_lon),'f',3,64)+","+strconv.FormatFloat(float64(airbox_json.Feeds[i].Gps_lat),'f',3,64)+")"+"\n"
							txtmessage=txtmessage+"Timestamp: "+airbox_json.Feeds[i].Timestamp+"\n"
							txtmessage=txtmessage+"PM2.5: "+strconv.FormatFloat(float64(airbox_json.Feeds[i].S_d0),'f',0,64)+"\n"
							txtmessage=txtmessage+"Humidity: "+strconv.FormatFloat(float64(airbox_json.Feeds[i].S_h0),'f',0,64)+"\n"
							txtmessage=txtmessage+"Temperature: "+strconv.FormatFloat(float64(airbox_json.Feeds[i].S_t0),'f',0,64)
							break
						}
					}
				}
				if len(txtmessage)==0{
					txtmessage="Sorry! No this device ID, please check again."
				}
				if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(txtmessage)).Do(); err != nil {
					log.Print(err)
				}
			}
		}
	}
}

func stringInSlice(a string, list []string) bool {
    for _, b := range list {
        if b == a {
            return true
        }
    }
    return false
}
<commit_msg>test push<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"io/ioutil"
	"encoding/json"
	"strings"
	"github.com/line/line-bot-sdk-go/linebot"
	"github.com/go-redis/redis"
	"strconv"
	"time"
)

type device struct {
	// Gps_num float32 `json:"gps_num"`
	// App string `json:"app"`
	// Gps_alt float32 `json:"gps_alt"`
	// Fmt_opt int `json:"fmt_opt"`
	// Device string `json:"device"`
	// S_d2 float32 `json:"s_d2"`
	S_d0 float32 `json:"s_d0"`
	// S_d1 float32 `json:"s_d1"`
	S_h0 float32 `json:"s_h0"`
	SiteName string `json:"SiteName"`
	// Gps_fix float32 `json:"gps_fix"`
	// Ver_app string `json:"ver_app"`
	Gps_lat float32 `json:"gps_lat"`
	S_t0 float32 `json:"s_t0"`
	Timestamp string `json:"timestamp"`
	Gps_lon float32 `json:"gps_lon"`
	// Date string `json:"date"`
	// Tick float32 `json:"tick"`
	Device_id string `json:"device_id"`
	// S_1 float32 `json:"s_1"`
	// S_0 float32 `json:"s_0"`
	// S_3 float32 `json:"s_3"`
	// S_2 float32 `json:"s_2"`
	// Ver_format string `json:"ver_format"`
	// Time string `json:"time"`
}

type airbox struct {
	Source string `json:"source"`
	Feeds []device `json:"feeds"`
	Version string `json:"version"`
	Num_of_records int `json:"num_of_records"`
}

var bot *linebot.Client
var airbox_json airbox
var	client=redis.NewClient(&redis.Options{
		Addr:"hipposerver.ddns.net:6379",
		Password:"",
		DB:0,
	})

func main() {
	url := "https://data.lass-net.org/data/last-all-airbox.json"
	req, _ := http.NewRequest("GET", url, nil)
	res, _ := http.DefaultClient.Do(req)
	defer res.Body.Close()
	body, _ := ioutil.ReadAll(res.Body)
	errs := json.Unmarshal(body, &airbox_json)
	if errs != nil {
		fmt.Println(errs)
	}

	// fmt.Println(airbox_json)
	var err error
	bot, err = linebot.New(os.Getenv("ChannelSecret"), os.Getenv("ChannelAccessToken"))
	log.Println("Bot:", bot, " err:", err)
	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)

}
func pushmessage(){
	_,err:=bot.PushMessage("U3617adbdd46283d7e859f36302f4f471", linebot.NewTextMessage("hi!")).Do()
	if err!=nil{
		panic(err)
	}
}
func callbackHandler(w http.ResponseWriter, r *http.Request) {
	t:=time.Now()
	_, min, _:=t.Clock()
	if min==0{
		pushmessage()
	}
	events, err := bot.ParseRequest(r)

	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}

	for _, event := range events {
		if event.Type == linebot.EventTypeMessage {
			switch message := event.Message.(type) {
			case *linebot.TextMessage:
				var txtmessage string
				inText := strings.ToLower(message.Text)
				if strings.Contains(inText,"訂閱"){
					userID:=event.Source.UserID
					// pong, _ := client.Ping().Result()
					// txtmessage=pong
					for i:=0; i<len(airbox_json.Feeds); i++ {
						if strings.Contains(inText,strings.ToLower(airbox_json.Feeds[i].Device_id)) {
							val, err:=client.Get(airbox_json.Feeds[i].Device_id).Result()
							if err!=nil{
								client.Set(airbox_json.Feeds[i].Device_id,userID,0)
								txtmessage="訂閱成功!"
								break
							}
							stringSlice:=strings.Split(val,",")
							if stringInSlice(userID,stringSlice){
								txtmessage="您已訂閱過此ID!"
								break
							} else{
								val=val+","+userID
								client.Set(airbox_json.Feeds[i].Device_id,val,0)
								txtmessage="訂閱成功!"
								break
							}
						}
					}
				} else{
					for i:=0; i<len(airbox_json.Feeds); i++ {
						if strings.Contains(inText,strings.ToLower(airbox_json.Feeds[i].Device_id)) {
							txtmessage="Device_id: "+airbox_json.Feeds[i].Device_id+"\n"
							txtmessage=txtmessage+"Site Name: "+airbox_json.Feeds[i].SiteName+"\n"
							txtmessage=txtmessage+"Location: ("+strconv.FormatFloat(float64(airbox_json.Feeds[i].Gps_lon),'f',3,64)+","+strconv.FormatFloat(float64(airbox_json.Feeds[i].Gps_lat),'f',3,64)+")"+"\n"
							txtmessage=txtmessage+"Timestamp: "+airbox_json.Feeds[i].Timestamp+"\n"
							txtmessage=txtmessage+"PM2.5: "+strconv.FormatFloat(float64(airbox_json.Feeds[i].S_d0),'f',0,64)+"\n"
							txtmessage=txtmessage+"Humidity: "+strconv.FormatFloat(float64(airbox_json.Feeds[i].S_h0),'f',0,64)+"\n"
							txtmessage=txtmessage+"Temperature: "+strconv.FormatFloat(float64(airbox_json.Feeds[i].S_t0),'f',0,64)
							break
						}
					}
				}
				if len(txtmessage)==0{
					txtmessage="Sorry! No this device ID, please check again."
				}
				if _, err = bot.ReplyMessage(event.ReplyToken, linebot.NewTextMessage(txtmessage)).Do(); err != nil {
					log.Print(err)
				}
			}
		}
	}
}

func stringInSlice(a string, list []string) bool {
    for _, b := range list {
        if b == a {
            return true
        }
    }
    return false
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"
)

const (
	PollNum  = 5
	PollTime = 1000
)

var st state

type state struct {
	files map[string]file
}

type file struct {
	Data     string
	Modified time.Time
}

type response struct {
	Updates  map[string]file
	Requests []string
}

func sync(w http.ResponseWriter, r *http.Request) {
	files := make(map[string]file)

	err := json.NewDecoder(r.Body).Decode(&files)
	if err != nil {
		log.Fatal(err)
	}

	updates, requests := poll(files)

	resp := response{Updates: updates, Requests: requests}
	js, err := json.Marshal(resp)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	fmt.Println(resp)

	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func poll(files map[string]file) (map[string]file, []string) {
	updates := make(map[string]file)
	requests := make([]string, 0)

	for i := 0; i < PollNum; i++ {
		for key, serverFile := range st.files {
			clientFile, prs := files[key]
			if !prs {
				clientFile = file{Data: "", Modified: time.Unix(0, 0)}
			}

			if serverFile.Modified.After(clientFile.Modified) {
				updates[key] = serverFile
			} else if serverFile.Modified.Before(clientFile.Modified) {
				requests = append(requests, key)
			}
		}

		for key, _ := range files {
			_, prs := st.files[key]
			if !prs {
				requests = append(requests, key)
			}
		}

		if len(updates) > 0 || len(requests) > 0 {
			break
		} else {
			time.Sleep(PollTime * time.Millisecond)
		}
	}

	return updates, requests
}

func main() {
	st = state{files: make(map[string]file)}
	//http.HandleFunc(`/save/{rest:[a-zA-Z0-9=\-\/\.]+}`, saved)
	http.HandleFunc("/sync", sync)
	http.ListenAndServe(":8000", nil)
}
<commit_msg>update url<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"
)

const (
	PollNum  = 5
	PollTime = 1000
)

var st state

type state struct {
	files map[string]file
}

type file struct {
	Data     string
	Modified time.Time
}

type response struct {
	Updates  map[string]file
	Requests []string
}

func sync(w http.ResponseWriter, r *http.Request) {
	files := make(map[string]file)
	err := json.NewDecoder(r.Body).Decode(&files)
	if err != nil {
		log.Fatal(err)
	}

	updates, requests := poll(files)

	resp := response{Updates: updates, Requests: requests}
	js, err := json.Marshal(resp)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	fmt.Println(resp)

	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func poll(files map[string]file) (map[string]file, []string) {
	updates := make(map[string]file)
	requests := make([]string, 0)

	for i := 0; i < PollNum; i++ {
		for key, serverFile := range st.files {
			clientFile, prs := files[key]
			if !prs {
				clientFile = file{Data: "", Modified: time.Unix(0, 0)}
			}

			if serverFile.Modified.After(clientFile.Modified) {
				updates[key] = serverFile
			} else if serverFile.Modified.Before(clientFile.Modified) {
				requests = append(requests, key)
			}
		}

		for key, _ := range files {
			_, prs := st.files[key]
			if !prs {
				requests = append(requests, key)
			}
		}

		if len(updates) > 0 || len(requests) > 0 {
			break
		} else {
			time.Sleep(PollTime * time.Millisecond)
		}
	}

	return updates, requests
}

func update(w http.ResponseWriter, r *http.Request) {
	files := make(map[string]file)
	err := json.NewDecoder(r.Body).Decode(&files)
	if err != nil {
		log.Fatal(err)
	}

	for key, value := range files {
		st.files[key] = value
	}

	w.Write([]byte(""))
}

func main() {
	st = state{files: make(map[string]file)}
	http.HandleFunc("/sync", sync)
	http.HandleFunc("/update", update)
	http.ListenAndServe(":8000", nil)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
	"time"
	"net/http"
	"encoding/json"
	"github.com/rancher/go-rancher/client"
	"os"
)

type Rancher struct {
	Url        string `json:"url"`
	AccessKey  string `json:"access_key"`
	SecretKey  string `json:"secret_key"`
	Service    string `json:"service"`
	Image      string `json:"docker_image"`
	StartFirst bool   `json:"start_first"`
	Confirm    bool   `json:"confirm"`
	Timeout    int    `json:"timeout"`
}

func main() {
	fmt.Printf("Started upgrade server\n")
	http.HandleFunc("/upgrade/", requestSuccess)
	http.ListenAndServe(":8080", nil)
}

func requestSuccess(w http.ResponseWriter, r *http.Request){
	vargs := Rancher{
		StartFirst: false,
		Timeout:    30,
	}

	err := json.NewDecoder(r.Body).Decode(&vargs)
	if err != nil {
		fmt.Printf(err.Error())
		Error(w, err.Error(), 400)
		return
	} else {
		if len(vargs.Url) == 0 || len(vargs.AccessKey) == 0 || len(vargs.SecretKey) == 0 {
			Error(w, "Rancher credentials not set", 400)
			return
		}
		go upgradeRancher(vargs)
		w.WriteHeader(200)
		return
	}
}

func upgradeRancher(vargs Rancher) {
	var wantedImage, wantedVer string
	if !strings.HasPrefix(vargs.Image, "docker:") {
		vargs.Image = fmt.Sprintf("docker:%s", vargs.Image)
		parts := strings.Split(vargs.Image, ":")
		wantedImage = parts[1]
		wantedVer = parts[2]
	}

	var wantedService, wantedStack string
	if strings.Contains(vargs.Service, "/") {
		parts := strings.SplitN(vargs.Service, "/", 2)
		wantedStack = parts[0]
		wantedService = parts[1]
	} else if vargs.Service != "" {
		wantedService = vargs.Service
	}

	rancher, err := client.NewRancherClient(&client.ClientOpts{
		Url:       vargs.Url,
		AccessKey: vargs.AccessKey,
		SecretKey: vargs.SecretKey,
	})

	if err != nil {
		fmt.Printf("Failed to create rancher client: %s\n", err)
		return
	}

	var stackId string
	if wantedStack != "" {
		environments, err := rancher.Environment.List(&client.ListOpts{})
		if err != nil {
			fmt.Printf("Failed to list rancher environments: %s\n", err)
			return
		}

		for _, env := range environments.Data {
			if env.Name == wantedStack {
				stackId = env.Id
			}
		}

		if stackId == "" {
			fmt.Printf("Unable to find stack %s\n", wantedStack)
			return
		}
	}

	services, err := rancher.Service.List(&client.ListOpts{})
	if err != nil {
		fmt.Printf("Failed to list rancher services: %s\n", err)
		return
	}

	var upgradeLabel = os.Getenv("UPGRADE_LABEL")

	var foundImage, foundVer string
	for services != nil {
		for _, svc := range services.Data {
			if svc.LaunchConfig != nil {
				if _, ok := svc.LaunchConfig.Labels[upgradeLabel]; ok {
					parts := strings.Split(svc.LaunchConfig.ImageUuid, ":")
					foundImage = parts[1]
					foundVer = parts[2]
					if foundImage == wantedImage && ((foundVer < wantedVer) || (wantedVer == "latest")) {
						fmt.Printf("Trying to upgrade...\n")
						err := doUpgrade(vargs, svc, rancher)
						if err != nil {
							fmt.Printf(err.Error())
						}
						if (vargs.Confirm) {
							fmt.Printf("Trying to confirm...\n")
							err := confirmUpgrade(vargs, svc, rancher)
							if err != nil {
								fmt.Printf(err.Error())
							}
						}
						continue
					}
					if svc.Name == wantedService && ((wantedStack != "" && svc.EnvironmentId == stackId) || wantedStack == "") {
						err := doUpgrade(vargs, svc, rancher)
						if err != nil {
							fmt.Printf(err.Error())
						}
						if vargs.Confirm {
							err := confirmUpgrade(vargs, svc, rancher)
							if err != nil {
								fmt.Printf(err.Error())
							}
						}
					}
				}
			}
		}
		services, err = services.Next()
	}
}

func doUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	service.LaunchConfig.ImageUuid = vargs.Image
	upgrade := &client.ServiceUpgrade{}
	upgrade.InServiceStrategy = &client.InServiceUpgradeStrategy{
		LaunchConfig:           service.LaunchConfig,
		SecondaryLaunchConfigs: service.SecondaryLaunchConfigs,
		StartFirst:             vargs.StartFirst,
	}
	upgrade.ToServiceStrategy = &client.ToServiceUpgradeStrategy{}

	_, err := rancher.Service.ActionUpgrade(&service, upgrade)
	if err != nil {
		fmt.Printf("Unable to upgrade service %s: %s\n", vargs.Service, err)
	} else {
		fmt.Printf("Upgraded %s to %s\n", service.Name, vargs.Image)
	}
	return err
}

func confirmUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	srv, err := retry(func() (interface{}, error) {
		s, e := rancher.Service.ById(service.Id)
		if e != nil {
			return nil, e
		}
		if s.State != "upgraded" {
			return nil, fmt.Errorf("Service not upgraded: %s", s.State)
		}
		return s, nil
	}, time.Duration(vargs.Timeout)*time.Second, 3*time.Second)

	if err != nil {
		fmt.Printf("Error waiting for service upgrade to complete: %s", err)
		return err
	}

	_, err = rancher.Service.ActionFinishupgrade(srv.(*client.Service))
	if err != nil {
		fmt.Printf("Unable to finish upgrade %s: %s\n", vargs.Service, err)
		return err
	}
	fmt.Printf("Finished upgrade %s\n", vargs.Service)
	return err
}

type retryFunc func() (interface{}, error)

func retry(f retryFunc, timeout time.Duration, interval time.Duration) (interface{}, error) {
	finish := time.After(timeout)
	for {
		result, err := f()
		if err == nil {
			return result, nil
		}
		select {
		case <-finish:
			return nil, err
		case <-time.After(interval):
		}
	}
}

func Error(w http.ResponseWriter, error string, code int) {
	w.Header().Set("Content Type,", "text/plain; charset=UTF-8")
	w.WriteHeader(code)
	fmt.Fprint(w, error)
}<commit_msg>removing one-off support, adding slack notifications<commit_after>package main

import (
	"fmt"
	"strings"
	"time"
	"net/http"
	"encoding/json"
	"github.com/rancher/go-rancher/client"
	"os"
	"github.com/ashwanthkumar/slack-go-webhook"
)

type Rancher struct {
	Url        string `json:"url"`
	AccessKey  string `json:"access_key"`
	SecretKey  string `json:"secret_key"`
	Service    string `json:"service"`
	Image      string `json:"docker_image"`
	StartFirst bool   `json:"start_first"`
	Confirm    bool   `json:"confirm"`
	Timeout    int    `json:"timeout"`
}

func main() {
	fmt.Println("Started upgrade server")
	http.HandleFunc("/upgrade/", requestSuccess)
	http.ListenAndServe(":8080", nil)
}

func requestSuccess(w http.ResponseWriter, r *http.Request){
	vargs := Rancher{
		StartFirst: false,
		Timeout:    30,
	}

	err := json.NewDecoder(r.Body).Decode(&vargs)
	if err != nil {
		fmt.Printf(err.Error())
		Error(w, err.Error(), 400)
		return
	} else {
		if len(vargs.Url) == 0 || len(vargs.AccessKey) == 0 || len(vargs.SecretKey) == 0 {
			Error(w, "Rancher credentials not set", 400)
			return
		}
		go upgradeRancher(vargs)
		w.WriteHeader(200)
		return
	}
}

func upgradeRancher(vargs Rancher) {
	var wantedImage, wantedVer string
	if !strings.HasPrefix(vargs.Image, "docker:") {
		vargs.Image = fmt.Sprintf("docker:%s", vargs.Image)
		parts := strings.Split(vargs.Image, ":")
		wantedImage = parts[1]
		wantedVer = parts[2]
	}
	parts := strings.Split(wantedImage, "/")
	vargs.Service = parts[1]

	rancher, err := client.NewRancherClient(&client.ClientOpts{
		Url:       vargs.Url,
		AccessKey: vargs.AccessKey,
		SecretKey: vargs.SecretKey,
	})

	if err != nil {
		fmt.Printf("Failed to create rancher client: %s\n", err)
		return
	}

	services, err := rancher.Service.List(&client.ListOpts{})
	if err != nil {
		fmt.Printf("Failed to list rancher services: %s\n", err)
		return
	}

	var upgradeLabel = os.Getenv("UPGRADE_LABEL")

	var foundImage, foundVer string
	for services != nil {
		for _, svc := range services.Data {
			if svc.LaunchConfig != nil {
				if _, ok := svc.LaunchConfig.Labels[upgradeLabel]; ok {
					parts := strings.Split(svc.LaunchConfig.ImageUuid, ":")
					foundImage = parts[1]
					foundVer = parts[2]
					if foundImage == wantedImage && ((foundVer < wantedVer) || (wantedVer == "latest")) {
						fmt.Println("Trying to upgrade...")
						err := doUpgrade(vargs, svc, rancher)
						if err != nil {
							fmt.Println(err.Error())
						} else {
							if (vargs.Confirm) {
								fmt.Println("Trying to confirm...")
								err := confirmUpgrade(vargs, svc, rancher)
								if err != nil {
									fmt.Println("Unable to upgrade service %s: %s\n", vargs.Service, err.Error())
									message := fmt.Sprintf("Unable to confirm upgrade to `%s`. \nCheck status at https://rancher.connectedfleet.io", vargs.Service)
									slackMessage("danger", message)
								} else {
									fmt.Printf("Upgraded %s to %s\n", svc.Name, vargs.Image)
									message := fmt.Sprintf("Successfully upgraded `%v`", vargs.Service)
									slackMessage("good", message)

								}
							}
						}
						continue
					}
				}
			}
		}
		services, err = services.Next()
	}
}

func doUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	service.LaunchConfig.ImageUuid = vargs.Image
	upgrade := &client.ServiceUpgrade{}
	upgrade.InServiceStrategy = &client.InServiceUpgradeStrategy{
		LaunchConfig:           service.LaunchConfig,
		SecondaryLaunchConfigs: service.SecondaryLaunchConfigs,
		StartFirst:             vargs.StartFirst,
	}
	upgrade.ToServiceStrategy = &client.ToServiceUpgradeStrategy{}
	_, err := rancher.Service.ActionUpgrade(&service, upgrade)
	return err
}

func confirmUpgrade(vargs Rancher, service client.Service, rancher *client.RancherClient) (error) {
	srv, err := retry(func() (interface{}, error) {
		s, e := rancher.Service.ById(service.Id)
		if e != nil {
			return nil, e
		}
		if s.State != "upgraded" {
			return nil, fmt.Errorf("Service not upgraded: %s\n", s.State)
		}
		return s, nil
	}, time.Duration(vargs.Timeout)*time.Second, 3*time.Second)
	if err != nil {
		return err
	}

	_, err = rancher.Service.ActionFinishupgrade(srv.(*client.Service))
	if err != nil {
		return err
	}
	fmt.Printf("Finished upgrade %s\n", vargs.Service)
	return err
}

type retryFunc func() (interface{}, error)

func retry(f retryFunc, timeout time.Duration, interval time.Duration) (interface{}, error) {
	finish := time.After(timeout)
	for {
		result, err := f()
		if err == nil {
			return result, nil
		}
		select {
		case <-finish:
			return nil, err
		case <-time.After(interval):
		}
	}
}

func Error(w http.ResponseWriter, error string, code int) {
	w.Header().Set("Content Type,", "text/plain; charset=UTF-8")
	w.WriteHeader(code)
	fmt.Fprint(w, error)
}

func slackMessage(status string, message string) {
	var webhookUrl = os.Getenv("SLACK_WEBHOOK")
	mkdwn := "text"
	var array[] *string
	array = append(array, &mkdwn )
	attachment := slack.Attachment {Color: &status, Text: &message, Mkdwn: array}
	payload := slack.Payload {
		Username: "rancher-updater-service",
		Attachments: []slack.Attachment{attachment},
	}
	printable, _ := json.Marshal(payload)
	fmt.Println(string(printable))
	err := slack.Send(webhookUrl, "", payload)
	if len(err) > 0 {
		fmt.Printf("error: %s\n", err)
	}
}<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/golang/groupcache"
	"github.com/jessevdk/go-flags"
	_ "github.com/mattn/go-sqlite3"
	"github.com/zenazn/goji"
	"github.com/zenazn/goji/web"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

var options struct {
	Port      uint16 `long:"port" description:"Server port" default:"8080"`
	Tilesets  string `long:"tilesets" description:"path to tilesets" default:"./tilesets"`
	CachePort uint16 `long:"cacheport" description:"GroupCache port" default:"8000"`
	CacheSize int64  `long:"cachesize" description:"Size of Cache (MB)" default:"10"`
	MaxAge    uint   `long:"max_age" description:"Response max-age duration (seconds)" default:"3600"`
}

var (
	pool         *groupcache.HTTPPool
	cache        *groupcache.Group
	connections  map[string]*sql.DB
	imageQueries map[string]*sql.Stmt
	contentTypes map[string]string
	blankPNG     []byte
	cacheSince   = time.Now().Format(http.TimeFormat)
)

func main() {
	_, err := flags.ParseArgs(&options, os.Args)
	if err != nil {
		log.Fatal(err)
	}
	blankPNG, _ = ioutil.ReadFile("blank.png")

	connections = make(map[string]*sql.DB)
	imageQueries = make(map[string]*sql.Stmt)
	contentTypes = make(map[string]string)
	tilesets, _ := filepath.Glob(path.Join(options.Tilesets, "*.mbtiles"))

	for _, filename := range tilesets {
		_, service := filepath.Split(filename)
		service = strings.Split(service, ".")[0]

		fmt.Println("Service: ", service)

		db, err := sql.Open("sqlite3", filename)
		if err != nil {
			log.Fatal(err)
		}
		defer db.Close()
		connections[service] = db

		stmt, err := db.Prepare("select tile_data from tiles where zoom_level = ? and tile_column = ? and tile_row = ?")
		if err != nil {
			log.Fatal(err)
		}
		defer stmt.Close()
		imageQueries[service] = stmt

		//query a sample tile to determine if png or jpg, since metadata from tilemill doesn't give this to us
		var tileData []byte
		err = db.QueryRow("select tile_data from images limit 1").Scan(&tileData)
		if err != nil {
			log.Fatal(err)
		}
		contentTypes[service] = http.DetectContentType(tileData)
	}

	pool = groupcache.NewHTTPPool(fmt.Sprintf("http://127.0.0.1:%v", options.CachePort))
	cache = groupcache.NewGroup("TileCache", options.CacheSize*1048576, groupcache.GetterFunc(
		func(ctx groupcache.Context, key string, dest groupcache.Sink) error {
			pathParams := strings.Split(key, "/")
			service := pathParams[1]
			yParams := strings.Split(pathParams[4], ".")
			z, _ := strconv.ParseUint(pathParams[2], 0, 64)
			x, _ := strconv.ParseUint(pathParams[3], 0, 64)
			y, _ := strconv.ParseUint(yParams[0], 0, 64)

			//flip y to match TMS spec
			y = (1 << z) - 1 - y

			var stmt *sql.Stmt
			stmt = imageQueries[service]

			var tileData []byte
			err := stmt.QueryRow(uint8(z), uint16(x), uint16(y)).Scan(&tileData)
			if err != nil {
				if err != sql.ErrNoRows {
					log.Fatal(err)
				}
			}
			dest.SetBytes(tileData)
			return nil
		}))

	goji.Get("/services", ListServices)
	goji.Get("/:service/:z/:x/:filename", GetTile)
	flag.Set("bind", fmt.Sprintf(":%v", options.Port))
	goji.Serve()
}

type ServiceInfo struct {
	URI       string `json:"uri"`
	ImageType string `json:"imageType"`
}

func ListServices(c web.C, w http.ResponseWriter, r *http.Request) {
	services := make([]ServiceInfo, len(imageQueries))
	i := 0
	for service, _ := range imageQueries {
		services[i] = ServiceInfo{
			URI:       fmt.Sprintf("/%s", service),
			ImageType: strings.Split(contentTypes[service], "/")[1],
		}
		i++
	}
	json, _ := json.Marshal(services)
	w.Write(json)
}

func GetTile(c web.C, w http.ResponseWriter, r *http.Request) {
	var (
		data        []byte
		contentType string
	)
	//TODO: validate x, y, z

	service := c.URLParams["service"]
	if _, exists := imageQueries[service]; !exists {
		http.Error(w, fmt.Sprintf("Service not found: %s", service), 404)
		return
	}

	key := r.URL.String()

	err := cache.Get(nil, key, groupcache.AllocatingByteSliceSink(&data))
	if err != nil {
		log.Println("Error fetching key", key)
		http.Error(w, fmt.Sprintf("Cache get failed for key: %s", key), 500)
		return
	}
	etag := fmt.Sprintf("%x", md5.Sum(data))

	if r.Header.Get("If-None-Match") == etag {
		w.WriteHeader(304)
		return
	}

	if len(data) <= 1 {
		data = blankPNG
		contentType = "image/png"
	} else {
		contentType = contentTypes[service]
	}

	w.Header().Add("Cache-Control", fmt.Sprintf("max-age=%v", options.MaxAge))
	w.Header().Add("Last-Modified", cacheSince)
	w.Header().Add("Content-Type", contentType)
	w.Header().Add("ETag", etag)
	w.Write(data)

	//TODO: gzip response
}
<commit_msg>Added stub for service tilejson<commit_after>package main

import (
	"crypto/md5"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/golang/groupcache"
	"github.com/jessevdk/go-flags"
	_ "github.com/mattn/go-sqlite3"
	"github.com/zenazn/goji"
	"github.com/zenazn/goji/web"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

var options struct {
	Port      uint16 `long:"port" description:"Server port" default:"8080"`
	Tilesets  string `long:"tilesets" description:"path to tilesets" default:"./tilesets"`
	CachePort uint16 `long:"cacheport" description:"GroupCache port" default:"8000"`
	CacheSize int64  `long:"cachesize" description:"Size of Cache (MB)" default:"10"`
	MaxAge    uint   `long:"max_age" description:"Response max-age duration (seconds)" default:"3600"`
}

var (
	pool  *groupcache.HTTPPool
	cache *groupcache.Group
	//TODO: consolidate these into a single map of structs!
	connections  map[string]*sql.DB
	imageQueries map[string]*sql.Stmt
	contentTypes map[string]string
	blankPNG     []byte
	cacheSince   = time.Now().Format(http.TimeFormat)
)

func main() {
	_, err := flags.ParseArgs(&options, os.Args)
	if err != nil {
		log.Fatal(err)
	}
	blankPNG, _ = ioutil.ReadFile("blank.png")

	connections = make(map[string]*sql.DB)
	imageQueries = make(map[string]*sql.Stmt)
	contentTypes = make(map[string]string)
	tilesets, _ := filepath.Glob(path.Join(options.Tilesets, "*.mbtiles"))

	for _, filename := range tilesets {
		_, service := filepath.Split(filename)
		service = strings.Split(service, ".")[0]

		fmt.Println("Service: ", service)

		db, err := sql.Open("sqlite3", filename)
		if err != nil {
			log.Fatal(err)
		}
		defer db.Close()
		connections[service] = db

		stmt, err := db.Prepare("select tile_data from tiles where zoom_level = ? and tile_column = ? and tile_row = ?")
		if err != nil {
			log.Fatal(err)
		}
		defer stmt.Close()
		imageQueries[service] = stmt

		//query a sample tile to determine if png or jpg, since metadata from tilemill doesn't give this to us
		var tileData []byte
		err = db.QueryRow("select tile_data from images limit 1").Scan(&tileData)
		if err != nil {
			log.Fatal(err)
		}
		contentTypes[service] = http.DetectContentType(tileData)
	}

	pool = groupcache.NewHTTPPool(fmt.Sprintf("http://127.0.0.1:%v", options.CachePort))
	cache = groupcache.NewGroup("TileCache", options.CacheSize*1048576, groupcache.GetterFunc(
		func(ctx groupcache.Context, key string, dest groupcache.Sink) error {
			pathParams := strings.Split(key, "/")
			service := pathParams[0]
			yParams := strings.Split(pathParams[3], ".")
			z, _ := strconv.ParseUint(pathParams[1], 0, 64)
			x, _ := strconv.ParseUint(pathParams[2], 0, 64)
			y, _ := strconv.ParseUint(yParams[0], 0, 64)
			//flip y to match TMS spec
			y = (1 << z) - 1 - y

			var stmt *sql.Stmt
			stmt = imageQueries[service]

			var tileData []byte
			err := stmt.QueryRow(uint8(z), uint16(x), uint16(y)).Scan(&tileData)
			if err != nil {
				if err != sql.ErrNoRows {
					log.Fatal(err)
				}
			}
			dest.SetBytes(tileData)
			return nil
		}))

	//TODO: add gzip
	//TODO: add 301s for all non-slash terminated routes
	goji.Get("/services/", ListServices)
	goji.Get("/services", http.RedirectHandler("/services/", http.StatusMovedPermanently))
	goji.Get("/services/:service", GetService)
	goji.Get("/services/:service/tiles/:z/:x/:filename", GetTile)
	//TODO:  goji.Get("/:service/grids/:z/:x/:filename", GetGrid) //return UTF8 grid

	flag.Set("bind", fmt.Sprintf(":%v", options.Port))
	goji.Serve()
}

type ServiceInfo struct {
	URI       string `json:"uri"`
	ImageType string `json:"imageType"`
}

func ListServices(c web.C, w http.ResponseWriter, r *http.Request) {
	services := make([]ServiceInfo, len(imageQueries))
	i := 0
	for service, _ := range imageQueries {
		services[i] = ServiceInfo{
			URI:       fmt.Sprintf("/%s", service),
			ImageType: strings.Split(contentTypes[service], "/")[1],
		}
		i++
	}
	json, _ := json.Marshal(services)
	w.Write(json)
}

func GetService(c web.C, w http.ResponseWriter, r *http.Request) {
	fmt.Println(c.URLParams)
	service := c.URLParams["service"]
	if _, exists := imageQueries[service]; !exists {
		http.Error(w, fmt.Sprintf("Service not found: %s", service), http.StatusNotFound)
		return
	}
	w.Write([]byte("TODO: Service info"))
}

func GetTile(c web.C, w http.ResponseWriter, r *http.Request) {
	var (
		data        []byte
		contentType string
	)
	//TODO: validate x, y, z

	service := c.URLParams["service"]
	if _, exists := imageQueries[service]; !exists {
		http.Error(w, fmt.Sprintf("Service not found: %s", service), http.StatusNotFound)
		return
	}

	// key := strings.TrimPrefix(r.URL.String(), "/services/")
	key := strings.Join([]string{c.URLParams["service"], c.URLParams["z"], c.URLParams["x"], c.URLParams["filename"]}, "/")
	fmt.Println("URL key: ", key)

	err := cache.Get(nil, key, groupcache.AllocatingByteSliceSink(&data))
	if err != nil {
		log.Println("Error fetching key", key)
		http.Error(w, fmt.Sprintf("Cache get failed for key: %s", key), http.StatusInternalServerError)
		return
	}
	etag := fmt.Sprintf("%x", md5.Sum(data))

	if r.Header.Get("If-None-Match") == etag {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	if len(data) <= 1 {
		data = blankPNG
		contentType = "image/png"
	} else {
		contentType = contentTypes[service]
	}

	w.Header().Add("Cache-Control", fmt.Sprintf("max-age=%v", options.MaxAge))
	w.Header().Add("Last-Modified", cacheSince)
	w.Header().Add("Content-Type", contentType)
	w.Header().Add("ETag", etag)
	w.Write(data)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/trace"
	"cloud.google.com/go/trace/traceutil"
	"golang.org/x/net/context"
)

func helloHandler(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "Hello world!")
}

type MetadataResponse struct {
	ProjectID          string
	NumericProjectID   string
	SysHostname        string
	Hostname           string
	InternalIP         string
	ExternalIP         string
	InstanceID         string
	InstanceName       string
	Zone               string
	InstanceAttributes map[string]string
	InstanceTags       []string
	ProjectAttributes  map[string]string
	Scopes             []string
}

var traceClient *trace.Client

func inspectHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte(fmt.Sprintf("Request from %s\n\n", r.RemoteAddr)))

	dump, err := httputil.DumpRequest(r, true)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write(dump)
}

func envHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Environment:\n\n%s", html.EscapeString(strings.Join(os.Environ(), "\n")))
}

func curlHandler(w http.ResponseWriter, r *http.Request) {
	url := r.PostFormValue("url")

	if url == "" {
		fmt.Fprint(w, `
<html>
<form method="POST">
URL: <input type="text" name="url" size="80">
<input type="submit" value="Fire!">
</form>
</html>`)
		return
	}

	span := traceClient.SpanFromRequest(r)
	defer span.Finish()
	span.SetLabel("custom/url", url)

	hc := traceutil.NewHTTPClient(traceClient, nil)
	req, _ := http.NewRequest("GET", url, nil)
	resp, err := hc.Do(req)
	if err != nil {
		fmt.Fprintf(w, "GET failed: %s", err)
		return
	}

	fmt.Fprintf(w, "Response to GET %s\n\n", url)
	dump, err := httputil.DumpResponse(resp, true)
	if err != nil {
		fmt.Fprintf(w, "DumpResponse failed: %s", err)
		return
	}
	fmt.Fprint(w, html.EscapeString(string(dump)))
}

func metadataHandler(w http.ResponseWriter, r *http.Request) {
	if !metadata.OnGCE() {
		io.WriteString(w, "not on gce")
		return
	}

	response := &MetadataResponse{
		InstanceAttributes: make(map[string]string),
		ProjectAttributes:  make(map[string]string),
	}

	v, _ := os.Hostname()
	response.SysHostname = v

	v, _ = metadata.ProjectID()
	response.ProjectID = v

	v, _ = metadata.NumericProjectID()
	response.NumericProjectID = v

	v, _ = metadata.ExternalIP()
	response.ExternalIP = v

	v, _ = metadata.InternalIP()
	response.InternalIP = v

	v, _ = metadata.Hostname()
	response.Hostname = v

	v, _ = metadata.InstanceID()
	response.InstanceID = v

	v, _ = metadata.InstanceName()
	response.InstanceName = v

	v, _ = metadata.Zone()
	response.Zone = v

	vs, _ := metadata.InstanceAttributes()
	for _, attr := range vs {
		if val, err := metadata.InstanceAttributeValue(attr); err == nil {
			response.InstanceAttributes[attr] = val
		}
	}

	vs, _ = metadata.InstanceTags()
	response.InstanceTags = vs

	vs, _ = metadata.ProjectAttributes()
	for _, attr := range vs {
		if val, err := metadata.ProjectAttributeValue(attr); err == nil {
			response.ProjectAttributes[attr] = val
		}
	}

	vs, _ = metadata.Scopes("default")
	response.Scopes = vs

	js, err := json.Marshal(response)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func main() {

	log.Print("starting")

	ctx := context.Background()
	if projectId, err := metadata.ProjectID(); err == nil {
		if traceClient, err = trace.NewClient(ctx, projectId); err != nil {
			log.Fatalf("cant create trace client: %s", err)
		}
		log.Printf("enabled tracing, projectid: %s", projectId)
	} else {
		log.Print("cant enable tracing, no metadata available (not in gce?)")
	}

	log.Print("start serving")

	http.Handle("/", traceutil.HTTPHandler(traceClient, helloHandler))
	http.Handle("/metadata", traceutil.HTTPHandler(traceClient, metadataHandler))
	http.Handle("/inspect", traceutil.HTTPHandler(traceClient, inspectHandler))
	http.Handle("/curl", traceutil.HTTPHandler(traceClient, curlHandler))
	http.Handle("/env", traceutil.HTTPHandler(traceClient, envHandler))
	http.ListenAndServe(":8000", nil)
}
<commit_msg>span from request 2<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"os"
	"strings"

	"cloud.google.com/go/compute/metadata"
	"cloud.google.com/go/trace"
	"cloud.google.com/go/trace/traceutil"
	"golang.org/x/net/context"
)

func helloHandler(w http.ResponseWriter, r *http.Request) {
	io.WriteString(w, "Hello world!")
}

type MetadataResponse struct {
	ProjectID          string
	NumericProjectID   string
	SysHostname        string
	Hostname           string
	InternalIP         string
	ExternalIP         string
	InstanceID         string
	InstanceName       string
	Zone               string
	InstanceAttributes map[string]string
	InstanceTags       []string
	ProjectAttributes  map[string]string
	Scopes             []string
}

var traceClient *trace.Client

func inspectHandler(w http.ResponseWriter, r *http.Request) {
	w.Write([]byte(fmt.Sprintf("Request from %s\n\n", r.RemoteAddr)))

	dump, err := httputil.DumpRequest(r, true)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Write(dump)
}

func envHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Environment:\n\n%s", html.EscapeString(strings.Join(os.Environ(), "\n")))
}

func curlHandler(w http.ResponseWriter, r *http.Request) {
	span := traceClient.SpanFromRequest(r)
	defer span.Finish()

	url := r.PostFormValue("url")

	if url == "" {
		fmt.Fprint(w, `
<html>
<form method="POST">
URL: <input type="text" name="url" size="80">
<input type="submit" value="Fire!">
</form>
</html>`)
		return
	}

	span.SetLabel("custom/url", url)

	hc := traceutil.NewHTTPClient(traceClient, nil)
	req, _ := http.NewRequest("GET", url, nil)

	r_span := span.NewRemoteChild(req)

	resp, err := hc.Do(req)
	r_span.Finish(trace.WithResponse(resp))

	if err != nil {
		fmt.Fprintf(w, "GET failed: %s", err)
		return
	}

	fmt.Fprintf(w, "Response to GET %s\n\n", url)
	dump, err := httputil.DumpResponse(resp, true)
	if err != nil {
		fmt.Fprintf(w, "DumpResponse failed: %s", err)
		return
	}
	fmt.Fprint(w, html.EscapeString(string(dump)))
}

func metadataHandler(w http.ResponseWriter, r *http.Request) {
	if !metadata.OnGCE() {
		io.WriteString(w, "not on gce")
		return
	}

	response := &MetadataResponse{
		InstanceAttributes: make(map[string]string),
		ProjectAttributes:  make(map[string]string),
	}

	v, _ := os.Hostname()
	response.SysHostname = v

	v, _ = metadata.ProjectID()
	response.ProjectID = v

	v, _ = metadata.NumericProjectID()
	response.NumericProjectID = v

	v, _ = metadata.ExternalIP()
	response.ExternalIP = v

	v, _ = metadata.InternalIP()
	response.InternalIP = v

	v, _ = metadata.Hostname()
	response.Hostname = v

	v, _ = metadata.InstanceID()
	response.InstanceID = v

	v, _ = metadata.InstanceName()
	response.InstanceName = v

	v, _ = metadata.Zone()
	response.Zone = v

	vs, _ := metadata.InstanceAttributes()
	for _, attr := range vs {
		if val, err := metadata.InstanceAttributeValue(attr); err == nil {
			response.InstanceAttributes[attr] = val
		}
	}

	vs, _ = metadata.InstanceTags()
	response.InstanceTags = vs

	vs, _ = metadata.ProjectAttributes()
	for _, attr := range vs {
		if val, err := metadata.ProjectAttributeValue(attr); err == nil {
			response.ProjectAttributes[attr] = val
		}
	}

	vs, _ = metadata.Scopes("default")
	response.Scopes = vs

	js, err := json.Marshal(response)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func main() {

	log.Print("starting")

	ctx := context.Background()
	if projectId, err := metadata.ProjectID(); err == nil {
		if traceClient, err = trace.NewClient(ctx, projectId); err != nil {
			log.Fatalf("cant create trace client: %s", err)
		}
		log.Printf("enabled tracing, projectid: %s", projectId)
	} else {
		log.Print("cant enable tracing, no metadata available (not in gce?)")
	}

	log.Print("start serving")

	http.Handle("/", traceutil.HTTPHandler(traceClient, helloHandler))
	http.Handle("/metadata", traceutil.HTTPHandler(traceClient, metadataHandler))
	http.Handle("/inspect", traceutil.HTTPHandler(traceClient, inspectHandler))
	http.Handle("/curl", traceutil.HTTPHandler(traceClient, curlHandler))
	http.Handle("/env", traceutil.HTTPHandler(traceClient, envHandler))
	http.ListenAndServe(":8000", nil)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"path/filepath"
	"strings"

	"github.com/facebookgo/grace/gracehttp"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	listen   string
	quotaDir string
	users    string
	timeout  time.Duration

	startTime      = time.Now()
	lastReloadTime = time.Now()
)

func loadQuotaFiles(quotaDir string) error {
	log.Printf("Loading configuration files from [%s]\n", quotaDir)

	glob := fmt.Sprintf("%s%c%s", quotaDir, filepath.Separator, "*.xml")
	files, _ := filepath.Glob(glob)
	if len(files) == 0 {
		return fmt.Errorf("no quota XML files found in [%s] - exiting", quotaDir)
	}

	for _, file := range files {
		loadQuotaFile(file)
	}
	return nil
}

func loadQuotaFile(file string) {
	fileName := filepath.Base(file)
	quotaName := strings.TrimSuffix(fileName, filepath.Ext(fileName))
	var browsers Browsers
	err := readConfig(file, &browsers)
	if err != nil {
		log.Printf("Failed to load configuration from [%s]: %v", fileName, err)
		return
	}
	updateQuota(quotaName, browsers)
	log.Printf("Loaded configuration from [%s]:\n%v\n", file, browsers)
}

func updateQuota(quotaName string, browsers Browsers) {
	confLock.Lock()
	defer confLock.Unlock()
	quota[quotaName] = browsers
	routes = appendRoutes(routes, &browsers)
	lastReloadTime = time.Now()
}

func init() {
	flag.StringVar(&listen, "listen", ":4444", "host and port to listen to")
	flag.StringVar(&quotaDir, "quotaDir", "quota", "quota directory")
	flag.StringVar(&users, "users", ".htpasswd", "htpasswd auth file path")
	flag.DurationVar(&timeout, "timeout", 300*time.Second, "session creation timeout in time.Duration format, e.g. 300s or 500ms")
	flag.Parse()
	log.Printf("Users file is [%s]\n", users)
	if err := loadQuotaFiles(quotaDir); err != nil {
		log.Fatalf("%v\n", err)
	}
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGHUP)
	go func() {
		for {
			<-sig
			loadQuotaFiles(quotaDir)
		}
	}()
}

func main() {
	gracehttp.Serve([]*http.Server{
		{Addr: listen, Handler: mux()},
	}...)
}
<commit_msg>Showing server startup error (fixes #61)<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"path/filepath"
	"strings"

	"github.com/facebookgo/grace/gracehttp"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	listen   string
	quotaDir string
	users    string
	timeout  time.Duration

	startTime      = time.Now()
	lastReloadTime = time.Now()
)

func loadQuotaFiles(quotaDir string) error {
	log.Printf("Loading configuration files from [%s]\n", quotaDir)

	glob := fmt.Sprintf("%s%c%s", quotaDir, filepath.Separator, "*.xml")
	files, _ := filepath.Glob(glob)
	if len(files) == 0 {
		return fmt.Errorf("no quota XML files found in [%s] - exiting", quotaDir)
	}

	for _, file := range files {
		loadQuotaFile(file)
	}
	return nil
}

func loadQuotaFile(file string) {
	fileName := filepath.Base(file)
	quotaName := strings.TrimSuffix(fileName, filepath.Ext(fileName))
	var browsers Browsers
	err := readConfig(file, &browsers)
	if err != nil {
		log.Printf("Failed to load configuration from [%s]: %v", fileName, err)
		return
	}
	updateQuota(quotaName, browsers)
	log.Printf("Loaded configuration from [%s]:\n%v\n", file, browsers)
}

func updateQuota(quotaName string, browsers Browsers) {
	confLock.Lock()
	defer confLock.Unlock()
	quota[quotaName] = browsers
	routes = appendRoutes(routes, &browsers)
	lastReloadTime = time.Now()
}

func init() {
	flag.StringVar(&listen, "listen", ":4444", "host and port to listen to")
	flag.StringVar(&quotaDir, "quotaDir", "quota", "quota directory")
	flag.StringVar(&users, "users", ".htpasswd", "htpasswd auth file path")
	flag.DurationVar(&timeout, "timeout", 300*time.Second, "session creation timeout in time.Duration format, e.g. 300s or 500ms")
	flag.Parse()
	log.Printf("Users file is [%s]\n", users)
	if err := loadQuotaFiles(quotaDir); err != nil {
		log.Fatalf("%v\n", err)
	}
	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGHUP)
	go func() {
		for {
			<-sig
			loadQuotaFiles(quotaDir)
		}
	}()
}

func main() {
	log.Fatal(gracehttp.Serve([]*http.Server{
		{Addr: listen, Handler: mux()},
	}...))
}
<|endoftext|>


<commit_before>package main

import (
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "runtime"
    "strconv"
    "strings"
    "time"

    "github.com/BurntSushi/toml"
    "github.com/akrennmair/gopcap"
    "github.com/nranchev/go-libGeoIP"
)

const Version = "0.2.0"

type Packet struct {
    ts      time.Time
    tuple   IpPortTuple
    payload []byte
}

type protocolType uint16

const (
    UnknownProtocol protocolType = iota
    HttpProtocol
    MysqlProtocol
    RedisProtocol
    PgsqlProtocol
)

var protocolNames = []string{"unknown", "http", "mysql", "redis", "pgsql"}

type tomlConfig struct {
    Interfaces    tomlInterfaces
    RunOptions    tomlRunOptions
    Protocols     map[string]tomlProtocol
    Procs         tomlProcs
    Elasticsearch tomlMothership
    Agent         tomlAgent
    Logging       tomlLogging
}

type tomlInterfaces struct {
    Device string
}

type tomlRunOptions struct {
    Uid int
    Gid int
}

type tomlLogging struct {
    Selectors []string
}

var _Config tomlConfig
var _ConfigMeta toml.MetaData
var _GeoLite *libgeo.GeoIP

func Bytes_Ipv4_Ntoa(bytes []byte) string {
    var strarr []string = make([]string, 4)
    for i, b := range bytes {
        strarr[i] = strconv.Itoa(int(b))
    }
    return strings.Join(strarr, ".")
}

func decodePktEth(datalink int, pkt *pcap.Packet) {
    defer RECOVER("decodePktEth exception")

    packet := new(Packet)
    var l2hlen int
    var eth_type uint16

    if datalink != pcap.LINKTYPE_LINUX_SLL {
        // we are just assuming ETH
        l2hlen = 14

        // bytes 12 and 13 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[12:14])

    } else {
        l2hlen = 16

        // bytes 14 and 15 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[14:16])
    }

    if eth_type != 0x800 {
        DEBUG("ip", "Not ipv4 packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+20 {
        DEBUG("ip", "Packet too short to be ethernet")
        return
    }

    // IP header
    iphl := int((uint16(pkt.Data[l2hlen]) & 0x0f) * 4)
    if len(pkt.Data) < l2hlen+iphl {
        DEBUG("ip", "Packet too short to be IP")
        return
    }
    iphdr := pkt.Data[l2hlen : l2hlen+iphl]

    //DEBUG("Packet timestamp: %s", pkt.Time)
    packet.ts = pkt.Time

    packet.tuple.Src_ip = Bytes_Ntohl(iphdr[12:16])
    packet.tuple.Dst_ip = Bytes_Ntohl(iphdr[16:20])

    mf := (uint8(iphdr[6]&0x20) != 0)
    frag_offset := (uint16(iphdr[7]) << 3) | (uint16(iphdr[6]) & 0x1F)

    if mf || frag_offset > 0 {
        DEBUG("ip", "Fragmented packets not yet supported")
        return
    }

    ip_length := int(Bytes_Ntohs(iphdr[2:4]))

    protocol := uint8(iphdr[9])
    if protocol != 6 {
        DEBUG("ip", "Not TCP packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+iphl+20 {
        DEBUG("ip", "Packet too short to be TCP")
        return
    }

    tcphl := int((uint16(pkt.Data[l2hlen+iphl+12]) >> 4) * 4)
    if tcphl > 20 && len(pkt.Data) < l2hlen+iphl+tcphl {
        DEBUG("ip", "Packet too short to contain TCP header")
        return
    }

    tcphdr := pkt.Data[l2hlen+iphl : l2hlen+iphl+tcphl]
    packet.tuple.Src_port = Bytes_Ntohs(tcphdr[0:2])
    packet.tuple.Dst_port = Bytes_Ntohs(tcphdr[2:4])

    //DEBUG(" %s:%d -> %s:%d", packet.tuple.src_ip, packet.tuple.src_port,
    //      packet.tuple.dst_ip, packet.tuple.dst_port)

    data_offset := (tcphdr[12] >> 4) * 4

    if l2hlen+iphl+int(data_offset) > l2hlen+ip_length {
        DEBUG("ip", "data_offset pointing outside of packet")
        return
    }

    packet.payload = pkt.Data[l2hlen+iphl+int(data_offset) : l2hlen+ip_length]

    FollowTcp(tcphdr, packet)
}

func main() {

    configfile := flag.String("c", "packetbeat.conf", "Configuration file")
    file := flag.String("I", "", "file")
    loop := flag.Int("l", 1, "Loop file. 0 - loop forever")
    debugSelectorsStr := flag.String("d", "", "Enable certain debug selectors")
    oneAtAtime := flag.Bool("O", false, "Read packets one at a time (press Enter)")
    toStdout := flag.Bool("e", false, "Output to stdout instead of syslog")
    topSpeed := flag.Bool("t", false, "Read packets as fast as possible, without sleeping")
    publishDisabled := flag.Bool("N", false, "Disable actual publishing for testing")
    printVersion := flag.Bool("version", false, "Print version and exit")

    flag.Parse()

    if *printVersion {
        fmt.Printf("Packetbeat version %s (%s)\n", Version, runtime.GOARCH)
    }

    debugSelectors := []string{}
    logLevel := LOG_DEBUG
    if len(*debugSelectorsStr) > 0 {
        debugSelectors = strings.Split(*debugSelectorsStr, ",")
    }

    var h *pcap.Pcap
    var err error

    if _ConfigMeta, err = toml.DecodeFile(*configfile, &_Config); err != nil {
        fmt.Printf("TOML config parsing failed on %s: %s. Exiting.\n", *configfile, err)
        return
    }

    if len(debugSelectors) == 0 {
        debugSelectors = _Config.Logging.Selectors
    }
    LogInit(logLevel, "", !*toStdout, debugSelectors)

    if !IS_DEBUG("stdlog") {
        // disable standard logging by default
        log.SetOutput(ioutil.Discard)
    }

    if *file != "" {
        h, err = pcap.Openoffline(*file)
        if h == nil {
            ERR("Openoffline(%s) failed: %s", *file, err)
            return
        }
    } else {
        device := _Config.Interfaces.Device

        h, err = pcap.Openlive(device, 65535, true, 0)
        if h == nil {
            ERR("pcap.Openlive failed(%s): %s", device, err)
            return
        }
    }
    defer h.Close()

    if err = DropPrivileges(); err != nil {
        CRIT(err.Error())
        return
    }

    filter := configToFilter(&_Config)
    if filter != "" {
        DEBUG("pcapfilter", "Installing filter '%s'", filter)
        err := h.Setfilter(filter)
        if err != nil {
            ERR("pcap.Setfilter failed: %s", err)
            return
        }
    }

    if err = Publisher.Init(*publishDisabled); err != nil {
        CRIT(err.Error())
        return
    }

    if err = procWatcher.Init(&_Config.Procs); err != nil {
        CRIT(err.Error())
        return
    }

    if err = TcpInit(); err != nil {
        CRIT(err.Error())
        return
    }

    datalink := h.Datalink()
    if datalink != pcap.LINKTYPE_ETHERNET && datalink != pcap.LINKTYPE_LINUX_SLL {
        WARN("Unsupported link type: %d", datalink)
    }

    _GeoLite, err = libgeo.Load("/usr/share/GeoIP/GeoIP.dat")
    if err != nil {
        WARN("Could not load GeoIP data: %s", err.Error())
    }

    counter := 0
    live := true
    loopCount := 1
    var lastPktTime *time.Time = nil
    for live {
        var pkt *pcap.Packet
        var res int32

        if *oneAtAtime {
            fmt.Println("Press enter to read packet")
            fmt.Scanln()
        }

        pkt, res = h.NextEx()
        switch res {
        case -1:
            ERR("pcap.NextEx() error: %s", h.Geterror())
            live = false
            continue
        case -2:
            DEBUG("pcapread", "End of file")
            loopCount += 1
            if *loop > 0 && loopCount > *loop {
                // give a bit of time to the publish goroutine
                // to flush
                time.Sleep(300 * time.Millisecond)
                live = false
            }
            // reopen the file
            h, err = pcap.Openoffline(*file)
            if h == nil {
                ERR("Openoffline(%s) failed: %s", *file, err)
                return
            }
            lastPktTime = nil
            continue
        case 0:
            // timeout
            continue
        }
        if res != 1 {
            panic(fmt.Sprintf("Unexpected return code from pcap.NextEx: %d", res))
        }

        if pkt == nil {
            panic("Nil packet despite res=1")
        }

        if *file != "" {
            if lastPktTime != nil && !*topSpeed {
                sleep := pkt.Time.Sub(*lastPktTime)
                if sleep > 0 {
                    time.Sleep(sleep)
                } else {
                    WARN("Time in pcap went backwards: %d", sleep)
                }
            }
            _lastPktTime := pkt.Time
            lastPktTime = &_lastPktTime
            pkt.Time = time.Now() // overwrite what we get from the pcap
        }
        counter++

        decodePktEth(datalink, pkt)
    }
    INFO("Input finish. Processed %d packets. Have a nice day!", counter)
}
<commit_msg>Check packet length before saving it.<commit_after>package main

import (
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "runtime"
    "strconv"
    "strings"
    "time"

    "github.com/BurntSushi/toml"
    "github.com/akrennmair/gopcap"
    "github.com/nranchev/go-libGeoIP"
)

const Version = "0.2.0"

type Packet struct {
    ts      time.Time
    tuple   IpPortTuple
    payload []byte
}

type protocolType uint16

const (
    UnknownProtocol protocolType = iota
    HttpProtocol
    MysqlProtocol
    RedisProtocol
    PgsqlProtocol
)

var protocolNames = []string{"unknown", "http", "mysql", "redis", "pgsql"}

type tomlConfig struct {
    Interfaces    tomlInterfaces
    RunOptions    tomlRunOptions
    Protocols     map[string]tomlProtocol
    Procs         tomlProcs
    Elasticsearch tomlMothership
    Agent         tomlAgent
    Logging       tomlLogging
}

type tomlInterfaces struct {
    Device string
}

type tomlRunOptions struct {
    Uid int
    Gid int
}

type tomlLogging struct {
    Selectors []string
}

var _Config tomlConfig
var _ConfigMeta toml.MetaData
var _GeoLite *libgeo.GeoIP

func Bytes_Ipv4_Ntoa(bytes []byte) string {
    var strarr []string = make([]string, 4)
    for i, b := range bytes {
        strarr[i] = strconv.Itoa(int(b))
    }
    return strings.Join(strarr, ".")
}

func decodePktEth(datalink int, pkt *pcap.Packet) {
    defer RECOVER("decodePktEth exception")

    packet := new(Packet)
    var l2hlen int
    var eth_type uint16

    if datalink != pcap.LINKTYPE_LINUX_SLL {
        // we are just assuming ETH
        l2hlen = 14

        // bytes 12 and 13 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[12:14])

    } else {
        l2hlen = 16

        // bytes 14 and 15 are the pkt type
        eth_type = Bytes_Ntohs(pkt.Data[14:16])
    }

    if eth_type != 0x800 {
        DEBUG("ip", "Not ipv4 packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+20 {
        DEBUG("ip", "Packet too short to be ethernet")
        return
    }

    // IP header
    iphl := int((uint16(pkt.Data[l2hlen]) & 0x0f) * 4)
    if len(pkt.Data) < l2hlen+iphl {
        DEBUG("ip", "Packet too short to be IP")
        return
    }
    iphdr := pkt.Data[l2hlen : l2hlen+iphl]

    //DEBUG("Packet timestamp: %s", pkt.Time)
    packet.ts = pkt.Time

    packet.tuple.Src_ip = Bytes_Ntohl(iphdr[12:16])
    packet.tuple.Dst_ip = Bytes_Ntohl(iphdr[16:20])

    mf := (uint8(iphdr[6]&0x20) != 0)
    frag_offset := (uint16(iphdr[7]) << 3) | (uint16(iphdr[6]) & 0x1F)

    if mf || frag_offset > 0 {
        DEBUG("ip", "Fragmented packets not yet supported")
        return
    }

    ip_length := int(Bytes_Ntohs(iphdr[2:4]))

    protocol := uint8(iphdr[9])
    if protocol != 6 {
        DEBUG("ip", "Not TCP packet. Ignoring")
        return
    }

    if len(pkt.Data) < l2hlen+iphl+20 {
        DEBUG("ip", "Packet too short to be TCP")
        return
    }

    tcphl := int((uint16(pkt.Data[l2hlen+iphl+12]) >> 4) * 4)
    if tcphl > 20 && len(pkt.Data) < l2hlen+iphl+tcphl {
        DEBUG("ip", "Packet too short to contain TCP header")
        return
    }

    tcphdr := pkt.Data[l2hlen+iphl : l2hlen+iphl+tcphl]
    packet.tuple.Src_port = Bytes_Ntohs(tcphdr[0:2])
    packet.tuple.Dst_port = Bytes_Ntohs(tcphdr[2:4])

    //DEBUG(" %s:%d -> %s:%d", packet.tuple.src_ip, packet.tuple.src_port,
    //      packet.tuple.dst_ip, packet.tuple.dst_port)

    data_offset := (tcphdr[12] >> 4) * 4

    if l2hlen+iphl+int(data_offset) > l2hlen+ip_length {
        DEBUG("ip", "data_offset pointing outside of packet")
        return
    }

    if len(pkt.Data) < l2hlen + ip_length {
        DEBUG("ip", "Captured packet smaller then advertised in IP layer")
        return
    }

    packet.payload = pkt.Data[l2hlen+iphl+int(data_offset) : l2hlen+ip_length]

    FollowTcp(tcphdr, packet)
}

func main() {

    configfile := flag.String("c", "packetbeat.conf", "Configuration file")
    file := flag.String("I", "", "file")
    loop := flag.Int("l", 1, "Loop file. 0 - loop forever")
    debugSelectorsStr := flag.String("d", "", "Enable certain debug selectors")
    oneAtAtime := flag.Bool("O", false, "Read packets one at a time (press Enter)")
    toStdout := flag.Bool("e", false, "Output to stdout instead of syslog")
    topSpeed := flag.Bool("t", false, "Read packets as fast as possible, without sleeping")
    publishDisabled := flag.Bool("N", false, "Disable actual publishing for testing")
    printVersion := flag.Bool("version", false, "Print version and exit")

    flag.Parse()

    if *printVersion {
        fmt.Printf("Packetbeat version %s (%s)\n", Version, runtime.GOARCH)
    }

    debugSelectors := []string{}
    logLevel := LOG_DEBUG
    if len(*debugSelectorsStr) > 0 {
        debugSelectors = strings.Split(*debugSelectorsStr, ",")
    }

    var h *pcap.Pcap
    var err error

    if _ConfigMeta, err = toml.DecodeFile(*configfile, &_Config); err != nil {
        fmt.Printf("TOML config parsing failed on %s: %s. Exiting.\n", *configfile, err)
        return
    }

    if len(debugSelectors) == 0 {
        debugSelectors = _Config.Logging.Selectors
    }
    LogInit(logLevel, "", !*toStdout, debugSelectors)

    if !IS_DEBUG("stdlog") {
        // disable standard logging by default
        log.SetOutput(ioutil.Discard)
    }

    if *file != "" {
        h, err = pcap.Openoffline(*file)
        if h == nil {
            ERR("Openoffline(%s) failed: %s", *file, err)
            return
        }
    } else {
        device := _Config.Interfaces.Device

        h, err = pcap.Openlive(device, 65535, true, 0)
        if h == nil {
            ERR("pcap.Openlive failed(%s): %s", device, err)
            return
        }
    }
    defer h.Close()

    if err = DropPrivileges(); err != nil {
        CRIT(err.Error())
        return
    }

    filter := configToFilter(&_Config)
    if filter != "" {
        DEBUG("pcapfilter", "Installing filter '%s'", filter)
        err := h.Setfilter(filter)
        if err != nil {
            ERR("pcap.Setfilter failed: %s", err)
            return
        }
    }

    if err = Publisher.Init(*publishDisabled); err != nil {
        CRIT(err.Error())
        return
    }

    if err = procWatcher.Init(&_Config.Procs); err != nil {
        CRIT(err.Error())
        return
    }

    if err = TcpInit(); err != nil {
        CRIT(err.Error())
        return
    }

    datalink := h.Datalink()
    if datalink != pcap.LINKTYPE_ETHERNET && datalink != pcap.LINKTYPE_LINUX_SLL {
        WARN("Unsupported link type: %d", datalink)
    }

    _GeoLite, err = libgeo.Load("/usr/share/GeoIP/GeoIP.dat")
    if err != nil {
        WARN("Could not load GeoIP data: %s", err.Error())
    }

    counter := 0
    live := true
    loopCount := 1
    var lastPktTime *time.Time = nil
    for live {
        var pkt *pcap.Packet
        var res int32

        if *oneAtAtime {
            fmt.Println("Press enter to read packet")
            fmt.Scanln()
        }

        pkt, res = h.NextEx()
        switch res {
        case -1:
            ERR("pcap.NextEx() error: %s", h.Geterror())
            live = false
            continue
        case -2:
            DEBUG("pcapread", "End of file")
            loopCount += 1
            if *loop > 0 && loopCount > *loop {
                // give a bit of time to the publish goroutine
                // to flush
                time.Sleep(300 * time.Millisecond)
                live = false
            }
            // reopen the file
            h, err = pcap.Openoffline(*file)
            if h == nil {
                ERR("Openoffline(%s) failed: %s", *file, err)
                return
            }
            lastPktTime = nil
            continue
        case 0:
            // timeout
            continue
        }
        if res != 1 {
            panic(fmt.Sprintf("Unexpected return code from pcap.NextEx: %d", res))
        }

        if pkt == nil {
            panic("Nil packet despite res=1")
        }

        if *file != "" {
            if lastPktTime != nil && !*topSpeed {
                sleep := pkt.Time.Sub(*lastPktTime)
                if sleep > 0 {
                    time.Sleep(sleep)
                } else {
                    WARN("Time in pcap went backwards: %d", sleep)
                }
            }
            _lastPktTime := pkt.Time
            lastPktTime = &_lastPktTime
            pkt.Time = time.Now() // overwrite what we get from the pcap
        }
        counter++

        decodePktEth(datalink, pkt)
    }
    INFO("Input finish. Processed %d packets. Have a nice day!", counter)
}
<|endoftext|>