<commit_before>package sched

import (
	"bytes"
	"fmt"
	"log"
	"text/template"
	"time"

	"bosun.org/cmd/bosun/conf"
	"bosun.org/cmd/bosun/expr"
)

// Poll dispatches notification checks when needed.
func (s *Schedule) Poll() {
	for {
		rh := s.NewRunHistory(time.Now())
		timeout := s.CheckNotifications(rh)
		s.Save()
		// Wait for one of these two.
		select {
		case <-time.After(timeout):
		case <-s.nc:
		}
	}
}

func (s *Schedule) Notify(st *State, n *conf.Notification) {
	if s.notifications == nil {
		s.notifications = make(map[*conf.Notification][]*State)
	}
	s.notifications[n] = append(s.notifications[n], st)
}

// CheckNotifications processes past notification events. It returns the
// duration until the soonest notification triggers.
func (s *Schedule) CheckNotifications(rh *RunHistory) time.Duration {
	silenced := s.Silenced()
	s.Lock()
	defer s.Unlock()
	notifications := s.Notifications
	s.Notifications = nil
	for ak, ns := range notifications {
		if _, present := silenced[ak]; present {
			log.Println("silencing", ak)
			continue
		}
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(time.Now())
			if remaining > 0 {
				s.AddNotification(ak, n, t)
				continue
			}
			st := s.status[ak]
			if st == nil {
				continue
			}
			s.Notify(st, n)
		}
	}
	s.sendNotifications(rh, silenced)
	s.notifications = nil
	timeout := time.Hour
	now := time.Now()
	for _, ns := range s.Notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(now)
			if remaining < timeout {
				timeout = remaining
			}
		}
	}
	return timeout
}

func (s *Schedule) sendNotifications(rh *RunHistory, silenced map[expr.AlertKey]Silence) {
	if s.Conf.Quiet {
		log.Println("quiet mode prevented", len(s.notifications), "notifications")
		return
	}
	for n, states := range s.notifications {
		ustates := make(States)
		for _, st := range states {
			ak := st.AlertKey()
			if st.Last().Status == StUnknown {
				if _, ok := silenced[ak]; ok {
					log.Println("silencing unknown", ak)
					continue
				}
				ustates[ak] = st
			} else {
				s.notify(rh, st, n)
			}
			if n.Next != nil {
				s.AddNotification(ak, n, time.Now().UTC())
			}
		}
		var c int
		tHit := false
		oTSets := make(map[string]expr.AlertKeys)
		groupSets := ustates.GroupSets()
		for name, group := range groupSets {
			c++
			if c >= s.Conf.UnknownThreshold && s.Conf.UnknownThreshold > 0 {
				if !tHit && len(groupSets) == 0 {
					// If the threshold is hit but only 1 email remains, just send the normal unknown
					s.unotify(name, group, n)
					break
				}
				tHit = true
				oTSets[name] = group
			} else {
				s.unotify(name, group, n)
			}
		}
		if len(oTSets) > 0 {
			s.utnotify(oTSets, n)
		}
	}
}

var unknownMultiGroup = template.Must(template.New("unknownMultiGroup").Parse(`

Threshold of {{ .Threshold }} reached for unknown notifications. The following unknown
	group emails were not sent.

	{{ range $group, $alertKeys := .Groups }}

			{{ $group }}

				{{ range $ak := $alertKeys }}

{{ $ak }}

				{{ end }}

	{{ end }}

	`))

func (s *Schedule) notify(rh *RunHistory, st *State, n *conf.Notification) {
	n.Notify([]byte(st.Subject), st.EmailBody, s.Conf, string(st.AlertKey()), st.Attachments...)
}

// utnotify is single notification for N unknown groups into a single notification
func (s *Schedule) utnotify(groups map[string]expr.AlertKeys, n *conf.Notification) {
	var total int
	now := time.Now().UTC()
	for _, group := range groups {
		// Don't know what the following line does, just copied from unotify
		s.Group[now] = group
		total += len(group)
	}
	subject := fmt.Sprintf("%v unknown alert instances suppressed", total)
	body := new(bytes.Buffer)
	if err := unknownMultiGroup.Execute(body, struct {
		Groups    map[string]expr.AlertKeys
		Threshold int
	}{
		groups,
		s.Conf.UnknownThreshold,
	}); err != nil {
		log.Println(err)
	}
	n.Notify([]byte(subject), body.Bytes(), s.Conf, "unknown_treshold")
}

func (s *Schedule) unotify(name string, group expr.AlertKeys, n *conf.Notification) {
	subject := new(bytes.Buffer)
	body := new(bytes.Buffer)
	now := time.Now().UTC()
	s.Group[now] = group
	if t := s.Conf.UnknownTemplate; t != nil {
		data := s.unknownData(now, name, group)
		if t.Body != nil {
			if err := t.Body.Execute(body, &data); err != nil {
				log.Println("unknown template error:", err)
			}
		}
		if t.Subject != nil {
			if err := t.Subject.Execute(subject, &data); err != nil {
				log.Println("unknown template error:", err)
			}
		}
	}
	n.Notify(subject.Bytes(), body.Bytes(), s.Conf, name)
}

func (s *Schedule) AddNotification(ak expr.AlertKey, n *conf.Notification, started time.Time) {
	if s.Notifications == nil {
		s.Notifications = make(map[expr.AlertKey]map[string]time.Time)
	}
	if s.Notifications[ak] == nil {
		s.Notifications[ak] = make(map[string]time.Time)
	}
	s.Notifications[ak][n.Name] = started
}
<commit_msg>Making sure to queue next notification instead of requeueing  same notification again.<commit_after>package sched

import (
	"bytes"
	"fmt"
	"log"
	"text/template"
	"time"

	"bosun.org/cmd/bosun/conf"
	"bosun.org/cmd/bosun/expr"
)

// Poll dispatches notification checks when needed.
func (s *Schedule) Poll() {
	for {
		rh := s.NewRunHistory(time.Now())
		timeout := s.CheckNotifications(rh)
		s.Save()
		// Wait for one of these two.
		select {
		case <-time.After(timeout):
		case <-s.nc:
		}
	}
}

func (s *Schedule) Notify(st *State, n *conf.Notification) {
	if s.notifications == nil {
		s.notifications = make(map[*conf.Notification][]*State)
	}
	s.notifications[n] = append(s.notifications[n], st)
}

// CheckNotifications processes past notification events. It returns the
// duration until the soonest notification triggers.
func (s *Schedule) CheckNotifications(rh *RunHistory) time.Duration {
	silenced := s.Silenced()
	s.Lock()
	defer s.Unlock()
	notifications := s.Notifications
	s.Notifications = nil
	for ak, ns := range notifications {
		if _, present := silenced[ak]; present {
			log.Println("silencing", ak)
			continue
		}
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(time.Now())
			if remaining > 0 {
				s.AddNotification(ak, n, t)
				continue
			}
			st := s.status[ak]
			if st == nil {
				continue
			}
			s.Notify(st, n)
		}
	}
	s.sendNotifications(rh, silenced)
	s.notifications = nil
	timeout := time.Hour
	now := time.Now()
	for _, ns := range s.Notifications {
		for name, t := range ns {
			n, present := s.Conf.Notifications[name]
			if !present {
				continue
			}
			remaining := t.Add(n.Timeout).Sub(now)
			if remaining < timeout {
				timeout = remaining
			}
		}
	}
	return timeout
}

func (s *Schedule) sendNotifications(rh *RunHistory, silenced map[expr.AlertKey]Silence) {
	if s.Conf.Quiet {
		log.Println("quiet mode prevented", len(s.notifications), "notifications")
		return
	}
	for n, states := range s.notifications {
		ustates := make(States)
		for _, st := range states {
			ak := st.AlertKey()
			if st.Last().Status == StUnknown {
				if _, ok := silenced[ak]; ok {
					log.Println("silencing unknown", ak)
					continue
				}
				ustates[ak] = st
			} else {
				s.notify(rh, st, n)
			}
			if n.Next != nil {
				s.AddNotification(ak, n.Next, time.Now().UTC())
			}
		}
		var c int
		tHit := false
		oTSets := make(map[string]expr.AlertKeys)
		groupSets := ustates.GroupSets()
		for name, group := range groupSets {
			c++
			if c >= s.Conf.UnknownThreshold && s.Conf.UnknownThreshold > 0 {
				if !tHit && len(groupSets) == 0 {
					// If the threshold is hit but only 1 email remains, just send the normal unknown
					s.unotify(name, group, n)
					break
				}
				tHit = true
				oTSets[name] = group
			} else {
				s.unotify(name, group, n)
			}
		}
		if len(oTSets) > 0 {
			s.utnotify(oTSets, n)
		}
	}
}

var unknownMultiGroup = template.Must(template.New("unknownMultiGroup").Parse(`

Threshold of {{ .Threshold }} reached for unknown notifications. The following unknown
	group emails were not sent.

	{{ range $group, $alertKeys := .Groups }}

			{{ $group }}

				{{ range $ak := $alertKeys }}

{{ $ak }}

				{{ end }}

	{{ end }}

	`))

func (s *Schedule) notify(rh *RunHistory, st *State, n *conf.Notification) {
	n.Notify([]byte(st.Subject), st.EmailBody, s.Conf, string(st.AlertKey()), st.Attachments...)
}

// utnotify is single notification for N unknown groups into a single notification
func (s *Schedule) utnotify(groups map[string]expr.AlertKeys, n *conf.Notification) {
	var total int
	now := time.Now().UTC()
	for _, group := range groups {
		// Don't know what the following line does, just copied from unotify
		s.Group[now] = group
		total += len(group)
	}
	subject := fmt.Sprintf("%v unknown alert instances suppressed", total)
	body := new(bytes.Buffer)
	if err := unknownMultiGroup.Execute(body, struct {
		Groups    map[string]expr.AlertKeys
		Threshold int
	}{
		groups,
		s.Conf.UnknownThreshold,
	}); err != nil {
		log.Println(err)
	}
	n.Notify([]byte(subject), body.Bytes(), s.Conf, "unknown_treshold")
}

func (s *Schedule) unotify(name string, group expr.AlertKeys, n *conf.Notification) {
	subject := new(bytes.Buffer)
	body := new(bytes.Buffer)
	now := time.Now().UTC()
	s.Group[now] = group
	if t := s.Conf.UnknownTemplate; t != nil {
		data := s.unknownData(now, name, group)
		if t.Body != nil {
			if err := t.Body.Execute(body, &data); err != nil {
				log.Println("unknown template error:", err)
			}
		}
		if t.Subject != nil {
			if err := t.Subject.Execute(subject, &data); err != nil {
				log.Println("unknown template error:", err)
			}
		}
	}
	n.Notify(subject.Bytes(), body.Bytes(), s.Conf, name)
}

func (s *Schedule) AddNotification(ak expr.AlertKey, n *conf.Notification, started time.Time) {
	if s.Notifications == nil {
		s.Notifications = make(map[expr.AlertKey]map[string]time.Time)
	}
	if s.Notifications[ak] == nil {
		s.Notifications[ak] = make(map[string]time.Time)
	}
	s.Notifications[ak][n.Name] = started
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"github.com/GeoNet/weft"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// supported query parameters for the event service from http://www.fdsn.org/webservices/FDSN-WS-Specifications-1.1.pdf
type fdsnEventV1 struct {
	PublicID             string  `schema:"eventid"`      // select a specific event by ID; event identifiers are data center specific.
	MinLatitude          float64 `schema:"minlatitude"`  // limit to events with a latitude larger than or equal to the specified minimum.
	MaxLatitude          float64 `schema:"maxlatitude"`  // limit to events with a latitude smaller than or equal to the specified maximum.
	MinLongitude         float64 `schema:"minlongitude"` // limit to events with a longitude larger than or equal to the specified minimum.
	MaxLongitude         float64 `schema:"maxlongitude"` // limit to events with a longitude smaller than or equal to the specified maximum.
	MinDepth             float64 `schema:"mindepth"`     // limit to events with depth more than the specified minimum.
	MaxDepth             float64 `schema:"maxdepth"`     // limit to events with depth less than the specified maximum.
	MinMagnitude         float64 `schema:"minmagnitude"` // limit to events with a magnitude larger than the specified minimum.
	MaxMagnitude         float64 `schema:"maxmagnitude"` // limit to events with a magnitude smaller than the specified maximum.
	OrderBy              string  `schema:"orderby"`      // order the result by time or magnitude with the following possibilities: time, time-asc, magnitude, magnitude-asc
	StartTime            Time    `schema:"starttime"`    // limit to events on or after the specified start time.
	EndTime              Time    `schema:"endtime"`      // limit to events on or before the specified end time.
	IncludeAllOrigins    bool    `schema:"includeallorigins"`
	IncludeAllMagnitudes bool    `schema:"includeallmagnitudes"`
	IncludeArrivals      bool    `schema:"includearrivals"`
	Format               string  `schema:"format"`
}

type Time struct {
	time.Time
}

var fdsnEventWadlFile []byte
var fdsnEventIndex []byte
var eventNotSupported = map[string]bool{
	"latitude":      true,
	"longitude":     true,
	"minradius":     true,
	"maxraduis":     true,
	"magnitudetype": true,
	"limit":         true,
	"offset":        true,
	"catalog":       true,
	"contributor":   true,
	"updateafter":   true,
	"nodata":        true,
}

func init() {
	var err error
	fdsnEventWadlFile, err = ioutil.ReadFile("assets/fdsn-ws-event.wadl")
	if err != nil {
		log.Printf("error reading assets/fdsn-ws-event.wadl: %s", err.Error())
	}

	fdsnEventIndex, err = ioutil.ReadFile("assets/fdsn-ws-event.html")
	if err != nil {
		log.Printf("error reading assets/fdsn-ws-event.html: %s", err.Error())
	}
}

/*
parses the time in text as per the FDSN spec.  Pads text for parsing with
time.RFC3339Nano.  Accepted formats are (UTC):
   YYYY-MM-DDTHH:MM:SS.ssssss
   YYYY-MM-DDTHH:MM:SS
   YYYY-MM-DD

Implements the encoding.TextUnmarshaler interface.
*/
func (t *Time) UnmarshalText(text []byte) (err error) {
	s := string(text)
	l := len(s)
	if len(s) < 10 {
		return fmt.Errorf("invalid time format: %s", s)
	}

	if l >= 19 && l <= 26 && l!=20 {	// length 20: "YYYY-MM-DDTHH:MM:SS." invalid
		s = s + ".000000000Z"[(l-19):] // "YYYY-MM-DDTHH:MM:SS" append to nano
	} else if l == 10 {
		s = s + "T00:00:00.000000000Z" // YYYY-MM-DD
	} else {
		return fmt.Errorf("invalid time format: %s", s)
	}
	t.Time, err = time.Parse(time.RFC3339Nano, s)
	return
}

func parseEventV1(v url.Values) (fdsnEventV1, error) {
	// All query parameters are optional and float zero values overlap
	// with possible request ranges so the default is set to the max float val.
	e := fdsnEventV1{
		MinLatitude:  math.MaxFloat64,
		MaxLatitude:  math.MaxFloat64,
		MinLongitude: math.MaxFloat64,
		MaxLongitude: math.MaxFloat64,
		MinDepth:     math.MaxFloat64,
		MaxDepth:     math.MaxFloat64,
		MinMagnitude: math.MaxFloat64,
		MaxMagnitude: math.MaxFloat64,
	}

	for key, val := range v {
		if _, ok := eventNotSupported[key]; ok {
			return e, fmt.Errorf("\"%s\" is not supported", key)
		}
		if len(val[0]) == 0 {
			return e, fmt.Errorf("Invalid %s value", key)
		}
	}

	err := decoder.Decode(&e, v)
	if err != nil {
		return e, err
	}

	if e.IncludeAllMagnitudes {
		return e, errors.New("include all magnitudes is not supported.")
	}

	if e.IncludeAllOrigins {
		return e, errors.New("include all origins is not supported.")
	}

	if e.IncludeArrivals {
		return e, errors.New("include arrivals is not supported.")
	}

	// geometry bounds checking
	if e.MinLatitude != math.MaxFloat64 && e.MinLatitude < -90.0 {
		err = fmt.Errorf("minlatitude < -90.0: %f", e.MinLatitude)
		return e, err
	}

	if e.MaxLatitude != math.MaxFloat64 && e.MaxLatitude > 90.0 {
		err = fmt.Errorf("maxlatitude > 90.0: %f", e.MaxLatitude)
		return e, err
	}

	if e.MinLongitude != math.MaxFloat64 && e.MinLongitude < -180.0 {
		err = fmt.Errorf("minlongitude < -180.0: %f", e.MinLongitude)
		return e, err
	}

	if e.MaxLongitude != math.MaxFloat64 && e.MaxLongitude > 180.0 {
		err = fmt.Errorf("maxlongitude > 180.0: %f", e.MaxLongitude)
		return e, err
	}

	switch e.OrderBy {
	case "", "time", "time-asc", "magnitude", "magnitude-asc":
	default:
		err = fmt.Errorf("invalid option for orderby: %s", e.OrderBy)
	}

	return e, err
}

// query queries the DB for events matching e.
// The caller must close sql.Rows.
func (e *fdsnEventV1) query() (*sql.Rows, error) {
	q := "SELECT Quakeml12Event FROM fdsn.event WHERE deleted != true"

	qq, args := e.filter()

	if qq != "" {
		q = q + " AND " + qq
	}

	switch e.OrderBy {
	case "":
	case "time":
		q += " ORDER BY origintime desc"
	case "time-asc":
		q += " ORDER BY origintime asc"
	case "magnitude":
		q += " ORDER BY magnitude desc"
	case "magnitude-asc":
		q += " ORDER BY magnitude desc"
	}

	return db.Query(q, args...)
}

// query returns a count of events in the DB for e.
func (e *fdsnEventV1) count() (int, error) {
	q := "SELECT count(*) FROM fdsn.event WHERE deleted != true"

	qq, args := e.filter()

	if qq != "" {
		q = q + " AND " + qq
	}

	var c int
	err := db.QueryRow(q, args...).Scan(&c)

	return c, err
}

func (e *fdsnEventV1) filter() (q string, args []interface{}) {
	i := 1

	if e.PublicID != "" {
		q = fmt.Sprintf("%s publicid = $%d AND", q, i)
		args = append(args, e.PublicID)
		i++
	}

	if e.MinLatitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s latitude >= $%d AND", q, i)
		args = append(args, e.MinLatitude)
		i++
	}

	if e.MaxLatitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s latitude <= $%d AND", q, i)
		args = append(args, e.MaxLatitude)
		i++
	}

	if e.MinLongitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s longitude >= $%d AND", q, i)
		args = append(args, e.MinLongitude)
		i++
	}

	if e.MaxLongitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s longitude <= $%d AND", q, i)
		args = append(args, e.MaxLongitude)
		i++
	}

	if e.MinDepth != math.MaxFloat64 {
		q = fmt.Sprintf("%s depth > $%d AND", q, i)
		args = append(args, e.MinDepth)
		i++
	}

	if e.MaxDepth != math.MaxFloat64 {
		q = fmt.Sprintf("%s depth < $%d AND", q, i)
		args = append(args, e.MaxDepth)
		i++
	}

	if e.MinMagnitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s magnitude > $%d AND", q, i)
		args = append(args, e.MinMagnitude)
		i++
	}

	if e.MaxMagnitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s magnitude < $%d AND", q, i)
		args = append(args, e.MaxMagnitude)
		i++
	}

	if !e.StartTime.Time.IsZero() {
		q = fmt.Sprintf("%s origintime >= $%d AND", q, i)
		args = append(args, e.StartTime.Time)
		i++
	}

	if !e.EndTime.Time.IsZero() {
		q = fmt.Sprintf("%s origintime <= $%d AND", q, i)
		args = append(args, e.EndTime.Time)
		i++
	}

	q = strings.TrimSuffix(q, " AND")

	return
}

/*
eventV1Handler assembles QuakeML event fragments from the DB into a complete
QuakeML event.  The result set is limited to 10,000 events which will be ~1.2GB.
*/
func fdsnEventV1Handler(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	if r.Method != "GET" {
		return &weft.MethodNotAllowed
	}

	e, err := parseEventV1(r.URL.Query())
	if err != nil {
		return weft.BadRequest(err.Error())
	}

	c, err := e.count()
	if err != nil {
		return weft.ServiceUnavailableError(err)
	}

	if c > 10000 {
		return &weft.Result{
			Code: http.StatusRequestEntityTooLarge,
			Msg:  fmt.Sprintf("result to large found %d events, limit is 10,000", c),
		}
	}

	rows, err := e.query()
	if err != nil {
		return weft.ServiceUnavailableError(err)
	}
	defer rows.Close()

	b.WriteString(`<?xml version="1.0" encoding="UTF-8"?>
	<q:quakeml xmlns:q="http://quakeml.org/xmlns/quakeml/1.2" xmlns="http://quakeml.org/xmlns/bed/1.2">
	  <eventParameters publicID="smi:nz.org.geonet/NA">`)

	var xml string

	for rows.Next() {
		err = rows.Scan(&xml)
		if err != nil {
			return weft.ServiceUnavailableError(err)
		}

		b.WriteString(xml)
	}

	b.WriteString(`</eventParameters></q:quakeml>`)

	log.Printf("%s found %d events, result size %.1f (MB)", r.RequestURI, c, float64(b.Len())/1000000.0)

	h.Set("Content-Type", "application/xml")

	return &weft.StatusOK
}

func fdsnEventVersion(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "text/plain")
		b.WriteString("1.1")
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventContributors(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "application/xml")
		b.WriteString(`<Contributors><Contributor>WEL</Contributor></Contributors>`)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventCatalogs(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "application/xml")
		b.WriteString(`<Catalogs><Catalog>GeoNet</Catalog></Catalogs>`)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventWadl(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "application/xml")
		b.Write(fdsnEventWadlFile)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventV1Index(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "text/html")
		b.Write(fdsnEventIndex)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}
<commit_msg>gofmt.<commit_after>package main

import (
	"bytes"
	"database/sql"
	"errors"
	"fmt"
	"github.com/GeoNet/weft"
	"io/ioutil"
	"log"
	"math"
	"net/http"
	"net/url"
	"strings"
	"time"
)

// supported query parameters for the event service from http://www.fdsn.org/webservices/FDSN-WS-Specifications-1.1.pdf
type fdsnEventV1 struct {
	PublicID             string  `schema:"eventid"`      // select a specific event by ID; event identifiers are data center specific.
	MinLatitude          float64 `schema:"minlatitude"`  // limit to events with a latitude larger than or equal to the specified minimum.
	MaxLatitude          float64 `schema:"maxlatitude"`  // limit to events with a latitude smaller than or equal to the specified maximum.
	MinLongitude         float64 `schema:"minlongitude"` // limit to events with a longitude larger than or equal to the specified minimum.
	MaxLongitude         float64 `schema:"maxlongitude"` // limit to events with a longitude smaller than or equal to the specified maximum.
	MinDepth             float64 `schema:"mindepth"`     // limit to events with depth more than the specified minimum.
	MaxDepth             float64 `schema:"maxdepth"`     // limit to events with depth less than the specified maximum.
	MinMagnitude         float64 `schema:"minmagnitude"` // limit to events with a magnitude larger than the specified minimum.
	MaxMagnitude         float64 `schema:"maxmagnitude"` // limit to events with a magnitude smaller than the specified maximum.
	OrderBy              string  `schema:"orderby"`      // order the result by time or magnitude with the following possibilities: time, time-asc, magnitude, magnitude-asc
	StartTime            Time    `schema:"starttime"`    // limit to events on or after the specified start time.
	EndTime              Time    `schema:"endtime"`      // limit to events on or before the specified end time.
	IncludeAllOrigins    bool    `schema:"includeallorigins"`
	IncludeAllMagnitudes bool    `schema:"includeallmagnitudes"`
	IncludeArrivals      bool    `schema:"includearrivals"`
	Format               string  `schema:"format"`
}

type Time struct {
	time.Time
}

var fdsnEventWadlFile []byte
var fdsnEventIndex []byte
var eventNotSupported = map[string]bool{
	"latitude":      true,
	"longitude":     true,
	"minradius":     true,
	"maxraduis":     true,
	"magnitudetype": true,
	"limit":         true,
	"offset":        true,
	"catalog":       true,
	"contributor":   true,
	"updateafter":   true,
	"nodata":        true,
}

func init() {
	var err error
	fdsnEventWadlFile, err = ioutil.ReadFile("assets/fdsn-ws-event.wadl")
	if err != nil {
		log.Printf("error reading assets/fdsn-ws-event.wadl: %s", err.Error())
	}

	fdsnEventIndex, err = ioutil.ReadFile("assets/fdsn-ws-event.html")
	if err != nil {
		log.Printf("error reading assets/fdsn-ws-event.html: %s", err.Error())
	}
}

/*
parses the time in text as per the FDSN spec.  Pads text for parsing with
time.RFC3339Nano.  Accepted formats are (UTC):
   YYYY-MM-DDTHH:MM:SS.ssssss
   YYYY-MM-DDTHH:MM:SS
   YYYY-MM-DD

Implements the encoding.TextUnmarshaler interface.
*/
func (t *Time) UnmarshalText(text []byte) (err error) {
	s := string(text)
	l := len(s)
	if len(s) < 10 {
		return fmt.Errorf("invalid time format: %s", s)
	}

	if l >= 19 && l <= 26 && l != 20 { // length 20: "YYYY-MM-DDTHH:MM:SS." invalid
		s = s + ".000000000Z"[(l-19):] // "YYYY-MM-DDTHH:MM:SS" append to nano
	} else if l == 10 {
		s = s + "T00:00:00.000000000Z" // YYYY-MM-DD
	} else {
		return fmt.Errorf("invalid time format: %s", s)
	}
	t.Time, err = time.Parse(time.RFC3339Nano, s)
	return
}

func parseEventV1(v url.Values) (fdsnEventV1, error) {
	// All query parameters are optional and float zero values overlap
	// with possible request ranges so the default is set to the max float val.
	e := fdsnEventV1{
		MinLatitude:  math.MaxFloat64,
		MaxLatitude:  math.MaxFloat64,
		MinLongitude: math.MaxFloat64,
		MaxLongitude: math.MaxFloat64,
		MinDepth:     math.MaxFloat64,
		MaxDepth:     math.MaxFloat64,
		MinMagnitude: math.MaxFloat64,
		MaxMagnitude: math.MaxFloat64,
	}

	for key, val := range v {
		if _, ok := eventNotSupported[key]; ok {
			return e, fmt.Errorf("\"%s\" is not supported", key)
		}
		if len(val[0]) == 0 {
			return e, fmt.Errorf("Invalid %s value", key)
		}
	}

	err := decoder.Decode(&e, v)
	if err != nil {
		return e, err
	}

	if e.IncludeAllMagnitudes {
		return e, errors.New("include all magnitudes is not supported.")
	}

	if e.IncludeAllOrigins {
		return e, errors.New("include all origins is not supported.")
	}

	if e.IncludeArrivals {
		return e, errors.New("include arrivals is not supported.")
	}

	// geometry bounds checking
	if e.MinLatitude != math.MaxFloat64 && e.MinLatitude < -90.0 {
		err = fmt.Errorf("minlatitude < -90.0: %f", e.MinLatitude)
		return e, err
	}

	if e.MaxLatitude != math.MaxFloat64 && e.MaxLatitude > 90.0 {
		err = fmt.Errorf("maxlatitude > 90.0: %f", e.MaxLatitude)
		return e, err
	}

	if e.MinLongitude != math.MaxFloat64 && e.MinLongitude < -180.0 {
		err = fmt.Errorf("minlongitude < -180.0: %f", e.MinLongitude)
		return e, err
	}

	if e.MaxLongitude != math.MaxFloat64 && e.MaxLongitude > 180.0 {
		err = fmt.Errorf("maxlongitude > 180.0: %f", e.MaxLongitude)
		return e, err
	}

	switch e.OrderBy {
	case "", "time", "time-asc", "magnitude", "magnitude-asc":
	default:
		err = fmt.Errorf("invalid option for orderby: %s", e.OrderBy)
	}

	return e, err
}

// query queries the DB for events matching e.
// The caller must close sql.Rows.
func (e *fdsnEventV1) query() (*sql.Rows, error) {
	q := "SELECT Quakeml12Event FROM fdsn.event WHERE deleted != true"

	qq, args := e.filter()

	if qq != "" {
		q = q + " AND " + qq
	}

	switch e.OrderBy {
	case "":
	case "time":
		q += " ORDER BY origintime desc"
	case "time-asc":
		q += " ORDER BY origintime asc"
	case "magnitude":
		q += " ORDER BY magnitude desc"
	case "magnitude-asc":
		q += " ORDER BY magnitude desc"
	}

	return db.Query(q, args...)
}

// query returns a count of events in the DB for e.
func (e *fdsnEventV1) count() (int, error) {
	q := "SELECT count(*) FROM fdsn.event WHERE deleted != true"

	qq, args := e.filter()

	if qq != "" {
		q = q + " AND " + qq
	}

	var c int
	err := db.QueryRow(q, args...).Scan(&c)

	return c, err
}

func (e *fdsnEventV1) filter() (q string, args []interface{}) {
	i := 1

	if e.PublicID != "" {
		q = fmt.Sprintf("%s publicid = $%d AND", q, i)
		args = append(args, e.PublicID)
		i++
	}

	if e.MinLatitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s latitude >= $%d AND", q, i)
		args = append(args, e.MinLatitude)
		i++
	}

	if e.MaxLatitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s latitude <= $%d AND", q, i)
		args = append(args, e.MaxLatitude)
		i++
	}

	if e.MinLongitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s longitude >= $%d AND", q, i)
		args = append(args, e.MinLongitude)
		i++
	}

	if e.MaxLongitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s longitude <= $%d AND", q, i)
		args = append(args, e.MaxLongitude)
		i++
	}

	if e.MinDepth != math.MaxFloat64 {
		q = fmt.Sprintf("%s depth > $%d AND", q, i)
		args = append(args, e.MinDepth)
		i++
	}

	if e.MaxDepth != math.MaxFloat64 {
		q = fmt.Sprintf("%s depth < $%d AND", q, i)
		args = append(args, e.MaxDepth)
		i++
	}

	if e.MinMagnitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s magnitude > $%d AND", q, i)
		args = append(args, e.MinMagnitude)
		i++
	}

	if e.MaxMagnitude != math.MaxFloat64 {
		q = fmt.Sprintf("%s magnitude < $%d AND", q, i)
		args = append(args, e.MaxMagnitude)
		i++
	}

	if !e.StartTime.Time.IsZero() {
		q = fmt.Sprintf("%s origintime >= $%d AND", q, i)
		args = append(args, e.StartTime.Time)
		i++
	}

	if !e.EndTime.Time.IsZero() {
		q = fmt.Sprintf("%s origintime <= $%d AND", q, i)
		args = append(args, e.EndTime.Time)
		i++
	}

	q = strings.TrimSuffix(q, " AND")

	return
}

/*
eventV1Handler assembles QuakeML event fragments from the DB into a complete
QuakeML event.  The result set is limited to 10,000 events which will be ~1.2GB.
*/
func fdsnEventV1Handler(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	if r.Method != "GET" {
		return &weft.MethodNotAllowed
	}

	e, err := parseEventV1(r.URL.Query())
	if err != nil {
		return weft.BadRequest(err.Error())
	}

	c, err := e.count()
	if err != nil {
		return weft.ServiceUnavailableError(err)
	}

	if c > 10000 {
		return &weft.Result{
			Code: http.StatusRequestEntityTooLarge,
			Msg:  fmt.Sprintf("result to large found %d events, limit is 10,000", c),
		}
	}

	rows, err := e.query()
	if err != nil {
		return weft.ServiceUnavailableError(err)
	}
	defer rows.Close()

	b.WriteString(`<?xml version="1.0" encoding="UTF-8"?>
	<q:quakeml xmlns:q="http://quakeml.org/xmlns/quakeml/1.2" xmlns="http://quakeml.org/xmlns/bed/1.2">
	  <eventParameters publicID="smi:nz.org.geonet/NA">`)

	var xml string

	for rows.Next() {
		err = rows.Scan(&xml)
		if err != nil {
			return weft.ServiceUnavailableError(err)
		}

		b.WriteString(xml)
	}

	b.WriteString(`</eventParameters></q:quakeml>`)

	log.Printf("%s found %d events, result size %.1f (MB)", r.RequestURI, c, float64(b.Len())/1000000.0)

	h.Set("Content-Type", "application/xml")

	return &weft.StatusOK
}

func fdsnEventVersion(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "text/plain")
		b.WriteString("1.1")
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventContributors(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "application/xml")
		b.WriteString(`<Contributors><Contributor>WEL</Contributor></Contributors>`)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventCatalogs(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "application/xml")
		b.WriteString(`<Catalogs><Catalog>GeoNet</Catalog></Catalogs>`)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventWadl(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "application/xml")
		b.Write(fdsnEventWadlFile)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}

func fdsnEventV1Index(r *http.Request, h http.Header, b *bytes.Buffer) *weft.Result {
	switch r.Method {
	case "GET":
		if res := weft.CheckQuery(r, []string{}, []string{}); !res.Ok {
			return res
		}

		h.Set("Content-Type", "text/html")
		b.Write(fdsnEventIndex)
		return &weft.StatusOK
	default:
		return &weft.MethodNotAllowed
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/labstack/echo"
	"github.com/monochromegane/gannoy"
)

var (
	dataDir string
)

func init() {
	flag.StringVar(&dataDir, "d", ".", "Data directory.")
	flag.Parse()
}

type Feature struct {
	W []float64 `json:"features"`
}

func main() {
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	databases := map[string]gannoy.GannoyIndex{}
	for _, file := range files {
		if file.IsDir() || filepath.Ext(file.Name()) != ".meta" {
			continue
		}
		key := strings.TrimSuffix(file.Name(), ".meta")
		gannoy, err := gannoy.NewGannoyIndex("hoge.meta", gannoy.Angular{}, gannoy.RandRandom{})
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		databases[key] = gannoy
	}

	e := echo.New()
	e.GET("/search", func(c echo.Context) error {
		database := c.QueryParam("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusNotFound)
		}
		id, err := strconv.Atoi(c.QueryParam("id"))
		if err != nil {
			id = -1
		}
		limit, err := strconv.Atoi(c.QueryParam("limit"))
		if err != nil {
			limit = 10
		}

		gannoy := databases[database]
		r, err := gannoy.GetNnsByKey(id, limit, -1)
		if err != nil || len(r) == 0 {
			return c.NoContent(http.StatusNotFound)
		}

		return c.JSON(http.StatusOK, r)
	})

	e.PUT("/databases/:database/features/:id", func(c echo.Context) error {
		database := c.Param("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		id, err := strconv.Atoi(c.Param("id"))
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		feature := new(Feature)
		if err := c.Bind(feature); err != nil {
			return err
		}

		gannoy := databases[database]
		err = gannoy.AddItem(id, feature.W)
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		return c.NoContent(http.StatusOK)
	})

	e.DELETE("/databases/:database/features/:id", func(c echo.Context) error {
		database := c.Param("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		id, err := strconv.Atoi(c.Param("id"))
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		gannoy := databases[database]
		err = gannoy.RemoveItem(id)
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}

		return c.NoContent(http.StatusOK)
	})

	e.Start(":1323")
}
<commit_msg>Use key instead of id in gannoy-server.<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/labstack/echo"
	"github.com/monochromegane/gannoy"
)

var (
	dataDir string
)

func init() {
	flag.StringVar(&dataDir, "d", ".", "Data directory.")
	flag.Parse()
}

type Feature struct {
	W []float64 `json:"features"`
}

func main() {
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	databases := map[string]gannoy.GannoyIndex{}
	for _, file := range files {
		if file.IsDir() || filepath.Ext(file.Name()) != ".meta" {
			continue
		}
		key := strings.TrimSuffix(file.Name(), ".meta")
		gannoy, err := gannoy.NewGannoyIndex("hoge.meta", gannoy.Angular{}, gannoy.RandRandom{})
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		databases[key] = gannoy
	}

	e := echo.New()
	e.GET("/search", func(c echo.Context) error {
		database := c.QueryParam("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusNotFound)
		}
		key, err := strconv.Atoi(c.QueryParam("key"))
		if err != nil {
			key = -1
		}
		limit, err := strconv.Atoi(c.QueryParam("limit"))
		if err != nil {
			limit = 10
		}

		gannoy := databases[database]
		r, err := gannoy.GetNnsByKey(key, limit, -1)
		if err != nil || len(r) == 0 {
			return c.NoContent(http.StatusNotFound)
		}

		return c.JSON(http.StatusOK, r)
	})

	e.PUT("/databases/:database/features/:key", func(c echo.Context) error {
		database := c.Param("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		key, err := strconv.Atoi(c.Param("key"))
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		feature := new(Feature)
		if err := c.Bind(feature); err != nil {
			return err
		}

		gannoy := databases[database]
		err = gannoy.AddItem(key, feature.W)
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		return c.NoContent(http.StatusOK)
	})

	e.DELETE("/databases/:database/features/:key", func(c echo.Context) error {
		database := c.Param("database")
		if _, ok := databases[database]; !ok {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		key, err := strconv.Atoi(c.Param("key"))
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}
		gannoy := databases[database]
		err = gannoy.RemoveItem(key)
		if err != nil {
			return c.NoContent(http.StatusUnprocessableEntity)
		}

		return c.NoContent(http.StatusOK)
	})

	e.Start(":1323")
}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/pipestream"
)

const (
	reassignNodeFilename = "reassignment-node.json"
)

type Migrate struct {
	Ui  cli.Ui
	Cmd string

	zkcluster  *zk.ZkCluster
	zone       string
	cluster    string
	topic      string
	brokerId   string
	partition  string
	verifyMode bool
}

func (this *Migrate) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("migrate", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", "", "")
	cmdFlags.StringVar(&this.cluster, "c", "", "")
	cmdFlags.StringVar(&this.topic, "t", "", "")
	cmdFlags.StringVar(&this.partition, "p", "", "")
	cmdFlags.StringVar(&this.brokerId, "brokers", "", "")
	cmdFlags.BoolVar(&this.verifyMode, "verify", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-z", "-c", "-t").
		invalid(args) {
		return 2
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	this.zkcluster = zkzone.NewCluster(this.cluster)

	if this.verifyMode {
		this.Ui.Info(fmt.Sprintf("You MUST manually remove the %s after migration is done.", reassignNodeFilename))
		this.Ui.Info(fmt.Sprintf("After verify ok, modify producer/consumer to point to new brokers!"))

		for {
			this.Ui.Output(fmt.Sprintf("%s", time.Now().String()))
			this.verify()

			time.Sleep(time.Second * 2)
		}

		return
	}

	if validateArgs(this, this.Ui).
		require("-z", "-c", "-t", "-p", "-brokers").
		requireAdminRights("-z").
		invalid(args) {
		return 2
	}

	//this.ensureBrokersAreAlive()
	data := this.generateReassignFile()
	this.Ui.Output(data)
	yes, _ := this.Ui.Ask("Are you sure to execute the migration? [Y/N]")
	if yes == "Y" {
		this.executeReassignment()
	} else {
		this.Ui.Output("bye")
	}

	return
}

func (this *Migrate) ensureBrokersAreAlive() {
	brokerIds := strings.Split(this.brokerId, ",")
	liveBrokers := this.zkcluster.Brokers()
	for _, id := range brokerIds {
		if _, present := liveBrokers[id]; !present {
			this.Ui.Error(fmt.Sprintf("broker:%s not alive", id))
			os.Exit(1)
		}
	}

}

// generate reassignment-node.json
func (this *Migrate) generateReassignFile() string {
	// {"version":1,"partitions":[{"topic":"fortest1","partition":0,"replicas":[3,4]}

	type PartitionMeta struct {
		Topic     string `json:"topic"`
		Partition int    `json:"partition"`
		Replicas  []int  `json:"replicas"`
	}
	type ReassignMeta struct {
		Version    int             `json:"version"`
		Partitions []PartitionMeta `json:"partitions"`
	}

	var js ReassignMeta
	js.Version = 1
	js.Partitions = make([]PartitionMeta, 0)
	for _, p := range strings.Split(this.partition, ",") {
		p = strings.TrimSpace(p)
		pid, err := strconv.Atoi(p)
		swallow(err)

		pmeta := PartitionMeta{
			Topic:     this.topic,
			Partition: pid,
			Replicas:  make([]int, 0),
		}
		for _, b := range strings.Split(this.brokerId, ",") {
			b = strings.TrimSpace(b)
			bid, err := strconv.Atoi(b)
			swallow(err)

			pmeta.Replicas = append(pmeta.Replicas, bid)
		}

		js.Partitions = append(js.Partitions, pmeta)
	}

	b, err := json.Marshal(js)
	swallow(err)
	swallow(ioutil.WriteFile(reassignNodeFilename, b, 0644))
	return string(b)
}

func (this *Migrate) executeReassignment() {
	/*
		1. kafka-reassign-partitions.sh write /admin/reassign_partitions
		2. controller listens to the path above
		3. For each topic partition, the controller does the following:
		  3.1. Start new replicas in RAR – AR (RAR = Reassigned Replicas, AR = original list of Assigned Replicas)
		  3.2. Wait until new replicas are in sync with the leader
		  3.3. If the leader is not in RAR, elect a new leader from RAR
		  3.4 4. Stop old replicas AR – RAR
		  3.5. Write new AR
		  3.6. Remove partition from the /admin/reassign_partitions path

	*/
	cmd := pipestream.New(fmt.Sprintf("%s/bin/kafka-reassign-partitions.sh", ctx.KafkaHome()),
		fmt.Sprintf("--zookeeper %s", this.zkcluster.ZkConnectAddr()),
		fmt.Sprintf("--reassignment-json-file %s", reassignNodeFilename),
		fmt.Sprintf("--execute"),
	)
	err := cmd.Open()
	if err != nil {
		return
	}
	defer cmd.Close()

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		this.Ui.Output(color.Yellow(scanner.Text()))
	}
}

func (this *Migrate) verify() {
	cmd := pipestream.New(fmt.Sprintf("%s/bin/kafka-reassign-partitions.sh", ctx.KafkaHome()),
		fmt.Sprintf("--zookeeper %s", this.zkcluster.ZkConnectAddr()),
		fmt.Sprintf("--reassignment-json-file %s", reassignNodeFilename),
		fmt.Sprintf("--verify"),
	)
	err := cmd.Open()
	if err != nil {
		return
	}
	defer cmd.Close()

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		this.Ui.Output(color.Yellow(scanner.Text()))
	}

}

func (*Migrate) Synopsis() string {
	return "Migrate given topic partition to specified broker ids"
}

func (this *Migrate) Help() string {
	help := fmt.Sprintf(`
Usage: %s migrate -z zone -c cluster [options]

    Migrate given topic partition to specified broker ids.

    e,g. migrate partition 0 of order to broker 2 as master and 3 as replica, then verify
      gk migrate -z prod -c trade -t order -p 0 -brokers 2,3
      gk migrate -z prod -c trade -t order -p 0 -brokers 2,3 -verify

Options:

    -t topic

    -p partitionId
      Multiple partition ids seperated by comma.
      e,g. -p 0,1

    -brokers id1,id2,idN
      Migrate the topic to given broker ids.

    -verify
      Verify the migration ongoing.

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>notes about kafka preferred leader<commit_after>package command

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/pipestream"
)

const (
	reassignNodeFilename = "reassignment-node.json"
)

type Migrate struct {
	Ui  cli.Ui
	Cmd string

	zkcluster  *zk.ZkCluster
	zone       string
	cluster    string
	topic      string
	brokerId   string
	partition  string
	verifyMode bool
}

func (this *Migrate) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("migrate", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", "", "")
	cmdFlags.StringVar(&this.cluster, "c", "", "")
	cmdFlags.StringVar(&this.topic, "t", "", "")
	cmdFlags.StringVar(&this.partition, "p", "", "")
	cmdFlags.StringVar(&this.brokerId, "brokers", "", "")
	cmdFlags.BoolVar(&this.verifyMode, "verify", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-z", "-c", "-t").
		invalid(args) {
		return 2
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	this.zkcluster = zkzone.NewCluster(this.cluster)

	if this.verifyMode {
		this.Ui.Info(fmt.Sprintf("You MUST manually remove the %s after migration is done.", reassignNodeFilename))
		this.Ui.Info(fmt.Sprintf("After verify ok, modify producer/consumer to point to new brokers!"))

		for {
			this.Ui.Output(fmt.Sprintf("%s", time.Now().String()))
			this.verify()

			time.Sleep(time.Second * 2)
		}

		return
	}

	if validateArgs(this, this.Ui).
		require("-z", "-c", "-t", "-p", "-brokers").
		requireAdminRights("-z").
		invalid(args) {
		return 2
	}

	//this.ensureBrokersAreAlive()
	data := this.generateReassignFile()
	this.Ui.Output(data)
	yes, _ := this.Ui.Ask("Are you sure to execute the migration? [Y/N]")
	if yes == "Y" {
		this.executeReassignment()
	} else {
		this.Ui.Output("bye")
	}

	return
}

func (this *Migrate) ensureBrokersAreAlive() {
	brokerIds := strings.Split(this.brokerId, ",")
	liveBrokers := this.zkcluster.Brokers()
	for _, id := range brokerIds {
		if _, present := liveBrokers[id]; !present {
			this.Ui.Error(fmt.Sprintf("broker:%s not alive", id))
			os.Exit(1)
		}
	}

}

// generate reassignment-node.json
func (this *Migrate) generateReassignFile() string {
	// {"version":1,"partitions":[{"topic":"fortest1","partition":0,"replicas":[3,4]}

	type PartitionMeta struct {
		Topic     string `json:"topic"`
		Partition int    `json:"partition"`
		Replicas  []int  `json:"replicas"`
	}
	type ReassignMeta struct {
		Version    int             `json:"version"`
		Partitions []PartitionMeta `json:"partitions"`
	}

	var js ReassignMeta
	js.Version = 1
	js.Partitions = make([]PartitionMeta, 0)
	for _, p := range strings.Split(this.partition, ",") {
		p = strings.TrimSpace(p)
		pid, err := strconv.Atoi(p)
		swallow(err)

		pmeta := PartitionMeta{
			Topic:     this.topic,
			Partition: pid,
			Replicas:  make([]int, 0),
		}
		for _, b := range strings.Split(this.brokerId, ",") {
			b = strings.TrimSpace(b)
			bid, err := strconv.Atoi(b)
			swallow(err)

			pmeta.Replicas = append(pmeta.Replicas, bid)
		}

		js.Partitions = append(js.Partitions, pmeta)
	}

	b, err := json.Marshal(js)
	swallow(err)
	swallow(ioutil.WriteFile(reassignNodeFilename, b, 0644))
	return string(b)
}

func (this *Migrate) executeReassignment() {
	/*
		1. kafka-reassign-partitions.sh write /admin/reassign_partitions
		2. controller listens to the path above
		3. For each topic partition, the controller does the following:
		  3.1. Start new replicas in RAR – AR (RAR = Reassigned Replicas, AR = original list of Assigned Replicas)
		  3.2. Wait until new replicas are in sync with the leader
		  3.3. If the leader is not in RAR, elect a new leader from RAR
		  3.4 4. Stop old replicas AR – RAR
		  3.5. Write new AR
		  3.6. Remove partition from the /admin/reassign_partitions path

	*/
	cmd := pipestream.New(fmt.Sprintf("%s/bin/kafka-reassign-partitions.sh", ctx.KafkaHome()),
		fmt.Sprintf("--zookeeper %s", this.zkcluster.ZkConnectAddr()),
		fmt.Sprintf("--reassignment-json-file %s", reassignNodeFilename),
		fmt.Sprintf("--execute"),
	)
	err := cmd.Open()
	if err != nil {
		return
	}
	defer cmd.Close()

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		this.Ui.Output(color.Yellow(scanner.Text()))
	}
}

func (this *Migrate) verify() {
	cmd := pipestream.New(fmt.Sprintf("%s/bin/kafka-reassign-partitions.sh", ctx.KafkaHome()),
		fmt.Sprintf("--zookeeper %s", this.zkcluster.ZkConnectAddr()),
		fmt.Sprintf("--reassignment-json-file %s", reassignNodeFilename),
		fmt.Sprintf("--verify"),
	)
	err := cmd.Open()
	if err != nil {
		return
	}
	defer cmd.Close()

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		this.Ui.Output(color.Yellow(scanner.Text()))
	}

}

func (*Migrate) Synopsis() string {
	return "Migrate given topic partition to specified broker ids"
}

func (this *Migrate) Help() string {
	help := fmt.Sprintf(`
Usage: %s migrate -z zone -c cluster [options]

    Migrate given topic partition to specified broker ids.

    e,g. migrate partition 0 of order to broker 2 as master and 3 as replica, then verify
      gk migrate -z prod -c trade -t order -p 0 -brokers 2,3
      gk migrate -z prod -c trade -t order -p 0 -brokers 2,3 -verify

Options:

    -t topic

    -p partitionId
      Multiple partition ids seperated by comma.
      e,g. -p 0,1

    -brokers id1,id2,idN
      Migrate the topic to given broker ids.
      brokers order is IMPORTANT! The 1st is the preferred leader.

    -verify
      Verify the migration ongoing.

`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package main

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju/testing"
)

type MachineSuite struct {
	testing.JujuConnSuite
}

var _ = Suite(&MachineSuite{})

func (s *MachineSuite) TestParseSuccess(c *C) {
	create := func() (cmd.Command, *AgentConf) {
		a := &MachineAgent{}
		return a, &a.Conf
	}
	a := CheckAgentCommand(c, create, []string{"--machine-id", "42"})
	c.Assert(a.(*MachineAgent).MachineId, Equals, 42)
}

func (s *MachineSuite) TestParseNonsense(c *C) {
	for _, args := range [][]string{
		[]string{},
		[]string{"--machine-id", "-4004"},
	} {
		err := ParseAgentCommand(&MachineAgent{}, args)
		c.Assert(err, ErrorMatches, "--machine-id option must be set, and expects a non-negative integer")
	}
}

func (s *MachineSuite) TestParseUnknown(c *C) {
	a := &MachineAgent{}
	err := ParseAgentCommand(a, []string{"--machine-id", "42", "blistering barnacles"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["blistering barnacles"\]`)
}

func (s *MachineSuite) TestRunInvalidMachineId(c *C) {
	c.Skip("agents don't yet distinguish between temporary and permanent errors")
	a := &MachineAgent{
		Conf: AgentConf{
			JujuDir:   environs.VarDir,
			StateInfo: *s.StateInfo(c),
		},
		MachineId: 2,
	}
	err := a.Run(nil)
	c.Assert(err, ErrorMatches, "some error")
}

func (s *MachineSuite) TestRunStop(c *C) {
	m, err := s.State.AddMachine()
	c.Assert(err, IsNil)
	a := &MachineAgent{
		Conf: AgentConf{
			JujuDir:   environs.VarDir,
			StateInfo: *s.StateInfo(c),
		},
		MachineId: m.Id(),
	}
	done := make(chan error)
	go func() {
		done <- a.Run(nil)
	}()
}
<commit_msg>cmd/jujud: add MachineSuite.TestRunStop<commit_after>package main

import (
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju/testing"
)

type MachineSuite struct {
	testing.JujuConnSuite
}

var _ = Suite(&MachineSuite{})

func (s *MachineSuite) TestParseSuccess(c *C) {
	create := func() (cmd.Command, *AgentConf) {
		a := &MachineAgent{}
		return a, &a.Conf
	}
	a := CheckAgentCommand(c, create, []string{"--machine-id", "42"})
	c.Assert(a.(*MachineAgent).MachineId, Equals, 42)
}

func (s *MachineSuite) TestParseNonsense(c *C) {
	for _, args := range [][]string{
		[]string{},
		[]string{"--machine-id", "-4004"},
	} {
		err := ParseAgentCommand(&MachineAgent{}, args)
		c.Assert(err, ErrorMatches, "--machine-id option must be set, and expects a non-negative integer")
	}
}

func (s *MachineSuite) TestParseUnknown(c *C) {
	a := &MachineAgent{}
	err := ParseAgentCommand(a, []string{"--machine-id", "42", "blistering barnacles"})
	c.Assert(err, ErrorMatches, `unrecognized args: \["blistering barnacles"\]`)
}

func (s *MachineSuite) TestRunInvalidMachineId(c *C) {
	c.Skip("agents don't yet distinguish between temporary and permanent errors")
	a := &MachineAgent{
		Conf: AgentConf{
			JujuDir:   environs.VarDir,
			StateInfo: *s.StateInfo(c),
		},
		MachineId: 2,
	}
	err := a.Run(nil)
	c.Assert(err, ErrorMatches, "some error")
}

func (s *MachineSuite) TestRunStop(c *C) {
	m, err := s.State.AddMachine()
	c.Assert(err, IsNil)
	a := &MachineAgent{
		Conf: AgentConf{
			JujuDir:   environs.VarDir,
			StateInfo: *s.StateInfo(c),
		},
		MachineId: m.Id(),
	}
	done := make(chan error)
	go func() {
		done <- a.Run(nil)
	}()
	err = a.Stop()
	c.Assert(err, IsNil)
	c.Assert(<-done, IsNil)
}

<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/environs"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju-core/juju/environs/dummy"
	_ "launchpad.net/juju-core/juju/environs/ec2"
)

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	// TODO(dfc) place the logic in a loop with a suitable delay
	p, err := NewProvisioner(&a.Conf.StateInfo)
	if err != nil {
		return err
	}
	return p.Wait()
}

type Provisioner struct {
	st      *state.State
	info    *state.Info
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher

	// machine.Id => environs.Instance
	instances map[int]environs.Instance
	// instance.Id => *state.Machine
	machines map[string]*state.Machine
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(info *state.Info) (*Provisioner, error) {
	st, err := state.Open(info)
	if err != nil {
		return nil, err
	}
	p := &Provisioner{
		st:        st,
		info:      info,
		instances: make(map[int]environs.Instance),
		machines:  make(map[string]*state.Machine),
	}
	go p.loop()
	return p, nil
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	defer p.st.Close()
	p.environWatcher = p.st.WatchEnvironConfig()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			log.Printf("provisioner loaded new environment configuration")

			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	// call processMachines to stop any unknown instances before watching machines.
	if err := p.processMachines(nil, nil); err != nil {
		p.tomb.Kill(err)
	}
	p.machinesWatcher = p.st.WatchMachines()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner loaded new environment configuration")
		case machines, ok := <-p.machinesWatcher.Changes():
			if !ok {
				err := p.machinesWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			// TODO(dfc) fire process machines periodically to shut down unknown
			// instances.
			if err := p.processMachines(machines.Added, machines.Deleted); err != nil {
				p.tomb.Kill(err)
			}
		}
	}
}

// Wait waits for the Provisioner to exit.
func (p *Provisioner) Wait() error {
	return p.tomb.Wait()
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

func (p *Provisioner) processMachines(added, removed []*state.Machine) error {
	// step 1. find which of the added machines have not
	// yet been allocated a started instance.
	notstarted, err := p.findNotStarted(added)
	if err != nil {
		return err
	}

	// step 2. start an instance for any machines we found.
	if err := p.startMachines(notstarted); err != nil {
		return err
	}

	// step 3. stop all machines that were removed from the state.
	stopping, err := p.instancesForMachines(removed)
	if err != nil {
		return err
	}

	// step 4. find instances which are running but have no machine
	// associated with them.
	unknown, err := p.findUnknownInstances()

	return p.stopInstances(append(stopping, unknown...))
}

// findUnknownInstances finds instances which are not associated with a machine.
func (p *Provisioner) findUnknownInstances() ([]environs.Instance, error) {
	all, err := p.environ.AllInstances()
	if err != nil {
		return nil, err
	}
	instances := make(map[string]environs.Instance)
	for _, i := range all {
		instances[i.Id()] = i
	}
	// TODO(dfc) this is very inefficient, p.machines cache may help.
	machines, err := p.st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			continue
		}
		delete(instances, id)
	}
	var unknown []environs.Instance
	for _, i := range instances {
		unknown = append(unknown, i)
	}
	return unknown, nil
}

// findNotStarted finds machines without an InstanceId set, these are defined as not started.
func (p *Provisioner) findNotStarted(machines []*state.Machine) ([]*state.Machine, error) {
	var notstarted []*state.Machine
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			notstarted = append(notstarted, m)
		} else {
			log.Printf("machine %s already started as instance %q", m, id)
		}
	}
	return notstarted, nil
}

func (p *Provisioner) startMachines(machines []*state.Machine) error {
	for _, m := range machines {
		if err := p.startMachine(m); err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) startMachine(m *state.Machine) error {
	// TODO(dfc) the state.Info passed to environ.StartInstance remains contentious
	// however as the PA only knows one state.Info, and that info is used by MAs and
	// UAs to locate the ZK for this environment, it is logical to use the same
	// state.Info as the PA.
	inst, err := p.environ.StartInstance(m.Id(), p.info)
	if err != nil {
		log.Printf("provisioner can't start machine %s: %v", m, err)
		return err
	}

	// assign the instance id to the machine
	if err := m.SetInstanceId(inst.Id()); err != nil {
		return err
	}

	// populate the local cache
	p.instances[m.Id()] = inst
	p.machines[inst.Id()] = m
	log.Printf("provisioner started machine %s as instance %s", m, inst.Id())
	return nil
}

func (p *Provisioner) stopInstances(instances []environs.Instance) error {
	// Although calling StopInstance with an empty slice should produce no change in the
	// provider, environs like dummy do not consider this a noop.
	if len(instances) == 0 {
		return nil
	}
	if err := p.environ.StopInstances(instances); err != nil {
		return err
	}

	// cleanup cache
	for _, i := range instances {
		if m, ok := p.machines[i.Id()]; ok {
			delete(p.machines, i.Id())
			delete(p.instances, m.Id())
		}
	}
	return nil
}

// instanceForMachine returns the environs.Instance that represents this machine's instance.
func (p *Provisioner) instanceForMachine(m *state.Machine) (environs.Instance, error) {
	inst, ok := p.instances[m.Id()]
	if !ok {
		// not cached locally, ask the environ.
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			return nil, fmt.Errorf("machine %s not found", m)
		}
		// TODO(dfc) this should be batched, or the cache preloaded at startup to
		// avoid N calls to the envirion.
		insts, err := p.environ.Instances([]string{id})
		if err != nil {
			// the provider doesn't know about this instance, give up.
			return nil, err
		}
		inst = insts[0]
	}
	return inst, nil
}

// instancesForMachines returns a list of environs.Instance that represent the list of machines running
// in the provider.
func (p *Provisioner) instancesForMachines(machines []*state.Machine) ([]environs.Instance, error) {
	var insts []environs.Instance
	for _, m := range machines {
		inst, err := p.instanceForMachine(m)
		if err != nil {
			return nil, err
		}
		insts = append(insts, inst)
	}
	return insts, nil
}
<commit_msg>cmd/jujud: remove dummy environ import from provisioning agent<commit_after>package main

import (
	"fmt"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/environs"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju-core/juju/environs/ec2"
)

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	// TODO(dfc) place the logic in a loop with a suitable delay
	p, err := NewProvisioner(&a.Conf.StateInfo)
	if err != nil {
		return err
	}
	return p.Wait()
}

type Provisioner struct {
	st      *state.State
	info    *state.Info
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher

	// machine.Id => environs.Instance
	instances map[int]environs.Instance
	// instance.Id => *state.Machine
	machines map[string]*state.Machine
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(info *state.Info) (*Provisioner, error) {
	st, err := state.Open(info)
	if err != nil {
		return nil, err
	}
	p := &Provisioner{
		st:        st,
		info:      info,
		instances: make(map[int]environs.Instance),
		machines:  make(map[string]*state.Machine),
	}
	go p.loop()
	return p, nil
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	defer p.st.Close()
	p.environWatcher = p.st.WatchEnvironConfig()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			log.Printf("provisioner loaded new environment configuration")

			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	// call processMachines to stop any unknown instances before watching machines.
	if err := p.processMachines(nil, nil); err != nil {
		p.tomb.Kill(err)
	}
	p.machinesWatcher = p.st.WatchMachines()
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environWatcher.Changes():
			if !ok {
				err := p.environWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner loaded invalid environment configuration: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner loaded new environment configuration")
		case machines, ok := <-p.machinesWatcher.Changes():
			if !ok {
				err := p.machinesWatcher.Stop()
				if err != nil {
					p.tomb.Kill(err)
				}
				return
			}
			// TODO(dfc) fire process machines periodically to shut down unknown
			// instances.
			if err := p.processMachines(machines.Added, machines.Deleted); err != nil {
				p.tomb.Kill(err)
			}
		}
	}
}

// Wait waits for the Provisioner to exit.
func (p *Provisioner) Wait() error {
	return p.tomb.Wait()
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	return p.tomb.Wait()
}

func (p *Provisioner) processMachines(added, removed []*state.Machine) error {
	// step 1. find which of the added machines have not
	// yet been allocated a started instance.
	notstarted, err := p.findNotStarted(added)
	if err != nil {
		return err
	}

	// step 2. start an instance for any machines we found.
	if err := p.startMachines(notstarted); err != nil {
		return err
	}

	// step 3. stop all machines that were removed from the state.
	stopping, err := p.instancesForMachines(removed)
	if err != nil {
		return err
	}

	// step 4. find instances which are running but have no machine
	// associated with them.
	unknown, err := p.findUnknownInstances()

	return p.stopInstances(append(stopping, unknown...))
}

// findUnknownInstances finds instances which are not associated with a machine.
func (p *Provisioner) findUnknownInstances() ([]environs.Instance, error) {
	all, err := p.environ.AllInstances()
	if err != nil {
		return nil, err
	}
	instances := make(map[string]environs.Instance)
	for _, i := range all {
		instances[i.Id()] = i
	}
	// TODO(dfc) this is very inefficient, p.machines cache may help.
	machines, err := p.st.AllMachines()
	if err != nil {
		return nil, err
	}
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			continue
		}
		delete(instances, id)
	}
	var unknown []environs.Instance
	for _, i := range instances {
		unknown = append(unknown, i)
	}
	return unknown, nil
}

// findNotStarted finds machines without an InstanceId set, these are defined as not started.
func (p *Provisioner) findNotStarted(machines []*state.Machine) ([]*state.Machine, error) {
	var notstarted []*state.Machine
	for _, m := range machines {
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			notstarted = append(notstarted, m)
		} else {
			log.Printf("machine %s already started as instance %q", m, id)
		}
	}
	return notstarted, nil
}

func (p *Provisioner) startMachines(machines []*state.Machine) error {
	for _, m := range machines {
		if err := p.startMachine(m); err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) startMachine(m *state.Machine) error {
	// TODO(dfc) the state.Info passed to environ.StartInstance remains contentious
	// however as the PA only knows one state.Info, and that info is used by MAs and
	// UAs to locate the ZK for this environment, it is logical to use the same
	// state.Info as the PA.
	inst, err := p.environ.StartInstance(m.Id(), p.info)
	if err != nil {
		log.Printf("provisioner can't start machine %s: %v", m, err)
		return err
	}

	// assign the instance id to the machine
	if err := m.SetInstanceId(inst.Id()); err != nil {
		return err
	}

	// populate the local cache
	p.instances[m.Id()] = inst
	p.machines[inst.Id()] = m
	log.Printf("provisioner started machine %s as instance %s", m, inst.Id())
	return nil
}

func (p *Provisioner) stopInstances(instances []environs.Instance) error {
	// Although calling StopInstance with an empty slice should produce no change in the
	// provider, environs like dummy do not consider this a noop.
	if len(instances) == 0 {
		return nil
	}
	if err := p.environ.StopInstances(instances); err != nil {
		return err
	}

	// cleanup cache
	for _, i := range instances {
		if m, ok := p.machines[i.Id()]; ok {
			delete(p.machines, i.Id())
			delete(p.instances, m.Id())
		}
	}
	return nil
}

// instanceForMachine returns the environs.Instance that represents this machine's instance.
func (p *Provisioner) instanceForMachine(m *state.Machine) (environs.Instance, error) {
	inst, ok := p.instances[m.Id()]
	if !ok {
		// not cached locally, ask the environ.
		id, err := m.InstanceId()
		if err != nil {
			return nil, err
		}
		if id == "" {
			// TODO(dfc) InstanceId should return an error if the id isn't set.
			return nil, fmt.Errorf("machine %s not found", m)
		}
		// TODO(dfc) this should be batched, or the cache preloaded at startup to
		// avoid N calls to the envirion.
		insts, err := p.environ.Instances([]string{id})
		if err != nil {
			// the provider doesn't know about this instance, give up.
			return nil, err
		}
		inst = insts[0]
	}
	return inst, nil
}

// instancesForMachines returns a list of environs.Instance that represent the list of machines running
// in the provider.
func (p *Provisioner) instancesForMachines(machines []*state.Machine) ([]environs.Instance, error) {
	var insts []environs.Instance
	for _, m := range machines {
		inst, err := p.instanceForMachine(m)
		if err != nil {
			return nil, err
		}
		insts = append(insts, inst)
	}
	return insts, nil
}
<|endoftext|>


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

package main

import (
	"io/ioutil"
	"os"
	"path"
	"testing"

	"github.com/stretchr/testify/require"
)

const (
	validMetadata = `
name: metricreceiver
metrics:
  system.cpu.time:
    description: Total CPU seconds broken down by different states.
    extended_description: Additional information on CPU Time can be found [here](https://en.wikipedia.org/wiki/CPU_time).
    unit: s
    sum:
      aggregation: cumulative
    attributes: []
`
)

func Test_runContents(t *testing.T) {
	type args struct {
		yml string
	}
	tests := []struct {
		name    string
		args    args
		want    string
		wantErr string
	}{
		{
			name: "valid metadata",
			args: args{validMetadata},
			want: "",
		},
		{
			name:    "invalid yaml",
			args:    args{"invalid"},
			want:    "",
			wantErr: "cannot unmarshal",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tmpdir, err := ioutil.TempDir("", "metadata-test-*")
			require.NoError(t, err)
			t.Cleanup(func() {
				require.NoError(t, os.RemoveAll(tmpdir))
			})

			metadataFile := path.Join(tmpdir, "metadata.yaml")
			require.NoError(t, ioutil.WriteFile(metadataFile, []byte(tt.args.yml), 0600))

			err = run(metadataFile)

			if tt.wantErr != "" {
				require.Regexp(t, tt.wantErr, err)
			} else {
				require.NoError(t, err)
				require.FileExists(t, path.Join(tmpdir, "internal/metadata/generated_metrics.go"))
				require.FileExists(t, path.Join(tmpdir, "documentation.md"))
			}
		})
	}
}

func Test_run(t *testing.T) {
	type args struct {
		ymlPath string
	}
	tests := []struct {
		name    string
		args    args
		wantErr bool
	}{
		{
			name:    "no argument",
			args:    args{""},
			wantErr: true,
		},
		{
			name:    "no such file",
			args:    args{"/no/such/file"},
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if err := run(tt.args.ymlPath); (err != nil) != tt.wantErr {
				t.Errorf("run() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}
<commit_msg>[mdatagen] Fix metadata yaml in metricdata_test.go (#6473)<commit_after>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"io/ioutil"
	"os"
	"path"
	"testing"

	"github.com/stretchr/testify/require"
)

const (
	validMetadata = `
name: metricreceiver
metrics:
  system.cpu.time:
    description: Total CPU seconds broken down by different states.
    extended_documentation: Additional information on CPU Time can be found [here](https://en.wikipedia.org/wiki/CPU_time).
    unit: s
    sum:
      aggregation: cumulative
    attributes: []
`
)

func Test_runContents(t *testing.T) {
	type args struct {
		yml string
	}
	tests := []struct {
		name    string
		args    args
		want    string
		wantErr string
	}{
		{
			name: "valid metadata",
			args: args{validMetadata},
			want: "",
		},
		{
			name:    "invalid yaml",
			args:    args{"invalid"},
			want:    "",
			wantErr: "cannot unmarshal",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tmpdir, err := ioutil.TempDir("", "metadata-test-*")
			require.NoError(t, err)
			t.Cleanup(func() {
				require.NoError(t, os.RemoveAll(tmpdir))
			})

			metadataFile := path.Join(tmpdir, "metadata.yaml")
			require.NoError(t, ioutil.WriteFile(metadataFile, []byte(tt.args.yml), 0600))

			err = run(metadataFile)

			if tt.wantErr != "" {
				require.Regexp(t, tt.wantErr, err)
			} else {
				require.NoError(t, err)
				require.FileExists(t, path.Join(tmpdir, "internal/metadata/generated_metrics.go"))
				require.FileExists(t, path.Join(tmpdir, "documentation.md"))
			}
		})
	}
}

func Test_run(t *testing.T) {
	type args struct {
		ymlPath string
	}
	tests := []struct {
		name    string
		args    args
		wantErr bool
	}{
		{
			name:    "no argument",
			args:    args{""},
			wantErr: true,
		},
		{
			name:    "no such file",
			args:    args{"/no/such/file"},
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if err := run(tt.args.ymlPath); (err != nil) != tt.wantErr {
				t.Errorf("run() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}
<|endoftext|>


<commit_before>package sshttp

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

	"github.com/emirozer/exposq/osquery"
)

// gets http.NewServeMux from main and sets the routes
func SetMux(mux http.ServeMux) {

	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		fmt.Fprintf(w, "Main")
	})

	mux.HandleFunc("/kernel_info", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["kernel_info"] + "\""
		sout := dispatchCmd(cmd)

		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/kernel_integrity", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["kernel_integrity"] + "\""
		sout := dispatchCmd(cmd)

		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/kernel_modules", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["kernel_modules"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/mounts", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["mounts"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/mounts_device_alias_none", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["mounts_device_alias_none"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/block_devices", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["block_devices"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/block_device_sz0", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["block_device_sz0"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/block_device_ata", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["block_device_ata"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/acpi_tables", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["acpi_tables"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/cpuid", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["cpuid"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/crontab", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["crontab"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/disk_encryption", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["disk_encryption"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/etc_hosts", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["etc_hosts"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/etc_protocols", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["etc_protocols"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/etc_services", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["etc_services"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/groups", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["groups"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/interface_addresses", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["interface_addresses"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/interface_details", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["interface_details"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/iptables", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["iptables"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/last", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["last"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/listening_ports", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["listening_ports"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/logged_in_users", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["logged_in_users"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/memory_map", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["memory_map"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/pci_devices", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["pci_devices"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/passwd_changes", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["passwd_changes"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/processes", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["processes"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/processes_root", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["processes_root"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_envs", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["process_envs"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_memory_map", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["process_memory_map"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_open_files", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["process_open_files"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_open_sockets", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["process_open_sockets"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/routes", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["routes"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/shared_memory", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["shared_memory"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/shell_history", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["shell_history"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/users", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["users"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/user_groups", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["user_groups"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/uptime", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["uptime"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/relay", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["relay"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/mitm", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["mitm"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/setuid_enabled", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["setuid_enabled"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/q_scan_ps_bin", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["q_scan_ps_bin"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/ps_lst_tcp_udp", func(w http.ResponseWriter, req *http.Request) {
		oq := osquery.GenericOsQueries()

		cmd := "osqueryi " + "\"" + oq["ps_lst_tcp_udp"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

}

func dispatchCmd(cmd string) string {
	out, err := exec.Command("sh", "-c", cmd).Output()
	if err != nil {
		log.Println(err)
	}
	sout := string(out[:])

	if len(sout) == 0 {
		sout = fmt.Sprintf("No response for the following query from this machine : %v", cmd)
	}
	return sout
}
<commit_msg>added temp table to root route<commit_after>package sshttp

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

	"github.com/emirozer/exposq/osquery"
)

// gets http.NewServeMux from main and sets the routes
func SetMux(mux http.ServeMux) {

	oq := osquery.GenericOsQueries()
	coq := osquery.CentOsQueries()
	doq := osquery.DebUbOsQueries()

	mux.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		fmt.Fprintln(w, "Generic Os Query Routes:")
		for k, _ := range oq {
			fmt.Fprintf(w, "/"+k+"\n")
		}

		fmt.Fprintln(w, " RedHat Based Os Query Routes:")
		for k, _ := range coq {
			fmt.Fprintf(w, "/"+k+"\n")
		}

		fmt.Fprintln(w, "Debian Based Os Query Routes:")
		for k, _ := range doq {
			fmt.Fprintf(w, "/"+k+"\n")
		}

	})

	mux.HandleFunc("/kernel_info", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["kernel_info"] + "\""
		sout := dispatchCmd(cmd)

		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/kernel_integrity", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["kernel_integrity"] + "\""
		sout := dispatchCmd(cmd)

		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/kernel_modules", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["kernel_modules"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/mounts", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["mounts"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/mounts_device_alias_none", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["mounts_device_alias_none"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/block_devices", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["block_devices"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/block_device_sz0", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["block_device_sz0"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/block_device_ata", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["block_device_ata"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/acpi_tables", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["acpi_tables"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/cpuid", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["cpuid"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/crontab", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["crontab"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/disk_encryption", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["disk_encryption"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/etc_hosts", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["etc_hosts"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/etc_protocols", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["etc_protocols"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/etc_services", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["etc_services"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/groups", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["groups"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/interface_addresses", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["interface_addresses"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/interface_details", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["interface_details"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/iptables", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["iptables"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/last", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["last"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/listening_ports", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["listening_ports"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/logged_in_users", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["logged_in_users"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/memory_map", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["memory_map"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/pci_devices", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["pci_devices"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/passwd_changes", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["passwd_changes"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/processes", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["processes"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/processes_root", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["processes_root"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_envs", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["process_envs"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_memory_map", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["process_memory_map"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_open_files", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["process_open_files"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/process_open_sockets", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["process_open_sockets"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/routes", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["routes"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/shared_memory", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["shared_memory"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/shell_history", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["shell_history"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/users", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["users"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/user_groups", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["user_groups"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/uptime", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["uptime"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/relay", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["relay"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/mitm", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["mitm"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/setuid_enabled", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["setuid_enabled"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/q_scan_ps_bin", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["q_scan_ps_bin"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

	mux.HandleFunc("/ps_lst_tcp_udp", func(w http.ResponseWriter, req *http.Request) {

		cmd := "osqueryi " + "\"" + oq["ps_lst_tcp_udp"] + "\""
		sout := dispatchCmd(cmd)
		fmt.Fprintf(w, sout)
	})

}

func dispatchCmd(cmd string) string {

	out, err := exec.Command("sh", "-c", cmd).Output()
	if err != nil {
		log.Println(err)
	}
	sout := string(out[:])

	if len(sout) == 0 {
		sout = fmt.Sprintf("No response for the following query from this machine : %v", cmd)
	}
	return sout
}
<|endoftext|>