<commit_before>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"github.com/mackerelio/mackerel-agent/logging"
	"net/http"
	"os"
)

var logger = logging.GetLogger("metrics.plugin.elasticsearch")

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"elasticsearch.http": mp.Graphs{
		Label: "Elasticsearch HTTP",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "http_opened", Label: "Opened", Diff: true},
		},
	},
	"elasticsearch.indices": mp.Graphs{
		Label: "Elasticsearch Indices",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "total_indexing_index", Label: "Indexing-Index", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_indexing_delete", Label: "Indexing-Delete", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_get", Label: "Get", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_search_query", Label: "Search-Query", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_search_fetch", Label: "Search-fetch", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_merges", Label: "Merges", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_refresh", Label: "Refresh", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_flush", Label: "Flush", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_warmer", Label: "Warmer", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_percolate", Label: "Percolate", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_suggest", Label: "Suggest", Diff: true, Stacked: true},
		},
	},
	"elasticsearch.indices.docs": mp.Graphs{
		Label: "Elasticsearch Indices Docs",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "docs_count", Label: "Count", Stacked: true},
			mp.Metrics{Name: "docs_deleted", Label: "Deleted", Stacked: true},
		},
	},
	"elasticsearch.indices.memory_size": mp.Graphs{
		Label: "Elasticsearch Indices Memory Size",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "fielddata_size", Label: "Fielddata", Stacked: true},
			mp.Metrics{Name: "filter_cache_size", Label: "Filter Cache", Stacked: true},
			mp.Metrics{Name: "segments_size", Label: "Lucene Segments", Stacked: true},
		},
	},
	"elasticsearch.indices.evictions": mp.Graphs{
		Label: "Elasticsearch Indices Evictions",
		Unit: "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "evictions_fielddata", Label: "Fielddata", Diff: true},
			mp.Metrics{Name: "evictions_filter_cache", Label: "Filter Cache", Diff: true},
		},

	},
	"elasticsearch.jvm.heap": mp.Graphs{
		Label: "Elasticsearch JVM Heap Mem",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "heap_used", Label: "Used"},
			mp.Metrics{Name: "heap_max", Label: "Max"},
		},
	},
	"elasticsearch.thread_pool.threads": mp.Graphs{
		Label: "Elasticsearch Thread-Pool Threads",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "threads_generic", Label: "Generic", Stacked: true},
			mp.Metrics{Name: "threads_index", Label: "Index", Stacked: true},
			mp.Metrics{Name: "threads_snapshot_data", Label: "Snapshot Data", Stacked: true},
			mp.Metrics{Name: "threads_get", Label: "Get", Stacked: true},
			mp.Metrics{Name: "threads_bench", Label: "Bench", Stacked: true},
			mp.Metrics{Name: "threads_snapshot", Label: "Snapshot", Stacked: true},
			mp.Metrics{Name: "threads_merge", Label: "Merge", Stacked: true},
			mp.Metrics{Name: "threads_suggest", Label: "Suggest", Stacked: true},
			mp.Metrics{Name: "threads_bulk", Label: "Bulk", Stacked: true},
			mp.Metrics{Name: "threads_optimize", Label: "Optimize", Stacked: true},
			mp.Metrics{Name: "threads_warmer", Label: "Warmer", Stacked: true},
			mp.Metrics{Name: "threads_flush", Label: "Flush", Stacked: true},
			mp.Metrics{Name: "threads_search", Label: "Search", Stacked: true},
			mp.Metrics{Name: "threads_percolate", Label: "Percolate", Stacked: true},
			mp.Metrics{Name: "threads_refresh", Label: "Refresh", Stacked: true},
			mp.Metrics{Name: "threads_management", Label: "Management", Stacked: true},
		},
	},
	"elasticsearch.transport.count": mp.Graphs{
		Label: "Elasticsearch Transport Count",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "count_rx", Label: "TX", Diff: true},
			mp.Metrics{Name: "count_tx", Label: "RX", Diff: true},
		},
	},
}

var metricPlace map[string][]string = map[string][]string{
	"http_opened":           []string{"http", "total_opened"},
	"total_indexing_index":  []string{"indices", "indexing", "index_total"},
	"total_indexing_delete": []string{"indices", "indexing", "delete_total"},
	"total_get":             []string{"indices", "get", "total"},
	"total_search_query":    []string{"indices", "search", "query_total"},
	"total_search_fetch":    []string{"indices", "search", "fetch_total"},
	"total_merges":          []string{"indices", "merges", "total"},
	"total_refresh":         []string{"indices", "refresh", "total"},
	"total_flush":           []string{"indices", "flush", "total"},
	"total_warmer":          []string{"indices", "warmer", "total"},
	"total_percolate":       []string{"indices", "percolate", "total"},
	"total_suggest":         []string{"indices", "suggest", "total"},
	"docs_count":            []string{"indices", "docs", "count"},
	"docs_deleted":          []string{"indices", "docs", "deleted"},
	"fielddata_size":        []string{"indices", "fielddata", "memory_size_in_bytes"},
	"filter_cache_size":     []string{"indices", "filter_cache", "memory_size_in_bytes"},
	"segments_size":         []string{"indices", "segments", "memory_in_bytes"},
	"evictions_fielddata":   []string{"indices", "fielddata", "evictions"},
	"evictions_filter_cache":[]string{"indices", "filter_cache", "evictions"},
	"heap_used":             []string{"jvm", "mem", "heap_used_in_bytes"},
	"heap_max":              []string{"jvm", "mem", "heap_max_in_bytes"},
	"threads_generic":       []string{"thread_pool", "generic", "threads"},
	"threads_index":         []string{"thread_pool", "index", "threads"},
	"threads_snapshot_data": []string{"thread_pool", "snapshot_data", "threads"},
	"threads_get":           []string{"thread_pool", "get", "threads"},
	"threads_bench":         []string{"thread_pool", "bench", "threads"},
	"threads_snapshot":      []string{"thread_pool", "snapshot", "threads"},
	"threads_merge":         []string{"thread_pool", "merge", "threads"},
	"threads_suggest":       []string{"thread_pool", "suggest", "threads"},
	"threads_bulk":          []string{"thread_pool", "bulk", "threads"},
	"threads_optimize":      []string{"thread_pool", "optimize", "threads"},
	"threads_warmer":        []string{"thread_pool", "warmer", "threads"},
	"threads_flush":         []string{"thread_pool", "flush", "threads"},
	"threads_search":        []string{"thread_pool", "search", "threads"},
	"threads_percolate":     []string{"thread_pool", "percolate", "threads"},
	"threads_refresh":       []string{"thread_pool", "refresh", "threads"},
	"threads_management":    []string{"thread_pool", "management", "threads"},
	"count_rx":              []string{"transport", "rx_count"},
	"count_tx":              []string{"transport", "tx_count"},
}

func GetFloatValue(s map[string]interface{}, keys []string) (float64, error) {
	var val float64
	sm := s
	for i, k := range keys {
		if i+1 < len(keys) {
			switch sm[k].(type) {
			case map[string]interface{}:
				sm = sm[k].(map[string]interface{})
			default:
				return 0, errors.New("Cannot handle as a hash")
			}
		} else {
			switch sm[k].(type) {
			case float64:
				val = sm[k].(float64)
			default:
				return 0, errors.New("Not float64")
			}
		}
	}

	return val, nil
}

type ElasticsearchPlugin struct {
	Uri string
}

func (p ElasticsearchPlugin) FetchMetrics() (map[string]float64, error) {
	resp, err := http.Get(p.Uri + "/_nodes/_local/stats")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	stat := make(map[string]float64)
	decoder := json.NewDecoder(resp.Body)

	var s map[string]interface{}
	err = decoder.Decode(&s)
	if err != nil {
		return nil, err
	}

	nodes := s["nodes"].(map[string]interface{})
	n := ""
	for k, _ := range nodes {
		if n != "" {
			return nil, errors.New("Multiple node found")
		}
		n = k
	}
	node := nodes[n].(map[string]interface{})

	for k, v := range metricPlace {
		val, err := GetFloatValue(node, v)
		if err != nil {
			logger.Errorf("Failed to find '%s': %s", k, err)
			continue
		}

		stat[k] = val
	}

	return stat, nil
}

func (n ElasticsearchPlugin) GraphDefinition() map[string](mp.Graphs) {
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Host")
	optPort := flag.String("port", "9200", "Port")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var elasticsearch ElasticsearchPlugin
	elasticsearch.Uri = fmt.Sprintf("http://%s:%s", *optHost, *optPort)

	helper := mp.NewMackerelPlugin(elasticsearch)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-elasticsearch-%s-%s", *optHost, *optPort)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>elasticsearch: add memory size used by lucene segments such as fixed bit set, which were added at Elasticsearch 1.4<commit_after>package main

import (
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"github.com/mackerelio/mackerel-agent/logging"
	"net/http"
	"os"
)

var logger = logging.GetLogger("metrics.plugin.elasticsearch")

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"elasticsearch.http": mp.Graphs{
		Label: "Elasticsearch HTTP",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "http_opened", Label: "Opened", Diff: true},
		},
	},
	"elasticsearch.indices": mp.Graphs{
		Label: "Elasticsearch Indices",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "total_indexing_index", Label: "Indexing-Index", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_indexing_delete", Label: "Indexing-Delete", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_get", Label: "Get", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_search_query", Label: "Search-Query", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_search_fetch", Label: "Search-fetch", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_merges", Label: "Merges", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_refresh", Label: "Refresh", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_flush", Label: "Flush", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_warmer", Label: "Warmer", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_percolate", Label: "Percolate", Diff: true, Stacked: true},
			mp.Metrics{Name: "total_suggest", Label: "Suggest", Diff: true, Stacked: true},
		},
	},
	"elasticsearch.indices.docs": mp.Graphs{
		Label: "Elasticsearch Indices Docs",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "docs_count", Label: "Count", Stacked: true},
			mp.Metrics{Name: "docs_deleted", Label: "Deleted", Stacked: true},
		},
	},
	"elasticsearch.indices.memory_size": mp.Graphs{
		Label: "Elasticsearch Indices Memory Size",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "fielddata_size", Label: "Fielddata", Stacked: true},
			mp.Metrics{Name: "filter_cache_size", Label: "Filter Cache", Stacked: true},
			mp.Metrics{Name: "segments_size", Label: "Lucene Segments", Stacked: true},
			mp.Metrics{Name: "segments_index_writer_size", Label: "Lucene Segments Index Writer", Stacked: true},
			mp.Metrics{Name: "segments_version_map_size", Label: "Lucene Segments Version Map", Stacked: true},
			mp.Metrics{Name: "segments_fixed_bit_set_size", Label: "Lucene Segments Fixed Bit Set", Stacked: true},
		},
	},
	"elasticsearch.indices.evictions": mp.Graphs{
		Label: "Elasticsearch Indices Evictions",
		Unit: "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "evictions_fielddata", Label: "Fielddata", Diff: true},
			mp.Metrics{Name: "evictions_filter_cache", Label: "Filter Cache", Diff: true},
		},

	},
	"elasticsearch.jvm.heap": mp.Graphs{
		Label: "Elasticsearch JVM Heap Mem",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "heap_used", Label: "Used"},
			mp.Metrics{Name: "heap_max", Label: "Max"},
		},
	},
	"elasticsearch.thread_pool.threads": mp.Graphs{
		Label: "Elasticsearch Thread-Pool Threads",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "threads_generic", Label: "Generic", Stacked: true},
			mp.Metrics{Name: "threads_index", Label: "Index", Stacked: true},
			mp.Metrics{Name: "threads_snapshot_data", Label: "Snapshot Data", Stacked: true},
			mp.Metrics{Name: "threads_get", Label: "Get", Stacked: true},
			mp.Metrics{Name: "threads_bench", Label: "Bench", Stacked: true},
			mp.Metrics{Name: "threads_snapshot", Label: "Snapshot", Stacked: true},
			mp.Metrics{Name: "threads_merge", Label: "Merge", Stacked: true},
			mp.Metrics{Name: "threads_suggest", Label: "Suggest", Stacked: true},
			mp.Metrics{Name: "threads_bulk", Label: "Bulk", Stacked: true},
			mp.Metrics{Name: "threads_optimize", Label: "Optimize", Stacked: true},
			mp.Metrics{Name: "threads_warmer", Label: "Warmer", Stacked: true},
			mp.Metrics{Name: "threads_flush", Label: "Flush", Stacked: true},
			mp.Metrics{Name: "threads_search", Label: "Search", Stacked: true},
			mp.Metrics{Name: "threads_percolate", Label: "Percolate", Stacked: true},
			mp.Metrics{Name: "threads_refresh", Label: "Refresh", Stacked: true},
			mp.Metrics{Name: "threads_management", Label: "Management", Stacked: true},
		},
	},
	"elasticsearch.transport.count": mp.Graphs{
		Label: "Elasticsearch Transport Count",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "count_rx", Label: "TX", Diff: true},
			mp.Metrics{Name: "count_tx", Label: "RX", Diff: true},
		},
	},
}

var metricPlace map[string][]string = map[string][]string{
	"http_opened":           []string{"http", "total_opened"},
	"total_indexing_index":  []string{"indices", "indexing", "index_total"},
	"total_indexing_delete": []string{"indices", "indexing", "delete_total"},
	"total_get":             []string{"indices", "get", "total"},
	"total_search_query":    []string{"indices", "search", "query_total"},
	"total_search_fetch":    []string{"indices", "search", "fetch_total"},
	"total_merges":          []string{"indices", "merges", "total"},
	"total_refresh":         []string{"indices", "refresh", "total"},
	"total_flush":           []string{"indices", "flush", "total"},
	"total_warmer":          []string{"indices", "warmer", "total"},
	"total_percolate":       []string{"indices", "percolate", "total"},
	"total_suggest":         []string{"indices", "suggest", "total"},
	"docs_count":            []string{"indices", "docs", "count"},
	"docs_deleted":          []string{"indices", "docs", "deleted"},
	"fielddata_size":        []string{"indices", "fielddata", "memory_size_in_bytes"},
	"filter_cache_size":     []string{"indices", "filter_cache", "memory_size_in_bytes"},
	"segments_size":         []string{"indices", "segments", "memory_in_bytes"},
	"segments_index_writer_size":  []string{"indices", "segments", "index_writer_memory_in_bytes"},
	"segments_version_map_size":   []string{"indices", "segments", "version_map_memory_in_bytes"},
	"segments_fixed_bit_set_size": []string{"indices", "segments", "fixed_bit_set_memory_in_bytes"},
	"evictions_fielddata":   []string{"indices", "fielddata", "evictions"},
	"evictions_filter_cache":[]string{"indices", "filter_cache", "evictions"},
	"heap_used":             []string{"jvm", "mem", "heap_used_in_bytes"},
	"heap_max":              []string{"jvm", "mem", "heap_max_in_bytes"},
	"threads_generic":       []string{"thread_pool", "generic", "threads"},
	"threads_index":         []string{"thread_pool", "index", "threads"},
	"threads_snapshot_data": []string{"thread_pool", "snapshot_data", "threads"},
	"threads_get":           []string{"thread_pool", "get", "threads"},
	"threads_bench":         []string{"thread_pool", "bench", "threads"},
	"threads_snapshot":      []string{"thread_pool", "snapshot", "threads"},
	"threads_merge":         []string{"thread_pool", "merge", "threads"},
	"threads_suggest":       []string{"thread_pool", "suggest", "threads"},
	"threads_bulk":          []string{"thread_pool", "bulk", "threads"},
	"threads_optimize":      []string{"thread_pool", "optimize", "threads"},
	"threads_warmer":        []string{"thread_pool", "warmer", "threads"},
	"threads_flush":         []string{"thread_pool", "flush", "threads"},
	"threads_search":        []string{"thread_pool", "search", "threads"},
	"threads_percolate":     []string{"thread_pool", "percolate", "threads"},
	"threads_refresh":       []string{"thread_pool", "refresh", "threads"},
	"threads_management":    []string{"thread_pool", "management", "threads"},
	"count_rx":              []string{"transport", "rx_count"},
	"count_tx":              []string{"transport", "tx_count"},
}

func GetFloatValue(s map[string]interface{}, keys []string) (float64, error) {
	var val float64
	sm := s
	for i, k := range keys {
		if i+1 < len(keys) {
			switch sm[k].(type) {
			case map[string]interface{}:
				sm = sm[k].(map[string]interface{})
			default:
				return 0, errors.New("Cannot handle as a hash")
			}
		} else {
			switch sm[k].(type) {
			case float64:
				val = sm[k].(float64)
			default:
				return 0, errors.New("Not float64")
			}
		}
	}

	return val, nil
}

type ElasticsearchPlugin struct {
	Uri string
}

func (p ElasticsearchPlugin) FetchMetrics() (map[string]float64, error) {
	resp, err := http.Get(p.Uri + "/_nodes/_local/stats")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	stat := make(map[string]float64)
	decoder := json.NewDecoder(resp.Body)

	var s map[string]interface{}
	err = decoder.Decode(&s)
	if err != nil {
		return nil, err
	}

	nodes := s["nodes"].(map[string]interface{})
	n := ""
	for k, _ := range nodes {
		if n != "" {
			return nil, errors.New("Multiple node found")
		}
		n = k
	}
	node := nodes[n].(map[string]interface{})

	for k, v := range metricPlace {
		val, err := GetFloatValue(node, v)
		if err != nil {
			logger.Errorf("Failed to find '%s': %s", k, err)
			continue
		}

		stat[k] = val
	}

	return stat, nil
}

func (n ElasticsearchPlugin) GraphDefinition() map[string](mp.Graphs) {
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Host")
	optPort := flag.String("port", "9200", "Port")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var elasticsearch ElasticsearchPlugin
	elasticsearch.Uri = fmt.Sprintf("http://%s:%s", *optHost, *optPort)

	helper := mp.NewMackerelPlugin(elasticsearch)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-elasticsearch-%s-%s", *optHost, *optPort)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/maliceio/go-plugin-utils/utils"
	"github.com/maliceio/malice/malice/database"
	"github.com/maliceio/malice/malice/persist"
	util "github.com/maliceio/malice/utils"
	elastic "gopkg.in/olivere/elastic.v3"
)

// ElasticAddr ElasticSearch address to user for connections
var ElasticAddr string

// InitElasticSearch initalizes ElasticSearch for use with malice
func InitElasticSearch() error {
	client, err := elastic.NewSimpleClient()
	utils.Assert(err)

	exists, err := client.IndexExists("malice").Do()
	utils.Assert(err)

	if !exists {
		// Index does not exist yet.
		createIndex, err := client.CreateIndex("malice").BodyString(mapping).Do()
		utils.Assert(err)
		if !createIndex.Acknowledged {
			// Not acknowledged
			log.Error("Couldn't create Index.")
		} else {
			log.Info("Created Index: ", "malice")
		}
	} else {
		log.Info("Index malice already exists.")
	}

	return err
}

// TestConnection tests the ElasticSearch connection
func TestConnection(addr string) error {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	// connect to ElasticSearch where --link elastic was using via malice in Docker
	log.Debugf("Attempting to connect to: %s", ElasticAddr)
	_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

	// Ping the Elasticsearch server to get e.g. the version number
	// info, code, err := client.Ping(ElasticAddr).Do()
	// utils.Assert(err)
	// fmt.Printf("Elasticsearch returned with code %d and version %s", code, info.Version.Number)

	if err != nil {
		// connect to ElasticSearch via malice in Docker
		ElasticAddr = fmt.Sprintf("%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", addr))
		log.Debugf("Attempting to connect to: %s", ElasticAddr)
		_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

		if err != nil {
			// connect to ElasticSearch using Docker for Mac
			ElasticAddr = fmt.Sprintf("%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "localhost"))
			log.Debugf("Attempting to connect to: %s", ElasticAddr)
			_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
			return err
		}
		return err
	}
	return err
}

// WriteFileToDatabase inserts sample into Database
func WriteFileToDatabase(sample persist.File) elastic.IndexResponse {
	client, err := elastic.NewSimpleClient()
	utils.Assert(err)

	// getSample, err := client.Get().
	// 	Index("malice").
	// 	Type("samples").
	// 	Id("1").
	// 	Do()

	// fmt.Println(getSample)
	// fmt.Println(err)
	// if err != nil {

	// }

	// if getSample.Found {
	// 	fmt.Printf("Got document %s in version %d from index %s, type %s\n", getSample.Id, getSample.Version, getSample.Index, getSample.Type)
	// } else {

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file":      sample,
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	update, err := client.Update().Index("malice").Type("samples").Id(newScan.Id).
		Doc(map[string]interface{}{
			"plugins": map[string]interface{}{
				"intel": map[string]interface{}{
					"nsrl": "UPDATED",
				},
			},
		}).
		Do()
	utils.Assert(err)
	log.Debugf("New version of sample %q is now %d\n", update.Id, update.Version)

	// }

	return *newScan
}

// WriteHashToDatabase inserts sample into Database
func WriteHashToDatabase(hash string) elastic.IndexResponse {

	hashType, err := util.GetHashType(hash)
	utils.Assert(err)

	client, err := elastic.NewSimpleClient()
	utils.Assert(err)

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file": map[string]interface{}{
			hashType: hash,
		},
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	return *newScan
}
<commit_msg>update elasticsearch.go<commit_after>package elasticsearch

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/maliceio/go-plugin-utils/utils"
	"github.com/maliceio/malice/malice/database"
	"github.com/maliceio/malice/malice/persist"
	util "github.com/maliceio/malice/utils"
	elastic "gopkg.in/olivere/elastic.v3"
)

// ElasticAddr ElasticSearch address to user for connections
var ElasticAddr string

// InitElasticSearch initalizes ElasticSearch for use with malice
func InitElasticSearch() error {
	client, err := elastic.NewSimpleClient()
	utils.Assert(err)

	exists, err := client.IndexExists("malice").Do()
	utils.Assert(err)

	if !exists {
		// Index does not exist yet.
		createIndex, err := client.CreateIndex("malice").BodyString(mapping).Do()
		utils.Assert(err)
		if !createIndex.Acknowledged {
			// Not acknowledged
			log.Error("Couldn't create Index.")
		} else {
			log.Info("Created Index: ", "malice")
		}
	} else {
		log.Info("Index malice already exists.")
	}

	return err
}

// TestConnection tests the ElasticSearch connection
func TestConnection(addr string) error {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	// connect to ElasticSearch where --link elastic was using via malice in Docker
	log.Debugf("Attempting to connect to: %s", ElasticAddr)
	_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

	// Ping the Elasticsearch server to get e.g. the version number
	// info, code, err := client.Ping(ElasticAddr).Do()
	// utils.Assert(err)
	// fmt.Printf("Elasticsearch returned with code %d and version %s", code, info.Version.Number)

	if err != nil {
		// connect to ElasticSearch via malice in Docker
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", addr))
		log.Debugf("Attempting to connect to: %s", ElasticAddr)
		_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

		if err != nil {
			// connect to ElasticSearch using Docker for Mac
			ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "localhost"))
			log.Debugf("Attempting to connect to: %s", ElasticAddr)
			_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
			return err
		}
		return err
	}
	return err
}

// WriteFileToDatabase inserts sample into Database
func WriteFileToDatabase(sample persist.File) elastic.IndexResponse {
	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	// getSample, err := client.Get().
	// 	Index("malice").
	// 	Type("samples").
	// 	Id("1").
	// 	Do()

	// fmt.Println(getSample)
	// fmt.Println(err)
	// if err != nil {

	// }

	// if getSample.Found {
	// 	fmt.Printf("Got document %s in version %d from index %s, type %s\n", getSample.Id, getSample.Version, getSample.Index, getSample.Type)
	// } else {

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file":      sample,
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	update, err := client.Update().Index("malice").Type("samples").Id(newScan.Id).
		Doc(map[string]interface{}{
			"plugins": map[string]interface{}{
				"intel": map[string]interface{}{
					"nsrl": "UPDATED",
				},
			},
		}).
		Do()
	utils.Assert(err)
	log.Debugf("New version of sample %q is now %d\n", update.Id, update.Version)

	// }

	return *newScan
}

// WriteHashToDatabase inserts sample into Database
func WriteHashToDatabase(hash string) elastic.IndexResponse {

	hashType, err := util.GetHashType(hash)
	utils.Assert(err)

	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file": map[string]interface{}{
			hashType: hash,
		},
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	return *newScan
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"owl/common/types"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

var mydb *Storage

type Storage struct {
	*sqlx.DB
}

func InitMysqlConnPool() error {
	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8&parseTime=true&loc=Local",
		GlobalConfig.MySQLUser, GlobalConfig.MySQLPassword, GlobalConfig.MySQLAddr, GlobalConfig.MySQLDBName)
	db, err := sqlx.Open("mysql", dsn)
	if err != nil {
		return err
	}
	db.SetMaxIdleConns(GlobalConfig.MySQLMaxIdleConn)
	db.SetMaxOpenConns(GlobalConfig.MySQLMaxConn)
	mydb = &Storage{db}
	return nil
}

func (s *Storage) createHost(host *types.Host) error {
	now := time.Now().Format(timeFomart)
	sqlString := fmt.Sprintf("insert into `host`(`id`, `ip`, `hostname`, `uptime`, `idle_pct`, `agent_version`, `create_at`, `update_at`) "+
		" values('%s', '%s', '%s', %0.2f, %0.2f, '%s','%s','%s')", host.ID, host.IP, host.Hostname, host.Uptime, host.IdlePct, host.AgentVersion, now, now)
	lg.Debug("create host:%s", sqlString)
	_, err := s.Exec(sqlString)
	return err
}

func (s *Storage) updateHost(host *types.Host) error {
	sqlString := fmt.Sprintf("update `host` set `ip`='%s', `uptime`=%0.2f, `idle_pct`=%0.2f, `hostname`='%s', `agent_version`='%s', `update_at`='%s' where id='%s'",
		host.IP, host.Uptime, host.IdlePct, host.Hostname, host.AgentVersion, time.Now().Format(timeFomart), host.ID)
	lg.Debug("update host:%s", sqlString)
	_, err := s.Exec(sqlString)
	return err
}

func (s *Storage) getHost(hostID string) (*types.Host, error) {
	host := &types.Host{}
	sqlString := fmt.Sprintf("select id, ip, hostname, agent_version,status,create_at, update_at  from `host` where id='%s'", hostID)
	lg.Debug("getHost:%s", sqlString)
	err := s.Get(host, sqlString)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return host, err
}

func (s *Storage) getAllHosts() []*types.Host {
	hosts := []*types.Host{}
	sqlString := fmt.Sprintf("select id, ip, hostname, agent_version,status,create_at, update_at  from `host`")
	if err := s.Select(&hosts, sqlString); err != nil {
		lg.Error("getNoMaintainHost %s", err)
		return nil
	}
	return hosts
}

func (s *Storage) setHostAlive(hostID string, status string) {
	sqlString := fmt.Sprintf("update `host` set `status` = '%s' where `id`='%s'", status, hostID)
	lg.Debug("setHostAlive:%s", sqlString)
	s.Exec(sqlString)
}

func (s *Storage) getHostPlugins(hostID string) ([]types.Plugin, error) {
	plugins := []types.Plugin{}
	idMap := make(map[int]struct{})
	sqlString := fmt.Sprintf("select `id`, `name`, `path`, `args`, `checksum`, `interval`, `timeout` from `plugin` where"+
		" id in (select `plugin_id` from `host_plugin` where `host_id`='%s')", hostID)
	lg.Debug("getHostPlugins:%s", sqlString)
	rows, err := s.Query(sqlString)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		plugin := types.Plugin{}
		if err := rows.Scan(&plugin.ID, &plugin.Name, &plugin.Path, &plugin.Args, &plugin.Checksum, &plugin.Interval, &plugin.Timeout); err != nil {
			fmt.Println(err)
			continue
		}
		plugins = append(plugins, plugin)
		idMap[plugin.ID] = struct{}{}
	}
	//获取主机组所有的插件
	sqlString = fmt.Sprintf("select `id`, `name`, `path`, `args`, `checksum`, `interval`, `timeout` from `plugin` where "+
		"id in (select `plugin_id` from `host_group_plugin` where group_id "+
		"in(select `host_group_id` from `host_group_host` where host_id='%s'))", hostID)
	lg.Debug("getHostGroupPlugins:%s", sqlString)
	rows, err = s.Query(sqlString)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		plugin := types.Plugin{}
		if err := rows.Scan(&plugin.ID, &plugin.Name, &plugin.Path, &plugin.Args, &plugin.Checksum, &plugin.Interval, &plugin.Timeout); err != nil {
			fmt.Println(err)
			continue
		}
		if _, ok := idMap[plugin.ID]; ok {
			continue
		}
		plugins = append(plugins, plugin)
		idMap[plugin.ID] = struct{}{}
	}
	return plugins, nil
}

func (s *Storage) metricIsExists(hostID, metric string, tags string) bool {
	sqlString := fmt.Sprintf("select `id` from `metric` where `host_id`= '%s' and `metric`='%s' and `tags`='%s'", hostID, metric, tags)
	lg.Debug("metricIsExists:%s", sqlString)
	var id int
	if err := s.Get(&id, sqlString); err != nil {
		if err == sql.ErrNoRows {
			// no row
			return false
		}
		// error
		lg.Error("metricIsExists:%s", err)
		return false
	}
	// exists
	return true
}

func (s *Storage) createMetric(hostID string, tsd types.TimeSeriesData) error {
	now := time.Now().Format(timeFomart)
	sqlString := fmt.Sprintf("insert into `metric` (`host_id`, `metric`, `tags`, `dt`, `cycle`, `create_at`, `update_at`) "+
		"values('%s', '%s', '%s', '%s', %d, '%s', '%s')",
		hostID, tsd.Metric, tsd.Tags2String(), tsd.DataType, tsd.Cycle, now, now)
	lg.Debug("createMetric:%s", sqlString)
	_, err := s.Exec(sqlString)
	return err
}
<commit_msg>更新服务端获取主机插件列表方法<commit_after>package main

import (
	"database/sql"
	"fmt"
	"owl/common/types"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

var mydb *Storage

type Storage struct {
	*sqlx.DB
}

func InitMysqlConnPool() error {
	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=utf8&parseTime=true&loc=Local",
		GlobalConfig.MySQLUser, GlobalConfig.MySQLPassword, GlobalConfig.MySQLAddr, GlobalConfig.MySQLDBName)
	db, err := sqlx.Open("mysql", dsn)
	if err != nil {
		return err
	}
	db.SetMaxIdleConns(GlobalConfig.MySQLMaxIdleConn)
	db.SetMaxOpenConns(GlobalConfig.MySQLMaxConn)
	mydb = &Storage{db}
	return nil
}

func (s *Storage) createHost(host *types.Host) error {
	now := time.Now().Format(timeFomart)
	sqlString := fmt.Sprintf("insert into `host`(`id`, `ip`, `hostname`, `uptime`, `idle_pct`, `agent_version`, `create_at`, `update_at`) "+
		" values('%s', '%s', '%s', %0.2f, %0.2f, '%s','%s','%s')", host.ID, host.IP, host.Hostname, host.Uptime, host.IdlePct, host.AgentVersion, now, now)
	lg.Debug("create host:%s", sqlString)
	_, err := s.Exec(sqlString)
	return err
}

func (s *Storage) updateHost(host *types.Host) error {
	sqlString := fmt.Sprintf("update `host` set `ip`='%s', `uptime`=%0.2f, `idle_pct`=%0.2f, `hostname`='%s', `agent_version`='%s', `update_at`='%s' where id='%s'",
		host.IP, host.Uptime, host.IdlePct, host.Hostname, host.AgentVersion, time.Now().Format(timeFomart), host.ID)
	lg.Debug("update host:%s", sqlString)
	_, err := s.Exec(sqlString)
	return err
}

func (s *Storage) getHost(hostID string) (*types.Host, error) {
	host := &types.Host{}
	sqlString := fmt.Sprintf("select id, ip, hostname, agent_version,status,create_at, update_at  from `host` where id='%s'", hostID)
	lg.Debug("getHost:%s", sqlString)
	err := s.Get(host, sqlString)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	return host, err
}

func (s *Storage) getAllHosts() []*types.Host {
	hosts := []*types.Host{}
	sqlString := fmt.Sprintf("select id, ip, hostname, agent_version,status,create_at, update_at  from `host`")
	if err := s.Select(&hosts, sqlString); err != nil {
		lg.Error("getNoMaintainHost %s", err)
		return nil
	}
	return hosts
}

func (s *Storage) setHostAlive(hostID string, status string) {
	sqlString := fmt.Sprintf("update `host` set `status` = '%s' where `id`='%s'", status, hostID)
	lg.Debug("setHostAlive:%s", sqlString)
	s.Exec(sqlString)
}

func (s *Storage) getHostPlugins(hostID string) ([]types.Plugin, error) {
	plugins := []types.Plugin{}
	idMap := make(map[string]types.Plugin)
	sqlString := fmt.Sprintf("select hp.id, p.name, p.path, hp.args, p.checksum, hp.interval, hp.timeout from host_plugin as hp "+
		" left join plugin as p on hp.plugin_id = p.id where host_id='%s'", hostID)
	lg.Debug("getHostPlugins:%s", sqlString)
	rows, err := s.Query(sqlString)
	if err != nil {
		lg.Error("getHostPlugins error:%v", err)
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		plugin := types.Plugin{}
		if err := rows.Scan(&plugin.ID, &plugin.Name, &plugin.Path, &plugin.Args, &plugin.Checksum, &plugin.Interval, &plugin.Timeout); err != nil {
			lg.Error("getHostPlugins error:%v", err)
			continue
		}
		plugins = append(plugins, plugin)
		idMap[plugin.UniqueKey()] = plugin
	}
	//获取主机组所有的插件
	sqlString = fmt.Sprintf("select hgp.id, p.name, p.path, hgp.args, p.checksum, hgp.interval, hgp.timeout from plugin as p "+
		" left join host_group_plugin as hgp on p.id = hgp.plugin_id where hgp.group_id in (select host_group_id from host_group_host where host_id='%s')", hostID)
	lg.Debug("getHostGroupPlugins:%s", sqlString)
	rows, err = s.Query(sqlString)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		plugin := types.Plugin{}
		if err := rows.Scan(&plugin.ID, &plugin.Name, &plugin.Path, &plugin.Args, &plugin.Checksum, &plugin.Interval, &plugin.Timeout); err != nil {
			lg.Error("getHostPlugins error:%v", err)
			continue
		}
		uniqueKey := plugin.UniqueKey()
		if p, ok := idMap[uniqueKey]; ok {
			lg.Warn("getHostPlugins: duplicate host group plugin (%v, %v)", plugin, p)
			continue
		}
		plugins = append(plugins, plugin)
		idMap[uniqueKey] = plugin
	}
	return plugins, nil
}

func (s *Storage) metricIsExists(hostID, metric string, tags string) bool {
	sqlString := fmt.Sprintf("select `id` from `metric` where `host_id`= '%s' and `metric`='%s' and `tags`='%s'", hostID, metric, tags)
	lg.Debug("metricIsExists:%s", sqlString)
	var id int
	if err := s.Get(&id, sqlString); err != nil {
		if err == sql.ErrNoRows {
			// no row
			return false
		}
		// error
		lg.Error("metricIsExists:%s", err)
		return false
	}
	// exists
	return true
}

func (s *Storage) createMetric(hostID string, tsd types.TimeSeriesData) error {
	now := time.Now().Format(timeFomart)
	sqlString := fmt.Sprintf("insert into `metric` (`host_id`, `metric`, `tags`, `dt`, `cycle`, `create_at`, `update_at`) "+
		"values('%s', '%s', '%s', '%s', %d, '%s', '%s')",
		hostID, tsd.Metric, tsd.Tags2String(), tsd.DataType, tsd.Cycle, now, now)
	lg.Debug("createMetric:%s", sqlString)
	_, err := s.Exec(sqlString)
	return err
}
<|endoftext|>


<commit_before>/*
Copyright 2022 Gravitational, Inc.

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 tshwrap

import (
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"

	"github.com/coreos/go-semver/semver"
	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/api/constants"
	"github.com/gravitational/teleport/api/identityfile"
	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/lib/client"
	"github.com/gravitational/teleport/lib/tlsca"
	"github.com/gravitational/teleport/tool/tbot/config"
	"github.com/gravitational/teleport/tool/tbot/identity"
	"github.com/gravitational/trace"
	"github.com/sirupsen/logrus"
)

const (
	// TSHVarName is the name of the environment variable that can override the
	// tsh path that would otherwise be located on the $PATH.
	TSHVarName = "TSH"

	// TSHMinVersion is the minimum version of tsh that supports Machine ID
	// proxies.
	TSHMinVersion = "9.3.0"
)

var log = logrus.WithFields(logrus.Fields{
	trace.Component: teleport.ComponentTBot,
})

// capture runs a command (presumably tsh) with the given arguments and
// returns it's captured stdout. Stderr is ignored. Errors are returned per
// exec.Command().Output() semantics.
func capture(tshPath string, args ...string) ([]byte, error) {
	out, err := exec.Command(tshPath, args...).Output()
	if err != nil {
		return nil, trace.Wrap(err, "error executing tsh")
	}

	return out, nil
}

// Wrapper is a wrapper to execute `tsh` commands via a subprocess.
type Wrapper struct {
	// path is a path to the tsh executable
	path string

	// capture is the function for capturing a command's output. It may be
	// overridden by tests for mocking purposes, but by default is expected to
	// execute an actual tsh binary on the host system.
	capture func(tshPath string, args ...string) ([]byte, error)
}

// New creates a new tsh wrapper. If a $TSH var is set it uses that path,
// otherwise looks for tsh on the OS path.
func New() (*Wrapper, error) {
	if val, ok := os.LookupEnv(TSHVarName); ok {
		return &Wrapper{
			path:    val,
			capture: capture,
		}, nil
	}

	binary := "tsh"
	if runtime.GOOS == constants.WindowsOS {
		binary = "tsh.exe"
	}

	path, err := exec.LookPath(binary)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	return &Wrapper{
		path:    path,
		capture: capture,
	}, nil
}

// Exec runs tsh with the given environment variables and arguments. The child
// process inherits stdin/stdout/stderr and runs until completion. Errors are
// returned per `exec.Command().Run()` semantics.
func (w *Wrapper) Exec(env map[string]string, args ...string) error {
	// The subprocess should inherit the environment plus our vars. Our env
	// vars will safely overwrite those from the environment, per `exec.Cmd`
	// docs.
	environ := os.Environ()
	for k, v := range env {
		// In case of similar keys, last env var wins.
		environ = append(environ, k+"="+v)
	}

	log.Debugf("executing %s with env=%+v and args=%+v", w.path, env, args)

	child := exec.Command(w.path, args...)
	child.Env = environ
	child.Stdin = os.Stdin
	child.Stdout = os.Stdout
	child.Stderr = os.Stderr

	return trace.Wrap(child.Run(), "unable to execute tsh")
}

// GetTSHVersion queries the system tsh for its version.
func GetTSHVersion(w *Wrapper) (*semver.Version, error) {
	rawVersion, err := w.capture("version", "-f", "json")
	if err != nil {
		return nil, trace.Wrap(err, "querying tsh version")
	}

	versionInfo := struct {
		Version string `json:"version"`
	}{}
	if err := json.Unmarshal(rawVersion, &versionInfo); err != nil {
		return nil, trace.Wrap(err, "error deserializing tsh version from string: %s", rawVersion)
	}

	sv, err := semver.NewVersion(versionInfo.Version)
	if err != nil {
		return nil, trace.Wrap(err, "error parsing tsh version: %s", versionInfo.Version)
	}

	return sv, nil
}

// CheckTSHSupported checks if the current tsh supports Machine ID.
func CheckTSHSupported(w *Wrapper) error {
	version, err := GetTSHVersion(w)
	if err != nil {
		return trace.Wrap(err, "unable to determine tsh version")
	}

	minVersion := semver.New(TSHMinVersion)
	if version.LessThan(*minVersion) {
		return trace.BadParameter(
			"installed tsh version %s does not support Machine ID proxies, "+
				"please upgrade to at least %s",
			version, minVersion,
		)
	}

	log.Debugf("tsh version %s is supported", version)

	return nil
}

// GetDestination attempts to select an unambiguous destination, either from
// CLI or YAML config. It returns an error if the selected destination is
// invalid.
func GetDestination(botConfig *config.BotConfig, cf *config.CLIConf) (*config.DestinationConfig, error) {
	// Note: this only supports filesystem destinations.
	if cf.DestinationDir != "" {
		dest, err := botConfig.GetDestinationByPath(cf.DestinationDir)
		if err != nil {
			return nil, trace.Wrap(err, "unable to find destination %s in the "+
				"configuration; has the configuration file been "+
				"specified with `-c <path>`?", cf.DestinationDir)
		}

		return dest, nil
	}

	if len(botConfig.Destinations) == 0 {
		return nil, trace.BadParameter("either --destination-dir or a config file must be specified")
	} else if len(botConfig.Destinations) > 1 {
		return nil, trace.BadParameter("the config file contains multiple destinations; a --destination-dir must be specified")
	}

	return botConfig.Destinations[0], nil
}

// GetDestinationPath returns a path to a filesystem destination.
func GetDestinationPath(destination *config.DestinationConfig) (string, error) {
	destinationImpl, err := destination.GetDestination()
	if err != nil {
		return "", trace.Wrap(err)
	}

	destinationDir, ok := destinationImpl.(*config.DestinationDirectory)
	if !ok {
		return "", trace.BadParameter("destination %s must be a directory", destinationImpl)
	}

	return destinationDir.Path, nil
}

// GetTLSCATemplate returns the TLS CA template for the given destination. It's
// a required template so this should never fail.
func GetTLSCATemplate(destination *config.DestinationConfig) (*config.TemplateTLSCAs, error) {
	tpl := destination.GetConfigByName(config.TemplateTLSCAsName)
	if tpl == nil {
		return nil, trace.NotFound("no template with name %s found, this is a bug", config.TemplateTLSCAsName)
	}

	tlsCAs, ok := tpl.(*config.TemplateTLSCAs)
	if !ok {
		return nil, trace.BadParameter("invalid TLS CA template")
	}

	return tlsCAs, nil
}

// GetIdentityTemplate returns the identity template for the given destination.
// This is a required template so it _should_ never fail.
func GetIdentityTemplate(destination *config.DestinationConfig) (*config.TemplateIdentity, error) {
	tpl := destination.GetConfigByName(config.TemplateIdentityName)
	if tpl == nil {
		return nil, trace.NotFound("no template with name %s found, this is a bug", config.TemplateIdentityName)
	}

	identity, ok := tpl.(*config.TemplateIdentity)
	if !ok {
		return nil, trace.BadParameter("invalid identity template")
	}

	return identity, nil
}

// mergeEnv applies the given value to each key inside the specified map.
func mergeEnv(m map[string]string, value string, keys []string) {
	for _, key := range keys {
		m[key] = value
	}
}

// GetEnvForTSH returns a map of environment variables needed to properly wrap
// tsh so that it uses our Machine ID certificates where necessary.
func GetEnvForTSH(destination *config.DestinationConfig) (map[string]string, error) {
	tlsCAs, err := GetTLSCATemplate(destination)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	destPath, err := GetDestinationPath(destination)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	// The env var interface does allow us to set specific resource names for
	// everything but also has generic fallbacks. We'll use the fallbacks for
	// now but could eventually communicate more info to tsh if desired.
	env := make(map[string]string)
	mergeEnv(env, filepath.Join(destPath, identity.PrivateKeyKey), client.VirtualPathEnvNames(client.VirtualPathKey, nil))

	// Database certs are a bit awkward since a few databases (cockroach) have
	// special naming requirements. We can document around these for now and
	// automate later. (I don't think tsh handles this perfectly today anyway).
	mergeEnv(env, filepath.Join(destPath, identity.TLSCertKey), client.VirtualPathEnvNames(client.VirtualPathDatabase, nil))

	mergeEnv(env, filepath.Join(destPath, identity.TLSCertKey), client.VirtualPathEnvNames(client.VirtualPathApp, nil))

	// We don't want to provide a fallback for CAs since it would be ambiguous,
	// so we'll specify them exactly.
	env[client.VirtualPathEnvName(client.VirtualPathCA, client.VirtualPathCAParams(types.UserCA))] =
		filepath.Join(destPath, tlsCAs.UserCAPath)
	env[client.VirtualPathEnvName(client.VirtualPathCA, client.VirtualPathCAParams(types.HostCA))] =
		filepath.Join(destPath, tlsCAs.HostCAPath)
	env[client.VirtualPathEnvName(client.VirtualPathCA, client.VirtualPathCAParams(types.DatabaseCA))] =
		filepath.Join(destPath, tlsCAs.DatabaseCAPath)

	// TODO(timothyb89): Kubernetes support. We don't generate kubeconfigs yet, so we have
	// nothing to give tsh for now.

	return env, nil
}

// LoadIdentity loads a Teleport identity from an identityfile. Secondary bot
// identities are not loadable, so we'll just read the Teleport identity (which
// is required for tsh to function anyway).
func LoadIdentity(identityPath string) (*tlsca.Identity, error) {
	f, err := os.Open(identityPath)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer f.Close()

	idFile, err := identityfile.Read(f)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	cert, err := tlsca.ParseCertificatePEM(idFile.Certs.TLS)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	parsed, err := tlsca.FromSubject(cert.Subject, cert.NotAfter)
	return parsed, trace.Wrap(err)
}
<commit_msg>Fix broken version check in tbot's `tshwrap` (#13034)<commit_after>/*
Copyright 2022 Gravitational, Inc.

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 tshwrap

import (
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"

	"github.com/coreos/go-semver/semver"
	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/api/constants"
	"github.com/gravitational/teleport/api/identityfile"
	"github.com/gravitational/teleport/api/types"
	"github.com/gravitational/teleport/lib/client"
	"github.com/gravitational/teleport/lib/tlsca"
	"github.com/gravitational/teleport/tool/tbot/config"
	"github.com/gravitational/teleport/tool/tbot/identity"
	"github.com/gravitational/trace"
	"github.com/sirupsen/logrus"
)

const (
	// TSHVarName is the name of the environment variable that can override the
	// tsh path that would otherwise be located on the $PATH.
	TSHVarName = "TSH"

	// TSHMinVersion is the minimum version of tsh that supports Machine ID
	// proxies.
	TSHMinVersion = "9.3.0"
)

var log = logrus.WithFields(logrus.Fields{
	trace.Component: teleport.ComponentTBot,
})

// capture runs a command (presumably tsh) with the given arguments and
// returns it's captured stdout. Stderr is ignored. Errors are returned per
// exec.Command().Output() semantics.
func capture(tshPath string, args ...string) ([]byte, error) {
	out, err := exec.Command(tshPath, args...).Output()
	if err != nil {
		return nil, trace.Wrap(err, "error executing tsh")
	}

	return out, nil
}

// Wrapper is a wrapper to execute `tsh` commands via a subprocess.
type Wrapper struct {
	// path is a path to the tsh executable
	path string

	// capture is the function for capturing a command's output. It may be
	// overridden by tests for mocking purposes, but by default is expected to
	// execute an actual tsh binary on the host system.
	capture func(tshPath string, args ...string) ([]byte, error)
}

// New creates a new tsh wrapper. If a $TSH var is set it uses that path,
// otherwise looks for tsh on the OS path.
func New() (*Wrapper, error) {
	if val, ok := os.LookupEnv(TSHVarName); ok {
		return &Wrapper{
			path:    val,
			capture: capture,
		}, nil
	}

	binary := "tsh"
	if runtime.GOOS == constants.WindowsOS {
		binary = "tsh.exe"
	}

	path, err := exec.LookPath(binary)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	return &Wrapper{
		path:    path,
		capture: capture,
	}, nil
}

// Exec runs tsh with the given environment variables and arguments. The child
// process inherits stdin/stdout/stderr and runs until completion. Errors are
// returned per `exec.Command().Run()` semantics.
func (w *Wrapper) Exec(env map[string]string, args ...string) error {
	// The subprocess should inherit the environment plus our vars. Our env
	// vars will safely overwrite those from the environment, per `exec.Cmd`
	// docs.
	environ := os.Environ()
	for k, v := range env {
		// In case of similar keys, last env var wins.
		environ = append(environ, k+"="+v)
	}

	log.Debugf("executing %s with env=%+v and args=%+v", w.path, env, args)

	child := exec.Command(w.path, args...)
	child.Env = environ
	child.Stdin = os.Stdin
	child.Stdout = os.Stdout
	child.Stderr = os.Stderr

	return trace.Wrap(child.Run(), "unable to execute tsh")
}

// GetTSHVersion queries the system tsh for its version.
func GetTSHVersion(w *Wrapper) (*semver.Version, error) {
	rawVersion, err := w.capture(w.path, "version", "-f", "json")
	if err != nil {
		return nil, trace.Wrap(err, "querying tsh version")
	}

	versionInfo := struct {
		Version string `json:"version"`
	}{}
	if err := json.Unmarshal(rawVersion, &versionInfo); err != nil {
		return nil, trace.Wrap(err, "error deserializing tsh version from string: %s", rawVersion)
	}

	sv, err := semver.NewVersion(versionInfo.Version)
	if err != nil {
		return nil, trace.Wrap(err, "error parsing tsh version: %s", versionInfo.Version)
	}

	return sv, nil
}

// CheckTSHSupported checks if the current tsh supports Machine ID.
func CheckTSHSupported(w *Wrapper) error {
	version, err := GetTSHVersion(w)
	if err != nil {
		return trace.Wrap(err, "unable to determine tsh version")
	}

	minVersion := semver.New(TSHMinVersion)
	if version.LessThan(*minVersion) {
		return trace.BadParameter(
			"installed tsh version %s does not support Machine ID proxies, "+
				"please upgrade to at least %s",
			version, minVersion,
		)
	}

	log.Debugf("tsh version %s is supported", version)

	return nil
}

// GetDestination attempts to select an unambiguous destination, either from
// CLI or YAML config. It returns an error if the selected destination is
// invalid.
func GetDestination(botConfig *config.BotConfig, cf *config.CLIConf) (*config.DestinationConfig, error) {
	// Note: this only supports filesystem destinations.
	if cf.DestinationDir != "" {
		dest, err := botConfig.GetDestinationByPath(cf.DestinationDir)
		if err != nil {
			return nil, trace.Wrap(err, "unable to find destination %s in the "+
				"configuration; has the configuration file been "+
				"specified with `-c <path>`?", cf.DestinationDir)
		}

		return dest, nil
	}

	if len(botConfig.Destinations) == 0 {
		return nil, trace.BadParameter("either --destination-dir or a config file must be specified")
	} else if len(botConfig.Destinations) > 1 {
		return nil, trace.BadParameter("the config file contains multiple destinations; a --destination-dir must be specified")
	}

	return botConfig.Destinations[0], nil
}

// GetDestinationPath returns a path to a filesystem destination.
func GetDestinationPath(destination *config.DestinationConfig) (string, error) {
	destinationImpl, err := destination.GetDestination()
	if err != nil {
		return "", trace.Wrap(err)
	}

	destinationDir, ok := destinationImpl.(*config.DestinationDirectory)
	if !ok {
		return "", trace.BadParameter("destination %s must be a directory", destinationImpl)
	}

	return destinationDir.Path, nil
}

// GetTLSCATemplate returns the TLS CA template for the given destination. It's
// a required template so this should never fail.
func GetTLSCATemplate(destination *config.DestinationConfig) (*config.TemplateTLSCAs, error) {
	tpl := destination.GetConfigByName(config.TemplateTLSCAsName)
	if tpl == nil {
		return nil, trace.NotFound("no template with name %s found, this is a bug", config.TemplateTLSCAsName)
	}

	tlsCAs, ok := tpl.(*config.TemplateTLSCAs)
	if !ok {
		return nil, trace.BadParameter("invalid TLS CA template")
	}

	return tlsCAs, nil
}

// GetIdentityTemplate returns the identity template for the given destination.
// This is a required template so it _should_ never fail.
func GetIdentityTemplate(destination *config.DestinationConfig) (*config.TemplateIdentity, error) {
	tpl := destination.GetConfigByName(config.TemplateIdentityName)
	if tpl == nil {
		return nil, trace.NotFound("no template with name %s found, this is a bug", config.TemplateIdentityName)
	}

	identity, ok := tpl.(*config.TemplateIdentity)
	if !ok {
		return nil, trace.BadParameter("invalid identity template")
	}

	return identity, nil
}

// mergeEnv applies the given value to each key inside the specified map.
func mergeEnv(m map[string]string, value string, keys []string) {
	for _, key := range keys {
		m[key] = value
	}
}

// GetEnvForTSH returns a map of environment variables needed to properly wrap
// tsh so that it uses our Machine ID certificates where necessary.
func GetEnvForTSH(destination *config.DestinationConfig) (map[string]string, error) {
	tlsCAs, err := GetTLSCATemplate(destination)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	destPath, err := GetDestinationPath(destination)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	// The env var interface does allow us to set specific resource names for
	// everything but also has generic fallbacks. We'll use the fallbacks for
	// now but could eventually communicate more info to tsh if desired.
	env := make(map[string]string)
	mergeEnv(env, filepath.Join(destPath, identity.PrivateKeyKey), client.VirtualPathEnvNames(client.VirtualPathKey, nil))

	// Database certs are a bit awkward since a few databases (cockroach) have
	// special naming requirements. We can document around these for now and
	// automate later. (I don't think tsh handles this perfectly today anyway).
	mergeEnv(env, filepath.Join(destPath, identity.TLSCertKey), client.VirtualPathEnvNames(client.VirtualPathDatabase, nil))

	mergeEnv(env, filepath.Join(destPath, identity.TLSCertKey), client.VirtualPathEnvNames(client.VirtualPathApp, nil))

	// We don't want to provide a fallback for CAs since it would be ambiguous,
	// so we'll specify them exactly.
	env[client.VirtualPathEnvName(client.VirtualPathCA, client.VirtualPathCAParams(types.UserCA))] =
		filepath.Join(destPath, tlsCAs.UserCAPath)
	env[client.VirtualPathEnvName(client.VirtualPathCA, client.VirtualPathCAParams(types.HostCA))] =
		filepath.Join(destPath, tlsCAs.HostCAPath)
	env[client.VirtualPathEnvName(client.VirtualPathCA, client.VirtualPathCAParams(types.DatabaseCA))] =
		filepath.Join(destPath, tlsCAs.DatabaseCAPath)

	// TODO(timothyb89): Kubernetes support. We don't generate kubeconfigs yet, so we have
	// nothing to give tsh for now.

	return env, nil
}

// LoadIdentity loads a Teleport identity from an identityfile. Secondary bot
// identities are not loadable, so we'll just read the Teleport identity (which
// is required for tsh to function anyway).
func LoadIdentity(identityPath string) (*tlsca.Identity, error) {
	f, err := os.Open(identityPath)
	if err != nil {
		return nil, trace.Wrap(err)
	}
	defer f.Close()

	idFile, err := identityfile.Read(f)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	cert, err := tlsca.ParseCertificatePEM(idFile.Certs.TLS)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	parsed, err := tlsca.FromSubject(cert.Subject, cert.NotAfter)
	return parsed, trace.Wrap(err)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	"github.com/bradfitz/slice"
)

// Character is the type representing a role playing character
type Character struct {
	Name            string
	Backgrounds     map[string]Background
	Aptitudes       map[string]Aptitude
	Characteristics map[string]Characteristic
	Skills          map[string]Skill
	Talents         map[string]Talent
	Gauges          map[string]Gauge
	Rules           map[string]Rule
	Experience      int
	Spent           int
}

// NewCharacter creates a new character from the given sheet and universe.
func NewCharacter(universe Universe, sheet Sheet) (Character, error) {

	// Create a character
	c := Character{
		Name:            sheet.Header.Name,
		Backgrounds:     make(map[string]Background),
		Aptitudes:       make(map[string]Aptitude),
		Characteristics: make(map[string]Characteristic),
		Skills:          make(map[string]Skill),
		Talents:         make(map[string]Talent),
		Gauges:          make(map[string]Gauge),
		Rules:           make(map[string]Rule),
		Experience:      0,
		Spent:           0,
	}

	// The characteristics described in the header of the sheet are parsed as upgrades
	for _, upgrade := range sheet.Characteristics {

		// Get the characteristic from the universe
		characteristic, found := universe.FindCharacteristic(upgrade.Name)
		if !found {
			return c, NewError(UndefinedCharacteristic, upgrade.Line)
		}

		// Check it is not already applied
		_, found = c.Characteristics[characteristic.Name]
		if found {
			return c, NewError(DuplicateCharacteristic, upgrade.Line)
		}

		// Apply the upgrade
		err := c.ApplyCharacteristicUpgrade(characteristic, upgrade)
		if err != nil {
			return c, err
		}
	}

	// Next are the backgrounds
	for typ, metas := range sheet.Header.Metas {

		for _, meta := range metas {

			// Find the background corresponding to the meta
			background, err := universe.FindBackground(typ, meta.Label)
			if err != nil {
				return c, err
			}

			err = c.ApplyBackground(background, universe)
			if err != nil {
				return c, err
			}
		}
	}

	// Next are the sessions
	for _, session := range sheet.Sessions {

		// Apply the experience gain if needed
		if session.Reward != nil {
			c.Experience += *session.Reward
		}

		// Apply each upgrade in order
		for _, upgrade := range session.Upgrades {
			err := c.ApplyUpgrade(upgrade, universe)
			if err != nil {
				return c, err
			}
		}
	}

	return c, nil
}

// CountMatchingAptitudes return the number of aptitudes of the given slice
// that are in the character's aptitudes.
func (character Character) CountMatchingAptitudes(aptitudes []Aptitude) int {

	count := 0
	for _, aptitude := range aptitudes {
		if _, found := character.Aptitudes[string(aptitude)]; found {
			count++
		}
	}
	return count
}

// ApplyBackground changes the character's trait according to the history values
func (character *Character) ApplyBackground(background Background, universe Universe) error {

	cost := 0

	// For each upgrade associated to the history, apply each option.
	for _, upgrade := range background.Upgrades {
		err := character.ApplyUpgrade(Upgrade{
			Mark: MarkSpecial,
			Name: upgrade,
			Cost: &cost,
		}, universe)
		if err != nil {
			return err
		}
	}

	// Add the background to the character's backgrounds
	character.Backgrounds[background.Name] = background

	return nil
}

// ApplyUpgrade changes the character's attributes according to the given upgrade.
func (character *Character) ApplyUpgrade(upgrade Upgrade, universe Universe) error {

	var err error

	// Find the attribute corresponding to the upgrade., and initialize a new
	// rule if there isn't any.
	coster, found := universe.FindCoster(upgrade.Name)
	if !found {
		coster = Rule{
			Name: upgrade.Name,
		}
	}

	// If no cost is defined, compute it on the fly.
	if upgrade.Cost == nil {
		cost, err := coster.Cost(universe, *character)
		if err != nil {
			return err
		}

		upgrade.Cost = &cost
	}

	// Update the spent experience.
	character.Spent += *upgrade.Cost

	// Apply the upgrade depending on the target attribute.
	switch attribute := coster.(type) {
	case Characteristic:
		err = character.ApplyCharacteristicUpgrade(attribute, upgrade)
	case Skill:
		err = character.ApplySkillUpgrade(attribute, upgrade)
	case Talent:
		err = character.ApplyTalentUpgrade(attribute, upgrade)
	case Aptitude:
		err = character.ApplyAptitudeUpgrade(attribute, upgrade)
	case Gauge:
		err = character.ApplyGaugeUpgrade(attribute, upgrade)
	case Rule:
		err = character.ApplyRuleUpgrade(attribute, upgrade)
	}

	return err
}

func (character *Character) ApplyCharacteristicUpgrade(characteristic Characteristic, upgrade Upgrade) error {

	// Get the attribute from the character's characteristic map.
	c, found := character.Characteristics[characteristic.Name]
	if !found {
		c = characteristic
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		c.Tier++
	}

	// Parse the characteristic's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, characteristic.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidCharacteristicValue)
	}

	// Update the characteristic value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		c.Value += value
	} else {
		c.Value = value
	}

	character.Characteristics[c.Name] = c

	return nil
}

func (character *Character) ApplySkillUpgrade(skill Skill, upgrade Upgrade) error {

	// Get the skill from the character's skill map.
	s, found := character.Skills[skill.FullName()]
	if !found {
		s = skill
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		s.Tier++
	}

	// Put the skill back on the map.
	character.Skills[skill.FullName()] = s

	return nil
}

func (character *Character) ApplyTalentUpgrade(talent Talent, upgrade Upgrade) error {

	// Get the talent from the character.
	t, found := character.Talents[talent.FullName()]
	if !found {
		t = talent
	}

	// Increment the value of the talent.
	t.Value++

	// Put it back on the map.
	character.Talents[talent.FullName()] = t

	return nil
}

func (character *Character) ApplyAptitudeUpgrade(aptitude Aptitude, upgrade Upgrade) error {

	// Add the aptitude to the character's aptitudes.
	character.Aptitudes[string(aptitude)] = aptitude

	return nil
}

func (character *Character) ApplyGaugeUpgrade(gauge Gauge, upgrade Upgrade) error {

	// Get the gauge from the character.
	g, found := character.Gauges[gauge.Name]
	if !found {
		g = gauge
	}

	// Parse the gauge's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, g.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidGaugeValue)
	}

	// Update the gauge value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		g.Value += value
	} else {
		g.Value = value
	}

	// Set the gauge back on the map.
	character.Gauges[g.Name] = g

	return nil
}

func (character *Character) ApplyRuleUpgrade(rule Rule, upgrade Upgrade) error {

	// Add the rule to the character's rules.
	character.Rules[rule.Name] = rule

	return nil
}

// Print the character sheet on the screen
func (character Character) Print() {
	// Print the name
	fmt.Printf("%s\t%s\n", theme.Title("Name"), character.Name)

	// Print the backgrounds
	backgrounds := []Background{}

	for _, background := range character.Backgrounds {
		backgrounds = append(backgrounds, background)
	}

	slice.Sort(backgrounds, func(i, j int) bool {
		if backgrounds[i].Type != backgrounds[j].Type {
			return backgrounds[i].Type < backgrounds[j].Type
		}

		return backgrounds[i].Name < backgrounds[j].Name
	})

	for _, background := range backgrounds {
		fmt.Printf("%s\t%s\n", theme.Title(strings.Title(background.Type)), strings.Title(background.Name))
	}

	// Print the experience
	fmt.Printf("\n%s\t%d/%d\n", theme.Title("Experience"), character.Spent, character.Experience)

	// Print the characteristics
	fmt.Printf("\n%s\n", theme.Title("Characteristics"))

	characteristics := []Characteristic{}
	for _, characteristic := range character.Characteristics {
		characteristics = append(characteristics, characteristic)
	}

	slice.Sort(characteristics, func(i, j int) bool {
		return characteristics[i].Name < characteristics[j].Name
	})

	for _, characteristic := range characteristics {
		fmt.Printf("%s\t%s\n", characteristic.Name, theme.Value(characteristic.Value))
	}

	// Print the gauges
	fmt.Printf("\n%s\n", theme.Title("Gauges"))

	gauges := []Gauge{}
	for _, gauge := range character.Gauges {
		gauges = append(gauges, gauge)
	}

	slice.Sort(gauges, func(i, j int) bool {
		return gauges[i].Name < gauges[j].Name
	})

	for _, gauge := range gauges {
		fmt.Printf("%s\t%s\n", gauge.Name, theme.Value(gauge.Value))
	}

	// Print the talents
	fmt.Printf("\n%s\n", theme.Title("Talents"))

	talents := []Talent{}
	for _, talent := range character.Talents {
		talents = append(talents, talent)
	}

	slice.Sort(talents, func(i, j int) bool {
		return talents[i].Name < talents[j].Name
	})

	for _, talent := range talents {
		if talent.Value != 1 {
			fmt.Printf("%s (%s)\n", strings.Title(talent.Name), theme.Value(talent.Value))
		} else {
			fmt.Printf("%s\n", strings.Title(talent.Name))
		}
	}

	// Print the skills using a tabwriter
	fmt.Printf("\n%s\n", theme.Title("Skills"))

	skills := []Skill{}
	for _, skill := range character.Skills {
		skills = append(skills, skill)
	}

	slice.Sort(skills, func(i, j int) bool {
		return skills[i].Name < skills[j].Name
	})

	w := tabwriter.NewWriter(os.Stdout, 10, 1, 2, ' ', 0)
	for _, skill := range skills {
		fmt.Fprintf(w, "%s\t%s\n", strings.Title(skill.Name), theme.Value(skill.Tier))
	}
	w.Flush()

	// Print the special rules
	fmt.Printf("\n%s\n", theme.Title("Rules"))

	rules := []Rule{}

	for _, rule := range character.Rules {
		rules = append(rules, rule)
	}

	slice.Sort(rules, func(i, j int) bool {
		return rules[i].Name < rules[j].Name
	})

	for _, rule := range rules {
		fmt.Printf("%s\t%s\n", strings.Title(rule.Name), rule.Description)
	}
}
<commit_msg>Print skills before talents<commit_after>package main

import (
	"fmt"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"

	"github.com/bradfitz/slice"
)

// Character is the type representing a role playing character
type Character struct {
	Name            string
	Backgrounds     map[string]Background
	Aptitudes       map[string]Aptitude
	Characteristics map[string]Characteristic
	Skills          map[string]Skill
	Talents         map[string]Talent
	Gauges          map[string]Gauge
	Rules           map[string]Rule
	Experience      int
	Spent           int
}

// NewCharacter creates a new character from the given sheet and universe.
func NewCharacter(universe Universe, sheet Sheet) (Character, error) {

	// Create a character
	c := Character{
		Name:            sheet.Header.Name,
		Backgrounds:     make(map[string]Background),
		Aptitudes:       make(map[string]Aptitude),
		Characteristics: make(map[string]Characteristic),
		Skills:          make(map[string]Skill),
		Talents:         make(map[string]Talent),
		Gauges:          make(map[string]Gauge),
		Rules:           make(map[string]Rule),
		Experience:      0,
		Spent:           0,
	}

	// The characteristics described in the header of the sheet are parsed as upgrades
	for _, upgrade := range sheet.Characteristics {

		// Get the characteristic from the universe
		characteristic, found := universe.FindCharacteristic(upgrade.Name)
		if !found {
			return c, NewError(UndefinedCharacteristic, upgrade.Line)
		}

		// Check it is not already applied
		_, found = c.Characteristics[characteristic.Name]
		if found {
			return c, NewError(DuplicateCharacteristic, upgrade.Line)
		}

		// Apply the upgrade
		err := c.ApplyCharacteristicUpgrade(characteristic, upgrade)
		if err != nil {
			return c, err
		}
	}

	// Next are the backgrounds
	for typ, metas := range sheet.Header.Metas {

		for _, meta := range metas {

			// Find the background corresponding to the meta
			background, err := universe.FindBackground(typ, meta.Label)
			if err != nil {
				return c, err
			}

			err = c.ApplyBackground(background, universe)
			if err != nil {
				return c, err
			}
		}
	}

	// Next are the sessions
	for _, session := range sheet.Sessions {

		// Apply the experience gain if needed
		if session.Reward != nil {
			c.Experience += *session.Reward
		}

		// Apply each upgrade in order
		for _, upgrade := range session.Upgrades {
			err := c.ApplyUpgrade(upgrade, universe)
			if err != nil {
				return c, err
			}
		}
	}

	return c, nil
}

// CountMatchingAptitudes return the number of aptitudes of the given slice
// that are in the character's aptitudes.
func (character Character) CountMatchingAptitudes(aptitudes []Aptitude) int {

	count := 0
	for _, aptitude := range aptitudes {
		if _, found := character.Aptitudes[string(aptitude)]; found {
			count++
		}
	}
	return count
}

// ApplyBackground changes the character's trait according to the history values
func (character *Character) ApplyBackground(background Background, universe Universe) error {

	cost := 0

	// For each upgrade associated to the history, apply each option.
	for _, upgrade := range background.Upgrades {
		err := character.ApplyUpgrade(Upgrade{
			Mark: MarkSpecial,
			Name: upgrade,
			Cost: &cost,
		}, universe)
		if err != nil {
			return err
		}
	}

	// Add the background to the character's backgrounds
	character.Backgrounds[background.Name] = background

	return nil
}

// ApplyUpgrade changes the character's attributes according to the given upgrade.
func (character *Character) ApplyUpgrade(upgrade Upgrade, universe Universe) error {

	var err error

	// Find the attribute corresponding to the upgrade., and initialize a new
	// rule if there isn't any.
	coster, found := universe.FindCoster(upgrade.Name)
	if !found {
		coster = Rule{
			Name: upgrade.Name,
		}
	}

	// If no cost is defined, compute it on the fly.
	if upgrade.Cost == nil {
		cost, err := coster.Cost(universe, *character)
		if err != nil {
			return err
		}

		upgrade.Cost = &cost
	}

	// Update the spent experience.
	character.Spent += *upgrade.Cost

	// Apply the upgrade depending on the target attribute.
	switch attribute := coster.(type) {
	case Characteristic:
		err = character.ApplyCharacteristicUpgrade(attribute, upgrade)
	case Skill:
		err = character.ApplySkillUpgrade(attribute, upgrade)
	case Talent:
		err = character.ApplyTalentUpgrade(attribute, upgrade)
	case Aptitude:
		err = character.ApplyAptitudeUpgrade(attribute, upgrade)
	case Gauge:
		err = character.ApplyGaugeUpgrade(attribute, upgrade)
	case Rule:
		err = character.ApplyRuleUpgrade(attribute, upgrade)
	}

	return err
}

func (character *Character) ApplyCharacteristicUpgrade(characteristic Characteristic, upgrade Upgrade) error {

	// Get the attribute from the character's characteristic map.
	c, found := character.Characteristics[characteristic.Name]
	if !found {
		c = characteristic
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		c.Tier++
	}

	// Parse the characteristic's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, characteristic.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidCharacteristicValue)
	}

	// Update the characteristic value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		c.Value += value
	} else {
		c.Value = value
	}

	character.Characteristics[c.Name] = c

	return nil
}

func (character *Character) ApplySkillUpgrade(skill Skill, upgrade Upgrade) error {

	// Get the skill from the character's skill map.
	s, found := character.Skills[skill.FullName()]
	if !found {
		s = skill
	}

	// Increment the tier if the mark is default.
	if upgrade.Mark == MarkDefault {
		s.Tier++
	}

	// Put the skill back on the map.
	character.Skills[skill.FullName()] = s

	return nil
}

func (character *Character) ApplyTalentUpgrade(talent Talent, upgrade Upgrade) error {

	// Get the talent from the character.
	t, found := character.Talents[talent.FullName()]
	if !found {
		t = talent
	}

	// Increment the value of the talent.
	t.Value++

	// Put it back on the map.
	character.Talents[talent.FullName()] = t

	return nil
}

func (character *Character) ApplyAptitudeUpgrade(aptitude Aptitude, upgrade Upgrade) error {

	// Add the aptitude to the character's aptitudes.
	character.Aptitudes[string(aptitude)] = aptitude

	return nil
}

func (character *Character) ApplyGaugeUpgrade(gauge Gauge, upgrade Upgrade) error {

	// Get the gauge from the character.
	g, found := character.Gauges[gauge.Name]
	if !found {
		g = gauge
	}

	// Parse the gauge's upgrade value.
	raw := strings.TrimSpace(strings.TrimLeft(upgrade.Name, g.Name))
	value, err := strconv.Atoi(raw)
	if err != nil {
		return NewError(InvalidGaugeValue)
	}

	// Update the gauge value.
	if strings.HasPrefix(raw, "+") || strings.HasPrefix(raw, "-") {
		g.Value += value
	} else {
		g.Value = value
	}

	// Set the gauge back on the map.
	character.Gauges[g.Name] = g

	return nil
}

func (character *Character) ApplyRuleUpgrade(rule Rule, upgrade Upgrade) error {

	// Add the rule to the character's rules.
	character.Rules[rule.Name] = rule

	return nil
}

// Print the character sheet on the screen
func (character Character) Print() {
	// Print the name
	fmt.Printf("%s\t%s\n", theme.Title("Name"), character.Name)

	// Print the backgrounds
	backgrounds := []Background{}

	for _, background := range character.Backgrounds {
		backgrounds = append(backgrounds, background)
	}

	slice.Sort(backgrounds, func(i, j int) bool {
		if backgrounds[i].Type != backgrounds[j].Type {
			return backgrounds[i].Type < backgrounds[j].Type
		}

		return backgrounds[i].Name < backgrounds[j].Name
	})

	for _, background := range backgrounds {
		fmt.Printf("%s\t%s\n", theme.Title(strings.Title(background.Type)), strings.Title(background.Name))
	}

	// Print the experience
	fmt.Printf("\n%s\t%d/%d\n", theme.Title("Experience"), character.Spent, character.Experience)

	// Print the characteristics
	fmt.Printf("\n%s\n", theme.Title("Characteristics"))

	characteristics := []Characteristic{}
	for _, characteristic := range character.Characteristics {
		characteristics = append(characteristics, characteristic)
	}

	slice.Sort(characteristics, func(i, j int) bool {
		return characteristics[i].Name < characteristics[j].Name
	})

	for _, characteristic := range characteristics {
		fmt.Printf("%s\t%s\n", characteristic.Name, theme.Value(characteristic.Value))
	}

	// Print the gauges
	fmt.Printf("\n%s\n", theme.Title("Gauges"))

	gauges := []Gauge{}
	for _, gauge := range character.Gauges {
		gauges = append(gauges, gauge)
	}

	slice.Sort(gauges, func(i, j int) bool {
		return gauges[i].Name < gauges[j].Name
	})

	for _, gauge := range gauges {
		fmt.Printf("%s\t%s\n", gauge.Name, theme.Value(gauge.Value))
	}

	// Print the skills using a tabwriter
	fmt.Printf("\n%s\n", theme.Title("Skills"))

	skills := []Skill{}
	for _, skill := range character.Skills {
		skills = append(skills, skill)
	}

	slice.Sort(skills, func(i, j int) bool {
		return skills[i].FullName() < skills[j].FullName()
	})

	w := tabwriter.NewWriter(os.Stdout, 10, 1, 2, ' ', 0)
	for _, skill := range skills {
		fmt.Fprintf(w, "%s\t+%s\n", strings.Title(skill.FullName()), theme.Value(skill.Tier*10))
	}
	w.Flush()

	// Print the talents
	fmt.Printf("\n%s\n", theme.Title("Talents"))

	talents := []Talent{}
	for _, talent := range character.Talents {
		talents = append(talents, talent)
	}

	slice.Sort(talents, func(i, j int) bool {
		return talents[i].FullName() < talents[j].FullName()
	})

	for _, talent := range talents {
		if talent.Value != 1 {
			fmt.Printf("%s (%s)\n", strings.Title(talent.FullName()), theme.Value(talent.Value))
		} else {
			fmt.Printf("%s\n", strings.Title(talent.FullName()))
		}
	}

	// Print the special rules
	fmt.Printf("\n%s\n", theme.Title("Rules"))

	rules := []Rule{}

	for _, rule := range character.Rules {
		rules = append(rules, rule)
	}

	slice.Sort(rules, func(i, j int) bool {
		return rules[i].Name < rules[j].Name
	})

	for _, rule := range rules {
		fmt.Printf("%s\t%s\n", strings.Title(rule.Name), rule.Description)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/abbot/go-http-auth"
	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

var config Config

type Config struct {
	HttpBasicUsername string
	HttpBasicPassword string
	HttpBasicRealm    string
	Port              string
	Location          string
	ChecksUrl         string
	MeasurerCount     int
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
	SizeDownload      float64 `json:"size_download,omitempty"`
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	size_download, _ := easy.Getinfo(curl.INFO_SIZE_DOWNLOAD)
	m.SizeDownload = size_download.(float64)

	return m
}

func measurer(config Config, toMeasurer chan Check, toStreamer chan Measurement) {
	for {
		c := <-toMeasurer
		m := c.Measure(config)

		toStreamer <- m
	}
}

func streamer(config Config, toStreamer chan Measurement) {
	a := func(user, realm string) string {
		if user == config.HttpBasicUsername {
			return config.HttpBasicPassword
		}
		return ""
	}

	h := func(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
		w.Header().Set("Content-Type", "application/json")
		enc := json.NewEncoder(w)
		for {
			err := enc.Encode(<-toStreamer)
			if err != nil {
				return
			}

			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		}
	}

	authenticator := auth.NewBasicAuthenticator(config.HttpBasicRealm, a)

	http.HandleFunc("/measurements", authenticator.Wrap(h))

	log.Printf("fn=streamer listening=true port=%s\n", config.Port)
	err := http.ListenAndServe(":"+config.Port, nil)
	if err != nil {
		log.Fatal(err)
	}
}

func getChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func scheduler(check Check, toMeasurer chan Check) {
	for {
		toMeasurer <- check
		time.Sleep(1000 * time.Millisecond)
	}
}

func init() {
	flag.StringVar(&config.HttpBasicUsername, "http_basic_username", "", "HTTP basic authentication username")
	flag.StringVar(&config.HttpBasicPassword, "http_basic_password", "", "HTTP basic authentication password")
	flag.StringVar(&config.HttpBasicRealm, "http_basic_realm", "", "HTTP basic authentication realm")
	flag.StringVar(&config.Port, "port", "5000", "port the HTTP server should listen on")
	flag.StringVar(&config.Location, "location", "undefined", "location of this sensor")
	flag.StringVar(&config.ChecksUrl, "checks_url", "https://s3.amazonaws.com/canary-public-data/checks.json", "URL for check data")
	flag.IntVar(&config.MeasurerCount, "measurer_count", 1, "number of measurers to run")
}

func main() {
	flag.Parse()

	if len(config.HttpBasicUsername) == 0 && len(config.HttpBasicPassword) == 0 {
		log.Fatal("fatal - HTTP basic auth not set correctly")
	}

	check_list := getChecks(config)

	toMeasurer := make(chan Check)
	toStreamer := make(chan Measurement)

	// spawn one scheduler per check
	for _, c := range check_list {
		go scheduler(c, toMeasurer)
	}

	// spawn N measurers
	for i := 0; i < config.MeasurerCount; i++ {
		go measurer(config, toMeasurer, toStreamer)
	}

	// stream measurements to clients over HTTP
	go streamer(config, toStreamer)

	select {}
}
<commit_msg>lean on log.Fatal rather than panic<commit_after>package main

import (
	"encoding/json"
	"flag"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/abbot/go-http-auth"
	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

var config Config

type Config struct {
	HttpBasicUsername string
	HttpBasicPassword string
	HttpBasicRealm    string
	Port              string
	Location          string
	ChecksUrl         string
	MeasurerCount     int
}

type Check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type Measurement struct {
	Check             Check   `json:"check"`
	Id                string  `json:"id"`
	Location          string  `json:"location"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
	SizeDownload      float64 `json:"size_download,omitempty"`
}

func (c *Check) Measure(config Config) Measurement {
	var m Measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.Check = *c
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	size_download, _ := easy.Getinfo(curl.INFO_SIZE_DOWNLOAD)
	m.SizeDownload = size_download.(float64)

	return m
}

func measurer(config Config, toMeasurer chan Check, toStreamer chan Measurement) {
	for {
		c := <-toMeasurer
		m := c.Measure(config)

		toStreamer <- m
	}
}

func streamer(config Config, toStreamer chan Measurement) {
	a := func(user, realm string) string {
		if user == config.HttpBasicUsername {
			return config.HttpBasicPassword
		}
		return ""
	}

	h := func(w http.ResponseWriter, r *auth.AuthenticatedRequest) {
		w.Header().Set("Content-Type", "application/json")
		enc := json.NewEncoder(w)
		for {
			err := enc.Encode(<-toStreamer)
			if err != nil {
				return
			}

			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		}
	}

	authenticator := auth.NewBasicAuthenticator(config.HttpBasicRealm, a)

	http.HandleFunc("/measurements", authenticator.Wrap(h))

	log.Printf("fn=streamer listening=true port=%s\n", config.Port)
	err := http.ListenAndServe(":"+config.Port, nil)
	if err != nil {
		log.Fatal(err)
	}
}

func getChecks(config Config) []Check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		log.Fatal(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.Fatal(err)
	}

	var checks []Check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		log.Fatal(err)
	}

	return checks
}

func scheduler(check Check, toMeasurer chan Check) {
	for {
		toMeasurer <- check
		time.Sleep(1000 * time.Millisecond)
	}
}

func init() {
	flag.StringVar(&config.HttpBasicUsername, "http_basic_username", "", "HTTP basic authentication username")
	flag.StringVar(&config.HttpBasicPassword, "http_basic_password", "", "HTTP basic authentication password")
	flag.StringVar(&config.HttpBasicRealm, "http_basic_realm", "", "HTTP basic authentication realm")
	flag.StringVar(&config.Port, "port", "5000", "port the HTTP server should listen on")
	flag.StringVar(&config.Location, "location", "undefined", "location of this sensor")
	flag.StringVar(&config.ChecksUrl, "checks_url", "https://s3.amazonaws.com/canary-public-data/checks.json", "URL for check data")
	flag.IntVar(&config.MeasurerCount, "measurer_count", 1, "number of measurers to run")
}

func main() {
	flag.Parse()

	if len(config.HttpBasicUsername) == 0 && len(config.HttpBasicPassword) == 0 {
		log.Fatal("fatal - HTTP basic auth not set correctly")
	}

	check_list := getChecks(config)

	toMeasurer := make(chan Check)
	toStreamer := make(chan Measurement)

	// spawn one scheduler per check
	for _, c := range check_list {
		go scheduler(c, toMeasurer)
	}

	// spawn N measurers
	for i := 0; i < config.MeasurerCount; i++ {
		go measurer(config, toMeasurer, toStreamer)
	}

	// stream measurements to clients over HTTP
	go streamer(config, toStreamer)

	select {}
}
<|endoftext|>


<commit_before>package main

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

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location        string
	ChecksUrl       string
	MeasurementsUrl string
}

type check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type measurement struct {
	Id                string  `json:"id"`
	CheckId           string  `json:"check_id"`
	Location          string  `json:"location"`
	Url               string  `json:"url"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func GetEnvWithDefault(env string, def string) string {
	tmp := os.Getenv(env)

	if tmp == "" {
		return def
	}

	return tmp
}

func measure(config Config, c check) measurement {
	var m measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.CheckId = c.Id
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	m.Url = c.Url

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func measurer(config Config, checks chan check, measurements chan measurement) {
	for {
		c := <-checks
		m := measure(config, c)

		measurements <- m
	}
}

func recorder(config Config, measurements chan measurement) {
	payload := make([]measurement, 0, 100)
	for {
		m := <-measurements
		payload = append(payload, m)

		s, err := json.Marshal(&payload)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			panic(err)
		}
		resp.Body.Close()
		payload = make([]measurement, 0, 100)

		fmt.Println(resp)
	}
}

func get_checks(config Config) []check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func main() {
	var config Config
	config.Location = GetEnvWithDefault("LOCATION", "undefined")
	config.ChecksUrl = GetEnvWithDefault("CHECKS_URL", "https://s3.amazonaws.com/canary-public-data/data.json")
	config.MeasurementsUrl = GetEnvWithDefault("MEASUREMENTS_URL", "http://localhost:5000/measurements")

	fmt.Printf("%s\n", config.MeasurementsUrl)

	check_list := get_checks(config)

	checks := make(chan check)
	measurements := make(chan measurement)

	go measurer(config, checks, measurements)
	go recorder(config, measurements)

	for {
		for _, c := range check_list {
			checks <- c
		}

		time.Sleep(1000 * time.Millisecond)
	}
}
<commit_msg>set connection and transfer timeouts<commit_after>package main

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

	"github.com/andelf/go-curl"
	"github.com/nu7hatch/gouuid"
)

type Config struct {
	Location        string
	ChecksUrl       string
	MeasurementsUrl string
}

type check struct {
	Id  string `json:"id"`
	Url string `json:"url"`
}

type measurement struct {
	Id                string  `json:"id"`
	CheckId           string  `json:"check_id"`
	Location          string  `json:"location"`
	Url               string  `json:"url"`
	T                 int     `json:"t"`
	ExitStatus        int     `json:"exit_status"`
	ConnectTime       float64 `json:"connect_time,omitempty"`
	StartTransferTime float64 `json:"starttransfer_time,omitempty"`
	LocalIp           string  `json:"local_ip,omitempty"`
	PrimaryIp         string  `json:"primary_ip,omitempty"`
	TotalTime         float64 `json:"total_time,omitempty"`
	HttpStatus        int     `json:"http_status,omitempty"`
	NameLookupTime    float64 `json:"namelookup_time,omitempty"`
}

func GetEnvWithDefault(env string, def string) string {
	tmp := os.Getenv(env)

	if tmp == "" {
		return def
	}

	return tmp
}

func measure(config Config, c check) measurement {
	var m measurement

	id, _ := uuid.NewV4()
	m.Id = id.String()
	m.CheckId = c.Id
	m.Location = config.Location

	easy := curl.EasyInit()
	defer easy.Cleanup()

	easy.Setopt(curl.OPT_URL, c.Url)

	m.Url = c.Url

	// dummy func for curl output
	noOut := func(buf []byte, userdata interface{}) bool {
		return true
	}

	easy.Setopt(curl.OPT_WRITEFUNCTION, noOut)
	easy.Setopt(curl.OPT_CONNECTTIMEOUT, 10)
	easy.Setopt(curl.OPT_TIMEOUT, 10)

	now := time.Now()
	m.T = int(now.Unix())

	if err := easy.Perform(); err != nil {
		if e, ok := err.(curl.CurlError); ok {
			m.ExitStatus = (int(e))
			return m
		}
		os.Exit(1)
	}

	m.ExitStatus = 0
	http_status, _ := easy.Getinfo(curl.INFO_RESPONSE_CODE)
	m.HttpStatus = http_status.(int)

	connect_time, _ := easy.Getinfo(curl.INFO_CONNECT_TIME)
	m.ConnectTime = connect_time.(float64)

	namelookup_time, _ := easy.Getinfo(curl.INFO_NAMELOOKUP_TIME)
	m.NameLookupTime = namelookup_time.(float64)

	starttransfer_time, _ := easy.Getinfo(curl.INFO_STARTTRANSFER_TIME)
	m.StartTransferTime = starttransfer_time.(float64)

	total_time, _ := easy.Getinfo(curl.INFO_TOTAL_TIME)
	m.TotalTime = total_time.(float64)

	local_ip, _ := easy.Getinfo(curl.INFO_LOCAL_IP)
	m.LocalIp = local_ip.(string)

	primary_ip, _ := easy.Getinfo(curl.INFO_PRIMARY_IP)
	m.PrimaryIp = primary_ip.(string)

	return m
}

func measurer(config Config, checks chan check, measurements chan measurement) {
	for {
		c := <-checks
		m := measure(config, c)

		measurements <- m
	}
}

func recorder(config Config, measurements chan measurement) {
	payload := make([]measurement, 0, 100)
	for {
		m := <-measurements
		payload = append(payload, m)

		s, err := json.Marshal(&payload)
		if err != nil {
			panic(err)
		}

		body := bytes.NewBuffer(s)
		req, err := http.NewRequest("POST", config.MeasurementsUrl, body)
		if err != nil {
			panic(err)
		}

		req.Header.Add("Content-Type", "application/json")
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			panic(err)
		}
		resp.Body.Close()
		payload = make([]measurement, 0, 100)

		fmt.Println(resp)
	}
}

func get_checks(config Config) []check {
	url := config.ChecksUrl

	res, err := http.Get(url)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		panic(err)
	}

	var checks []check
	err = json.Unmarshal(body, &checks)
	if err != nil {
		panic(err)
	}

	return checks
}

func main() {
	var config Config
	config.Location = GetEnvWithDefault("LOCATION", "undefined")
	config.ChecksUrl = GetEnvWithDefault("CHECKS_URL", "https://s3.amazonaws.com/canary-public-data/data.json")
	config.MeasurementsUrl = GetEnvWithDefault("MEASUREMENTS_URL", "http://localhost:5000/measurements")

	fmt.Printf("%s\n", config.MeasurementsUrl)

	check_list := get_checks(config)

	checks := make(chan check)
	measurements := make(chan measurement)

	go measurer(config, checks, measurements)
	go recorder(config, measurements)

	for {
		for _, c := range check_list {
			checks <- c
		}

		time.Sleep(1000 * time.Millisecond)
	}
}
<|endoftext|>