<commit_before>/*
	Go Language Raspberry Pi Interface
	(c) Copyright David Thorpe 2016-2017
	All Rights Reserved

	Documentation http://djthorpe.github.io/gopi/
	For Licensing and Usage information, please see LICENSE.md
*/

package gopi

import (
	"flag"
	"fmt"
	"strings"
	"time"
)

type Flags struct {
	flagset *flag.FlagSet
	flagmap map[string]bool
	params  map[AppParam]interface{}
	name    string
}

////////////////////////////////////////////////////////////////////////////////
// PUBLIC METHODS

// Create a new flags object
func NewFlags(name string) *Flags {
	this := new(Flags)
	this.flagset = flag.NewFlagSet(name, flag.ContinueOnError)
	this.flagmap = nil
	this.name = name
	this.params = make(map[AppParam]interface{}, 10)
	return this
}

// Parse command line argumentsinto flags and pure arguments
func (this *Flags) Parse(args []string) error {

	// parse flags
	err := this.flagset.Parse(args)
	if err != nil {
		return err
	}

	// set hash of flags that were set
	this.flagmap = make(map[string]bool)
	this.flagset.Visit(func(f *flag.Flag) {
		this.flagmap[f.Name] = true
	})

	// return success
	return nil
}

// Parsed reports whether the command-line flags have been parsed
func (this *Flags) Parsed() bool {
	return this.flagset.Parsed()
}

// Name returns the name of the flagset (usually same as application)
func (this *Flags) Name() string {
	return this.name
}

// Args returns the command line arguments as an array which aren't flags
func (this *Flags) Args() []string {
	return this.flagset.Args()
}

// Flags returns the array of flags which were set on the command line
func (this *Flags) Flags() []string {
	if this.flagmap == nil {
		return []string{}
	}
	flags := make([]string, 0)
	for k := range this.flagmap {
		flags = append(flags, k)
	}
	return flags
}

// HasFlag returns a boolean indicating if a flag was set on the command line
func (this *Flags) HasFlag(name string) bool {
	if this.flagmap == nil {
		return false
	}
	_, exists := this.flagmap[name]
	return exists
}

// SetUsageFunc sets the usage function which prints
// usage information to stderr
func (this *Flags) SetUsageFunc(usage_func func(flags *Flags)) {
	this.flagset.Usage = func() {
		usage_func(this)
	}
}

// PrintUsage will call the usage function
func (this *Flags) PrintUsage() {
	this.flagset.Usage()
}

// PrintDefaults will output the flags to stderr
func (this *Flags) PrintDefaults() {
	this.flagset.PrintDefaults()
}

func (this *Flags) PrintVersion() {
	writer := this.flagset.Output()
	for k, v := range this.params {
		if k == PARAM_TIMESTAMP { // ignore timestamp
			continue
		} else if v_ := fmt.Sprint(v); v_ != "" {
			k_ := strings.TrimPrefix(fmt.Sprint(k), "PARAM_")
			fmt.Fprintf(writer, "%20s: %s\n", k_, v_)
		}
	}
}

// String returns a human-readable form of the Flags object
func (this *Flags) String() string {
	return fmt.Sprintf("<app.Flags>{ parsed=%v name=%v flags=%v args=%v }", this.Parsed(), this.Name(), this.Flags(), this.Args())
}

////////////////////////////////////////////////////////////////////////////////
// DEFINE FLAGS

// FlagString defines string flag and return pointer to the flag value
func (this *Flags) FlagString(name string, value string, usage string) *string {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.String(name, value, usage)
	}
}

// FlagBool defines a boolean flag and return pointer to the flag value
func (this *Flags) FlagBool(name string, value bool, usage string) *bool {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Bool(name, value, usage)
	}
}

// FlagDuration defines duration flag and return pointer to the flag value
func (this *Flags) FlagDuration(name string, value time.Duration, usage string) *time.Duration {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Duration(name, value, usage)
	}
}

// FlagInt defines integer flag and return pointer to the flag value
func (this *Flags) FlagInt(name string, value int, usage string) *int {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Int(name, value, usage)
	}
}

// FlagUint defines unsigned integer flag and return pointer to the flag value
func (this *Flags) FlagUint(name string, value uint, usage string) *uint {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Uint(name, value, usage)
	}
}

// FlagFloat64 defines float64 flag and return pointer to the flag value
func (this *Flags) FlagFloat64(name string, value float64, usage string) *float64 {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Float64(name, value, usage)
	}
}

////////////////////////////////////////////////////////////////////////////////
// GET FLAGS

// Get boolean value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetBool(name string) (bool, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return false, false
	}
	return value.Value.(flag.Getter).Get().(bool), this.HasFlag(name)
}

// Get string value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetString(name string) (string, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return "", false
	}
	return value.Value.(flag.Getter).Get().(string), this.HasFlag(name)
}

// Get duration value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetDuration(name string) (time.Duration, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return time.Duration(0), false
	}
	return value.Value.(flag.Getter).Get().(time.Duration), this.HasFlag(name)
}

// Get integer value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetInt(name string) (int, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(int), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint(name string) (uint, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(uint), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint16(name string) (uint16, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	uint_value := value.Value.(flag.Getter).Get().(uint)
	return uint16(uint_value), this.HasFlag(name)
}

// Get float64 value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetFloat64(name string) (float64, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0.0, false
	}
	return value.Value.(flag.Getter).Get().(float64), this.HasFlag(name)
}

////////////////////////////////////////////////////////////////////////////////
// SET FLAGS

// Set a flag string value. The flag must have previously been configured
// using FlagXX method. Will return an error if the value couldn't be parsed
func (this *Flags) SetString(name, value string) error {
	if flag := this.flagset.Lookup(name); flag == nil {
		return fmt.Errorf("SetString: No such flag: %v", name)
	} else {
		return flag.Value.Set(value)
	}
}

// Set a flag uint value
func (this *Flags) SetUint(name string, value uint) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag int value
func (this *Flags) SetInt(name string, value int) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag bool value
func (this *Flags) SetBool(name string, value bool) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag float64 value
func (this *Flags) SetFloat64(name string, value float64) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag duration value
func (this *Flags) SetDuration(name string, value time.Duration) error {
	return this.SetString(name, fmt.Sprint(value))
}

// SetParam set a parameter to an opaque type
func (this *Flags) SetParam(key AppParam, value interface{}) {
	this.params[key] = value
}
<commit_msg>Added GetParam method for flags<commit_after>/*
	Go Language Raspberry Pi Interface
	(c) Copyright David Thorpe 2016-2017
	All Rights Reserved

	Documentation http://djthorpe.github.io/gopi/
	For Licensing and Usage information, please see LICENSE.md
*/

package gopi

import (
	"flag"
	"fmt"
	"strings"
	"time"
)

type Flags struct {
	flagset *flag.FlagSet
	flagmap map[string]bool
	params  map[AppParam]interface{}
	name    string
}

////////////////////////////////////////////////////////////////////////////////
// PUBLIC METHODS

// Create a new flags object
func NewFlags(name string) *Flags {
	this := new(Flags)
	this.flagset = flag.NewFlagSet(name, flag.ContinueOnError)
	this.flagmap = nil
	this.name = name
	this.params = make(map[AppParam]interface{}, 10)
	return this
}

// Parse command line argumentsinto flags and pure arguments
func (this *Flags) Parse(args []string) error {

	// parse flags
	err := this.flagset.Parse(args)
	if err != nil {
		return err
	}

	// set hash of flags that were set
	this.flagmap = make(map[string]bool)
	this.flagset.Visit(func(f *flag.Flag) {
		this.flagmap[f.Name] = true
	})

	// return success
	return nil
}

// Parsed reports whether the command-line flags have been parsed
func (this *Flags) Parsed() bool {
	return this.flagset.Parsed()
}

// Name returns the name of the flagset (usually same as application)
func (this *Flags) Name() string {
	return this.name
}

// Args returns the command line arguments as an array which aren't flags
func (this *Flags) Args() []string {
	return this.flagset.Args()
}

// Flags returns the array of flags which were set on the command line
func (this *Flags) Flags() []string {
	if this.flagmap == nil {
		return []string{}
	}
	flags := make([]string, 0)
	for k := range this.flagmap {
		flags = append(flags, k)
	}
	return flags
}

// HasFlag returns a boolean indicating if a flag was set on the command line
func (this *Flags) HasFlag(name string) bool {
	if this.flagmap == nil {
		return false
	}
	_, exists := this.flagmap[name]
	return exists
}

// SetUsageFunc sets the usage function which prints
// usage information to stderr
func (this *Flags) SetUsageFunc(usage_func func(flags *Flags)) {
	this.flagset.Usage = func() {
		usage_func(this)
	}
}

// PrintUsage will call the usage function
func (this *Flags) PrintUsage() {
	this.flagset.Usage()
}

// PrintDefaults will output the flags to stderr
func (this *Flags) PrintDefaults() {
	this.flagset.PrintDefaults()
}

func (this *Flags) PrintVersion() {
	writer := this.flagset.Output()
	for k, v := range this.params {
		if k == PARAM_TIMESTAMP { // ignore timestamp
			continue
		} else if v_ := fmt.Sprint(v); v_ != "" {
			k_ := strings.TrimPrefix(fmt.Sprint(k), "PARAM_")
			fmt.Fprintf(writer, "%20s: %s\n", k_, v_)
		}
	}
}

// String returns a human-readable form of the Flags object
func (this *Flags) String() string {
	return fmt.Sprintf("<app.Flags>{ parsed=%v name=%v flags=%v args=%v }", this.Parsed(), this.Name(), this.Flags(), this.Args())
}

////////////////////////////////////////////////////////////////////////////////
// DEFINE FLAGS

// FlagString defines string flag and return pointer to the flag value
func (this *Flags) FlagString(name string, value string, usage string) *string {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.String(name, value, usage)
	}
}

// FlagBool defines a boolean flag and return pointer to the flag value
func (this *Flags) FlagBool(name string, value bool, usage string) *bool {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Bool(name, value, usage)
	}
}

// FlagDuration defines duration flag and return pointer to the flag value
func (this *Flags) FlagDuration(name string, value time.Duration, usage string) *time.Duration {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Duration(name, value, usage)
	}
}

// FlagInt defines integer flag and return pointer to the flag value
func (this *Flags) FlagInt(name string, value int, usage string) *int {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Int(name, value, usage)
	}
}

// FlagUint defines unsigned integer flag and return pointer to the flag value
func (this *Flags) FlagUint(name string, value uint, usage string) *uint {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Uint(name, value, usage)
	}
}

// FlagFloat64 defines float64 flag and return pointer to the flag value
func (this *Flags) FlagFloat64(name string, value float64, usage string) *float64 {
	if this.flagset == nil {
		return nil
	} else {
		return this.flagset.Float64(name, value, usage)
	}
}

////////////////////////////////////////////////////////////////////////////////
// GET FLAGS

// Get boolean value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetBool(name string) (bool, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return false, false
	}
	return value.Value.(flag.Getter).Get().(bool), this.HasFlag(name)
}

// Get string value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetString(name string) (string, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return "", false
	}
	return value.Value.(flag.Getter).Get().(string), this.HasFlag(name)
}

// Get duration value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetDuration(name string) (time.Duration, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return time.Duration(0), false
	}
	return value.Value.(flag.Getter).Get().(time.Duration), this.HasFlag(name)
}

// Get integer value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetInt(name string) (int, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(int), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint(name string) (uint, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(uint), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint16(name string) (uint16, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	uint_value := value.Value.(flag.Getter).Get().(uint)
	return uint16(uint_value), this.HasFlag(name)
}

// Get float64 value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetFloat64(name string) (float64, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0.0, false
	}
	return value.Value.(flag.Getter).Get().(float64), this.HasFlag(name)
}

////////////////////////////////////////////////////////////////////////////////
// SET FLAGS

// Set a flag string value. The flag must have previously been configured
// using FlagXX method. Will return an error if the value couldn't be parsed
func (this *Flags) SetString(name, value string) error {
	if flag := this.flagset.Lookup(name); flag == nil {
		return fmt.Errorf("SetString: No such flag: %v", name)
	} else {
		return flag.Value.Set(value)
	}
}

// Set a flag uint value
func (this *Flags) SetUint(name string, value uint) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag int value
func (this *Flags) SetInt(name string, value int) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag bool value
func (this *Flags) SetBool(name string, value bool) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag float64 value
func (this *Flags) SetFloat64(name string, value float64) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag duration value
func (this *Flags) SetDuration(name string, value time.Duration) error {
	return this.SetString(name, fmt.Sprint(value))
}

////////////////////////////////////////////////////////////////////////////////
// GET AND SET PARAMETERS

// SetParam set a parameter to an opaque type
func (this *Flags) SetParam(key AppParam, value interface{}) {
	this.params[key] = value
}

// GetParam returns a parameter value as an opaque type, or nil
func (this *Flags) GetParam(key AppParam) interface{} {
	if value, exists := this.params[key]; exists {
		return value
	} else {
		return nil
	}
}
<|endoftext|>


<commit_before>/*
	Go Language Raspberry Pi Interface
	(c) Copyright David Thorpe 2016-2017
	All Rights Reserved

	Documentation http://djthorpe.github.io/gopi/
	For Licensing and Usage information, please see LICENSE.md
*/

package gopi

import (
	"flag"
	"fmt"
	"time"
)

type Flags struct {
	flagset *flag.FlagSet
	flagmap map[string]bool
	name    string
}

////////////////////////////////////////////////////////////////////////////////
// PUBLIC METHODS

// Create a new flags object
func NewFlags(name string) *Flags {
	this := new(Flags)
	this.flagset = flag.NewFlagSet(name, flag.ContinueOnError)
	this.flagmap = nil
	this.name = name
	return this
}

// Parse command line argumentsinto flags and pure arguments
func (this *Flags) Parse(args []string) error {

	// parse flags
	err := this.flagset.Parse(args)
	if err != nil {
		return err
	}

	// set hash of flags that were set
	this.flagmap = make(map[string]bool)
	this.flagset.Visit(func(f *flag.Flag) {
		this.flagmap[f.Name] = true
	})

	// return success
	return nil
}

// Parsed reports whether the command-line flags have been parsed
func (this *Flags) Parsed() bool {
	return this.flagset.Parsed()
}

// Name returns the name of the flagset (usually same as application)
func (this *Flags) Name() string {
	return this.name
}

// Args returns the command line arguments as an array which aren't flags
func (this *Flags) Args() []string {
	return this.flagset.Args()
}

// Flags returns the array of flags which were set on the command line
func (this *Flags) Flags() []string {
	if this.flagmap == nil {
		return []string{}
	}
	flags := make([]string, 0)
	for k := range this.flagmap {
		flags = append(flags, k)
	}
	return flags
}

// HasFlag returns a boolean indicating if a flag was set on the command line
func (this *Flags) HasFlag(name string) bool {
	if this.flagmap == nil {
		return false
	}
	_, exists := this.flagmap[name]
	return exists
}

// SetUsageFunc sets the usage function which prints
// usage information to stderr
func (this *Flags) SetUsageFunc(usage_func func()) {
	this.flagset.Usage = usage_func
}

// PrintUsage will call the usage function
func (this *Flags) PrintUsage() {
	this.flagset.Usage()
}

// PrintDefaults will output the flags to stderr
func (this *Flags) PrintDefaults() {
	this.flagset.PrintDefaults()
}

// String returns a human-readable form of the Flags object
func (this *Flags) String() string {
	return fmt.Sprintf("<app.Flags>{ parsed=%v name=%v flags=%v args=%v }", this.Parsed(), this.Name(), this.Flags(), this.Args())
}

////////////////////////////////////////////////////////////////////////////////
// DEFINE FLAGS

// FlagString defines string flag and return pointer to the flag value
func (this *Flags) FlagString(name string, value string, usage string) *string {
	return this.flagset.String(name, value, usage)
}

// FlagBool defines a boolean flag and return pointer to the flag value
func (this *Flags) FlagBool(name string, value bool, usage string) *bool {
	return this.flagset.Bool(name, value, usage)
}

// FlagDuration defines duration flag and return pointer to the flag value
func (this *Flags) FlagDuration(name string, value time.Duration, usage string) *time.Duration {
	return this.flagset.Duration(name, value, usage)
}

// FlagInt defines integer flag and return pointer to the flag value
func (this *Flags) FlagInt(name string, value int, usage string) *int {
	return this.flagset.Int(name, value, usage)
}

// FlagUint defines unsigned integer flag and return pointer to the flag value
func (this *Flags) FlagUint(name string, value uint, usage string) *uint {
	return this.flagset.Uint(name, value, usage)
}

// FlagFloat64 defines float64 flag and return pointer to the flag value
func (this *Flags) FlagFloat64(name string, value float64, usage string) *float64 {
	return this.flagset.Float64(name, value, usage)
}

////////////////////////////////////////////////////////////////////////////////
// GET FLAGS

// Get boolean value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetBool(name string) (bool, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return false, false
	}
	return value.Value.(flag.Getter).Get().(bool), this.HasFlag(name)
}

// Get string value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetString(name string) (string, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return "", false
	}
	return value.Value.(flag.Getter).Get().(string), this.HasFlag(name)
}

// Get duration value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetDuration(name string) (time.Duration, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return time.Duration(0), false
	}
	return value.Value.(flag.Getter).Get().(time.Duration), this.HasFlag(name)
}

// Get integer value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetInt(name string) (int, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(int), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint(name string) (uint, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(uint), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint16(name string) (uint16, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	uint_value := value.Value.(flag.Getter).Get().(uint)
	return uint16(uint_value), this.HasFlag(name)
}

// Get float64 value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetFloat64(name string) (float64, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0.0, false
	}
	return value.Value.(flag.Getter).Get().(float64), this.HasFlag(name)
}
<commit_msg>Updated flags to add SetString method and other types<commit_after>/*
	Go Language Raspberry Pi Interface
	(c) Copyright David Thorpe 2016-2017
	All Rights Reserved

	Documentation http://djthorpe.github.io/gopi/
	For Licensing and Usage information, please see LICENSE.md
*/

package gopi

import (
	"flag"
	"fmt"
	"time"
)

type Flags struct {
	flagset *flag.FlagSet
	flagmap map[string]bool
	name    string
}

////////////////////////////////////////////////////////////////////////////////
// PUBLIC METHODS

// Create a new flags object
func NewFlags(name string) *Flags {
	this := new(Flags)
	this.flagset = flag.NewFlagSet(name, flag.ContinueOnError)
	this.flagmap = nil
	this.name = name
	return this
}

// Parse command line argumentsinto flags and pure arguments
func (this *Flags) Parse(args []string) error {

	// parse flags
	err := this.flagset.Parse(args)
	if err != nil {
		return err
	}

	// set hash of flags that were set
	this.flagmap = make(map[string]bool)
	this.flagset.Visit(func(f *flag.Flag) {
		this.flagmap[f.Name] = true
	})

	// return success
	return nil
}

// Parsed reports whether the command-line flags have been parsed
func (this *Flags) Parsed() bool {
	return this.flagset.Parsed()
}

// Name returns the name of the flagset (usually same as application)
func (this *Flags) Name() string {
	return this.name
}

// Args returns the command line arguments as an array which aren't flags
func (this *Flags) Args() []string {
	return this.flagset.Args()
}

// Flags returns the array of flags which were set on the command line
func (this *Flags) Flags() []string {
	if this.flagmap == nil {
		return []string{}
	}
	flags := make([]string, 0)
	for k := range this.flagmap {
		flags = append(flags, k)
	}
	return flags
}

// HasFlag returns a boolean indicating if a flag was set on the command line
func (this *Flags) HasFlag(name string) bool {
	if this.flagmap == nil {
		return false
	}
	_, exists := this.flagmap[name]
	return exists
}

// SetUsageFunc sets the usage function which prints
// usage information to stderr
func (this *Flags) SetUsageFunc(usage_func func()) {
	this.flagset.Usage = usage_func
}

// PrintUsage will call the usage function
func (this *Flags) PrintUsage() {
	this.flagset.Usage()
}

// PrintDefaults will output the flags to stderr
func (this *Flags) PrintDefaults() {
	this.flagset.PrintDefaults()
}

// String returns a human-readable form of the Flags object
func (this *Flags) String() string {
	return fmt.Sprintf("<app.Flags>{ parsed=%v name=%v flags=%v args=%v }", this.Parsed(), this.Name(), this.Flags(), this.Args())
}

////////////////////////////////////////////////////////////////////////////////
// DEFINE FLAGS

// FlagString defines string flag and return pointer to the flag value
func (this *Flags) FlagString(name string, value string, usage string) *string {
	return this.flagset.String(name, value, usage)
}

// FlagBool defines a boolean flag and return pointer to the flag value
func (this *Flags) FlagBool(name string, value bool, usage string) *bool {
	return this.flagset.Bool(name, value, usage)
}

// FlagDuration defines duration flag and return pointer to the flag value
func (this *Flags) FlagDuration(name string, value time.Duration, usage string) *time.Duration {
	return this.flagset.Duration(name, value, usage)
}

// FlagInt defines integer flag and return pointer to the flag value
func (this *Flags) FlagInt(name string, value int, usage string) *int {
	return this.flagset.Int(name, value, usage)
}

// FlagUint defines unsigned integer flag and return pointer to the flag value
func (this *Flags) FlagUint(name string, value uint, usage string) *uint {
	return this.flagset.Uint(name, value, usage)
}

// FlagFloat64 defines float64 flag and return pointer to the flag value
func (this *Flags) FlagFloat64(name string, value float64, usage string) *float64 {
	return this.flagset.Float64(name, value, usage)
}

////////////////////////////////////////////////////////////////////////////////
// GET FLAGS

// Get boolean value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetBool(name string) (bool, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return false, false
	}
	return value.Value.(flag.Getter).Get().(bool), this.HasFlag(name)
}

// Get string value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetString(name string) (string, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return "", false
	}
	return value.Value.(flag.Getter).Get().(string), this.HasFlag(name)
}

// Get duration value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetDuration(name string) (time.Duration, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return time.Duration(0), false
	}
	return value.Value.(flag.Getter).Get().(time.Duration), this.HasFlag(name)
}

// Get integer value for a flag, and a boolean which indicates if the flag
// was set
func (this *Flags) GetInt(name string) (int, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(int), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint(name string) (uint, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	return value.Value.(flag.Getter).Get().(uint), this.HasFlag(name)
}

// Get unsigned integer value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetUint16(name string) (uint16, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0, false
	}
	uint_value := value.Value.(flag.Getter).Get().(uint)
	return uint16(uint_value), this.HasFlag(name)
}

// Get float64 value for a flag, and a boolean which indicates if
// the flag was set
func (this *Flags) GetFloat64(name string) (float64, bool) {
	value := this.flagset.Lookup(name)
	if value == nil {
		return 0.0, false
	}
	return value.Value.(flag.Getter).Get().(float64), this.HasFlag(name)
}

////////////////////////////////////////////////////////////////////////////////
// SET FLAGS

// Set a flag string value. The flag must have previously been configured
// using FlagXX method. Will return an error if the value couldn't be parsed
func (this *Flags) SetString(name, value string) error {
	if flag := this.flagset.Lookup(name); flag == nil {
		return fmt.Errorf("SetString: No such flag: %v", name)
	} else {
		return flag.Value.Set(value)
	}
}

// Set a flag uint value
func (this *Flags) SetUint(name string, value uint) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag int value
func (this *Flags) SetInt(name string, value int) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag bool value
func (this *Flags) SetBool(name string, value bool) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag float64 value
func (this *Flags) SetFloat64(name string, value float64) error {
	return this.SetString(name, fmt.Sprint(value))
}

// Set a flag duration value
func (this *Flags) SetDuration(name string, value time.Duration) error {
	return this.SetString(name, fmt.Sprint(value))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/gob"
	"encoding/json"
	"encoding/xml"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/bemasher/rtlamr/csv"
)

var logFilename = flag.String("logfile", "/dev/stdout", "log statement dump file")
var logFile *os.File

var sampleFilename = flag.String("samplefile", os.DevNull, "raw signal dump file")
var sampleFile *os.File

var msgType = flag.String("msgtype", "scm", "message type to receive: scm or idm")

var symbolLength = flag.Int("symbollength", 73, "symbol length in samples, see -help for valid lengths")

var timeLimit = flag.Duration("duration", 0, "time to run for, 0 for infinite")
var meterID = flag.Uint("filterid", 0, "display only messages matching given id")
var meterType = flag.Uint("filtertype", 0, "display only messages matching given type")

var encoder Encoder
var format = flag.String("format", "plain", "format to write log messages in: plain, csv, json, xml or gob")
var gobUnsafe = flag.Bool("gobunsafe", false, "allow gob output to stdout")

var quiet = flag.Bool("quiet", false, "suppress printing state information at startup")
var single = flag.Bool("single", false, "one shot execution")

func RegisterFlags() {
	// Override default center frequency.
	centerFreqFlag := flag.CommandLine.Lookup("centerfreq")
	centerFreqString := strconv.FormatUint(CenterFreq, 10)
	centerFreqFlag.DefValue = centerFreqString
	centerFreqFlag.Value.Set(centerFreqString)

	rtlamrFlags := map[string]bool{
		"logfile":      true,
		"samplefile":   true,
		"msgtype":      true,
		"symbollength": true,
		"duration":     true,
		"filterid":     true,
		"filtertype":   true,
		"format":       true,
		"gobunsafe":    true,
		"quiet":        true,
		"single":       true,
		// "help":         true,
	}

	printDefaults := func(validFlags map[string]bool, inclusion bool) {
		flag.CommandLine.VisitAll(func(f *flag.Flag) {
			if validFlags[f.Name] != inclusion {
				return
			}

			format := "  -%s=%s: %s\n"
			fmt.Fprintf(os.Stderr, format, f.Name, f.DefValue, f.Usage)
		})
	}

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		printDefaults(rtlamrFlags, true)

		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, "rtltcp specific:")
		printDefaults(rtlamrFlags, false)
	}
}

func HandleFlags() {
	var err error

	if *logFilename == "/dev/stdout" {
		logFile = os.Stdout
	} else {
		logFile, err = os.Create(*logFilename)
		if err != nil {
			log.Fatal("Error creating log file:", err)
		}
	}
	log.SetOutput(logFile)

	sampleFile, err = os.Create(*sampleFilename)
	if err != nil {
		log.Fatal("Error creating sample file:", err)
	}

	*format = strings.ToLower(*format)
	switch *format {
	case "plain":
		break
	case "csv":
		encoder = csv.NewEncoder(logFile)
	case "json":
		encoder = json.NewEncoder(logFile)
	case "xml":
		encoder = xml.NewEncoder(logFile)
	case "gob":
		encoder = gob.NewEncoder(logFile)
		if !*gobUnsafe && *logFilename == "/dev/stdout" {
			fmt.Println("Gob encoded messages are not stdout safe, specify non-stdout -logfile or use -gobunsafe.")
			os.Exit(1)
		}
	}
}

// JSON, XML and GOB all implement this interface so we can simplify log
// output formatting.
type Encoder interface {
	Encode(interface{}) error
}
<commit_msg>Add cpuprofile to rtlamr-specific flags.<commit_after>package main

import (
	"encoding/gob"
	"encoding/json"
	"encoding/xml"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"

	"github.com/bemasher/rtlamr/csv"
)

var logFilename = flag.String("logfile", "/dev/stdout", "log statement dump file")
var logFile *os.File

var sampleFilename = flag.String("samplefile", os.DevNull, "raw signal dump file")
var sampleFile *os.File

var msgType = flag.String("msgtype", "scm", "message type to receive: scm or idm")

var symbolLength = flag.Int("symbollength", 73, "symbol length in samples, see -help for valid lengths")

var timeLimit = flag.Duration("duration", 0, "time to run for, 0 for infinite, ex. 1h5m10s")
var meterID = flag.Uint("filterid", 0, "display only messages matching given id")
var meterType = flag.Uint("filtertype", 0, "display only messages matching given type")

var encoder Encoder
var format = flag.String("format", "plain", "format to write log messages in: plain, csv, json, xml or gob")
var gobUnsafe = flag.Bool("gobunsafe", false, "allow gob output to stdout")

var quiet = flag.Bool("quiet", false, "suppress printing state information at startup")
var single = flag.Bool("single", false, "one shot execution")

func RegisterFlags() {
	// Override default center frequency.
	centerFreqFlag := flag.CommandLine.Lookup("centerfreq")
	centerFreqString := strconv.FormatUint(CenterFreq, 10)
	centerFreqFlag.DefValue = centerFreqString
	centerFreqFlag.Value.Set(centerFreqString)

	rtlamrFlags := map[string]bool{
		"logfile":      true,
		"samplefile":   true,
		"msgtype":      true,
		"symbollength": true,
		"duration":     true,
		"filterid":     true,
		"filtertype":   true,
		"format":       true,
		"gobunsafe":    true,
		"quiet":        true,
		"single":       true,
		"cpuprofile":   true,
	}

	printDefaults := func(validFlags map[string]bool, inclusion bool) {
		flag.CommandLine.VisitAll(func(f *flag.Flag) {
			if validFlags[f.Name] != inclusion {
				return
			}

			format := "  -%s=%s: %s\n"
			fmt.Fprintf(os.Stderr, format, f.Name, f.DefValue, f.Usage)
		})
	}

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		printDefaults(rtlamrFlags, true)

		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, "rtltcp specific:")
		printDefaults(rtlamrFlags, false)
	}
}

func HandleFlags() {
	var err error

	if *logFilename == "/dev/stdout" {
		logFile = os.Stdout
	} else {
		logFile, err = os.Create(*logFilename)
		if err != nil {
			log.Fatal("Error creating log file:", err)
		}
	}
	log.SetOutput(logFile)

	sampleFile, err = os.Create(*sampleFilename)
	if err != nil {
		log.Fatal("Error creating sample file:", err)
	}

	*format = strings.ToLower(*format)
	switch *format {
	case "plain":
		break
	case "csv":
		encoder = csv.NewEncoder(logFile)
	case "json":
		encoder = json.NewEncoder(logFile)
	case "xml":
		encoder = xml.NewEncoder(logFile)
	case "gob":
		encoder = gob.NewEncoder(logFile)
		if !*gobUnsafe && *logFilename == "/dev/stdout" {
			fmt.Println("Gob encoded messages are not stdout safe, specify non-stdout -logfile or use -gobunsafe.")
			os.Exit(1)
		}
	}
}

// JSON, XML and GOB all implement this interface so we can simplify log
// output formatting.
type Encoder interface {
	Encode(interface{}) error
}
<|endoftext|>


<commit_before>package golib

import (
	"bytes"
	"fmt"
	"strings"
)

// Flags is a bit-mask type used in the RegisterFlags() function.
type Flags uint32

const (
	// FlagsAll makes RegisterFlags() enable all available flags.
	FlagsAll = 0xffffffff

	// FlagsLog enables flags that configure the logger (Package github.com/sirupsen/logrus).
	FlagsLog = 1 << iota

	// FlagsProfile enables flags the configure profiling of CPU and memory.
	FlagsProfile

	// FlagsTasks enables flags that help debugging the shutdown sequence Tasks and TaskGroups.
	FlagsTasks
)

// RegisterFlags registers various flags provided by the golib package, controlled
// by the bit-mask parameter.
func RegisterFlags(flags Flags) {
	if flags&FlagsLog != 0 {
		RegisterLogFlags()
	}
	if flags&FlagsProfile != 0 {
		RegisterProfileFlags()
	}
	if flags&FlagsTasks != 0 {
		RegisterTaskFlags()
	}
}

// StringSlice implements the flag.Value interface and stores every occurrence
// of the according flag in one string slice.
type StringSlice []string

// String implements the flag.Value interface by printing the contents of the underlying string slice.
func (i *StringSlice) String() string {
	return fmt.Sprintf("%v", *i)
}

// Set implements the flag.Value interface by adding the given string to the underlying string slice.
func (i *StringSlice) Set(value string) error {
	*i = append(*i, value)
	return nil
}

// KeyValueSeparator is used by KeyValueStringSlice to split 'key=value' parameters.
const KeyValueSeparator = "="

// KeyValueStringSlice implements the flag.Value interface. It expects value of the form 'key=value'
// and splits them into the corresponding parts.
type KeyValueStringSlice struct {
	Keys   []string
	Values []string
}

// String implements the flag.Value interface by printing all contains key-value pairs.
func (k *KeyValueStringSlice) String() string {
	return FormatOrderedMap(k.Keys, k.Values)
}

// Set implements the flag.Value interface by splitting the 'key=value' string
// and returning an error if the format is wrong.
func (k *KeyValueStringSlice) Set(value string) error {
	parts := strings.SplitN(value, KeyValueSeparator, 2)
	if len(parts) != 2 {
		return fmt.Errorf("Wrong format. Need key=value, got " + value)
	}
	k.Keys = append(k.Keys, parts[0])
	k.Values = append(k.Values, parts[1])
	return nil
}

// Map returns a map-representation of the contained key-value pairs.
func (k *KeyValueStringSlice) Map() map[string]string {
	result := make(map[string]string, len(k.Keys))
	for i, key := range k.Keys {
		result[key] = k.Values[i]
	}
	return result
}

// Put sets the given value to the first instance of the given key. All other instances of the
// given key remain unchanged. If the key is not yet present in the receiver, the new key-value pair
// is appended.
func (k *KeyValueStringSlice) Put(key, value string) {
	for i, storedKey := range k.Keys {
		if storedKey == key {
			k.Values[i] = value
			return
		}
	}
	k.Keys = append(k.Keys, key)
	k.Values = append(k.Values, value)
}

// Delete deletes all instances of the given key from the receiving KeyValueStringSlice. If the key
// is not present, the receiver remains unchanged.
func (k *KeyValueStringSlice) Delete(key string) {
	for i := 0; i < len(k.Keys); i++ {
		if k.Keys[i] == key {
			k.Keys = k.deleteIndex(i, k.Keys)
			k.Values = k.deleteIndex(i, k.Values)
		}
	}
}

func (k *KeyValueStringSlice) deleteIndex(i int, slice []string) []string {
	copy(slice[i:], slice[i+1:])
	return slice[:len(slice)-1]
}

// FormatMap returns a readable representation of the given string map.
func FormatMap(m map[string]string) string {
	keys := make([]string, 0, len(m))
	values := make([]string, 0, len(m))
	for k, v := range m {
		keys = append(keys, k)
		values = append(values, v)
	}
	return FormatOrderedMap(keys, values)
}

// FormatOrderedMap returns a readable representation of the given key-value pairs.
func FormatOrderedMap(keys []string, values []string) string {
	var buf bytes.Buffer
	started := false
	for i, val := range values {
		key := keys[i]
		if started {
			buf.WriteString(", ")
		}
		buf.WriteString(key)
		buf.WriteString("=")
		buf.WriteString(val)
		started = true
	}
	return buf.String()
}
<commit_msg>Added convenience functions to handle flags<commit_after>package golib

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"strings"
)

// Flags is a bit-mask type used in the RegisterFlags() function.
type Flags uint32

const (
	// FlagsAll makes RegisterFlags() enable all available flags.
	FlagsAll = 0xffffffff

	// FlagsLog enables flags that configure the logger (Package github.com/sirupsen/logrus).
	FlagsLog = 1 << iota

	// FlagsProfile enables flags the configure profiling of CPU and memory.
	FlagsProfile

	// FlagsTasks enables flags that help debugging the shutdown sequence Tasks and TaskGroups.
	FlagsTasks
)

// RegisterFlags registers various flags provided by the golib package, controlled
// by the bit-mask parameter.
func RegisterFlags(flags Flags) {
	if flags&FlagsLog != 0 {
		RegisterLogFlags()
	}
	if flags&FlagsProfile != 0 {
		RegisterProfileFlags()
	}
	if flags&FlagsTasks != 0 {
		RegisterTaskFlags()
	}
}

// StringSlice implements the flag.Value interface and stores every occurrence
// of the according flag in one string slice.
type StringSlice []string

// String implements the flag.Value interface by printing the contents of the underlying string slice.
func (i *StringSlice) String() string {
	return fmt.Sprintf("%v", *i)
}

// Set implements the flag.Value interface by adding the given string to the underlying string slice.
func (i *StringSlice) Set(value string) error {
	*i = append(*i, value)
	return nil
}

// KeyValueSeparator is used by KeyValueStringSlice to split 'key=value' parameters.
const KeyValueSeparator = "="

// KeyValueStringSlice implements the flag.Value interface. It expects value of the form 'key=value'
// and splits them into the corresponding parts.
type KeyValueStringSlice struct {
	Keys   []string
	Values []string
}

// String implements the flag.Value interface by printing all contains key-value pairs.
func (k *KeyValueStringSlice) String() string {
	return FormatOrderedMap(k.Keys, k.Values)
}

// Set implements the flag.Value interface by splitting the 'key=value' string
// and returning an error if the format is wrong.
func (k *KeyValueStringSlice) Set(value string) error {
	parts := strings.SplitN(value, KeyValueSeparator, 2)
	if len(parts) != 2 {
		return fmt.Errorf("Wrong format. Need key=value, got " + value)
	}
	k.Keys = append(k.Keys, parts[0])
	k.Values = append(k.Values, parts[1])
	return nil
}

// Map returns a map-representation of the contained key-value pairs.
func (k *KeyValueStringSlice) Map() map[string]string {
	result := make(map[string]string, len(k.Keys))
	for i, key := range k.Keys {
		result[key] = k.Values[i]
	}
	return result
}

// Put sets the given value to the first instance of the given key. All other instances of the
// given key remain unchanged. If the key is not yet present in the receiver, the new key-value pair
// is appended.
func (k *KeyValueStringSlice) Put(key, value string) {
	for i, storedKey := range k.Keys {
		if storedKey == key {
			k.Values[i] = value
			return
		}
	}
	k.Keys = append(k.Keys, key)
	k.Values = append(k.Values, value)
}

// Delete deletes all instances of the given key from the receiving KeyValueStringSlice. If the key
// is not present, the receiver remains unchanged.
func (k *KeyValueStringSlice) Delete(key string) {
	for i := 0; i < len(k.Keys); i++ {
		if k.Keys[i] == key {
			k.Keys = k.deleteIndex(i, k.Keys)
			k.Values = k.deleteIndex(i, k.Values)
		}
	}
}

func (k *KeyValueStringSlice) deleteIndex(i int, slice []string) []string {
	copy(slice[i:], slice[i+1:])
	return slice[:len(slice)-1]
}

// FormatMap returns a readable representation of the given string map.
func FormatMap(m map[string]string) string {
	keys := make([]string, 0, len(m))
	values := make([]string, 0, len(m))
	for k, v := range m {
		keys = append(keys, k)
		values = append(values, v)
	}
	return FormatOrderedMap(keys, values)
}

// FormatOrderedMap returns a readable representation of the given key-value pairs.
func FormatOrderedMap(keys []string, values []string) string {
	var buf bytes.Buffer
	started := false
	for i, val := range values {
		key := keys[i]
		if started {
			buf.WriteString(", ")
		}
		buf.WriteString(key)
		buf.WriteString("=")
		buf.WriteString(val)
		started = true
	}
	return buf.String()
}

// EscapeExistingFlags can be used before defining new flags to escape existing flags that have been defined
// by other packages or modules. This can be used to avoid collisions of flag names.
func EscapeExistingFlags(prefix string) {
	oldCommandLine := flag.CommandLine
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	oldCommandLine.VisitAll(func(f *flag.Flag) {
		flag.Var(f.Value, prefix+f.Name, f.Usage)
	})
}

// When packages or modules are loaded AFTER parsing flags, avoid collisions when flags are re-defined.
func ParseFlags() {
	flag.Parse()
	flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError)
}
<|endoftext|>


<commit_before>package signatures

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"math/big"

	"github.com/NebulousLabs/Andromeda/encoding"
)

type (
	PublicKey ecdsa.PublicKey
	SecretKey *ecdsa.PrivateKey
)

type Signature struct {
	R, S *big.Int
}

// GenerateKeyPair creates a public-secret keypair that can be used to sign and
// verify messages.
func GenerateKeyPair() (sk SecretKey, pk PublicKey, err error) {
	ecdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	sk = SecretKey(ecdsaKey)
	if err != nil {
		return
	}
	pk = PublicKey(sk.PublicKey)
	return
}

// SignBytes signs a message using a secret key.
func SignBytes(data []byte, sk SecretKey) (sig Signature, err error) {
	sig.R, sig.S, err = ecdsa.Sign(rand.Reader, sk, data)
	return
}

// VerifyBytes uses a public key and input data to verify a signature.
func VerifyBytes(data []byte, pubKey PublicKey, sig Signature) bool {
	return ecdsa.Verify((*ecdsa.PublicKey)(&pubKey), data, sig.R, sig.S)
}

// Signature.MarshalSia implements the Marshaler interface for Signatures.
func (s *Signature) MarshalSia() []byte {
	if s.R == nil || s.S == nil {
		return []byte{0, 0}
	}
	// pretend Signature is a tuple of []bytes
	// this lets us use Marshal instead of doing manual length-prefixing
	return encoding.Marshal(struct{ R, S []byte }{s.R.Bytes(), s.S.Bytes()})
}

// Signature.UnmarshalSia implements the Unmarshaler interface for Signatures.
func (s *Signature) UnmarshalSia(b []byte) int {
	// inverse of the struct trick used in Signature.MarshalSia
	str := struct{ R, S []byte }{}
	if encoding.Unmarshal(b, &str) != nil {
		return 0
	}
	s.R = new(big.Int).SetBytes(str.R)
	s.S = new(big.Int).SetBytes(str.S)
	return len(str.R) + len(str.S) + 2
}

// PublicKey.MarshalSia implements the Marshaler interface for PublicKeys.
func (pk PublicKey) MarshalSia() []byte {
	if pk.X == nil || pk.Y == nil {
		return []byte{0, 0}
	}
	// see Signature.MarshalSia
	return encoding.Marshal(struct{ X, Y []byte }{pk.X.Bytes(), pk.Y.Bytes()})
}

// PublicKey.UnmarshalSia implements the Unmarshaler interface for PublicKeys.
func (pk *PublicKey) UnmarshalSia(b []byte) int {
	// see Signature.UnmarshalSia
	str := struct{ X, Y []byte }{}
	if encoding.Unmarshal(b, &str) != nil {
		return 0
	}
	pk.X = new(big.Int).SetBytes(str.X)
	pk.Y = new(big.Int).SetBytes(str.Y)
	return len(str.X) + len(str.Y) + 2
}
<commit_msg>refactor signatures into using byte arrays<commit_after>package signatures

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"math/big"
)

type (
	SecretKey [32]byte
	PublicKey [64]byte
	Signature [64]byte
)

// GenerateKeyPair creates a public-secret keypair that can be used to sign and
// verify messages.
func GenerateKeyPair() (sk SecretKey, pk PublicKey, err error) {
	// Get the ecdsa keys.
	ecdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

	// Copy the secret key into a byte array.
	skBytes := ecdsaKey.D.Bytes()
	copy(sk[:], skBytes)

	// Copy the public key into a byte array.
	pkBytes := ecdsaKey.PublicKey.X.Bytes()
	pkBytes = append(pkBytes, ecdsaKey.PublicKey.Y.Bytes()...)
	copy(pk[:], pkBytes)

	return
}

// SignBytes signs a message using a secret key.
func SignBytes(data []byte, sk SecretKey) (sig Signature, err error) {
	// Convert sk to an ecdsa.PrivateKey
	ecdsaKey := new(ecdsa.PrivateKey)
	ecdsaKey.PublicKey.Curve = elliptic.P256()
	ecdsaKey.D = new(big.Int).SetBytes(sk[:])

	// Get the signature.
	r, s, err := ecdsa.Sign(rand.Reader, ecdsaKey, data)
	if err != nil {
		return
	}

	// Convert the signature to a byte array.
	sigBytes := r.Bytes()
	sigBytes = append(sigBytes, s.Bytes()...)
	copy(sig[:], sigBytes)

	return
}

// VerifyBytes uses a public key and input data to verify a signature.
func VerifyBytes(data []byte, pk PublicKey, sig Signature) bool {
	// Get the public key.
	ecdsaKey := new(ecdsa.PublicKey)
	ecdsaKey.Curve = elliptic.P256()
	ecdsaKey.X = new(big.Int).SetBytes(pk[:32])
	ecdsaKey.Y = new(big.Int).SetBytes(pk[32:])
	r := new(big.Int).SetBytes(sig[:32])
	s := new(big.Int).SetBytes(sig[32:])
	return ecdsa.Verify(ecdsaKey, data, r, s)
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/output"
	"gopkg.in/olivere/elastic.v2"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 5
	esMaxConns         = 20
	esRecvBuffer       = 100
	esSendBuffer       = 100
)

type Indexer struct {
	bulkService *elastic.BulkService
	indexPrefix string
	indexType   string
}

type Config struct {
	Hosts       []string `json:"hosts"`
	IndexPrefix string   `json:"index"`
	IndexType   string   `json:"indexType"`
}

type ESServer struct {
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
}

func init() {
	output.Register("elasticsearch", &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	})
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func indexDoc(ev *buffer.Event) *map[string]interface{} {
	return &*ev.Fields
}

func (i *Indexer) flush() error {
	numEvents := i.bulkService.NumberOfActions()

	if numEvents > 0 {
		log.Printf("Flushing %d event(s) to elasticsearch", numEvents)
		_, err := i.bulkService.Do()

		if err != nil {
			log.Printf("Unable to flush events: %s", err)
		}

		return err
	}

	return nil
}

func (i *Indexer) index(ev *buffer.Event) error {
	doc := indexDoc(ev)
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	request := elastic.NewBulkIndexRequest().Index(idx).Type(typ).Doc(doc)
	i.bulkService.Add(request)

	numEvents := i.bulkService.NumberOfActions()

	if numEvents < esSendBuffer {
		return nil
	}

	return i.flush()
}

func (e *ESServer) Init(config json.RawMessage, b buffer.Sender) error {
	var esConfig *Config
	if err := json.Unmarshal(config, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	return nil
}

func readInputChannel(idx *Indexer, receiveChan chan *buffer.Event) {
	for {
		// Drain the channel only if we have room
		if idx.bulkService.NumberOfActions() < esSendBuffer {
			select {
			case ev := <-receiveChan:
				idx.index(ev)
			}
		} else {
			log.Printf("Internal Elasticsearch buffer is full, waiting")
			time.Sleep(1 * time.Second)
		}
	}
}

func (es *ESServer) Start() error {
	var client *elastic.Client
	var err error

	for {
		client, err = elastic.NewClient(elastic.SetURL(es.hosts...))
		if err != nil {
			log.Printf("Error starting Elasticsearch: %s, will retry", err)
			time.Sleep(2 * time.Second)
			continue
		}

		break
	}

	log.Printf("Connected to Elasticsarch")

	service := elastic.NewBulkService(client)

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	// Create indexer
	idx := &Indexer{service, es.config.IndexPrefix, es.config.IndexType}

	// Loop events and publish to elasticsearch
	tick := time.NewTicker(time.Duration(esFlushInterval) * time.Second)

	go readInputChannel(idx, receiveChan)

	for {
		select {
		case <-tick.C:
			idx.flush()
		case <-es.term:
			tick.Stop()
			log.Println("Elasticsearch received term signal")
			break
		}
	}

	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<commit_msg>Measure Elasticsearch output rate<commit_after>package elasticsearch

import (
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/output"
	"github.com/paulbellamy/ratecounter"
	"gopkg.in/olivere/elastic.v2"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 5
	esMaxConns         = 20
	esRecvBuffer       = 100
	esSendBuffer       = 100
)

type Indexer struct {
	bulkService *elastic.BulkService
	indexPrefix string
	indexType   string
	RateCounter *ratecounter.RateCounter
}

type Config struct {
	Hosts       []string `json:"hosts"`
	IndexPrefix string   `json:"index"`
	IndexType   string   `json:"indexType"`
}

type ESServer struct {
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
}

func init() {
	output.Register("elasticsearch", &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	})
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func indexDoc(ev *buffer.Event) *map[string]interface{} {
	return &*ev.Fields
}

func (i *Indexer) flush() error {
	numEvents := i.bulkService.NumberOfActions()

	if numEvents > 0 {
		//log.Printf("Flushing %d event(s) to elasticsearch", numEvents)
		_, err := i.bulkService.Do()

		if err != nil {
			log.Printf("Unable to flush events: %s", err)
		}

		return err
	}

	return nil
}

func (i *Indexer) index(ev *buffer.Event) error {
	doc := indexDoc(ev)
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	request := elastic.NewBulkIndexRequest().Index(idx).Type(typ).Doc(doc)
	i.bulkService.Add(request)
	i.RateCounter.Incr(1)

	numEvents := i.bulkService.NumberOfActions()

	if numEvents < esSendBuffer {
		return nil
	}

	return i.flush()
}

func (e *ESServer) Init(config json.RawMessage, b buffer.Sender) error {
	var esConfig *Config
	if err := json.Unmarshal(config, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	return nil
}

func readInputChannel(idx *Indexer, receiveChan chan *buffer.Event) {
	for {
		// Drain the channel only if we have room
		if idx.bulkService.NumberOfActions() < esSendBuffer {
			select {
			case ev := <-receiveChan:
				idx.index(ev)
			}
		} else {
			log.Printf("Internal Elasticsearch buffer is full, waiting")
			time.Sleep(1 * time.Second)
		}
	}
}

func (es *ESServer) Start() error {
	var client *elastic.Client
	var err error

	for {
		client, err = elastic.NewClient(elastic.SetURL(es.hosts...))
		if err != nil {
			log.Printf("Error starting Elasticsearch: %s, will retry", err)
			time.Sleep(2 * time.Second)
			continue
		}

		break
	}

	log.Printf("Connected to Elasticsarch")

	service := elastic.NewBulkService(client)

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	rateCounter := ratecounter.NewRateCounter(1 * time.Second)

	// Create indexer
	idx := &Indexer{service, es.config.IndexPrefix, es.config.IndexType, rateCounter}

	// Loop events and publish to elasticsearch
	tick := time.NewTicker(time.Duration(esFlushInterval) * time.Second)

	go readInputChannel(idx, receiveChan)

	for {
		select {
		case <-tick.C:
			log.Printf("Current Elasticsearch output rate: %d/s\n", rateCounter.Rate())
			idx.flush()
		case <-es.term:
			tick.Stop()
			log.Println("Elasticsearch received term signal")
			break
		}
	}

	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<|endoftext|>


<commit_before>package similarity

import (
	"bytes"
	"encoding/json"
	"io"
	"math"
	"sync"
)

// An item that has been rated.
type Item struct {
	Name  string
	Value float64
}

// The result of a similarity search. Name can be used to retrieve full information
// from the store and Similarity is the similarity score.
type Result struct {
	Name       string
	Similarity float64
}

// A function used to compare the similarity of two keys in the Similarity engine.
type comparison func(key1 string, key2 string) (distance float64)

// Similarity is a similarity storage and retrieval engine.
type Similarity struct {
	mutex sync.RWMutex
	data  map[string]map[string]Item
}

// Create a new Similarity engine.
func NewSimilarity() *Similarity {
	return &Similarity{data: make(map[string]map[string]Item)}
}

// Add an Item to the engine with a key.
func (sim *Similarity) Add(key string, item Item) {
	sim.mutex.Lock()
	defer sim.mutex.Unlock()
	m, ok := sim.data[key]
	if !ok {
		m = make(map[string]Item)
	}
	m[item.Name] = item
	sim.data[key] = m
}

// Get all the keys in this Similarity.
func (sim *Similarity) Keys() []string {
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	keys := make([]string, 0, len(sim.data))
	for k, _ := range sim.data {
		keys = append(keys, k)
	}
	return keys
}

// Dump the data backing a Similarity engine to a Writer.
func (sim *Similarity) WriteJson(w io.Writer) error {
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	b, err := json.Marshal(sim.data)
	if err != nil {
		return err
	}
	_, err = w.Write(b)
	if err != nil {
		return err
	}
	return nil
}

// Read the data for a Similarity engine from a Writer and load it
func (sim *Similarity) ReadJson(r io.Reader) error {
	// TODO: Is it idomatic to return underlying errors such as those encountered by ReadFrom or Unmarshal?
	b := new(bytes.Buffer)
	sim.mutex.Lock()
	defer sim.mutex.Unlock()
	_, err := b.ReadFrom(r)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b.Bytes(), &sim.data)
	if err != nil {
		return err
	}
	return nil
}

// Returns the Euclidean distance of two keys in our Similarity engine.
func (sim *Similarity) EuclideanDistance(key1 string, key2 string) float64 {
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	// Don't compute if either key is missing.
	if sim.data[key1] == nil || sim.data[key2] == nil {
		return -1
	}
	firstItems := sim.data[key1]
	secondItems := sim.data[key2]
	// Find common Items for the two keys
	sum := 0.0
	for _, item := range firstItems {
		secondItem, found := secondItems[item.Name]
		if found {
			sum += math.Pow(item.Value-secondItem.Value, 2)
		}
	}
	return 1 / (1 + math.Sqrt(sum))
}

// Find similar keys using Euclidean distance comparison.
func (sim *Similarity) SimilarEuclidean(key string, limit int) []Result {
	return sim.Similar(key, limit, sim.EuclideanDistance)
}

// Find similar keys using the provided distance comparison.
func (sim *Similarity) Similar(key string, limit int, distance comparison) []Result {
	results := make([]Result, 0)
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	for k, _ := range sim.data {
		if k == key {
			// Don't check ourselves.
			continue
		}
		score := distance(key, k)
		// TODO: replace -1 return value with an error in EuclideanDistance.
		if score != -1 {
			results = append(results, Result{k, score})
		}
	}
	if len(results) > limit {
		return results[:limit]
	} else {
		return results
	}

}
<commit_msg>Sort results before returning, highest scores first.<commit_after>package similarity

import (
	"bytes"
	"encoding/json"
	"io"
	"math"
	"sync"
	"sort"
)

// An item that has been rated.
type Item struct {
	Name  string
	Value float64
}

// The result of a similarity search. Name can be used to retrieve full information
// from the store and Similarity is the similarity score.
type Result struct {
	Name       string
	Similarity float64
}

// A function used to compare the similarity of two keys in the Similarity engine.
type comparison func(key1 string, key2 string) (distance float64)

// Similarity is a similarity storage and retrieval engine.
type Similarity struct {
	mutex sync.RWMutex
	data  map[string]map[string]Item
}

// Create a new Similarity engine.
func NewSimilarity() *Similarity {
	return &Similarity{data: make(map[string]map[string]Item)}
}

// Add an Item to the engine with a key.
func (sim *Similarity) Add(key string, item Item) {
	sim.mutex.Lock()
	defer sim.mutex.Unlock()
	m, ok := sim.data[key]
	if !ok {
		m = make(map[string]Item)
	}
	m[item.Name] = item
	sim.data[key] = m
}

// Get all the keys in this Similarity.
func (sim *Similarity) Keys() []string {
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	keys := make([]string, 0, len(sim.data))
	for k, _ := range sim.data {
		keys = append(keys, k)
	}
	return keys
}

// Dump the data backing a Similarity engine to a Writer.
func (sim *Similarity) WriteJson(w io.Writer) error {
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	b, err := json.Marshal(sim.data)
	if err != nil {
		return err
	}
	_, err = w.Write(b)
	if err != nil {
		return err
	}
	return nil
}

// Read the data for a Similarity engine from a Writer and load it
func (sim *Similarity) ReadJson(r io.Reader) error {
	// TODO: Is it idomatic to return underlying errors such as those encountered by ReadFrom or Unmarshal?
	b := new(bytes.Buffer)
	sim.mutex.Lock()
	defer sim.mutex.Unlock()
	_, err := b.ReadFrom(r)
	if err != nil {
		return err
	}
	err = json.Unmarshal(b.Bytes(), &sim.data)
	if err != nil {
		return err
	}
	return nil
}

// Returns the Euclidean distance of two keys in our Similarity engine.
func (sim *Similarity) EuclideanDistance(key1 string, key2 string) float64 {
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	// Don't compute if either key is missing.
	if sim.data[key1] == nil || sim.data[key2] == nil {
		return -1
	}
	firstItems := sim.data[key1]
	secondItems := sim.data[key2]
	// Find common Items for the two keys
	sum := 0.0
	for _, item := range firstItems {
		secondItem, found := secondItems[item.Name]
		if found {
			sum += math.Pow(item.Value-secondItem.Value, 2)
		}
	}
	return 1 / (1 + math.Sqrt(sum))
}

// Find similar keys using Euclidean distance comparison.
func (sim *Similarity) SimilarEuclidean(key string, limit int) []Result {
	return sim.Similar(key, limit, sim.EuclideanDistance)
}

// Find similar keys using the provided distance comparison.
func (sim *Similarity) Similar(key string, limit int, distance comparison) []Result {
	results := make([]Result, 0)
	sim.mutex.RLock()
	defer sim.mutex.RUnlock()
	for k, _ := range sim.data {
		if k == key {
			// Don't check ourselves.
			continue
		}
		score := distance(key, k)
		// TODO: replace -1 return value with an error in EuclideanDistance.
		if score != -1 {
			results = append(results, Result{k, score})
		}
	}
	sort.Sort(bySimilarity(results))
	if len(results) > limit {
		return results[:limit]
	} else {
		return results
	}

}

// byScore implements sort.Interface for []Result allowing us to sort results by score,
// sorting higher scoring results ahead of lower scoring results.
type bySimilarity []Result

func (a bySimilarity) Len() int           { return len(a) }
func (a bySimilarity) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a bySimilarity) Less(i, j int) bool { return a[i].Similarity > a[j].Similarity }
<|endoftext|>


<commit_before>package elasticsearch

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"
	"golang.org/x/net/context"

	"github.com/packetzoom/logzoom/buffer"
	"github.com/packetzoom/logzoom/output"
	"github.com/packetzoom/logzoom/route"
	"github.com/packetzoom/logzoom/server"
	"github.com/paulbellamy/ratecounter"
	"gopkg.in/olivere/elastic.v5"
	"gopkg.in/yaml.v2"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 10
	esRecvBuffer       = 10000
	esSendBuffer       = 10000
	esWorker           = 20
	esBulkLimit        = 10000
)

type Indexer struct {
	bulkProcessor     *elastic.BulkProcessor
	indexPrefix       string
	indexType         string
	RateCounter       *ratecounter.RateCounter
	lastDisplayUpdate time.Time
}

type Config struct {
	Hosts           []string `yaml:"hosts"`
	IndexPrefix     string   `yaml:"index"`
	IndexType       string   `yaml:"index_type"`
	Timeout         int      `yaml:"timeout"`
	GzipEnabled     bool     `yaml:"gzip_enabled"`
	InfoLogEnabled  bool     `yaml:"info_log_enabled"`
	ErrorLogEnabled bool     `yaml:"error_log_enabled"`
	SampleSize      *int      `yaml:"sample_size,omitempty"`
}

type ESServer struct {
	name   string
	fields map[string]string
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
	idx    *Indexer
}

func init() {
	output.Register("elasticsearch", New)
}

func New() (output.Output) {
	return &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	}
}

// Dummy discard, satisfies io.Writer without importing io or os.
type DevNull struct{}

func (DevNull) Write(p []byte) (int, error) {
	return len(p), nil
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func (i *Indexer) index(ev *buffer.Event) error {
	doc := *ev.Text
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	request := elastic.NewBulkIndexRequest().Index(idx).Type(typ).Doc(doc)
	i.bulkProcessor.Add(request)
	i.RateCounter.Incr(1)

	return nil
}

func (e *ESServer) ValidateConfig(config *Config) error {
	if len(config.Hosts) == 0 {
		return errors.New("Missing hosts")
	}

	if len(config.IndexPrefix) == 0 {
		return errors.New("Missing index prefix (e.g. logstash)")
	}

	if len(config.IndexType) == 0 {
		return errors.New("Missing index type (e.g. logstash)")
	}

	if e.config.SampleSize == nil {
		i := 100
		e.config.SampleSize = &i
	}
	log.Printf("[%s] Setting Sample Size to %d", e.name, *e.config.SampleSize)

	return nil
}

func (e *ESServer) Init(name string, config yaml.MapSlice, b buffer.Sender, route route.Route) error {
	var esConfig *Config

	// go-yaml doesn't have a great way to partially unmarshal YAML data
	// See https://github.com/go-yaml/yaml/issues/13
	yamlConfig, _ := yaml.Marshal(config)

	if err := yaml.Unmarshal(yamlConfig, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.name = name
	e.fields = route.Fields
	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	if err := e.ValidateConfig(esConfig); err != nil {
		return fmt.Errorf("Error in config: %v", err)
	}

	return nil
}

func readInputChannel(sampleSize int, idx *Indexer, receiveChan chan *buffer.Event) {
	select {
		case ev := <-receiveChan:
			if (server.RandInt(0, 100) < sampleSize) {
				idx.index(ev)
			}
	}
}

func (es *ESServer) insertIndexTemplate(client *elastic.Client) error {
	var template map[string]interface{}
	err := json.Unmarshal([]byte(IndexTemplate), &template)

	if err != nil {
		return err
	}

	template["template"] = es.config.IndexPrefix + "-*"

	inserter := elastic.NewIndicesPutTemplateService(client)
	inserter.Name(es.config.IndexPrefix)
	inserter.Create(true)
	inserter.BodyJson(template)

	response, err := inserter.Do(context.Background())

	if response != nil {
		log.Println("Inserted template response:", response.Acknowledged)
	}

	return err
}

func (es *ESServer) afterCommit(id int64, requests []elastic.BulkableRequest, response *elastic.BulkResponse, err error) {
	if (es.idx.RateCounter.Rate() > 0) {
		log.Printf("Flushed events to Elasticsearch, current rate: %d/s", es.idx.RateCounter.Rate())
	}
}

func (es *ESServer) Start() error {
	if (es.b == nil) {
		log.Printf("[%s] No Route is specified for this output", es.name)
		return nil
	}
	var client *elastic.Client
	var err error

	for {
		httpClient := http.DefaultClient
		timeout := 60 * time.Second

		if es.config.Timeout > 0 {
			timeout = time.Duration(es.config.Timeout) * time.Second
		}

		log.Printf("[%s] Setting HTTP timeout to %v", es.name, timeout)
		log.Printf("[%s] Setting GZIP enabled: %v", es.name, es.config.GzipEnabled)

		httpClient.Timeout = timeout

		var infoLogger, errorLogger *log.Logger

		if es.config.InfoLogEnabled {
			infoLogger = log.New(os.Stdout, "", log.LstdFlags)
		} else {
			infoLogger = log.New(new(DevNull), "", log.LstdFlags)
		}

		if es.config.ErrorLogEnabled {
			errorLogger = log.New(os.Stderr, "", log.LstdFlags)
		} else {
			errorLogger = log.New(new(DevNull), "", log.LstdFlags)
		}

		client, err = elastic.NewClient(elastic.SetURL(es.hosts...),
			elastic.SetHttpClient(httpClient),
			elastic.SetGzip(es.config.GzipEnabled),
			elastic.SetInfoLog(infoLogger),
			elastic.SetErrorLog(errorLogger))

		if err != nil {
			log.Printf("Error starting Elasticsearch: %s, will retry", err)
			time.Sleep(2 * time.Second)
			continue
		}

		es.insertIndexTemplate(client)

		break
	}

	log.Printf("Connected to Elasticsearch")

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	rateCounter := ratecounter.NewRateCounter(1 * time.Second)

	// Create bulk processor
        bulkProcessor, err := client.BulkProcessor().
		After(es.afterCommit).                        // Function to call after commit
		Workers(esWorker).                            // # of workers
		BulkActions(esBulkLimit).                     // # of queued requests before committed
		BulkSize(-1).                                 // No limit
		FlushInterval(esFlushInterval * time.Second). // autocommit every # seconds
		Stats(true).                                  // gather statistics
		Do()

        if err != nil {
            log.Println(err)
        }

	idx := &Indexer{bulkProcessor, es.config.IndexPrefix, es.config.IndexType, rateCounter, time.Now()}
	es.idx = idx

	for {

		readInputChannel(*es.config.SampleSize, idx, receiveChan)

		if len(es.term) > 0 {
			select {
			case <-es.term:
				log.Println("Elasticsearch received term signal")
				break
			}
		}
	}

	log.Println("Shutting down. Flushing existing events.")
	defer bulkProcessor.Close()
	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<commit_msg>Updated for compatibility with ElasticSearch v5.0.24.<commit_after>package elasticsearch

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"
	"golang.org/x/net/context"

	"github.com/packetzoom/logzoom/buffer"
	"github.com/packetzoom/logzoom/output"
	"github.com/packetzoom/logzoom/route"
	"github.com/packetzoom/logzoom/server"
	"github.com/paulbellamy/ratecounter"
	"gopkg.in/olivere/elastic.v5"
	"gopkg.in/yaml.v2"
)

const (
	defaultHost        = "127.0.0.1"
	defaultIndexPrefix = "logstash"
	esFlushInterval    = 10
	esRecvBuffer       = 10000
	esSendBuffer       = 10000
	esWorker           = 20
	esBulkLimit        = 10000
)

type Indexer struct {
	bulkProcessor     *elastic.BulkProcessor
	indexPrefix       string
	indexType         string
	RateCounter       *ratecounter.RateCounter
	lastDisplayUpdate time.Time
}

type Config struct {
	Hosts           []string `yaml:"hosts"`
	IndexPrefix     string   `yaml:"index"`
	IndexType       string   `yaml:"index_type"`
	Timeout         int      `yaml:"timeout"`
	GzipEnabled     bool     `yaml:"gzip_enabled"`
	InfoLogEnabled  bool     `yaml:"info_log_enabled"`
	ErrorLogEnabled bool     `yaml:"error_log_enabled"`
	SampleSize      *int      `yaml:"sample_size,omitempty"`
}

type ESServer struct {
	name   string
	fields map[string]string
	config Config
	host   string
	hosts  []string
	b      buffer.Sender
	term   chan bool
	idx    *Indexer
}

func init() {
	output.Register("elasticsearch", New)
}

func New() (output.Output) {
	return &ESServer{
		host: fmt.Sprintf("%s:%d", defaultHost, time.Now().Unix()),
		term: make(chan bool, 1),
	}
}

// Dummy discard, satisfies io.Writer without importing io or os.
type DevNull struct{}

func (DevNull) Write(p []byte) (int, error) {
	return len(p), nil
}

func indexName(idx string) string {
	if len(idx) == 0 {
		idx = defaultIndexPrefix
	}

	return fmt.Sprintf("%s-%s", idx, time.Now().Format("2006.01.02"))
}

func (i *Indexer) index(ev *buffer.Event) error {
	doc := *ev.Text
	idx := indexName(i.indexPrefix)
	typ := i.indexType

	request := elastic.NewBulkIndexRequest().Index(idx).Type(typ).Doc(doc)
	i.bulkProcessor.Add(request)
	i.RateCounter.Incr(1)

	return nil
}

func (e *ESServer) ValidateConfig(config *Config) error {
	if len(config.Hosts) == 0 {
		return errors.New("Missing hosts")
	}

	if len(config.IndexPrefix) == 0 {
		return errors.New("Missing index prefix (e.g. logstash)")
	}

	if len(config.IndexType) == 0 {
		return errors.New("Missing index type (e.g. logstash)")
	}

	if e.config.SampleSize == nil {
		i := 100
		e.config.SampleSize = &i
	}
	log.Printf("[%s] Setting Sample Size to %d", e.name, *e.config.SampleSize)

	return nil
}

func (e *ESServer) Init(name string, config yaml.MapSlice, b buffer.Sender, route route.Route) error {
	var esConfig *Config

	// go-yaml doesn't have a great way to partially unmarshal YAML data
	// See https://github.com/go-yaml/yaml/issues/13
	yamlConfig, _ := yaml.Marshal(config)

	if err := yaml.Unmarshal(yamlConfig, &esConfig); err != nil {
		return fmt.Errorf("Error parsing elasticsearch config: %v", err)
	}

	e.name = name
	e.fields = route.Fields
	e.config = *esConfig
	e.hosts = esConfig.Hosts
	e.b = b

	if err := e.ValidateConfig(esConfig); err != nil {
		return fmt.Errorf("Error in config: %v", err)
	}

	return nil
}

func readInputChannel(sampleSize int, idx *Indexer, receiveChan chan *buffer.Event) {
	select {
		case ev := <-receiveChan:
			if (server.RandInt(0, 100) < sampleSize) {
				idx.index(ev)
			}
	}
}

func (es *ESServer) insertIndexTemplate(client *elastic.Client) error {
	var template map[string]interface{}
	err := json.Unmarshal([]byte(IndexTemplate), &template)

	if err != nil {
		return err
	}

	template["template"] = es.config.IndexPrefix + "-*"

	inserter := elastic.NewIndicesPutTemplateService(client)
	inserter.Name(es.config.IndexPrefix)
	inserter.Create(true)
	inserter.BodyJson(template)

	response, err := inserter.Do(context.Background())

	if response != nil {
		log.Println("Inserted template response:", response.Acknowledged)
	}

	return err
}

func (es *ESServer) afterCommit(id int64, requests []elastic.BulkableRequest, response *elastic.BulkResponse, err error) {
	if (es.idx.RateCounter.Rate() > 0) {
		log.Printf("Flushed events to Elasticsearch, current rate: %d/s", es.idx.RateCounter.Rate())
	}
}

func (es *ESServer) Start() error {
	if (es.b == nil) {
		log.Printf("[%s] No Route is specified for this output", es.name)
		return nil
	}
	var client *elastic.Client
	var err error

	for {
		httpClient := http.DefaultClient
		timeout := 60 * time.Second

		if es.config.Timeout > 0 {
			timeout = time.Duration(es.config.Timeout) * time.Second
		}

		log.Printf("[%s] Setting HTTP timeout to %v", es.name, timeout)
		log.Printf("[%s] Setting GZIP enabled: %v", es.name, es.config.GzipEnabled)

		httpClient.Timeout = timeout

		var infoLogger, errorLogger *log.Logger

		if es.config.InfoLogEnabled {
			infoLogger = log.New(os.Stdout, "", log.LstdFlags)
		} else {
			infoLogger = log.New(new(DevNull), "", log.LstdFlags)
		}

		if es.config.ErrorLogEnabled {
			errorLogger = log.New(os.Stderr, "", log.LstdFlags)
		} else {
			errorLogger = log.New(new(DevNull), "", log.LstdFlags)
		}

		client, err = elastic.NewClient(elastic.SetURL(es.hosts...),
			elastic.SetHttpClient(httpClient),
			elastic.SetGzip(es.config.GzipEnabled),
			elastic.SetInfoLog(infoLogger),
			elastic.SetErrorLog(errorLogger))

		if err != nil {
			log.Printf("Error starting Elasticsearch: %s, will retry", err)
			time.Sleep(2 * time.Second)
			continue
		}

		es.insertIndexTemplate(client)

		break
	}

	log.Printf("Connected to Elasticsearch")

	// Add the client as a subscriber
	receiveChan := make(chan *buffer.Event, esRecvBuffer)
	es.b.AddSubscriber(es.host, receiveChan)
	defer es.b.DelSubscriber(es.host)

	rateCounter := ratecounter.NewRateCounter(1 * time.Second)

	// Create bulk processor
        bulkProcessor, err := client.BulkProcessor().
		After(es.afterCommit).                        // Function to call after commit
		Workers(esWorker).                            // # of workers
		BulkActions(esBulkLimit).                     // # of queued requests before committed
		BulkSize(-1).                                 // No limit
		FlushInterval(esFlushInterval * time.Second). // autocommit every # seconds
		Stats(true).                                  // gather statistics
		Do(context.Background())

        if err != nil {
            log.Println(err)
        }

	idx := &Indexer{bulkProcessor, es.config.IndexPrefix, es.config.IndexType, rateCounter, time.Now()}
	es.idx = idx

	for {

		readInputChannel(*es.config.SampleSize, idx, receiveChan)

		if len(es.term) > 0 {
			select {
			case <-es.term:
				log.Println("Elasticsearch received term signal")
				break
			}
		}
	}

	log.Println("Shutting down. Flushing existing events.")
	defer bulkProcessor.Close()
	return nil
}

func (es *ESServer) Stop() error {
	es.term <- true
	return nil
}
<|endoftext|>