<commit_before>package gunfish

import (
	"encoding/json"
	"fmt"
	"math"
	"net"
	"net/http"
	"os"
	"os/signal"
	"reflect"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fukata/golang-stats-api-handler"
	"github.com/kayac/Gunfish/apns"
	"github.com/kayac/Gunfish/gcm"
	"github.com/lestrrat/go-server-starter/listener"
	"github.com/shogo82148/go-gracedown"
	"golang.org/x/net/netutil"
)

// Provider defines Gunfish httpHandler and has a state
// of queue which is shared by the supervisor.
type Provider struct {
	sup Supervisor
}

// ResponseHandler provides you to implement handling on success or on error response from apns.
// Therefore, you can specifies hook command which is set at toml file.
type ResponseHandler interface {
	OnResponse(Request, Response, error)
	HookCmd() string
}

// DefaultResponseHandler is the default ResponseHandler if not specified.
type DefaultResponseHandler struct {
	hook string
}

// OnResponse is performed when to receive response from APNS.
func (rh DefaultResponseHandler) OnResponse(req Request, res Response, err error) {
}

// HookCmd returns hook command to execute after getting response from APNS
// only when to get error response.
func (rh DefaultResponseHandler) HookCmd() string {
	return rh.hook
}

// StartServer starts an apns provider server on http.
func StartServer(conf Config, env Environment) {
	// Initialize DefaultResponseHandler if response handlers are not defined.
	if successResponseHandler == nil {
		InitSuccessResponseHandler(DefaultResponseHandler{})
	}

	if errorResponseHandler == nil {
		InitErrorResponseHandler(DefaultResponseHandler{hook: conf.Apns.ErrorHook})
	}

	// Init Provider
	srvStats = NewStats()
	prov := &Provider{}

	srvStats.DebugPort = conf.Provider.DebugPort
	LogWithFields(logrus.Fields{
		"type": "provider",
	}).Infof("Size of POST request queue is %d", conf.Provider.QueueSize)

	// Set APNS host addr according of environment
	if env == Production {
		conf.Apns.Host = ProdServer
	} else if env == Development {
		conf.Apns.Host = DevServer
	} else if env == Test {
		conf.Apns.Host = MockServer
	}

	// start supervisor
	sup, err := StartSupervisor(&conf)
	if err != nil {
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Fatalf("Failed to start Gunfish: %s", err.Error())
	}
	prov.sup = sup

	LogWithFields(logrus.Fields{
		"type": "supervisor",
	}).Infof("Starts supervisor at %s", Production.String())

	// StartServer listener
	listeners, err := listener.ListenAll()
	if err != nil {
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Warnf("%s. If you want graceful to restart Gunfish, you should use 'starter_server' (github.com/lestrrat/go-server-starter).", err)
	}

	// Start gunfish provider server
	var lis net.Listener
	if err == listener.ErrNoListeningTarget {
		// Fallback if not running under ServerStarter
		service := fmt.Sprintf(":%d", conf.Provider.Port)
		lis, err = net.Listen("tcp", service)
		if err != nil {
			LogWithFields(logrus.Fields{
				"type": "provider",
			}).Error(err)
			sup.Shutdown()
			return
		}
	} else {
		if l, ok := listeners[0].Addr().(*net.TCPAddr); ok && l.Port != conf.Provider.Port {
			LogWithFields(logrus.Fields{
				"type": "provider",
			}).Infof("'start_server' starts on :%d", l.Port)
		}
		// Starts Gunfish under ServerStarter.
		conf.Provider.Port = listeners[0].Addr().(*net.TCPAddr).Port
		lis = listeners[0]
	}

	// If many connections establishs between Gunfish provider and your application,
	// Gunfish provider would be overload, and decrease performance.
	llis := netutil.LimitListener(lis, conf.Provider.MaxConnections)

	// signal handling
	go startSignalReciever(lis)

	// Start Gunfish provider
	LogWithFields(logrus.Fields{
		"type": "provider",
	}).Infof("Starts provider on :%d ...", conf.Provider.Port)

	mux := http.NewServeMux()
	mux.HandleFunc("/push/apns", prov.pushAPNsHandler())
	mux.HandleFunc("/push/gcm", prov.pushGCMHandler())
	mux.HandleFunc("/stats/app", prov.statsHandler())
	mux.HandleFunc("/stats/profile", stats_api.Handler)

	if err := gracedown.Serve(llis, mux); err != nil {
		LogWithFields(logrus.Fields{}).Error(err)
	}

	// if Gunfish server stop, Close queue
	LogWithFields(logrus.Fields{
		"type": "provider",
	}).Info("Stopping server")

	// if Gunfish server stop, Close queue
	sup.Shutdown()
}

func (prov *Provider) pushAPNsHandler() http.HandlerFunc {
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		atomic.AddInt64(&(srvStats.RequestCount), 1)

		// Method Not Alllowed
		if err := validateMethod(res, req); err != nil {
			logrus.Warn(err)
			return
		}

		// Parse request body
		c := req.Header.Get("Content-Type")
		var ps []PostedData
		switch c {
		case ApplicationXW3FormURLEncoded:
			body := req.FormValue("json")
			if err := json.Unmarshal([]byte(body), &ps); err != nil {
				LogWithFields(logrus.Fields{}).Warnf("%s: %s", err, body)
				res.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(res, `{"reason": "%s"}`, err.Error())
				return
			}
		case ApplicationJSON:
			decoder := json.NewDecoder(req.Body)
			if err := decoder.Decode(&ps); err != nil {
				LogWithFields(logrus.Fields{}).Warnf("%s: %v", err, ps)
				res.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(res, `{"reason": "%s"}`, err.Error())
				return
			}
		default:
			// Unsupported Media Type
			logrus.Warnf("Unsupported Media Type: %s", c)
			res.WriteHeader(http.StatusUnsupportedMediaType)
			fmt.Fprintf(res, `{"reason":"Unsupported Media Type"}`)
			return
		}

		// Validates posted data
		if err := validatePostedData(ps); err != nil {
			res.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(res, `{"reason":"%s"}`, err.Error())
			return
		}

		// Create requests
		reqs := make([]Request, len(ps))
		for i, p := range ps {
			switch t := p.Payload.Alert.(type) {
			case map[string]interface{}:
				var alert apns.Alert
				mapToAlert(t, &alert)
				p.Payload.Alert = alert
			}

			req := Request{
				Notification: apns.Notification{
					Header:  p.Header,
					Token:   p.Token,
					Payload: p.Payload,
				},
				Tries: 0,
			}

			reqs[i] = req
		}

		// enqueues one request into supervisor's queue.
		if err := prov.sup.EnqueueClientRequest(&reqs); err != nil {
			setRetryAfter(res, req, err.Error())
			return
		}

		// success
		res.WriteHeader(http.StatusOK)
		fmt.Fprint(res, "{\"result\": \"ok\"}")
	})
}

func (prov *Provider) pushGCMHandler() http.HandlerFunc {
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		atomic.AddInt64(&(srvStats.RequestCount), 1)

		// Method Not Alllowed
		if err := validateMethod(res, req); err != nil {
			logrus.Warn(err)
			return
		}

		// only Content-Type application/json
		c := req.Header.Get("Content-Type")
		if c != ApplicationJSON {
			// Unsupported Media Type
			logrus.Warnf("Unsupported Media Type: %s", c)
			res.WriteHeader(http.StatusUnsupportedMediaType)
			fmt.Fprintf(res, `{"reason":"Unsupported Media Type"}`)
			return
		}

		// create request for gcm
		var payload gcm.Payload
		dec := json.NewDecoder(req.Body)
		if err := dec.Decode(&payload); err != nil {
			logrus.Warnf("Internal Server Error: %s", err)
			res.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(res, "{\"reason\":\"%s\"}", err.Error())
		}
		grs := []Request{
			Request{
				Notification: payload,
				Tries:        0,
			},
		}

		// enqueues one request into supervisor's queue.
		if err := prov.sup.EnqueueClientRequest(&grs); err != nil {
			setRetryAfter(res, req, err.Error())
			return
		}

		// success
		res.WriteHeader(http.StatusOK)
		fmt.Fprint(res, "{\"result\": \"ok\"}")
	})
}

func validateMethod(res http.ResponseWriter, req *http.Request) error {
	if req.Method != "POST" {
		res.WriteHeader(http.StatusMethodNotAllowed)
		fmt.Fprintf(res, "{\"reason\":\"Method Not Allowed.\"}")
		return fmt.Errorf("Method Not Allowed: %s", req.Method)
	}
	return nil
}

func setRetryAfter(res http.ResponseWriter, req *http.Request, reason string) {
	atomic.StoreInt64(&(srvStats.ServiceUnavailableAt), time.Now().Unix())
	// Retry-After is set seconds
	res.Header().Set("Retry-After", fmt.Sprintf("%d", srvStats.RetryAfter))
	res.WriteHeader(http.StatusServiceUnavailable)
	fmt.Fprintf(res, fmt.Sprintf(`{"reason":"%s"}`, reason))
}

func (prov *Provider) statsHandler() http.HandlerFunc {
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if ok := validateStatsHandler(res, req); ok != true {
			return
		}

		wqs := 0
		for _, w := range prov.sup.workers {
			wqs += len(w.queue)
		}

		atomic.StoreInt64(&(srvStats.QueueSize), int64(len(prov.sup.queue)))
		atomic.StoreInt64(&(srvStats.RetryQueueSize), int64(len(prov.sup.retryq)))
		atomic.StoreInt64(&(srvStats.WorkersQueueSize), int64(wqs))
		atomic.StoreInt64(&(srvStats.CommandQueueSize), int64(len(prov.sup.cmdq)))
		res.WriteHeader(http.StatusOK)
		encoder := json.NewEncoder(res)
		err := encoder.Encode(srvStats.GetStats())
		if err != nil {
			res.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(res, `{"reason":"Internal Server Error"}`)
			return
		}
	})
}

func validatePostedData(ps []PostedData) error {
	if len(ps) == 0 {
		return fmt.Errorf("PostedData must not be empty: %v", ps)
	}

	if len(ps) > MaxRequestSize {
		return fmt.Errorf("PostedData was too long. Be less than %d: %v", MaxRequestSize, len(ps))
	}

	for _, p := range ps {
		if p.Payload.APS == nil || p.Token == "" {
			return fmt.Errorf("Payload format was malformed: %v", p.Payload)
		}
	}
	return nil
}

func validateStatsHandler(res http.ResponseWriter, req *http.Request) bool {
	// Method Not Alllowed
	if req.Method != "GET" {
		res.WriteHeader(http.StatusMethodNotAllowed)
		fmt.Fprintf(res, `{"reason":"Method Not Allowed."}`)
		logrus.Warnf("Method Not Allowed: %s", req.Method)
		return false
	}

	return true
}

func mapToAlert(mapVal map[string]interface{}, alert *apns.Alert) {
	a := reflect.ValueOf(alert).Elem()
	for k, v := range mapVal {
		newk, ok := AlertKeyToField[k]
		if ok == true {
			a.FieldByName(newk).Set(reflect.ValueOf(v))
		} else {
			logrus.Warnf("\"%s\" is not supported key for Alert struct.", k)
		}
	}
}

func startSignalReciever(lis net.Listener) {
	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, syscall.SIGTERM)
	s := <-sigChan
	switch s {
	case syscall.SIGHUP:
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Info("Gunfish recieved SIGHUP signal.")
		gracedown.Close()
	case syscall.SIGTERM:
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Info("Gunfish recieved SIGTERM signal.")
		gracedown.Close()
	case syscall.SIGINT:
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Info("Gunfish recieved SIGINT signal. Stopping server now...")
		gracedown.Close()
	}
}

func updateRetryAfterStat(x int64) {
	var nxtRA int64
	if x > int64(ResetRetryAfterSecond/time.Second) {
		nxtRA = int64(RetryAfterSecond / time.Second)
	} else {
		a := int64(math.Log(float64(10/(x+1) + 1)))
		if srvStats.RetryAfter+2*a < int64(ResetRetryAfterSecond/time.Second) {
			nxtRA = srvStats.RetryAfter + 2*a
		} else {
			nxtRA = int64(ResetRetryAfterSecond / time.Second)
		}
	}

	atomic.StoreInt64(&(srvStats.RetryAfter), nxtRA)
}
<commit_msg>call updateRetryAfterStat() at setRetryAfter()<commit_after>package gunfish

import (
	"encoding/json"
	"fmt"
	"math"
	"net"
	"net/http"
	"os"
	"os/signal"
	"reflect"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fukata/golang-stats-api-handler"
	"github.com/kayac/Gunfish/apns"
	"github.com/kayac/Gunfish/gcm"
	"github.com/lestrrat/go-server-starter/listener"
	"github.com/shogo82148/go-gracedown"
	"golang.org/x/net/netutil"
)

// Provider defines Gunfish httpHandler and has a state
// of queue which is shared by the supervisor.
type Provider struct {
	sup Supervisor
}

// ResponseHandler provides you to implement handling on success or on error response from apns.
// Therefore, you can specifies hook command which is set at toml file.
type ResponseHandler interface {
	OnResponse(Request, Response, error)
	HookCmd() string
}

// DefaultResponseHandler is the default ResponseHandler if not specified.
type DefaultResponseHandler struct {
	hook string
}

// OnResponse is performed when to receive response from APNS.
func (rh DefaultResponseHandler) OnResponse(req Request, res Response, err error) {
}

// HookCmd returns hook command to execute after getting response from APNS
// only when to get error response.
func (rh DefaultResponseHandler) HookCmd() string {
	return rh.hook
}

// StartServer starts an apns provider server on http.
func StartServer(conf Config, env Environment) {
	// Initialize DefaultResponseHandler if response handlers are not defined.
	if successResponseHandler == nil {
		InitSuccessResponseHandler(DefaultResponseHandler{})
	}

	if errorResponseHandler == nil {
		InitErrorResponseHandler(DefaultResponseHandler{hook: conf.Apns.ErrorHook})
	}

	// Init Provider
	srvStats = NewStats()
	prov := &Provider{}

	srvStats.DebugPort = conf.Provider.DebugPort
	LogWithFields(logrus.Fields{
		"type": "provider",
	}).Infof("Size of POST request queue is %d", conf.Provider.QueueSize)

	// Set APNS host addr according of environment
	if env == Production {
		conf.Apns.Host = ProdServer
	} else if env == Development {
		conf.Apns.Host = DevServer
	} else if env == Test {
		conf.Apns.Host = MockServer
	}

	// start supervisor
	sup, err := StartSupervisor(&conf)
	if err != nil {
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Fatalf("Failed to start Gunfish: %s", err.Error())
	}
	prov.sup = sup

	LogWithFields(logrus.Fields{
		"type": "supervisor",
	}).Infof("Starts supervisor at %s", Production.String())

	// StartServer listener
	listeners, err := listener.ListenAll()
	if err != nil {
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Warnf("%s. If you want graceful to restart Gunfish, you should use 'starter_server' (github.com/lestrrat/go-server-starter).", err)
	}

	// Start gunfish provider server
	var lis net.Listener
	if err == listener.ErrNoListeningTarget {
		// Fallback if not running under ServerStarter
		service := fmt.Sprintf(":%d", conf.Provider.Port)
		lis, err = net.Listen("tcp", service)
		if err != nil {
			LogWithFields(logrus.Fields{
				"type": "provider",
			}).Error(err)
			sup.Shutdown()
			return
		}
	} else {
		if l, ok := listeners[0].Addr().(*net.TCPAddr); ok && l.Port != conf.Provider.Port {
			LogWithFields(logrus.Fields{
				"type": "provider",
			}).Infof("'start_server' starts on :%d", l.Port)
		}
		// Starts Gunfish under ServerStarter.
		conf.Provider.Port = listeners[0].Addr().(*net.TCPAddr).Port
		lis = listeners[0]
	}

	// If many connections establishs between Gunfish provider and your application,
	// Gunfish provider would be overload, and decrease performance.
	llis := netutil.LimitListener(lis, conf.Provider.MaxConnections)

	// signal handling
	go startSignalReciever(lis)

	// Start Gunfish provider
	LogWithFields(logrus.Fields{
		"type": "provider",
	}).Infof("Starts provider on :%d ...", conf.Provider.Port)

	mux := http.NewServeMux()
	mux.HandleFunc("/push/apns", prov.pushAPNsHandler())
	mux.HandleFunc("/push/gcm", prov.pushGCMHandler())
	mux.HandleFunc("/stats/app", prov.statsHandler())
	mux.HandleFunc("/stats/profile", stats_api.Handler)

	if err := gracedown.Serve(llis, mux); err != nil {
		LogWithFields(logrus.Fields{}).Error(err)
	}

	// if Gunfish server stop, Close queue
	LogWithFields(logrus.Fields{
		"type": "provider",
	}).Info("Stopping server")

	// if Gunfish server stop, Close queue
	sup.Shutdown()
}

func (prov *Provider) pushAPNsHandler() http.HandlerFunc {
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		atomic.AddInt64(&(srvStats.RequestCount), 1)

		// Method Not Alllowed
		if err := validateMethod(res, req); err != nil {
			logrus.Warn(err)
			return
		}

		// Parse request body
		c := req.Header.Get("Content-Type")
		var ps []PostedData
		switch c {
		case ApplicationXW3FormURLEncoded:
			body := req.FormValue("json")
			if err := json.Unmarshal([]byte(body), &ps); err != nil {
				LogWithFields(logrus.Fields{}).Warnf("%s: %s", err, body)
				res.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(res, `{"reason": "%s"}`, err.Error())
				return
			}
		case ApplicationJSON:
			decoder := json.NewDecoder(req.Body)
			if err := decoder.Decode(&ps); err != nil {
				LogWithFields(logrus.Fields{}).Warnf("%s: %v", err, ps)
				res.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(res, `{"reason": "%s"}`, err.Error())
				return
			}
		default:
			// Unsupported Media Type
			logrus.Warnf("Unsupported Media Type: %s", c)
			res.WriteHeader(http.StatusUnsupportedMediaType)
			fmt.Fprintf(res, `{"reason":"Unsupported Media Type"}`)
			return
		}

		// Validates posted data
		if err := validatePostedData(ps); err != nil {
			res.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(res, `{"reason":"%s"}`, err.Error())
			return
		}

		// Create requests
		reqs := make([]Request, len(ps))
		for i, p := range ps {
			switch t := p.Payload.Alert.(type) {
			case map[string]interface{}:
				var alert apns.Alert
				mapToAlert(t, &alert)
				p.Payload.Alert = alert
			}

			req := Request{
				Notification: apns.Notification{
					Header:  p.Header,
					Token:   p.Token,
					Payload: p.Payload,
				},
				Tries: 0,
			}

			reqs[i] = req
		}

		// enqueues one request into supervisor's queue.
		if err := prov.sup.EnqueueClientRequest(&reqs); err != nil {
			setRetryAfter(res, req, err.Error())
			return
		}

		// success
		res.WriteHeader(http.StatusOK)
		fmt.Fprint(res, "{\"result\": \"ok\"}")
	})
}

func (prov *Provider) pushGCMHandler() http.HandlerFunc {
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		atomic.AddInt64(&(srvStats.RequestCount), 1)

		// Method Not Alllowed
		if err := validateMethod(res, req); err != nil {
			logrus.Warn(err)
			return
		}

		// only Content-Type application/json
		c := req.Header.Get("Content-Type")
		if c != ApplicationJSON {
			// Unsupported Media Type
			logrus.Warnf("Unsupported Media Type: %s", c)
			res.WriteHeader(http.StatusUnsupportedMediaType)
			fmt.Fprintf(res, `{"reason":"Unsupported Media Type"}`)
			return
		}

		// create request for gcm
		var payload gcm.Payload
		dec := json.NewDecoder(req.Body)
		if err := dec.Decode(&payload); err != nil {
			logrus.Warnf("Internal Server Error: %s", err)
			res.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(res, "{\"reason\":\"%s\"}", err.Error())
		}
		grs := []Request{
			Request{
				Notification: payload,
				Tries:        0,
			},
		}

		// enqueues one request into supervisor's queue.
		if err := prov.sup.EnqueueClientRequest(&grs); err != nil {
			setRetryAfter(res, req, err.Error())
			return
		}

		// success
		res.WriteHeader(http.StatusOK)
		fmt.Fprint(res, "{\"result\": \"ok\"}")
	})
}

func validateMethod(res http.ResponseWriter, req *http.Request) error {
	if req.Method != "POST" {
		res.WriteHeader(http.StatusMethodNotAllowed)
		fmt.Fprintf(res, "{\"reason\":\"Method Not Allowed.\"}")
		return fmt.Errorf("Method Not Allowed: %s", req.Method)
	}
	return nil
}

func setRetryAfter(res http.ResponseWriter, req *http.Request, reason string) {
	now := time.Now().Unix()
	atomic.StoreInt64(&(srvStats.ServiceUnavailableAt), now)
	updateRetryAfterStat(now - srvStats.ServiceUnavailableAt)
	// Retry-After is set seconds
	res.Header().Set("Retry-After", fmt.Sprintf("%d", srvStats.RetryAfter))
	res.WriteHeader(http.StatusServiceUnavailable)
	fmt.Fprintf(res, fmt.Sprintf(`{"reason":"%s"}`, reason))
}

func (prov *Provider) statsHandler() http.HandlerFunc {
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		if ok := validateStatsHandler(res, req); ok != true {
			return
		}

		wqs := 0
		for _, w := range prov.sup.workers {
			wqs += len(w.queue)
		}

		atomic.StoreInt64(&(srvStats.QueueSize), int64(len(prov.sup.queue)))
		atomic.StoreInt64(&(srvStats.RetryQueueSize), int64(len(prov.sup.retryq)))
		atomic.StoreInt64(&(srvStats.WorkersQueueSize), int64(wqs))
		atomic.StoreInt64(&(srvStats.CommandQueueSize), int64(len(prov.sup.cmdq)))
		res.WriteHeader(http.StatusOK)
		encoder := json.NewEncoder(res)
		err := encoder.Encode(srvStats.GetStats())
		if err != nil {
			res.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(res, `{"reason":"Internal Server Error"}`)
			return
		}
	})
}

func validatePostedData(ps []PostedData) error {
	if len(ps) == 0 {
		return fmt.Errorf("PostedData must not be empty: %v", ps)
	}

	if len(ps) > MaxRequestSize {
		return fmt.Errorf("PostedData was too long. Be less than %d: %v", MaxRequestSize, len(ps))
	}

	for _, p := range ps {
		if p.Payload.APS == nil || p.Token == "" {
			return fmt.Errorf("Payload format was malformed: %v", p.Payload)
		}
	}
	return nil
}

func validateStatsHandler(res http.ResponseWriter, req *http.Request) bool {
	// Method Not Alllowed
	if req.Method != "GET" {
		res.WriteHeader(http.StatusMethodNotAllowed)
		fmt.Fprintf(res, `{"reason":"Method Not Allowed."}`)
		logrus.Warnf("Method Not Allowed: %s", req.Method)
		return false
	}

	return true
}

func mapToAlert(mapVal map[string]interface{}, alert *apns.Alert) {
	a := reflect.ValueOf(alert).Elem()
	for k, v := range mapVal {
		newk, ok := AlertKeyToField[k]
		if ok == true {
			a.FieldByName(newk).Set(reflect.ValueOf(v))
		} else {
			logrus.Warnf("\"%s\" is not supported key for Alert struct.", k)
		}
	}
}

func startSignalReciever(lis net.Listener) {
	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, syscall.SIGTERM)
	s := <-sigChan
	switch s {
	case syscall.SIGHUP:
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Info("Gunfish recieved SIGHUP signal.")
		gracedown.Close()
	case syscall.SIGTERM:
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Info("Gunfish recieved SIGTERM signal.")
		gracedown.Close()
	case syscall.SIGINT:
		LogWithFields(logrus.Fields{
			"type": "provider",
		}).Info("Gunfish recieved SIGINT signal. Stopping server now...")
		gracedown.Close()
	}
}

func updateRetryAfterStat(x int64) {
	var nxtRA int64
	if x > int64(ResetRetryAfterSecond/time.Second) {
		nxtRA = int64(RetryAfterSecond / time.Second)
	} else {
		a := int64(math.Log(float64(10/(x+1) + 1)))
		if srvStats.RetryAfter+2*a < int64(ResetRetryAfterSecond/time.Second) {
			nxtRA = srvStats.RetryAfter + 2*a
		} else {
			nxtRA = int64(ResetRetryAfterSecond / time.Second)
		}
	}

	atomic.StoreInt64(&(srvStats.RetryAfter), nxtRA)
}
<|endoftext|>


<commit_before>// Package coap provides a CoAP client and server.
package coap

import (
	"log"
	"net"
	"time"
)

const maxPktLen = 1500

// Handler is a type that handles CoAP messages.
type Handler interface {
	// Handle the message and optionally return a response message.
	ServeCOAP(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message
}

type funcHandler func(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message

func (f funcHandler) ServeCOAP(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message {
	return f(l, a, m)
}

// FuncHandler builds a handler from a function.
func FuncHandler(f func(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message) Handler {
	return funcHandler(f)
}

func handlePacket(l *net.UDPConn, data []byte, u *net.UDPAddr,
	rh Handler) {

	msg, err := parseMessage(data)
	if err != nil {
		log.Printf("Error parsing %v", err)
		return
	}

	rv := rh.ServeCOAP(l, u, &msg)
	if rv != nil {
		log.Printf("Transmitting %#v", rv)
		Transmit(l, u, *rv)
	}
}

// Transmit a message.
func Transmit(l *net.UDPConn, a *net.UDPAddr, m Message) error {
	d, err := m.MarshalBinary()
	if err != nil {
		return err
	}

	if a == nil {
		_, err = l.Write(d)
	} else {
		_, err = l.WriteTo(d, a)
	}
	return err
}

// Receive a message.
func Receive(l *net.UDPConn, buf []byte) (Message, error) {
	l.SetReadDeadline(time.Now().Add(ResponseTimeout))

	nr, _, err := l.ReadFromUDP(buf)
	if err != nil {
		return Message{}, err
	}
	return parseMessage(buf[:nr])
}

// ListenAndServe binds to the given address and serve requests forever.
func ListenAndServe(n, addr string, rh Handler) error {
	uaddr, err := net.ResolveUDPAddr(n, addr)
	if err != nil {
		return err
	}

	l, err := net.ListenUDP(n, uaddr)
	if err != nil {
		return err
	}

	return Serve(l, rh)
}

// Serve processes incoming UDP packets on the given listener, and processes
// these requests forever (or until the listener is closed).
func Serve(listener *net.UDPConn, rh Handler) error {
	buf := make([]byte, maxPktLen)
	for {
		nr, addr, err := listener.ReadFromUDP(buf)
		if err != nil {
			if neterr, ok := err.(net.Error); ok && (neterr.Temporary() || neterr.Timeout()) {
				time.Sleep(5 * time.Millisecond)
				continue
			}
			return err
		}
		tmp := make([]byte, nr)
		copy(tmp, buf)
		go handlePacket(listener, tmp, addr, rh)
	}
}
<commit_msg>Remove "Transmitting ..." log print from server library code<commit_after>// Package coap provides a CoAP client and server.
package coap

import (
	"log"
	"net"
	"time"
)

const maxPktLen = 1500

// Handler is a type that handles CoAP messages.
type Handler interface {
	// Handle the message and optionally return a response message.
	ServeCOAP(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message
}

type funcHandler func(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message

func (f funcHandler) ServeCOAP(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message {
	return f(l, a, m)
}

// FuncHandler builds a handler from a function.
func FuncHandler(f func(l *net.UDPConn, a *net.UDPAddr, m *Message) *Message) Handler {
	return funcHandler(f)
}

func handlePacket(l *net.UDPConn, data []byte, u *net.UDPAddr,
	rh Handler) {

	msg, err := parseMessage(data)
	if err != nil {
		log.Printf("Error parsing %v", err)
		return
	}

	rv := rh.ServeCOAP(l, u, &msg)
	if rv != nil {
		Transmit(l, u, *rv)
	}
}

// Transmit a message.
func Transmit(l *net.UDPConn, a *net.UDPAddr, m Message) error {
	d, err := m.MarshalBinary()
	if err != nil {
		return err
	}

	if a == nil {
		_, err = l.Write(d)
	} else {
		_, err = l.WriteTo(d, a)
	}
	return err
}

// Receive a message.
func Receive(l *net.UDPConn, buf []byte) (Message, error) {
	l.SetReadDeadline(time.Now().Add(ResponseTimeout))

	nr, _, err := l.ReadFromUDP(buf)
	if err != nil {
		return Message{}, err
	}
	return parseMessage(buf[:nr])
}

// ListenAndServe binds to the given address and serve requests forever.
func ListenAndServe(n, addr string, rh Handler) error {
	uaddr, err := net.ResolveUDPAddr(n, addr)
	if err != nil {
		return err
	}

	l, err := net.ListenUDP(n, uaddr)
	if err != nil {
		return err
	}

	return Serve(l, rh)
}

// Serve processes incoming UDP packets on the given listener, and processes
// these requests forever (or until the listener is closed).
func Serve(listener *net.UDPConn, rh Handler) error {
	buf := make([]byte, maxPktLen)
	for {
		nr, addr, err := listener.ReadFromUDP(buf)
		if err != nil {
			if neterr, ok := err.(net.Error); ok && (neterr.Temporary() || neterr.Timeout()) {
				time.Sleep(5 * time.Millisecond)
				continue
			}
			return err
		}
		tmp := make([]byte, nr)
		copy(tmp, buf)
		go handlePacket(listener, tmp, addr, rh)
	}
}
<|endoftext|>


<commit_before>package main

import (
    "net/http"
    "log"

    "code.google.com/p/go.net/websocket"
)

type Server struct {
    Config  *Configuration
    Store   *Storage
}

func main() {
    conf  := initConfig()
    store := initStore(&conf)
    //initLogger()
    server := Server{&conf, &store}

    //server.initAppListner()
    go server.initSocketListener()

    http.HandleFunc("/", rootHandler)
    err := http.ListenAndServe(listenAddr, nil)
    if err != nil {
        log.Fatal(err)
    }
}

func (this *Server) initSocketListener() {
    Connect := func(ws *websocket.Conn) {
        sock := Socket{ws, "", make(chan *Message, 1000), make(chan bool), this}

        log.Printf("Connected via %s\n", ws.RemoteAddr());
        if err := Authenticate(&sock); err != nil {
            log.Printf("Error: %s\n", err.Error())
            return
        }

        go listenForMessages(&sock)
        go listenForWrites(&sock)

        <-sock.done
        sock.Close()
    }

    http.Handle("/socket", websocket.Handler(Connect))
}<commit_msg>add app listener that listens for redis messages<commit_after>package main

import (
    "net/http"
    "log"
    "encoding/json"

    "code.google.com/p/go.net/websocket"
)

type Server struct {
    Config  *Configuration
    Store   *Storage
}

func main() {
    conf  := initConfig()
    store := initStore(&conf)
    //initLogger()
    server := Server{&conf, &store}

    go server.initAppListner()
    go server.initSocketListener()

    http.HandleFunc("/", rootHandler)
    err := http.ListenAndServe(listenAddr, nil)
    if err != nil {
        log.Fatal(err)
    }
}

func (this *Server) initSocketListener() {
    Connect := func(ws *websocket.Conn) {
        sock := Socket{ws, "", make(chan *Message, 1000), make(chan bool), this}

        log.Printf("Connected via %s\n", ws.RemoteAddr());
        if err := Authenticate(&sock); err != nil {
            log.Printf("Error: %s\n", err.Error())
            return
        }

        go listenForMessages(&sock)
        go listenForWrites(&sock)

        <-sock.done
        sock.Close()
    }

    http.Handle("/socket", websocket.Handler(Connect))
}

func (this *Server) initAppListner() {

    rec := make(chan []string)

    consumer, err := this.Store.redis.Subscribe(rec, "Message")
    if err != nil {
        log.Fatal("Couldn't subscribe to redis channel")
    }
    defer consumer.Quit()

    var ms []string
    for {
        var msg Message
        log.Println("LISENING FOR REDIS MESSAGE")
        ms = <- rec
        json.Unmarshal([]byte(ms[2]), &msg)
        log.Printf("Received %v\n", msg.Name)

        go msg.FromSocket(&Socket{nil, "", nil, nil, this})
    }

}
    <|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"os"
	"time"

	"github.com/garyburd/go-oauth/oauth"
	"github.com/garyburd/redigo/redis"
	"github.com/pjebs/tokbox"
	"gopkg.in/boj/redistore.v1"
)

const (
	Development = iota
	Production
)

// Environment Variables:

// TOKBOX_KEY
// TOKBOX_SECRET_KEY
// BLARG_ENV
// TWITTER_KEY
// TWITTER_SECRET_KEY
// REDIS_PORT_6379_TCP_ADDR
// REDIS_PORT_6379_TCP_POST
// REDIS_PASSWORD

type Server struct {
	RediStore     *redistore.RediStore
	RedisPool     *redis.Pool
	OAuthClient   oauth.Client
	State         State
	ClientBaseUrl string
	ForceHTTPS    bool
	TokBox        *tokbox.Tokbox
	TokBoxKey     string
}

func newPool(addr, password string) *redis.Pool {
	/* environment := Development
	if os.Getenv("TRUCAST_ENV") == "production" {
		environment = Production
	}*/
	return &redis.Pool{
		// Maximum number of idle connections in the pool.
		MaxIdle: 100,

		// Maximum number of connections allocated by the pool at a given time.
		// When zero, there is no limit on the number of connections in the pool.
		MaxActive: 100,

		// Close connections after remaining idle for this duration. If the value
		// is zero, then idle connections are not closed. Applications should set
		// the timeout to a value less than the server's timeout.
		IdleTimeout: 240 * time.Second,

		// Dial is an application supplied function for creating and configuring a
		// connection
		Dial: func() (redis.Conn, error) {
			// c, err := redis.Dial("tcp", "localhost:6379")
			c, err := redis.DialURL(os.Getenv("REDIS_URL"))
			if err != nil {
				return nil, err
			}

			/* if environment == Production {
				c, err = redis.DialURL(os.Getenv("REDIS_URL"))
				if err != nil {
				return nil, err
				}
			}*/

			if password != "" {
				if _, err := c.Do("AUTH", password); err != nil {
					c.Close()
					return nil, err
				}
			}

			return c, nil
		},
	}
}

func NewServer() Server {
	// Get environment
	environment := Development
	if os.Getenv("TRUCAST_ENV") == "production" {
		environment = Production
	}

	// Redis
	redisAddr := os.Getenv("REDIS_URL")
	/* redisAddr := os.Getenv("REDIS_PORT_6379_TCP_ADDR") +
		":" + os.Getenv("REDIS_PORT_6379_TCP_PORT")
	if environment == Production {
		redisAddr = os.Getenv("REDIS_URL")
	}*/
	redisPool := newPool(redisAddr, os.Getenv("REDIS_PASSWORD"))

	// Redis store
	rediStore, err := redistore.NewRediStoreWithPool(redisPool, []byte("secret string "))
	if err != nil {
		panic(err)
	}

	// TokBox
	tokBoxKey := os.Getenv("TOKBOX_KEY")
	tb := tokbox.New(
		tokBoxKey,
		os.Getenv("TOKBOX_SECRET_KEY"),
	)

	// OAuth
	gob.Register(&oauth.Credentials{})
	oauthClient := oauth.Client{
		TemporaryCredentialRequestURI: "https://api.twitter.com/oauth/request_token",
		ResourceOwnerAuthorizationURI: "https://api.twitter.com/oauth/authenticate",
		TokenRequestURI:               "https://api.twitter.com/oauth/access_token",
	}
	oauthClient.Credentials.Token = os.Getenv("TWITTER_KEY")
	oauthClient.Credentials.Secret = os.Getenv("TWITTER_SECRET_KEY")

	// State
	state := NewLocalState()

	// Config
	clientBaseUrl := "http://localhost:8000"
	forceHTTPS := false
	if environment == Production {
		clientBaseUrl = "https://peer-cast.herokuapp.com"
		forceHTTPS = true
	}

	return Server{
		RedisPool:     redisPool,
		RediStore:     rediStore,
		OAuthClient:   oauthClient,
		State:         &state,
		ClientBaseUrl: clientBaseUrl,
		ForceHTTPS:    forceHTTPS,
		TokBox:        tb,
		TokBoxKey:     tokBoxKey,
	}
}
<commit_msg>local to prod<commit_after>package main

import (
	"encoding/gob"
	"os"
	"time"

	"github.com/garyburd/go-oauth/oauth"
	"github.com/garyburd/redigo/redis"
	"github.com/pjebs/tokbox"
	"gopkg.in/boj/redistore.v1"
)

const (
	Development = iota
	Production
)

// Environment Variables:

// TOKBOX_KEY
// TOKBOX_SECRET_KEY
// BLARG_ENV
// TWITTER_KEY
// TWITTER_SECRET_KEY
// REDIS_PORT_6379_TCP_ADDR
// REDIS_PORT_6379_TCP_POST
// REDIS_PASSWORD

type Server struct {
	RediStore     *redistore.RediStore
	RedisPool     *redis.Pool
	OAuthClient   oauth.Client
	State         State
	ClientBaseUrl string
	ForceHTTPS    bool
	TokBox        *tokbox.Tokbox
	TokBoxKey     string
}

func newPool(addr, password string) *redis.Pool {
	/* environment := Development
	if os.Getenv("TRUCAST_ENV") == "production" {
		environment = Production
	}*/
	return &redis.Pool{
		// Maximum number of idle connections in the pool.
		MaxIdle: 100,

		// Maximum number of connections allocated by the pool at a given time.
		// When zero, there is no limit on the number of connections in the pool.
		MaxActive: 100,

		// Close connections after remaining idle for this duration. If the value
		// is zero, then idle connections are not closed. Applications should set
		// the timeout to a value less than the server's timeout.
		IdleTimeout: 240 * time.Second,

		// Dial is an application supplied function for creating and configuring a
		// connection
		Dial: func() (redis.Conn, error) {
			// c, err := redis.Dial("tcp", "localhost:6379")
			c, err := redis.DialURL(os.Getenv("REDIS_URL"))
			if err != nil {
				return nil, err
			}

			/* if environment == Production {
				c, err = redis.DialURL(os.Getenv("REDIS_URL"))
				if err != nil {
				return nil, err
				}
			}*/

			if password != "" {
				if _, err := c.Do("AUTH", password); err != nil {
					c.Close()
					return nil, err
				}
			}

			return c, nil
		},
	}
}

func NewServer() Server {
	// Get environment
	environment := Development
	if os.Getenv("TRUCAST_ENV") == "production" {
		environment = Production
	}

	// Redis
	redisAddr := os.Getenv("REDIS_URL")
	/* redisAddr := os.Getenv("REDIS_PORT_6379_TCP_ADDR") +
		":" + os.Getenv("REDIS_PORT_6379_TCP_PORT")
	if environment == Production {
		redisAddr = os.Getenv("REDIS_URL")
	}*/
	redisPool := newPool(redisAddr, os.Getenv("REDIS_PASSWORD"))

	// Redis store
	rediStore, err := redistore.NewRediStoreWithPool(redisPool, []byte("secret string "))
	if err != nil {
		panic(err)
	}

	// TokBox
	tokBoxKey := os.Getenv("TOKBOX_KEY")
	tb := tokbox.New(
		tokBoxKey,
		os.Getenv("TOKBOX_SECRET_KEY"),
	)

	// OAuth
	gob.Register(&oauth.Credentials{})
	oauthClient := oauth.Client{
		TemporaryCredentialRequestURI: "https://api.twitter.com/oauth/request_token",
		ResourceOwnerAuthorizationURI: "https://api.twitter.com/oauth/authenticate",
		TokenRequestURI:               "https://api.twitter.com/oauth/access_token",
	}
	oauthClient.Credentials.Token = os.Getenv("TWITTER_KEY")
	oauthClient.Credentials.Secret = os.Getenv("TWITTER_SECRET_KEY")

	// State
	state := NewLocalState()

	// Config
	clientBaseUrl := "https://peer-cast.herokuapp.com"
	forceHTTPS := false
	if environment == Production {
		clientBaseUrl = "https://peer-cast.herokuapp.com"
		forceHTTPS = true
	}

	return Server{
		RedisPool:     redisPool,
		RediStore:     rediStore,
		OAuthClient:   oauthClient,
		State:         &state,
		ClientBaseUrl: clientBaseUrl,
		ForceHTTPS:    forceHTTPS,
		TokBox:        tb,
		TokBoxKey:     tokBoxKey,
	}
}
<|endoftext|>


<commit_before>package pkg

// the line directive should not affect the line ignores

//line random-file:1
func fn1() {} // MATCH "test problem"

//lint:ignore TEST1000 This should be ignored, because ...
func fn2() {}

//lint:ignore TEST1000
func fn3() {} // MATCH "test problem"

//lint:ignore TEST1000 ignore
func fn4() {
	//lint:ignore TEST1000 ignore
	var _ int
}

// MATCH:11 "malformed linter directive"
// MATCH:16 "this linter directive didn't match anything"
<commit_msg>lint: add test case for multiple linter directives<commit_after>package pkg

// the line directive should not affect the line ignores

//line random-file:1
func fn1() {} // MATCH "test problem"

//lint:ignore TEST1000 This should be ignored, because ...
//lint:ignore XXX1000 Testing that multiple linter directives work correctly
func fn2() {}

//lint:ignore TEST1000
func fn3() {} // MATCH "test problem"

//lint:ignore TEST1000 ignore
func fn4() {
	//lint:ignore TEST1000 ignore
	var _ int
}

// MATCH:12 "malformed linter directive"
// MATCH:17 "this linter directive didn't match anything"
<|endoftext|>


<commit_before>// Package exec is the entry point for security automation Cloud Functions.
package exec

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

import (
	"context"
	"log"

	"cloud.google.com/go/pubsub"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/cloud-sql/removepublic"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/cloud-sql/requiressl"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/disabledashboard"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gce/createsnapshot"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gce/openfirewall"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gce/removepublicip"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gcs/closebucket"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gke/disabledashboard"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/iam/removenonorgmembers"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/iam/revokeiam"
	"github.com/googlecloudplatform/threat-automation/entities"
)

var ent *entities.Entity

func init() {
	ctx := context.Background()
	var err error
	ent, err = entities.New(ctx)
	if err != nil {
		log.Fatalf("failed to initialize entities: %q", err)
	}
}

// RevokeIAM is the entry point for the IAM revoker Cloud Function.
//
// This Cloud Function will be triggered when Event Threat Detection
// detects an anomalous IAM grant. Once triggered this function will
// attempt to revoke the external members added to the policy if they match the provided
// list of disallowed domains. Additionally this method will only remove members if the
// project they were added to is within the specified folders. This configuration allows
// you to take a remediation action only on specific members and folders. For example,
// you may have a folder "development" where users can experiment without strict policies.
// However in your "production" folder you may want to revoke any grants that ETD finds as
// long as they match the domains you specify.
//
// Permissions required
// 	- roles/resourcemanager.folderAdmin to revoke IAM grants.
//	- roles/viewer to verify the affected project is within the enforced folder.
//
func RevokeIAM(ctx context.Context, m pubsub.Message) error {
	r, err := revokeiam.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return revokeiam.Execute(ctx, r, ent)
}

// SnapshotDisk is the entry point for the auto creation of GCE snapshots Cloud Function.
//
// This Cloud Function will respond to Event Threat Detection **bad IP** findings. Once a bad IP
// finding is received this Cloud Function will look for any existing disk snapshots for the
// affected instance. If there are recent snapshots then no action is taken. If we have not
// taken a snapshot recently, take a new snapshot for each disk within the instance.
//
// Permissions required
//	- roles/compute.instanceAdmin.v1 to manage disk snapshots.
//
func SnapshotDisk(ctx context.Context, m pubsub.Message) error {
	r, err := createsnapshot.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return createsnapshot.Execute(ctx, r, ent)
}

// CloseBucket will remove any public users from buckets found within the provided folders.
//
// Permissions required
//	- roles/viewer to retrieve ancestry.
//	- roles/storeage.admin to modify buckets.
//
func CloseBucket(ctx context.Context, m pubsub.Message) error {
	r, err := closebucket.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return closebucket.Execute(ctx, r, ent)
}

// OpenFirewall will remediate an open firewall.
//
// Permissions required
//	- roles/viewer to retrieve ancestry.
//	- roles/compute.securityAdmin to modify firewall rules.
//
func OpenFirewall(ctx context.Context, m pubsub.Message) error {
	r, err := openfirewall.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return openfirewall.Execute(ctx, r, ent)
}

// RemoveNonOrganizationMember removes all members that do not match the organization domain.
//
// This Cloud Function will respond to Security Health Analytics **NON_ORG_IAM_MEMBER** findings from **IAM Scanner**.
// All user member types (user:) that do not correspond to the organization will be removed from policy binding.
//
// Permissions required
//	- roles/resourcemanager.organizationAdmin to get org info and policies and set policies.
//
func RemoveNonOrganizationMember(ctx context.Context, m pubsub.Message) error {
	r, err := removenonorgmembers.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return removenonorgmembers.Execute(ctx, r, ent)
}

// RemovePublicIP removes all the external IP addresses of a GCE instance.
//
// This Cloud Function will respond to Security Health Analytics **Public IP Address** findings
// from **Compute Instance Scanner**. All public IP addresses of the affected instance will be
// deleted when this function is activated.
//
// Permissions required
//	- roles/compute.instanceAdmin.v1 to get instance data and delete access config.
//
func RemovePublicIP(ctx context.Context, m pubsub.Message) error {
	r, err := removepublicip.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return removepublicip.Execute(ctx, r, ent)
}

// CloseCloudSQL removes public IP for a Cloud SQL instance.
//
// This Cloud Function will respond to Security Health Analytics **Public SQL Instance** findings
// from **SQL Scanner**. All public IP addresses of the affected instance will be
// deleted when this function is activated.
//
// Permissions required
//	- roles/cloudsql.editor to get instance data and delete access config.
//
func CloseCloudSQL(ctx context.Context, m pubsub.Message) error {
	r, err := removepublic.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return removepublic.Execute(ctx, r, ent)
}

// CloudSQLRequireSSL enables the SSL requirement for a Cloud SQL instance.
//
// This Cloud Function will respond to Security Health Analytics **Public SQL Instance** findings
// from **SQL Scanner**. All public IP addresses of the affected instance will be
// deleted when this function is activated.
//
// Permissions required
//	- roles/cloudsql.editor to get instance data and delete access config.
//
func CloudSQLRequireSSL(ctx context.Context, m pubsub.Message) error {
	r, err := requiressl.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return requiressl.Execute(ctx, r, ent)
}

// DisableDashboard will disable the Kubernetes dashboard addon.
//
// This Cloud Function will respond to Security Health Analytics **Web UI Enabled** findings
// from **Container Scanner**. The Kubernetes dashboard addon will be disabled when this
// function is activated.
//
// Permissions required
//	- roles/container.clusterAdmin update cluster addon.
//
func DisableDashboard(ctx context.Context, m pubsub.Message) error {
	r, err := disabledashboard.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return disabledashboard.Execute(ctx, r, ent)
}
<commit_msg>correct disabledashboard<commit_after>// Package exec is the entry point for security automation Cloud Functions.
package exec

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

import (
	"context"
	"log"

	"cloud.google.com/go/pubsub"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/cloud-sql/removepublic"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/cloud-sql/requiressl"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gce/createsnapshot"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gce/openfirewall"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gce/removepublicip"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gcs/closebucket"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/gke/disabledashboard"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/iam/removenonorgmembers"
	"github.com/googlecloudplatform/threat-automation/cloudfunctions/iam/revokeiam"
	"github.com/googlecloudplatform/threat-automation/entities"
)

var ent *entities.Entity

func init() {
	ctx := context.Background()
	var err error
	ent, err = entities.New(ctx)
	if err != nil {
		log.Fatalf("failed to initialize entities: %q", err)
	}
}

// RevokeIAM is the entry point for the IAM revoker Cloud Function.
//
// This Cloud Function will be triggered when Event Threat Detection
// detects an anomalous IAM grant. Once triggered this function will
// attempt to revoke the external members added to the policy if they match the provided
// list of disallowed domains. Additionally this method will only remove members if the
// project they were added to is within the specified folders. This configuration allows
// you to take a remediation action only on specific members and folders. For example,
// you may have a folder "development" where users can experiment without strict policies.
// However in your "production" folder you may want to revoke any grants that ETD finds as
// long as they match the domains you specify.
//
// Permissions required
// 	- roles/resourcemanager.folderAdmin to revoke IAM grants.
//	- roles/viewer to verify the affected project is within the enforced folder.
//
func RevokeIAM(ctx context.Context, m pubsub.Message) error {
	r, err := revokeiam.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return revokeiam.Execute(ctx, r, ent)
}

// SnapshotDisk is the entry point for the auto creation of GCE snapshots Cloud Function.
//
// This Cloud Function will respond to Event Threat Detection **bad IP** findings. Once a bad IP
// finding is received this Cloud Function will look for any existing disk snapshots for the
// affected instance. If there are recent snapshots then no action is taken. If we have not
// taken a snapshot recently, take a new snapshot for each disk within the instance.
//
// Permissions required
//	- roles/compute.instanceAdmin.v1 to manage disk snapshots.
//
func SnapshotDisk(ctx context.Context, m pubsub.Message) error {
	r, err := createsnapshot.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return createsnapshot.Execute(ctx, r, ent)
}

// CloseBucket will remove any public users from buckets found within the provided folders.
//
// Permissions required
//	- roles/viewer to retrieve ancestry.
//	- roles/storeage.admin to modify buckets.
//
func CloseBucket(ctx context.Context, m pubsub.Message) error {
	r, err := closebucket.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return closebucket.Execute(ctx, r, ent)
}

// OpenFirewall will remediate an open firewall.
//
// Permissions required
//	- roles/viewer to retrieve ancestry.
//	- roles/compute.securityAdmin to modify firewall rules.
//
func OpenFirewall(ctx context.Context, m pubsub.Message) error {
	r, err := openfirewall.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return openfirewall.Execute(ctx, r, ent)
}

// RemoveNonOrganizationMember removes all members that do not match the organization domain.
//
// This Cloud Function will respond to Security Health Analytics **NON_ORG_IAM_MEMBER** findings from **IAM Scanner**.
// All user member types (user:) that do not correspond to the organization will be removed from policy binding.
//
// Permissions required
//	- roles/resourcemanager.organizationAdmin to get org info and policies and set policies.
//
func RemoveNonOrganizationMember(ctx context.Context, m pubsub.Message) error {
	r, err := removenonorgmembers.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return removenonorgmembers.Execute(ctx, r, ent)
}

// RemovePublicIP removes all the external IP addresses of a GCE instance.
//
// This Cloud Function will respond to Security Health Analytics **Public IP Address** findings
// from **Compute Instance Scanner**. All public IP addresses of the affected instance will be
// deleted when this function is activated.
//
// Permissions required
//	- roles/compute.instanceAdmin.v1 to get instance data and delete access config.
//
func RemovePublicIP(ctx context.Context, m pubsub.Message) error {
	r, err := removepublicip.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return removepublicip.Execute(ctx, r, ent)
}

// CloseCloudSQL removes public IP for a Cloud SQL instance.
//
// This Cloud Function will respond to Security Health Analytics **Public SQL Instance** findings
// from **SQL Scanner**. All public IP addresses of the affected instance will be
// deleted when this function is activated.
//
// Permissions required
//	- roles/cloudsql.editor to get instance data and delete access config.
//
func CloseCloudSQL(ctx context.Context, m pubsub.Message) error {
	r, err := removepublic.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return removepublic.Execute(ctx, r, ent)
}

// CloudSQLRequireSSL enables the SSL requirement for a Cloud SQL instance.
//
// This Cloud Function will respond to Security Health Analytics **Public SQL Instance** findings
// from **SQL Scanner**. All public IP addresses of the affected instance will be
// deleted when this function is activated.
//
// Permissions required
//	- roles/cloudsql.editor to get instance data and delete access config.
//
func CloudSQLRequireSSL(ctx context.Context, m pubsub.Message) error {
	r, err := requiressl.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return requiressl.Execute(ctx, r, ent)
}

// DisableDashboard will disable the Kubernetes dashboard addon.
//
// This Cloud Function will respond to Security Health Analytics **Web UI Enabled** findings
// from **Container Scanner**. The Kubernetes dashboard addon will be disabled when this
// function is activated.
//
// Permissions required
//	- roles/container.clusterAdmin update cluster addon.
//
func DisableDashboard(ctx context.Context, m pubsub.Message) error {
	r, err := disabledashboard.ReadFinding(m.Data)
	if err != nil {
		return err
	}
	return disabledashboard.Execute(ctx, r, ent)
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws/awserr"
)

type ExecCommandInput struct {
	Profile     string
	Command     string
	Args        []string
	Keyring     keyring.Keyring
	Duration    time.Duration
	MfaToken    string
	StartServer bool
	Signals     chan os.Signal
}

func ExecCommand(ui Ui, input ExecCommandInput) {
	creds, err := NewVaultCredentials(input.Keyring, input.Profile, VaultOptions{
		SessionDuration: input.Duration,
		MfaToken:        input.MfaToken,
	})
	if err != nil {
		ui.Error.Fatal(err)
	}

	val, err := creds.Get()
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
			ui.Error.Fatalf("No credentials found for profile %q", input.Profile)
		} else {
			ui.Error.Fatal(err)
		}
	}

	profs, err := parseProfiles()
	if err != nil {
		ui.Error.Fatal(err)
	}

	cfg, err := writeTempConfig(input.Profile, profs)
	if err != nil {
		ui.Error.Fatal(err)
	}

	env := environ(os.Environ())
	env.Set("AWS_CONFIG_FILE", cfg.Name())
	env.Set("AWS_DEFAULT_PROFILE", input.Profile)
	env.Set("AWS_PROFILE", input.Profile)

	env.Unset("AWS_ACCESS_KEY_ID")
	env.Unset("AWS_SECRET_ACCESS_KEY")
	env.Unset("AWS_CREDENTIAL_FILE")

	if region, ok := profs[input.Profile]["region"]; ok {
		env.Set("AWS_DEFAULT_REGION", region)
		env.Set("AWS_REGION", region)
	}

	writeEnv := true

	if input.StartServer {
		if err := startCredentialsServer(ui, creds); err != nil {
			ui.Error.Fatal(err)
		} else {
			writeEnv = false
		}
	}

	if writeEnv {
		ui.Debug.Println("Writing temporary credentials to ENV")

		env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
		env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

		if val.SessionToken != "" {
			env.Set("AWS_SESSION_TOKEN", val.SessionToken)
			env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
		}
	}

	cmd := exec.Command(input.Command, input.Args...)
	cmd.Env = env
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	go func() {
		sig := <-input.Signals
		if cmd.Process != nil {
			cmd.Process.Signal(sig)
		}
	}()

	var waitStatus syscall.WaitStatus
	if err := cmd.Run(); err != nil {
		if err != nil {
			ui.Error.Println(err)
		}
		if exitError, ok := err.(*exec.ExitError); ok {
			waitStatus = exitError.Sys().(syscall.WaitStatus)
			os.Exit(waitStatus.ExitStatus())
		}
	}
}

// write out a config excluding role switching keys
func writeTempConfig(profile string, conf profiles) (*os.File, error) {
	tmpConfig, err := ioutil.TempFile(os.TempDir(), "aws-vault")
	if err != nil {
		return nil, err
	}

	newConfig := map[string]string{}

	for k, v := range conf[profile] {
		if k != "source_profile" && k != "role_arn" {
			newConfig[k] = v
		}
	}

	return tmpConfig, writeProfiles(tmpConfig, profiles{profile: newConfig})
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}
<commit_msg>Rely on ENV only for profile and config in subproc<commit_after>package main

import (
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"

	"github.com/99designs/aws-vault/keyring"
	"github.com/aws/aws-sdk-go/aws/awserr"
)

type ExecCommandInput struct {
	Profile     string
	Command     string
	Args        []string
	Keyring     keyring.Keyring
	Duration    time.Duration
	MfaToken    string
	StartServer bool
	Signals     chan os.Signal
}

func ExecCommand(ui Ui, input ExecCommandInput) {
	creds, err := NewVaultCredentials(input.Keyring, input.Profile, VaultOptions{
		SessionDuration: input.Duration,
		MfaToken:        input.MfaToken,
	})
	if err != nil {
		ui.Error.Fatal(err)
	}

	val, err := creds.Get()
	if err != nil {
		if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
			ui.Error.Fatalf("No credentials found for profile %q", input.Profile)
		} else {
			ui.Error.Fatal(err)
		}
	}

	profs, err := parseProfiles()
	if err != nil {
		ui.Error.Fatal(err)
	}

	env := environ(os.Environ())
	env.Set("AWS_CONFIG_FILE", "/dev/null")
	env.Set("AWS_VAULT", input.Profile)

	env.Unset("AWS_ACCESS_KEY_ID")
	env.Unset("AWS_SECRET_ACCESS_KEY")
	env.Unset("AWS_CREDENTIAL_FILE")
	env.Unset("AWS_DEFAULT_PROFILE")
	env.Unset("AWS_PROFILE")

	if region, ok := profs[input.Profile]["region"]; ok {
		env.Set("AWS_DEFAULT_REGION", region)
		env.Set("AWS_REGION", region)
	}

	writeEnv := true

	if input.StartServer {
		if err := startCredentialsServer(ui, creds); err != nil {
			ui.Error.Fatal(err)
		} else {
			writeEnv = false
		}
	}

	if writeEnv {
		ui.Debug.Println("Writing temporary credentials to ENV")

		env.Set("AWS_ACCESS_KEY_ID", val.AccessKeyID)
		env.Set("AWS_SECRET_ACCESS_KEY", val.SecretAccessKey)

		if val.SessionToken != "" {
			env.Set("AWS_SESSION_TOKEN", val.SessionToken)
			env.Set("AWS_SECURITY_TOKEN", val.SessionToken)
		}
	}

	cmd := exec.Command(input.Command, input.Args...)
	cmd.Env = env
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	go func() {
		sig := <-input.Signals
		if cmd.Process != nil {
			cmd.Process.Signal(sig)
		}
	}()

	var waitStatus syscall.WaitStatus
	if err := cmd.Run(); err != nil {
		if err != nil {
			ui.Error.Println(err)
		}
		if exitError, ok := err.(*exec.ExitError); ok {
			waitStatus = exitError.Sys().(syscall.WaitStatus)
			os.Exit(waitStatus.ExitStatus())
		}
	}
}

// environ is a slice of strings representing the environment, in the form "key=value".
type environ []string

// Unset an environment variable by key
func (e *environ) Unset(key string) {
	for i := range *e {
		if strings.HasPrefix((*e)[i], key+"=") {
			(*e)[i] = (*e)[len(*e)-1]
			*e = (*e)[:len(*e)-1]
			break
		}
	}
}

// Set adds an environment variable, replacing any existing ones of the same key
func (e *environ) Set(key, val string) {
	e.Unset(key)
	*e = append(*e, key+"="+val)
}
<|endoftext|>


<commit_before>package eval

import (
	"errors"
	"fmt"
	"reflect"

        "go/ast"
)

// EvalExpr is the main function to call to evaluate an ast-parsed
// expression, expr.  Parameter ctx contains a string representation
// of expr. Parameter env, contains an evaluation environment from
// which to get reflect.Values from. Note however that env can be
// subverted somewhat by supplying callback hooks routines which
// access variables and by supplying user-defined conversion routines.
func EvalExpr(ctx *Ctx, expr Expr, env *Env) (*[]reflect.Value, bool, error) {
	switch node := expr.(type) {
	case *Ident:
		v, typed, err := evalIdentExprCallback(ctx, node, env)
		if v == nil {
			return nil, false, err
		}
		ret := []reflect.Value{*v}
		return &ret, typed, err
	case *Ellipsis:
	case *BasicLit:
		v, err := evalBasicLit(ctx, node)
		return &[]reflect.Value{v}, true, err
	case *FuncLit:
	case *CompositeLit:
		v, err := evalCompositeLit(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *ParenExpr:
		return EvalExpr(ctx, node.X.(Expr), env)
	case *SelectorExpr:
		v, typed, err := evalSelectorExprCallback(ctx, node, env)
		if v == nil {
			return nil, typed, err
		}
		return &[]reflect.Value{*v}, typed, err
	case *IndexExpr:
		v, err := evalIndexExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *SliceExpr:
		v, err := evalSliceExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *TypeAssertExpr:
	case *CallExpr:
		vs, err := evalCallExpr(ctx, node, env)
		return &vs, true, err
	case *StarExpr:
		v, err := evalStarExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *UnaryExpr:
		v, err := evalUnaryExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *BinaryExpr:
		v, typed, err := evalBinaryExpr(ctx, node, env)
		return &[]reflect.Value{v}, typed, err
	case *KeyValueExpr:
	default:
		panic(node)
		return nil , false, errors.New("undefined type")
	}
	return &[]reflect.Value{reflect.ValueOf("Alice")}, true, nil
}

func evalType(ctx *Ctx, expr ast.Expr, env *Env) (reflect.Type, error) {
	switch node := expr.(type) {
	case *ast.Ident:
		if t, ok := env.Types[node.Name]; ok {
			return t, nil
		} else if t, ok := builtinTypes[node.Name]; ok {
			return t, nil
		} else {
			return t, errors.New("undefined type: " + node.Name)
		}
	case *ast.ArrayType:
		if node.Len != nil {
			return nil, errors.New("array types not implemented")
		} else if eltT, err := evalType(ctx, node.Elt, env); err != nil {
			return nil, err
		} else {
			return reflect.SliceOf(eltT), nil
		}
	case *ast.StructType:
		return nil, errors.New("struct types not implemented")
	case *ast.FuncType:
		return nil, errors.New("func types not implemented")
	case *ast.InterfaceType:
		return nil, errors.New("interface types not implemented")
	case *ast.MapType:
		return nil, errors.New("map types not implemented")
	case *ast.ChanType:
		return nil, errors.New("chan types not implemented")
	default:
		return nil, errors.New(fmt.Sprintf("Type: Bad type (%+v)", node))
	}
}
<commit_msg>Fixed parenthesised types<commit_after>package eval

import (
	"errors"
	"fmt"
	"reflect"

        "go/ast"
)

// EvalExpr is the main function to call to evaluate an ast-parsed
// expression, expr.  Parameter ctx contains a string representation
// of expr. Parameter env, contains an evaluation environment from
// which to get reflect.Values from. Note however that env can be
// subverted somewhat by supplying callback hooks routines which
// access variables and by supplying user-defined conversion routines.
func EvalExpr(ctx *Ctx, expr Expr, env *Env) (*[]reflect.Value, bool, error) {
	switch node := expr.(type) {
	case *Ident:
		v, typed, err := evalIdentExprCallback(ctx, node, env)
		if v == nil {
			return nil, false, err
		}
		ret := []reflect.Value{*v}
		return &ret, typed, err
	case *Ellipsis:
	case *BasicLit:
		v, err := evalBasicLit(ctx, node)
		return &[]reflect.Value{v}, true, err
	case *FuncLit:
	case *CompositeLit:
		v, err := evalCompositeLit(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *ParenExpr:
		return EvalExpr(ctx, node.X.(Expr), env)
	case *SelectorExpr:
		v, typed, err := evalSelectorExprCallback(ctx, node, env)
		if v == nil {
			return nil, typed, err
		}
		return &[]reflect.Value{*v}, typed, err
	case *IndexExpr:
		v, err := evalIndexExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *SliceExpr:
		v, err := evalSliceExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *TypeAssertExpr:
	case *CallExpr:
		vs, err := evalCallExpr(ctx, node, env)
		return &vs, true, err
	case *StarExpr:
		v, err := evalStarExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *UnaryExpr:
		v, err := evalUnaryExpr(ctx, node, env)
		return &[]reflect.Value{v}, true, err
	case *BinaryExpr:
		v, typed, err := evalBinaryExpr(ctx, node, env)
		return &[]reflect.Value{v}, typed, err
	case *KeyValueExpr:
	default:
		panic(node)
		return nil , false, errors.New("undefined type")
	}
	return &[]reflect.Value{reflect.ValueOf("Alice")}, true, nil
}

func evalType(ctx *Ctx, expr ast.Expr, env *Env) (reflect.Type, error) {
	for parens, ok := expr.(*ast.ParenExpr); ok; parens, ok = expr.(*ast.ParenExpr) {
		expr = parens.X
	}
	switch node := expr.(type) {
	case *ast.Ident:
		if t, ok := env.Types[node.Name]; ok {
			return t, nil
		} else if t, ok := builtinTypes[node.Name]; ok {
			return t, nil
		} else {
			return t, errors.New("undefined type: " + node.Name)
		}
	case *ast.ArrayType:
		if node.Len != nil {
			return nil, errors.New("array types not implemented")
		} else if eltT, err := evalType(ctx, node.Elt, env); err != nil {
			return nil, err
		} else {
			return reflect.SliceOf(eltT), nil
		}
	case *ast.StructType:
		return nil, errors.New("struct types not implemented")
	case *ast.FuncType:
		return nil, errors.New("func types not implemented")
	case *ast.InterfaceType:
		return nil, errors.New("interface types not implemented")
	case *ast.MapType:
		return nil, errors.New("map types not implemented")
	case *ast.ChanType:
		return nil, errors.New("chan types not implemented")
	default:
		return nil, errors.New(fmt.Sprintf("Type: Bad type (%+v)", node))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/ActiveState/tail"
	"github.com/foize/go.fifo"
	"time"
	"encoding/json"
	"os"
	"net/http"
	"bytes"
)

type Configuration struct {
    Api_endpoint, Api_key, Project, Env string
    Shipping_interval int
    Components []Component
}

type Component struct {
	Name, File string
}

type Message struct {
	Project, Env, Component, Text string
	Timestamp int64
}

func main() {
	// declare our outgoing queue and http client
	outgoing := fifo.NewQueue()
	client := &http.Client{}

	// load the configuration
	file, _ := os.Open("config.json")
	decoder := json.NewDecoder(file)
	configuration := Configuration{}
	err := decoder.Decode(&configuration)
	if err != nil {
		fmt.Println("error:", err)
	}

    // start Filewatchers for the files
    for _,comp := range configuration.Components {
    	fmt.Println("watching", comp.Name, comp.File)
    	go Filewatcher(comp, configuration, outgoing)
    }

    // loop
    for {
		outlen := outgoing.Len()
		if outlen > 0 {
			// create a request body
			var data = bytes.NewBuffer([]byte(``))

			// add queue contents as payload
			for {
				item := outgoing.Next()
				if item == nil {
					break
				}

				fmt.Println("shipping item:", item)

			}

			// send request
			req, _ := http.NewRequest("POST", configuration.Api_endpoint, data)
			req.Header.Set("x-api-key", configuration.Api_key)

			fmt.Println("sending request of", outlen, "items")
			resp, err := client.Do(req)
			if err != nil {
				fmt.Println("error requesting:", err)
			} else {
				fmt.Println(req)
				fmt.Println(resp)
			}
		}
		time.Sleep(time.Duration(configuration.Shipping_interval)*time.Second)
	}
}

func Filewatcher(comp Component, conf Configuration, outgoing *fifo.Queue) {
	// tail a given file, add appended lines into a Message queue
    t, _ := tail.TailFile(comp.File, tail.Config{Follow: true})
    for line := range t.Lines {
    	outgoing.Add(&Message {
        	Project: conf.Project,
        	Env: conf.Env,
        	Component: comp.Name,
        	Text: line.Text,
        	Timestamp: time.Now().Unix(),
    	})
    }
}<commit_msg>chore: formatting<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/ActiveState/tail"
	"github.com/foize/go.fifo"
	"net/http"
	"os"
	"time"
)

type Configuration struct {
	Api_endpoint, Api_key, Project, Env string
	Shipping_interval                   int
	Components                          []Component
}

type Component struct {
	Name, File string
}

type Message struct {
	Project, Env, Component, Text string
	Timestamp                     int64
}

func main() {
	// declare our outgoing queue and http client
	outgoing := fifo.NewQueue()
	client := &http.Client{}

	// load the configuration
	file, _ := os.Open("config.json")
	decoder := json.NewDecoder(file)
	configuration := Configuration{}
	err := decoder.Decode(&configuration)
	if err != nil {
		fmt.Println("error:", err)
	}

	// start Filewatchers for the files
	for _, comp := range configuration.Components {
		fmt.Println("watching", comp.Name, comp.File)
		go Filewatcher(comp, configuration, outgoing)
	}

	// loop
	for {
		outlen := outgoing.Len()
		if outlen > 0 {
			// create a request body
			var data = bytes.NewBuffer([]byte(``))

			// add queue contents as payload
			for {
				item := outgoing.Next()
				if item == nil {
					break
				}

				fmt.Println("shipping item:", item)

			}

			// send request
			req, _ := http.NewRequest("POST", configuration.Api_endpoint, data)
			req.Header.Set("x-api-key", configuration.Api_key)

			fmt.Println("sending request of", outlen, "items")
			resp, err := client.Do(req)
			if err != nil {
				fmt.Println("error requesting:", err)
			} else {
				fmt.Println(req)
				fmt.Println(resp)
			}
		}
		time.Sleep(time.Duration(configuration.Shipping_interval) * time.Second)
	}
}

func Filewatcher(comp Component, conf Configuration, outgoing *fifo.Queue) {
	// tail a given file, add appended lines into a Message queue
	t, _ := tail.TailFile(comp.File, tail.Config{Follow: true})
	for line := range t.Lines {
		outgoing.Add(&Message{
			Project:   conf.Project,
			Env:       conf.Env,
			Component: comp.Name,
			Text:      line.Text,
			Timestamp: time.Now().Unix(),
		})
	}
}
<|endoftext|>


<commit_before>package tview

import (
	"github.com/gdamore/tcell"
)

// Configuration values.
const (
	FlexRow = iota
	FlexColumn
)

// flexItem holds layout options for one item.
type flexItem struct {
	Item       Primitive // The item to be positioned. May be nil for an empty item.
	FixedSize  int       // The item's fixed size which may not be changed, 0 if it has no fixed size.
	Proportion int       // The item's proportion.
	Focus      bool      // Whether or not this item attracts the layout's focus.
}

// Flex is a basic implementation of the Flexbox layout. The contained
// primitives are arranged horizontally or vertically. The way they are
// distributed along that dimension depends on their layout settings, which is
// either a fixed length or a proportional length. See AddItem() for details.
//
// See https://github.com/rivo/tview/wiki/Flex for an example.
type Flex struct {
	*Box

	// The items to be positioned.
	items []flexItem

	// FlexRow or FlexColumn.
	direction int

	// If set to true, Flex will use the entire screen as its available space
	// instead its box dimensions.
	fullScreen bool
}

// NewFlex returns a new flexbox layout container with no primitives and its
// direction set to FlexColumn. To add primitives to this layout, see AddItem().
// To change the direction, see SetDirection().
//
// Note that Box, the superclass of Flex, will have its background color set to
// transparent so that any nil flex items will leave their background unchanged.
// To clear a Flex's background before any items are drawn, set it to the
// desired color:
//
//   flex.SetBackgroundColor(tview.Styles.PrimitiveBackgroundColor)
func NewFlex() *Flex {
	f := &Flex{
		Box:       NewBox().SetBackgroundColor(tcell.ColorDefault),
		direction: FlexColumn,
	}
	f.focus = f
	return f
}

// SetDirection sets the direction in which the contained primitives are
// distributed. This can be either FlexColumn (default) or FlexRow.
func (f *Flex) SetDirection(direction int) *Flex {
	f.direction = direction
	return f
}

// SetFullScreen sets the flag which, when true, causes the flex layout to use
// the entire screen space instead of whatever size it is currently assigned to.
func (f *Flex) SetFullScreen(fullScreen bool) *Flex {
	f.fullScreen = fullScreen
	return f
}

// AddItem adds a new item to the container. The "fixedSize" argument is a width
// or height that may not be changed by the layout algorithm. A value of 0 means
// that its size is flexible and may be changed. The "proportion" argument
// defines the relative size of the item compared to other flexible-size items.
// For example, items with a proportion of 2 will be twice as large as items
// with a proportion of 1. The proportion must be at least 1 if fixedSize == 0
// (ignored otherwise).
//
// If "focus" is set to true, the item will receive focus when the Flex
// primitive receives focus. If multiple items have the "focus" flag set to
// true, the first one will receive focus.
//
// You can provide a nil value for the primitive. This will still consume screen
// space but nothing will be drawn.
func (f *Flex) AddItem(item Primitive, fixedSize, proportion int, focus bool) *Flex {
	f.items = append(f.items, flexItem{Item: item, FixedSize: fixedSize, Proportion: proportion, Focus: focus})
	return f
}

// RemoveItem removes all items for the given primitive from the container,
// keeping the order of the remaining items intact.
func (f *Flex) RemoveItem(p Primitive) *Flex {
	for index := len(f.items) - 1; index >= 0; index-- {
		if f.items[index].Item == p {
			f.items = append(f.items[:index], f.items[index+1:]...)
		}
	}
	return f
}

// Draw draws this primitive onto the screen.
func (f *Flex) Draw(screen tcell.Screen) {
	// Calculate size and position of the items.

	// Do we use the entire screen?
	if f.fullScreen {
		width, height := screen.Size()
		f.SetRect(0, 0, width, height)
	}

	// How much space can we distribute?
	x, y, width, height := f.GetInnerRect()
	var proportionSum int
	distSize := width
	if f.direction == FlexRow {
		distSize = height
	}
	for _, item := range f.items {
		if item.FixedSize > 0 {
			distSize -= item.FixedSize
		} else {
			proportionSum += item.Proportion
		}
	}

	// Calculate positions and draw items.
	pos := x
	if f.direction == FlexRow {
		pos = y
	}
	for _, item := range f.items {
		size := item.FixedSize
		if size <= 0 {
			size = distSize * item.Proportion / proportionSum
			distSize -= size
			proportionSum -= item.Proportion
		}
		if item.Item != nil {
			if f.direction == FlexColumn {
				item.Item.SetRect(pos, y, size, height)
			} else {
				item.Item.SetRect(x, pos, width, size)
			}
		}
		pos += size

		if item.Item != nil {
			if item.Item.GetFocusable().HasFocus() {
				defer item.Item.Draw(screen)
			} else {
				item.Item.Draw(screen)
			}
		}
	}
}

// Focus is called when this primitive receives focus.
func (f *Flex) Focus(delegate func(p Primitive)) {
	for _, item := range f.items {
		if item.Item != nil && item.Focus {
			delegate(item.Item)
			return
		}
	}
}

// HasFocus returns whether or not this primitive has focus.
func (f *Flex) HasFocus() bool {
	for _, item := range f.items {
		if item.Item != nil && item.Item.GetFocusable().HasFocus() {
			return true
		}
	}
	return false
}
<commit_msg>Bugfix: Flex will call Box.Draw() again. Fixes #104<commit_after>package tview

import (
	"github.com/gdamore/tcell"
)

// Configuration values.
const (
	FlexRow = iota
	FlexColumn
)

// flexItem holds layout options for one item.
type flexItem struct {
	Item       Primitive // The item to be positioned. May be nil for an empty item.
	FixedSize  int       // The item's fixed size which may not be changed, 0 if it has no fixed size.
	Proportion int       // The item's proportion.
	Focus      bool      // Whether or not this item attracts the layout's focus.
}

// Flex is a basic implementation of the Flexbox layout. The contained
// primitives are arranged horizontally or vertically. The way they are
// distributed along that dimension depends on their layout settings, which is
// either a fixed length or a proportional length. See AddItem() for details.
//
// See https://github.com/rivo/tview/wiki/Flex for an example.
type Flex struct {
	*Box

	// The items to be positioned.
	items []flexItem

	// FlexRow or FlexColumn.
	direction int

	// If set to true, Flex will use the entire screen as its available space
	// instead its box dimensions.
	fullScreen bool
}

// NewFlex returns a new flexbox layout container with no primitives and its
// direction set to FlexColumn. To add primitives to this layout, see AddItem().
// To change the direction, see SetDirection().
//
// Note that Box, the superclass of Flex, will have its background color set to
// transparent so that any nil flex items will leave their background unchanged.
// To clear a Flex's background before any items are drawn, set it to the
// desired color:
//
//   flex.SetBackgroundColor(tview.Styles.PrimitiveBackgroundColor)
func NewFlex() *Flex {
	f := &Flex{
		Box:       NewBox().SetBackgroundColor(tcell.ColorDefault),
		direction: FlexColumn,
	}
	f.focus = f
	return f
}

// SetDirection sets the direction in which the contained primitives are
// distributed. This can be either FlexColumn (default) or FlexRow.
func (f *Flex) SetDirection(direction int) *Flex {
	f.direction = direction
	return f
}

// SetFullScreen sets the flag which, when true, causes the flex layout to use
// the entire screen space instead of whatever size it is currently assigned to.
func (f *Flex) SetFullScreen(fullScreen bool) *Flex {
	f.fullScreen = fullScreen
	return f
}

// AddItem adds a new item to the container. The "fixedSize" argument is a width
// or height that may not be changed by the layout algorithm. A value of 0 means
// that its size is flexible and may be changed. The "proportion" argument
// defines the relative size of the item compared to other flexible-size items.
// For example, items with a proportion of 2 will be twice as large as items
// with a proportion of 1. The proportion must be at least 1 if fixedSize == 0
// (ignored otherwise).
//
// If "focus" is set to true, the item will receive focus when the Flex
// primitive receives focus. If multiple items have the "focus" flag set to
// true, the first one will receive focus.
//
// You can provide a nil value for the primitive. This will still consume screen
// space but nothing will be drawn.
func (f *Flex) AddItem(item Primitive, fixedSize, proportion int, focus bool) *Flex {
	f.items = append(f.items, flexItem{Item: item, FixedSize: fixedSize, Proportion: proportion, Focus: focus})
	return f
}

// RemoveItem removes all items for the given primitive from the container,
// keeping the order of the remaining items intact.
func (f *Flex) RemoveItem(p Primitive) *Flex {
	for index := len(f.items) - 1; index >= 0; index-- {
		if f.items[index].Item == p {
			f.items = append(f.items[:index], f.items[index+1:]...)
		}
	}
	return f
}

// Draw draws this primitive onto the screen.
func (f *Flex) Draw(screen tcell.Screen) {
	f.Box.Draw(screen)

	// Calculate size and position of the items.

	// Do we use the entire screen?
	if f.fullScreen {
		width, height := screen.Size()
		f.SetRect(0, 0, width, height)
	}

	// How much space can we distribute?
	x, y, width, height := f.GetInnerRect()
	var proportionSum int
	distSize := width
	if f.direction == FlexRow {
		distSize = height
	}
	for _, item := range f.items {
		if item.FixedSize > 0 {
			distSize -= item.FixedSize
		} else {
			proportionSum += item.Proportion
		}
	}

	// Calculate positions and draw items.
	pos := x
	if f.direction == FlexRow {
		pos = y
	}
	for _, item := range f.items {
		size := item.FixedSize
		if size <= 0 {
			size = distSize * item.Proportion / proportionSum
			distSize -= size
			proportionSum -= item.Proportion
		}
		if item.Item != nil {
			if f.direction == FlexColumn {
				item.Item.SetRect(pos, y, size, height)
			} else {
				item.Item.SetRect(x, pos, width, size)
			}
		}
		pos += size

		if item.Item != nil {
			if item.Item.GetFocusable().HasFocus() {
				defer item.Item.Draw(screen)
			} else {
				item.Item.Draw(screen)
			}
		}
	}
}

// Focus is called when this primitive receives focus.
func (f *Flex) Focus(delegate func(p Primitive)) {
	for _, item := range f.items {
		if item.Item != nil && item.Focus {
			delegate(item.Item)
			return
		}
	}
}

// HasFocus returns whether or not this primitive has focus.
func (f *Flex) HasFocus() bool {
	for _, item := range f.items {
		if item.Item != nil && item.Item.GetFocusable().HasFocus() {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package immortal

import (
	"os"
	"os/exec"
	"syscall"
)

func Fork() (int, error) {
	args := os.Args[1:]
	cmd := exec.Command(os.Args[0], args...)
	cmd.Env = append(cmd.Env, os.Environ()...)
	cmd.Stdin = nil
	cmd.Stdout = nil
	cmd.Stderr = nil
	cmd.ExtraFiles = nil
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid: true,
	}
	if err := cmd.Start(); err != nil {
		return 0, err
	}
	return cmd.Process.Pid, nil
}
<commit_msg>Setctty: true Set controlling terminal to fd Ctty<commit_after>package immortal

import (
	"os"
	"os/exec"
	"syscall"
)

func Fork() (int, error) {
	args := os.Args[1:]
	cmd := exec.Command(os.Args[0], args...)
	cmd.Env = append(cmd.Env, os.Environ()...)
	cmd.Stdin = nil
	cmd.Stdout = nil
	cmd.Stderr = nil
	cmd.ExtraFiles = nil
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Setsid:  true,
		Setctty: true,
	}
	if err := cmd.Start(); err != nil {
		return 0, err
	}
	return cmd.Process.Pid, nil
}
<|endoftext|>


<commit_before>// Package gami provites primitives for interacting with Asterisk AMI
/*

Basic Usage

	ami, err := gami.Dial("127.0.0.1:5038")
	if err != nil {
		fmt.Print(err)
		os.Exit(1)
	}
	ami.Run()
	defer ami.Close()

	//install manager
	go func() {
		for {
			select {
			//handle network errors
			case err := <-ami.NetError:
				log.Println("Network Error:", err)
				//try new connection every second
				for {
					<-time.After(time.Second)
					if err := ami.Reconnect(); err == nil {
						//call start actions
						if _, err := ami.Action("Events", gami.Params{"EventMask": "on"}); err != nil {
							break
						}
					}
				}

			case err := <-ami.Error:
				log.Println("error:", err)
			//wait events and process
			case ev := <-ami.Events:
				log.Println("Event Detect: %v", *ev)
				//if want type of events
				log.Println("EventType:", event.New(ev))
			}
		}
	}()

	if err := ami.Login("admin", "root"); err != nil {
		log.Fatal(err)
	}

	if rs, err = ami.Action("Ping", nil); err == nil {
		log.Fatal(rs)
	}

	//or with can do async
	pingResp, pingErr := ami.AsyncAction("Ping", gami.Params{"ActionID": "miping"})
	if pingErr != nil {
		log.Fatal(pingErr)
	}

	if rs, err = ami.Action("Events", ami.Params{"EventMask":"on"}); err != nil {
		fmt.Print(err)
	}

	log.Println("future ping:", <-pingResp)

*/
package gami

import (
	"errors"
	"io"
	"net"
	"net/textproto"
	"strings"
	"syscall"
)

const responseChanGamiID = "gamigeneral"

var errNotEvent = errors.New("Not Event")

// Raise when not response expected protocol AMI
var ErrNotAMI = errors.New("Server not AMI interface")

// Params for the actions
type Params map[string]string

// AMIClient a connection to AMI server
type AMIClient struct {
	conn    *textproto.Conn
	connRaw *net.Conn

	address string
	amiUser string
	amiPass string

	// network wait for a new connection
	waitNewConnection chan struct{}

	response map[string]chan *AMIResponse

	// Events for client parse
	Events chan *AMIEvent

	// Error Raise on logic
	Error chan error

	//NetError a network error
	NetError chan error
}

// AMIResponse from action
type AMIResponse struct {
	ID     string
	Status string
	Params map[string]string
}

// AMIEvent it's a representation of Event readed
type AMIEvent struct {
	//Identification of event Event: xxxx
	ID string

	Privilege []string

	// Params  of arguments received
	Params map[string]string
}

// Login authenticate to AMI
func (client *AMIClient) Login(username, password string) error {
	response, err := client.Action("Login", Params{"Username": username, "Secret": password})
	if err != nil {
		return err
	}

	if (*response).Status == "Error" {
		return errors.New((*response).Params["Message"])
	}

	client.amiUser = username
	client.amiPass = password
	return nil
}

// Reconnect the session, autologin if a new network error it put on client.NetError
func (client *AMIClient) Reconnect() error {
	client.conn.Close()
	reconnect, err := Dial(client.address)

	if err != nil {
		client.NetError <- err
		return err
	}

	//new connection
	client.conn = reconnect.conn
	client.connRaw = reconnect.connRaw
	client.waitNewConnection <- struct{}{}

	if err := client.Login(client.amiUser, client.amiPass); err != nil {
		return err
	}

	return nil
}

// AsyncAction return chan for wait response of action with parameter *ActionID* this can be helpful for
// massive actions,
func (client *AMIClient) AsyncAction(action string, params Params) (<-chan *AMIResponse, error) {
	if err := client.conn.PrintfLine("Action: %s", strings.TrimSpace(action)); err != nil {
		return nil, err
	}

	if _, ok := params["ActionID"]; ok {
		params["ActionID"] = responseChanGamiID
	}

	if _, ok := client.response[params["ActionID"]]; !ok {
		client.response[params["ActionID"]] = make(chan *AMIResponse, 1)
	}

	for k, v := range params {
		if err := client.conn.PrintfLine("%s: %s", k, strings.TrimSpace(v)); err != nil {
			return nil, err
		}
	}

	if err := client.conn.PrintfLine(""); err != nil {
		return nil, err
	}

	return client.response[params["ActionID"]], nil
}

// Action send with params
func (client *AMIClient) Action(action string, params Params) (*AMIResponse, error) {
	resp, err := client.AsyncAction(action, params)
	if err != nil {
		return nil, err
	}
	response := <-resp

	return response, nil
}

// Run process socket waiting events and responses
func (client *AMIClient) Run() {
	go func() {
		for {
			data, err := client.conn.ReadMIMEHeader()

			if err != nil {
				switch err {
				case syscall.ECONNABORTED:
					fallthrough
				case syscall.ECONNRESET:
					fallthrough
				case syscall.ECONNREFUSED:
					fallthrough
				case io.EOF:
					client.NetError <- err
					<-client.waitNewConnection
				default:
					client.Error <- err
				}
				continue
			}

			if ev, err := newEvent(&data); err != nil {
				if err != errNotEvent {
					client.Error <- err
				}
			} else {
				client.Events <- ev
			}

			if response, err := newResponse(&data); err == nil {
				client.response[response.ID] <- response
			}

		}
	}()
}

// Close the connection to AMI
func (client *AMIClient) Close() {
	client.Action("Logoff", nil)
	(*client.connRaw).Close()
}

//newResponse build a response for action
func newResponse(data *textproto.MIMEHeader) (*AMIResponse, error) {
	if data.Get("Response") == "" {
		return nil, errors.New("Not Response")
	}
	response := &AMIResponse{"", "", make(map[string]string)}
	for k, v := range *data {
		if k == "Response" {
			continue
		}
		response.Params[k] = v[0]
	}
	response.ID = data.Get("Actionid")
	response.Status = data.Get("Response")
	return response, nil
}

//newEvent build event
func newEvent(data *textproto.MIMEHeader) (*AMIEvent, error) {
	if data.Get("Event") == "" {
		return nil, errNotEvent
	}
	ev := &AMIEvent{data.Get("Event"), strings.Split(data.Get("Privilege"), ","), make(map[string]string)}
	for k, v := range *data {
		if k == "Event" || k == "Privilege" {
			continue
		}
		ev.Params[k] = v[0]
	}
	return ev, nil
}

// Dial create a new connection to AMI
func Dial(address string) (*AMIClient, error) {
	connRaw, err := net.Dial("tcp", address)

	if err != nil {
		return nil, err
	}
	conn := textproto.NewConn(connRaw)
	label, err := conn.ReadLine()
	if err != nil {
		return nil, err
	}

	if strings.Contains(label, "Asterisk Call Manager") != true {
		return nil, ErrNotAMI
	}

	client := &AMIClient{
		conn:              conn,
		connRaw:           &connRaw,
		address:           address,
		amiUser:           "",
		amiPass:           "",
		waitNewConnection: make(chan struct{}),
		response:          make(map[string]chan *AMIResponse),
		Events:            make(chan *AMIEvent, 100),
		Error:             make(chan error, 1),
		NetError:          make(chan error, 1),
	}
	return client, nil
}
<commit_msg>NewConn reuse connection<commit_after>// Package gami provites primitives for interacting with Asterisk AMI
/*

Basic Usage

	ami, err := gami.Dial("127.0.0.1:5038")
	if err != nil {
		fmt.Print(err)
		os.Exit(1)
	}
	ami.Run()
	defer ami.Close()

	//install manager
	go func() {
		for {
			select {
			//handle network errors
			case err := <-ami.NetError:
				log.Println("Network Error:", err)
				//try new connection every second
				for {
					<-time.After(time.Second)
					if err := ami.Reconnect(); err == nil {
						//call start actions
						if _, err := ami.Action("Events", gami.Params{"EventMask": "on"}); err != nil {
							break
						}
					}
				}

			case err := <-ami.Error:
				log.Println("error:", err)
			//wait events and process
			case ev := <-ami.Events:
				log.Println("Event Detect: %v", *ev)
				//if want type of events
				log.Println("EventType:", event.New(ev))
			}
		}
	}()

	if err := ami.Login("admin", "root"); err != nil {
		log.Fatal(err)
	}

	if rs, err = ami.Action("Ping", nil); err == nil {
		log.Fatal(rs)
	}

	//or with can do async
	pingResp, pingErr := ami.AsyncAction("Ping", gami.Params{"ActionID": "miping"})
	if pingErr != nil {
		log.Fatal(pingErr)
	}

	if rs, err = ami.Action("Events", ami.Params{"EventMask":"on"}); err != nil {
		fmt.Print(err)
	}

	log.Println("future ping:", <-pingResp)

*/
package gami

import (
	"errors"
	"io"
	"net"
	"net/textproto"
	"strings"
	"syscall"
)

const responseChanGamiID = "gamigeneral"

var errNotEvent = errors.New("Not Event")

// Raise when not response expected protocol AMI
var ErrNotAMI = errors.New("Server not AMI interface")

// Params for the actions
type Params map[string]string

// AMIClient a connection to AMI server
type AMIClient struct {
	conn    *textproto.Conn
	connRaw io.ReadWriteCloser

	address string
	amiUser string
	amiPass string

	// network wait for a new connection
	waitNewConnection chan struct{}

	response map[string]chan *AMIResponse

	// Events for client parse
	Events chan *AMIEvent

	// Error Raise on logic
	Error chan error

	//NetError a network error
	NetError chan error
}

// AMIResponse from action
type AMIResponse struct {
	ID     string
	Status string
	Params map[string]string
}

// AMIEvent it's a representation of Event readed
type AMIEvent struct {
	//Identification of event Event: xxxx
	ID string

	Privilege []string

	// Params  of arguments received
	Params map[string]string
}

// Login authenticate to AMI
func (client *AMIClient) Login(username, password string) error {
	response, err := client.Action("Login", Params{"Username": username, "Secret": password})
	if err != nil {
		return err
	}

	if (*response).Status == "Error" {
		return errors.New((*response).Params["Message"])
	}

	client.amiUser = username
	client.amiPass = password
	return nil
}

// Reconnect the session, autologin if a new network error it put on client.NetError
func (client *AMIClient) Reconnect() error {
	client.conn.Close()
	reconnect, err := Dial(client.address)

	if err != nil {
		client.NetError <- err
		return err
	}

	//new connection
	client.conn = reconnect.conn
	client.connRaw = reconnect.connRaw
	client.waitNewConnection <- struct{}{}

	if err := client.Login(client.amiUser, client.amiPass); err != nil {
		return err
	}

	return nil
}

// AsyncAction return chan for wait response of action with parameter *ActionID* this can be helpful for
// massive actions,
func (client *AMIClient) AsyncAction(action string, params Params) (<-chan *AMIResponse, error) {
	if err := client.conn.PrintfLine("Action: %s", strings.TrimSpace(action)); err != nil {
		return nil, err
	}

	if _, ok := params["ActionID"]; ok {
		params["ActionID"] = responseChanGamiID
	}

	if _, ok := client.response[params["ActionID"]]; !ok {
		client.response[params["ActionID"]] = make(chan *AMIResponse, 1)
	}

	for k, v := range params {
		if err := client.conn.PrintfLine("%s: %s", k, strings.TrimSpace(v)); err != nil {
			return nil, err
		}
	}

	if err := client.conn.PrintfLine(""); err != nil {
		return nil, err
	}

	return client.response[params["ActionID"]], nil
}

// Action send with params
func (client *AMIClient) Action(action string, params Params) (*AMIResponse, error) {
	resp, err := client.AsyncAction(action, params)
	if err != nil {
		return nil, err
	}
	response := <-resp

	return response, nil
}

// Run process socket waiting events and responses
func (client *AMIClient) Run() {
	go func() {
		for {
			data, err := client.conn.ReadMIMEHeader()

			if err != nil {
				switch err {
				case syscall.ECONNABORTED:
					fallthrough
				case syscall.ECONNRESET:
					fallthrough
				case syscall.ECONNREFUSED:
					fallthrough
				case io.EOF:
					client.NetError <- err
					<-client.waitNewConnection
				default:
					client.Error <- err
				}
				continue
			}

			if ev, err := newEvent(&data); err != nil {
				if err != errNotEvent {
					client.Error <- err
				}
			} else {
				client.Events <- ev
			}

			if response, err := newResponse(&data); err == nil {
				client.response[response.ID] <- response
			}

		}
	}()
}

// Close the connection to AMI
func (client *AMIClient) Close() {
	client.Action("Logoff", nil)
	(client.connRaw).Close()
}

//newResponse build a response for action
func newResponse(data *textproto.MIMEHeader) (*AMIResponse, error) {
	if data.Get("Response") == "" {
		return nil, errors.New("Not Response")
	}
	response := &AMIResponse{"", "", make(map[string]string)}
	for k, v := range *data {
		if k == "Response" {
			continue
		}
		response.Params[k] = v[0]
	}
	response.ID = data.Get("Actionid")
	response.Status = data.Get("Response")
	return response, nil
}

//newEvent build event
func newEvent(data *textproto.MIMEHeader) (*AMIEvent, error) {
	if data.Get("Event") == "" {
		return nil, errNotEvent
	}
	ev := &AMIEvent{data.Get("Event"), strings.Split(data.Get("Privilege"), ","), make(map[string]string)}
	for k, v := range *data {
		if k == "Event" || k == "Privilege" {
			continue
		}
		ev.Params[k] = v[0]
	}
	return ev, nil
}

// Dial create a new connection to AMI
func Dial(address string) (*AMIClient, error) {
	connRaw, err := net.Dial("tcp", address)

	if err != nil {
		return nil, err
	}
	conn := textproto.NewConn(connRaw)
	label, err := conn.ReadLine()
	if err != nil {
		return nil, err
	}

	if strings.Contains(label, "Asterisk Call Manager") != true {
		return nil, ErrNotAMI
	}

	client := &AMIClient{
		conn:              conn,
		connRaw:           connRaw,
		address:           address,
		amiUser:           "",
		amiPass:           "",
		waitNewConnection: make(chan struct{}),
		response:          make(map[string]chan *AMIResponse),
		Events:            make(chan *AMIEvent, 100),
		Error:             make(chan error, 1),
		NetError:          make(chan error, 1),
	}
	return client, nil
}

// NewConn create a new connection to AMI
func NewConn(connRaw io.ReadWriteCloser, address string) (*AMIClient, error) {
	conn := textproto.NewConn(connRaw)
	label, err := conn.ReadLine()
	if err != nil {
		return nil, err
	}

	if strings.Contains(label, "Asterisk Call Manager") != true {
		return nil, ErrNotAMI
	}

	client := &AMIClient{
		conn:              conn,
		connRaw:           connRaw,
		address:           address,
		amiUser:           "",
		amiPass:           "",
		waitNewConnection: make(chan struct{}),
		response:          make(map[string]chan *AMIResponse),
		Events:            make(chan *AMIEvent, 100),
		Error:             make(chan error, 1),
		NetError:          make(chan error, 1),
	}
	return client, nil
}
<|endoftext|>