<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/phonyphonecall/turnip"
	"github.com/tylerharter/open-lambda/worker/container"
	"github.com/tylerharter/open-lambda/worker/handler"
)

type Server struct {
	manager  container.ContainerManager
	handlers *handler.HandlerSet

	// config options
	registry_host string
	registry_port string
	docker_host   string

	lambdaTimer *turnip.Turnip
}

type httpErr struct {
	msg  string
	code int
}

func newHttpErr(msg string, code int) *httpErr {
	return &httpErr{msg: msg, code: code}
}

func NewServer(
	registry_host string,
	registry_port string,
	docker_host string) (*Server, error) {

	// registry
	if registry_host == "" {
		registry_host = "localhost"
		log.Printf("Using '%v' for registry_host", registry_host)
	}

	if registry_port == "" {
		registry_port = "5000"
		log.Printf("Using '%v' for registry_port", registry_port)
	}

	// daemon
	cm := container.NewDockerManager(registry_host, registry_port)
	if docker_host == "" {
		endpoint := cm.Client().Endpoint()
		local := "unix://"
		nonLocal := "https://"
		if strings.HasPrefix(endpoint, local) {
			docker_host = "localhost"
		} else if strings.HasPrefix(endpoint, nonLocal) {
			start := strings.Index(endpoint, nonLocal) + len([]rune(nonLocal))
			end := strings.LastIndex(endpoint, ":")
			docker_host = endpoint[start:end]
		} else {
			return nil, fmt.Errorf("please specify a docker host!")
		}
		log.Printf("Using '%v' for docker_host", docker_host)
	}

	// create server
	opts := handler.HandlerSetOpts{
		Cm:  cm,
		Lru: handler.NewHandlerLRU(100), // TODO(tyler)
	}
	server := &Server{
		registry_host: registry_host,
		registry_port: registry_port,
		docker_host:   docker_host,
		manager:       cm,
		handlers:      handler.NewHandlerSet(opts),
		lambdaTimer:   turnip.NewTurnip(),
	}

	return server, nil
}

func (s *Server) Manager() container.ContainerManager {
	return s.manager
}

func (s *Server) ForwardToContainer(handler *handler.Handler, r *http.Request, input []byte) ([]byte, *http.Response, *httpErr) {
	port, err := handler.RunStart()
	if err != nil {
		if strings.Contains(err.Error(), "is not paused") {
			err = errors.New("Error: Container stopped when it should be paused. Execution of lambda handler likely failed.")
		}
		return nil, nil, newHttpErr(
			err.Error(),
			http.StatusInternalServerError)
	}
	defer handler.RunFinish()

	// forward request to container.  r and w are the server
	// request and response respectively.  r2 and w2 are the
	// container request and response respectively.
	host := fmt.Sprintf("%s:%s", s.docker_host, port)
	url := fmt.Sprintf("http://%s%s", host, r.URL.Path)
	// log.Printf("proxying request to %s\n", url)

	// TODO(tyler): some sort of smarter backoff.  Or, a better
	// way to detect a started container.
	max_tries := 10
	for tries := 1; ; tries++ {
		r2, err := http.NewRequest(r.Method, url, bytes.NewReader(input))
		if err != nil {
			return nil, nil, newHttpErr(
				err.Error(),
				http.StatusInternalServerError)
		}

		r2.Header.Set("Content-Type", r.Header.Get("Content-Type"))
		client := &http.Client{}
		w2, err := client.Do(r2)
		if err != nil {
			log.Printf("request to container failed with %v\n", err)
			if tries == max_tries {
				return nil, nil, newHttpErr(
					err.Error(),
					http.StatusInternalServerError)
			}
			log.Printf("retry request\n")
			time.Sleep(time.Duration(tries*10) * time.Millisecond)
			continue
		}

		defer w2.Body.Close()
		wbody, err := ioutil.ReadAll(w2.Body)
		if err != nil {
			return nil, nil, newHttpErr(
				err.Error(),
				http.StatusInternalServerError)
		}
		return wbody, w2, nil
	}
}

func (s *Server) RunLambdaErr(w http.ResponseWriter, r *http.Request) *httpErr {
	// components represent runLambda[0]/<name_of_container>[1]/<extra_things>...
	// ergo we want [1] for name of container
	urlParts := getUrlComponents(r)
	if len(urlParts) < 2 {
		return newHttpErr(
			"Name of image to run required",
			http.StatusBadRequest)
	}
	img := urlParts[1]
	i := strings.Index(img, "?")
	if i >= 0 {
		img = img[:i-1]
	}

	// read request
	rbody := []byte{}
	if r.Body != nil {
		defer r.Body.Close()
		var err error
		rbody, err = ioutil.ReadAll(r.Body)
		if err != nil {
			return newHttpErr(
				err.Error(),
				http.StatusInternalServerError)
		}
	}

	// forward to container
	handler := s.handlers.Get(img)
	wbody, w2, err := s.ForwardToContainer(handler, r, rbody)
	if err != nil {
		return err
	}

	// write response
	// TODO(tyler): origins should be configurable
	// w.Header().Set("Access-Control-Allow-Origin", "*")
	// w.Header().Set("Access-Control-Allow-Methods",
	//	"GET, PUT, POST, DELETE, OPTIONS")
	//w.Header().Set("Access-Control-Allow-Headers",
	//	"Content-Type, Content-Range, Content-Disposition, Content-Description")

	w.WriteHeader(w2.StatusCode)

	if _, err := w.Write(wbody); err != nil {
		return newHttpErr(
			err.Error(),
			http.StatusInternalServerError)
	}

	return nil
}

// RunLambda expects POST requests like this:
//
// curl -X POST localhost:8080/runLambda/<lambda-name> -d '{}'
func (s *Server) RunLambda(w http.ResponseWriter, r *http.Request) {
	// write response headers
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods",
		"GET, PUT, POST, DELETE, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers",
		"Content-Type, Content-Range, Content-Disposition, Content-Description, X-Requested-With")

	if r.Method == "OPTIONS" {
		w.WriteHeader(200)
	} else {
		s.lambdaTimer.Start()
		if err := s.RunLambdaErr(w, r); err != nil {
			log.Printf("could not handle request: %s\n", err.msg)
			http.Error(w, err.msg, err.code)
		}
		s.lambdaTimer.Stop()

	}

}

func (s *Server) Dump() {
	log.Printf("============ Server Stats ===========\n")
	log.Printf("\tlambda: \t%fms\n", s.lambdaTimer.AverageMs())
	log.Printf("=====================================\n")
}

// Parses request URL into its "/" delimated components
func getUrlComponents(r *http.Request) []string {
	path := r.URL.Path

	// trim prefix
	if strings.HasPrefix(path, "/") {
		path = path[1:]
	}

	// trim trailing "/"
	if strings.HasSuffix(path, "/") {
		path = path[:len(path)-1]
	}

	components := strings.Split(path, "/")
	return components
}

func main() {
	if len(os.Args) < 3 {
		log.Fatalf("usage: %s <registry hostname> <registry port>\n", os.Args[0])
	}

	docker_host, ok := os.LookupEnv("OL_DOCKER_HOST")
	if !ok {
		docker_host = ""
	}
	server, err := NewServer(os.Args[1], os.Args[2], docker_host)
	if err != nil {
		log.Fatal(err)
	}

	http.HandleFunc("/runLambda/", server.RunLambda)
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<commit_msg>handling OPTIONS requests (issue #16)<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"
	"time"

	"github.com/phonyphonecall/turnip"
	"github.com/tylerharter/open-lambda/worker/container"
	"github.com/tylerharter/open-lambda/worker/handler"
)

type Server struct {
	manager  container.ContainerManager
	handlers *handler.HandlerSet

	// config options
	registry_host string
	registry_port string
	docker_host   string

	lambdaTimer *turnip.Turnip
}

type httpErr struct {
	msg  string
	code int
}

func newHttpErr(msg string, code int) *httpErr {
	return &httpErr{msg: msg, code: code}
}

func NewServer(
	registry_host string,
	registry_port string,
	docker_host string) (*Server, error) {

	// registry
	if registry_host == "" {
		registry_host = "localhost"
		log.Printf("Using '%v' for registry_host", registry_host)
	}

	if registry_port == "" {
		registry_port = "5000"
		log.Printf("Using '%v' for registry_port", registry_port)
	}

	// daemon
	cm := container.NewDockerManager(registry_host, registry_port)
	if docker_host == "" {
		endpoint := cm.Client().Endpoint()
		local := "unix://"
		nonLocal := "https://"
		if strings.HasPrefix(endpoint, local) {
			docker_host = "localhost"
		} else if strings.HasPrefix(endpoint, nonLocal) {
			start := strings.Index(endpoint, nonLocal) + len([]rune(nonLocal))
			end := strings.LastIndex(endpoint, ":")
			docker_host = endpoint[start:end]
		} else {
			return nil, fmt.Errorf("please specify a docker host!")
		}
		log.Printf("Using '%v' for docker_host", docker_host)
	}

	// create server
	opts := handler.HandlerSetOpts{
		Cm:  cm,
		Lru: handler.NewHandlerLRU(100), // TODO(tyler)
	}
	server := &Server{
		registry_host: registry_host,
		registry_port: registry_port,
		docker_host:   docker_host,
		manager:       cm,
		handlers:      handler.NewHandlerSet(opts),
		lambdaTimer:   turnip.NewTurnip(),
	}

	return server, nil
}

func (s *Server) Manager() container.ContainerManager {
	return s.manager
}

func (s *Server) ForwardToContainer(handler *handler.Handler, r *http.Request, input []byte) ([]byte, *http.Response, *httpErr) {
	port, err := handler.RunStart()
	if err != nil {
		return nil, nil, newHttpErr(
			err.Error(),
			http.StatusInternalServerError)
	}
	defer handler.RunFinish()

	// forward request to container.  r and w are the server
	// request and response respectively.  r2 and w2 are the
	// container request and response respectively.
	host := fmt.Sprintf("%s:%s", s.docker_host, port)
	url := fmt.Sprintf("http://%s%s", host, r.URL.Path)
	// log.Printf("proxying request to %s\n", url)

	// TODO(tyler): some sort of smarter backoff.  Or, a better
	// way to detect a started container.
	max_tries := 10
	for tries := 1; ; tries++ {
		r2, err := http.NewRequest(r.Method, url, bytes.NewReader(input))
		if err != nil {
			return nil, nil, newHttpErr(
				err.Error(),
				http.StatusInternalServerError)
		}

		r2.Header.Set("Content-Type", r.Header.Get("Content-Type"))
		client := &http.Client{}
		w2, err := client.Do(r2)
		if err != nil {
			log.Printf("request to container failed with %v\n", err)
			if tries == max_tries {
				return nil, nil, newHttpErr(
					err.Error(),
					http.StatusInternalServerError)
			}
			log.Printf("retry request\n")
			time.Sleep(time.Duration(tries*10) * time.Millisecond)
			continue
		}

		defer w2.Body.Close()
		wbody, err := ioutil.ReadAll(w2.Body)
		if err != nil {
			return nil, nil, newHttpErr(
				err.Error(),
				http.StatusInternalServerError)
		}
		return wbody, w2, nil
	}
}

func (s *Server) RunLambdaErr(w http.ResponseWriter, r *http.Request) *httpErr {
	// components represent runLambda[0]/<name_of_container>[1]/<extra_things>...
	// ergo we want [1] for name of container
	urlParts := getUrlComponents(r)
	if len(urlParts) < 2 {
		return newHttpErr(
			"Name of image to run required",
			http.StatusBadRequest)
	}
	img := urlParts[1]
	i := strings.Index(img, "?")
	if i >= 0 {
		img = img[:i-1]
	}

	// read request
	rbody := []byte{}
	if r.Body != nil {
		defer r.Body.Close()
		var err error
		rbody, err = ioutil.ReadAll(r.Body)
		if err != nil {
			return newHttpErr(
				err.Error(),
				http.StatusInternalServerError)
		}
	}

	// forward to container
	handler := s.handlers.Get(img)
	wbody, w2, err := s.ForwardToContainer(handler, r, rbody)
	if err != nil {
		return err
	}

	w.WriteHeader(w2.StatusCode)

	if _, err := w.Write(wbody); err != nil {
		return newHttpErr(
			err.Error(),
			http.StatusInternalServerError)
	}

	return nil
}

// RunLambda expects POST requests like this:
//
// curl -X POST localhost:8080/runLambda/<lambda-name> -d '{}'
func (s *Server) RunLambda(w http.ResponseWriter, r *http.Request) {
	// write response headers
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Allow-Methods",
		"GET, PUT, POST, DELETE, OPTIONS")
	w.Header().Set("Access-Control-Allow-Headers",
		"Content-Type, Content-Range, Content-Disposition, Content-Description, X-Requested-With")

	if r.Method == "OPTIONS" {
		w.WriteHeader(200)
	} else {
		s.lambdaTimer.Start()
		if err := s.RunLambdaErr(w, r); err != nil {
			log.Printf("could not handle request: %s\n", err.msg)
			http.Error(w, err.msg, err.code)
		}
		s.lambdaTimer.Stop()

	}

}

func (s *Server) Dump() {
	log.Printf("============ Server Stats ===========\n")
	log.Printf("\tlambda: \t%fms\n", s.lambdaTimer.AverageMs())
	log.Printf("=====================================\n")
}

// Parses request URL into its "/" delimated components
func getUrlComponents(r *http.Request) []string {
	path := r.URL.Path

	// trim prefix
	if strings.HasPrefix(path, "/") {
		path = path[1:]
	}

	// trim trailing "/"
	if strings.HasSuffix(path, "/") {
		path = path[:len(path)-1]
	}

	components := strings.Split(path, "/")
	return components
}

func main() {
	if len(os.Args) < 3 {
		log.Fatalf("usage: %s <registry hostname> <registry port>\n", os.Args[0])
	}

	docker_host, ok := os.LookupEnv("OL_DOCKER_HOST")
	if !ok {
		docker_host = ""
	}
	server, err := NewServer(os.Args[1], os.Args[2], docker_host)
	if err != nil {
		log.Fatal(err)
	}

	http.HandleFunc("/runLambda/", server.RunLambda)
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<|endoftext|>


<commit_before>package worker

import (
	"github.com/oklahomer/go-sarah/log"
	"golang.org/x/net/context"
	"runtime"
	"time"
)

type Config struct {
	WorkerNum         uint          `json:"worker_num" yaml:"worker_num"`
	QueueSize         uint          `json:"queue_size" yaml:"queue_size"`
	SuperviseInterval time.Duration `json:"supervise_interval" yaml:"supervise_interval"`
}

// NewConfig returns Config instance with default configuration values.
// To Override with desired value, pass the returned instance to json.Unmarshal or yaml.Unmarshal.
func NewConfig() *Config {
	// Set default values.
	return &Config{
		WorkerNum:         100,
		QueueSize:         10,
		SuperviseInterval: 60 * time.Second,
	}
}

// Run creates as many child workers as specified and start those child workers.
// First argument, cancel channel, can be context.Context.Done to propagate upstream status change.
func Run(ctx context.Context, config *Config) chan<- func() {
	log.Infof("start workers")

	job := make(chan func(), config.QueueSize)

	var i uint
	for i = 1; i <= config.WorkerNum; i++ {
		go runChild(ctx, job, i)
	}

	if config.SuperviseInterval > 0 {
		go superviseQueueLength(ctx, job, config.SuperviseInterval)
	}

	return job
}

func runChild(ctx context.Context, job <-chan func(), workerID uint) {
	log.Infof("start worker id: %d.", workerID)

	for {
		select {
		case <-ctx.Done():
			log.Infof("stopping worker id: %d", workerID)
			return
		case job := <-job:
			log.Debugf("receiving job on worker: %d", workerID)
			// To avoid given job's panic affect later jobs, wrap them with recover.
			func() {
				defer func() {
					if r := recover(); r != nil {
						log.Warnf("panic in given job. recovered: %#v", r)

						// Display stack trace
						for depth := 0; ; depth++ {
							_, src, line, ok := runtime.Caller(depth)
							if !ok {
								break
							}
							log.Warnf(" -> depth:%d. file:%s. line:%d.", depth, src, line)
						}
					}

				}()
				job()
			}()
		}
	}
}

func superviseQueueLength(ctx context.Context, job chan<- func(), interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			log.Infof("worker queue length: %d", len(job))
		}
	}
}
<commit_msg>Update obsolete comment<commit_after>package worker

import (
	"github.com/oklahomer/go-sarah/log"
	"golang.org/x/net/context"
	"runtime"
	"time"
)

type Config struct {
	WorkerNum         uint          `json:"worker_num" yaml:"worker_num"`
	QueueSize         uint          `json:"queue_size" yaml:"queue_size"`
	SuperviseInterval time.Duration `json:"supervise_interval" yaml:"supervise_interval"`
}

// NewConfig returns Config instance with default configuration values.
// To Override with desired value, pass the returned instance to json.Unmarshal or yaml.Unmarshal.
func NewConfig() *Config {
	// Set default values.
	return &Config{
		WorkerNum:         100,
		QueueSize:         10,
		SuperviseInterval: 60 * time.Second,
	}
}

// Run creates as many child workers as specified and start those child workers.
func Run(ctx context.Context, config *Config) chan<- func() {
	log.Infof("start workers")

	job := make(chan func(), config.QueueSize)

	var i uint
	for i = 1; i <= config.WorkerNum; i++ {
		go runChild(ctx, job, i)
	}

	if config.SuperviseInterval > 0 {
		go superviseQueueLength(ctx, job, config.SuperviseInterval)
	}

	return job
}

func runChild(ctx context.Context, job <-chan func(), workerID uint) {
	log.Infof("start worker id: %d.", workerID)

	for {
		select {
		case <-ctx.Done():
			log.Infof("stopping worker id: %d", workerID)
			return
		case job := <-job:
			log.Debugf("receiving job on worker: %d", workerID)
			// To avoid given job's panic affect later jobs, wrap them with recover.
			func() {
				defer func() {
					if r := recover(); r != nil {
						log.Warnf("panic in given job. recovered: %#v", r)

						// Display stack trace
						for depth := 0; ; depth++ {
							_, src, line, ok := runtime.Caller(depth)
							if !ok {
								break
							}
							log.Warnf(" -> depth:%d. file:%s. line:%d.", depth, src, line)
						}
					}

				}()
				job()
			}()
		}
	}
}

func superviseQueueLength(ctx context.Context, job chan<- func(), interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			log.Infof("worker queue length: %d", len(job))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 Xing Xing <mikespook@gmail.com> All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package worker

import (
	"encoding/binary"
	"fmt"
	"sync"
	"time"
)

const (
	Unlimited = 0
	OneByOne  = 1

	Immediately = 0
)

/*
Worker side api for gearman

usage:
w = worker.New(worker.Unlimited)
w.AddFunction("foobar", foobar)
w.AddServer("127.0.0.1:4730")
w.Work() // Enter the worker's main loop

The definition of the callback function 'foobar' should suit for the type 'JobFunction'.
It looks like this:

func foobar(job *Job) (data []byte, err os.Error) {
    //sth. here
    //plaplapla...
    return
}
*/
type Worker struct {
	agents  map[string]*agent
	funcs   JobFuncs
	in      chan *inPack
	running bool
	limit   chan bool

	Id string
	// assign a ErrFunc to handle errors
	ErrorHandler ErrorHandler
	JobHandler   JobHandler
	mutex        sync.Mutex
}

// Get a new worker
func New(l int) (worker *Worker) {
	worker = &Worker{
		agents: make(map[string]*agent, QUEUE_SIZE),
		funcs:  make(JobFuncs),
		in:     make(chan *inPack, QUEUE_SIZE),
	}
	if l != Unlimited {
		worker.limit = make(chan bool, l)
	}
	return
}

//
func (worker *Worker) err(e error) {
	if worker.ErrorHandler != nil {
		worker.ErrorHandler(e)
	}
}

// Add a server. The addr should be 'host:port' format.
// The connection is established at this time.
func (worker *Worker) AddServer(net, addr string) (err error) {
	// Create a new job server's client as a agent of server
	a, err := newAgent(net, addr, worker)
	if err != nil {
		return err
	}
	worker.agents[net+addr] = a
	return
}

// Write a job to job server.
// Here, the job's mean is not the oraginal mean.
// Just looks like a network package for job's result or tell job server, there was a fail.
func (worker *Worker) broadcast(outpack *outPack) {
	for _, v := range worker.agents {
		v.write(outpack)
	}
}

// Add a function.
// Plz added job servers first, then functions.
// The API will tell every connected job server that 'I can do this'
func (worker *Worker) AddFunc(funcname string,
	f JobFunc, timeout uint32) (err error) {
	worker.mutex.Lock()
	defer worker.mutex.Unlock()
	if _, ok := worker.funcs[funcname]; ok {
		return fmt.Errorf("The function already exists: %s", funcname)
	}
	worker.funcs[funcname] = &jobFunc{f: f, timeout: timeout}
	if worker.running {
		worker.addFunc(funcname, timeout)
	}
	return
}

// inner add function
func (worker *Worker) addFunc(funcname string, timeout uint32) {
	outpack := getOutPack()
	if timeout == 0 {
		outpack.dataType = CAN_DO
		outpack.data = []byte(funcname)
	} else {
		outpack.dataType = CAN_DO_TIMEOUT
		l := len(funcname)
		outpack.data = getBuffer(l + 5)
		copy(outpack.data, []byte(funcname))
		outpack.data[l] = '\x00'
		binary.BigEndian.PutUint32(outpack.data[l+1:], timeout)
	}
	worker.broadcast(outpack)
}

// Remove a function.
func (worker *Worker) RemoveFunc(funcname string) (err error) {
	worker.mutex.Lock()
	defer worker.mutex.Unlock()
	if _, ok := worker.funcs[funcname]; !ok {
		return fmt.Errorf("The function does not exist: %s", funcname)
	}
	delete(worker.funcs, funcname)
	if worker.running {
		worker.removeFunc(funcname)
	}
	return
}

// inner remove function
func (worker *Worker) removeFunc(funcname string) {
	outpack := getOutPack()
	outpack.dataType = CANT_DO
	outpack.data = []byte(funcname)
	worker.broadcast(outpack)
}

func (worker *Worker) handleInPack(inpack *inPack) {
	defer func() {
		if worker.running && worker.limit != nil {
			<-worker.limit
		}
	}()
	switch inpack.dataType {
	case NO_JOB:
		inpack.a.PreSleep()
	case NOOP:
		inpack.a.Grab()
	case ERROR:
		worker.err(GetError(inpack.data))
	case JOB_ASSIGN, JOB_ASSIGN_UNIQ:
		if err := worker.exec(inpack); err != nil {
			worker.err(err)
		}
	default:
		worker.customeHandler(inpack)
	}
}

func (worker *Worker) Ready() (err error) {
	for _, v := range worker.agents {
		if err = v.Connect(); err != nil {
			return
		}
	}
	for funcname, f := range worker.funcs {
		worker.addFunc(funcname, f.timeout)
	}
	return
}

// Main loop
func (worker *Worker) Work() {
	worker.running = true
	for _, v := range worker.agents {
		v.Grab()
	}
	var inpack *inPack
	for inpack = range worker.in {
		go worker.handleInPack(inpack)
	}
}

// job handler
func (worker *Worker) customeHandler(inpack *inPack) {
	if worker.JobHandler != nil {
		if err := worker.JobHandler(inpack); err != nil {
			worker.err(err)
		}
	}
}

// Close.
func (worker *Worker) Close() {
	worker.running = false
	close(worker.in)
	if worker.limit != nil {
		close(worker.limit)
	}
}

// Send a something out, get the samething back.
func (worker *Worker) Echo(data []byte) {
	outpack := getOutPack()
	outpack.dataType = ECHO_REQ
	outpack.data = data
	worker.broadcast(outpack)
}

// Remove all of functions.
// Both from the worker or job servers.
func (worker *Worker) Reset() {
	outpack := getOutPack()
	outpack.dataType = RESET_ABILITIES
	worker.broadcast(outpack)
	worker.funcs = make(JobFuncs)
}

// Set the worker's unique id.
func (worker *Worker) SetId(id string) {
	worker.Id = id
	outpack := getOutPack()
	outpack.dataType = SET_CLIENT_ID
	outpack.data = []byte(id)
	worker.broadcast(outpack)
}

// Execute the job. And send back the result.
func (worker *Worker) exec(inpack *inPack) (err error) {
	defer func() {
		if r := recover(); r != nil {
			if e, ok := r.(error); ok {
				err = e
			} else {
				err = ErrUnknown
			}
		}
	}()
	f, ok := worker.funcs[inpack.fn]
	if !ok {
		return fmt.Errorf("The function does not exist: %s", inpack.fn)
	}
	var r *result
	if f.timeout == 0 {
		d, e := f.f(inpack)
		r = &result{data: d, err: e}
	} else {
		r = execTimeout(f.f, inpack, time.Duration(f.timeout)*time.Second)
	}
	if worker.running {
		outpack := getOutPack()
		if r.err == nil {
			outpack.dataType = WORK_COMPLETE
		} else {
			if len(r.data) == 0 {
				outpack.dataType = WORK_FAIL
			} else {
				outpack.dataType = WORK_EXCEPTION
			}
			err = r.err
		}
		outpack.handle = inpack.handle
		outpack.data = r.data
		inpack.a.write(outpack)
		inpack.a.Grab()
	}
	return
}

type result struct {
	data []byte
	err  error
}

func execTimeout(f JobFunc, job Job, timeout time.Duration) (r *result) {
	rslt := make(chan *result)
	defer close(rslt)
	go func() {
		defer func() { recover() }()
		d, e := f(job)
		rslt <- &result{data: d, err: e}
	}()
	select {
	case r = <-rslt:
	case <-time.After(timeout):
		return &result{err: ErrTimeOut}
	}
	return r
}
<commit_msg>The limitation param was removed<commit_after>// Copyright 2011 Xing Xing <mikespook@gmail.com> All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package worker

import (
	"encoding/binary"
	"fmt"
	"sync"
	"time"
)

const (
	Immediately = 0
)

/*
Worker side api for gearman

usage:
w = worker.New()
w.AddFunction("foobar", foobar)
w.AddServer("127.0.0.1:4730")
w.Work() // Enter the worker's main loop

The definition of the callback function 'foobar' should suit for the type 'JobFunction'.
It looks like this:

func foobar(job *Job) (data []byte, err os.Error) {
    //sth. here
    //plaplapla...
    return
}
*/
type Worker struct {
	agents  []*agent
	funcs   JobFuncs
	in      chan *inPack
	running bool

	Id string
	// assign a ErrFunc to handle errors
	ErrorHandler ErrorHandler
	JobHandler   JobHandler
	mutex        sync.Mutex
}

// Get a new worker
func New() (worker *Worker) {
	worker = &Worker{
		agents: make([]*agent, 0),
		funcs:  make(JobFuncs),
		in:     make(chan *inPack, QUEUE_SIZE),
	}
	return
}

//
func (worker *Worker) err(e error) {
	if worker.ErrorHandler != nil {
		worker.ErrorHandler(e)
	}
}

// Add a server. The addr should be 'host:port' format.
// The connection is established at this time.
func (worker *Worker) AddServer(net, addr string) (err error) {
	// Create a new job server's client as a agent of server
	a, err := newAgent(net, addr, worker)
	if err != nil {
		return err
	}
	worker.agents = append(worker.agents, a)
	return
}

// Write a job to job server.
// Here, the job's mean is not the oraginal mean.
// Just looks like a network package for job's result or tell job server, there was a fail.
func (worker *Worker) broadcast(outpack *outPack) {
	for _, v := range worker.agents {
		v.write(outpack)
	}
}

// Add a function.
// Plz added job servers first, then functions.
// The API will tell every connected job server that 'I can do this'
func (worker *Worker) AddFunc(funcname string,
	f JobFunc, timeout uint32) (err error) {
	worker.mutex.Lock()
	defer worker.mutex.Unlock()
	if _, ok := worker.funcs[funcname]; ok {
		return fmt.Errorf("The function already exists: %s", funcname)
	}
	worker.funcs[funcname] = &jobFunc{f: f, timeout: timeout}
	if worker.running {
		worker.addFunc(funcname, timeout)
	}
	return
}

// inner add function
func (worker *Worker) addFunc(funcname string, timeout uint32) {
	outpack := getOutPack()
	if timeout == 0 {
		outpack.dataType = CAN_DO
		outpack.data = []byte(funcname)
	} else {
		outpack.dataType = CAN_DO_TIMEOUT
		l := len(funcname)
		outpack.data = getBuffer(l + 5)
		copy(outpack.data, []byte(funcname))
		outpack.data[l] = '\x00'
		binary.BigEndian.PutUint32(outpack.data[l+1:], timeout)
	}
	worker.broadcast(outpack)
}

// Remove a function.
func (worker *Worker) RemoveFunc(funcname string) (err error) {
	worker.mutex.Lock()
	defer worker.mutex.Unlock()
	if _, ok := worker.funcs[funcname]; !ok {
		return fmt.Errorf("The function does not exist: %s", funcname)
	}
	delete(worker.funcs, funcname)
	if worker.running {
		worker.removeFunc(funcname)
	}
	return
}

// inner remove function
func (worker *Worker) removeFunc(funcname string) {
	outpack := getOutPack()
	outpack.dataType = CANT_DO
	outpack.data = []byte(funcname)
	worker.broadcast(outpack)
}

func (worker *Worker) handleInPack(inpack *inPack) {
	switch inpack.dataType {
	case NO_JOB:
		inpack.a.PreSleep()
	case NOOP:
		inpack.a.Grab()
	case ERROR:
		worker.err(GetError(inpack.data))
	case JOB_ASSIGN, JOB_ASSIGN_UNIQ:
		if err := worker.exec(inpack); err != nil {
			worker.err(err)
		}
	default:
		worker.customeHandler(inpack)
	}
}

func (worker *Worker) Ready() (err error) {
	for _, v := range worker.agents {
		if err = v.Connect(); err != nil {
			return
		}
	}
	for funcname, f := range worker.funcs {
		worker.addFunc(funcname, f.timeout)
	}
	return
}

// Main loop
func (worker *Worker) Work() {
	worker.running = true
	for _, v := range worker.agents {
		v.Grab()
	}
	var inpack *inPack
	for inpack = range worker.in {
		go worker.handleInPack(inpack)
	}
}

// job handler
func (worker *Worker) customeHandler(inpack *inPack) {
	if worker.JobHandler != nil {
		if err := worker.JobHandler(inpack); err != nil {
			worker.err(err)
		}
	}
}

// Close.
func (worker *Worker) Close() {
	worker.running = false
	close(worker.in)
}

// Send a something out, get the samething back.
func (worker *Worker) Echo(data []byte) {
	outpack := getOutPack()
	outpack.dataType = ECHO_REQ
	outpack.data = data
	worker.broadcast(outpack)
}

// Remove all of functions.
// Both from the worker or job servers.
func (worker *Worker) Reset() {
	outpack := getOutPack()
	outpack.dataType = RESET_ABILITIES
	worker.broadcast(outpack)
	worker.funcs = make(JobFuncs)
}

// Set the worker's unique id.
func (worker *Worker) SetId(id string) {
	worker.Id = id
	outpack := getOutPack()
	outpack.dataType = SET_CLIENT_ID
	outpack.data = []byte(id)
	worker.broadcast(outpack)
}

// Execute the job. And send back the result.
func (worker *Worker) exec(inpack *inPack) (err error) {
	defer func() {
		if r := recover(); r != nil {
			if e, ok := r.(error); ok {
				err = e
			} else {
				err = ErrUnknown
			}
		}
	}()
	f, ok := worker.funcs[inpack.fn]
	if !ok {
		return fmt.Errorf("The function does not exist: %s", inpack.fn)
	}
	var r *result
	if f.timeout == 0 {
		d, e := f.f(inpack)
		r = &result{data: d, err: e}
	} else {
		r = execTimeout(f.f, inpack, time.Duration(f.timeout)*time.Second)
	}
	if worker.running {
		outpack := getOutPack()
		if r.err == nil {
			outpack.dataType = WORK_COMPLETE
		} else {
			if len(r.data) == 0 {
				outpack.dataType = WORK_FAIL
			} else {
				outpack.dataType = WORK_EXCEPTION
			}
			err = r.err
		}
		outpack.handle = inpack.handle
		outpack.data = r.data
		inpack.a.write(outpack)
		inpack.a.Grab()
	}
	return
}

type result struct {
	data []byte
	err  error
}

func execTimeout(f JobFunc, job Job, timeout time.Duration) (r *result) {
	rslt := make(chan *result)
	defer close(rslt)
	go func() {
		defer func() { recover() }()
		d, e := f(job)
		rslt <- &result{data: d, err: e}
	}()
	select {
	case r = <-rslt:
	case <-time.After(timeout):
		return &result{err: ErrTimeOut}
	}
	return r
}
<|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 gg

import (
	"fmt"

	"github.com/aclements/go-gg/table"
)

func defaultCols(p *Plot, cols ...*string) {
	dcols := p.Data().Columns()
	for i, colp := range cols {
		if *colp == "" {
			if i >= len(dcols) {
				panic(fmt.Sprintf("cannot get default column %d; table has only %d columns", i, len(dcols)))
			}
			*colp = dcols[i]
		}
	}
}

// LayerLines is like LayerPaths, but connects data points in order by
// the "x" property.
type LayerLines LayerPaths

func (l LayerLines) Apply(p *Plot) {
	LayerPaths(l).apply(p, true)
}

//go:generate stringer -type StepMode

// StepMode controls how LayerSteps connects subsequent points.
type StepMode int

const (
	// StepHV makes LayerSteps connect subsequent points with a
	// horizontal segment and then a vertical segment.
	StepHV StepMode = iota

	// StepVH makes LayerSteps connect subsequent points with a
	// vertical segment and then a horizontal segment.
	StepVH

	// StepHMid makes LayerSteps connect subsequent points A and B
	// with three segments: a horizontal segment from A to the
	// midpoint between A and B, followed by vertical segment,
	// followed by a horizontal segment from the midpoint to B.
	StepHMid

	// StepVMid makes LayerSteps connect subsequent points A and B
	// with three segments: a vertical segment from A to the
	// midpoint between A and B, followed by horizontal segment,
	// followed by a vertical segment from the midpoint to B.
	StepVMid
)

// LayerSteps is like LayerPaths, but connects data points with a path
// consisting only of horizontal and vertical segments.
type LayerSteps struct {
	LayerPaths

	Step StepMode
}

func (l LayerSteps) Apply(p *Plot) {
	// TODO: Should this also support only showing horizontal or
	// vertical segments?
	//
	// TODO: This could be a data transform instead of a layer.
	// Then it could be used in conjunction with, for example,
	// ribbons.

	defaultCols(p, &l.X, &l.Y)
	p.marks = append(p.marks, plotMark{&markSteps{
		l.Step,
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPaths groups by Color and Fill, and then connects successive
// data points in each group with a path and/or a filled polygon.
type LayerPaths struct {
	// X and Y name columns that define the input and response of
	// each point on the path. If these are empty, they default to
	// the first and second columns, respectively.
	X, Y string

	// Color names a column that defines the stroke color of each
	// path. If Color is "", it defaults to constant black.
	// Otherwise, the data is grouped by Color.
	Color string

	// Fill names a column that defines the fill color of each
	// path. If Fill is "", it defaults to none. Otherwise, the
	// data is grouped by Fill.
	Fill string

	// XXX Perhaps the theme should provide default values for
	// things like "color". That would suggest we need to resolve
	// defaults like that at render time. Possibly a special scale
	// that gets values from the theme could be used to resolve
	// them.
	//
	// XXX strokeOpacity, fillOpacity, strokeWidth, what other
	// properties do SVG strokes have?
	//
	// XXX Should the set of known styling bindings be fixed, and
	// all possible rendering targets have to know what to do with
	// them, or should the rendering target be able to have
	// different styling bindings they understand (presumably with
	// some reasonable base set)? If the renderer can determine
	// the known bindings, we would probably just capture the
	// environment here (and make it so a captured environment
	// does not change) and hand that to the renderer later.
}

func (l LayerPaths) Apply(p *Plot) {
	l.apply(p, false)
}

func (l LayerPaths) apply(p *Plot, sort bool) {
	defaultCols(p, &l.X, &l.Y)
	if l.Color != "" {
		p.GroupBy(l.Color)
	}
	if l.Fill != "" {
		p.GroupBy(l.Fill)
	}
	if sort {
		defer p.Save().Restore()
		p = p.SortBy(l.X)
	}

	p.marks = append(p.marks, plotMark{&markPath{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPoints layers a point mark at each data point.
type LayerPoints struct {
	// X and Y name columns that define input and response of each
	// point. If these are empty, they default to the first and
	// second columns, respectively.
	X, Y string

	// Color names the column that defines the fill color of each
	// point. If Color is "", it defaults to constant black.
	Color string

	// Opacity names the column that defines the opacity of each
	// point. If Opacity is "", it defaults to fully opaque. This
	// is multiplied by any alpha value specified by Color.
	Opacity string

	// Size names the column that defines the size of each point.
	// If Size is "", it defaults to 1% of the smallest plot
	// dimension.
	Size string

	// XXX fill vs stroke, shape
}

func (l LayerPoints) Apply(p *Plot) {
	defaultCols(p, &l.X, &l.Y)
	p.marks = append(p.marks, plotMark{&markPoint{
		p.use("x", l.X),
		p.use("y", l.Y),
		// TODO: It's actually the fill color, but I generally
		// want it to match things that are stroke colors.
		// Maybe I should have a "color" aesthetic for the
		// "primary" color? Or I could have a hierarchy of
		// aesthetics, in which this uses "stroke" if it has a
		// scale, but otherwise uses "color".
		p.use("stroke", l.Color),
		// TODO: What scale for opacity? Or should I assume
		// callers will use PreScaled values if they want
		// specific opacities? What's the physical type?
		p.use("opacity", l.Opacity),
		p.use("size", l.Size),
	}, p.Data().Tables()})
}

// LayerTiles layers a rectangle at each data point. The rectangle is
// specified by its center, width, and height.
type LayerTiles struct {
	// X and Y name columns that define the input and response at
	// the center of each rectangle. If they are "", they default
	// to the first and second columns, respectively.
	X, Y string

	// Width and Height name columns that define the width and
	// height of each rectangle. If they are "", the width and/or
	// height are automatically determined from the smallest
	// spacing between distinct X and Y points.
	Width, Height string

	// Fill names a column that defines the fill color of each
	// rectangle. If it is "", the default fill is black.
	Fill string

	// XXX Stroke color/width, opacity, center adjustment.
}

func (l LayerTiles) Apply(p *Plot) {
	defaultCols(p, &l.X, &l.Y)
	if l.Width != "" || l.Height != "" {
		// TODO: What scale are these in? (x+width) is in the
		// X scale, but width itself is not. It doesn't make
		// sense to train the X scale on width, and if there's
		// a scale transform, (x+width) has to happen before
		// the transform. OTOH, if x is discrete, I can't do
		// (x+width); maybe in that case you just can't
		// specify a width. OTOOH, if width is specified and
		// the value is unscaled, I could still do something
		// reasonable with that if x is discrete.
		panic("not implemented: non-default width/height")
	}
	p.marks = append(p.marks, plotMark{&markTiles{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerTags attaches text annotations to data points.
//
// TODO: Currently this makes one annotation per group. This should be
// an option.
type LayerTags struct {
	// X and Y name columns that define the input and response
	// each tag is attached to. If they are "", they default to
	// the first and second columns, respectively.
	X, Y string

	// Label names the column that gives the text to put in the
	// tag at X, Y. Label is required.
	Label string
}

func (l LayerTags) Apply(p *Plot) {
	// TODO: Should there be special "annotation marks" that are
	// always on top and can perhaps extend outside the plot area?

	defaultCols(p, &l.X, &l.Y)
	// TODO: I keep wanting an abstraction for a column across
	// groups like this.
	labels := make(map[table.GroupID]table.Slice)
	for _, gid := range p.Data().Tables() {
		labels[gid] = p.Data().Table(gid).MustColumn(l.Label)
	}

	p.marks = append(p.marks, plotMark{&markTags{
		p.use("x", l.X),
		p.use("y", l.Y),
		labels,
	}, p.Data().Tables()})
}

// LayerTooltips attaches hover tooltips to data points.
type LayerTooltips struct {
	// X and Y name columns that define locations of tooltips. If
	// they are "", they default to the first and second columns,
	// respectively.
	X, Y string

	// Label names the column that gives the text of the tooltip.
	Label string

	// TODO: Text styling, closest X or closest point, multiple
	// tooltips if there are multiple points at the same X with
	// different Ys?
}

func (l LayerTooltips) Apply(p *Plot) {
	defer p.Save().Restore()

	defaultCols(p, &l.X, &l.Y)

	// Split up by subplot and flatten each subplot.
	tables := map[*subplot][]*table.Table{}
	for _, gid := range p.Data().Tables() {
		s := subplotOf(gid)
		tables[s] = append(tables[s], p.Data().Table(gid))
	}
	var ng table.GroupingBuilder
	for k, ts := range tables {
		var subg table.GroupingBuilder
		for i, t := range ts {
			subg.Add(table.RootGroupID.Extend(i), t)
		}
		ng.Add(table.RootGroupID.Extend(k), table.Flatten(subg.Done()))
	}
	p.SetData(ng.Done())

	labels := make(map[table.GroupID]table.Slice)
	for _, gid := range p.Data().Tables() {
		labels[gid] = p.Data().Table(gid).MustColumn(l.Label)
	}
	p.marks = append(p.marks, plotMark{&markTooltips{
		p.use("x", l.X),
		p.use("y", l.Y),
		labels,
	}, p.Data().Tables()})
}
<commit_msg>gg: make LayerTags group by label<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 gg

import (
	"fmt"

	"github.com/aclements/go-gg/table"
)

func defaultCols(p *Plot, cols ...*string) {
	dcols := p.Data().Columns()
	for i, colp := range cols {
		if *colp == "" {
			if i >= len(dcols) {
				panic(fmt.Sprintf("cannot get default column %d; table has only %d columns", i, len(dcols)))
			}
			*colp = dcols[i]
		}
	}
}

// LayerLines is like LayerPaths, but connects data points in order by
// the "x" property.
type LayerLines LayerPaths

func (l LayerLines) Apply(p *Plot) {
	LayerPaths(l).apply(p, true)
}

//go:generate stringer -type StepMode

// StepMode controls how LayerSteps connects subsequent points.
type StepMode int

const (
	// StepHV makes LayerSteps connect subsequent points with a
	// horizontal segment and then a vertical segment.
	StepHV StepMode = iota

	// StepVH makes LayerSteps connect subsequent points with a
	// vertical segment and then a horizontal segment.
	StepVH

	// StepHMid makes LayerSteps connect subsequent points A and B
	// with three segments: a horizontal segment from A to the
	// midpoint between A and B, followed by vertical segment,
	// followed by a horizontal segment from the midpoint to B.
	StepHMid

	// StepVMid makes LayerSteps connect subsequent points A and B
	// with three segments: a vertical segment from A to the
	// midpoint between A and B, followed by horizontal segment,
	// followed by a vertical segment from the midpoint to B.
	StepVMid
)

// LayerSteps is like LayerPaths, but connects data points with a path
// consisting only of horizontal and vertical segments.
type LayerSteps struct {
	LayerPaths

	Step StepMode
}

func (l LayerSteps) Apply(p *Plot) {
	// TODO: Should this also support only showing horizontal or
	// vertical segments?
	//
	// TODO: This could be a data transform instead of a layer.
	// Then it could be used in conjunction with, for example,
	// ribbons.

	defaultCols(p, &l.X, &l.Y)
	p.marks = append(p.marks, plotMark{&markSteps{
		l.Step,
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPaths groups by Color and Fill, and then connects successive
// data points in each group with a path and/or a filled polygon.
type LayerPaths struct {
	// X and Y name columns that define the input and response of
	// each point on the path. If these are empty, they default to
	// the first and second columns, respectively.
	X, Y string

	// Color names a column that defines the stroke color of each
	// path. If Color is "", it defaults to constant black.
	// Otherwise, the data is grouped by Color.
	Color string

	// Fill names a column that defines the fill color of each
	// path. If Fill is "", it defaults to none. Otherwise, the
	// data is grouped by Fill.
	Fill string

	// XXX Perhaps the theme should provide default values for
	// things like "color". That would suggest we need to resolve
	// defaults like that at render time. Possibly a special scale
	// that gets values from the theme could be used to resolve
	// them.
	//
	// XXX strokeOpacity, fillOpacity, strokeWidth, what other
	// properties do SVG strokes have?
	//
	// XXX Should the set of known styling bindings be fixed, and
	// all possible rendering targets have to know what to do with
	// them, or should the rendering target be able to have
	// different styling bindings they understand (presumably with
	// some reasonable base set)? If the renderer can determine
	// the known bindings, we would probably just capture the
	// environment here (and make it so a captured environment
	// does not change) and hand that to the renderer later.
}

func (l LayerPaths) Apply(p *Plot) {
	l.apply(p, false)
}

func (l LayerPaths) apply(p *Plot, sort bool) {
	defaultCols(p, &l.X, &l.Y)
	if l.Color != "" {
		p.GroupBy(l.Color)
	}
	if l.Fill != "" {
		p.GroupBy(l.Fill)
	}
	if sort {
		defer p.Save().Restore()
		p = p.SortBy(l.X)
	}

	p.marks = append(p.marks, plotMark{&markPath{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("stroke", l.Color),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerPoints layers a point mark at each data point.
type LayerPoints struct {
	// X and Y name columns that define input and response of each
	// point. If these are empty, they default to the first and
	// second columns, respectively.
	X, Y string

	// Color names the column that defines the fill color of each
	// point. If Color is "", it defaults to constant black.
	Color string

	// Opacity names the column that defines the opacity of each
	// point. If Opacity is "", it defaults to fully opaque. This
	// is multiplied by any alpha value specified by Color.
	Opacity string

	// Size names the column that defines the size of each point.
	// If Size is "", it defaults to 1% of the smallest plot
	// dimension.
	Size string

	// XXX fill vs stroke, shape
}

func (l LayerPoints) Apply(p *Plot) {
	defaultCols(p, &l.X, &l.Y)
	p.marks = append(p.marks, plotMark{&markPoint{
		p.use("x", l.X),
		p.use("y", l.Y),
		// TODO: It's actually the fill color, but I generally
		// want it to match things that are stroke colors.
		// Maybe I should have a "color" aesthetic for the
		// "primary" color? Or I could have a hierarchy of
		// aesthetics, in which this uses "stroke" if it has a
		// scale, but otherwise uses "color".
		p.use("stroke", l.Color),
		// TODO: What scale for opacity? Or should I assume
		// callers will use PreScaled values if they want
		// specific opacities? What's the physical type?
		p.use("opacity", l.Opacity),
		p.use("size", l.Size),
	}, p.Data().Tables()})
}

// LayerTiles layers a rectangle at each data point. The rectangle is
// specified by its center, width, and height.
type LayerTiles struct {
	// X and Y name columns that define the input and response at
	// the center of each rectangle. If they are "", they default
	// to the first and second columns, respectively.
	X, Y string

	// Width and Height name columns that define the width and
	// height of each rectangle. If they are "", the width and/or
	// height are automatically determined from the smallest
	// spacing between distinct X and Y points.
	Width, Height string

	// Fill names a column that defines the fill color of each
	// rectangle. If it is "", the default fill is black.
	Fill string

	// XXX Stroke color/width, opacity, center adjustment.
}

func (l LayerTiles) Apply(p *Plot) {
	defaultCols(p, &l.X, &l.Y)
	if l.Width != "" || l.Height != "" {
		// TODO: What scale are these in? (x+width) is in the
		// X scale, but width itself is not. It doesn't make
		// sense to train the X scale on width, and if there's
		// a scale transform, (x+width) has to happen before
		// the transform. OTOH, if x is discrete, I can't do
		// (x+width); maybe in that case you just can't
		// specify a width. OTOOH, if width is specified and
		// the value is unscaled, I could still do something
		// reasonable with that if x is discrete.
		panic("not implemented: non-default width/height")
	}
	p.marks = append(p.marks, plotMark{&markTiles{
		p.use("x", l.X),
		p.use("y", l.Y),
		p.use("fill", l.Fill),
	}, p.Data().Tables()})
}

// LayerTags attaches text annotations to data points.
//
// TODO: Currently this groups by label and makes one annotation per
// group. This should be a controllable.
type LayerTags struct {
	// X and Y name columns that define the input and response
	// each tag is attached to. If they are "", they default to
	// the first and second columns, respectively.
	X, Y string

	// Label names the column that gives the text to put in the
	// tag at X, Y. Label is required.
	Label string
}

func (l LayerTags) Apply(p *Plot) {
	// TODO: Should there be special "annotation marks" that are
	// always on top and can perhaps extend outside the plot area?

	defaultCols(p, &l.X, &l.Y)
	defer p.Save().Restore()
	p.GroupBy(l.Label)
	// TODO: I keep wanting an abstraction for a column across
	// groups like this.
	labels := make(map[table.GroupID]table.Slice)
	for _, gid := range p.Data().Tables() {
		labels[gid] = p.Data().Table(gid).MustColumn(l.Label)
	}

	p.marks = append(p.marks, plotMark{&markTags{
		p.use("x", l.X),
		p.use("y", l.Y),
		labels,
	}, p.Data().Tables()})
}

// LayerTooltips attaches hover tooltips to data points.
type LayerTooltips struct {
	// X and Y name columns that define locations of tooltips. If
	// they are "", they default to the first and second columns,
	// respectively.
	X, Y string

	// Label names the column that gives the text of the tooltip.
	Label string

	// TODO: Text styling, closest X or closest point, multiple
	// tooltips if there are multiple points at the same X with
	// different Ys?
}

func (l LayerTooltips) Apply(p *Plot) {
	defer p.Save().Restore()

	defaultCols(p, &l.X, &l.Y)

	// Split up by subplot and flatten each subplot.
	tables := map[*subplot][]*table.Table{}
	for _, gid := range p.Data().Tables() {
		s := subplotOf(gid)
		tables[s] = append(tables[s], p.Data().Table(gid))
	}
	var ng table.GroupingBuilder
	for k, ts := range tables {
		var subg table.GroupingBuilder
		for i, t := range ts {
			subg.Add(table.RootGroupID.Extend(i), t)
		}
		ng.Add(table.RootGroupID.Extend(k), table.Flatten(subg.Done()))
	}
	p.SetData(ng.Done())

	labels := make(map[table.GroupID]table.Slice)
	for _, gid := range p.Data().Tables() {
		labels[gid] = p.Data().Table(gid).MustColumn(l.Label)
	}
	p.marks = append(p.marks, plotMark{&markTooltips{
		p.use("x", l.X),
		p.use("y", l.Y),
		labels,
	}, p.Data().Tables()})
}
<|endoftext|>


<commit_before>package godotenv

import (
	"bufio"
	"errors"
	"os"
	"strings"
)

func Load(filenames ...string) (err error) {
	for _, filename := range filenames {
		err = loadFile(filename)
		if err != nil {
			return // return early on a spazout
		}
	}
	return
}

func loadFile(filename string) (err error) {
	file, err := os.Open(filename)
	if err != nil {
		return
	}

	bufferSize := 20
	lines := make([]string, bufferSize)
	lineReader := bufio.NewReaderSize(file, bufferSize)
	for line, isPrefix, e := lineReader.ReadLine(); e == nil; line, isPrefix, e = lineReader.ReadLine() {
		fullLine := string(line)
		if isPrefix {
			for {
				line, isPrefix, _ = lineReader.ReadLine()
				fullLine += string(line)
				if !isPrefix {
					break
				}
			}
		}
		// add a line to the game/parse
		lines = append(lines, string(line))
	}

	for _, fullLine := range lines {
		key, value, err := parseLine(fullLine)

		if err == nil {
			os.Setenv(key, value)
		}
	}

	return
}

func parseLine(line string) (key string, value string, err error) {
	if len(line) == 0 {
		err = errors.New("zero length string")
		return
	}

	splitString := strings.Split(line, "=")

	if len(splitString) != 2 {
		// try yaml mode!
		splitString = strings.Split(line, ":")
	}

	if len(splitString) != 2 {
		err = errors.New("Can't separate key from value")
		return
	}

	key = splitString[0]
	if strings.HasPrefix(key, "export") {
		key = strings.TrimPrefix(key, "export")
	}
	key = strings.Trim(key, " ")

	value = splitString[1]

	// ditch the comments
	if strings.Contains(value, "#") {
		value = strings.Trim(strings.Split(value, "#")[0], " ")
	}

	// check if we've got quoted values
	if strings.Count(value, "\"") == 2 || strings.Count(value, "'") == 2 {
		// pull the quotes off the edge
		value = strings.Trim(value, "\"'")

		// expand quotes
		value = strings.Replace(value, "\\\"", "\"", -1)
		// expand newlines
		value = strings.Replace(value, "\\n", "\n", -1)
	}
	// trim
	value = strings.Trim(value, " ")

	return
}
<commit_msg>World's most naive comment parser...<commit_after>package godotenv

import (
	"bufio"
	"errors"
	"os"
	"strings"
)

func Load(filenames ...string) (err error) {
	for _, filename := range filenames {
		err = loadFile(filename)
		if err != nil {
			return // return early on a spazout
		}
	}
	return
}

func loadFile(filename string) (err error) {
	file, err := os.Open(filename)
	if err != nil {
		return
	}

	bufferSize := 20
	lines := make([]string, bufferSize)
	lineReader := bufio.NewReaderSize(file, bufferSize)
	for line, isPrefix, e := lineReader.ReadLine(); e == nil; line, isPrefix, e = lineReader.ReadLine() {
		fullLine := string(line)
		if isPrefix {
			for {
				line, isPrefix, _ = lineReader.ReadLine()
				fullLine += string(line)
				if !isPrefix {
					break
				}
			}
		}
		// add a line to the game/parse
		lines = append(lines, string(line))
	}

	for _, fullLine := range lines {
		key, value, err := parseLine(fullLine)

		if err == nil {
			os.Setenv(key, value)
		}
	}

	return
}

func parseLine(line string) (key string, value string, err error) {
	if len(line) == 0 {
		err = errors.New("zero length string")
		return
	}

	splitString := strings.Split(line, "=")

	if len(splitString) != 2 {
		// try yaml mode!
		splitString = strings.Split(line, ":")
	}

	if len(splitString) != 2 {
		err = errors.New("Can't separate key from value")
		return
	}

	key = splitString[0]
	if strings.HasPrefix(key, "export") {
		key = strings.TrimPrefix(key, "export")
	}
	key = strings.Trim(key, " ")

	value = splitString[1]

	// ditch the comments
	if strings.Contains(value, "#") {
		segmentsBetweenHashes := strings.Split(value, "#")
		value = segmentsBetweenHashes[0]
		// open quote in leftmost segment
		if strings.Count(value, "\"") == 1 {
			value = value + "#" + segmentsBetweenHashes[1]
		}
	}

	// check if we've got quoted values
	if strings.Count(value, "\"") == 2 || strings.Count(value, "'") == 2 {
		// pull the quotes off the edges
		value = strings.Trim(value, "\"' ")

		// expand quotes
		value = strings.Replace(value, "\\\"", "\"", -1)
		// expand newlines
		value = strings.Replace(value, "\\n", "\n", -1)
	}
	// trim
	value = strings.Trim(value, " ")

	return
}
<|endoftext|>


<commit_before>/*
package goptions implements a flexible parser for command line options.

Key targets were the support for both long and short flag versions, mutually
exclusive flags, and verbs. Flags and their corresponding variables are defined
by the tags in a (possibly anonymous) struct.

    var options struct {
    	Name string `goptions:"-n, --name"`
    	Force bool `goptions:"-f, --force"`
    	Verbosity int `goptions:"-v, --verbose, accumulate"`
    }

Short flags can be combined (e.g. `-nfv`). Long flags take their value after a
separating space. The equals notation (`--long-flag=value`) is NOT supported
right now.

Every member of the struct, which is supposed to catch a command line value
has to have a "goptions" tag. Multiple short and long flag names can be specified.
Each tag can also list any number of the following options:

    accumulate        - (Only valid for `int`) Counts how of then the flag has been
                        specified in the short version. The long version simply
                        accepts an int.
    obligatory        - Flag must be specified. Otherwise an error will be returned
                        when Parse() is called.
    description='...' - Set the description for this particular flag. Will be
                        used by the HelpFunc.
    mutexgroup='...'  - Sets the name of the MutexGroup. Only one flag of the
                        ones sharing a MutexGroup can be set. Otherwise an error
                        will be returned when Parse() is called.

goptions also has support for verbs. Each verb accepts its own set of flags which
take exactly the same tag format as global options. For an usage example of verbs
see the PrintHelp() example.
*/
package goptions

import (
	"io"
	"os"
	"sync"
	"text/tabwriter"
	"text/template"
)

const (
	VERSION = "1.1.0"
)

var (
	globalFlagSet *FlagSet
)

// Parse parses the command-line flags from os.Args[1:].
func Parse(v interface{}) error {
	globalFlagSet = NewFlagSet(os.Args[0], v)
	return globalFlagSet.Parse(os.Args[1:])
}

// PrintHelp renders the default help to os.Stderr.
func PrintHelp() {
	if globalFlagSet == nil {
		panic("Must call Parse() before PrintHelp()")
	}
	globalFlagSet.PrintHelp(os.Stderr)
}

// Generates a new HelpFunc taking a `text/template.Template`-formatted
// string as an argument. The resulting template will be executed with the FlagSet
// as its data.
func NewTemplatedHelpFunc(tpl string) HelpFunc {
	var once sync.Once
	var t *template.Template
	return func(w io.Writer, fs *FlagSet) {
		once.Do(func() {
			t = template.Must(template.New("helpTemplate").Parse(tpl))
		})
		err := t.Execute(w, fs)
		if err != nil {
			panic(err)
		}
	}
}

const (
	_DEFAULT_HELP = `
Usage: {{.Name}} [global options] {{with .Verbs}}<verb> [verb options]{{end}}

Global options:{{range .Flags}}
	{{if len .Short}}-{{index .Short 0}},{{end}}	{{if len .Long}}--{{index .Long 0}}{{end}}	{{.Description}}{{if .Obligatory}} (*){{end}}{{end}}

{{if .Verbs}}Verbs:{{range .Verbs}}
	{{.Name}}:{{range .Flags}}
		{{if len .Short}}-{{index .Short 0}},{{end}}	{{if len .Long}}--{{index .Long 0}}{{end}}	{{.Description}}{{if .Obligatory}} (*){{end}}{{end}}{{end}}{{end}}
`
)

// DefaultHelpFunc is a HelpFunc which renders the default help template and pipes
// the output through a text/tabwriter.Writer before flushing it to the output.
func DefaultHelpFunc(w io.Writer, fs *FlagSet) {
	tw := &tabwriter.Writer{}
	tw.Init(w, 4, 4, 1, ' ', 0)
	NewTemplatedHelpFunc(_DEFAULT_HELP)(tw, fs)
	tw.Flush()
}
<commit_msg>Use filename without path for help message<commit_after>/*
package goptions implements a flexible parser for command line options.

Key targets were the support for both long and short flag versions, mutually
exclusive flags, and verbs. Flags and their corresponding variables are defined
by the tags in a (possibly anonymous) struct.

    var options struct {
    	Name string `goptions:"-n, --name"`
    	Force bool `goptions:"-f, --force"`
    	Verbosity int `goptions:"-v, --verbose, accumulate"`
    }

Short flags can be combined (e.g. `-nfv`). Long flags take their value after a
separating space. The equals notation (`--long-flag=value`) is NOT supported
right now.

Every member of the struct, which is supposed to catch a command line value
has to have a "goptions" tag. Multiple short and long flag names can be specified.
Each tag can also list any number of the following options:

    accumulate        - (Only valid for `int`) Counts how of then the flag has been
                        specified in the short version. The long version simply
                        accepts an int.
    obligatory        - Flag must be specified. Otherwise an error will be returned
                        when Parse() is called.
    description='...' - Set the description for this particular flag. Will be
                        used by the HelpFunc.
    mutexgroup='...'  - Sets the name of the MutexGroup. Only one flag of the
                        ones sharing a MutexGroup can be set. Otherwise an error
                        will be returned when Parse() is called.

goptions also has support for verbs. Each verb accepts its own set of flags which
take exactly the same tag format as global options. For an usage example of verbs
see the PrintHelp() example.
*/
package goptions

import (
	"io"
	"os"
	"path/filepath"
	"sync"
	"text/tabwriter"
	"text/template"
)

const (
	VERSION = "1.1.0"
)

var (
	globalFlagSet *FlagSet
)

// Parse parses the command-line flags from os.Args[1:].
func Parse(v interface{}) error {
	globalFlagSet = NewFlagSet(filepath.Base(os.Args[0]), v)
	return globalFlagSet.Parse(os.Args[1:])
}

// PrintHelp renders the default help to os.Stderr.
func PrintHelp() {
	if globalFlagSet == nil {
		panic("Must call Parse() before PrintHelp()")
	}
	globalFlagSet.PrintHelp(os.Stderr)
}

// Generates a new HelpFunc taking a `text/template.Template`-formatted
// string as an argument. The resulting template will be executed with the FlagSet
// as its data.
func NewTemplatedHelpFunc(tpl string) HelpFunc {
	var once sync.Once
	var t *template.Template
	return func(w io.Writer, fs *FlagSet) {
		once.Do(func() {
			t = template.Must(template.New("helpTemplate").Parse(tpl))
		})
		err := t.Execute(w, fs)
		if err != nil {
			panic(err)
		}
	}
}

const (
	_DEFAULT_HELP = `
Usage: {{.Name}} [global options] {{with .Verbs}}<verb> [verb options]{{end}}

Global options:{{range .Flags}}
	{{if len .Short}}-{{index .Short 0}},{{end}}	{{if len .Long}}--{{index .Long 0}}{{end}}	{{.Description}}{{if .Obligatory}} (*){{end}}{{end}}

{{if .Verbs}}Verbs:{{range .Verbs}}
	{{.Name}}:{{range .Flags}}
		{{if len .Short}}-{{index .Short 0}},{{end}}	{{if len .Long}}--{{index .Long 0}}{{end}}	{{.Description}}{{if .Obligatory}} (*){{end}}{{end}}{{end}}{{end}}
`
)

// DefaultHelpFunc is a HelpFunc which renders the default help template and pipes
// the output through a text/tabwriter.Writer before flushing it to the output.
func DefaultHelpFunc(w io.Writer, fs *FlagSet) {
	tw := &tabwriter.Writer{}
	tw.Init(w, 4, 4, 1, ' ', 0)
	NewTemplatedHelpFunc(_DEFAULT_HELP)(tw, fs)
	tw.Flush()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
)

var (
	inputDir    string         // user input; top-level path to search
	searchText  string         // user input; keyword to search
	verbose     bool           // user input; if true displays all paths
	numFound    int            // # of files matching keyword
	fileVisit   int            // # of files visited by search
	dirFound    int            // # of directories matching keyword
	folderVisit int            // # of folders visited by search
	wg          sync.WaitGroup // sync goroutines / channels
	lock        sync.Mutex     // control access to counters (race prevention)
	maxSize     int64          // max file size
)

type walkresult struct {
	path    string
	name    string
	found   bool
	isDir   bool
	size    int64
	modTime time.Time
}

func usage() {
	fmt.Println("Usage:")
	fmt.Println("    gosearch [OPTIONS] -p path -k keyword")
	flag.PrintDefaults()
}

func init() {
	// log set to JSON format
	log.SetFormatter(&log.JSONFormatter{})

	// max file size to process
	maxSize = 100 * 1024 * 1024

	// flag init
	flag.StringVar(&inputDir, "p", "", "Path to directory to search")
	flag.StringVar(&searchText, "k", "", "Keyword to search")
	flag.BoolVar(&verbose, "v", false, "Verbose (prints all files searched)")
}

// duration keeps track of function elapsed time
func duration(start time.Time, name string) {
	elapsed := time.Since(start)
	fmt.Printf("func %s elapsed %s\n", name, elapsed)
}

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

func errorOut(message string) bool {
	fmt.Fprintln(os.Stderr, message)
	return false
}

// check path exists
func exists(path string) bool {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

// walkFiles walks all files and sub-directory paths
func walkFiles(directory string, keyword string, filesFound chan walkresult, done chan bool) {

	// launch goroutine to walk path; add wait count
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := filepath.Walk(directory, func(path string, f os.FileInfo, err error) error {
			errorCheck(err)

			// if file launch main search process
			if !f.IsDir() {
				fileCount()

				// only launch search if file is under size limit,
				if f.Size() < maxSize {
					wg.Add(1)
					go readFile(path, f, filesFound)
				} else {
					fmt.Printf("%s skipped. File too large.", path)
				}
			}

			// folder path, increment count
			folderCount()
			wg.Add(1)
			go searchPath(path, f, filesFound)
			return nil
		})

		// launch cleanup, but sync wait until goroutines complete
		go cleanup(filesFound, done)

		// check errors for walk func
		errorCheck(err)
		return
	}()
	return
}

// readFile puts contents of file in memory
func readFile(path string, f os.FileInfo, filesFound chan walkresult) {
	defer wg.Done()
	content, err := ioutil.ReadFile(path)
	if err != nil {
		if !verbose {
			return
		}
		fmt.Printf("%s FILE cannot be read\n", path)
		return
	}
	wg.Add(1)
	go searchFile(path, content, f, filesFound)
}

// searchFile parses the contents of file looking for keyword
func searchFile(path string, content []byte, f os.FileInfo, filesFound chan walkresult) {
	defer wg.Done()
	x := string(content)
	search := strings.Contains(x, searchText)
	switch search {
	case true:
		lock.Lock()
		numFound++
		lock.Unlock()
		found := true
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	case false:
		found := false
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	}
}

// searchPath searches match in file or folder name
func searchPath(path string, f os.FileInfo, filesFound chan walkresult) {
	defer wg.Done()
	search := strings.Contains(f.Name(), searchText)
	switch search {
	case true:
		if f.IsDir() {
			lock.Lock()
			dirFound++
			lock.Unlock()
		} else {
			lock.Lock()
			numFound++
			lock.Unlock()
		}
		found := true
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	case false:
		found := false
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	}
}

// folderCount keeps count of folders visited during search
func folderCount() {
	lock.Lock()
	folderVisit++
	lock.Unlock()
}

// fileCount keeps count of files visited during search
func fileCount() {
	lock.Lock()
	fileVisit++
	lock.Unlock()
}

// waits for goroutines to complete, sets done signal and closes channels
func cleanup(filesFound chan walkresult, done chan bool) {
	wg.Wait()
	close(filesFound)
	done <- true
	<-done
	close(done)
	return
}

// summary prints results, counts, lets user know search is done
func summary() {
	fmt.Println("==================================")
	log.Printf("Done searching for %s\n", searchText)
	log.Printf("Path: %s\n", inputDir)
	log.Printf("Checked %d files in %d folders\n", fileVisit, folderVisit)
	log.Printf("Found %d files containing %s\n", numFound, searchText)
	log.Printf("Found %d folders containing %s\n", dirFound, searchText)
	fmt.Println("==================================")
}

func main() {
	// main timer
	defer duration(time.Now(), "main")

	// user messaging
	fmt.Println("==================================")
	fmt.Println("gosearch: A search in text utility written in Go.")
	fmt.Println("searching...")
	fmt.Println("==================================")

	// check args provided
	flag.Parse()
	ok := true
	if inputDir == "" {
		ok = errorOut("ERROR: Missing path to directory")
	} else {
		// check path exists
		verify := exists(inputDir)
		if !verify {
			ok = errorOut("ERROR: Path provided does not exist.")
		}
	}
	if searchText == "" {
		ok = errorOut("ERROR: Missing keyword to search")
	}

	if !ok {
		usage()
		os.Exit(1)
	}

	// create channels
	filesFound := make(chan walkresult)
	done := make(chan bool)

	// start search work
	go walkFiles(inputDir, searchText, filesFound, done)

	// receive channel results and print
loop:
	for {
		select {
		case print := <-filesFound:
			if (len(print.path) > 0) && verbose && (print.found == false) {
				switch print.isDir {
				case true:
					log.WithFields(log.Fields{
						"type": "folder",
						"name": print.name,
						"path": print.path,
					}).Info("Match NOT found.")
				case false:
					log.WithFields(log.Fields{
						"type": "file",
						"name": print.name,
						"path": print.path,
					}).Info("Match NOT found.")
				}
			}
			if print.found == true {
				switch print.isDir {
				case true:
					log.WithFields(log.Fields{
						"type": "folder",
						"name": print.name,
						"path": print.path,
					}).Info("Match found.")
				case false:
					log.WithFields(log.Fields{
						"type": "file",
						"name": print.name,
						"path": print.path,
					}).Info("Match found.")
				}

			}
		case <-done:
			fmt.Println("==================================")
			log.Println("Search complete.")
			done <- true
			break loop
		}
	}

	// print search summary, file counts
	summary()
}
<commit_msg>Add logging options.<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
)

var (
	inputDir    string         // user input; top-level path to search
	searchText  string         // user input; keyword to search
	verbose     bool           // user input; if true displays all paths
	numFound    int            // # of files matching keyword
	fileVisit   int            // # of files visited by search
	dirFound    int            // # of directories matching keyword
	folderVisit int            // # of folders visited by search
	wg          sync.WaitGroup // sync goroutines / channels
	lock        sync.Mutex     // control access to counters (race prevention)
	maxSize     int64          // max file size
	json        bool           // output in json if true
	help        bool           // display help if true
	// output      string         // output file name
)

// walkresult struct for result document
type walkresult struct {
	path    string
	name    string
	found   bool
	isDir   bool
	size    int64
	modTime time.Time
}

func usage() {
	// user messaging
	fmt.Println("==================================")
	fmt.Println("gosearch: A search-in-text utility written in Go.")
	fmt.Println("==================================")
	fmt.Println("Usage:")
	fmt.Println("    gosearch [OPTIONS] -p path -k keyword")
	flag.PrintDefaults()
}

func init() {
	// flag init
	flag.StringVar(&inputDir, "p", "", "Path of directory to search")
	flag.StringVar(&searchText, "k", "", "Keyword to search")
	flag.Int64Var(&maxSize, "s", 100, "Max file size to search in MB - optional")
	flag.BoolVar(&json, "j", false, "Output in JSON - optional")
	// flag.StringVar(&output, "o", "out.log", "Output file name - optional")
	flag.BoolVar(&verbose, "v", false, "Verbose = optional (prints all files searched)")
	flag.BoolVar(&help, "h", false, "Print help menu")
}

// duration keeps track of function elapsed time
func duration(start time.Time, name string) {
	elapsed := time.Since(start)
	log.Printf("func %s elapsed %s\n", name, elapsed)
}

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

func errorOut(message string) bool {
	fmt.Fprintln(os.Stderr, message)
	return false
}

// check path exists
func exists(path string) bool {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

// walkFiles walks all files and sub-directory paths
func walkFiles(directory string, keyword string, filesFound chan walkresult, done chan bool) {

	// launch goroutine to walk path; add wait count
	wg.Add(1)
	go func() {
		defer wg.Done()
		err := filepath.Walk(directory, func(path string, f os.FileInfo, err error) error {
			errorCheck(err)

			// if file launch main search process
			if !f.IsDir() {
				fileCount()

				// only launch search if file is under size limit,
				if f.Size() < maxSize*1024*1024 {
					wg.Add(1)
					go readFile(path, f, filesFound)
				} else {
					log.WithFields(log.Fields{
						"type": "file",
						"name": f.Name(),
						"path": path,
					}).Warn("Skip file too large: ", f.Size())
				}
			}

			// folder path, increment count
			folderCount()
			wg.Add(1)
			go searchPath(path, f, filesFound)
			return nil
		})

		// launch cleanup, but sync wait until goroutines complete
		go cleanup(filesFound, done)

		// check errors for walk func
		errorCheck(err)
		return
	}()
	return
}

// readFile puts contents of file in memory, starts search
func readFile(path string, f os.FileInfo, filesFound chan walkresult) {
	defer wg.Done()
	content, err := ioutil.ReadFile(path)
	if err != nil {
		if !verbose {
			return
		}
		log.WithFields(log.Fields{
			"type": "file",
			"name": f.Name(),
			"path": path,
		}).Warn("File cannot be read", f.Size())
		return
	}
	wg.Add(1)
	go searchFile(path, content, f, filesFound)
}

// searchFile parses the contents of file looking for keyword
func searchFile(path string, content []byte, f os.FileInfo, filesFound chan walkresult) {
	defer wg.Done()
	x := string(content)
	search := strings.Contains(x, searchText)
	switch search {
	case true:
		lock.Lock()
		numFound++
		lock.Unlock()
		found := true
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	case false:
		found := false
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	}
}

// searchPath searches match in file or folder name
func searchPath(path string, f os.FileInfo, filesFound chan walkresult) {
	defer wg.Done()
	search := strings.Contains(f.Name(), searchText)
	switch search {
	case true:
		if f.IsDir() {
			lock.Lock()
			dirFound++
			lock.Unlock()
		} else {
			lock.Lock()
			numFound++
			lock.Unlock()
		}
		found := true
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	case false:
		found := false
		filesFound <- walkresult{path, f.Name(), found, f.IsDir(), f.Size(), f.ModTime()}
		return
	}
}

// folderCount keeps count of folders visited during search
func folderCount() {
	lock.Lock()
	folderVisit++
	lock.Unlock()
}

// fileCount keeps count of files visited during search
func fileCount() {
	lock.Lock()
	fileVisit++
	lock.Unlock()
}

// waits for goroutines to complete, sets done signal and closes channels
func cleanup(filesFound chan walkresult, done chan bool) {
	wg.Wait()
	close(filesFound)
	done <- true
	<-done
	close(done)
	return
}

// summary prints results, counts, lets user know search is done
func summary(searchText string, path string) {
	log.WithFields(log.Fields{
		"searchString":   searchText,  // text to search
		"path":           path,        // file path requeted to search
		"filesChecked":   fileVisit,   // num of files visited during search
		"foldersChecked": folderVisit, // num of folders visited during search
		"filesFound":     numFound,    // num of files that contain match for search string
		"foldersFound":   dirFound,    // num of folders that contain match for search string
	}).Info("Search completed")
}

func main() {
	// main timer
	defer duration(time.Now(), "main")

	// check args provided
	flag.Parse()
	ok := true

	if help == true {
		usage()
		os.Exit(1)
	}
	if inputDir == "" {
		ok = errorOut("ERROR: Missing path to directory")
	} else {
		// check path exists
		verify := exists(inputDir)
		if !verify {
			ok = errorOut("ERROR: Path provided does not exist.")
		}
	}
	if searchText == "" {
		ok = errorOut("ERROR: Missing keyword to search")
	}

	if !ok {
		usage()
		os.Exit(1)
	}

	// log set to JSON format
	if json == true {
		log.SetFormatter(&log.JSONFormatter{})
	} else {
		// The TextFormatter is default, you don't actually have to do this.
		log.SetFormatter(&log.TextFormatter{})
	}

	// output file definition
	/*
		out, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
		if err != nil {
			log.Fatalf("error opening log file: %v", err)
		}
		defer out.Close()
		multi := io.MultiWriter(out, os.Stdout)
		log.SetOutput(multi)
	*/

	// create channels
	filesFound := make(chan walkresult)
	done := make(chan bool)

	// notify user search started
	log.WithFields(log.Fields{
		"searchString": searchText,
		"path":         inputDir,
	}).Info("Search started")

	// start search work
	go walkFiles(inputDir, searchText, filesFound, done)

	// receive channel results and print
loop:
	for {
		select {
		case print := <-filesFound:
			if (len(print.path) > 0) && verbose && (print.found == false) {
				switch print.isDir {
				case true:
					log.WithFields(log.Fields{
						"type": "folder",
						"name": print.name,
						"path": print.path,
					}).Info("Match not found")
				case false:
					log.WithFields(log.Fields{
						"type": "file",
						"name": print.name,
						"path": print.path,
					}).Info("Match not found")
				}
			}
			if print.found == true {
				switch print.isDir {
				case true:
					log.WithFields(log.Fields{
						"type": "folder",
						"name": print.name,
						"path": print.path,
					}).Info("Match found")
				case false:
					log.WithFields(log.Fields{
						"type": "file",
						"name": print.name,
						"path": print.path,
					}).Info("Match found")
				}

			}
		case <-done:
			done <- true
			break loop
		}
	}

	// print search summary, file counts
	summary(searchText, inputDir)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"image"
	"image/color"
	"image/png"
	"log"
	"os"
	"os/signal"
	"strconv"
	"time"
)

const (
	SOLDER     = 0x01
	TERMINAL   = 0x02
	TIMESLIZES = 0x04
)

// Version number of this tool
const Version = "0.0.2"

// Data is a struct for each network packet
type Data struct {
	toa     int64  // Timestamp of arrival in microseconds
	payload []byte // Copied network packet
}

// configs represents all the configuration data
type configs struct {
	bpP   uint // Bits per Pixel
	ppI   uint // Number of packets per Image
	ts    uint // "Duration" for one Image
	limit uint // Number of network packets to process
	stil  uint // Type of illustration
}

func getBitsFromPacket(packet []byte, byteP, bitP *int, bpP uint) uint8 {
	var c uint8
	for i := 0; i < (int(bpP) / 3); i++ {
		if *byteP >= len(packet) {
			break
		}
		c |= (packet[*byteP] & (1 << uint8(7-*bitP)))
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	}
	return c
}

func createPixel(packet []byte, byteP, bitP *int, bpP uint) (c color.Color) {
	var r, g, b uint8

	if bpP == 1 {
		if (packet[*byteP] & (1 << uint8(7-*bitP))) == 0 {
			c = color.NRGBA{R: 0,
				G: 0,
				B: 0,
				A: 255}
		} else {
			c = color.NRGBA{R: 255,
				G: 255,
				B: 255,
				A: 255}
		}
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	} else {
		r = getBitsFromPacket(packet, byteP, bitP, bpP)
		g = getBitsFromPacket(packet, byteP, bitP, bpP)
		b = getBitsFromPacket(packet, byteP, bitP, bpP)

		c = color.NRGBA{R: r,
			G: g,
			B: b,
			A: 255}
	}
	return
}

func createTerminalVisualization(data []Data, bitsPerPixel uint) {
	var bitPos int
	var bytePos int
	var packetLen int

	for i := range data {
		packetLen = len(data[i].payload)
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[i].payload, &bytePos, &bitPos, bitsPerPixel)
			r, g, b, _ := c.RGBA()
			fmt.Printf("\x1B[0m\x1B[38;2;%d;%d;%dm\u2588", uint8(r), uint8(g), uint8(b))
			if bytePos >= packetLen {
				break
			}
		}
		fmt.Printf("\x1B[m\n")

	}

}
func createTimeVisualization(data []Data, xMax int, prefix string, ts uint, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int
	var firstPkg time.Time

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, int(ts)))

	for pkg := range data {
		if firstPkg.IsZero() {
			firstPkg = time.Unix(0, data[pkg].toa*int64(time.Microsecond))
		}
		packetLen = len(data[pkg].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[pkg].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, int(data[pkg].toa%int64(ts)), c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}
	}

	filename := prefix
	filename += "-"
	filename += firstPkg.Format(time.RFC3339Nano)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func createFixedVisualization(data []Data, xMax int, prefix string, num int, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, len(data)))

	for yPos := range data {
		packetLen = len(data[yPos].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[yPos].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, yPos, c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}

	}

	filename := prefix
	filename += strconv.Itoa(num)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func handlePackets(ps *gopacket.PacketSource, num uint, ch chan<- Data, done <-chan bool) {
	var count uint
	for packet := range ps.Packets() {
		var k Data
		select {
		case <-done:
			close(ch)
			return
		default:
		}
		count++
		if num != 0 && count > num {
			break
		}

		elements := packet.Data()
		if len(elements) == 0 {
			continue
		}
		k = Data{toa: (packet.Metadata().CaptureInfo.Timestamp.UnixNano() / int64(time.Microsecond)), payload: packet.Data()}
		ch <- k
	}
	close(ch)
	return
}

func availableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	for _, device := range devices {
		if len(device.Addresses) == 0 {
			continue
		}
		fmt.Println("Interface: ", device.Name)
		for _, address := range device.Addresses {
			fmt.Println("   IP address:  ", address.IP)
			fmt.Println("   Subnet mask: ", address.Netmask)
		}
		fmt.Println("")
	}
}

func initSource(dev, file *string) (handle *pcap.Handle, err error) {
	if len(*dev) > 0 {
		handle, err = pcap.OpenLive(*dev, 4096, true, pcap.BlockForever)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else if len(*file) > 0 {
		handle, err = pcap.OpenOffline(*file)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else {
		return nil, fmt.Errorf("Source is missing\n")
	}
	return
}

func checkConfig(cfg configs) error {
	if cfg.bpP%3 != 0 && cfg.bpP != 1 {
		return fmt.Errorf("%d must be divisible by three or should be one", cfg.bpP)
	} else if cfg.bpP > 25 {
		return fmt.Errorf("%d must be smaller than 25", cfg.bpP)
	}

	if cfg.ts > 0 {
		cfg.stil |= TIMESLIZES
	}

	if cfg.stil == (TIMESLIZES | TERMINAL) {
		return fmt.Errorf("-timeslize and -terminal can't be combined")
	} else if cfg.stil == 0 {
		// If way of stil is provided, we will stick to the default one
		cfg.stil |= SOLDER
	}
	return nil
}

func main() {
	var err error
	var handle *pcap.Handle
	var data []Data
	var xMax int
	var index int = 1
	osSig := make(chan os.Signal, 1)
	done := make(chan bool, 1)
	signal.Notify(osSig, os.Interrupt)
	var slicer int64
	var cfg configs
	ch := make(chan Data)

	go func() {
		<-osSig     // Blocking till interrupt signal is received
		osSig = nil // ignore further signals
		done <- true
	}()

	dev := flag.String("interface", "", "Choose an interface for online processing")
	file := flag.String("file", "", "Choose a file for offline processing")
	filter := flag.String("filter", "", "Set a specific filter")
	lst := flag.Bool("list_interfaces", false, "List available interfaces")
	vers := flag.Bool("version", false, "Show version")
	help := flag.Bool("help", false, "Show this help")
	terminalOut := flag.Bool("terminal", false, "Visualize on terminal")
	num := flag.Uint("count", 25, "Number of packets to process.\n\tIf argument is 0 the limit is removed")
	output := flag.String("prefix", "image", "Prefix of the resulting image")
	size := flag.Uint("size", 25, "Number of packets per image")
	bits := flag.Uint("bits", 24, "Number of bits per pixel.\n\tIt must be divisible by three and smaller than 25\n\tTo get black/white results, choose 1 as input.")
	ts := flag.Uint("timeslize", 0, "Number of microseconds per resulting image.\n\tSo each pixel of the height of the resulting image represents one microsecond")
	flag.Parse()

	if flag.NFlag() < 1 {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	if *lst {
		availableInterfaces()
		return
	}

	if *vers {
		fmt.Println("Version:", Version)
		return
	}

	if *help {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	cfg.bpP = *bits
	cfg.ppI = *size
	cfg.ts = *ts
	cfg.limit = *num
	cfg.stil = 0

	if *terminalOut == true {
		cfg.stil |= TERMINAL
	}
	if *ts != 0 {
		cfg.stil |= TIMESLIZES
	}

	if err = checkConfig(cfg); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	handle, err = initSource(dev, file)
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
	defer handle.Close()

	if len(*filter) != 0 {
		err = handle.SetBPFFilter(*filter)
		if err != nil {
			log.Fatal(err, "\tInvalid filter: ", *filter)
			os.Exit(1)
		}
	}

	packetSource := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	packetSource.DecodeOptions = gopacket.Lazy

	go handlePackets(packetSource, cfg.limit, ch, done)

	switch cfg.stil {
	case SOLDER:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
			if len(data) >= int(*size) {
				xMax++
				createFixedVisualization(data, xMax, *output, index, cfg.bpP)
				xMax = 0
				index++
				data = data[:0]
			}
		}
	case TERMINAL:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			createTerminalVisualization(data, cfg.bpP)
			data = data[:0]
		}
	case TIMESLIZES:
		for i, ok := <-ch; ok; i, ok = <-ch {
			if slicer == 0 {
				slicer = i.toa + int64(*ts)
			}
			if slicer < i.toa {
				xMax++
				createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
				xMax = 0
				data = data[:0]
				slicer = i.toa + int64(*ts)
			}
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
		}
	}

	if len(data) > 0 {
		xMax++
		switch cfg.stil {
		case SOLDER:
			createFixedVisualization(data, xMax, *output, index, cfg.bpP)
		case TERMINAL:
			createTerminalVisualization(data, cfg.bpP)
		case TIMESLIZES:
			createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
		}
	}

}
<commit_msg>Restructuring<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"image"
	"image/color"
	"image/png"
	"log"
	"os"
	"os/signal"
	"strconv"
	"time"
)

const (
	SOLDER     = 0x01
	TERMINAL   = 0x02
	TIMESLIZES = 0x04
)

// Version number of this tool
const Version = "0.0.2"

// Data is a struct for each network packet
type Data struct {
	toa     int64  // Timestamp of arrival in microseconds
	payload []byte // Copied network packet
}

// configs represents all the configuration data
type configs struct {
	bpP   uint // Bits per Pixel
	ppI   uint // Number of packets per Image
	ts    uint // "Duration" for one Image
	limit uint // Number of network packets to process
	stil  uint // Type of illustration
}

func getBitsFromPacket(packet []byte, byteP, bitP *int, bpP uint) uint8 {
	var c uint8
	for i := 0; i < (int(bpP) / 3); i++ {
		if *byteP >= len(packet) {
			break
		}
		c |= (packet[*byteP] & (1 << uint8(7-*bitP)))
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	}
	return c
}

func createPixel(packet []byte, byteP, bitP *int, bpP uint) (c color.Color) {
	var r, g, b uint8

	if bpP == 1 {
		if (packet[*byteP] & (1 << uint8(7-*bitP))) == 0 {
			c = color.NRGBA{R: 0,
				G: 0,
				B: 0,
				A: 255}
		} else {
			c = color.NRGBA{R: 255,
				G: 255,
				B: 255,
				A: 255}
		}
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	} else {
		r = getBitsFromPacket(packet, byteP, bitP, bpP)
		g = getBitsFromPacket(packet, byteP, bitP, bpP)
		b = getBitsFromPacket(packet, byteP, bitP, bpP)

		c = color.NRGBA{R: r,
			G: g,
			B: b,
			A: 255}
	}
	return
}

func createTerminalVisualization(data []Data, bitsPerPixel uint) {
	var bitPos int
	var bytePos int
	var packetLen int

	for i := range data {
		packetLen = len(data[i].payload)
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[i].payload, &bytePos, &bitPos, bitsPerPixel)
			r, g, b, _ := c.RGBA()
			fmt.Printf("\x1B[0m\x1B[38;2;%d;%d;%dm\u2588", uint8(r), uint8(g), uint8(b))
			if bytePos >= packetLen {
				break
			}
		}
		fmt.Printf("\x1B[m\n")

	}

}
func createTimeVisualization(data []Data, xMax int, prefix string, ts uint, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int
	var firstPkg time.Time

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, int(ts)))

	for pkg := range data {
		if firstPkg.IsZero() {
			firstPkg = time.Unix(0, data[pkg].toa*int64(time.Microsecond))
		}
		packetLen = len(data[pkg].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[pkg].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, int(data[pkg].toa%int64(ts)), c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}
	}

	filename := prefix
	filename += "-"
	filename += firstPkg.Format(time.RFC3339Nano)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func createFixedVisualization(data []Data, xMax int, prefix string, num int, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, len(data)))

	for yPos := range data {
		packetLen = len(data[yPos].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[yPos].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, yPos, c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}

	}

	filename := prefix
	filename += strconv.Itoa(num)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		log.Fatal(err)
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	return
}

func handlePackets(ps *gopacket.PacketSource, num uint, ch chan<- Data, done <-chan bool) {
	var count uint
	for packet := range ps.Packets() {
		var k Data
		select {
		case <-done:
			close(ch)
			return
		default:
		}
		count++
		if num != 0 && count > num {
			break
		}

		elements := packet.Data()
		if len(elements) == 0 {
			continue
		}
		k = Data{toa: (packet.Metadata().CaptureInfo.Timestamp.UnixNano() / int64(time.Microsecond)), payload: packet.Data()}
		ch <- k
	}
	close(ch)
	return
}

func availableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	for _, device := range devices {
		if len(device.Addresses) == 0 {
			continue
		}
		fmt.Println("Interface: ", device.Name)
		for _, address := range device.Addresses {
			fmt.Println("   IP address:  ", address.IP)
			fmt.Println("   Subnet mask: ", address.Netmask)
		}
		fmt.Println("")
	}
}

func initSource(dev, file *string) (handle *pcap.Handle, err error) {
	if len(*dev) > 0 {
		handle, err = pcap.OpenLive(*dev, 4096, true, pcap.BlockForever)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else if len(*file) > 0 {
		handle, err = pcap.OpenOffline(*file)
		if err != nil {
			log.Fatal(err)
			os.Exit(1)
		}
	} else {
		return nil, fmt.Errorf("Source is missing\n")
	}
	return
}

func checkConfig(cfg configs) error {
	if cfg.bpP%3 != 0 && cfg.bpP != 1 {
		return fmt.Errorf("%d must be divisible by three or should be one", cfg.bpP)
	} else if cfg.bpP > 25 {
		return fmt.Errorf("%d must be smaller than 25", cfg.bpP)
	}

	if cfg.ts > 0 {
		cfg.stil |= TIMESLIZES
	}

	if cfg.stil == (TIMESLIZES | TERMINAL) {
		return fmt.Errorf("-timeslize and -terminal can't be combined")
	} else if cfg.stil == 0 {
		// If way of stil is provided, we will stick to the default one
		cfg.stil |= SOLDER
	}
	return nil
}

func init() {
	if len(os.Args) < 2 {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		os.Exit(1)
	}
}

func main() {
	var err error
	var handle *pcap.Handle
	var data []Data
	var xMax int
	var index int = 1
	osSig := make(chan os.Signal, 1)
	done := make(chan bool, 1)
	signal.Notify(osSig, os.Interrupt)
	var slicer int64
	var cfg configs
	ch := make(chan Data)

	go func() {
		<-osSig     // Blocking till interrupt signal is received
		osSig = nil // ignore further signals
		done <- true
	}()

	dev := flag.String("interface", "", "Choose an interface for online processing")
	file := flag.String("file", "", "Choose a file for offline processing")
	filter := flag.String("filter", "", "Set a specific filter")
	lst := flag.Bool("list_interfaces", false, "List available interfaces")
	vers := flag.Bool("version", false, "Show version")
	help := flag.Bool("help", false, "Show this help")
	terminalOut := flag.Bool("terminal", false, "Visualize on terminal")
	num := flag.Uint("count", 25, "Number of packets to process.\n\tIf argument is 0 the limit is removed")
	output := flag.String("prefix", "image", "Prefix of the resulting image")
	size := flag.Uint("size", 25, "Number of packets per image")
	bits := flag.Uint("bits", 24, "Number of bits per pixel.\n\tIt must be divisible by three and smaller than 25\n\tTo get black/white results, choose 1 as input.")
	ts := flag.Uint("timeslize", 0, "Number of microseconds per resulting image.\n\tSo each pixel of the height of the resulting image represents one microsecond")
	flag.Parse()

	if *lst {
		availableInterfaces()
		return
	}

	if *vers {
		fmt.Println("Version:", Version)
		return
	}

	if *help {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	cfg.bpP = *bits
	cfg.ppI = *size
	cfg.ts = *ts
	cfg.limit = *num
	cfg.stil = 0

	if *terminalOut == true {
		cfg.stil |= TERMINAL
	}
	if *ts != 0 {
		cfg.stil |= TIMESLIZES
	}

	if err = checkConfig(cfg); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	handle, err = initSource(dev, file)
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}
	defer handle.Close()

	if len(*filter) != 0 {
		err = handle.SetBPFFilter(*filter)
		if err != nil {
			log.Fatal(err, "\tInvalid filter: ", *filter)
			os.Exit(1)
		}
	}

	packetSource := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	packetSource.DecodeOptions = gopacket.Lazy

	go handlePackets(packetSource, cfg.limit, ch, done)

	switch cfg.stil {
	case SOLDER:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
			if len(data) >= int(*size) {
				xMax++
				createFixedVisualization(data, xMax, *output, index, cfg.bpP)
				xMax = 0
				index++
				data = data[:0]
			}
		}
	case TERMINAL:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			createTerminalVisualization(data, cfg.bpP)
			data = data[:0]
		}
	case TIMESLIZES:
		for i, ok := <-ch; ok; i, ok = <-ch {
			if slicer == 0 {
				slicer = i.toa + int64(*ts)
			}
			if slicer < i.toa {
				xMax++
				createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
				xMax = 0
				data = data[:0]
				slicer = i.toa + int64(*ts)
			}
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
		}
	}

	if len(data) > 0 {
		xMax++
		switch cfg.stil {
		case SOLDER:
			createFixedVisualization(data, xMax, *output, index, cfg.bpP)
		case TERMINAL:
			createTerminalVisualization(data, cfg.bpP)
		case TIMESLIZES:
			createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
		}
	}

}
<|endoftext|>