<commit_before>package mqplan

import (
	"fmt"
	"meqa/mqswag"
	"meqa/mqutil"
	"sort"
	"strings"

	"github.com/go-openapi/spec"
)

func createInitTask() *Test {
	initTask := &Test{}
	initTask.Name = MeqaInit
	return initTask
}

func addInitTestSuite(testPlan *TestPlan) {
	testSuite := CreateTestSuite(MeqaInit, nil, testPlan)
	testSuite.comment = "The meqa_init section initializes parameters (e.g. pathParams) that are applied to all suites"
	testSuite.Tests = append(testSuite.Tests, createInitTask())
	testPlan.Add(testSuite)
}

// Given a path name, retrieve the last entry that is not a path param.
func GetLastPathElement(name string) string {
	nameArray := strings.Split(name, "/")
	for i := len(nameArray) - 1; i >= 0; i-- {
		if len(nameArray[i]) > 0 && nameArray[i][0] != '{' {
			return nameArray[i]
		}
	}
	return ""
}

// If the last entry on path is a parameter, return it. Otherwise return ""
func GetLastPathParam(name string) string {
	nameArray := strings.Split(name, "/")
	var last string
	for i := len(nameArray) - 1; i >= 0; i-- {
		if len(nameArray[i]) > 0 {
			last = nameArray[i]
			break
		}
	}
	if last[0] == '{' && last[len(last)-1] == '}' {
		return last[1 : len(last)-1]
	}
	return ""
}

func CreateTestFromOp(opNode *mqswag.DAGNode, testId int) *Test {
	op := opNode.Data.((*spec.Operation))
	t := &Test{}
	t.Path = opNode.GetName()
	t.Method = opNode.GetMethod()
	opId := op.ID
	if len(opId) == 0 {
		opId = GetLastPathElement(t.Path)
	}
	t.Name = fmt.Sprintf("%s_%s_%d", t.Method, opId, testId)

	return t
}

func OperationMatches(node *mqswag.DAGNode, method string) bool {
	op, ok := node.Data.(*spec.Operation)
	if ok && op != nil {
		tag := mqswag.GetMeqaTag(op.Description)
		if (tag != nil && tag.Operation == method) || ((tag == nil || len(tag.Operation) == 0) && node.GetMethod() == method) {
			return true
		}
	}
	return false
}

// GenerateTestsForObject for the obj that we traversed to from create. Add the test suites
// generated to plan.
func GenerateTestsForObject(create *mqswag.DAGNode, obj *mqswag.DAGNode, plan *TestPlan) error {
	if obj.GetType() != mqswag.TypeDef {
		return nil
	}
	if create.GetType() != mqswag.TypeOp {
		return nil
	}
	createPath := create.GetName()
	objName := obj.GetName()

	// A loop where we go through all the child operations
	testId := 1
	testSuite := CreateTestSuite(fmt.Sprintf("%s -- %s -- all", createPath, objName), nil, plan)
	testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(create, testId))
	for _, child := range obj.Children {
		if child.GetType() != mqswag.TypeOp {
			continue
		}
		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(child, testId))
		if OperationMatches(child, mqswag.MethodDelete) {
			testId++
			testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(create, testId))
		}
	}
	if len(testSuite.Tests) > 0 {
		plan.Add(testSuite)
	}

	// a loop where we pick random operations and pair it with the create operation.
	// This would generate a few objects.
	/* disable random stuff during development
	testId = 0
	testSuite = &TestSuite{nil, fmt.Sprintf("%s -- %s -- random", createPath, objName)}
	for i := 0; i < 2*len(obj.Children); i++ {
		j := rand.Intn(len(obj.Children))
		child := obj.Children[j]
		if child.GetType() != mqswag.TypeOp {
			mqutil.Logger.Printf("unexpected: (%s) has a child (%s) that's not an operation", obj.Name, child.Name)
			continue
		}
		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(create, testId))
		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(child, testId))
	}
	if len(testSuite.Tests) > 0 {
		plan.Add(testSuite)
	}
	*/

	return nil
}

func GenerateTestPlan(swagger *mqswag.Swagger, dag *mqswag.DAG) (*TestPlan, error) {
	testPlan := &TestPlan{}
	testPlan.Init(swagger, nil)
	testPlan.comment = `
This test plan has test suites that are about objects. Each test suite create an object,
then exercise REST calls that use that object as an input.
`
	addInitTestSuite(testPlan)

	genFunc := func(previous *mqswag.DAGNode, current *mqswag.DAGNode) error {
		if current.GetType() != mqswag.TypeOp {
			return nil
		}

		// Exercise the function by itself.
		/*
			testSuite := CreateTestSuite(current.GetName()+" "+current.GetMethod(), nil, testPlan)
			testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(current, 1))
			testPlan.Add(testSuite)
		*/

		// When iterating by weight previous is always nil.
		for _, c := range current.Children {
			err := GenerateTestsForObject(current, c, testPlan)
			if err != nil {
				return err
			}
		}

		return nil
	}
	err := dag.IterateByWeight(genFunc)
	if err != nil {
		return nil, err
	}
	return testPlan, nil
}

// All the operations have the same path. We generate one test suite, with the
// tests of ascending weight and priority among the operations
func GeneratePathTestSuite(operations mqswag.NodeList, plan *TestPlan) {
	if len(operations) == 0 {
		return
	}

	pathName := operations[0].GetName()
	sort.Sort(operations)
	testId := 0
	testSuite := CreateTestSuite(fmt.Sprintf("%s", pathName), nil, plan)
	for _, o := range operations {
		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(o, testId))

		if OperationMatches(o, mqswag.MethodDelete) {
			lastTest := testSuite.Tests[len(testSuite.Tests)-1]
			// Find an operation that takes the same last path param.
			lastParam := GetLastPathParam(o.GetName())
			if len(lastParam) > 0 {
				for _, repeatOp := range operations {
					if lastParam == GetLastPathParam(repeatOp.GetName()) && !OperationMatches(repeatOp, mqswag.MethodDelete) {
						testId++
						repeatTest := CreateTestFromOp(repeatOp, testId)
						repeatTest.PathParams = make(map[string]interface{})
						repeatTest.Expect = make(map[string]interface{})
						repeatTest.PathParams[lastParam] = fmt.Sprintf("{{%s.pathParams.%s}}", lastTest.Name, lastParam)
						repeatTest.Expect["status"] = "fail"
						testSuite.Tests = append(testSuite.Tests, repeatTest)
						break
					}
				}
			}
		}
	}
	if len(testSuite.Tests) > 0 {
		plan.Add(testSuite)
	}
}

type PathWeight struct {
	path   string
	weight int
}

type PathWeightList []PathWeight

func (n PathWeightList) Len() int {
	return len(n)
}

func (n PathWeightList) Swap(i, j int) {
	n[i], n[j] = n[j], n[i]
}

func (n PathWeightList) Less(i, j int) bool {
	return n[i].weight < n[j].weight || (n[i].weight == n[j].weight && n[i].path < n[j].path)
}

// Go through all the paths in swagger, and generate the tests for all the operations under
// the path.
func GeneratePathTestPlan(swagger *mqswag.Swagger, dag *mqswag.DAG) (*TestPlan, error) {
	testPlan := &TestPlan{}
	testPlan.Init(swagger, nil)
	testPlan.comment = `
In this test plan, the test suites are the REST paths, and the tests are the different
operations under the path. The tests under the same suite will share each others'
parameters by default.
	`
	addInitTestSuite(testPlan)

	pathMap := make(map[string]mqswag.NodeList)
	pathWeight := make(map[string]int)

	addFunc := func(previous *mqswag.DAGNode, current *mqswag.DAGNode) error {
		if current.GetType() != mqswag.TypeOp {
			return nil
		}
		name := current.GetName()

		// if the last path element is a {..} path param we remove it. Also remove the ending "/"
		// because it has no effect.
		nameArray := strings.Split(name, "/")
		if len(nameArray) > 0 && len(nameArray[len(nameArray)-1]) == 0 {
			nameArray = nameArray[:len(nameArray)-1]
		}
		if len(nameArray) > 0 {
			if last := nameArray[len(nameArray)-1]; len(last) > 0 && last[0] == '{' && last[len(last)-1] == '}' {
				nameArray = nameArray[:len(nameArray)-1]
			}
		}
		name = strings.Join(nameArray, "/")

		pathMap[name] = append(pathMap[name], current)

		currentWeight := current.Weight*mqswag.DAGDepth + current.Priority
		if pathWeight[name] <= currentWeight {
			pathWeight[name] = currentWeight
		}

		return nil
	}

	dag.IterateByWeight(addFunc)

	var pathWeightList PathWeightList
	// Sort the path by weight
	for k, v := range pathWeight {
		p := PathWeight{k, v}
		pathWeightList = append(pathWeightList, p)
	}
	sort.Sort(pathWeightList)

	for _, p := range pathWeightList {
		GeneratePathTestSuite(pathMap[p.path], testPlan)
	}
	return testPlan, nil
}

// Go through all the paths in swagger, and generate the tests for all the operations under
// the path.
func GenerateSimpleTestPlan(swagger *mqswag.Swagger, dag *mqswag.DAG) (*TestPlan, error) {
	testPlan := &TestPlan{}
	testPlan.Init(swagger, nil)
	addInitTestSuite(testPlan)

	testId := 0
	testSuite := CreateTestSuite(fmt.Sprintf("simple test suite"), nil, testPlan)
	testSuite.comment = "The meqa_init task within a test suite initializes parameters that are applied to all tests within this suite"
	testSuite.Tests = append(testSuite.Tests, createInitTask())
	addFunc := func(previous *mqswag.DAGNode, current *mqswag.DAGNode) error {
		if testId >= 10 {
			return mqutil.NewError(mqutil.ErrOK, "done")
		}

		if current.GetType() != mqswag.TypeOp {
			return nil
		}

		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(current, testId))

		return nil
	}

	dag.IterateByWeight(addFunc)
	testPlan.Add(testSuite)
	testPlan.comment = "\nThis is a simple and short test plan. We just sampled up to 10 REST calls into one test suite.\n"

	return testPlan, nil
}
<commit_msg>Pass objectID to subsequent reqs once created<commit_after>package mqplan

import (
	"fmt"
	"meqa/mqswag"
	"meqa/mqutil"
	"sort"
	"strings"

	"github.com/go-openapi/spec"
)

func createInitTask() *Test {
	initTask := &Test{}
	initTask.Name = MeqaInit
	return initTask
}

func addInitTestSuite(testPlan *TestPlan) {
	testSuite := CreateTestSuite(MeqaInit, nil, testPlan)
	testSuite.comment = "The meqa_init section initializes parameters (e.g. pathParams) that are applied to all suites"
	testSuite.Tests = append(testSuite.Tests, createInitTask())
	testPlan.Add(testSuite)
}

// Given a path name, retrieve the last entry that is not a path param.
func GetLastPathElement(name string) string {
	nameArray := strings.Split(name, "/")
	for i := len(nameArray) - 1; i >= 0; i-- {
		if len(nameArray[i]) > 0 && nameArray[i][0] != '{' {
			return nameArray[i]
		}
	}
	return ""
}

// If the last entry on path is a parameter, return it. Otherwise return ""
func GetLastPathParam(name string) string {
	nameArray := strings.Split(name, "/")
	var last string
	for i := len(nameArray) - 1; i >= 0; i-- {
		if len(nameArray[i]) > 0 {
			last = nameArray[i]
			break
		}
	}
	if last[0] == '{' && last[len(last)-1] == '}' {
		return last[1 : len(last)-1]
	}
	return ""
}

func CreateTestFromOp(opNode *mqswag.DAGNode, testId int) *Test {
	op := opNode.Data.((*spec.Operation))
	t := &Test{}
	t.Path = opNode.GetName()
	t.Method = opNode.GetMethod()
	opId := op.ID
	if len(opId) == 0 {
		opId = GetLastPathElement(t.Path)
	}
	t.Name = fmt.Sprintf("%s_%s_%d", t.Method, opId, testId)

	return t
}

func OperationMatches(node *mqswag.DAGNode, method string) bool {
	op, ok := node.Data.(*spec.Operation)
	if ok && op != nil {
		tag := mqswag.GetMeqaTag(op.Description)
		if (tag != nil && tag.Operation == method) || ((tag == nil || len(tag.Operation) == 0) && node.GetMethod() == method) {
			return true
		}
	}
	return false
}

// GenerateTestsForObject for the obj that we traversed to from create. Add the test suites
// generated to plan.
func GenerateTestsForObject(create *mqswag.DAGNode, obj *mqswag.DAGNode, plan *TestPlan) error {
	if obj.GetType() != mqswag.TypeDef {
		return nil
	}
	if create.GetType() != mqswag.TypeOp {
		return nil
	}
	createPath := create.GetName()
	objName := obj.GetName()

	// A loop where we go through all the child operations
	testId := 1
	testSuite := CreateTestSuite(fmt.Sprintf("%s -- %s -- all", createPath, objName), nil, plan)
	testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(create, testId))
	for _, child := range obj.Children {
		if child.GetType() != mqswag.TypeOp {
			continue
		}
		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(child, testId))
		if OperationMatches(child, mqswag.MethodDelete) {
			testId++
			testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(create, testId))
		}
	}
	if len(testSuite.Tests) > 0 {
		plan.Add(testSuite)
	}

	// a loop where we pick random operations and pair it with the create operation.
	// This would generate a few objects.
	/* disable random stuff during development
	testId = 0
	testSuite = &TestSuite{nil, fmt.Sprintf("%s -- %s -- random", createPath, objName)}
	for i := 0; i < 2*len(obj.Children); i++ {
		j := rand.Intn(len(obj.Children))
		child := obj.Children[j]
		if child.GetType() != mqswag.TypeOp {
			mqutil.Logger.Printf("unexpected: (%s) has a child (%s) that's not an operation", obj.Name, child.Name)
			continue
		}
		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(create, testId))
		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(child, testId))
	}
	if len(testSuite.Tests) > 0 {
		plan.Add(testSuite)
	}
	*/

	return nil
}

func GenerateTestPlan(swagger *mqswag.Swagger, dag *mqswag.DAG) (*TestPlan, error) {
	testPlan := &TestPlan{}
	testPlan.Init(swagger, nil)
	testPlan.comment = `
This test plan has test suites that are about objects. Each test suite create an object,
then exercise REST calls that use that object as an input.
`
	addInitTestSuite(testPlan)

	genFunc := func(previous *mqswag.DAGNode, current *mqswag.DAGNode) error {
		if current.GetType() != mqswag.TypeOp {
			return nil
		}

		// Exercise the function by itself.
		/*
			testSuite := CreateTestSuite(current.GetName()+" "+current.GetMethod(), nil, testPlan)
			testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(current, 1))
			testPlan.Add(testSuite)
		*/

		// When iterating by weight previous is always nil.
		for _, c := range current.Children {
			err := GenerateTestsForObject(current, c, testPlan)
			if err != nil {
				return err
			}
		}

		return nil
	}
	err := dag.IterateByWeight(genFunc)
	if err != nil {
		return nil, err
	}
	return testPlan, nil
}

// All the operations have the same path. We generate one test suite, with the
// tests of ascending weight and priority among the operations
func GeneratePathTestSuite(operations mqswag.NodeList, plan *TestPlan) {
	if len(operations) == 0 {
		return
	}

	pathName := operations[0].GetName()
	sort.Sort(operations)
	testId := 0
	testSuite := CreateTestSuite(fmt.Sprintf("%s", pathName), nil, plan)
	createTest := &Test{}
	idTag := "id"
	for _, o := range operations {
		testId++
		currentTest := CreateTestFromOp(o, testId)
		testSuite.Tests = append(testSuite.Tests, currentTest)
		if OperationMatches(o, mqswag.MethodPost) {
			createTest = currentTest
		} else if strings.Contains(o.GetName(), idTag) {
			currentTest.PathParams = make(map[string]interface{})
			currentTest.PathParams[idTag] = fmt.Sprintf("{{%s.outputs.%s}}", createTest.Name, idTag)
		}
		if OperationMatches(o, mqswag.MethodDelete) {
			lastTest := testSuite.Tests[len(testSuite.Tests)-1]
			// Find an operation that takes the same last path param.
			lastParam := GetLastPathParam(o.GetName())
			if len(lastParam) > 0 {
				for _, repeatOp := range operations {
					if lastParam == GetLastPathParam(repeatOp.GetName()) && !OperationMatches(repeatOp, mqswag.MethodDelete) {
						testId++
						repeatTest := CreateTestFromOp(repeatOp, testId)
						repeatTest.PathParams = make(map[string]interface{})
						repeatTest.Expect = make(map[string]interface{})
						repeatTest.PathParams[lastParam] = fmt.Sprintf("{{%s.pathParams.%s}}", lastTest.Name, lastParam)
						repeatTest.Expect["status"] = "fail"
						testSuite.Tests = append(testSuite.Tests, repeatTest)
						break
					}
				}
			}
		}
	}
	if len(testSuite.Tests) > 0 {
		plan.Add(testSuite)
	}
}

type PathWeight struct {
	path   string
	weight int
}

type PathWeightList []PathWeight

func (n PathWeightList) Len() int {
	return len(n)
}

func (n PathWeightList) Swap(i, j int) {
	n[i], n[j] = n[j], n[i]
}

func (n PathWeightList) Less(i, j int) bool {
	return n[i].weight < n[j].weight || (n[i].weight == n[j].weight && n[i].path < n[j].path)
}

// Go through all the paths in swagger, and generate the tests for all the operations under
// the path.
func GeneratePathTestPlan(swagger *mqswag.Swagger, dag *mqswag.DAG) (*TestPlan, error) {
	testPlan := &TestPlan{}
	testPlan.Init(swagger, nil)
	testPlan.comment = `
In this test plan, the test suites are the REST paths, and the tests are the different
operations under the path. The tests under the same suite will share each others'
parameters by default.
	`
	addInitTestSuite(testPlan)

	pathMap := make(map[string]mqswag.NodeList)
	pathWeight := make(map[string]int)

	addFunc := func(previous *mqswag.DAGNode, current *mqswag.DAGNode) error {
		if current.GetType() != mqswag.TypeOp {
			return nil
		}
		name := current.GetName()

		// if the last path element is a {..} path param we remove it. Also remove the ending "/"
		// because it has no effect.
		nameArray := strings.Split(name, "/")
		if len(nameArray) > 0 && len(nameArray[len(nameArray)-1]) == 0 {
			nameArray = nameArray[:len(nameArray)-1]
		}
		if len(nameArray) > 0 {
			if last := nameArray[len(nameArray)-1]; len(last) > 0 && last[0] == '{' && last[len(last)-1] == '}' {
				nameArray = nameArray[:len(nameArray)-1]
			}
		}
		name = strings.Join(nameArray, "/")

		pathMap[name] = append(pathMap[name], current)

		currentWeight := current.Weight*mqswag.DAGDepth + current.Priority
		if pathWeight[name] <= currentWeight {
			pathWeight[name] = currentWeight
		}

		return nil
	}

	dag.IterateByWeight(addFunc)

	var pathWeightList PathWeightList
	// Sort the path by weight
	for k, v := range pathWeight {
		p := PathWeight{k, v}
		pathWeightList = append(pathWeightList, p)
	}
	sort.Sort(pathWeightList)

	for _, p := range pathWeightList {
		GeneratePathTestSuite(pathMap[p.path], testPlan)
	}
	return testPlan, nil
}

// Go through all the paths in swagger, and generate the tests for all the operations under
// the path.
func GenerateSimpleTestPlan(swagger *mqswag.Swagger, dag *mqswag.DAG) (*TestPlan, error) {
	testPlan := &TestPlan{}
	testPlan.Init(swagger, nil)
	addInitTestSuite(testPlan)

	testId := 0
	testSuite := CreateTestSuite(fmt.Sprintf("simple test suite"), nil, testPlan)
	testSuite.comment = "The meqa_init task within a test suite initializes parameters that are applied to all tests within this suite"
	testSuite.Tests = append(testSuite.Tests, createInitTask())
	addFunc := func(previous *mqswag.DAGNode, current *mqswag.DAGNode) error {
		if testId >= 10 {
			return mqutil.NewError(mqutil.ErrOK, "done")
		}

		if current.GetType() != mqswag.TypeOp {
			return nil
		}

		testId++
		testSuite.Tests = append(testSuite.Tests, CreateTestFromOp(current, testId))

		return nil
	}

	dag.IterateByWeight(addFunc)
	testPlan.Add(testSuite)
	testPlan.comment = "\nThis is a simple and short test plan. We just sampled up to 10 REST calls into one test suite.\n"

	return testPlan, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net"
	"github.com/boltdb/bolt"
	//"encoding/json"
	"seqdb"
	"log"
	"os"
	"os/signal"
	"syscall"
)

const (
	SEQDB_VERSION = "0.0.2"
	CONN_HOST = "localhost"
	CONN_PORT = "3333"
	CONN_TYPE = "tcp"
)

type Message struct {
	BucketName string
	SequenceName string
	Value uint64
}

func main() {
	fmt.Println("SeqDB v.", SEQDB_VERSION)
	fmt.Println("Written by Daniel Fekete <daniel.fekete@voov.hu>")

	signalCh := make(chan os.Signal)

	signal.Notify(signalCh, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM)

	// TODO: make the database file variable based on nodename
	db, err := bolt.Open("seq.db", 0600, nil)

	if err != nil {
		log.Fatalf("Cannot connect to database: %v\r\n", err)
	}

	seqdb.SetDB(db)

	laddr, err := net.ResolveTCPAddr(CONN_TYPE, fmt.Sprintf("%s:%s", CONN_HOST, CONN_PORT))

	l, err := net.ListenTCP(CONN_TYPE, laddr)

	service := seqdb.NewService()
	go service.Serve(l)

	<-signalCh

	db.Close()
	service.Stop()

	log.Println("Terminating SeqDB")
}
<commit_msg>host and db file flags<commit_after>package main

import (
	"fmt"
	"net"
	"github.com/boltdb/bolt"
	//"encoding/json"
	"seqdb"
	"log"
	"os"
	"os/signal"
	"syscall"
	"flag"
)

/*
	Setup command line flags
 */

var (
	host = flag.String("h", "localhost:3333", "Which IP and port should the server listen on [ipaddr:port]")
	dbFile = flag.String("d", "seq.db", "The path to the SeqDB database file")
)

const (
	SEQDB_VERSION = "0.1.0"
	CONN_TYPE = "tcp"
)

type Message struct {
	BucketName string
	SequenceName string
	Value uint64
}

func main() {
	fmt.Println("SeqDB v.", SEQDB_VERSION)
	fmt.Println("Written by Daniel Fekete <daniel.fekete@voov.hu>")
	flag.Parse()

	signalCh := make(chan os.Signal)

	signal.Notify(signalCh, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM)

	// TODO: make the database file variable based on nodename
	db, err := bolt.Open(*dbFile, 0600, nil)

	if err != nil {
		log.Fatalf("Cannot connect to database: %v\r\n", err)
	}
	log.Printf("Database file %s opened\r\n", *dbFile)
	seqdb.SetDB(db)

	laddr, err := net.ResolveTCPAddr(CONN_TYPE, *host)

	l, err := net.ListenTCP(CONN_TYPE, laddr)
	log.Printf("Listening on %v\r\n", laddr)

	service := seqdb.NewService()
	go service.Serve(l)

	<-signalCh

	db.Close()
	service.Stop()

	log.Println("Terminating SeqDB")
}
<|endoftext|>


<commit_before>package ai

import (
	"flag"
	"testing"
	"time"

	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

var size = flag.Int("size", 5, "board size to benchmark")
var depth = flag.Int("depth", 4, "minimax search depth")

func BenchmarkMinimax(b *testing.B) {
	var cfg = tak.Config{Size: *size}
	p := tak.New(cfg)
	p, _ = p.Move(&tak.Move{X: 0, Y: 0, Type: tak.PlaceFlat})
	p, _ = p.Move(&tak.Move{X: *size - 1, Y: *size - 1, Type: tak.PlaceFlat})
	ai := NewMinimax(MinimaxConfig{Size: *size, Depth: *depth})

	b.ReportAllocs()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		var e error
		m := ai.GetMove(p, time.Minute)
		p, e = p.Move(&m)
		if e != nil {
			b.Fatal("bad move", e)
		}
		if over, _ := p.GameOver(); over {
			p = tak.New(cfg)
			p, _ = p.Move(&tak.Move{X: 0, Y: 0, Type: tak.PlaceFlat})
			p, _ = p.Move(&tak.Move{X: *size - 1, Y: *size - 1, Type: tak.PlaceFlat})
		}
	}
}

func TestRegression(t *testing.T) {
	game, err := ptn.ParseTPS(
		`2,x4/x2,2,x2/x,2,2,x2/x2,12,2,1/1,1,21,2,1 1 9`,
	)
	if err != nil {
		panic(err)
	}
	ai := NewMinimax(MinimaxConfig{Size: game.Size(), Depth: 3})
	m := ai.GetMove(game, time.Minute)
	_, e := game.Move(&m)
	if e != nil {
		t.Fatalf("ai returned illegal move: %s: %s", ptn.FormatMove(&m), e)
	}
}
<commit_msg>Add a fixed seed to benchmark<commit_after>package ai

import (
	"flag"
	"testing"
	"time"

	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

var size = flag.Int("size", 5, "board size to benchmark")
var depth = flag.Int("depth", 4, "minimax search depth")
var seed = flag.Int64("seed", 0, "random seed")

func BenchmarkMinimax(b *testing.B) {
	var cfg = tak.Config{Size: *size}
	p := tak.New(cfg)
	p, _ = p.Move(&tak.Move{X: 0, Y: 0, Type: tak.PlaceFlat})
	p, _ = p.Move(&tak.Move{X: *size - 1, Y: *size - 1, Type: tak.PlaceFlat})
	ai := NewMinimax(MinimaxConfig{
		Size:  *size,
		Depth: *depth,
		Seed:  *seed,
	})

	b.ReportAllocs()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		var e error
		m := ai.GetMove(p, time.Minute)
		p, e = p.Move(&m)
		if e != nil {
			b.Fatal("bad move", e)
		}
		if over, _ := p.GameOver(); over {
			p = tak.New(cfg)
			p, _ = p.Move(&tak.Move{X: 0, Y: 0, Type: tak.PlaceFlat})
			p, _ = p.Move(&tak.Move{X: *size - 1, Y: *size - 1, Type: tak.PlaceFlat})
		}
	}
}

func TestRegression(t *testing.T) {
	game, err := ptn.ParseTPS(
		`2,x4/x2,2,x2/x,2,2,x2/x2,12,2,1/1,1,21,2,1 1 9`,
	)
	if err != nil {
		panic(err)
	}
	ai := NewMinimax(MinimaxConfig{Size: game.Size(), Depth: 3})
	m := ai.GetMove(game, time.Minute)
	_, e := game.Move(&m)
	if e != nil {
		t.Fatalf("ai returned illegal move: %s: %s", ptn.FormatMove(&m), e)
	}
}
<|endoftext|>


<commit_before>// problem60.go
//
// The primes 3, 7, 109, and 673, are quite remarkable. By taking any two
// primes and concatenating them in any order the result will always be prime.
// For example, taking 7 and 109, both 7109 and 1097 are prime. The sum of
// these four primes, 792, represents the lowest sum for a set of four primes
// with this property.
//
// Find the lowest sum for a set of five primes for which any two primes
// concatenate to produce another prime.

package main

import (
	"euler/tools"
	"fmt"
	"strconv"
	"sync"
)

type cache struct {
	sync.Mutex
	m map[string]bool
}

func newCache() *cache {
	return &cache{sync.Mutex{}, make(map[string]bool)}
}

func (c *cache) concatsToPrime(x, y int) bool {
	key := strconv.Itoa(x) + strconv.Itoa(y)
	c.Lock()
	defer c.Unlock()
	// Try to find the answer in the cache.
	if val, ok := c.m[key]; ok {
		return val
	}
	// Otherwise find it manually and add to the cache.
	xstr, ystr := strconv.Itoa(x), strconv.Itoa(y)
	a, _ := strconv.Atoi(xstr + ystr)
	b, _ := strconv.Atoi(ystr + xstr)
	val := tools.IsPrime(a) && tools.IsPrime(b)
	c.m[key] = val
	return val
}

func (c *cache) allConcatToPrime(n node) bool {
	for _, x := range n {
		for _, y := range n {
			if x != y && x < y {
				if c.concatsToPrime(x, y) == false {
					return false
				}
			}
		}
	}
	return true
}

// A node is a candidate solution to the problem.
type node []int

func (n node) max() int {
	if len(n) == 0 {
		panic("max: node has 0 elements")
	}
	x := n[0]
	for _, y := range n[1:] {
		if y > x {
			x = y
		}
	}
	return x
}

func (n node) sum() int {
	res := 0
	for _, x := range n {
		res += x
	}
	return res
}

// Stack is a simple stack implementation.
type stack []node

// Pop a node off the stack. Panics if s is empty.
func (s *stack) pop() node {
	x := (*s)[len(*s)-1]
	*s = (*s)[:len(*s)-1]

	// Shrink the underlying array if the slice length <= 1/4 its capacity.
	if len(*s) <= cap(*s)/4 {
		*s = append([]node{}, *s...)
	}
	return x
}

// Push a node onto the stack.
func (s *stack) push(x node) {
	*s = append(*s, x)
}

// We are going to use a concurrent depth-first search with a worker goroutine
// pool of 4. Each goroutine will search for a solution from a different
// starting prime. As soon as a solution is found, we return from the function.
// Otherwise, we wait for all starting primes to be checked, and return an
// error.
func problem60() (int, error) {
	// It's not clear how many primes to search through. Experimentation
	// suggests a limit of 9000 produces the correct answer: 26033. Note
	// our algorithm does not guarantee the solution is the smallest
	// possible, but as a matter of fact, it is. We could verify our
	// answer by raising the limit to 26033, searching exhaustively, and
	// observing that no smaller solutions are found.
	limit := 9000
	var primes []int
	for i := 0; i < limit; i++ {
		if tools.IsPrime(i) {
			primes = append(primes, i)
		}
	}

	c := newCache()
	ans := make(chan int)   // Used to send the answer.
	done := make(chan bool) // Used to signal that all worker goroutines are done.
	pchan := make(chan int) // Used to send worker goroutines a starting prime to search.
	var wg sync.WaitGroup

	go func() {
		for _, p := range primes {
			pchan <- p
		}
		close(pchan)
		wg.Wait()    // Wait for all workers to complete their search
		done <- true // before sending completion signal.
	}()

	// Woker goroutine pool of 4.
	for i := 0; i < 4; i++ {
		go func() {
			wg.Add(1)
			for {
				p, ok := <-pchan
				if !ok {
					wg.Done()
					return
				}

				// Perform depth-first search starting at the given prime.
				var frontier stack
				frontier.push(node{p})

				for len(frontier) != 0 {
					n := frontier.pop()
					if len(n) == 5 {
						ans <- n.sum()
					}
					for _, prime := range primes {
						child := append(append(*new(node), n...), prime)
						if prime > n.max() && c.allConcatToPrime(child) {
							frontier.push(child)
						}
					}
				}
			}
		}()
	}

	for {
		select {
		case x := <-ans:
			return x, nil
		case <-done:
			return -1, fmt.Errorf("problem60: no solution found with limit %v", limit)
		}
	}
}

func main() {
	ans, err := problem60()
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println(ans)
	}
}
<commit_msg>Move WaitGroup.Add call into main goroutine<commit_after>// problem60.go
//
// The primes 3, 7, 109, and 673, are quite remarkable. By taking any two
// primes and concatenating them in any order the result will always be prime.
// For example, taking 7 and 109, both 7109 and 1097 are prime. The sum of
// these four primes, 792, represents the lowest sum for a set of four primes
// with this property.
//
// Find the lowest sum for a set of five primes for which any two primes
// concatenate to produce another prime.

package main

import (
	"euler/tools"
	"fmt"
	"strconv"
	"sync"
)

type cache struct {
	sync.Mutex
	m map[string]bool
}

func newCache() *cache {
	return &cache{sync.Mutex{}, make(map[string]bool)}
}

func (c *cache) concatsToPrime(x, y int) bool {
	key := strconv.Itoa(x) + strconv.Itoa(y)
	c.Lock()
	defer c.Unlock()
	// Try to find the answer in the cache.
	if val, ok := c.m[key]; ok {
		return val
	}
	// Otherwise find it manually and add to the cache.
	xstr, ystr := strconv.Itoa(x), strconv.Itoa(y)
	a, _ := strconv.Atoi(xstr + ystr)
	b, _ := strconv.Atoi(ystr + xstr)
	val := tools.IsPrime(a) && tools.IsPrime(b)
	c.m[key] = val
	return val
}

func (c *cache) allConcatToPrime(n node) bool {
	for _, x := range n {
		for _, y := range n {
			if x != y && x < y {
				if c.concatsToPrime(x, y) == false {
					return false
				}
			}
		}
	}
	return true
}

// A node is a candidate solution to the problem.
type node []int

func (n node) max() int {
	if len(n) == 0 {
		panic("max: node has 0 elements")
	}
	x := n[0]
	for _, y := range n[1:] {
		if y > x {
			x = y
		}
	}
	return x
}

func (n node) sum() int {
	res := 0
	for _, x := range n {
		res += x
	}
	return res
}

// Stack is a simple stack implementation.
type stack []node

// Pop a node off the stack. Panics if s is empty.
func (s *stack) pop() node {
	x := (*s)[len(*s)-1]
	*s = (*s)[:len(*s)-1]

	// Shrink the underlying array if the slice length <= 1/4 its capacity.
	if len(*s) <= cap(*s)/4 {
		*s = append([]node{}, *s...)
	}
	return x
}

// Push a node onto the stack.
func (s *stack) push(x node) {
	*s = append(*s, x)
}

// We are going to use a concurrent depth-first search with a worker goroutine
// pool of 4. Each goroutine will search for a solution from a different
// starting prime. As soon as a solution is found, we return from the function.
// Otherwise, we wait for all starting primes to be checked, and return an
// error.
func problem60() (int, error) {
	// It's not clear how many primes to search through. Experimentation
	// suggests a limit of 9000 produces the correct answer: 26033. Note
	// our algorithm does not guarantee the solution is the smallest
	// possible, but as a matter of fact, it is. We could verify our
	// answer by raising the limit to 26033, searching exhaustively, and
	// observing that no smaller solutions are found.
	limit := 9000
	var primes []int
	for i := 0; i < limit; i++ {
		if tools.IsPrime(i) {
			primes = append(primes, i)
		}
	}

	c := newCache()
	ans := make(chan int)   // Used to send the answer.
	done := make(chan bool) // Used to signal that all worker goroutines are done.
	pchan := make(chan int) // Used to send worker goroutines a starting prime to search.
	var wg sync.WaitGroup

	// Woker goroutine pool of 4.
	for i := 0; i < 4; i++ {
		wg.Add(1)
		go func() {
			for {
				p, ok := <-pchan
				if !ok {
					wg.Done()
					return
				}

				// Perform depth-first search starting at the given prime.
				var frontier stack
				frontier.push(node{p})

				for len(frontier) != 0 {
					n := frontier.pop()
					if len(n) == 5 {
						ans <- n.sum()
					}
					for _, prime := range primes {
						child := append(append(*new(node), n...), prime)
						if prime > n.max() && c.allConcatToPrime(child) {
							frontier.push(child)
						}
					}
				}
			}
		}()
	}

	go func() {
		for _, p := range primes {
			pchan <- p
		}
		close(pchan)
		wg.Wait()    // Wait for all workers to complete their search
		done <- true // before sending completion signal.
	}()

	for {
		select {
		case x := <-ans:
			return x, nil
		case <-done:
			return -1, fmt.Errorf("problem60: no solution found with limit %v", limit)
		}
	}
}

func main() {
	ans, err := problem60()
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println(ans)
	}
}
<|endoftext|>


<commit_before>package main

// TODO: Add line numbers to the error messages and make them parseable by editors and IDEs

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"strings"
	"time"
)

// Global variables
var (
	// 32-bit (i686), 64-bit (x86_64) or 16-bit (i386)
	platformBits = 32

	// Is this a bootable kernel? (declared with "bootable" at the top)
	bootable_kernel = false

	// OS X or Linux
	osx = false
)

func main() {
	name := "Battlestar"
	version := "0.4"
	log.Println(name + " compiler")
	log.Println("Version " + version)
	log.Println("Alexander Rødseth")
	log.Println("2014")
	log.Println("MIT licensed")

	ps := NewProgramState()

	// TODO: Add an option for not adding an exit function
	// TODO: Automatically discover 32-bit/64-bit and Linux/OS X
	// TODO: ARM support

	// Check for -bits=32 or -bits=64 (default)
	platformBitsArg := flag.Int("bits", 64, "Output 32-bit or 64-bit x86 assembly")
	// Check for -osx=true or -osx=false (default)
	osxArg := flag.Bool("osx", false, "On OS X?")
	// Assembly output file
	asmfileArg := flag.String("o", "", "Assembly output file")
	// C output file
	cfileArg := flag.String("oc", "", "C output file")
	// Input file
	btsfileArg := flag.String("f", "", "BTS source file")
	// Is it not a standalone program, but a component? (just the .o file is needed)
	componentArg := flag.Bool("c", false, "Component, not a standalone program")

	flag.Parse()

	platformBits = *platformBitsArg
	osx = *osxArg
	asmfile := *asmfileArg
	cfile := *cfileArg
	btsfile := *btsfileArg
	component := *componentArg

	if flag.Arg(0) != "" {
		btsfile = flag.Arg(0)
	}

	if btsfile == "" {
		log.Fatalln("Abort: a source filename is needed. Provide one with -f or as the first argument.")
	}

	if asmfile == "" {
		asmfile = btsfile + ".asm"
	}

	if cfile == "" {
		cfile = btsfile + ".c"
	}

	// TODO: Consider adding an option for "start" as well, or a custom
	// start symbol

	if osx {
		linker_start_function = "_main"
	} else {
		linker_start_function = "_start"
	}

	// Assembly file contents
	asmdata := ""

	// C file contents
	cdata := ""

	// Read code from stdin and output 32-bit or 64-bit assembly code
	bytes, err := ioutil.ReadFile(btsfile)
	if err == nil {
		if len(strings.TrimSpace(string(bytes))) == 0 {
			// Empty program
			log.Fatalln("Error: Empty program")
		}

		t := time.Now()
		asmdata += fmt.Sprintf("; Generated with %s %s, at %s\n\n", name, version, t.String()[:16])

		// If "bootable" is the first token
		bootable := false
		if temptokens := tokenize(string(bytes), " "); (len(temptokens) > 2) && (temptokens[0].t == KEYWORD) && (temptokens[0].value == "bootable") && (temptokens[1].t == SEP) {
			bootable = true
			asmdata += fmt.Sprintf("bits %d\n", platformBits)
		} else {
			// Header for regular programs
			asmdata += fmt.Sprintf("bits %d\n", platformBits)
		}

		// Check if platformBits is valid
		if !hasi([]int{16, 32, 64}, platformBits) {
			log.Fatalln("Error: Unsupported bit size:", platformBits)
		}

		init_interrupt_parameter_registers(platformBits)

		btsCode := addExternMainIfMissing(string(bytes))
		tokens := addExitTokenIfMissing(tokenize(btsCode, " "))
		log.Println("--- Done tokenizing ---")
		constants, asmcode := TokensToAssembly(tokens, true, false, ps)
		if constants != "" {
			asmdata += "section .data\n"
			asmdata += constants + "\n"
		}
		if platformBits == 16 {
			asmdata += "org 0x100\n"
		}
		if !bootable {
			asmdata += "\n"
			asmdata += "section .text\n"
		}
		if platformBits == 16 {
			// If there are defined functions, jump over the definitions and start at
			// the main/_start function. If there is a main function, jump to the
			// linker start function. If not, just start at the top.
			// TODO: This is a quick fix. Don't depend on the comment, find a better way.
			if strings.Count(asmcode, "; name of the function") > 1 {
				if strings.Contains(asmcode, "\nmain:") {
					asmdata += "jmp " + linker_start_function + "\n"
				}
			}
		}
		if asmcode != "" {
			if component {
				asmdata += asmcode + "\n"
			} else {
				asmdata += addStartingPointIfMissing(asmcode, ps) + "\n"
			}
			if bootable {
				reg := "esp"
				if platformBits == 64 {
					reg = "rsp"
				}
				asmdata = strings.Replace(asmdata, "; starting point of the program\n", "; starting point of the program\n\tmov "+reg+", stack_top\t; set the "+reg+" register to the top of the stack (special case for bootable kernels)\n", 1)
			}
		}
		ccode := ExtractInlineC(strings.TrimSpace(string(bytes)), true)
		if ccode != "" {
			cdata += fmt.Sprintf("// Generated with %s %s, at %s\n\n", name, version, t.String()[:16])
			cdata += ccode
		}
	}

	log.Println("--- Finalizing ---")

	if asmdata != "" {
		err = ioutil.WriteFile(asmfile, []byte(asmdata), 0644)
		if err != nil {
			log.Fatalln("Error: Unable to write to", asmfile)
		}
		log.Printf("Wrote %s (%d bytes)\n", asmfile, len(asmdata))
	}

	if cdata != "" {
		err = ioutil.WriteFile(cfile, []byte(cdata), 0644)
		if err != nil {
			log.Fatalln("Error: Unable to write to", cfile)
		}
		log.Printf("Wrote %s (%d bytes)\n", cfile, len(cdata))
	}

	log.Println("Done.")
}
<commit_msg>Update version info<commit_after>package main

// TODO: Add line numbers to the error messages and make them parseable by editors and IDEs

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"strings"
	"time"
)

// Global variables
var (
	// 32-bit (i686), 64-bit (x86_64) or 16-bit (i386)
	platformBits = 32

	// Is this a bootable kernel? (declared with "bootable" at the top)
	bootable_kernel = false

	// OS X or Linux
	osx = false
)

func main() {
	name := "Battlestar"
	version := "0.4"
	log.Println(name + " compiler\nVersion " + version + "\nAlexander F Rødseth\n2017\nMIT licensed")

	ps := NewProgramState()

	// TODO: Add an option for not adding an exit function
	// TODO: Automatically discover 32-bit/64-bit and Linux/OS X
	// TODO: ARM support

	// Check for -bits=32 or -bits=64 (default)
	platformBitsArg := flag.Int("bits", 64, "Output 32-bit or 64-bit x86 assembly")
	// Check for -osx=true or -osx=false (default)
	osxArg := flag.Bool("osx", false, "On OS X?")
	// Assembly output file
	asmfileArg := flag.String("o", "", "Assembly output file")
	// C output file
	cfileArg := flag.String("oc", "", "C output file")
	// Input file
	btsfileArg := flag.String("f", "", "BTS source file")
	// Is it not a standalone program, but a component? (just the .o file is needed)
	componentArg := flag.Bool("c", false, "Component, not a standalone program")

	flag.Parse()

	platformBits = *platformBitsArg
	osx = *osxArg
	asmfile := *asmfileArg
	cfile := *cfileArg
	btsfile := *btsfileArg
	component := *componentArg

	if flag.Arg(0) != "" {
		btsfile = flag.Arg(0)
	}

	if btsfile == "" {
		log.Fatalln("Abort: a source filename is needed. Provide one with -f or as the first argument.")
	}

	if asmfile == "" {
		asmfile = btsfile + ".asm"
	}

	if cfile == "" {
		cfile = btsfile + ".c"
	}

	// TODO: Consider adding an option for "start" as well, or a custom
	// start symbol

	if osx {
		linker_start_function = "_main"
	} else {
		linker_start_function = "_start"
	}

	// Assembly file contents
	asmdata := ""

	// C file contents
	cdata := ""

	// Read code from stdin and output 32-bit or 64-bit assembly code
	bytes, err := ioutil.ReadFile(btsfile)
	if err == nil {
		if len(strings.TrimSpace(string(bytes))) == 0 {
			// Empty program
			log.Fatalln("Error: Empty program")
		}

		t := time.Now()
		asmdata += fmt.Sprintf("; Generated with %s %s, at %s\n\n", name, version, t.String()[:16])

		// If "bootable" is the first token
		bootable := false
		if temptokens := tokenize(string(bytes), " "); (len(temptokens) > 2) && (temptokens[0].t == KEYWORD) && (temptokens[0].value == "bootable") && (temptokens[1].t == SEP) {
			bootable = true
			asmdata += fmt.Sprintf("bits %d\n", platformBits)
		} else {
			// Header for regular programs
			asmdata += fmt.Sprintf("bits %d\n", platformBits)
		}

		// Check if platformBits is valid
		if !hasi([]int{16, 32, 64}, platformBits) {
			log.Fatalln("Error: Unsupported bit size:", platformBits)
		}

		init_interrupt_parameter_registers(platformBits)

		btsCode := addExternMainIfMissing(string(bytes))
		tokens := addExitTokenIfMissing(tokenize(btsCode, " "))
		log.Println("--- Done tokenizing ---")
		constants, asmcode := TokensToAssembly(tokens, true, false, ps)
		if constants != "" {
			asmdata += "section .data\n"
			asmdata += constants + "\n"
		}
		if platformBits == 16 {
			asmdata += "org 0x100\n"
		}
		if !bootable {
			asmdata += "\n"
			asmdata += "section .text\n"
		}
		if platformBits == 16 {
			// If there are defined functions, jump over the definitions and start at
			// the main/_start function. If there is a main function, jump to the
			// linker start function. If not, just start at the top.
			// TODO: This is a quick fix. Don't depend on the comment, find a better way.
			if strings.Count(asmcode, "; name of the function") > 1 {
				if strings.Contains(asmcode, "\nmain:") {
					asmdata += "jmp " + linker_start_function + "\n"
				}
			}
		}
		if asmcode != "" {
			if component {
				asmdata += asmcode + "\n"
			} else {
				asmdata += addStartingPointIfMissing(asmcode, ps) + "\n"
			}
			if bootable {
				reg := "esp"
				if platformBits == 64 {
					reg = "rsp"
				}
				asmdata = strings.Replace(asmdata, "; starting point of the program\n", "; starting point of the program\n\tmov "+reg+", stack_top\t; set the "+reg+" register to the top of the stack (special case for bootable kernels)\n", 1)
			}
		}
		ccode := ExtractInlineC(strings.TrimSpace(string(bytes)), true)
		if ccode != "" {
			cdata += fmt.Sprintf("// Generated with %s %s, at %s\n\n", name, version, t.String()[:16])
			cdata += ccode
		}
	}

	log.Println("--- Finalizing ---")

	if asmdata != "" {
		err = ioutil.WriteFile(asmfile, []byte(asmdata), 0644)
		if err != nil {
			log.Fatalln("Error: Unable to write to", asmfile)
		}
		log.Printf("Wrote %s (%d bytes)\n", asmfile, len(asmdata))
	}

	if cdata != "" {
		err = ioutil.WriteFile(cfile, []byte(cdata), 0644)
		if err != nil {
			log.Fatalln("Error: Unable to write to", cfile)
		}
		log.Printf("Wrote %s (%d bytes)\n", cfile, len(cdata))
	}

	log.Println("Done.")
}
<|endoftext|>


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

package as3

import (
	"fmt"

	. "github.com/F5Networks/k8s-bigip-ctlr/pkg/resource"
	log "github.com/F5Networks/k8s-bigip-ctlr/pkg/vlogger"
)

func (am *AS3Manager) prepareAS3ResourceConfig() as3ADC {
	adc := am.generateAS3ResourceDeclaration()
	// Support `Controls` class for TEEMs in user-defined AS3 configMap.
	controlObj := make(as3Control)
	controlObj.initDefault(am.userAgent)
	adc["controls"] = controlObj
	return adc
}

func (am *AS3Manager) generateAS3ResourceDeclaration() as3ADC {
	// Create Shared as3Application object for Routes
	adc := as3ADC{}
	adc.initDefault(DEFAULT_PARTITION)
	sharedApp := adc.getAS3SharedApp(DEFAULT_PARTITION)

	// Process CIS Resources to create AS3 Resources
	am.processResourcesForAS3(sharedApp)

	// Process CustomProfiles
	am.processCustomProfilesForAS3(sharedApp)

	// Process RouteProfiles
	am.processProfilesForAS3(sharedApp)

	// For Ingress process SecretName
	// Process IRules
	am.processIRulesForAS3(sharedApp)

	// Process DataGroup to be consumed by IRule
	am.processDataGroupForAS3(sharedApp)

	// Process F5 Resources
	am.processF5ResourcesForAS3(sharedApp)

	return adc
}

func (am *AS3Manager) processProfilesForAS3(sharedApp as3Application) {
	// Processes RouteProfs to create AS3 Declaration for Route annotations
	// Override/Set ServerTLS/ClientTLS in AS3 Service as annotation takes higher priority
	for svcName, cfg := range am.Resources.RsMap {
		if svc, ok := sharedApp[as3FormattedString(svcName, cfg.MetaData.ResourceType)].(*as3Service); ok {
			switch cfg.MetaData.ResourceType {
			case ResourceTypeRoute:
				processRouteTLSProfilesForAS3(&cfg.MetaData, svc)
			case ResourceTypeIngress:
				processIngressTLSProfilesForAS3(&cfg.Virtual, svc)
			default:
				log.Warningf("Unsupported resource type: %v", cfg.MetaData.ResourceType)
			}
		}
	}
}

func processIngressTLSProfilesForAS3(virtual *Virtual, svc *as3Service) {
	// lets discard BIGIP profile creation when there exists a custom profile.
	var serverTLS []as3ResourcePointer
	for _, profile := range virtual.Profiles {
		if profile.Partition == "Common" {
			switch profile.Context {
			case CustomProfileClient:
				// Incoming traffic (clientssl) from a web client will be handled by ServerTLS in AS3
				rsPointer := as3ResourcePointer{
					BigIP: fmt.Sprintf("/%v/%v", profile.Partition, profile.Name),
				}
				serverTLS = append(serverTLS, rsPointer)
				svc.ServerTLS = serverTLS
				updateVirtualToHTTPS(svc)
			case CustomProfileServer:
				// Outgoing traffic (serverssl) to BackEnd Servers from BigIP will be handled by ClientTLS in AS3
				svc.ClientTLS = &as3ResourcePointer{
					BigIP: fmt.Sprintf("/%v/%v", profile.Partition, profile.Name),
				}
				updateVirtualToHTTPS(svc)
			}
		}
	}
}

func processRouteTLSProfilesForAS3(metadata *MetaData, svc *as3Service) {
	var serverTLS []as3ResourcePointer
	existingProfile := map[string]struct{}{}
	// handle duplicate BIGIP pointers
	for key, val := range metadata.RouteProfs {
		if _, ok := existingProfile[val]; ok {
			continue
		}
		existingProfile[val] = struct{}{}
		switch key.Context {
		case CustomProfileClient:
			// Incoming traffic (clientssl) from a web client will be handled by ServerTLS in AS3
			rsPointer := as3ResourcePointer{BigIP: val}
			serverTLS = append(serverTLS, rsPointer)
			svc.ServerTLS = serverTLS
			updateVirtualToHTTPS(svc)
		case CustomProfileServer:
			// Outgoing traffic (serverssl) to BackEnd Servers from BigIP will be handled by ClientTLS in AS3
			svc.ClientTLS = &as3ResourcePointer{
				BigIP: val,
			}
			updateVirtualToHTTPS(svc)
		}
	}
}

// processF5ResourcesForAS3 does the following steps to implement WAF
// * Add WAF policy action to the corresponding rules
// * Add a default WAF disable Rule to corresponding policy
// * Add WAF disable action to all rules that do not handle WAF
func (am *AS3Manager) processF5ResourcesForAS3(sharedApp as3Application) {

	// Identify rules that do not handle waf and add waf disable action to that rule
	addWAFDisableAction := func(ep *as3EndpointPolicy) {
		enabled := false
		wafDisableAction := &as3Action{
			Type:    "waf",
			Enabled: &enabled,
		}

		for _, rule := range ep.Rules {
			isWAFRule := false
			for _, action := range rule.Actions {
				if action.Type == "waf" {
					isWAFRule = true
					break
				}
			}
			// BigIP requires a default WAF disable rule doesn't require WAF
			if !isWAFRule {
				rule.Actions = append(rule.Actions, wafDisableAction)
			}
		}
	}

	var isSecureWAF, isInsecureWAF bool
	var secureEP, insecureEP *as3EndpointPolicy

	secureEP, _ = sharedApp["openshift_secure_routes"].(*as3EndpointPolicy)
	insecureEP, _ = sharedApp["openshift_insecure_routes"].(*as3EndpointPolicy)

	// Update Rules with WAF action
	for _, resGroup := range am.IntF5Res {
		for rec, res := range resGroup {
			switch res.Virtual {
			case HTTPS:
				if secureEP != nil {
					isSecureWAF = true
					updatePolicyWithWAF(secureEP, rec, res)
				}
			case HTTPANDS:
				if secureEP != nil {
					isSecureWAF = true
					updatePolicyWithWAF(secureEP, rec, res)
				}
				fallthrough
			case HTTP:
				if insecureEP != nil {
					isInsecureWAF = true
					updatePolicyWithWAF(insecureEP, rec, res)
				}
			}
		}
	}

	enabled := false
	wafDisableAction := &as3Action{
		Type:    "waf",
		Enabled: &enabled,
	}

	wafDropAction := &as3Action{
		Type:  "drop",
		Event: "request",
	}

	wafDisableRule := &as3Rule{
		Name:    "openshift_route_waf_disable",
		Actions: []*as3Action{wafDropAction, wafDisableAction},
	}

	// Add a default WAF disable action to all non-WAF rules
	// BigIP requires a default WAF disable rule doesn't require WAF
	if isSecureWAF && secureEP != nil {
		secureEP.Rules = append(secureEP.Rules, wafDisableRule)
		addWAFDisableAction(secureEP)
	}

	if isInsecureWAF && insecureEP != nil {
		insecureEP.Rules = append(insecureEP.Rules, wafDisableRule)
		addWAFDisableAction(insecureEP)
	}
}
<commit_msg>Bug fix for continous posting<commit_after>/*-
 * Copyright (c) 2016-2019, F5 Networks, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package as3

import (
	"fmt"
	"sort"

	. "github.com/F5Networks/k8s-bigip-ctlr/pkg/resource"
	log "github.com/F5Networks/k8s-bigip-ctlr/pkg/vlogger"
)

func (am *AS3Manager) prepareAS3ResourceConfig() as3ADC {
	adc := am.generateAS3ResourceDeclaration()
	// Support `Controls` class for TEEMs in user-defined AS3 configMap.
	controlObj := make(as3Control)
	controlObj.initDefault(am.userAgent)
	adc["controls"] = controlObj
	return adc
}

func (am *AS3Manager) generateAS3ResourceDeclaration() as3ADC {
	// Create Shared as3Application object for Routes
	adc := as3ADC{}
	adc.initDefault(DEFAULT_PARTITION)
	sharedApp := adc.getAS3SharedApp(DEFAULT_PARTITION)

	// Process CIS Resources to create AS3 Resources
	am.processResourcesForAS3(sharedApp)

	// Process CustomProfiles
	am.processCustomProfilesForAS3(sharedApp)

	// Process RouteProfiles
	am.processProfilesForAS3(sharedApp)

	// For Ingress process SecretName
	// Process IRules
	am.processIRulesForAS3(sharedApp)

	// Process DataGroup to be consumed by IRule
	am.processDataGroupForAS3(sharedApp)

	// Process F5 Resources
	am.processF5ResourcesForAS3(sharedApp)

	return adc
}

func (am *AS3Manager) processProfilesForAS3(sharedApp as3Application) {
	// Processes RouteProfs to create AS3 Declaration for Route annotations
	// Override/Set ServerTLS/ClientTLS in AS3 Service as annotation takes higher priority
	for svcName, cfg := range am.Resources.RsMap {
		if svc, ok := sharedApp[as3FormattedString(svcName, cfg.MetaData.ResourceType)].(*as3Service); ok {
			switch cfg.MetaData.ResourceType {
			case ResourceTypeRoute:
				processRouteTLSProfilesForAS3(&cfg.MetaData, svc)
			case ResourceTypeIngress:
				processIngressTLSProfilesForAS3(&cfg.Virtual, svc)
			default:
				log.Warningf("Unsupported resource type: %v", cfg.MetaData.ResourceType)
			}
		}
	}
}

func processIngressTLSProfilesForAS3(virtual *Virtual, svc *as3Service) {
	// lets discard BIGIP profile creation when there exists a custom profile.
	var serverTLS []as3ResourcePointer
	for _, profile := range virtual.Profiles {
		if profile.Partition == "Common" {
			switch profile.Context {
			case CustomProfileClient:
				// Incoming traffic (clientssl) from a web client will be handled by ServerTLS in AS3
				rsPointer := as3ResourcePointer{
					BigIP: fmt.Sprintf("/%v/%v", profile.Partition, profile.Name),
				}
				serverTLS = append(serverTLS, rsPointer)
				// svc.ServerTLS = serverTLS

				updateVirtualToHTTPS(svc)
			case CustomProfileServer:
				// Outgoing traffic (serverssl) to BackEnd Servers from BigIP will be handled by ClientTLS in AS3
				svc.ClientTLS = &as3ResourcePointer{
					BigIP: fmt.Sprintf("/%v/%v", profile.Partition, profile.Name),
				}
				updateVirtualToHTTPS(svc)
			}
		}
	}
	if len(serverTLS) > 0 {
		sortedServerTLS := getSortedServerTLS(serverTLS)
		svc.ServerTLS = sortedServerTLS
	}
}

func processRouteTLSProfilesForAS3(metadata *MetaData, svc *as3Service) {
	var serverTLS []as3ResourcePointer
	existingProfile := map[string]struct{}{}
	// handle duplicate BIGIP pointers
	for key, val := range metadata.RouteProfs {
		if _, ok := existingProfile[val]; ok {
			continue
		}
		existingProfile[val] = struct{}{}
		switch key.Context {
		case CustomProfileClient:
			// Incoming traffic (clientssl) from a web client will be handled by ServerTLS in AS3
			rsPointer := as3ResourcePointer{BigIP: val}
			serverTLS = append(serverTLS, rsPointer)
			// svc.ServerTLS = serverTLS
			updateVirtualToHTTPS(svc)
		case CustomProfileServer:
			// Outgoing traffic (serverssl) to BackEnd Servers from BigIP will be handled by ClientTLS in AS3
			svc.ClientTLS = &as3ResourcePointer{
				BigIP: val,
			}
			updateVirtualToHTTPS(svc)
		}
	}
	if len(serverTLS) > 0 {
		sortedServerTLS := getSortedServerTLS(serverTLS)
		svc.ServerTLS = sortedServerTLS
	}

}

//Get sorted ServerTLS by value
func getSortedServerTLS(serverTLS []as3ResourcePointer) []as3ResourcePointer {
	if len(serverTLS) == 1 {
		return serverTLS
	}
	var ref []string
	for _, val := range serverTLS {
		ref = append(ref, val.BigIP)
	}
	sort.Strings(ref)
	var sortedServerTLS []as3ResourcePointer
	for _, val := range ref {
		rsPointer := as3ResourcePointer{
			BigIP: val,
		}
		sortedServerTLS = append(sortedServerTLS, rsPointer)
	}
	return sortedServerTLS
}

// processF5ResourcesForAS3 does the following steps to implement WAF
// * Add WAF policy action to the corresponding rules
// * Add a default WAF disable Rule to corresponding policy
// * Add WAF disable action to all rules that do not handle WAF
func (am *AS3Manager) processF5ResourcesForAS3(sharedApp as3Application) {

	// Identify rules that do not handle waf and add waf disable action to that rule
	addWAFDisableAction := func(ep *as3EndpointPolicy) {
		enabled := false
		wafDisableAction := &as3Action{
			Type:    "waf",
			Enabled: &enabled,
		}

		for _, rule := range ep.Rules {
			isWAFRule := false
			for _, action := range rule.Actions {
				if action.Type == "waf" {
					isWAFRule = true
					break
				}
			}
			// BigIP requires a default WAF disable rule doesn't require WAF
			if !isWAFRule {
				rule.Actions = append(rule.Actions, wafDisableAction)
			}
		}
	}

	var isSecureWAF, isInsecureWAF bool
	var secureEP, insecureEP *as3EndpointPolicy

	secureEP, _ = sharedApp["openshift_secure_routes"].(*as3EndpointPolicy)
	insecureEP, _ = sharedApp["openshift_insecure_routes"].(*as3EndpointPolicy)

	// Update Rules with WAF action
	for _, resGroup := range am.IntF5Res {
		for rec, res := range resGroup {
			switch res.Virtual {
			case HTTPS:
				if secureEP != nil {
					isSecureWAF = true
					updatePolicyWithWAF(secureEP, rec, res)
				}
			case HTTPANDS:
				if secureEP != nil {
					isSecureWAF = true
					updatePolicyWithWAF(secureEP, rec, res)
				}
				fallthrough
			case HTTP:
				if insecureEP != nil {
					isInsecureWAF = true
					updatePolicyWithWAF(insecureEP, rec, res)
				}
			}
		}
	}

	enabled := false
	wafDisableAction := &as3Action{
		Type:    "waf",
		Enabled: &enabled,
	}

	wafDropAction := &as3Action{
		Type:  "drop",
		Event: "request",
	}

	wafDisableRule := &as3Rule{
		Name:    "openshift_route_waf_disable",
		Actions: []*as3Action{wafDropAction, wafDisableAction},
	}

	// Add a default WAF disable action to all non-WAF rules
	// BigIP requires a default WAF disable rule doesn't require WAF
	if isSecureWAF && secureEP != nil {
		secureEP.Rules = append(secureEP.Rules, wafDisableRule)
		addWAFDisableAction(secureEP)
	}

	if isInsecureWAF && insecureEP != nil {
		insecureEP.Rules = append(insecureEP.Rules, wafDisableRule)
		addWAFDisableAction(insecureEP)
	}
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"encoding/json"
	"io"
	"os"
	"strings"
	"text/template"
)

var cmdList = &Command{
	UsageLine: "list [-e] [-f format] [-json] [-tags 'tag list'] [packages]",
	Short:     "list packages",
	Long: `
List lists the packages named by the import paths, one per line.

The default output shows the package import path:

    code.google.com/p/google-api-go-client/books/v1
    code.google.com/p/goauth2/oauth
    code.google.com/p/sqlite

The -f flag specifies an alternate format for the list, using the
syntax of package template.  The default output is equivalent to -f
'{{.ImportPath}}'.  One extra template function is available, "join",
which calls strings.Join. The struct being passed to the template is:

    type Package struct {
        Dir        string // directory containing package sources
        ImportPath string // import path of package in dir
        Name       string // package name
        Doc        string // package documentation string
        Target     string // install path
        Goroot     bool   // is this package in the Go root?
        Standard   bool   // is this package part of the standard Go library?
        Stale      bool   // would 'go install' do anything for this package?
        Root       string // Go root or Go path dir containing this package

        // Source files
        GoFiles  []string       // .go source files (excluding CgoFiles, TestGoFiles, XTestGoFiles)
        CgoFiles []string       // .go sources files that import "C"
        IgnoredGoFiles []string // .go sources ignored due to build constraints
        CFiles   []string       // .c source files
        CXXFiles []string       // .cc, .cxx and .cpp source files
        HFiles   []string       // .h, .hh, .hpp and .hxx source files
        SFiles   []string       // .s source files
        SwigFiles []string      // .swig files
        SwigCXXFiles []string   // .swigcxx files
        SysoFiles []string      // .syso object files to add to archive

        // Cgo directives
        CgoCFLAGS    []string // cgo: flags for C compiler
        CgoCPPFLAGS  []string // cgo: flags for C preprocessor
        CgoCXXFLAGS  []string // cgo: flags for C++ compiler
        CgoLDFLAGS   []string // cgo: flags for linker
        CgoPkgConfig []string // cgo: pkg-config names

        // Dependency information
        Imports []string // import paths used by this package
        Deps    []string // all (recursively) imported dependencies

        // Error information
        Incomplete bool            // this package or a dependency has an error
        Error      *PackageError   // error loading package
        DepsErrors []*PackageError // errors loading dependencies

        TestGoFiles  []string // _test.go files in package
        TestImports  []string // imports from TestGoFiles
        XTestGoFiles []string // _test.go files outside package
        XTestImports []string // imports from XTestGoFiles
    }

The -json flag causes the package data to be printed in JSON format
instead of using the template format.

The -e flag changes the handling of erroneous packages, those that
cannot be found or are malformed.  By default, the list command
prints an error to standard error for each erroneous package and
omits the packages from consideration during the usual printing.
With the -e flag, the list command never prints errors to standard
error and instead processes the erroneous packages with the usual
printing.  Erroneous packages will have a non-empty ImportPath and
a non-nil Error field; other information may or may not be missing
(zeroed).

The -tags flag specifies a list of build tags, like in the 'go build'
command.

For more about specifying packages, see 'go help packages'.
	`,
}

func init() {
	cmdList.Run = runList // break init cycle
	cmdList.Flag.Var(buildCompiler{}, "compiler", "")
	cmdList.Flag.Var((*stringsFlag)(&buildContext.BuildTags), "tags", "")
}

var listE = cmdList.Flag.Bool("e", false, "")
var listFmt = cmdList.Flag.String("f", "{{.ImportPath}}", "")
var listJson = cmdList.Flag.Bool("json", false, "")
var nl = []byte{'\n'}

func runList(cmd *Command, args []string) {
	out := newTrackingWriter(os.Stdout)
	defer out.w.Flush()

	var do func(*Package)
	if *listJson {
		do = func(p *Package) {
			b, err := json.MarshalIndent(p, "", "\t")
			if err != nil {
				out.Flush()
				fatalf("%s", err)
			}
			out.Write(b)
			out.Write(nl)
		}
	} else {
		tmpl, err := template.New("main").Funcs(template.FuncMap{"join": strings.Join}).Parse(*listFmt)
		if err != nil {
			fatalf("%s", err)
		}
		do = func(p *Package) {
			if err := tmpl.Execute(out, p); err != nil {
				out.Flush()
				fatalf("%s", err)
			}
			if out.NeedNL() {
				out.Write([]byte{'\n'})
			}
		}
	}

	load := packages
	if *listE {
		load = packagesAndErrors
	}

	for _, pkg := range load(args) {
		do(pkg)
	}
}

// TrackingWriter tracks the last byte written on every write so
// we can avoid printing a newline if one was already written or
// if there is no output at all.
type TrackingWriter struct {
	w    *bufio.Writer
	last byte
}

func newTrackingWriter(w io.Writer) *TrackingWriter {
	return &TrackingWriter{
		w:    bufio.NewWriter(w),
		last: '\n',
	}
}

func (t *TrackingWriter) Write(p []byte) (n int, err error) {
	n, err = t.w.Write(p)
	if n > 0 {
		t.last = p[n-1]
	}
	return
}

func (t *TrackingWriter) Flush() {
	t.w.Flush()
}

func (t *TrackingWriter) NeedNL() bool {
	return t.last != '\n'
}
<commit_msg>cmd/go: add -race flag to 'go list'<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"encoding/json"
	"io"
	"os"
	"strings"
	"text/template"
)

var cmdList = &Command{
	UsageLine: "list [-e] [-race] [-f format] [-json] [-tags 'tag list'] [packages]",
	Short:     "list packages",
	Long: `
List lists the packages named by the import paths, one per line.

The default output shows the package import path:

    code.google.com/p/google-api-go-client/books/v1
    code.google.com/p/goauth2/oauth
    code.google.com/p/sqlite

The -f flag specifies an alternate format for the list, using the
syntax of package template.  The default output is equivalent to -f
'{{.ImportPath}}'.  One extra template function is available, "join",
which calls strings.Join. The struct being passed to the template is:

    type Package struct {
        Dir        string // directory containing package sources
        ImportPath string // import path of package in dir
        Name       string // package name
        Doc        string // package documentation string
        Target     string // install path
        Goroot     bool   // is this package in the Go root?
        Standard   bool   // is this package part of the standard Go library?
        Stale      bool   // would 'go install' do anything for this package?
        Root       string // Go root or Go path dir containing this package

        // Source files
        GoFiles  []string       // .go source files (excluding CgoFiles, TestGoFiles, XTestGoFiles)
        CgoFiles []string       // .go sources files that import "C"
        IgnoredGoFiles []string // .go sources ignored due to build constraints
        CFiles   []string       // .c source files
        CXXFiles []string       // .cc, .cxx and .cpp source files
        HFiles   []string       // .h, .hh, .hpp and .hxx source files
        SFiles   []string       // .s source files
        SwigFiles []string      // .swig files
        SwigCXXFiles []string   // .swigcxx files
        SysoFiles []string      // .syso object files to add to archive

        // Cgo directives
        CgoCFLAGS    []string // cgo: flags for C compiler
        CgoCPPFLAGS  []string // cgo: flags for C preprocessor
        CgoCXXFLAGS  []string // cgo: flags for C++ compiler
        CgoLDFLAGS   []string // cgo: flags for linker
        CgoPkgConfig []string // cgo: pkg-config names

        // Dependency information
        Imports []string // import paths used by this package
        Deps    []string // all (recursively) imported dependencies

        // Error information
        Incomplete bool            // this package or a dependency has an error
        Error      *PackageError   // error loading package
        DepsErrors []*PackageError // errors loading dependencies

        TestGoFiles  []string // _test.go files in package
        TestImports  []string // imports from TestGoFiles
        XTestGoFiles []string // _test.go files outside package
        XTestImports []string // imports from XTestGoFiles
    }

The -json flag causes the package data to be printed in JSON format
instead of using the template format.

The -e flag changes the handling of erroneous packages, those that
cannot be found or are malformed.  By default, the list command
prints an error to standard error for each erroneous package and
omits the packages from consideration during the usual printing.
With the -e flag, the list command never prints errors to standard
error and instead processes the erroneous packages with the usual
printing.  Erroneous packages will have a non-empty ImportPath and
a non-nil Error field; other information may or may not be missing
(zeroed).

The -tags flag specifies a list of build tags, like in the 'go build'
command.

The -race flag causes the package data to include the dependencies
required by the race detector.

For more about specifying packages, see 'go help packages'.
	`,
}

func init() {
	cmdList.Run = runList // break init cycle
	cmdList.Flag.Var(buildCompiler{}, "compiler", "")
	cmdList.Flag.Var((*stringsFlag)(&buildContext.BuildTags), "tags", "")
}

var listE = cmdList.Flag.Bool("e", false, "")
var listFmt = cmdList.Flag.String("f", "{{.ImportPath}}", "")
var listJson = cmdList.Flag.Bool("json", false, "")
var listRace = cmdList.Flag.Bool("race", false, "")
var nl = []byte{'\n'}

func runList(cmd *Command, args []string) {
	out := newTrackingWriter(os.Stdout)
	defer out.w.Flush()

	if *listRace {
		buildRace = true
	}

	var do func(*Package)
	if *listJson {
		do = func(p *Package) {
			b, err := json.MarshalIndent(p, "", "\t")
			if err != nil {
				out.Flush()
				fatalf("%s", err)
			}
			out.Write(b)
			out.Write(nl)
		}
	} else {
		tmpl, err := template.New("main").Funcs(template.FuncMap{"join": strings.Join}).Parse(*listFmt)
		if err != nil {
			fatalf("%s", err)
		}
		do = func(p *Package) {
			if err := tmpl.Execute(out, p); err != nil {
				out.Flush()
				fatalf("%s", err)
			}
			if out.NeedNL() {
				out.Write([]byte{'\n'})
			}
		}
	}

	load := packages
	if *listE {
		load = packagesAndErrors
	}

	for _, pkg := range load(args) {
		do(pkg)
	}
}

// TrackingWriter tracks the last byte written on every write so
// we can avoid printing a newline if one was already written or
// if there is no output at all.
type TrackingWriter struct {
	w    *bufio.Writer
	last byte
}

func newTrackingWriter(w io.Writer) *TrackingWriter {
	return &TrackingWriter{
		w:    bufio.NewWriter(w),
		last: '\n',
	}
}

func (t *TrackingWriter) Write(p []byte) (n int, err error) {
	n, err = t.w.Write(p)
	if n > 0 {
		t.last = p[n-1]
	}
	return
}

func (t *TrackingWriter) Flush() {
	t.w.Flush()
}

func (t *TrackingWriter) NeedNL() bool {
	return t.last != '\n'
}
<|endoftext|>


<commit_before>// Copyright 2014 AKUALAB INC. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
optioner is a tool to generate functional options. Intended to be
used with go generate; see the invocation in example/example.go.

To learn about functional options, see:
http://dave.cheney.net/2014/10/17/functional-options-for-friendly-apis
http://commandcenter.blogspot.com.au/2014/01/self-referential-functions-and-design.html

This code was adapted from the stringer cmd (golang.org/x/tools/cmd/stringer/).

optioner will generate a file with code of the form:

  // N sets field N in Example.
  func N(n int) Option {
	  return func(ex *Example) {
		  ex.N = n
	  }
  }

The file is created in the package where the code is generated.

optioner will create options for all the fields in the struct except those that include
the tag `opt:"-"`.

For example, given this snippet,

  package example

  //go:generate optioner -type Example
  type Example struct {
	  N      int
	  FSlice []float64 `json:"float_slice"`
	  Map    map[string]int
	  Name   string `opt:"-" json:"name"`
	  ff     func(int) int
  }

  func NewExample(name string, options ...Option) *Example {

	  // Set required values and initialize optional fields with default values.
	  ex := &Example{
		  Name:   name,
		  N:      10,
		  FSlice: make([]float64, 0, 100),
		  Map:    make(map[string]int),
	  }

	  // Set options.
	  ex.init(options...)
  }

go generate will generate option functions for fields N, FSlice, Map, and ff. Your package
users can now set options as follows:

  myFunc := func(n int) int {return 2 * n}
  ex := example.NewExample("test", example.N(22), example.Ff(myFunc))

the new struct "ex" will use default values for "FSlice" and "Map", and custom values for
"N" and "ff". Because the argument "name" is required, the field "name" is excluded using a tag.

struct fields don't need to exported, however, the corresponding option will be exported by
capitalizing the first letter.
*/
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"reflect"
	"strings"

	"golang.org/x/tools/go/types"
)

var typeNameArg = flag.String("type", "", "type name of the options struct; must be set")
var output = flag.String("output", "", "output file name; default srcdir/<type>_gen_opt.go")

// Usage is a replacement usage function for the flags package.
func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\toptioner [flags] -type T\n")
	fmt.Fprintf(os.Stderr, "Use struct tag ```opt:\"-\"``` to exclude fields\n")
	fmt.Fprintf(os.Stderr, "For more information, see:\n")
	fmt.Fprintf(os.Stderr, "\thttp://github.com/akualab/optioner\n")
	fmt.Fprintf(os.Stderr, "Flags:\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("optioner: ")
	flag.Usage = Usage
	flag.Parse()
	if len(*typeNameArg) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	// Parse the package once.
	var (
		g Generator
	)

	g.typeName = *typeNameArg
	g.options = []option{}

	g.parsePackage()

	// Print the header and package clause.
	g.Printf("// generated by optioner %s; DO NOT EDIT\n", strings.Join(os.Args[1:], " "))
	g.Printf("\n")
	g.Printf(header)
	g.Printf("package %s", g.packageName)
	g.Printf("\n")
	g.Printf("// Option type is used to pass options to %s.\n", g.typeName)
	g.Printf("type Option func(*%s)\n", g.typeName)
	g.Printf("\n")
	g.Printf("// Package author: use this method inside func New%s()\n", g.typeName)
	g.Printf("// to set optional values.\n")
	g.Printf("func (t *%s) init(options ...Option) {\n", g.typeName)
	g.Printf("for _, option := range options {\n")
	g.Printf("option(t)\n")
	g.Printf("}\n")
	g.Printf("}\n")
	g.Printf("\n")

	for _, opt := range g.options {
		tname := strings.Title(opt.name)
		g.Printf("// %s sets optional value in %s.\n", tname, g.typeName)
		g.Printf("func %s(o %s) Option {\n", tname, opt.typ)
		g.Printf("return func(t *%s) {\n", g.typeName)
		g.Printf("t.%s = o\n", opt.name)
		g.Printf("}\n")
		g.Printf("}\n")
		g.Printf("\n")
	}

	// Format the output.
	src := g.format()

	// Write to file.
	outputName := *output
	if outputName == "" {
		baseName := fmt.Sprintf("%s_gen_opt.go", g.typeName)
		outputName = strings.ToLower(baseName)
	}
	err := ioutil.WriteFile(outputName, src, 0644)
	if err != nil {
		log.Fatalf("writing output: %s", err)
	}
}

func (g *Generator) parsePackage() {

	pkg, err := build.Default.ImportDir(".", 0)
	if err != nil {
		log.Fatalf("cannot process directory: %s", err)
	}

	for _, f := range pkg.GoFiles {
		if found := g.parseFields(f); found {
			break
		}
	}
}

func (g *Generator) parseFields(fn string) bool {

	log.Println("parse file: ", fn)
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, fn, nil, 0)
	if err != nil {
		log.Fatal(err)
	}
	g.packageName = f.Name.Name
	for _, v := range f.Decls {
		if genDecl, ok := v.(*ast.GenDecl); ok {
			if genDecl.Tok != token.TYPE {
				continue
			}
			for _, s := range genDecl.Specs {
				if typeSpec, ok := s.(*ast.TypeSpec); ok {
					if typeSpec.Name.String() != g.typeName {
						continue
					}
					if structDecl, ok := typeSpec.Type.(*ast.StructType); ok {
						fields := structDecl.Fields.List
						for _, field := range fields {
							id := field.Names[0]

							// check struct tags to exclude fields
							if field.Tag != nil {
								s := strings.Replace(field.Tag.Value, "`", "", -1)
								tag := reflect.StructTag(s).Get("opt")
								if tag == "-" {
									continue
								}
							}
							typeExpr := field.Type
							typ := types.ExprString(typeExpr)
							g.options = append(g.options, option{name: id.Name, typ: typ})
							log.Printf("generating option for field \"%s\" of type \"%s\")", id.Name, typ)
						}
						return true
					} else {
						log.Fatal("target type is not a struct")
					}
				}
			}
		}
	}
	return false
}

// Generator holds the state of the analysis. Primarily used to buffer
// the output for format.Source.
type Generator struct {
	buf         bytes.Buffer // Accumulated output.
	options     []option
	typeName    string
	packageName string
}

type option struct {
	name string
	typ  string
}

func (g *Generator) Printf(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, format, args...)
}

// format returns the gofmt-ed contents of the Generator's buffer.
func (g *Generator) format() []byte {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		// Should never happen, but can arise when developing this code.
		// The user can compile the output to see the error.
		log.Printf("warning: internal error: invalid Go generated: %s", err)
		log.Printf("warning: compile the package to analyze the error")
		return g.buf.Bytes()
	}
	return src
}

const header = `
// Please report issues and submit contributions at:
// http://github.com/akualab/optioner
// optioner is a project of AKUALAB INC.

`
<commit_msg>Edit comment.<commit_after>// Copyright 2014 AKUALAB INC. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
optioner is a tool to generate functional options. Intended to be
used with go generate; see the invocation in example/example.go.

To learn about functional options, see:
http://dave.cheney.net/2014/10/17/functional-options-for-friendly-apis
http://commandcenter.blogspot.com.au/2014/01/self-referential-functions-and-design.html

This code was adapted from the stringer cmd (golang.org/x/tools/cmd/stringer/).

optioner will generate a file with code of the form:

  // N sets field N in Example.
  func N(n int) Option {
	  return func(ex *Example) {
		  ex.N = n
	  }
  }

The file is created in the package where the code is generated.

optioner will create options for all the fields in the struct except those that include
the tag `opt:"-"`.

For example, given this snippet,

  package example

  //go:generate optioner -type Example
  type Example struct {
	  N      int
	  FSlice []float64 `json:"float_slice"`
	  Map    map[string]int
	  Name   string `opt:"-" json:"name"`
	  ff     func(int) int
  }

  func NewExample(name string, options ...Option) *Example {

	  // Set required values and initialize optional fields with default values.
	  ex := &Example{
		  Name:   name,
		  N:      10,
		  FSlice: make([]float64, 0, 100),
		  Map:    make(map[string]int),
		  ff:     func(n int) int { return n },
	  }

	  // Set options.
	  ex.init(options...)
  }

go generate will generate option functions for fields N, FSlice, Map, and ff. Your package
users can now set options as follows:

  myFunc := func(n int) int {return 2 * n}
  ex := example.NewExample("test", example.N(22), example.Ff(myFunc))

the new struct "ex" will use default values for "FSlice" and "Map", and custom values for
"N" and "ff". Because the argument "name" is required, the field "name" is excluded using a tag.

struct fields don't need to exported, however, the corresponding option will be exported by
capitalizing the first letter.
*/
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/build"
	"go/format"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"reflect"
	"strings"

	"golang.org/x/tools/go/types"
)

var typeNameArg = flag.String("type", "", "type name of the options struct; must be set")
var output = flag.String("output", "", "output file name; default srcdir/<type>_gen_opt.go")

// Usage is a replacement usage function for the flags package.
func Usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\toptioner [flags] -type T\n")
	fmt.Fprintf(os.Stderr, "Use struct tag ```opt:\"-\"``` to exclude fields\n")
	fmt.Fprintf(os.Stderr, "For more information, see:\n")
	fmt.Fprintf(os.Stderr, "\thttp://github.com/akualab/optioner\n")
	fmt.Fprintf(os.Stderr, "Flags:\n")
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("optioner: ")
	flag.Usage = Usage
	flag.Parse()
	if len(*typeNameArg) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	// Parse the package once.
	var (
		g Generator
	)

	g.typeName = *typeNameArg
	g.options = []option{}

	g.parsePackage()

	// Print the header and package clause.
	g.Printf("// generated by optioner %s; DO NOT EDIT\n", strings.Join(os.Args[1:], " "))
	g.Printf("\n")
	g.Printf(header)
	g.Printf("package %s", g.packageName)
	g.Printf("\n")
	g.Printf("// Option type is used to pass options to %s.\n", g.typeName)
	g.Printf("type Option func(*%s)\n", g.typeName)
	g.Printf("\n")
	g.Printf("// Package author: use this method inside func New%s()\n", g.typeName)
	g.Printf("// to set optional values.\n")
	g.Printf("func (t *%s) init(options ...Option) {\n", g.typeName)
	g.Printf("for _, option := range options {\n")
	g.Printf("option(t)\n")
	g.Printf("}\n")
	g.Printf("}\n")
	g.Printf("\n")

	for _, opt := range g.options {
		tname := strings.Title(opt.name)
		g.Printf("// %s sets optional value in %s.\n", tname, g.typeName)
		g.Printf("func %s(o %s) Option {\n", tname, opt.typ)
		g.Printf("return func(t *%s) {\n", g.typeName)
		g.Printf("t.%s = o\n", opt.name)
		g.Printf("}\n")
		g.Printf("}\n")
		g.Printf("\n")
	}

	// Format the output.
	src := g.format()

	// Write to file.
	outputName := *output
	if outputName == "" {
		baseName := fmt.Sprintf("%s_gen_opt.go", g.typeName)
		outputName = strings.ToLower(baseName)
	}
	err := ioutil.WriteFile(outputName, src, 0644)
	if err != nil {
		log.Fatalf("writing output: %s", err)
	}
}

func (g *Generator) parsePackage() {

	pkg, err := build.Default.ImportDir(".", 0)
	if err != nil {
		log.Fatalf("cannot process directory: %s", err)
	}

	for _, f := range pkg.GoFiles {
		if found := g.parseFields(f); found {
			break
		}
	}
}

func (g *Generator) parseFields(fn string) bool {

	log.Println("parse file: ", fn)
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, fn, nil, 0)
	if err != nil {
		log.Fatal(err)
	}
	g.packageName = f.Name.Name
	for _, v := range f.Decls {
		if genDecl, ok := v.(*ast.GenDecl); ok {
			if genDecl.Tok != token.TYPE {
				continue
			}
			for _, s := range genDecl.Specs {
				if typeSpec, ok := s.(*ast.TypeSpec); ok {
					if typeSpec.Name.String() != g.typeName {
						continue
					}
					if structDecl, ok := typeSpec.Type.(*ast.StructType); ok {
						fields := structDecl.Fields.List
						for _, field := range fields {
							id := field.Names[0]

							// check struct tags to exclude fields
							if field.Tag != nil {
								s := strings.Replace(field.Tag.Value, "`", "", -1)
								tag := reflect.StructTag(s).Get("opt")
								if tag == "-" {
									continue
								}
							}
							typeExpr := field.Type
							typ := types.ExprString(typeExpr)
							g.options = append(g.options, option{name: id.Name, typ: typ})
							log.Printf("generating option for field \"%s\" of type \"%s\")", id.Name, typ)
						}
						return true
					} else {
						log.Fatal("target type is not a struct")
					}
				}
			}
		}
	}
	return false
}

// Generator holds the state of the analysis. Primarily used to buffer
// the output for format.Source.
type Generator struct {
	buf         bytes.Buffer // Accumulated output.
	options     []option
	typeName    string
	packageName string
}

type option struct {
	name string
	typ  string
}

func (g *Generator) Printf(format string, args ...interface{}) {
	fmt.Fprintf(&g.buf, format, args...)
}

// format returns the gofmt-ed contents of the Generator's buffer.
func (g *Generator) format() []byte {
	src, err := format.Source(g.buf.Bytes())
	if err != nil {
		// Should never happen, but can arise when developing this code.
		// The user can compile the output to see the error.
		log.Printf("warning: internal error: invalid Go generated: %s", err)
		log.Printf("warning: compile the package to analyze the error")
		return g.buf.Bytes()
	}
	return src
}

const header = `
// Please report issues and submit contributions at:
// http://github.com/akualab/optioner
// optioner is a project of AKUALAB INC.

`
<|endoftext|>


<commit_before>// main.go
package main

import (
	"database/sql"
	"flag"
	"log"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

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

const (
	SLEEP_TIME          = 3 * time.Second
	KEYBOARD_BUFER_SIZE = 10000
	DATABASE_NAME       = "file:gokeystat.db?cache=shared&mode=rwc"
	CAPTURE_TIME        = 5 // time in seconds between capturing keyboard to db
)

type StatForTime struct {
	time int64
	keys map[uint8]int
}

func (stat *StatForTime) Init() {
	stat.time = time.Now().Unix()
	stat.keys = make(map[uint8]int)
}

// Return map from key numbers to key names like "F1", "Tab", "d"
func GetKeymap() map[uint8]string {
	return GetKeymapFromOutput(GetKeymapOutput())
}

// Return output of utility that prints system keymap
func GetKeymapOutput() []byte {
	cmd := exec.Command("xmodmap", "-pke")
	out, err := cmd.Output()
	if err != nil {
		log.Fatal(err)
	}
	return out
}

// Return map with keymap from text
func GetKeymapFromOutput(buf []byte) map[uint8]string {
	const KEY_NUM_STRING_RE = "\\d+[ ]*=[ ]*\\S+"
	re := regexp.MustCompile(KEY_NUM_STRING_RE)
	resByte := re.FindAll(buf, -1)
	keyMap := make(map[uint8]string)
	for _, line := range resByte {
		lineSpitted := strings.Split(string(line), " ")
		if key, err := strconv.Atoi(lineSpitted[0]); err == nil {
			keyMap[uint8(key)] = lineSpitted[2]
		}
	}
	return keyMap
}

// Extract pressed keys from bufer buf
// It returns slice with key numbers in the same order
func GetKeyNumsFromOutput(buf []byte) []uint8 {
	const KEY_NUM_STRING_RE = "press[ ]+(\\d+)"
	re := regexp.MustCompile(KEY_NUM_STRING_RE)
	resByte := re.FindAll(buf, -1)
	keyNums := make([]uint8, len(resByte))
	re = regexp.MustCompile("\\d+")
	for i, line := range resByte {
		numByte := re.Find(line)
		if num, err := strconv.Atoi(string(numByte)); err == nil {
			keyNums[i] = uint8(num)
		} else {
			log.Fatal(err)
		}
	}
	return keyNums
}

func GetKeyNumsFromKeyMap(keyMap map[uint8]string) []int {
	res := make([]int, 0, len(keyMap))
	for keyNum := range keyMap {
		res = append(res, int(keyNum))
	}
	sort.Ints(res)
	return res
}

// Creates tables, inserts keymap to db
func InitDb(db *sql.DB, keyMap map[uint8]string) {
	keyNums := GetKeyNumsFromKeyMap(keyMap)

	sqlInit := `CREATE TABLE IF NOT EXISTS keylog (
        time INTEGER primary key`

	for keyNum := range keyNums {
		sqlInit += ",\n" + "KEY" + strconv.Itoa(keyNum) + " INTEGER"
	}
	sqlInit += "\n);"

	// Inserting keymap to table
	sqlInit += ` CREATE TABLE IF NOT EXISTS keymap (
        num INTEGER primary key,
		value STRING
	);`

	_, err := db.Exec(sqlInit)
	if err != nil {
		log.Fatalf("%q: %s\n", err, sqlInit)
	}

	rows, err := db.Query("SELECT COUNT(*) FROM keymap")
	if err != nil {
		log.Fatal(err)
	}

	var rowsCount int
	rows.Next()
	rows.Scan(&rowsCount)
	if rowsCount > 0 {
		// already inserted keymap
		return
	}
	rows.Close()

	tx, err := db.Begin()
	if err != nil {
		log.Fatal(err)
	}
	stmt, err := tx.Prepare("INSERT INTO keymap(num, value) VALUES(?, ?)")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	for keyNum, keyName := range keyMap {
		_, err = stmt.Exec(keyNum, keyName)
		if err != nil {
			log.Fatal(err)
		}
	}
	tx.Commit()
}

func AddStatTimeToDb(db *sql.DB, statTime StatForTime, keyMap map[uint8]string) {
	keyNums := GetKeyNumsFromKeyMap(keyMap)
	sqlStmt := "insert into keylog(time"
	for keyNum := range keyNums {
		sqlStmt += ",\n" + "KEY" + strconv.Itoa(keyNum)
	}
	sqlStmt += ") values "
	sqlStmt += "(" + strconv.FormatInt(statTime.time, 10)
	for keyNum := range keyNums {
		keyNumber, _ := statTime.keys[uint8(keyNum)]
		sqlStmt += ",\n" + strconv.Itoa(keyNumber)
	}
	sqlStmt += ")"
	tx, err := db.Begin()
	if err != nil {
		log.Fatal(err)
	}

	_, err = tx.Exec(sqlStmt)
	if err != nil {
		log.Printf("%q: %s\n", err, sqlStmt)
	}

	tx.Commit()
}

// Returns slice with StatForTime objects that
func GetStatTimesFromDb(db *sql.DB, fromTime int64, keyMap map[uint8]string) []StatForTime {
	sqlStmt := "select * from keylog where time > " + strconv.FormatInt(fromTime, 10)
	rows, err := db.Query(sqlStmt)
	if err != nil {
		log.Fatalln("Error with query", sqlStmt, " is ", err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		log.Fatalln("Failed to get columns", err)
	}

	rawResult := make([][]byte, len(cols))
	result := make([]int64, len(cols))

	dest := make([]interface{}, len(cols)) // A temporary interface{} slice
	for i, _ := range rawResult {
		dest[i] = &rawResult[i] // Put pointers to each string in the interface slice
	}

	// keyNums[i] stores i-th keynum
	keyNums := GetKeyNumsFromKeyMap(keyMap)

	// result
	res := make([]StatForTime, 0)
	for rows.Next() {
		err = rows.Scan(dest...)
		if err != nil {
			log.Fatalln("Failed to scan row", err)
		}

		for i, raw := range rawResult {
			if raw == nil {
				result[i] = 0
			} else {
				// Only numbers in db: converting it to int64
				result[i], err = strconv.ParseInt(string(raw), 10, 64)
				if err != nil {
					log.Fatalln("Error when parsing ", raw, " from db:", err)
				}
			}
		}

		var resStatTime StatForTime
		resStatTime.time = result[0]
		resStatTime.keys = make(map[uint8]int)
		for index, val := range result[1:] {
			if val == 0 {
				continue
			}
			resStatTime.keys[uint8(keyNums[index])] = int(val)
		}
		res = append(res, resStatTime)
	}
	if err = rows.Err(); err != nil {
		log.Fatalln("Error when iterating over rows", err)
	}

	return res
}

func main() {

	keyboardId := flag.Int("id", -1, "Your keyboard id")
	outputPath := flag.String("o", "", "Path to export file")
	flag.Parse()
	log.Println("keyboardId =", *keyboardId, "outputPath =", *outputPath)
	switch {
	case *keyboardId == -1 && *outputPath == "":
		flag.PrintDefaults()
		return
	case *keyboardId != -1:
		// Opening database
		db, err := sql.Open("sqlite3", DATABASE_NAME)
		if err != nil {
			log.Fatal(err)
		}
		db.SetMaxIdleConns(5)
		db.SetMaxOpenConns(5)
		defer db.Close()

		keyMap := GetKeymap()

		InitDb(db, keyMap)

		cmd := exec.Command("xinput", "test", strconv.Itoa(*keyboardId))

		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		if err := cmd.Start(); err != nil {
			log.Fatal(err)
		}

		// output of xinput command
		buf := make([]byte, KEYBOARD_BUFER_SIZE)

		var curStat StatForTime
		curStat.Init()

		for {
			n, err := stdout.Read(buf)
			if err != nil {
				log.Fatal(err)
			}
			// processing buf here
			for _, keyNum := range GetKeyNumsFromOutput(buf[:n]) {
				oldKeyCount, _ := curStat.keys[keyNum]
				curStat.keys[keyNum] = oldKeyCount + 1
			}

			// Every CAPTURE_TIME seconds save to BD
			if time.Now().Unix()-curStat.time > CAPTURE_TIME {
				AddStatTimeToDb(db, curStat, keyMap)
				curStat.Init()
			}

			time.Sleep(SLEEP_TIME)
		}
	case *outputPath != "":
		//exporting here
	}
}
<commit_msg>Fixed some golint warnings<commit_after>// main.go
package main

import (
	"database/sql"
	"flag"
	"log"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

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

const (
	SLEEP_TIME          = 3 * time.Second // time between processing xinput output
	KEYBOARD_BUFER_SIZE = 10000
	DATABASE_NAME       = "file:gokeystat.db?cache=shared&mode=rwc"
	CAPTURE_TIME        = 5 // time in seconds between capturing keyboard to db
)

// StatForTime stotres pressed keys and beginning time
type StatForTime struct {
	time int64
	keys map[uint8]int
}

func (stat *StatForTime) Init() {
	stat.time = time.Now().Unix()
	stat.keys = make(map[uint8]int)
}

// GetKeymap returns map from key numbers to key names like "F1", "Tab", "d"
func GetKeymap() map[uint8]string {
	return GetKeymapFromOutput(GetKeymapOutput())
}

// GetKeymapOutput returns output of utility that prints system keymap
func GetKeymapOutput() []byte {
	cmd := exec.Command("xmodmap", "-pke")
	out, err := cmd.Output()
	if err != nil {
		log.Fatal(err)
	}
	return out
}

// GetKeymapFromOutput returns map with keymap from text
func GetKeymapFromOutput(buf []byte) map[uint8]string {
	const KEY_NUM_STRING_RE = "\\d+[ ]*=[ ]*\\S+"
	re := regexp.MustCompile(KEY_NUM_STRING_RE)
	resByte := re.FindAll(buf, -1)
	keyMap := make(map[uint8]string)
	for _, line := range resByte {
		lineSpitted := strings.Split(string(line), " ")
		if key, err := strconv.Atoi(lineSpitted[0]); err == nil {
			keyMap[uint8(key)] = lineSpitted[2]
		}
	}
	return keyMap
}

// GetKeyNumsFromOutput extract pressed keys from bufer buf
// It returns slice with key numbers in the same order
func GetKeyNumsFromOutput(buf []byte) []uint8 {
	const KEY_NUM_STRING_RE = "press[ ]+(\\d+)"
	re := regexp.MustCompile(KEY_NUM_STRING_RE)
	resByte := re.FindAll(buf, -1)
	keyNums := make([]uint8, len(resByte))
	re = regexp.MustCompile("\\d+")
	for i, line := range resByte {
		numByte := re.Find(line)
		if num, err := strconv.Atoi(string(numByte)); err == nil {
			keyNums[i] = uint8(num)
		} else {
			log.Fatal(err)
		}
	}
	return keyNums
}

func GetKeyNumsFromKeyMap(keyMap map[uint8]string) []int {
	res := make([]int, 0, len(keyMap))
	for keyNum := range keyMap {
		res = append(res, int(keyNum))
	}
	sort.Ints(res)
	return res
}

// InitDb creates tables, inserts keymap to db
func InitDb(db *sql.DB, keyMap map[uint8]string) {
	keyNums := GetKeyNumsFromKeyMap(keyMap)

	sqlInit := `CREATE TABLE IF NOT EXISTS keylog (
        time INTEGER primary key`

	for keyNum := range keyNums {
		sqlInit += ",\n" + "KEY" + strconv.Itoa(keyNum) + " INTEGER"
	}
	sqlInit += "\n);"

	// Inserting keymap to table
	sqlInit += ` CREATE TABLE IF NOT EXISTS keymap (
        num INTEGER primary key,
		value STRING
	);`

	_, err := db.Exec(sqlInit)
	if err != nil {
		log.Fatalf("%q: %s\n", err, sqlInit)
	}

	rows, err := db.Query("SELECT COUNT(*) FROM keymap")
	if err != nil {
		log.Fatal(err)
	}

	var rowsCount int
	rows.Next()
	rows.Scan(&rowsCount)
	if rowsCount > 0 {
		// already inserted keymap
		return
	}
	rows.Close()

	tx, err := db.Begin()
	if err != nil {
		log.Fatal(err)
	}
	stmt, err := tx.Prepare("INSERT INTO keymap(num, value) VALUES(?, ?)")
	if err != nil {
		log.Fatal(err)
	}
	defer stmt.Close()

	for keyNum, keyName := range keyMap {
		_, err = stmt.Exec(keyNum, keyName)
		if err != nil {
			log.Fatal(err)
		}
	}
	tx.Commit()
}

func AddStatTimeToDb(db *sql.DB, statTime StatForTime, keyMap map[uint8]string) {
	keyNums := GetKeyNumsFromKeyMap(keyMap)
	sqlStmt := "insert into keylog(time"
	for keyNum := range keyNums {
		sqlStmt += ",\n" + "KEY" + strconv.Itoa(keyNum)
	}
	sqlStmt += ") values "
	sqlStmt += "(" + strconv.FormatInt(statTime.time, 10)
	for keyNum := range keyNums {
		keyNumber, _ := statTime.keys[uint8(keyNum)]
		sqlStmt += ",\n" + strconv.Itoa(keyNumber)
	}
	sqlStmt += ")"
	tx, err := db.Begin()
	if err != nil {
		log.Fatal(err)
	}

	_, err = tx.Exec(sqlStmt)
	if err != nil {
		log.Printf("%q: %s\n", err, sqlStmt)
	}

	tx.Commit()
}

// GetStatTimesFromDb returns slice with StatForTime objects that
func GetStatTimesFromDb(db *sql.DB, fromTime int64, keyMap map[uint8]string) []StatForTime {
	sqlStmt := "select * from keylog where time > " + strconv.FormatInt(fromTime, 10)
	rows, err := db.Query(sqlStmt)
	if err != nil {
		log.Fatalln("Error with query", sqlStmt, " is ", err)
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		log.Fatalln("Failed to get columns", err)
	}

	rawResult := make([][]byte, len(cols))
	result := make([]int64, len(cols))

	dest := make([]interface{}, len(cols)) // A temporary interface{} slice
	for i := range rawResult {
		dest[i] = &rawResult[i] // Put pointers to each string in the interface slice
	}

	// keyNums[i] stores i-th keynum
	keyNums := GetKeyNumsFromKeyMap(keyMap)

	// result
	res := make([]StatForTime, 0)
	for rows.Next() {
		err = rows.Scan(dest...)
		if err != nil {
			log.Fatalln("Failed to scan row", err)
		}

		for i, raw := range rawResult {
			if raw == nil {
				result[i] = 0
			} else {
				// Only numbers in db: converting it to int64
				result[i], err = strconv.ParseInt(string(raw), 10, 64)
				if err != nil {
					log.Fatalln("Error when parsing ", raw, " from db:", err)
				}
			}
		}

		var resStatTime StatForTime
		resStatTime.time = result[0]
		resStatTime.keys = make(map[uint8]int)
		for index, val := range result[1:] {
			if val == 0 {
				continue
			}
			resStatTime.keys[uint8(keyNums[index])] = int(val)
		}
		res = append(res, resStatTime)
	}
	if err = rows.Err(); err != nil {
		log.Fatalln("Error when iterating over rows", err)
	}

	return res
}

func main() {

	keyboardID := flag.Int("id", -1, "Your keyboard id")
	outputPath := flag.String("o", "", "Path to export file")
	flag.Parse()
	log.Println("keyboardID =", *keyboardID, "outputPath =", *outputPath)
	switch {
	case *keyboardID == -1 && *outputPath == "":
		flag.PrintDefaults()
		return
	case *keyboardID != -1:
		// Opening database
		db, err := sql.Open("sqlite3", DATABASE_NAME)
		if err != nil {
			log.Fatal(err)
		}
		db.SetMaxIdleConns(5)
		db.SetMaxOpenConns(5)
		defer db.Close()

		keyMap := GetKeymap()

		InitDb(db, keyMap)

		cmd := exec.Command("xinput", "test", strconv.Itoa(*keyboardID))

		stdout, err := cmd.StdoutPipe()
		if err != nil {
			log.Fatal(err)
		}
		if err := cmd.Start(); err != nil {
			log.Fatal(err)
		}

		// output of xinput command
		buf := make([]byte, KEYBOARD_BUFER_SIZE)

		var curStat StatForTime
		curStat.Init()

		for {
			n, err := stdout.Read(buf)
			if err != nil {
				log.Fatal(err)
			}
			// processing buf here
			for _, keyNum := range GetKeyNumsFromOutput(buf[:n]) {
				oldKeyCount, _ := curStat.keys[keyNum]
				curStat.keys[keyNum] = oldKeyCount + 1
			}

			// Every CAPTURE_TIME seconds save to BD
			if time.Now().Unix()-curStat.time > CAPTURE_TIME {
				AddStatTimeToDb(db, curStat, keyMap)
				curStat.Init()
			}

			time.Sleep(SLEEP_TIME)
		}
	case *outputPath != "":
		//exporting here
	}
}
<|endoftext|>