<commit_before>package main

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type Any interface{}

type Env struct {
	m     map[string]Any
	outer *Env
}

func (e *Env) find(k Any) (*Env, bool) {
	if _, ok := e.m[k.(string)]; ok {
		return e, true
	} else if e.outer != nil {
		return e.outer.find(k)
	}
	return nil, false
}

func (e *Env) get(k Any) Any {
	return e.m[k.(string)]
}

func (e *Env) merge(m map[string]Any) {
	for k, v := range m {
		e.m[k] = v
	}
}

func (e *Env) set(k Any, v Any) {
	e.m[k.(string)] = v
}

func NewEnv(keys []Any, vals []Any, outer *Env) *Env {
	zipped := map[string]Any{}
	vlen := len(vals)
	for i, key := range keys {
		if i < vlen {
			zipped[key.(string)] = vals[i]
		}
	}
	return &Env{zipped, outer}
}

// Evaluate an expression
func eval(e Any, env *Env) Any {
	// fmt.Printf("[eval] expression: %v (%T)   env: %v\n", e, e, env)
	switch e := e.(type) {
	case string:
		if environ, ok := env.find(e); ok {
			return environ.get(e)
		}
		return nil
	case bool, int64, float64:
		return e
	case []Any:
		if len(e) > 0 {
			switch e[0] {
			case "quote": // (quote exp)
				return e[1]
			case "if": // (if test then else?)
				if eval(e[1], env).(bool) {
					return eval(e[2], env)
				} else {
					return eval(e[3], env)
				}
			case "set!": // (set! var exp)
				if environ, ok := env.find(e[1]); ok {
					environ.set(e[1], eval(e[2], env))
				}
			case "define", "def": // (define var exp)
				env.set(e[1], eval(e[2], env))
			case "lambda", "fn": // (lambda (var*) exp)
				fmt.Printf("[lambda] e[1]: %v (%T)\n", e[1], e[1])
				vars := e[1].([]Any)
				exp := e[2]
				// fmt.Printf("vars: %v\n", vars)
				// fmt.Printf("exp: %v\n", exp)
				return func(args ...Any) Any {
					fmt.Printf("[lamda] vars: %v\n", vars)
					fmt.Printf("[lamda] args: %v\n", args)
					fmt.Printf("[lamda] env: %v\n", env)
					fmt.Printf("[lamda] new-env: %v\n", NewEnv(vars, args, env))
					fmt.Printf("[lamda] exp: %v\n", exp)
					return eval(exp, NewEnv(vars, args, env))
				}
			case "begin": // (begin exp*)
				var val Any
				for _, exp := range e[1:] {
					val = eval(exp, env)
				}
				return val
			default: // (proc exp*)
				fmt.Printf("[proc] expr: %v\n", e)
				exprs := make([]Any, len(e))
				for i, exp := range e {
					fmt.Printf("[proc] exp: %v (%T)\n", exp, exp)
					fmt.Printf("[proc] env: %v\n", env)
					exprs[i] = eval(exp, env)
					fmt.Printf("[proc] result: %v\n", exprs[i])
				}
				fmt.Printf("[proc] exprs: %v\n", exprs)
				fn := exprs[0].(func(...Any) Any)
				fmt.Printf("[proc] fn: %v\n", fn)
				return fn(exprs[1:]...)
			}
		}
	}
	return e
}

// Read a Scheme expression from a string.
func read(s string) (Any, error) {
	parsed, _, err := read_from(tokenize(s))
	return parsed, err
}

// Converts a string into an array of tokens.
func tokenize(s string) []string {
	return regexp.MustCompile(`\s+`).Split(strings.TrimSpace(
		strings.Replace(strings.Replace(s, "(", " ( ", -1), ")", " ) ", -1)), -1)
}

// Read an expression from a sequence of tokens.
func read_from(tokens []string) (Any, []string, error) {
	if len(tokens) == 0 {
		return nil, nil, errors.New("unexpected EOF while reading")
	}
	token := tokens[0]
	tokens = tokens[1:]
	switch token {
	case "(":
		L := []Any{}
		for len(tokens) > 0 && tokens[0] != ")" {
			token, remaining, _ := read_from(tokens) // TODO handle err
			L = append(L, token)
			tokens = remaining
		}
		if len(tokens) > 0 {
			tokens = tokens[1:] // pop off ')'
		}
		return L, tokens, nil
	case ")":
		return token, tokens, errors.New("unexpected )")
	}
	return atom(token), tokens, nil
}

// Bools, ints, and floats are converted; every other token is a symbol.
func atom(s string) Any {
	if b, err := strconv.ParseBool(s); (s == "true" || s == "false") && err == nil {
		return b
	}
	if i, err := strconv.ParseInt(s, 0, 64); err == nil {
		return i
	}
	if f, err := strconv.ParseFloat(s, 64); err == nil {
		return f
	}
	return s
}

func prompt() {
	fmt.Print(">>> ")
}

func main() {
	env := NewEnv(nil, nil, nil)
	env.merge(map[string]Any{
		"+": func(xs ...Any) Any {
			sum := int64(0)
			for _, x := range xs {
				sum += x.(int64)
			}
			return sum
		},
		"-": func(xs ...Any) Any {
			switch len(xs) {
			case 0:
				return 0
			case 1:
				return -1 * xs[0].(int64)
			case 2:
				return xs[0].(int64) - xs[1].(int64)
			default:
				sum := xs[0].(int64) - xs[1].(int64)
				for _, x := range xs[2:] {
					sum -= x.(int64)
				}
				return sum
			}
		},
		"*": func(xs ...Any) Any {
			sum := int64(1)
			for _, x := range xs {
				sum *= x.(int64)
			}
			return sum
		},
		"<=": func(xs ...Any) Any {
			last := xs[0].(int64)
			for _, x := range xs {
				if x.(int64) < last {
					return false
				}
				last = x.(int64)
			}
			return true
		},
	})

	scanner := bufio.NewScanner(os.Stdin)
	for prompt(); scanner.Scan(); prompt() {
		in := scanner.Text()
		parsed, _ := read(in) // TODO handle err
		fmt.Printf("parsed: %v\n", parsed)
		evaled := eval(parsed, env)
		fmt.Printf("%v (%T)\n", evaled, evaled)
		// fmt.Printf("%v (%T) %v\n", evaled, evaled, env)
	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	}
}
<commit_msg>removed print stmts<commit_after>package main

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type Any interface{}

type Env struct {
	m     map[string]Any
	outer *Env
}

func (e *Env) find(k Any) (*Env, bool) {
	if _, ok := e.m[k.(string)]; ok {
		return e, true
	} else if e.outer != nil {
		return e.outer.find(k)
	}
	return nil, false
}

func (e *Env) get(k Any) Any {
	return e.m[k.(string)]
}

func (e *Env) merge(m map[string]Any) {
	for k, v := range m {
		e.m[k] = v
	}
}

func (e *Env) set(k Any, v Any) {
	e.m[k.(string)] = v
}

func NewEnv(keys []Any, vals []Any, outer *Env) *Env {
	zipped := map[string]Any{}
	vlen := len(vals)
	for i, key := range keys {
		if i < vlen {
			zipped[key.(string)] = vals[i]
		}
	}
	return &Env{zipped, outer}
}

// Evaluate an expression
func eval(e Any, env *Env) Any {
	switch e := e.(type) {
	case string:
		if environ, ok := env.find(e); ok {
			return environ.get(e)
		}
		return nil
	case bool, int64, float64:
		return e
	case []Any:
		if len(e) > 0 {
			switch e[0] {
			case "quote": // (quote exp)
				return e[1]
			case "if": // (if test then else?)
				if eval(e[1], env).(bool) {
					return eval(e[2], env)
				} else {
					return eval(e[3], env)
				}
			case "set!": // (set! var exp)
				if environ, ok := env.find(e[1]); ok {
					environ.set(e[1], eval(e[2], env))
				}
			case "define", "def": // (define var exp)
				env.set(e[1], eval(e[2], env))
			case "lambda", "fn": // (lambda (var*) exp)
				vars := e[1].([]Any)
				exp := e[2]
				return func(args ...Any) Any {
					return eval(exp, NewEnv(vars, args, env))
				}
			case "begin": // (begin exp*)
				var val Any
				for _, exp := range e[1:] {
					val = eval(exp, env)
				}
				return val
			default: // (proc exp*)
				exprs := make([]Any, len(e))
				for i, exp := range e {
					exprs[i] = eval(exp, env)
				}
				fn := exprs[0].(func(...Any) Any)
				return fn(exprs[1:]...)
			}
		}
	}
	return e
}

// Read a Scheme expression from a string.
func read(s string) (Any, error) {
	parsed, _, err := read_from(tokenize(s))
	return parsed, err
}

// Converts a string into an array of tokens.
func tokenize(s string) []string {
	return regexp.MustCompile(`\s+`).Split(strings.TrimSpace(
		strings.Replace(strings.Replace(s, "(", " ( ", -1), ")", " ) ", -1)), -1)
}

// Read an expression from a sequence of tokens.
func read_from(tokens []string) (Any, []string, error) {
	if len(tokens) == 0 {
		return nil, nil, errors.New("unexpected EOF while reading")
	}
	token := tokens[0]
	tokens = tokens[1:]
	switch token {
	case "(":
		L := []Any{}
		for len(tokens) > 0 && tokens[0] != ")" {
			token, remaining, _ := read_from(tokens) // TODO handle err
			L = append(L, token)
			tokens = remaining
		}
		if len(tokens) > 0 {
			tokens = tokens[1:] // pop off ')'
		}
		return L, tokens, nil
	case ")":
		return token, tokens, errors.New("unexpected )")
	}
	return atom(token), tokens, nil
}

// Bools, ints, and floats are converted; every other token is a symbol.
func atom(s string) Any {
	if b, err := strconv.ParseBool(s); (s == "true" || s == "false") && err == nil {
		return b
	}
	if i, err := strconv.ParseInt(s, 0, 64); err == nil {
		return i
	}
	if f, err := strconv.ParseFloat(s, 64); err == nil {
		return f
	}
	return s
}

func prompt() {
	fmt.Print(">>> ")
}

func main() {
	env := NewEnv(nil, nil, nil)
	env.merge(map[string]Any{
		"+": func(xs ...Any) Any {
			sum := int64(0)
			for _, x := range xs {
				sum += x.(int64)
			}
			return sum
		},
		"-": func(xs ...Any) Any {
			switch len(xs) {
			case 0:
				return 0
			case 1:
				return -1 * xs[0].(int64)
			case 2:
				return xs[0].(int64) - xs[1].(int64)
			default:
				sum := xs[0].(int64) - xs[1].(int64)
				for _, x := range xs[2:] {
					sum -= x.(int64)
				}
				return sum
			}
		},
		"*": func(xs ...Any) Any {
			sum := int64(1)
			for _, x := range xs {
				sum *= x.(int64)
			}
			return sum
		},
		"<=": func(xs ...Any) Any {
			last := xs[0].(int64)
			for _, x := range xs {
				if x.(int64) < last {
					return false
				}
				last = x.(int64)
			}
			return true
		},
	})

	scanner := bufio.NewScanner(os.Stdin)
	for prompt(); scanner.Scan(); prompt() {
		in := scanner.Text()
		parsed, _ := read(in) // TODO handle err
		evaled := eval(parsed, env)
		fmt.Printf("%v (%T)\n", evaled, evaled)
	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	}
}
<|endoftext|>


<commit_before>package goss

import "sort"

type (
	// Searcher is a interface for sort.Search
	Searcher interface {
		Target() int64
		Priority() string
	}

	// SortedSlice is a struct of sorted-slice
	SortedSlice struct {
		S    []Searcher
		DESC bool
	}
)

// Add append to slice
func (s *SortedSlice) Add(item Searcher) {

	if item == nil {
		return
	}

	i := sort.Search(len(s.S), func(i int) bool {
		if s.DESC {
			return s.S[i].Target() < item.Target() || (s.S[i].Target() == item.Target() && s.S[i].Priority() > item.Priority())
		}
		return s.S[i].Target() > item.Target() || (s.S[i].Target() == item.Target() && s.S[i].Priority() > item.Priority())
	})

	s.S = append(s.S, nil)
	copy(s.S[i+1:], s.S[i:])
	s.S[i] = item
}
<commit_msg>Add sync.Mutex<commit_after>package goss

import (
	"sort"
	"sync"
)

type (
	// Searcher is a interface for sort.Search
	Searcher interface {
		Target() int64
		Priority() string
	}

	// SortedSlice is a struct of sorted-slice
	SortedSlice struct {
		S    []Searcher
		DESC bool
		m sync.Mutex
	}
)

// Add append to slice
func (s *SortedSlice) Add(item Searcher) {

	if item == nil {
		return
	}

	s.m.Lock()
	defer s.m.Unlock()

	i := sort.Search(len(s.S), func(i int) bool {
		if s.DESC {
			return s.S[i].Target() < item.Target() || (s.S[i].Target() == item.Target() && s.S[i].Priority() > item.Priority())
		}
		return s.S[i].Target() > item.Target() || (s.S[i].Target() == item.Target() && s.S[i].Priority() > item.Priority())
	})

	s.S = append(s.S, nil)
	copy(s.S[i+1:], s.S[i:])
	s.S[i] = item
}

<|endoftext|>


<commit_before>package baseftrwapp

import (
	"fmt"
	standardLog "log"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/Financial-Times/http-handlers-go/httphandlers"
	"github.com/Financial-Times/service-status-go/gtg"
	status "github.com/Financial-Times/service-status-go/httphandlers"
	log "github.com/Sirupsen/logrus"
	graphite "github.com/cyberdelia/go-metrics-graphite"
	"github.com/gorilla/mux"
	metrics "github.com/rcrowley/go-metrics"
)

// RunServer will set up GET, PUT and DELETE endpoints for the specified path,
// calling the appropriate service functions:
// PUT -> Write
// GET -> Read
// DELETE -> Delete
// It will also setup the healthcheck and ping endpoints
// Endpoints are wrapped in a metrics timer and request loggin including transactionID, which is generated
// if not found on the request as X-Request-Id header
func RunServer(engs map[string]Service, healthHandler func(http.ResponseWriter, *http.Request), port int, serviceName string, env string) {
	for path, eng := range engs {
		err := eng.Initialise()
		if err != nil {
			log.Fatalf("Service for path %s could not startup, err=%s", path, err)
		}
	}

	m := router(engs, healthHandler)
	http.Handle("/", m)

	if env != "local" {
		f, err := os.OpenFile("/var/log/apps/"+serviceName+"-go-app.log", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
		if err == nil {
			log.SetOutput(f)
			log.SetFormatter(&log.TextFormatter{DisableColors: true})
		} else {
			log.Fatalf("Failed to initialise log file, %v", err)
		}

		defer f.Close()
	}

	log.Printf("listening on %d", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port),
		httphandlers.HTTPMetricsHandler(metrics.DefaultRegistry,
			httphandlers.TransactionAwareRequestLoggingHandler(log.StandardLogger(), m)))

	log.Printf("exiting on %s", serviceName)

}

//Router sets up the Router - extracted for testability
func router(engs map[string]Service, healthHandler func(http.ResponseWriter, *http.Request)) *mux.Router {
	m := mux.NewRouter()

	gtgChecker := make([]gtg.StatusChecker, 0)

	for path, eng := range engs {
		handlers := httpHandlers{eng}
		m.HandleFunc(fmt.Sprintf("/%s/__count", path), handlers.countHandler).Methods("GET")
		m.HandleFunc(fmt.Sprintf("/%s/{uuid}", path), handlers.getHandler).Methods("GET")
		m.HandleFunc(fmt.Sprintf("/%s/{uuid}", path), handlers.putHandler).Methods("PUT")
		m.HandleFunc(fmt.Sprintf("/%s/{uuid}", path), handlers.deleteHandler).Methods("DELETE")
		gtgChecker = append(gtgChecker, func() gtg.Status {
			if err := eng.Check(); err != nil {
				return gtg.Status{GoodToGo: false, Message: err.Error()}
			}

			return gtg.Status{GoodToGo: true}
		})
	}

	m.HandleFunc("/__health", healthHandler)
	// The top one of these feels more correct, but the lower one matches what we have in Dropwizard,
	// so it's what apps expect currently
	m.HandleFunc(status.PingPath, status.PingHandler)
	m.HandleFunc(status.PingPathDW, status.PingHandler)

	// The top one of these feels more correct, but the lower one matches what we have in Dropwizard,
	// so it's what apps expect currently same as ping, the content of build-info needs more definition
	m.HandleFunc(status.BuildInfoPath, status.BuildInfoHandler)
	m.HandleFunc(status.BuildInfoPathDW, status.BuildInfoHandler)

	m.HandleFunc(status.GTGPath, status.NewGoodToGoHandler(gtg.FailFastParallelCheck(gtgChecker)))

	return m
}

//OutputMetricsIfRequired will send metrics to Graphite if a non-empty graphiteTCPAddress is passed in, or to the standard log if logMetrics is true.
// Make sure a sensible graphitePrefix that will uniquely identify your service is passed in, e.g. "content.test.people.rw.neo4j.ftaps58938-law1a-eu-t
func OutputMetricsIfRequired(graphiteTCPAddress string, graphitePrefix string, logMetrics bool) {
	if graphiteTCPAddress != "" {
		addr, _ := net.ResolveTCPAddr("tcp", graphiteTCPAddress)
		go graphite.Graphite(metrics.DefaultRegistry, 5*time.Second, graphitePrefix, addr)
	}
	if logMetrics { //useful locally
		//messy use of the 'standard' log package here as this method takes the log struct, not an interface, so can't use logrus.Logger
		go metrics.Log(metrics.DefaultRegistry, 60*time.Second, standardLog.New(os.Stdout, "metrics", standardLog.Lmicroseconds))
	}
}
<commit_msg>Better routing<commit_after>package baseftrwapp

import (
	"fmt"
	standardLog "log"
	"net"
	"net/http"
	"os"
	"time"

	"github.com/Financial-Times/http-handlers-go/httphandlers"
	"github.com/Financial-Times/service-status-go/gtg"
	status "github.com/Financial-Times/service-status-go/httphandlers"
	log "github.com/Sirupsen/logrus"
	graphite "github.com/cyberdelia/go-metrics-graphite"
	"github.com/gorilla/mux"
	metrics "github.com/rcrowley/go-metrics"
)

// RunServer will set up GET, PUT and DELETE endpoints for the specified path,
// calling the appropriate service functions:
// PUT -> Write
// GET -> Read
// DELETE -> Delete
// It will also setup the healthcheck and ping endpoints
// Endpoints are wrapped in a metrics timer and request loggin including transactionID, which is generated
// if not found on the request as X-Request-Id header
func RunServer(engs map[string]Service, healthHandler func(http.ResponseWriter, *http.Request), port int, serviceName string, env string) {
	for path, eng := range engs {
		err := eng.Initialise()
		if err != nil {
			log.Fatalf("Service for path %s could not startup, err=%s", path, err)
		}
	}

	//	http.Handle("/", m)

	if env != "local" {
		f, err := os.OpenFile("/var/log/apps/"+serviceName+"-go-app.log", os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0644)
		if err == nil {
			log.SetOutput(f)
			log.SetFormatter(&log.TextFormatter{DisableColors: true})
		} else {
			log.Fatalf("Failed to initialise log file, %v", err)
		}

		defer f.Close()
	}

	var m http.Handler
	m = router(engs, healthHandler)
	m = httphandlers.TransactionAwareRequestLoggingHandler(log.StandardLogger(), m)
	m = httphandlers.HTTPMetricsHandler(metrics.DefaultRegistry, m)

	http.Handle("/", m)

	log.Printf("listening on %d", port)
	http.ListenAndServe(fmt.Sprintf(":%d", port), nil)

	log.Printf("exiting on %s", serviceName)

}

//Router sets up the Router - extracted for testability
func router(engs map[string]Service, healthHandler func(http.ResponseWriter, *http.Request)) *mux.Router {
	m := mux.NewRouter()

	gtgChecker := make([]gtg.StatusChecker, 0)

	for path, eng := range engs {
		handlers := httpHandlers{eng}
		m.HandleFunc(fmt.Sprintf("/%s/__count", path), handlers.countHandler).Methods("GET")
		m.HandleFunc(fmt.Sprintf("/%s/{uuid}", path), handlers.getHandler).Methods("GET")
		m.HandleFunc(fmt.Sprintf("/%s/{uuid}", path), handlers.putHandler).Methods("PUT")
		m.HandleFunc(fmt.Sprintf("/%s/{uuid}", path), handlers.deleteHandler).Methods("DELETE")
		gtgChecker = append(gtgChecker, func() gtg.Status {
			if err := eng.Check(); err != nil {
				return gtg.Status{GoodToGo: false, Message: err.Error()}
			}

			return gtg.Status{GoodToGo: true}
		})
	}

	m.HandleFunc("/__health", healthHandler)
	// The top one of these feels more correct, but the lower one matches what we have in Dropwizard,
	// so it's what apps expect currently
	m.HandleFunc(status.PingPath, status.PingHandler)
	m.HandleFunc(status.PingPathDW, status.PingHandler)

	// The top one of these feels more correct, but the lower one matches what we have in Dropwizard,
	// so it's what apps expect currently same as ping, the content of build-info needs more definition
	m.HandleFunc(status.BuildInfoPath, status.BuildInfoHandler)
	m.HandleFunc(status.BuildInfoPathDW, status.BuildInfoHandler)

	m.HandleFunc(status.GTGPath, status.NewGoodToGoHandler(gtg.FailFastParallelCheck(gtgChecker)))

	return m
}

//OutputMetricsIfRequired will send metrics to Graphite if a non-empty graphiteTCPAddress is passed in, or to the standard log if logMetrics is true.
// Make sure a sensible graphitePrefix that will uniquely identify your service is passed in, e.g. "content.test.people.rw.neo4j.ftaps58938-law1a-eu-t
func OutputMetricsIfRequired(graphiteTCPAddress string, graphitePrefix string, logMetrics bool) {
	if graphiteTCPAddress != "" {
		addr, _ := net.ResolveTCPAddr("tcp", graphiteTCPAddress)
		go graphite.Graphite(metrics.DefaultRegistry, 5*time.Second, graphitePrefix, addr)
	}
	if logMetrics { //useful locally
		//messy use of the 'standard' log package here as this method takes the log struct, not an interface, so can't use logrus.Logger
		go metrics.Log(metrics.DefaultRegistry, 60*time.Second, standardLog.New(os.Stdout, "metrics", standardLog.Lmicroseconds))
	}
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
// Copyright (c) 2014 Claes Mogren
// http://opensource.org/licenses/MIT

// Code to generate a weighted random identity based on weighted input files
// http://www.keithschwarz.com/darts-dice-coins/
// http://web.eecs.utk.edu/~vose/Publications/random.pdf
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"strconv"
	"strings"
	"time"
)

var langFlag string
var listLangFlag bool
var verbose bool
var nrLoopsFlag int

type Rig struct {
	firstname, lastname, street, city string
	streetNumber, zipCode             int
}

type RigFile struct {
	tot     float64
	weights []float64
	texts   [][]string
	// Vose
}

type RigDict struct {
	fnames, mnames, lnames, streets, zipcodes RigFile
}

func init() {
	flag.StringVar(&langFlag, "lang", "en", "Select ISO 639-1 language code")
	flag.BoolVar(&listLangFlag, "l", false, "List available ISO language codes")
	flag.BoolVar(&verbose, "v", false, "Verbose output")
	flag.IntVar(&nrLoopsFlag, "n", 1, "Number of identities to output")
}

func main() {
	flag.Parse()
	if listLangFlag {
		listLangs()
		os.Exit(0)
	}
	rand.Seed(time.Now().UnixNano())
	dict := loadData(langFlag)
	for i := 0; i < nrLoopsFlag; i++ {
		rig := getNext(dict)
		fmt.Println(rig.firstname, rig.lastname)
	}
}

func listLangs() {
	// for all dirs in data
	files, _ := ioutil.ReadDir("./data/")
	for _, f := range files {
		if f.IsDir() && !strings.HasPrefix(f.Name(), ",") {
			if validateDir(f.Name()) {
				fmt.Println(f.Name())
			}
		}
	}
}

func validateDir(iso string) bool {
	// Check for fnames, mnames, lnames, zipcodes and streets
	srcFileNames := []string{"fnames.grig", "lnames.grig", "mnames.grig", "streets.grig", "zipcodes.grig"}
	valid := true
	for _, srcFile := range srcFileNames {
		filename := "./data/" + iso + "/" + srcFile
		if _, err := os.Stat(filename); os.IsNotExist(err) {
			if verbose {
				fmt.Println("Data file", srcFile, "missing for", iso)
			}
			valid = false
		}
	}
	return valid
}

func loadData(iso string) RigDict {
	dict := RigDict{}
	dict.fnames = loadFile(iso, "fnames.grig")
	dict.mnames = loadFile(iso, "mnames.grig")
	dict.lnames = loadFile(iso, "lnames.grig")
	dict.streets = loadFile(iso, "streets.grig")
	dict.zipcodes = loadFile(iso, "zipcodes.grig")
	if verbose {
		fmt.Println("fname tot:", dict.fnames.tot, dict.fnames.texts[0])
		fmt.Println("mname tot:", dict.mnames.tot, dict.mnames.texts[0])
		fmt.Println("lname tot:", dict.lnames.tot, dict.lnames.texts[0])
	}
	return dict
}

func loadFile(iso string, srcFile string) RigFile {
	file, err := os.Open("data/" + iso + "/" + srcFile)
	rigFile := RigFile{}
	rigFile.weights = make([]float64, 0)
	rigFile.texts = make([][]string, 0)
	defer file.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}
	scanner := bufio.NewScanner(file)
	sum := 0.0
	for scanner.Scan() {
		dataStr := strings.Split(scanner.Text(), "\t")
		// string to float
		f, err := strconv.ParseFloat(dataStr[0], 64)
		if err != nil {
			// Ignore error
			fmt.Println(err)
		}
		sum += f
		str := dataStr[1:]
		rigFile.weights = append(rigFile.weights, f)
		rigFile.texts = append(rigFile.texts, str)
		if verbose {
			fmt.Println("P: ", f, str)
		}
	}
	rigFile.tot = sum
	return rigFile
}

func getNext(dict RigDict) Rig {
	rig := Rig{}
	if rand.Intn(2) == 0 {
		rig.firstname = dict.fnames.texts[rand.Intn(len(dict.fnames.texts))][0]
	} else {
		rig.firstname = dict.mnames.texts[rand.Intn(len(dict.mnames.texts))][0]
	}
	rig.lastname = dict.lnames.texts[rand.Intn(len(dict.lnames.texts))][0]
	rig.street = dict.streets.texts[rand.Intn(len(dict.streets.texts))][0]
	rig.streetNumber = rand.Intn(50)
	zip := dict.zipcodes.texts[rand.Intn(len(dict.zipcodes.texts))]
	rig.city = zip[1]
	rig.zipCode, _ = strconv.Atoi(zip[0])
	return rig
}
<commit_msg>Add json<commit_after>// The MIT License (MIT)
// Copyright (c) 2014 Claes Mogren
// http://opensource.org/licenses/MIT

// Code to generate a weighted random identity based on weighted input files
// http://www.keithschwarz.com/darts-dice-coins/
// http://web.eecs.utk.edu/~vose/Publications/random.pdf
package main

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

var langFlag string
var listLangFlag bool
var verbose bool
var nrLoopsFlag int

type Rig struct {
	Firstname, Lastname, Street, City string
	Streetnumber, Zipcode             int
}

func (r Rig) AsText() string {
	str := fmt.Sprintln(r.Firstname, r.Lastname)
	str += fmt.Sprintln(r.Street, r.Streetnumber)
	str += fmt.Sprintln(r.Zipcode, r.City)
	return str
}

func (r Rig) AsJson() string {
	b, err := json.MarshalIndent(r, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}
	return string(b)
}

type RigFile struct {
	tot     float64
	weights []float64
	texts   [][]string
	// Vose
}

type RigDict struct {
	fnames, mnames, lnames, Streets, zipcodes RigFile
}

func init() {
	flag.StringVar(&langFlag, "lang", "en", "Select ISO 639-1 language code")
	flag.BoolVar(&listLangFlag, "l", false, "List available ISO language codes")
	flag.BoolVar(&verbose, "v", false, "Verbose output")
	flag.IntVar(&nrLoopsFlag, "n", 1, "Number of identities to output")
}

func main() {
	flag.Parse()
	if listLangFlag {
		listLangs()
		os.Exit(0)
	}
	rand.Seed(time.Now().UnixNano())
	dict := loadData(langFlag)
	for i := 0; i < nrLoopsFlag; i++ {
		rig := getNext(dict)
		fmt.Println(rig.AsText())
		fmt.Println(rig.AsJson())
	}
}

func listLangs() {
	// for all dirs in data
	files, _ := ioutil.ReadDir("./data/")
	for _, f := range files {
		if f.IsDir() && !strings.HasPrefix(f.Name(), ",") {
			if validateDir(f.Name()) {
				fmt.Println(f.Name())
			}
		}
	}
}

func validateDir(iso string) bool {
	// Check for fnames, mnames, lnames, zipcodes and Streets
	srcFileNames := []string{"fnames.grig", "lnames.grig", "mnames.grig", "Streets.grig", "zipcodes.grig"}
	valid := true
	for _, srcFile := range srcFileNames {
		filename := "./data/" + iso + "/" + srcFile
		if _, err := os.Stat(filename); os.IsNotExist(err) {
			if verbose {
				fmt.Println("Data file", srcFile, "missing for", iso)
			}
			valid = false
		}
	}
	return valid
}

func loadData(iso string) RigDict {
	dict := RigDict{}
	dict.fnames = loadFile(iso, "fnames.grig")
	dict.mnames = loadFile(iso, "mnames.grig")
	dict.lnames = loadFile(iso, "lnames.grig")
	dict.Streets = loadFile(iso, "Streets.grig")
	dict.zipcodes = loadFile(iso, "zipcodes.grig")
	if verbose {
		fmt.Println("fname tot:", dict.fnames.tot, dict.fnames.texts[0])
		fmt.Println("mname tot:", dict.mnames.tot, dict.mnames.texts[0])
		fmt.Println("lname tot:", dict.lnames.tot, dict.lnames.texts[0])
	}
	return dict
}

func loadFile(iso string, srcFile string) RigFile {
	file, err := os.Open("data/" + iso + "/" + srcFile)
	rigFile := RigFile{}
	rigFile.weights = make([]float64, 0)
	rigFile.texts = make([][]string, 0)
	defer file.Close()
	if err != nil {
		fmt.Println(err)
		os.Exit(0)
	}
	scanner := bufio.NewScanner(file)
	sum := 0.0
	for scanner.Scan() {
		dataStr := strings.Split(scanner.Text(), "\t")
		// string to float
		f, err := strconv.ParseFloat(dataStr[0], 64)
		if err != nil {
			// Ignore error
			fmt.Println(err)
		}
		sum += f
		str := dataStr[1:]
		rigFile.weights = append(rigFile.weights, f)
		rigFile.texts = append(rigFile.texts, str)
		if verbose {
			fmt.Println("P: ", f, str)
		}
	}
	rigFile.tot = sum
	return rigFile
}

func getNext(dict RigDict) Rig {
	rig := Rig{}
	if rand.Intn(2) == 0 {
		rig.Firstname = dict.fnames.texts[rand.Intn(len(dict.fnames.texts))][0]
	} else {
		rig.Firstname = dict.mnames.texts[rand.Intn(len(dict.mnames.texts))][0]
	}
	rig.Lastname = dict.lnames.texts[rand.Intn(len(dict.lnames.texts))][0]
	rig.Street = dict.Streets.texts[rand.Intn(len(dict.Streets.texts))][0]
	rig.Streetnumber = rand.Intn(50)
	zip := dict.zipcodes.texts[rand.Intn(len(dict.zipcodes.texts))]
	rig.City = strings.TrimSpace(zip[1])
	rig.Zipcode, _ = strconv.Atoi(zip[0])
	return rig
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"os"
	"text/template"

	"github.com/flynn/flynn-controller/client"
)

var cmdHelp = &Command{
	NoClient: true,
	Usage:    "help [topic]",
	Long:     `Help shows usage for a command or other topic.`,
}

func init() {
	cmdHelp.Run = runHelp // break init loop
}

func runHelp(cmd *Command, args []string, client *controller.Client) error {
	if len(args) == 0 {
		printUsage()
		return nil // not os.Exit(2); success
	}
	if len(args) != 1 {
		return errors.New("too many arguments")
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] {
			cmd.printUsage(false)
			return nil
		}
	}

	fmt.Fprintf(os.Stderr, "Unknown help topic: %q. Run 'flynn help'.\n", args[0])
	os.Exit(2)
	return nil
}

var cmdVersion = &Command{
	Run:      runVersion,
	NoClient: true,
	Usage:    "version",
	Short:    "show flynn version",
	Long:     `Version shows the flynn client version string.`,
}

func runVersion(cmd *Command, args []string, client *controller.Client) error {
	fmt.Println(Version)
	return nil
}

var usageTemplate = template.Must(template.New("usage").Parse(`
Usage: flynn [-a app] [command] [options] [arguments]

Commands:
{{range .Commands}}{{if .Runnable}}{{if .List}}
    {{.Name | printf "%-8s"}}  {{.Short}}{{end}}{{end}}{{end}}

Run 'flynn help [command]' for details.`[1:]))

func printUsage() {
	usageTemplate.Execute(os.Stdout, struct {
		Commands []*Command
	}{
		commands,
	})
}

func usage() {
	printUsage()
	os.Exit(2)
}
<commit_msg>cli: Align usage command output<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"text/template"

	"github.com/flynn/flynn-controller/client"
)

var cmdHelp = &Command{
	NoClient: true,
	Usage:    "help [topic]",
	Long:     `Help shows usage for a command or other topic.`,
}

func init() {
	cmdHelp.Run = runHelp // break init loop
}

func runHelp(cmd *Command, args []string, client *controller.Client) error {
	if len(args) == 0 {
		printUsage()
		return nil // not os.Exit(2); success
	}
	if len(args) != 1 {
		return errors.New("too many arguments")
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] {
			cmd.printUsage(false)
			return nil
		}
	}

	fmt.Fprintf(os.Stderr, "Unknown help topic: %q. Run 'flynn help'.\n", args[0])
	os.Exit(2)
	return nil
}

var cmdVersion = &Command{
	Run:      runVersion,
	NoClient: true,
	Usage:    "version",
	Short:    "show flynn version",
	Long:     `Version shows the flynn client version string.`,
}

func runVersion(cmd *Command, args []string, client *controller.Client) error {
	fmt.Println(Version)
	return nil
}

var usageTemplate = template.Must(template.New("usage").Parse(`
Usage: flynn [-a app] [command] [options] [arguments]

Commands:
{{range .Commands}}{{if .Runnable}}{{if .List}}
    {{.Name | printf (print "%-" $.MaxCommandWidth "s")}}  {{.Short}}{{end}}{{end}}{{end}}

Run 'flynn help [command]' for details.`[1:]))

func printUsage() {
	data := &struct {
		Commands        []*Command
		MaxCommandWidth int
	}{
		commands,
		0,
	}

	for _, cmd := range commands {
		if len(cmd.Name()) > data.MaxCommandWidth {
			data.MaxCommandWidth = len(cmd.Name())
		}
	}

	usageTemplate.Execute(os.Stdout, data)
}

func usage() {
	printUsage()
	os.Exit(2)
}
<|endoftext|>


<commit_before>package db_test

import (
	. "db"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"database/sql"

	// sql drivers
	_ "github.com/mattn/go-sqlite3"
)

var _ = Describe("Database", func() {
	Describe("Connecting to the database", func() {
		Context("With an invalid DSN", func() {
			It("should fail", func() {
				db := &DB{
					Driver: "invalid",
					DSN:    "does-not-matter",
				}

				Ω(db.Connect()).Should(HaveOccurred())
				Ω(db.Connected()).Should(BeFalse())
				Ω(db.Disconnect()).ShouldNot(HaveOccurred())
			})
		})

		Context("With an in-memory SQLite database", func() {
			It("should succeed", func() {
				db := &DB{
					Driver: "sqlite3",
					DSN:    ":memory:",
				}

				Ω(db.Connect()).ShouldNot(HaveOccurred())
				Ω(db.Connected()).Should(BeTrue())
				Ω(db.Disconnect()).ShouldNot(HaveOccurred())
			})
		})
	})

	Describe("Registering (cached) SQL queries", func() {
		var db *DB

		BeforeEach(func() {
			db = &DB{
				Driver: "sqlite3",
				DSN:    ":memory:",
			}

			Ω(db.Connect()).ShouldNot(HaveOccurred())
			Ω(db.Connected()).Should(BeTrue())
		})

		Context("With an empty query cache", func() {
			It("has nothing registered", func() {
				Ω(db.Cached("my-query")).Should(BeFalse())
			})

			It("can register a SQL command", func() {
				Ω(db.Cache("my-query", "SELECT * FROM table")).ShouldNot(HaveOccurred())
				Ω(db.Cached("my-query")).Should(BeTrue())
			})

			It("can register the same name multiple times", func() {
				Ω(db.Cache("my-query", "SELECT * FROM table")).ShouldNot(HaveOccurred())
				Ω(db.Cache("my-query", "SELECT * FROM other_table")).ShouldNot(HaveOccurred())
				Ω(db.Cached("my-query")).Should(BeTrue())
			})
		})
	})

	Describe("Running SQL queries", func() {
		var db *DB

		BeforeEach(func() {
			db = &DB{
				Driver: "sqlite3",
				DSN:    ":memory:",
			}

			Ω(db.Connect()).ShouldNot(HaveOccurred())

			Ω(db.Cache("schema", `CREATE TABLE things (type TEXT, number INTEGER)`)).
				ShouldNot(HaveOccurred())
			Ω(db.Cache("new-thing", `INSERT INTO things (type, number) VALUES (?, 0)`)).
				ShouldNot(HaveOccurred())
			Ω(db.Cache("increase", `UPDATE things SET number = number + ? WHERE type = ?`)).
				ShouldNot(HaveOccurred())
			Ω(db.Cache("how-many?", `SELECT number FROM things WHERE type = ?`)).
				ShouldNot(HaveOccurred())
		})

		AfterEach(func() {
			db.Disconnect()
		})

		Context("With an empty database", func() {
			It("can create tables", func() {
				Ω(db.Exec("schema")).ShouldNot(HaveOccurred())
			})
		})

		numberOfThingsIn := func(r *sql.Rows) int {
			var n int

			Ω(r).ShouldNot(BeNil())
			Ω(r.Next()).Should(BeTrue())
			Ω(r.Scan(&n)).ShouldNot(HaveOccurred())
			return n
		}

		Context("With an empty table", func() {
			BeforeEach(func() {
				db.Disconnect()
				Ω(db.Connect()).ShouldNot(HaveOccurred())

				Ω(db.Exec("schema")).ShouldNot(HaveOccurred())
			})

			It("can insert records", func() {
				Ω(db.Exec("new-thing", "monkey")).ShouldNot(HaveOccurred())

				r, err := db.Query("how-many?", "monkey")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(numberOfThingsIn(r)).Should(Equal(0))
			})

			It("can update records", func() {
				Ω(db.Exec("new-thing", "monkey")).ShouldNot(HaveOccurred())
				Ω(db.Exec("increase", 42, "monkey")).ShouldNot(HaveOccurred())

				r, err := db.Query("how-many?", "monkey")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(numberOfThingsIn(r)).Should(Equal(42))
			})

			It("raises an error if an uncached query is requested for Exec()", func() {
				Ω(db.Exec("not-a-query")).Should(HaveOccurred())
			})

			It("raises an error if an uncached query is requested for Query()", func() {
				r, err := db.Query("not-a-query")
				Ω(r).Should(BeNil())
				Ω(err).Should(HaveOccurred())
			})

			It("can run arbitrary SQL", func() {
				Ω(db.ExecOnce("INSERT INTO things (type, number) VALUES (?, ?)", "lion", 3)).
					ShouldNot(HaveOccurred())

				r, err := db.Query("how-many?", "lion")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(numberOfThingsIn(r)).Should(Equal(3))
			})
		})

		Context("With malformed SQL queries", func() {
			BeforeEach(func() {
				Ω(db.Cache("error", `DO STUFF IN SQL`)).
					ShouldNot(HaveOccurred())
			})

			It("propagates errors from sql driver", func() {
				Ω(db.Exec("errors")).Should(HaveOccurred())
			})
		})
	})
})
<commit_msg>Test for SELECT queries without arguments<commit_after>package db_test

import (
	. "db"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"database/sql"

	// sql drivers
	_ "github.com/mattn/go-sqlite3"
)

var _ = Describe("Database", func() {
	Describe("Connecting to the database", func() {
		Context("With an invalid DSN", func() {
			It("should fail", func() {
				db := &DB{
					Driver: "invalid",
					DSN:    "does-not-matter",
				}

				Ω(db.Connect()).Should(HaveOccurred())
				Ω(db.Connected()).Should(BeFalse())
				Ω(db.Disconnect()).ShouldNot(HaveOccurred())
			})
		})

		Context("With an in-memory SQLite database", func() {
			It("should succeed", func() {
				db := &DB{
					Driver: "sqlite3",
					DSN:    ":memory:",
				}

				Ω(db.Connect()).ShouldNot(HaveOccurred())
				Ω(db.Connected()).Should(BeTrue())
				Ω(db.Disconnect()).ShouldNot(HaveOccurred())
			})
		})
	})

	Describe("Registering (cached) SQL queries", func() {
		var db *DB

		BeforeEach(func() {
			db = &DB{
				Driver: "sqlite3",
				DSN:    ":memory:",
			}

			Ω(db.Connect()).ShouldNot(HaveOccurred())
			Ω(db.Connected()).Should(BeTrue())
		})

		Context("With an empty query cache", func() {
			It("has nothing registered", func() {
				Ω(db.Cached("my-query")).Should(BeFalse())
			})

			It("can register a SQL command", func() {
				Ω(db.Cache("my-query", "SELECT * FROM table")).ShouldNot(HaveOccurred())
				Ω(db.Cached("my-query")).Should(BeTrue())
			})

			It("can register the same name multiple times", func() {
				Ω(db.Cache("my-query", "SELECT * FROM table")).ShouldNot(HaveOccurred())
				Ω(db.Cache("my-query", "SELECT * FROM other_table")).ShouldNot(HaveOccurred())
				Ω(db.Cached("my-query")).Should(BeTrue())
			})
		})
	})

	Describe("Running SQL queries", func() {
		var db *DB

		BeforeEach(func() {
			db = &DB{
				Driver: "sqlite3",
				DSN:    ":memory:",
			}

			Ω(db.Connect()).ShouldNot(HaveOccurred())

			Ω(db.Cache("schema", `CREATE TABLE things (type TEXT, number INTEGER)`)).
				ShouldNot(HaveOccurred())
			Ω(db.Cache("new-thing", `INSERT INTO things (type, number) VALUES (?, 0)`)).
				ShouldNot(HaveOccurred())
			Ω(db.Cache("increase", `UPDATE things SET number = number + ? WHERE type = ?`)).
				ShouldNot(HaveOccurred())
			Ω(db.Cache("how-many?", `SELECT number FROM things WHERE type = ?`)).
				ShouldNot(HaveOccurred())
			Ω(db.Cache("how-many-monkeys?", `SELECT number FROM things WHERE type = "monkey"`)).
				ShouldNot(HaveOccurred())
		})

		AfterEach(func() {
			db.Disconnect()
		})

		Context("With an empty database", func() {
			It("can create tables", func() {
				Ω(db.Exec("schema")).ShouldNot(HaveOccurred())
			})
		})

		numberOfThingsIn := func(r *sql.Rows) int {
			var n int

			Ω(r).ShouldNot(BeNil())
			Ω(r.Next()).Should(BeTrue())
			Ω(r.Scan(&n)).ShouldNot(HaveOccurred())
			return n
		}

		Context("With an empty table", func() {
			BeforeEach(func() {
				db.Disconnect()
				Ω(db.Connect()).ShouldNot(HaveOccurred())

				Ω(db.Exec("schema")).ShouldNot(HaveOccurred())
			})

			It("can insert records", func() {
				Ω(db.Exec("new-thing", "monkey")).ShouldNot(HaveOccurred())

				r, err := db.Query("how-many?", "monkey")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(numberOfThingsIn(r)).Should(Equal(0))
			})

			It("can update records", func() {
				Ω(db.Exec("new-thing", "monkey")).ShouldNot(HaveOccurred())
				Ω(db.Exec("increase", 42, "monkey")).ShouldNot(HaveOccurred())

				r, err := db.Query("how-many?", "monkey")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(numberOfThingsIn(r)).Should(Equal(42))
			})

			It("can handle queries without arguments", func() {
				Ω(db.Exec("new-thing", "monkey")).ShouldNot(HaveOccurred())
				Ω(db.Exec("increase", 13, "monkey")).ShouldNot(HaveOccurred())

				r, err := db.Query("how-many-monkeys?")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(numberOfThingsIn(r)).Should(Equal(13))
			})

			It("raises an error if an uncached query is requested for Exec()", func() {
				Ω(db.Exec("not-a-query")).Should(HaveOccurred())
			})

			It("raises an error if an uncached query is requested for Query()", func() {
				r, err := db.Query("not-a-query")
				Ω(r).Should(BeNil())
				Ω(err).Should(HaveOccurred())
			})

			It("can run arbitrary SQL", func() {
				Ω(db.ExecOnce("INSERT INTO things (type, number) VALUES (?, ?)", "lion", 3)).
					ShouldNot(HaveOccurred())

				r, err := db.Query("how-many?", "lion")
				Ω(err).ShouldNot(HaveOccurred())
				Ω(numberOfThingsIn(r)).Should(Equal(3))
			})
		})

		Context("With malformed SQL queries", func() {
			BeforeEach(func() {
				Ω(db.Cache("error", `DO STUFF IN SQL`)).
					ShouldNot(HaveOccurred())
			})

			It("propagates errors from sql driver", func() {
				Ω(db.Exec("errors")).Should(HaveOccurred())
			})
		})
	})
})
<|endoftext|>


<commit_before>package victims

import (
	"strings"
	"testing"

	"github.com/asobti/kube-monkey/config"
	"github.com/stretchr/testify/assert"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kube "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/fake"
)

const (
	NAMESPACE  = metav1.NamespaceDefault
	IDENTIFIER = "kube-monkey-id"
)

func newPod(name string, status v1.PodPhase) v1.Pod {

	return v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: NAMESPACE,
			Labels: map[string]string{
				"kube-monkey/identifier": IDENTIFIER,
			},
		},
		Status: v1.PodStatus{
			Phase: status,
		},
	}

}

func newVictimBase() *VictimBase {
	return New("Pod", "name", NAMESPACE, IDENTIFIER, 1)
}

func getPodList(client kube.Interface) *v1.PodList {
	podList, _ := client.CoreV1().Pods(NAMESPACE).List(metav1.ListOptions{})
	return podList
}

func TestVictimBaseTemplateGetters(t *testing.T) {

	v := newVictimBase()

	assert.Equal(t, "Pod", v.Kind())
	assert.Equal(t, "name", v.Name())
	assert.Equal(t, NAMESPACE, v.Namespace())
	assert.Equal(t, IDENTIFIER, v.Identifier())
	assert.Equal(t, 1, v.Mtbf())
}

func TestRunningPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	client := fake.NewSimpleClientset(&pod1, &pod2)

	podList, e := v.RunningPods(client)

	if e != nil {
		t.Errorf("Unexpected error %s while getting running pod", e)
	}

	if len(podList) != 1 {
		t.Errorf("Expected 1 item in podList, got %d", len(podList))
	}
}

func TestPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	client := fake.NewSimpleClientset(&pod1, &pod2)

	podList, _ := v.Pods(client)

	if len(podList) != 2 {
		t.Errorf("Expected 2 item in podList, got %d", len(podList))
	}
}

func TestDeletePod(t *testing.T) {

	v := newVictimBase()
	pod := newPod("app", v1.PodRunning)

	client := fake.NewSimpleClientset(&pod)

	if e := v.DeletePod(client, "app"); e != nil {
		t.Errorf("Unexpected error %s while deleting pod", e)
	}

	if podList := getPodList(client); len(podList.Items) != 0 {
		t.Errorf("Expected 0 items in podList, got %d", len(podList.Items))
	}
}

func TestDeleteRandomPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)
	pod3 := newPod("app3", v1.PodRunning)

	client := fake.NewSimpleClientset(&pod1, &pod2, &pod3)

	if e := v.DeleteRandomPods(client, 0); e == nil {
		t.Error("Expected an error if no termination was requested")
	}

	if e := v.DeleteRandomPods(client, -1); e == nil {
		t.Error("Expected an error if negative termination was requested")
	}

	_ = v.DeleteRandomPods(client, 1)

	if podList := getPodList(client); len(podList.Items) != 2 {
		t.Errorf("Expected 2 items in podList, got %d", len(podList.Items))
	}

	_ = v.DeleteRandomPods(client, 2)

	if podList := getPodList(client); len(podList.Items) != 1 {
		t.Errorf("Expected 1 items in podList, got %d", len(podList.Items))
	}

	if e := v.DeleteRandomPods(client, 2); e == nil {
		t.Error("Expected an error if no pods")
	}
}

func TestTerminateAllPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)
	pod3 := newPod("app3", v1.PodRunning)

	client := fake.NewSimpleClientset(&pod1, &pod2, &pod3)

	_ = v.TerminateAllPods(client)

	if podList := getPodList(client); len(podList.Items) != 0 {
		t.Errorf("Expected 0 items in podList, got %d", len(podList.Items))
	}

	if e := v.TerminateAllPods(client); e == nil {
		t.Error("Expected an error if no pods")
	}

}

func TestDeleteRandomPod(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	client := fake.NewSimpleClientset(&pod1, &pod2)

	_ = v.DeleteRandomPod(client)

	if podList := getPodList(client); len(podList.Items) != 1 {
		t.Errorf("Expected 1 items in podList, got %d", len(podList.Items))
	}

	if e := v.DeleteRandomPods(client, 2); e == nil {
		t.Error("Expected an error if no pods")
	}
}

func TestIsBlacklisted(t *testing.T) {

	v := newVictimBase()

	config.SetDefaults()

	if e := v.IsBlacklisted(); e != false {
		t.Errorf("%s namespace should not be blacklisted", NAMESPACE)
	}

	v = New("Pod", "name", metav1.NamespaceSystem, IDENTIFIER, 1)

	if e := v.IsBlacklisted(); e != true {
		t.Errorf("%s namespace should be blacklisted", metav1.NamespaceSystem)
	}

}

func TestIsWhitelisted(t *testing.T) {

	v := newVictimBase()

	config.SetDefaults()

	if e := v.IsWhitelisted(); e != true {
		t.Errorf("%s namespace should be whitelisted", NAMESPACE)
	}
}

func TestRandomPodName(t *testing.T) {

	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	if name := RandomPodName([]v1.Pod{pod1, pod2}); strings.HasPrefix(name, "pod") {
		t.Errorf("Pod name %s should start with 'app'", name)
	}

}
<commit_msg>Additional UT for victims package<commit_after>package victims

import (
	"strings"
	"testing"

	"github.com/asobti/kube-monkey/config"
	"github.com/stretchr/testify/assert"
	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kube "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/fake"
)

const (
	NAMESPACE  = metav1.NamespaceDefault
	IDENTIFIER = "kube-monkey-id"
)

func newPod(name string, status v1.PodPhase) v1.Pod {

	return v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: NAMESPACE,
			Labels: map[string]string{
				"kube-monkey/identifier": IDENTIFIER,
			},
		},
		Status: v1.PodStatus{
			Phase: status,
		},
	}

}

func newVictimBase() *VictimBase {
	return New("Pod", "name", NAMESPACE, IDENTIFIER, 1)
}

func getPodList(client kube.Interface) *v1.PodList {
	podList, _ := client.CoreV1().Pods(NAMESPACE).List(metav1.ListOptions{})
	return podList
}

func TestVictimBaseTemplateGetters(t *testing.T) {

	v := newVictimBase()

	assert.Equal(t, "Pod", v.Kind())
	assert.Equal(t, "name", v.Name())
	assert.Equal(t, NAMESPACE, v.Namespace())
	assert.Equal(t, IDENTIFIER, v.Identifier())
	assert.Equal(t, 1, v.Mtbf())
}

func TestRunningPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	client := fake.NewSimpleClientset(&pod1, &pod2)

	podList, e := v.RunningPods(client)

	if e != nil {
		t.Errorf("Unexpected error %s while getting running pod", e)
	}

	if len(podList) != 1 {
		t.Errorf("Expected 1 item in podList, got %d", len(podList))
	}

	if name := podList[0].GetName(); name != "app1" {
		t.Errorf("Unexpected pod name, got %s", name)
	}
}

func TestPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	client := fake.NewSimpleClientset(&pod1, &pod2)

	podList, _ := v.Pods(client)

	if len(podList) != 2 {
		t.Errorf("Expected 2 item in podList, got %d", len(podList))
	}
}

func TestDeletePod(t *testing.T) {

	v := newVictimBase()
	pod := newPod("app", v1.PodRunning)

	client := fake.NewSimpleClientset(&pod)

	if e := v.DeletePod(client, "app"); e != nil {
		t.Errorf("Unexpected error %s while deleting pod", e)
	}

	if podList := getPodList(client); len(podList.Items) != 0 {
		t.Errorf("Expected 0 items in podList, got %d", len(podList.Items))
	}
}

func TestDeleteRandomPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)
	pod3 := newPod("app3", v1.PodRunning)

	client := fake.NewSimpleClientset(&pod1, &pod2, &pod3)

	if e := v.DeleteRandomPods(client, 0); e == nil {
		t.Error("Expected an error if no termination was requested")
	}

	if e := v.DeleteRandomPods(client, -1); e == nil {
		t.Error("Expected an error if negative termination was requested")
	}

	_ = v.DeleteRandomPods(client, 1)

	if podList := getPodList(client); len(podList.Items) != 2 {
		t.Errorf("Expected 2 items in podList, got %d", len(podList.Items))
	}

	_ = v.DeleteRandomPods(client, 2)

	if podList := getPodList(client); len(podList.Items) != 1 {
		t.Errorf("Expected 1 items in podList, got %d", len(podList.Items))
	} else if name := podList.Items[0].GetName(); name != "app2" {
		t.Error("Expected not running pods not be deleted")
	}

	if e := v.DeleteRandomPods(client, 2); e == nil {
		t.Error("Expected an error if no pods")
	}
}

func TestTerminateAllPods(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)
	pod3 := newPod("app3", v1.PodRunning)

	client := fake.NewSimpleClientset(&pod1, &pod2, &pod3)

	_ = v.TerminateAllPods(client)

	if podList := getPodList(client); len(podList.Items) != 0 {
		t.Errorf("Expected 0 items in podList, got %d", len(podList.Items))
	}

	if e := v.TerminateAllPods(client); e == nil {
		t.Error("Expected an error if no pods")
	}

}

func TestDeleteRandomPod(t *testing.T) {

	v := newVictimBase()
	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	client := fake.NewSimpleClientset(&pod1, &pod2)

	_ = v.DeleteRandomPod(client)

	if podList := getPodList(client); len(podList.Items) != 1 {
		t.Errorf("Expected 1 items in podList, got %d", len(podList.Items))
	}

	if e := v.DeleteRandomPods(client, 2); e == nil {
		t.Error("Expected an error if no pods")
	}
}

func TestIsBlacklisted(t *testing.T) {

	v := newVictimBase()

	config.SetDefaults()

	if e := v.IsBlacklisted(); e != false {
		t.Errorf("%s namespace should not be blacklisted", NAMESPACE)
	}

	v = New("Pod", "name", metav1.NamespaceSystem, IDENTIFIER, 1)

	if e := v.IsBlacklisted(); e != true {
		t.Errorf("%s namespace should be blacklisted", metav1.NamespaceSystem)
	}

}

func TestIsWhitelisted(t *testing.T) {

	v := newVictimBase()

	config.SetDefaults()

	if e := v.IsWhitelisted(); e != true {
		t.Errorf("%s namespace should be whitelisted", NAMESPACE)
	}
}

func TestRandomPodName(t *testing.T) {

	pod1 := newPod("app1", v1.PodRunning)
	pod2 := newPod("app2", v1.PodPending)

	if name := RandomPodName([]v1.Pod{pod1, pod2}); strings.HasPrefix(name, "pod") {
		t.Errorf("Pod name %s should start with 'app'", name)
	}

}
<|endoftext|>


<commit_before>package hal

/*
 * Copyright 2016 Albert P. Tobey <atobey@netflix.com>
 *
 * 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.
 */

import (
	"fmt"
	"log"
	"sync"
)

// RouterCTX holds the router's context, including input/output chans.
type RouterCTX struct {
	brokers map[string]Broker
	in      chan *Evt     // messages from brokers --> plugins
	out     chan *Evt     // messages from plugins --> brokers
	update  chan struct{} // to notify the router that the instance list changed
	mut     sync.Mutex
	init    sync.Once
}

type fwdBroker struct {
	from Broker
	to   Broker
}

var routerSingleton RouterCTX

// Router returns the singleton router context. The router is initialized
// on the first call to this function.
func Router() *RouterCTX {
	routerSingleton.init.Do(func() {
		routerSingleton.in = make(chan *Evt, 1000)
		routerSingleton.out = make(chan *Evt, 1000)
		routerSingleton.update = make(chan struct{}, 1)
		routerSingleton.brokers = make(map[string]Broker)
	})

	return &routerSingleton
}

// forwardChan forwards events from one chan of to another.
// TODO: figure out if this needs to check for closed channels, etc.
func forwardChan(from, to chan *Evt) {
	for {
		select {
		case evt := <-from:
			to <- evt
		}
	}
}

// AddBroker adds a broker to the router and starts forwarding
// events between it and the router.
func (r *RouterCTX) AddBroker(b Broker) {
	r.mut.Lock()
	defer r.mut.Unlock()

	if _, exists := r.brokers[b.Name()]; exists {
		panic(fmt.Sprintf("BUG: broker '%s' added > 1 times.", b.Name()))
	}

	b2r := make(chan *Evt, 1000) // messages from the broker to the router

	// start the broker's event stream
	go b.Stream(b2r)

	// forward events from the broker to the router's input channel
	go forwardChan(b2r, r.in)

	r.brokers[b.Name()] = b
}

// GetBroker retrieves a broker handle by name.
func (r *RouterCTX) GetBroker(name string) Broker {
	r.mut.Lock()
	defer r.mut.Unlock()

	if broker, exists := r.brokers[name]; exists {
		return broker
	}

	return nil
}

// Brokers returns all brokers that have been added to the router.
// The returned list is not in any particular order.
func (r *RouterCTX) Brokers() []Broker {
	r.mut.Lock()
	defer r.mut.Unlock()

	out := make([]Broker, len(r.brokers))
	i := 0
	for _, b := range r.brokers {
		out[i] = b
		i++
	}

	return out
}

// Route is the main method for the router. It blocks and should be run in a
// goroutine exactly once. Running more than one router in the same process
// will result in shenanigans.
func (r *RouterCTX) Route() {
	for {
		select {
		case evt := <-r.in:
			// events are processed concurrently, plugins are not
			go r.processEvent(evt)
		}
	}
}

// processEvent processes one event and is intended to run in a goroutine.
func (r *RouterCTX) processEvent(evt *Evt) {
	var pname string // must be in the recovery handler's scope

	// get a snapshot of the instance list
	// TODO: keep an eye on the cost of copying this list for every message
	pr := PluginRegistry()
	instances := pr.InstanceList()

	// if a plugin panics, catch it & log it
	// TODO: report errors to a channel?
	defer func() {
		if r := recover(); r != nil {
			log.Printf("recovered panic in plugin %q\n", pname)
			log.Printf("panic: %q", r)
		}
	}()

	for _, inst := range instances {
		// the recovery handler will pick this up in a panic to provide
		// the name of the plugin that caused the panic
		pname = inst.Plugin.Name

		// check if it's the correct room
		if evt.RoomId != inst.RoomId {
			continue
		}

		// finally, check message text against the regex
		if inst.Regex == "" || inst.regex.MatchString(evt.Body) {
			// this will copy the struct twice. It's intentional to avoid
			// mutating the evt between calls. The plugin func signature
			// forces the second copy.
			evtcpy := *evt

			// pass the plugin instance pointer to the plugin function so
			// it can access its fields for settings, etc.
			evtcpy.instance = inst

			// call the plugin function
			// this may block other plugins from processing the same event but
			// since it's already in a goroutine, other events won't be blocked
			inst.Func(evtcpy)
		}
	}
}
<commit_msg>print stack traces after recovering plugin crashes<commit_after>package hal

/*
 * Copyright 2016 Albert P. Tobey <atobey@netflix.com>
 *
 * 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.
 */

import (
	"fmt"
	"log"
	"runtime/debug"
	"sync"
)

// RouterCTX holds the router's context, including input/output chans.
type RouterCTX struct {
	brokers map[string]Broker
	in      chan *Evt     // messages from brokers --> plugins
	out     chan *Evt     // messages from plugins --> brokers
	update  chan struct{} // to notify the router that the instance list changed
	mut     sync.Mutex
	init    sync.Once
}

type fwdBroker struct {
	from Broker
	to   Broker
}

var routerSingleton RouterCTX

// Router returns the singleton router context. The router is initialized
// on the first call to this function.
func Router() *RouterCTX {
	routerSingleton.init.Do(func() {
		routerSingleton.in = make(chan *Evt, 1000)
		routerSingleton.out = make(chan *Evt, 1000)
		routerSingleton.update = make(chan struct{}, 1)
		routerSingleton.brokers = make(map[string]Broker)
	})

	return &routerSingleton
}

// forwardChan forwards events from one chan of to another.
// TODO: figure out if this needs to check for closed channels, etc.
func forwardChan(from, to chan *Evt) {
	for {
		select {
		case evt := <-from:
			to <- evt
		}
	}
}

// AddBroker adds a broker to the router and starts forwarding
// events between it and the router.
func (r *RouterCTX) AddBroker(b Broker) {
	r.mut.Lock()
	defer r.mut.Unlock()

	if _, exists := r.brokers[b.Name()]; exists {
		panic(fmt.Sprintf("BUG: broker '%s' added > 1 times.", b.Name()))
	}

	b2r := make(chan *Evt, 1000) // messages from the broker to the router

	// start the broker's event stream
	go b.Stream(b2r)

	// forward events from the broker to the router's input channel
	go forwardChan(b2r, r.in)

	r.brokers[b.Name()] = b
}

// GetBroker retrieves a broker handle by name.
func (r *RouterCTX) GetBroker(name string) Broker {
	r.mut.Lock()
	defer r.mut.Unlock()

	if broker, exists := r.brokers[name]; exists {
		return broker
	}

	return nil
}

// Brokers returns all brokers that have been added to the router.
// The returned list is not in any particular order.
func (r *RouterCTX) Brokers() []Broker {
	r.mut.Lock()
	defer r.mut.Unlock()

	out := make([]Broker, len(r.brokers))
	i := 0
	for _, b := range r.brokers {
		out[i] = b
		i++
	}

	return out
}

// Route is the main method for the router. It blocks and should be run in a
// goroutine exactly once. Running more than one router in the same process
// will result in shenanigans.
func (r *RouterCTX) Route() {
	for {
		select {
		case evt := <-r.in:
			// events are processed concurrently, plugins are not
			go r.processEvent(evt)
		}
	}
}

// processEvent processes one event and is intended to run in a goroutine.
func (r *RouterCTX) processEvent(evt *Evt) {
	var pname string // must be in the recovery handler's scope

	// get a snapshot of the instance list
	// TODO: keep an eye on the cost of copying this list for every message
	pr := PluginRegistry()
	instances := pr.InstanceList()

	// if a plugin panics, catch it & log it
	// TODO: report errors to a channel?
	defer func() {
		if r := recover(); r != nil {
			log.Printf("recovered panic in plugin %q\n", pname)
			log.Printf("panic: %q", r)
			debug.PrintStack()
		}
	}()

	for _, inst := range instances {
		// the recovery handler will pick this up in a panic to provide
		// the name of the plugin that caused the panic
		pname = inst.Plugin.Name

		// check if it's the correct room
		if evt.RoomId != inst.RoomId {
			continue
		}

		// finally, check message text against the regex
		if inst.Regex == "" || inst.regex.MatchString(evt.Body) {
			// this will copy the struct twice. It's intentional to avoid
			// mutating the evt between calls. The plugin func signature
			// forces the second copy.
			evtcpy := *evt

			// pass the plugin instance pointer to the plugin function so
			// it can access its fields for settings, etc.
			evtcpy.instance = inst

			// call the plugin function
			// this may block other plugins from processing the same event but
			// since it's already in a goroutine, other events won't be blocked
			inst.Func(evtcpy)
		}
	}
}
<|endoftext|>


<commit_before>package rollrus

import (
	"fmt"
	"runtime"
	"strings"
	"time"

	"github.com/rollbar/rollbar-go"
	"github.com/sirupsen/logrus"
)

var _ logrus.Hook = &Hook{} //assert that *Hook is a logrus.Hook

// Hook is a wrapper for the Rollbar Client and is usable as a logrus.Hook.
type Hook struct {
	*rollbar.Client
	triggers        []logrus.Level
	ignoredErrors   []error
	ignoreErrorFunc func(error) bool
	ignoreFunc      func(error, map[string]interface{}) bool

	// only used for tests to verify whether or not a report happened.
	reported bool
}

// NewHookForLevels provided by the caller. Otherwise works like NewHook.
func NewHookForLevels(token string, env string, levels []logrus.Level) *Hook {
	return &Hook{
		Client:          rollbar.NewSync(token, env, "", "", ""),
		triggers:        levels,
		ignoredErrors:   make([]error, 0),
		ignoreErrorFunc: func(error) bool { return false },
		ignoreFunc:      func(error, map[string]interface{}) bool { return false },
	}
}

// ReportPanic attempts to report the panic to Rollbar using the provided
// client and then re-panic. If it can't report the panic it will print an
// error to stderr.
func (r *Hook) ReportPanic() {
	if p := recover(); p != nil {
		r.Client.ErrorWithLevel(rollbar.CRIT, fmt.Errorf("panic: %q", p))
		panic(p)
	}
}

// Levels returns the logrus log.Levels that this hook handles
func (r *Hook) Levels() []logrus.Level {
	if r.triggers == nil {
		return defaultTriggerLevels
	}
	return r.triggers
}

// Fire the hook. This is called by Logrus for entries that match the levels
// returned by Levels().
func (r *Hook) Fire(entry *logrus.Entry) error {
	cause := extractError(entry)
	for _, ie := range r.ignoredErrors {
		if ie == cause {
			return nil
		}
	}

	if r.ignoreErrorFunc(cause) {
		return nil
	}

	m := convertFields(entry.Data)
	if _, exists := m["time"]; !exists {
		m["time"] = entry.Time.Format(time.RFC3339)
	}

	if _, exists := m["msg"]; !exists && entry.Message != "" {
		m["msg"] = entry.Message
	}

	if r.ignoreFunc(cause, m) {
		return nil
	}

	r.report(entry, cause, m)

	return nil
}

func (r *Hook) report(entry *logrus.Entry, cause error, m map[string]interface{}) {
	level := entry.Level

	r.reported = true

	switch {
	case level == logrus.FatalLevel || level == logrus.PanicLevel:
		skip := framesToSkip(2)
		r.Client.ErrorWithStackSkipWithExtras(rollbar.CRIT, cause, skip, m)
	case level == logrus.ErrorLevel:
		skip := framesToSkip(2)
		r.Client.ErrorWithStackSkipWithExtras(rollbar.ERR, cause, skip, m)
	case level == logrus.WarnLevel:
		skip := framesToSkip(2)
		r.Client.ErrorWithStackSkipWithExtras(rollbar.WARN, cause, skip, m)
	case level == logrus.InfoLevel:
		r.Client.MessageWithExtras(rollbar.INFO, entry.Message, m)
	case level == logrus.DebugLevel:
		r.Client.MessageWithExtras(rollbar.DEBUG, entry.Message, m)
	}
}

// convertFields converts from log.Fields to map[string]interface{} so that we can
// report extra fields to Rollbar
func convertFields(fields logrus.Fields) map[string]interface{} {
	m := make(map[string]interface{})
	for k, v := range fields {
		switch t := v.(type) {
		case time.Time:
			m[k] = t.Format(time.RFC3339)
		case error:
			m[k] = t.Error()
		default:
			if s, ok := v.(fmt.Stringer); ok {
				m[k] = s.String()
			} else {
				m[k] = fmt.Sprintf("%+v", t)
			}
		}
	}

	return m
}

// extractError attempts to extract an error from a well known field, err or error
func extractError(entry *logrus.Entry) error {
	for _, f := range wellKnownErrorFields {
		e, ok := entry.Data[f]
		if !ok {
			continue
		}
		err, ok := e.(error)
		if !ok {
			continue
		}

		return err
	}

	// when no error found, default to the logged message.
	return fmt.Errorf(entry.Message)
}

// framesToSkip returns the number of caller frames to skip
// to get a stack trace that excludes rollrus and logrus.
func framesToSkip(rollrusSkip int) int {
	// skip 1 to get out of this function
	skip := rollrusSkip + 1

	// to get out of logrus, the amount can vary
	// depending on how the user calls the log functions
	// figure it out dynamically by skipping until
	// we're out of the logrus package
	for i := skip; ; i++ {
		_, file, _, ok := runtime.Caller(i)
		if !ok || !strings.Contains(file, "github.com/sirupsen/logrus") {
			skip = i
			break
		}
	}

	// rollbar-go is skipping too few frames (2)
	// subtract 1 since we're currently working from a function
	return skip + 2 - 1
}
<commit_msg>Fix ignore error tests<commit_after>package rollrus

import (
	"fmt"
	"runtime"
	"strings"
	"time"

	"github.com/rollbar/rollbar-go"
	"github.com/sirupsen/logrus"
)

var _ logrus.Hook = &Hook{} //assert that *Hook is a logrus.Hook

// Hook is a wrapper for the Rollbar Client and is usable as a logrus.Hook.
type Hook struct {
	*rollbar.Client
	triggers        []logrus.Level
	ignoredErrors   []error
	ignoreErrorFunc func(error) bool
	ignoreFunc      func(error, map[string]interface{}) bool

	// only used for tests to verify whether or not a report happened.
	reported bool
}

// NewHookForLevels provided by the caller. Otherwise works like NewHook.
func NewHookForLevels(token string, env string, levels []logrus.Level) *Hook {
	return &Hook{
		Client:          rollbar.NewSync(token, env, "", "", ""),
		triggers:        levels,
		ignoredErrors:   make([]error, 0),
		ignoreErrorFunc: func(error) bool { return false },
		ignoreFunc:      func(error, map[string]interface{}) bool { return false },
	}
}

// ReportPanic attempts to report the panic to Rollbar using the provided
// client and then re-panic. If it can't report the panic it will print an
// error to stderr.
func (r *Hook) ReportPanic() {
	if p := recover(); p != nil {
		r.Client.ErrorWithLevel(rollbar.CRIT, fmt.Errorf("panic: %q", p))
		panic(p)
	}
}

// Levels returns the logrus log.Levels that this hook handles
func (r *Hook) Levels() []logrus.Level {
	if r.triggers == nil {
		return defaultTriggerLevels
	}
	return r.triggers
}

// Fire the hook. This is called by Logrus for entries that match the levels
// returned by Levels().
func (r *Hook) Fire(entry *logrus.Entry) error {
	err := extractError(entry)
	cause := errorCause(err)
	for _, ie := range r.ignoredErrors {
		if ie == cause {
			return nil
		}
	}

	if r.ignoreErrorFunc(cause) {
		return nil
	}

	m := convertFields(entry.Data)
	if _, exists := m["time"]; !exists {
		m["time"] = entry.Time.Format(time.RFC3339)
	}

	if _, exists := m["msg"]; !exists && entry.Message != "" {
		m["msg"] = entry.Message
	}

	if r.ignoreFunc(cause, m) {
		return nil
	}

	r.report(entry, err, m)

	return nil
}

func (r *Hook) report(entry *logrus.Entry, cause error, m map[string]interface{}) {
	level := entry.Level

	r.reported = true

	switch {
	case level == logrus.FatalLevel || level == logrus.PanicLevel:
		skip := framesToSkip(2)
		r.Client.ErrorWithStackSkipWithExtras(rollbar.CRIT, cause, skip, m)
	case level == logrus.ErrorLevel:
		skip := framesToSkip(2)
		r.Client.ErrorWithStackSkipWithExtras(rollbar.ERR, cause, skip, m)
	case level == logrus.WarnLevel:
		skip := framesToSkip(2)
		r.Client.ErrorWithStackSkipWithExtras(rollbar.WARN, cause, skip, m)
	case level == logrus.InfoLevel:
		r.Client.MessageWithExtras(rollbar.INFO, entry.Message, m)
	case level == logrus.DebugLevel:
		r.Client.MessageWithExtras(rollbar.DEBUG, entry.Message, m)
	}
}

// convertFields converts from log.Fields to map[string]interface{} so that we can
// report extra fields to Rollbar
func convertFields(fields logrus.Fields) map[string]interface{} {
	m := make(map[string]interface{})
	for k, v := range fields {
		switch t := v.(type) {
		case time.Time:
			m[k] = t.Format(time.RFC3339)
		case error:
			m[k] = t.Error()
		default:
			if s, ok := v.(fmt.Stringer); ok {
				m[k] = s.String()
			} else {
				m[k] = fmt.Sprintf("%+v", t)
			}
		}
	}

	return m
}

// extractError attempts to extract an error from a well known field, err or error
func extractError(entry *logrus.Entry) error {
	for _, f := range wellKnownErrorFields {
		e, ok := entry.Data[f]
		if !ok {
			continue
		}
		err, ok := e.(error)
		if !ok {
			continue
		}

		return err
	}

	// when no error found, default to the logged message.
	return fmt.Errorf(entry.Message)
}

// framesToSkip returns the number of caller frames to skip
// to get a stack trace that excludes rollrus and logrus.
func framesToSkip(rollrusSkip int) int {
	// skip 1 to get out of this function
	skip := rollrusSkip + 1

	// to get out of logrus, the amount can vary
	// depending on how the user calls the log functions
	// figure it out dynamically by skipping until
	// we're out of the logrus package
	for i := skip; ; i++ {
		_, file, _, ok := runtime.Caller(i)
		if !ok || !strings.Contains(file, "github.com/sirupsen/logrus") {
			skip = i
			break
		}
	}

	// rollbar-go is skipping too few frames (2)
	// subtract 1 since we're currently working from a function
	return skip + 2 - 1
}

func errorCause(err error) error {
	type causer interface {
		Cause() error
	}

	for err != nil {
		cause, ok := err.(causer)
		if !ok {
			break
		}
		err = cause.Cause()
	}
	return err
}
<|endoftext|>


<commit_before>package hotp

import (
	"code.google.com/p/rsc/qr"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/asn1"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"math/big"
	"net/url"
	"strconv"
)

// RFC 4226 specifies the counter as being 8 bytes.
const ctrSize = 8

// ErrInvalidHOTPURL is returned via FromURL; it indicates a malformed
// HOTP otpauth URL.
var ErrInvalidHOTPURL = errors.New("hotp: invalid HOTP url")

// PRNG is the source of random data; this is used by GenerateHOTP
// and should be a cryptographically-secure PRNG.
var PRNG = rand.Reader

// HOTP represents a new key value for generating one-time passwords;
// it contains the key used to construct one-time passwords and the
// counter state used in the OTP generation. Digits contains the
// number of digits that generated OTPs should output. Key is a
// cryptographic secret, and should be treated as such.
type HOTP struct {
	Key     []byte
	counter *[ctrSize]byte
	Digits  int
}

// Counter returns the HOTP's 8-byte counter as an unsigned 64-bit
// integer.
func (otp HOTP) Counter() uint64 {
	counter := binary.BigEndian.Uint64(otp.counter[:])
	return counter
}

// Increment will increment an HOTP source's counter. This is useful
// for providers like the Google Authenticator app, which immediately
// increments the counter and uses the 0 counter value as an integrity
// check.
func (otp HOTP) Increment() {
	for i := ctrSize - 1; i >= 0; i-- {
		if otp.counter[i]++; otp.counter[i] != 0 {
			return
		}
	}
}

// OTP generates a new one-time password.
func (otp HOTP) OTP() string {
	h := hmac.New(sha1.New, otp.Key)
	h.Write(otp.counter[:])
	otp.Increment()
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64())
}

func (otp *HOTP) setCounter(counter uint64) bool {
	if otp.counter == nil {
		otp.counter = new([ctrSize]byte)
	}
	binary.BigEndian.PutUint64(otp.counter[:], counter)
	return true
}

// NewHOTP intialises a new HOTP instance with the key and counter
// values. No check is done on the digits, but typical values are 6
// and 8.
func NewHOTP(key []byte, counter uint64, digits int) *HOTP {
	otp := &HOTP{
		Key:    key,
		Digits: digits,
	}
	otp.counter = new([ctrSize]byte)
	binary.BigEndian.PutUint64(otp.counter[:], counter)

	return otp
}

// URL returns a suitable URL, such as for the Google Authenticator
// app. The label is used by these apps to identify the service to
// which this OTP belongs. The digits value is ignored by the Google
// authenticator app, and is therefore elided in the resulting URL.
func (otp *HOTP) URL(label string) string {
	secret := base32.StdEncoding.EncodeToString(otp.Key)
	u := url.URL{}
	v := url.Values{}
	u.Scheme = "otpauth"
	u.Host = "hotp"
	u.Path = label
	v.Add("secret", secret)
	v.Add("counter", fmt.Sprintf("%d", otp.Counter()))
	u.RawQuery = v.Encode()
	return u.String()
}

// QR generates a byte slice containing the a QR code encoded as a
// PNG with level Q error correction.
func (otp *HOTP) QR(label string) ([]byte, error) {
	u := otp.URL(label)
	code, err := qr.Encode(u, qr.Q)
	if err != nil {
		return nil, err
	}
	return code.PNG(), nil
}

// truncate contains the DT function from the RFC; this is used to
// deterministically select a sequence of 4 bytes from the HMAC
// counter hash.
func truncate(in []byte) int64 {
	offset := int(in[len(in)-1] & 0xF)
	p := in[offset : offset+4]
	var binCode int32
	binCode = int32((p[0] & 0x7f)) << 24
	binCode += int32((p[1] & 0xff)) << 16
	binCode += int32((p[2] & 0xff)) << 8
	binCode += int32((p[3] & 0xff))
	return int64(binCode) & 0x7FFFFFFF
}

// FromURL parses a new HOTP from a URL string. It returns the OTP,
// the label associated with the OTP, and any errors that occurred.
func FromURL(urlString string) (*HOTP, string, error) {
	u, err := url.Parse(urlString)
	if err != nil {
		return nil, "", err
	}

	if u.Scheme != "otpauth" {
		return nil, "", ErrInvalidHOTPURL
	} else if u.Host != "hotp" {
		return nil, "", ErrInvalidHOTPURL
	}

	v := u.Query()
	if len(v) == 0 {
		return nil, "", ErrInvalidHOTPURL
	}
	if v.Get("secret") == "" {
		return nil, "", ErrInvalidHOTPURL
	} else if algo := v.Get("algorithm"); algo != "" && algo != "SHA1" {
		return nil, "", ErrInvalidHOTPURL
	}

	var identity string
	if len(u.Path) > 1 {
		identity = u.Path[1:]
	}

	var counter uint64
	if ctr := v.Get("counter"); ctr != "" {
		counter, err = strconv.ParseUint(ctr, 10, 64)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	secret, err := base32.StdEncoding.DecodeString(v.Get("secret"))
	if err != nil {
		return nil, "", ErrInvalidHOTPURL
	}

	var digits int64 = 6
	if v.Get("digits") != "" {
		digits, err = strconv.ParseInt(v.Get("digits"), 10, 8)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	otp := NewHOTP(secret, counter, int(digits))
	return otp, identity, nil
}

// GenerateHOTP will generate a randomised HOTP source; if the
// randCounter parameter is true, the counter will be randomised.
func GenerateHOTP(digits int, randCounter bool) (*HOTP, error) {
	key := make([]byte, sha1.Size)
	_, err := io.ReadFull(PRNG, key)
	if err != nil {
		return nil, err
	}

	var counter uint64
	if randCounter {
		ctr, err := rand.Int(PRNG, big.NewInt(int64(math.MaxInt64)))
		if err != nil {
			return nil, err
		}
		counter = ctr.Uint64()
	}

	return NewHOTP(key, counter, digits), nil
}

// YubiKey reads an OATH-HOTP string as returned by a YubiKey, and
// returns three values. The first value contains the actual OTP, the
// second value contains the YubiKey's token identifier, and the final
// value indicates whether the input string was a valid YubiKey
// OTP. This does not check whether the code is correct or not, it
// only ensures that it is well-formed output from a token and
// splits the output into the code and the public identity.
func (otp *HOTP) YubiKey(in string) (string, string, bool) {
	if len(in) < otp.Digits {
		return "", "", false
	}

	otpStart := len(in) - otp.Digits
	code := in[otpStart:]
	pubid := in[:otpStart]
	return code, pubid, true
}

// IntegrityCheck returns two values, the base OTP and the current
// counter. This is used, for example, with the Google Authenticator
// app's "Check key value" function and can be used to verify that
// the application and the provider are in sync.
func (otp *HOTP) IntegrityCheck() (string, uint64) {
	h := hmac.New(sha1.New, otp.Key)
	counter := make([]byte, 8)
	h.Write(counter)
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64()), otp.Counter()
}

// Scan takes a code input (i.e. from the user), and scans ahead
// within a certain window of counter values. This can be used in the
// case where the server's counter and the user's counter have fallen
// out of sync.
func (otp *HOTP) Scan(code string, window int) bool {
	var valid bool
	codeBytes := []byte(code)
	counter := otp.Counter()

	for i := 0; i < window; i++ {
		genCode := []byte(otp.OTP())
		if subtle.ConstantTimeCompare(codeBytes, genCode) == 1 {
			valid = true
			break
		}
	}
	if !valid {
		otp.setCounter(counter)
	}
	return valid
}

// Check takes an input code and verifies it against the OTP. If
// successful, the counter is incremented.
func (otp *HOTP) Check(code string) bool {
	codeBytes := []byte(code)
	genCode := []byte(otp.OTP())
	if subtle.ConstantTimeCompare(codeBytes, genCode) != 1 {
		otp.setCounter(otp.Counter() - 1)
		return false
	}
	return true
}

// Marshal serialises an HOTP key value as a DER-encoded byte slice.
func Marshal(otp *HOTP) ([]byte, error) {
	var asnHOTP struct {
		Key     []byte
		Counter *big.Int
		Digits  int
	}
	asnHOTP.Key = otp.Key[:]
	asnHOTP.Counter = new(big.Int).SetUint64(otp.Counter())
	asnHOTP.Digits = otp.Digits
	return asn1.Marshal(asnHOTP)
}

// Unmarshal parses a DER-encoded serialised HOTP key value.
func Unmarshal(in []byte) (otp *HOTP, err error) {
	var asnHOTP struct {
		Key     []byte
		Counter *big.Int
		Digits  int
	}
	_, err = asn1.Unmarshal(in, &asnHOTP)
	if err != nil {
		return
	}

	otp = &HOTP{
		Key:    asnHOTP.Key[:],
		Digits: asnHOTP.Digits,
	}
	otp.setCounter(asnHOTP.Counter.Uint64())
	return
}
<commit_msg>replace code.google.com import<commit_after>package hotp

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha1"
	"crypto/subtle"
	"encoding/asn1"
	"encoding/base32"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math"
	"math/big"
	"net/url"
	"strconv"

	"github.com/mdp/rsc/qr"
)

// RFC 4226 specifies the counter as being 8 bytes.
const ctrSize = 8

// ErrInvalidHOTPURL is returned via FromURL; it indicates a malformed
// HOTP otpauth URL.
var ErrInvalidHOTPURL = errors.New("hotp: invalid HOTP url")

// PRNG is the source of random data; this is used by GenerateHOTP
// and should be a cryptographically-secure PRNG.
var PRNG = rand.Reader

// HOTP represents a new key value for generating one-time passwords;
// it contains the key used to construct one-time passwords and the
// counter state used in the OTP generation. Digits contains the
// number of digits that generated OTPs should output. Key is a
// cryptographic secret, and should be treated as such.
type HOTP struct {
	Key     []byte
	counter *[ctrSize]byte
	Digits  int
}

// Counter returns the HOTP's 8-byte counter as an unsigned 64-bit
// integer.
func (otp HOTP) Counter() uint64 {
	counter := binary.BigEndian.Uint64(otp.counter[:])
	return counter
}

// Increment will increment an HOTP source's counter. This is useful
// for providers like the Google Authenticator app, which immediately
// increments the counter and uses the 0 counter value as an integrity
// check.
func (otp HOTP) Increment() {
	for i := ctrSize - 1; i >= 0; i-- {
		if otp.counter[i]++; otp.counter[i] != 0 {
			return
		}
	}
}

// OTP generates a new one-time password.
func (otp HOTP) OTP() string {
	h := hmac.New(sha1.New, otp.Key)
	h.Write(otp.counter[:])
	otp.Increment()
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64())
}

func (otp *HOTP) setCounter(counter uint64) bool {
	if otp.counter == nil {
		otp.counter = new([ctrSize]byte)
	}
	binary.BigEndian.PutUint64(otp.counter[:], counter)
	return true
}

// NewHOTP intialises a new HOTP instance with the key and counter
// values. No check is done on the digits, but typical values are 6
// and 8.
func NewHOTP(key []byte, counter uint64, digits int) *HOTP {
	otp := &HOTP{
		Key:    key,
		Digits: digits,
	}
	otp.counter = new([ctrSize]byte)
	binary.BigEndian.PutUint64(otp.counter[:], counter)

	return otp
}

// URL returns a suitable URL, such as for the Google Authenticator
// app. The label is used by these apps to identify the service to
// which this OTP belongs. The digits value is ignored by the Google
// authenticator app, and is therefore elided in the resulting URL.
func (otp *HOTP) URL(label string) string {
	secret := base32.StdEncoding.EncodeToString(otp.Key)
	u := url.URL{}
	v := url.Values{}
	u.Scheme = "otpauth"
	u.Host = "hotp"
	u.Path = label
	v.Add("secret", secret)
	v.Add("counter", fmt.Sprintf("%d", otp.Counter()))
	u.RawQuery = v.Encode()
	return u.String()
}

// QR generates a byte slice containing the a QR code encoded as a
// PNG with level Q error correction.
func (otp *HOTP) QR(label string) ([]byte, error) {
	u := otp.URL(label)
	code, err := qr.Encode(u, qr.Q)
	if err != nil {
		return nil, err
	}
	return code.PNG(), nil
}

// truncate contains the DT function from the RFC; this is used to
// deterministically select a sequence of 4 bytes from the HMAC
// counter hash.
func truncate(in []byte) int64 {
	offset := int(in[len(in)-1] & 0xF)
	p := in[offset : offset+4]
	var binCode int32
	binCode = int32((p[0] & 0x7f)) << 24
	binCode += int32((p[1] & 0xff)) << 16
	binCode += int32((p[2] & 0xff)) << 8
	binCode += int32((p[3] & 0xff))
	return int64(binCode) & 0x7FFFFFFF
}

// FromURL parses a new HOTP from a URL string. It returns the OTP,
// the label associated with the OTP, and any errors that occurred.
func FromURL(urlString string) (*HOTP, string, error) {
	u, err := url.Parse(urlString)
	if err != nil {
		return nil, "", err
	}

	if u.Scheme != "otpauth" {
		return nil, "", ErrInvalidHOTPURL
	} else if u.Host != "hotp" {
		return nil, "", ErrInvalidHOTPURL
	}

	v := u.Query()
	if len(v) == 0 {
		return nil, "", ErrInvalidHOTPURL
	}
	if v.Get("secret") == "" {
		return nil, "", ErrInvalidHOTPURL
	} else if algo := v.Get("algorithm"); algo != "" && algo != "SHA1" {
		return nil, "", ErrInvalidHOTPURL
	}

	var identity string
	if len(u.Path) > 1 {
		identity = u.Path[1:]
	}

	var counter uint64
	if ctr := v.Get("counter"); ctr != "" {
		counter, err = strconv.ParseUint(ctr, 10, 64)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	secret, err := base32.StdEncoding.DecodeString(v.Get("secret"))
	if err != nil {
		return nil, "", ErrInvalidHOTPURL
	}

	var digits int64 = 6
	if v.Get("digits") != "" {
		digits, err = strconv.ParseInt(v.Get("digits"), 10, 8)
		if err != nil {
			return nil, "", ErrInvalidHOTPURL
		}
	}

	otp := NewHOTP(secret, counter, int(digits))
	return otp, identity, nil
}

// GenerateHOTP will generate a randomised HOTP source; if the
// randCounter parameter is true, the counter will be randomised.
func GenerateHOTP(digits int, randCounter bool) (*HOTP, error) {
	key := make([]byte, sha1.Size)
	_, err := io.ReadFull(PRNG, key)
	if err != nil {
		return nil, err
	}

	var counter uint64
	if randCounter {
		ctr, err := rand.Int(PRNG, big.NewInt(int64(math.MaxInt64)))
		if err != nil {
			return nil, err
		}
		counter = ctr.Uint64()
	}

	return NewHOTP(key, counter, digits), nil
}

// YubiKey reads an OATH-HOTP string as returned by a YubiKey, and
// returns three values. The first value contains the actual OTP, the
// second value contains the YubiKey's token identifier, and the final
// value indicates whether the input string was a valid YubiKey
// OTP. This does not check whether the code is correct or not, it
// only ensures that it is well-formed output from a token and
// splits the output into the code and the public identity.
func (otp *HOTP) YubiKey(in string) (string, string, bool) {
	if len(in) < otp.Digits {
		return "", "", false
	}

	otpStart := len(in) - otp.Digits
	code := in[otpStart:]
	pubid := in[:otpStart]
	return code, pubid, true
}

// IntegrityCheck returns two values, the base OTP and the current
// counter. This is used, for example, with the Google Authenticator
// app's "Check key value" function and can be used to verify that
// the application and the provider are in sync.
func (otp *HOTP) IntegrityCheck() (string, uint64) {
	h := hmac.New(sha1.New, otp.Key)
	counter := make([]byte, 8)
	h.Write(counter)
	hash := h.Sum(nil)
	result := truncate(hash)

	mod := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(otp.Digits)), nil)
	mod = mod.Mod(big.NewInt(result), mod)
	fmtStr := fmt.Sprintf("%%0%dd", otp.Digits)
	return fmt.Sprintf(fmtStr, mod.Uint64()), otp.Counter()
}

// Scan takes a code input (i.e. from the user), and scans ahead
// within a certain window of counter values. This can be used in the
// case where the server's counter and the user's counter have fallen
// out of sync.
func (otp *HOTP) Scan(code string, window int) bool {
	var valid bool
	codeBytes := []byte(code)
	counter := otp.Counter()

	for i := 0; i < window; i++ {
		genCode := []byte(otp.OTP())
		if subtle.ConstantTimeCompare(codeBytes, genCode) == 1 {
			valid = true
			break
		}
	}
	if !valid {
		otp.setCounter(counter)
	}
	return valid
}

// Check takes an input code and verifies it against the OTP. If
// successful, the counter is incremented.
func (otp *HOTP) Check(code string) bool {
	codeBytes := []byte(code)
	genCode := []byte(otp.OTP())
	if subtle.ConstantTimeCompare(codeBytes, genCode) != 1 {
		otp.setCounter(otp.Counter() - 1)
		return false
	}
	return true
}

// Marshal serialises an HOTP key value as a DER-encoded byte slice.
func Marshal(otp *HOTP) ([]byte, error) {
	var asnHOTP struct {
		Key     []byte
		Counter *big.Int
		Digits  int
	}
	asnHOTP.Key = otp.Key[:]
	asnHOTP.Counter = new(big.Int).SetUint64(otp.Counter())
	asnHOTP.Digits = otp.Digits
	return asn1.Marshal(asnHOTP)
}

// Unmarshal parses a DER-encoded serialised HOTP key value.
func Unmarshal(in []byte) (otp *HOTP, err error) {
	var asnHOTP struct {
		Key     []byte
		Counter *big.Int
		Digits  int
	}
	_, err = asn1.Unmarshal(in, &asnHOTP)
	if err != nil {
		return
	}

	otp = &HOTP{
		Key:    asnHOTP.Key[:],
		Digits: asnHOTP.Digits,
	}
	otp.setCounter(asnHOTP.Counter.Uint64())
	return
}
<|endoftext|>