<commit_before>package main

// +build -tags netgo -a

import (
	"bitbucket.org/strings/via/pkg"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/str1ngs/util"
	"github.com/str1ngs/util/console/command"
	"github.com/str1ngs/util/file"
	"github.com/str1ngs/util/json"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
)

var (
	elog     = log.New(os.Stderr, "", log.Lshortfile)
	root     = flag.String("r", "/", "root directory")
	verbose  = flag.Bool("v", false, "verbose output")
	finstall = flag.Bool("i", true, "install package after build (default true)")
	fdebug   = flag.Bool("d", false, "debug output")
	config   = via.GetConfig()
	fclean   = flag.Bool("c", false, "clean before build")
	fupdate  = flag.Bool("u", false, "force download source")
	fdeps    = flag.Bool("deps", false, "build depends if needed")
)

func main() {
	flag.Parse()
	via.Verbose(*verbose)
	via.Update(*fupdate)
	via.Deps(*fdeps)

	via.Root(*root)
	util.Verbose = *verbose
	via.Debug(*fdebug)
	command.Add("add", add, "add plan/s to git index")
	command.Add("branch", branch, "prints plan branch to stdout")
	command.Add("build", build, "build plan")
	command.Add("cd", cd, "returns a bash evaluable cd path")
	command.Add("checkout", checkout, "changes plan branch")
	command.Add("clean", clean, "clean build dir")
	command.Add("config", fnConfig, "prints config to stdout")
	command.Add("create", create, "create plan from URL")
	command.Add("diff", diff, "prints git diff for plan(s)")
	command.Add("edit", edit, "calls EDITOR to edit plan")
	command.Add("elf", elf, "prints elf information to stdout")
	command.Add("install", install, "install package")
	command.Add("ipfs", ipfs, "test ipfs connection")
	command.Add("lint", lint, "lint plans")
	command.Add("list", list, "lists files")
	command.Add("log", plog, "print config log for plan")
	command.Add("owns", owns, "finds which package owns a file")
	command.Add("pack", pack, "package plan")
	command.Add("remove", remove, "remove package")
	command.Add("repo", repo, "update repo")
	command.Add("search", search, "search for plans (currently lists all use grep)")
	command.Add("show", fnShow, "prints plan to stdout")
	command.Add("sync", sync, "fetch remote repo data")
	command.Add("test", via.Test, "test function")
	if *fdebug {
		path, _ := os.LookupEnv("PATH")
		fmt.Println("PATH", path)
		which("GCC", "gcc")
	}
	err := command.Run()
	if err != nil {
		elog.Fatal(err)
	}
	return
}

func which(label, path string) {
	fmt.Printf("GCC ")
	cmd := exec.Command("which", path)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		elog.Println(err)
	}
}

func ipfs() error {
	res, err := http.Get(config.Binary)
	io.Copy(os.Stdout, res.Body)
	return err
}

func cd() error {
	if len(command.Args()) < 1 {
		return errors.New("you need to specify a config path")
	}
	arg := command.Args()[0]
	switch arg {
	case "plans":
		fmt.Printf("cd %s", config.Plans)
	default:
		err := fmt.Sprintf("config path %s not found", arg)
		return errors.New(err)
	}
	return nil
}

func add() error {
	if len(command.Args()) < 1 {
		return errors.New("no plans specified")
	}
	for _, arg := range command.Args() {
		glob := filepath.Join(config.Plans, "*", arg+".json")
		res, err := filepath.Glob(glob)
		if err != nil {
			return err
		}
		git := exec.Command("git", "add", strings.Join(res, " "))
		git.Dir = config.Plans
		git.Stdout = os.Stdout
		git.Stderr = os.Stderr
		err = git.Run()
		if err != nil {
			return err
		}
	}
	return nil
}

func diff() error {
	if len(command.Args()) < 1 {
		return errors.New("no plans specified")
	}
	for _, arg := range command.Args() {
		glob := filepath.Join(config.Plans, "*", arg+".json")
		res, err := filepath.Glob(glob)
		if err != nil {
			return err
		}
		git := exec.Command("git", "diff", strings.Join(res, " "))
		git.Dir = config.Plans
		git.Stdout = os.Stdout
		git.Stderr = os.Stderr
		err = git.Run()
		if err != nil {
			return err
		}
	}
	return nil
}

func checkout() error {
	if len(command.Args()) < 1 {
		return errors.New("git branch needs to be specified")
	}
	arg := command.Args()[0]
	git := exec.Command("git", "checkout", arg)
	git.Dir = config.Plans
	git.Stdout = os.Stdout
	git.Stderr = os.Stderr
	return git.Run()
}

func branch() error {
	git := exec.Command("git", "branch")
	git.Dir = config.Plans
	git.Stdout = os.Stdout
	git.Stderr = os.Stderr
	return git.Run()

}
func edit() error {
	var (
		editor = os.Getenv("EDITOR")
		arg0   = command.Args()[0]
		p      = path.Join(config.Plans, "config.json")
		err    error
	)
	if arg0 != "config" {
		p, err = via.FindPlanPath(arg0)
		if err != nil {
			return err
		}
	}
	cmd := exec.Command(editor, p)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		return err
	}
	elog.Println("linting...")
	via.Verbose(false)
	return via.Lint()
}

func create() error {
	url := command.Args()[0]
	group := command.Args()[1]
	err := via.Create(url, group)
	if err != nil {
		return err
	}
	return nil
}

func plog() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		f := path.Join(plan.BuildDir(), "config.log")
		err = file.Cat(os.Stdout, f)
		if err != nil {
			log.Fatal(err)
		}
	}
	return nil
}

func build() error {
	for _, arg := range command.Args() {
		arg = strings.Replace(arg, ".json", "", 1)
		if *fclean {
			via.Clean(arg)
		}
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		if *fdeps {
			err := via.BuildDeps(plan)
			if err != nil {
				return err
			}
		}
		err = via.BuildSteps(plan)
		if err != nil {
			return err
		}
		if *finstall {
			err := via.Install(plan.Name)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func pack() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		err = via.Package("", plan)
		if err != nil {
			return err
		}
	}
	return nil
}

func list() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		for _, f := range plan.Files {
			fmt.Println(f)
		}
	}
	return nil
}

func install() error {
	return command.ArgsDo(via.Install)
}

func remove() error {
	return command.ArgsDo(via.Remove)
}

func lint() error {
	return via.Lint()
}

func fnShow() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			elog.Fatal(err)
		}
		buf := new(bytes.Buffer)
		less := exec.Command("less")
		less.Stdin = buf
		less.Stdout = os.Stdout
		less.Stderr = os.Stderr
		err = json.WritePretty(&plan, buf)
		if err != nil {
			fmt.Println(err)
		}
		less.Run()
	}
	return nil
}

func fnConfig() error {
	err := json.WritePretty(via.GetConfig(), os.Stdout)
	if err != nil {
		return err
	}
	return nil
}

func clean() error {
	return command.ArgsDo(via.Clean)
}

func elf() error {
	for _, arg := range command.Args() {
		err := via.Readelf(arg)
		if err != nil {
			return err
		}
	}
	return nil
}

func sync() error {
	return via.PlanSync()
}

func owns() error {
	rfiles, err := via.ReadRepoFiles()
	if err != nil {
		return err
	}
	for _, arg := range command.Args() {
		owner := rfiles.Owns(arg)
		if owner == "" {
			fmt.Println(arg+":", "owner not found.")
			continue
		}
		fmt.Println(owner)
	}
	return nil
}

func repo() error {
	return via.RepoCreate()
}

func search() error {
	plans, err := via.GetPlans()
	if err != nil {
		return err
	}
	plans.SortSize().Print()
	return nil
}
<commit_msg>output HOME when debugging<commit_after>package main

// +build -tags netgo -a

import (
	"bitbucket.org/strings/via/pkg"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/str1ngs/util"
	"github.com/str1ngs/util/console/command"
	"github.com/str1ngs/util/file"
	"github.com/str1ngs/util/json"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
)

var (
	elog     = log.New(os.Stderr, "", log.Lshortfile)
	root     = flag.String("r", "/", "root directory")
	verbose  = flag.Bool("v", false, "verbose output")
	finstall = flag.Bool("i", true, "install package after build (default true)")
	fdebug   = flag.Bool("d", false, "debug output")
	config   = via.GetConfig()
	fclean   = flag.Bool("c", false, "clean before build")
	fupdate  = flag.Bool("u", false, "force download source")
	fdeps    = flag.Bool("deps", false, "build depends if needed")
)

func main() {
	flag.Parse()
	via.Verbose(*verbose)
	via.Update(*fupdate)
	via.Deps(*fdeps)

	via.Root(*root)
	util.Verbose = *verbose
	via.Debug(*fdebug)
	command.Add("add", add, "add plan/s to git index")
	command.Add("branch", branch, "prints plan branch to stdout")
	command.Add("build", build, "build plan")
	command.Add("cd", cd, "returns a bash evaluable cd path")
	command.Add("checkout", checkout, "changes plan branch")
	command.Add("clean", clean, "clean build dir")
	command.Add("config", fnConfig, "prints config to stdout")
	command.Add("create", create, "create plan from URL")
	command.Add("diff", diff, "prints git diff for plan(s)")
	command.Add("edit", edit, "calls EDITOR to edit plan")
	command.Add("elf", elf, "prints elf information to stdout")
	command.Add("install", install, "install package")
	command.Add("ipfs", ipfs, "test ipfs connection")
	command.Add("lint", lint, "lint plans")
	command.Add("list", list, "lists files")
	command.Add("log", plog, "print config log for plan")
	command.Add("owns", owns, "finds which package owns a file")
	command.Add("pack", pack, "package plan")
	command.Add("remove", remove, "remove package")
	command.Add("repo", repo, "update repo")
	command.Add("search", search, "search for plans (currently lists all use grep)")
	command.Add("show", fnShow, "prints plan to stdout")
	command.Add("sync", sync, "fetch remote repo data")
	command.Add("test", via.Test, "test function")
	if *fdebug {
		path, _ := os.LookupEnv("PATH")
		home, _ := os.LookupEnv("HOME")
		fmt.Println("PATH", path)
		fmt.Println("HOME", home)
		which("GCC", "gcc")
	}
	err := command.Run()
	if err != nil {
		elog.Fatal(err)
	}
	return
}

func which(label, path string) {
	fmt.Printf("GCC ")
	cmd := exec.Command("which", path)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err := cmd.Run()
	if err != nil {
		elog.Println(err)
	}
}

func ipfs() error {
	res, err := http.Get(config.Binary)
	io.Copy(os.Stdout, res.Body)
	return err
}

func cd() error {
	if len(command.Args()) < 1 {
		return errors.New("you need to specify a config path")
	}
	arg := command.Args()[0]
	switch arg {
	case "plans":
		fmt.Printf("cd %s", config.Plans)
	default:
		err := fmt.Sprintf("config path %s not found", arg)
		return errors.New(err)
	}
	return nil
}

func add() error {
	if len(command.Args()) < 1 {
		return errors.New("no plans specified")
	}
	for _, arg := range command.Args() {
		glob := filepath.Join(config.Plans, "*", arg+".json")
		res, err := filepath.Glob(glob)
		if err != nil {
			return err
		}
		git := exec.Command("git", "add", strings.Join(res, " "))
		git.Dir = config.Plans
		git.Stdout = os.Stdout
		git.Stderr = os.Stderr
		err = git.Run()
		if err != nil {
			return err
		}
	}
	return nil
}

func diff() error {
	if len(command.Args()) < 1 {
		return errors.New("no plans specified")
	}
	for _, arg := range command.Args() {
		glob := filepath.Join(config.Plans, "*", arg+".json")
		res, err := filepath.Glob(glob)
		if err != nil {
			return err
		}
		git := exec.Command("git", "diff", strings.Join(res, " "))
		git.Dir = config.Plans
		git.Stdout = os.Stdout
		git.Stderr = os.Stderr
		err = git.Run()
		if err != nil {
			return err
		}
	}
	return nil
}

func checkout() error {
	if len(command.Args()) < 1 {
		return errors.New("git branch needs to be specified")
	}
	arg := command.Args()[0]
	git := exec.Command("git", "checkout", arg)
	git.Dir = config.Plans
	git.Stdout = os.Stdout
	git.Stderr = os.Stderr
	return git.Run()
}

func branch() error {
	git := exec.Command("git", "branch")
	git.Dir = config.Plans
	git.Stdout = os.Stdout
	git.Stderr = os.Stderr
	return git.Run()

}
func edit() error {
	var (
		editor = os.Getenv("EDITOR")
		arg0   = command.Args()[0]
		p      = path.Join(config.Plans, "config.json")
		err    error
	)
	if arg0 != "config" {
		p, err = via.FindPlanPath(arg0)
		if err != nil {
			return err
		}
	}
	cmd := exec.Command(editor, p)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		return err
	}
	elog.Println("linting...")
	via.Verbose(false)
	return via.Lint()
}

func create() error {
	url := command.Args()[0]
	group := command.Args()[1]
	err := via.Create(url, group)
	if err != nil {
		return err
	}
	return nil
}

func plog() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		f := path.Join(plan.BuildDir(), "config.log")
		err = file.Cat(os.Stdout, f)
		if err != nil {
			log.Fatal(err)
		}
	}
	return nil
}

func build() error {
	for _, arg := range command.Args() {
		arg = strings.Replace(arg, ".json", "", 1)
		if *fclean {
			via.Clean(arg)
		}
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		if *fdeps {
			err := via.BuildDeps(plan)
			if err != nil {
				return err
			}
		}
		err = via.BuildSteps(plan)
		if err != nil {
			return err
		}
		if *finstall {
			err := via.Install(plan.Name)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func pack() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		err = via.Package("", plan)
		if err != nil {
			return err
		}
	}
	return nil
}

func list() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			return err
		}
		for _, f := range plan.Files {
			fmt.Println(f)
		}
	}
	return nil
}

func install() error {
	return command.ArgsDo(via.Install)
}

func remove() error {
	return command.ArgsDo(via.Remove)
}

func lint() error {
	return via.Lint()
}

func fnShow() error {
	for _, arg := range command.Args() {
		plan, err := via.NewPlan(arg)
		if err != nil {
			elog.Fatal(err)
		}
		buf := new(bytes.Buffer)
		less := exec.Command("less")
		less.Stdin = buf
		less.Stdout = os.Stdout
		less.Stderr = os.Stderr
		err = json.WritePretty(&plan, buf)
		if err != nil {
			fmt.Println(err)
		}
		less.Run()
	}
	return nil
}

func fnConfig() error {
	err := json.WritePretty(via.GetConfig(), os.Stdout)
	if err != nil {
		return err
	}
	return nil
}

func clean() error {
	return command.ArgsDo(via.Clean)
}

func elf() error {
	for _, arg := range command.Args() {
		err := via.Readelf(arg)
		if err != nil {
			return err
		}
	}
	return nil
}

func sync() error {
	return via.PlanSync()
}

func owns() error {
	rfiles, err := via.ReadRepoFiles()
	if err != nil {
		return err
	}
	for _, arg := range command.Args() {
		owner := rfiles.Owns(arg)
		if owner == "" {
			fmt.Println(arg+":", "owner not found.")
			continue
		}
		fmt.Println(owner)
	}
	return nil
}

func repo() error {
	return via.RepoCreate()
}

func search() error {
	plans, err := via.GetPlans()
	if err != nil {
		return err
	}
	plans.SortSize().Print()
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/evolsnow/robot/conn"
	"golang.org/x/net/websocket"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
	"unicode"
)

func main() {
	var configFile string
	flag.StringVar(&configFile, "c", "config.json", "specify config file")
	config, err := ParseConfig(configFile)
	if err != nil {
		log.Fatal("a vailid json config file must exist")
	}

	redisPort := strconv.Itoa(config.RedisPort)
	redisServer := net.JoinHostPort(config.RedisAddress, redisPort)
	if !conn.Ping(redisServer, config.RedisPassword) {
		log.Fatal("connect to redis server failed")
	}
	conn.Pool = conn.NewPool(redisServer, config.RedisPassword, config.RedisDB)

	robot := newRobot(config.RobotToken, config.RobotName, config.WebHookAddress)
	go robot.run()

	http.Handle("/websocket", websocket.Handler(socketHandler))
	srvPort := strconv.Itoa(config.Port)
	log.Fatal(http.ListenAndServe(net.JoinHostPort(config.Server, srvPort), nil))
	//used for 104
	//go http.ListenAndServeTLS("0.0.0.0:8443", "server.crt", "server.key", nil)
}

//used for web samaritan robot
func socketHandler(ws *websocket.Conn) {
	for {
		var in, response string
		var ret []string
		sf := func(c rune) bool {
			return c == ',' || c == '，' || c == ';' || c == '。' || c == '.' || c == '？' || c == '?'
		}
		if err := websocket.Message.Receive(ws, &in); err != nil {
			log.Println(err)
			return
		}
		fmt.Printf("Received: %s\n", in)
		zh := false
		for _, r := range in {
			if unicode.Is(unicode.Scripts["Han"], r) {
				log.Println(in)
				zh = true
				break
			}
		}
		if zh {
			response = tlAI(in)
			// Separate into fields with func.
			ret = strings.FieldsFunc(response, sf)

		} else {
			response = mitAI(in)
			ret = strings.FieldsFunc(response, sf)
		}
		for i := range ret {
			websocket.Message.Send(ws, ret[i])
			time.Sleep(time.Second)
		}
		websocket.Message.Send(ws, "")
	}
}
<commit_msg>flag prase<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/evolsnow/robot/conn"
	"golang.org/x/net/websocket"
	"log"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"
	"unicode"
)

func main() {
	var configFile string
	flag.StringVar(&configFile, "c", "config.json", "specify config file")
	flag.Parse()
	config, err := ParseConfig(configFile)
	if err != nil {
		log.Fatal("a vailid json config file must exist")
	}

	redisPort := strconv.Itoa(config.RedisPort)
	redisServer := net.JoinHostPort(config.RedisAddress, redisPort)
	if !conn.Ping(redisServer, config.RedisPassword) {
		log.Fatal("connect to redis server failed")
	}
	conn.Pool = conn.NewPool(redisServer, config.RedisPassword, config.RedisDB)

	robot := newRobot(config.RobotToken, config.RobotName, config.WebHookAddress)
	go robot.run()

	http.Handle("/websocket", websocket.Handler(socketHandler))
	srvPort := strconv.Itoa(config.Port)
	log.Fatal(http.ListenAndServe(net.JoinHostPort(config.Server, srvPort), nil))
	//used for 104
	//go http.ListenAndServeTLS("0.0.0.0:8443", "server.crt", "server.key", nil)
}

//used for web samaritan robot
func socketHandler(ws *websocket.Conn) {
	for {
		var in, response string
		var ret []string
		sf := func(c rune) bool {
			return c == ',' || c == '，' || c == ';' || c == '。' || c == '.' || c == '？' || c == '?'
		}
		if err := websocket.Message.Receive(ws, &in); err != nil {
			log.Println(err)
			return
		}
		fmt.Printf("Received: %s\n", in)
		zh := false
		for _, r := range in {
			if unicode.Is(unicode.Scripts["Han"], r) {
				log.Println(in)
				zh = true
				break
			}
		}
		if zh {
			response = tlAI(in)
			// Separate into fields with func.
			ret = strings.FieldsFunc(response, sf)

		} else {
			response = mitAI(in)
			ret = strings.FieldsFunc(response, sf)
		}
		for i := range ret {
			websocket.Message.Send(ws, ret[i])
			time.Sleep(time.Second)
		}
		websocket.Message.Send(ws, "")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"os"
	"time"
)

const (
	Debug_AI                      = false
	Timeout                       = true
	N                             = 2
	FirstTurnTime                 = 100
	TimeLimit                     = 50
	extra_space_between_factories = 300
	W                             = 16000
	H                             = 6500
	Min_Production_Rate           = 4
	MOVE                          = moveType(0)
	BOMB                          = moveType(1)
	INCREASE                      = moveType(2)
)

var Move_Type_String = []string{"MOVE", "BOMB", "INC"}

type Vec struct {
	x int
	y int
}

func (a Vec) Sub(b Vec) Vec {
	return Vec{a.x - b.x, a.y - b.y}
}

type Factory struct {
	owner int
	units int
	prod  int
	turns int
	r     Vec
	l     []int
}

type Troop struct {
	owner  int
	source int
	target int
	units  int
	turns  int
}

type Bomb struct {
	owner  int
	source int
	target int
	turns  int
}

type State struct {
	F        []Factory
	T        []Troop
	B        map[int]Bomb
	N_Bombs  []int
	entityId int
}

type Play struct {
	moveType moveType
	from     int
	to       int
	amount   int
}

type moveType int

type strat []Vec

func main() {
	if !flag.Parsed() {
		flag.Parse()
	}

	if *aiFlag == "" {
		flag.Usage()
		os.Exit(1)
	}

	aiOne := NewAI()
	aiTwo := NewAI()

	for {
		PlayRound([]*AI{aiOne, aiTwo})
	}

}

func generator(min int, max int) int {
	if max == 0 {
		max = 1
	}
	rndSrc := rand.NewSource(time.Now().UnixNano())
	rnd := rand.New(rndSrc)
	if min > 0 {
		return rnd.Intn(max-min) + min
	} else {
		return rnd.Intn(max)
	}
}

func PlayRound(ai []*AI) int {
	swapDistrib := UniformIntDistribution{0, 1}
	playerSwap := swapDistrib.Gen() == 1
	if playerSwap {
		// TODO swap ai
	}
	state := State{}
	rndFactoryCount := generator(7, 15)
	if rndFactoryCount%2 == 0 { // must be odd
		rndFactoryCount++
	}
	state.F = make([]Factory, rndFactoryCount)

	factoryRadius := 700
	if rndFactoryCount > 10 {
		factoryRadius -= 100
	}
	minSpaceBetweenFactories := 2 * (factoryRadius + extra_space_between_factories)

	// center Factory
	state.F[0] = Factory{0, 0, 0, 0, Vec{W / 2, H / 2}, []int{}}

	// random distribution
	xDistrib := UniformIntDistribution{(0), (W/2 - 2*factoryRadius)}
	yDistrib := UniformIntDistribution{(0), (H - 2*factoryRadius)}
	productionDistrib := UniformIntDistribution{(0), (3)}
	initialUnitsDistrib := UniformIntDistribution{(15), (30)}

	for i := 1; i < len(state.F); i += 2 {
		var r Vec
		for !ValidSpawn(r, &state, i, minSpaceBetweenFactories) {
			r = Vec{xDistrib.Gen() + factoryRadius + extra_space_between_factories, yDistrib.Gen() + factoryRadius + extra_space_between_factories}
		}
		prod := productionDistrib.Gen()
		if i == 1 {
			units := initialUnitsDistrib.Gen()
			state.F[i] = Factory{1, units, prod, 0, r, make([]int, 0)}
			state.F[i+1] = Factory{-1, units, prod, 0, Vec{W, H}.Sub(r), make([]int, 0)}
		} else {
			initialNeutralsDistrib := UniformIntDistribution{(0), (5 * prod)}
			units := initialNeutralsDistrib.Gen()
			state.F[i] = Factory{0, units, prod, 0, r, make([]int, 0)}
			state.F[i+1] = Factory{0, units, prod, 0, Vec{W, H}.Sub(r), make([]int, 0)}
		}
	}

	total_prod := 0
	for i := 0; i < len(state.F); i++ {
		total_prod += state.F[i].prod
		state.F[i].l = make([]int, len(state.F))
	}
	for i := 0; i < len(state.F); i++ {
		for j := i + 1; j < len(state.F); j++ {
			d := int((Distance(state.F[i].r, state.F[j].r) - 2*factoryRadius) / 800.0)
			state.F[i].l[j] = d
			state.F[j].l[i] = d
		}
	}
	for i := 1; total_prod < Min_Production_Rate && i < len(state.F); i++ {
		f := state.F[i]
		if f.prod < 3 {
			f.prod++
			total_prod++
		}
	}
	state.N_Bombs = []int{2, 2}
	//println("PlayGame")
	winner := Play_Game(ai, &state)
	if playerSwap {
		if winner == -1 {
			return -1
		} else if winner == 0 {
			return 1
		}
		return 0
	} else {
		return winner
	}
}

func Play_Game(ai []*AI, state *State) int {
	for i := 0; i < N; i++ {
		ss := &bytes.Buffer{}
		fmt.Fprintf(ss, "%d %d\n", len(state.F), len(state.F)*len(state.F-1)/2)
		for j := 0; j < len(state.F); j++ {
			for k := j + 1; k < len(state.F); k++ {
				fmt.Fprintf(ss, "%d %d %d\n", j, k, state.F[j].l[k])
			}
		}
		ai[i].Feed_Inputs(ss)
	}
	state.entityId = 0
	for turn := 1; turn <= 400; turn++ {
		//println("Turn:", turn)
		M := []string{"WAIT", "WAIT"}
		for i := 0; i < N; i++ {
			if ai[i].Alive() {
				if Has_Won(ai, i) {
					//cerr << i << " has won in " << turn << " turns" << endl;
					return i
				}
				color := 1
				if i != 0 {
					color = -1
				}
				ss := &bytes.Buffer{}
				fmt.Fprintf(ss, "%d %d %d\n", len(state.F), len(state.T), len(state.B))
				for j := 0; j < len(state.F); j++ {
					fmt.Fprintf(ss, "%d FACTORY %d %d %d %d\n", color*state.F[j].owner, state.F[j].units, state.F[j].prod, state.F[j].turns, 0)
				}
				for j := 0; j < len(state.T); j++ {
					fmt.Fprintf(ss, "%d TROOP %d %d %d %d\n", color*state.T[j].owner, state.T[j].source, state.T[j].target, state.T[j].units, state.T[j].turns)
				}
				//for (const auto &it:S.B){
				//const bomb &b{it.second};
				//ss << it.first << " BOMB " << color*b.owner << " " << b.source << " " << (b.owner==color?b.target:-1) << " " << (b.owner==color?b.turns:-1) << " " << 0 << endl;
				//}
				for idx, b := range state.B {
					// TODO see above
					fmt.Fprintf(ss, "%d BOMB %d %d %d %d %d", idx, color*b.owner, b.source, -1, -1, 0)
				}
				ai[i].Feed_Inputs(ss)
				M[i] = GetMove(ai[i], turn)
			}
		}
		for i := 0; i < 2; i++ {
			Play_Move(state, ai[i], M[i])
		}
		Simulate(state)
		//for i := 0; i < N; i++ {
		//	if (!Player_Alive(state, i==0?1:-1)){
		//		ai[i].stop()
		//	}
		//}
		if turn == 200 {
			Units := []int{0, 0}
			for _, f := range state.F {
				if f.owner == 1 {
					Units[0] += f.units
				} else if f.owner == -1 {
					Units[1] += f.units
				}
			}
			for _, t := range state.T {
				if t.owner == 1 {
					Units[0] += t.units
				} else if t.owner == -1 {
					Units[1] += t.units
				}
			}
			if Units[0] == Units[1] {
				//cerr << "Draw" << endl;
				return -1
			} else if Units[0] > Units[1] {
				//cerr << 0 << " has won in " << turn << " turns" << endl;
				return 0
			} else {
				//cerr << 1 << " has won in " << turn << " turns" << endl;
				return 1
			}
		}

		if All_Dead(ai) {
			os.Exit(1)
			//return -1
		}
	}
	return -2
}

func Distance(a Vec, b Vec) int {
	sqrt := math.Sqrt(math.Pow(float64(a.x-b.x), 2.0) + math.Pow(float64(a.y-b.y), 2.0))
	return int(sqrt)
}

func ValidSpawn(r Vec, state *State, id int, minSpaceBetweenFactories int) bool {
	for j := 0; j < id; j++ {
		if Distance(r, state.F[j].r) < minSpaceBetweenFactories {
			return false
		}
	}
	return true
}

func Has_Won(bot []*AI, idx int) bool {
	for i := 0; i < N; i++ {
		if i != idx && bot[i].Alive() {
			return true
		}
	}
	return false
}

func GetMove(ai *AI, turn int) string {
	buf := &bytes.Buffer{}
	duration := TimeLimit * time.Second
	if turn == 1 {
		duration = FirstTurnTime * time.Millisecond
	}
	timer := time.NewTimer(duration)
	<-timer.C
	n, err := buf.ReadFrom(ai.outPipe)
	if err != nil {
		println(err)
		os.Exit(1)
	}

	//println("AI response:", buf.String())
	if n == 0 {
		println("Loss by Timeout of AI", ai.id, ai.name)
		os.Exit(1)
	}

	return buf.String()
}

func Simulate(s *State) {
	//println("Simulate++")
}

func Play_Move(state *State, ai *AI, m string) {
	//println("PlayMove")
}

func All_Dead(ai []*AI) bool {
	for _, bot := range ai {
		if bot.Alive() {
			return false
		}
	}
	return true
}
<commit_msg>Added entity id to round output<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"os"
	"time"
)

const (
	Debug_AI                      = false
	Timeout                       = true
	N                             = 2
	FirstTurnTime                 = 100
	TimeLimit                     = 50
	extra_space_between_factories = 300
	W                             = 16000
	H                             = 6500
	Min_Production_Rate           = 4
	MOVE                          = moveType(0)
	BOMB                          = moveType(1)
	INCREASE                      = moveType(2)
)

var Move_Type_String = []string{"MOVE", "BOMB", "INC"}

type Vec struct {
	x int
	y int
}

func (a Vec) Sub(b Vec) Vec {
	return Vec{a.x - b.x, a.y - b.y}
}

type Factory struct {
	owner int
	units int
	prod  int
	turns int
	r     Vec
	l     []int
}

type Troop struct {
	owner  int
	source int
	target int
	units  int
	turns  int
}

type Bomb struct {
	owner  int
	source int
	target int
	turns  int
}

type State struct {
	F        []Factory
	T        []Troop
	B        map[int]Bomb
	N_Bombs  []int
	entityId int
}

type Play struct {
	moveType moveType
	from     int
	to       int
	amount   int
}

type moveType int

type strat []Vec

func main() {
	if !flag.Parsed() {
		flag.Parse()
	}

	if *aiFlag == "" {
		flag.Usage()
		os.Exit(1)
	}

	aiOne := NewAI()
	aiTwo := NewAI()

	for {
		PlayRound([]*AI{aiOne, aiTwo})
	}

}

func generator(min int, max int) int {
	if max == 0 {
		max = 1
	}
	rndSrc := rand.NewSource(time.Now().UnixNano())
	rnd := rand.New(rndSrc)
	if min > 0 {
		return rnd.Intn(max-min) + min
	} else {
		return rnd.Intn(max)
	}
}

func PlayRound(ai []*AI) int {
	swapDistrib := UniformIntDistribution{0, 1}
	playerSwap := swapDistrib.Gen() == 1
	if playerSwap {
		// TODO swap ai
	}
	state := State{}
	rndFactoryCount := generator(7, 15)
	if rndFactoryCount%2 == 0 { // must be odd
		rndFactoryCount++
	}
	state.F = make([]Factory, rndFactoryCount)

	factoryRadius := 700
	if rndFactoryCount > 10 {
		factoryRadius -= 100
	}
	minSpaceBetweenFactories := 2 * (factoryRadius + extra_space_between_factories)

	// center Factory
	state.F[0] = Factory{0, 0, 0, 0, Vec{W / 2, H / 2}, []int{}}

	// random distribution
	xDistrib := UniformIntDistribution{(0), (W/2 - 2*factoryRadius)}
	yDistrib := UniformIntDistribution{(0), (H - 2*factoryRadius)}
	productionDistrib := UniformIntDistribution{(0), (3)}
	initialUnitsDistrib := UniformIntDistribution{(15), (30)}

	for i := 1; i < len(state.F); i += 2 {
		var r Vec
		for !ValidSpawn(r, &state, i, minSpaceBetweenFactories) {
			r = Vec{xDistrib.Gen() + factoryRadius + extra_space_between_factories, yDistrib.Gen() + factoryRadius + extra_space_between_factories}
		}
		prod := productionDistrib.Gen()
		if i == 1 {
			units := initialUnitsDistrib.Gen()
			state.F[i] = Factory{1, units, prod, 0, r, make([]int, 0)}
			state.F[i+1] = Factory{-1, units, prod, 0, Vec{W, H}.Sub(r), make([]int, 0)}
		} else {
			initialNeutralsDistrib := UniformIntDistribution{(0), (5 * prod)}
			units := initialNeutralsDistrib.Gen()
			state.F[i] = Factory{0, units, prod, 0, r, make([]int, 0)}
			state.F[i+1] = Factory{0, units, prod, 0, Vec{W, H}.Sub(r), make([]int, 0)}
		}
	}

	total_prod := 0
	for i := 0; i < len(state.F); i++ {
		total_prod += state.F[i].prod
		state.F[i].l = make([]int, len(state.F))
	}
	for i := 0; i < len(state.F); i++ {
		for j := i + 1; j < len(state.F); j++ {
			d := int((Distance(state.F[i].r, state.F[j].r) - 2*factoryRadius) / 800.0)
			state.F[i].l[j] = d
			state.F[j].l[i] = d
		}
	}
	for i := 1; total_prod < Min_Production_Rate && i < len(state.F); i++ {
		f := state.F[i]
		if f.prod < 3 {
			f.prod++
			total_prod++
		}
	}
	state.N_Bombs = []int{2, 2}
	//println("PlayGame")
	winner := Play_Game(ai, &state)
	if playerSwap {
		if winner == -1 {
			return -1
		} else if winner == 0 {
			return 1
		}
		return 0
	} else {
		return winner
	}
}

func Play_Game(ai []*AI, state *State) int {
	for i := 0; i < N; i++ {
		ss := &bytes.Buffer{}
		fmt.Fprintf(ss, "%d %d\n", len(state.F), len(state.F)*len(state.F-1)/2)
		for j := 0; j < len(state.F); j++ {
			for k := j + 1; k < len(state.F); k++ {
				fmt.Fprintf(ss, "%d %d %d\n", j, k, state.F[j].l[k])
			}
		}
		ai[i].Feed_Inputs(ss)
	}
	state.entityId = 0
	for turn := 1; turn <= 400; turn++ {
		//println("Turn:", turn)
		M := []string{"WAIT", "WAIT"}
		for i := 0; i < N; i++ {
			if ai[i].Alive() {
				if Has_Won(ai, i) {
					//cerr << i << " has won in " << turn << " turns" << endl;
					return i
				}
				color := 1
				if i != 0 {
					color = -1
				}
				ss := &bytes.Buffer{}
				fmt.Fprintf(ss, "%d %d %d\n", len(state.F), len(state.T), len(state.B))
				for j := 0; j < len(state.F); j++ {
					fmt.Fprintf(ss, "%d FACTORY %d %d %d %d %d\n", j, color*state.F[j].owner, state.F[j].units, state.F[j].prod, state.F[j].turns, 0)
				}
				for j := 0; j < len(state.T); j++ {
					fmt.Fprintf(ss, "%d TROOP %d %d %d %d %d\n", j, color*state.T[j].owner, state.T[j].source, state.T[j].target, state.T[j].units, state.T[j].turns)
				}
				//for (const auto &it:S.B){
				//const bomb &b{it.second};
				//ss << it.first << " BOMB " << color*b.owner << " " << b.source << " " << (b.owner==color?b.target:-1) << " " << (b.owner==color?b.turns:-1) << " " << 0 << endl;
				//}
				for idx, b := range state.B {
					// TODO see above
					fmt.Fprintf(ss, "%d BOMB %d %d %d %d %d", idx, color*b.owner, b.source, -1, -1, 0)
				}
				ai[i].Feed_Inputs(ss)
				M[i] = GetMove(ai[i], turn)
			}
		}
		for i := 0; i < 2; i++ {
			Play_Move(state, ai[i], M[i])
		}
		Simulate(state)
		//for i := 0; i < N; i++ {
		//	if (!Player_Alive(state, i==0?1:-1)){
		//		ai[i].stop()
		//	}
		//}
		if turn == 200 {
			Units := []int{0, 0}
			for _, f := range state.F {
				if f.owner == 1 {
					Units[0] += f.units
				} else if f.owner == -1 {
					Units[1] += f.units
				}
			}
			for _, t := range state.T {
				if t.owner == 1 {
					Units[0] += t.units
				} else if t.owner == -1 {
					Units[1] += t.units
				}
			}
			if Units[0] == Units[1] {
				//cerr << "Draw" << endl;
				return -1
			} else if Units[0] > Units[1] {
				//cerr << 0 << " has won in " << turn << " turns" << endl;
				return 0
			} else {
				//cerr << 1 << " has won in " << turn << " turns" << endl;
				return 1
			}
		}

		if All_Dead(ai) {
			os.Exit(1)
			//return -1
		}
	}
	return -2
}

func Distance(a Vec, b Vec) int {
	sqrt := math.Sqrt(math.Pow(float64(a.x-b.x), 2.0) + math.Pow(float64(a.y-b.y), 2.0))
	return int(sqrt)
}

func ValidSpawn(r Vec, state *State, id int, minSpaceBetweenFactories int) bool {
	for j := 0; j < id; j++ {
		if Distance(r, state.F[j].r) < minSpaceBetweenFactories {
			return false
		}
	}
	return true
}

func Has_Won(bot []*AI, idx int) bool {
	for i := 0; i < N; i++ {
		if i != idx && bot[i].Alive() {
			return true
		}
	}
	return false
}

func GetMove(ai *AI, turn int) string {
	buf := &bytes.Buffer{}
	duration := TimeLimit * time.Second
	if turn == 1 {
		duration = FirstTurnTime * time.Millisecond
	}
	timer := time.NewTimer(duration)
	<-timer.C
	n, err := buf.ReadFrom(ai.outPipe)
	if err != nil {
		println(err)
		os.Exit(1)
	}

	//println("AI response:", buf.String())
	if n == 0 {
		println("Loss by Timeout of AI", ai.id, ai.name)
		os.Exit(1)
	}

	return buf.String()
}

func Simulate(s *State) {
	//println("Simulate++")
}

func Play_Move(state *State, ai *AI, m string) {
	//println("PlayMove")
}

func All_Dead(ai []*AI) bool {
	for _, bot := range ai {
		if bot.Alive() {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"log"
	"os"
	"time"

	"github.com/fraenkel/candiedyaml"
	"github.com/garyburd/redigo/redis"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
	"github.com/tedsuo/router"
	proleroutes "github.com/winston-ci/prole/routes"

	"github.com/winston-ci/winston/api"
	"github.com/winston-ci/winston/api/drainer"
	apiroutes "github.com/winston-ci/winston/api/routes"
	"github.com/winston-ci/winston/builder"
	"github.com/winston-ci/winston/config"
	"github.com/winston-ci/winston/db"
	"github.com/winston-ci/winston/server"
	"github.com/winston-ci/winston/watcher"
	"github.com/winston-ci/winston/watchman"
)

var configPath = flag.String(
	"config",
	"",
	"path to winston server config .yml",
)

var templatesDir = flag.String(
	"templates",
	"",
	"path to directory containing the html templates",
)

var publicDir = flag.String(
	"public",
	"",
	"path to directory containing public resources (javascript, css, etc.)",
)

var proleURL = flag.String(
	"proleURL",
	"http://127.0.0.1:4637",
	"address denoting the prole service",
)

var redisAddr = flag.String(
	"redisAddr",
	"127.0.0.1:6379",
	"redis server address",
)

var peerAddr = flag.String(
	"peerAddr",
	"127.0.0.1:8081",
	"external address of the api server, used for callbacks",
)

var listenAddr = flag.String(
	"listenAddr",
	":8080",
	"port for the web server to listen on",
)

var apiListenAddr = flag.String(
	"apiListenAddr",
	":8081",
	"port for the api to listen on",
)

func main() {
	flag.Parse()

	if *configPath == "" {
		fatal(errors.New("must specify -config"))
	}

	if *templatesDir == "" {
		fatal(errors.New("must specify -templates"))
	}

	if *publicDir == "" {
		fatal(errors.New("must specify -public"))
	}

	configFile, err := os.Open(*configPath)
	if err != nil {
		fatal(err)
	}

	var conf config.Config
	err = candiedyaml.NewDecoder(configFile).Decode(&conf)
	if err != nil {
		fatal(err)
	}

	configFile.Close()

	redisDB := db.NewRedis(redis.NewPool(func() (redis.Conn, error) {
		return redis.DialTimeout("tcp", *redisAddr, 5*time.Second, 0, 0)
	}, 20))

	winstonEndpoint := router.NewRequestGenerator("http://"+*peerAddr, apiroutes.Routes)
	proleEndpoint := router.NewRequestGenerator(*proleURL, proleroutes.Routes)
	builder := builder.NewBuilder(redisDB, conf.Resources, proleEndpoint, winstonEndpoint)

	serverHandler, err := server.New(conf, redisDB, *templatesDir, *publicDir, *peerAddr, builder)
	if err != nil {
		fatal(err)
	}

	drainer := drainer.NewDrainer()

	apiHandler, err := api.New(redisDB, drainer)
	if err != nil {
		fatal(err)
	}

	watchman := watchman.NewWatchman(builder)

	group := grouper.EnvokeGroup(grouper.RunGroup{
		"web":     http_server.New(*listenAddr, serverHandler),
		"api":     http_server.New(*apiListenAddr, apiHandler),
		"watcher": watcher.NewWatcher(conf.Jobs, conf.Resources, redisDB, proleEndpoint, watchman),
		"drainer": ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
			close(ready)
			<-signals
			drainer.Drain()
			return nil
		}),
	})

	running := ifrit.Envoke(sigmon.New(group))

	log.Println("serving web on", *listenAddr)
	log.Println("serving api on", *apiListenAddr)

	err = <-running.Wait()
	if err != nil {
		fatal(err)
	}

	log.Println("exited")
}

func fatal(err error) {
	println(err.Error())
	os.Exit(1)
}
<commit_msg>change code around<commit_after>package main

import (
	"errors"
	"flag"
	"log"
	"os"
	"time"

	"github.com/fraenkel/candiedyaml"
	"github.com/garyburd/redigo/redis"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
	"github.com/tedsuo/router"
	proleroutes "github.com/winston-ci/prole/routes"

	"github.com/winston-ci/winston/api"
	"github.com/winston-ci/winston/api/drainer"
	apiroutes "github.com/winston-ci/winston/api/routes"
	"github.com/winston-ci/winston/builder"
	"github.com/winston-ci/winston/config"
	"github.com/winston-ci/winston/db"
	"github.com/winston-ci/winston/server"
	"github.com/winston-ci/winston/watcher"
	"github.com/winston-ci/winston/watchman"
)

var configPath = flag.String(
	"config",
	"",
	"path to winston server config .yml",
)

var templatesDir = flag.String(
	"templates",
	"",
	"path to directory containing the html templates",
)

var publicDir = flag.String(
	"public",
	"",
	"path to directory containing public resources (javascript, css, etc.)",
)

var proleURL = flag.String(
	"proleURL",
	"http://127.0.0.1:4637",
	"address denoting the prole service",
)

var redisAddr = flag.String(
	"redisAddr",
	"127.0.0.1:6379",
	"redis server address",
)

var peerAddr = flag.String(
	"peerAddr",
	"127.0.0.1:8081",
	"external address of the api server, used for callbacks",
)

var listenAddr = flag.String(
	"listenAddr",
	":8080",
	"port for the web server to listen on",
)

var apiListenAddr = flag.String(
	"apiListenAddr",
	":8081",
	"port for the api to listen on",
)

func main() {
	flag.Parse()

	if *configPath == "" {
		fatal(errors.New("must specify -config"))
	}

	if *templatesDir == "" {
		fatal(errors.New("must specify -templates"))
	}

	if *publicDir == "" {
		fatal(errors.New("must specify -public"))
	}

	configFile, err := os.Open(*configPath)
	if err != nil {
		fatal(err)
	}

	var conf config.Config
	err = candiedyaml.NewDecoder(configFile).Decode(&conf)
	if err != nil {
		fatal(err)
	}

	configFile.Close()

	redisDB := db.NewRedis(redis.NewPool(func() (redis.Conn, error) {
		return redis.DialTimeout("tcp", *redisAddr, 5*time.Second, 0, 0)
	}, 20))

	winstonEndpoint := router.NewRequestGenerator("http://"+*peerAddr, apiroutes.Routes)
	proleEndpoint := router.NewRequestGenerator(*proleURL, proleroutes.Routes)
	builder := builder.NewBuilder(redisDB, conf.Resources, proleEndpoint, winstonEndpoint)

	serverHandler, err := server.New(
		conf,
		redisDB,
		*templatesDir,
		*publicDir,
		*peerAddr,
		builder,
	)
	if err != nil {
		fatal(err)
	}

	drainer := drainer.NewDrainer()

	apiHandler, err := api.New(redisDB, drainer)
	if err != nil {
		fatal(err)
	}

	watchman := watchman.NewWatchman(builder)

	group := grouper.EnvokeGroup(grouper.RunGroup{
		"web":     http_server.New(*listenAddr, serverHandler),
		"api":     http_server.New(*apiListenAddr, apiHandler),
		"watcher": watcher.NewWatcher(conf.Jobs, conf.Resources, redisDB, proleEndpoint, watchman),
		"drainer": ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
			close(ready)
			<-signals
			drainer.Drain()
			return nil
		}),
	})

	running := ifrit.Envoke(sigmon.New(group))

	log.Println("serving web on", *listenAddr)
	log.Println("serving api on", *apiListenAddr)

	err = <-running.Wait()
	if err != nil {
		fatal(err)
	}

	log.Println("exited")
}

func fatal(err error) {
	println(err.Error())
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/anacrolix/utp"
	"io"
	"os"
)

func beNetcat(con io.ReadWriteCloser) {
	defer con.Close()

	go io.Copy(os.Stdout, con)
	io.Copy(con, os.Stdin)
}

func main() {
	list := flag.Bool("l", false, "listen on the given address")

	flag.Parse()

	if len(flag.Args()) < 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s [-l] <host> <port>\n", os.Args[0])
		os.Exit(1)
	}

	addr := fmt.Sprintf("%s:%s", flag.Arg(0), flag.Arg(1))

	var con io.ReadWriteCloser
	if *list {
		sock, err := utp.NewSocket("udp", addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "create socket failed: %s\n", err)
			os.Exit(1)
		}

		utpcon, err := sock.Accept()
		if err != nil {
			fmt.Fprintf(os.Stderr, "accept failed: %s\n", err)
			os.Exit(1)
		}

		con = utpcon
	} else {
		utpcon, err := utp.Dial(addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dial failed: %s\n", err)
			os.Exit(1)
		}

		con = utpcon
	}

	beNetcat(con)
}
<commit_msg>better error handling<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/anacrolix/utp"
	"io"
	"os"
)

func beNetcat(con io.ReadWriteCloser) {
	defer con.Close()

	go func() {
		_, err := io.Copy(os.Stdout, con)
		if err != io.EOF {
			fmt.Fprintf(os.Stderr, "Read error: %s\n", err)
			os.Exit(1)
		}
	}()
	_, err := io.Copy(con, os.Stdin)
	if err != io.EOF {
		fmt.Fprintf(os.Stderr, "Write error: %s\n", err)
		os.Exit(1)
	}
}

func main() {
	list := flag.Bool("l", false, "listen on the given address")

	flag.Parse()

	if len(flag.Args()) < 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s [-l] <host> <port>\n", os.Args[0])
		os.Exit(1)
	}

	addr := fmt.Sprintf("%s:%s", flag.Arg(0), flag.Arg(1))

	var con io.ReadWriteCloser
	if *list {
		sock, err := utp.NewSocket("udp", addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "create socket failed: %s\n", err)
			os.Exit(1)
		}

		defer sock.Close()

		utpcon, err := sock.Accept()
		if err != nil {
			fmt.Fprintf(os.Stderr, "accept failed: %s\n", err)
			os.Exit(1)
		}

		con = utpcon
	} else {
		utpcon, err := utp.Dial(addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "dial failed: %s\n", err)
			os.Exit(1)
		}

		con = utpcon
	}

	beNetcat(con)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
)

var cmd *cobra.Command

func init() {
	cmd = &cobra.Command{
		Use:   "hel",
		Short: "A mock generator for Go",
		Long: "A simple mock generator.  The origin of the name is the Norse goddess, Hel, " +
			"who guards over the souls of those unworthy to enter Valhalla.  You can probably " +
			"guess how much I like mocks.",
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) > 0 {
				fmt.Println("Invalid usage.\n")
				err := cmd.Help()
				if err != nil {
					panic(err)
				}
				os.Exit(1)
			}
			packages, err := cmd.Flags().GetStringSlice("package")
			if err != nil {
				panic(err)
			}
			types, err := cmd.Flags().GetStringSlice("type")
			if err != nil {
				panic(err)
			}
			output, err := cmd.Flags().GetString("output")
			if err != nil {
				panic(err)
			}
			fmt.Printf("Generating mocks for packages %v", packages)
			fmt.Printf("Generating mocks for types %v", types)
			fmt.Printf("Generating mocks in %s", output)
		},
	}
	cmd.Flags().StringSliceP("package", "p", []string{"."}, "The package(s) to generate mocks for.")
	cmd.Flags().StringSliceP("type", "t", []string{}, "The type(s) to generate mocks for.  If no types "+
 "are passed in, all exported interface types will be generated.")
	cmd.Flags().StringP("output", "o", "hel.go", "The file to write generated mocks to.  Since hel does "+
 "not generate exported types, this file will be saved directly in all packages with generated mocks.")
}

func main() {
	cmd.Execute()
}
<commit_msg>Fixed up main for current generation logic<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/nelsam/hel/mocks"
	"github.com/nelsam/hel/packages"
	"github.com/nelsam/hel/types"
	"github.com/spf13/cobra"
)

var cmd *cobra.Command

func init() {
	cmd = &cobra.Command{
		Use:   "hel",
		Short: "A mock generator for Go",
		Long: "A simple mock generator.  The origin of the name is the Norse goddess, Hel, " +
			"who guards over the souls of those unworthy to enter Valhalla.  You can probably " +
			"guess how much I like mocks.",
		Run: func(cmd *cobra.Command, args []string) {
			if len(args) > 0 {
				fmt.Println("Invalid usage.\n")
				err := cmd.Help()
				if err != nil {
					panic(err)
				}
				os.Exit(1)
			}
			packagePatterns, err := cmd.Flags().GetStringSlice("package")
			if err != nil {
				panic(err)
			}
			typePatterns, err := cmd.Flags().GetStringSlice("type")
			if err != nil {
				panic(err)
			}
			outputName, err := cmd.Flags().GetString("output")
			if err != nil {
				panic(err)
			}
			chanSize, err := cmd.Flags().GetInt("chan-size")
			if err != nil {
				panic(err)
			}
			fmt.Printf("Generating mocks for packages %v", packagePatterns)
			var packageList []packages.Package
			progress(func() {
				packageList = packages.Load(packagePatterns...)
			})
			fmt.Print("\n")
			fmt.Printf("Generating mocks for types %v", typePatterns)
			var typeList []types.Types
			progress(func() {
				typeList = types.Load(packageList...) //.Filter(typePatterns...)
			})
			fmt.Print("\n")
			fmt.Printf("Generating mocks in %s", outputName)
			progress(func() {
				for _, types := range typeList {
					if err := makeMocks(types, outputName, chanSize); err != nil {
						panic(err)
					}
				}
			})
			fmt.Print("\n")
		},
	}
	cmd.Flags().StringSliceP("package", "p", []string{"."}, "The package(s) to generate mocks for.")
	cmd.Flags().StringSliceP("type", "t", []string{}, "The type(s) to generate mocks for.  If no types "+
		"are passed in, all exported interface types will be generated.")
	cmd.Flags().StringP("output", "o", "hel_test.go", "The file to write generated mocks to.  Since hel does "+
		"not generate exported types, this file will be saved directly in all packages with generated mocks.  "+
		"Also note that, since the types are not exported, you will want the file to end in '_test.go'.")
	cmd.Flags().IntP("chan-size", "s", 100, "The size of channels used for method calls.")
}

func makeMocks(types types.Types, fileName string, chanSize int) error {
	mocks, err := mocks.Generate(types)
	if err != nil {
		return err
	}
	if len(mocks) == 0 {
		return nil
	}
	f, err := os.Create(filepath.Join(types.Dir(), fileName))
	if err != nil {
		return err
	}
	defer f.Close()
	return mocks.Output(types.TestPackage(), chanSize, f)
}

func progress(f func()) {
	stop, done := make(chan struct{}), make(chan struct{})
	defer func() {
		close(stop)
		<-done
	}()
	go showProgress(stop, done)
	f()
}

func showProgress(stop <-chan struct{}, done chan<- struct{}) {
	defer close(done)
	ticker := time.NewTicker(time.Second / 2)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			fmt.Print(".")
		case <-stop:
			return
		}
	}
}

func main() {
	cmd.Execute()
}
<|endoftext|>


<commit_before>package main

import (
	"bitbucket.org/kardianos/osext"
	"bufio"
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/fgrehm/devstep-cli/devstep"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

var (
	project devstep.Project
	client  devstep.DockerClient
)

var commands = []cli.Command{
	hackCmd,
	buildCmd,
	bootstrapCmd,
	infoCmd,
	runCmd,
	binstubsCmd,
	cleanCmd,
	pristineCmd,
}

func main() {
	app := cli.NewApp()
	app.Name = "devstep"
	app.Author = "Fábio Rehm"
	app.Email = "fgrehm@gmail.com"
	app.Usage = "development environments made easy"
	app.Version = "0.1.0"
	app.EnableBashCompletion = true
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "log-level, l", Usage: "log level", EnvVar: "DEVSTEP_LOG"},
	}
	app.Before = func(c *cli.Context) error {
		return devstep.SetLogLevel(c.GlobalString("log-level"))
	}
	app.Commands = commands

	app.RunAndExitOnError()
}

var dockerRunFlags = []cli.Flag{
	cli.StringFlag{Name: "w, working_dir", Usage: "Working directory inside the container"},
	cli.StringSliceFlag{Name: "p, publish", Value: &cli.StringSlice{}, Usage: "Publish a container's port to the host (hostPort:containerPort)"},
	cli.StringSliceFlag{Name: "link", Value: &cli.StringSlice{}, Usage: "Add link to another container (name:alias)"},
	cli.StringSliceFlag{Name: "e, env", Value: &cli.StringSlice{}, Usage: "Set environment variables"},
	cli.BoolFlag{Name: "privileged", Usage: "Give extended privileges to this container"},
}

func loadConfig() *devstep.ProjectConfig {
	projectRoot, err := os.Getwd()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	homeDir := os.Getenv("HOME")
	client = devstep.NewClient("unix:///var/run/docker.sock")
	loader := devstep.NewConfigLoader(client, homeDir, projectRoot)

	config, err := loader.Load()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	pluginsToLoad, err := filepath.Glob(homeDir + "/devstep/plugins/*/plugin.js")
	if err != nil {
		fmt.Printf("Error searching for plugins under '%s'\n%s\n", homeDir, err.Error())
		os.Exit(1)
	}

	if len(pluginsToLoad) > 0 {
		runtime := devstep.NewPluginRuntime(config)
		for _, pluginPath := range pluginsToLoad {
			runtime.Load(pluginPath)
		}
		runtime.Trigger("configLoaded")
	}

	if devstep.LogLevel != "" {
		config.Defaults.Env["DEVSTEP_LOG"] = devstep.LogLevel
	}

	return config
}

func newProject() devstep.Project {
	config := loadConfig()
	project, _ = devstep.NewProject(config)
	return project
}

var buildCmd = cli.Command{
	Name:  "build",
	Usage: "build a docker image for the current project",
	Action: func(c *cli.Context) {
		err := newProject().Build(client)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var bootstrapCmd = cli.Command{
	Name:  "bootstrap",
	Usage: "bootstrap an environment for the current project",
	Action: func(c *cli.Context) {
		err := newProject().Bootstrap(client)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var hackCmd = cli.Command{
	Name:  "hack",
	Usage: "start a hacking session for the current project",
	Flags: dockerRunFlags,
	BashComplete: func(c *cli.Context) {
		args := c.Args()
		if len(args) == 0 {
			fmt.Println("-p")
			fmt.Println("--publish")
			fmt.Println("--link")
			fmt.Println("-w")
			fmt.Println("--working_dir")
			fmt.Println("-e")
			fmt.Println("--env")
			fmt.Println("--privileged")
		}
	},
	Action: func(c *cli.Context) {
		project := newProject()
		runOpts := parseRunOpts(c)
		err := project.Hack(client, runOpts)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var runCmd = cli.Command{
	Name:  "run",
	Usage: "Run a one off command against the current base image",
	Flags: dockerRunFlags,
	Action: func(c *cli.Context) {
		project := newProject()
		commands := project.Config().Commands

		runOpts := parseRunOpts(c)
		runOpts.Cmd = c.Args()

		// Validate command
		if len(runOpts.Cmd) == 0 {
			fmt.Printf("No command provided to `devstep run`\n\n")
			cli.ShowCommandHelp(c, "run")
			os.Exit(1)
		}

		if cmd, ok := commands[runOpts.Cmd[0]]; ok {
			runOpts = cmd.Merge(runOpts)
			runOpts.Cmd = append(cmd.Cmd, runOpts.Cmd[1:]...)
		}

		// Prepend a `--` so that it doesn't interfere with the current init
		// process args
		runOpts.Cmd = append([]string{"--"}, runOpts.Cmd...)

		err := project.Run(client, runOpts)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var binstubsCmd = cli.Command{
	Name:  "binstubs",
	Usage: "Generate binstubs for the commands specified on devstep.yml",
	Action: func(c *cli.Context) {
		project := newProject()
		commands := project.Config().Commands

		if len(commands) == 0 {
			fmt.Println("No binstubs specified!")
			os.Exit(0)
		}

		binstubsPath := ".devstep/bin"
		os.MkdirAll("./"+binstubsPath, 0700)

		executable, _ := osext.Executable()

		for _, cmd := range commands {
			script := []byte("#!/usr/bin/env bash\neval \"" + executable + " run -- " + cmd.Name + " $@\"")
			err := ioutil.WriteFile(binstubsPath+"/"+cmd.Name, script, 0755)
			if err != nil {
				fmt.Printf("Error creating binstub '%s'\n%s\n", binstubsPath+"/"+cmd.Name, err)
				os.Exit(1)
			}
			fmt.Printf("Generated binstub for '%s' in '%s'\n", cmd.Name, binstubsPath)
		}
	},
}

var cleanCmd = cli.Command{
	Name:  "clean",
	Usage: "remove previously built images for the current environment",
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "force, f", Usage: "skip confirmation"},
	},
	Action: func(c *cli.Context) {
		if !c.Bool("force") {
			fmt.Print("Are you sure? [N/y] ")

			reader := bufio.NewReader(os.Stdin)
			line, _, err := reader.ReadLine()
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}

			if strings.ToUpper(string(line)) != "Y" {
				fmt.Println("Aborting")
				os.Exit(1)
			}
		}

		err := newProject().Clean(client)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var pristineCmd = cli.Command{
	Name:  "pristine",
	Usage: "rebuild project image from scratch",
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "force, f", Usage: "skip clean confirmation"},
	},
	Action: func(c *cli.Context) {
		// TODO: Figure out if this is the right way to invoke other CLI actions
		cleanCmd.Action(c)
		buildCmd.Action(c)
	},
}

var infoCmd = cli.Command{
	Name:  "info",
	Usage: "show information about the current environment",
	Action: func(c *cli.Context) {
		config := loadConfig()
		printConfig(config)
	},
}

func printConfig(config *devstep.ProjectConfig) {
	fmt.Println("==> Project info")
	fmt.Printf("Repository:   %s\n", config.RepositoryName)
	fmt.Printf("Source image: %s\n", config.SourceImage)
	fmt.Printf("Base image:   %s\n", config.BaseImage)
	fmt.Printf("Host dir:     %s\n", config.HostDir)
	fmt.Printf("Guest dir:    %s\n", config.GuestDir)
	fmt.Printf("Cache dir:    %s\n", config.CacheDir)

	if config.Defaults != nil {
		fmt.Println("\n==> Default options:")
		printDockerRunOpts(config.Defaults, "")
	}

	if config.HackOpts != nil {
		fmt.Println("\n==> Hack options:")
		printDockerRunOpts(config.HackOpts, "")
	}
	if config.Commands != nil {
		fmt.Println("\n==> Commands:")
		for _, cmd := range config.Commands {
			fmt.Printf("* %s\n", cmd.Name)
			fmt.Printf("  Cmd:        %v\n", cmd.Cmd)
			fmt.Printf("  Publish:    %v\n", cmd.Publish)
			printDockerRunOpts(&cmd.DockerRunOpts, "  ")
		}
	}
}

func printDockerRunOpts(opts *devstep.DockerRunOpts, prefix string) {
	fmt.Printf("%sPrivileged: %v\n", prefix, opts.Privileged)
	fmt.Printf("%sLinks:      %v\n", prefix, opts.Links)
	fmt.Printf("%sVolumes:    %v\n", prefix, opts.Volumes)
	fmt.Printf("%sEnv:        %v\n", prefix, opts.Env)
}

func parseRunOpts(c *cli.Context) *devstep.DockerRunOpts {
	runOpts := &devstep.DockerRunOpts{
		Publish: c.StringSlice("publish"),
		Links:   c.StringSlice("link"),
		Env:     make(map[string]string),
	}

	// Only set the privileged config if it was provided
	if c.IsSet("privileged") {
		privileged := c.Bool("privileged")
		runOpts.Privileged = &privileged
	}

	// Set working dir directly on the project object so that it get passed
	// along properly
	if workingDir := c.String("working_dir"); workingDir != "" {
		project.Config().GuestDir = workingDir
	}

	// Env vars
	for _, envVar := range c.StringSlice("env") {
		varAndValue := strings.Split(envVar, "=")
		runOpts.Env[varAndValue[0]] = varAndValue[1]
	}

	// Validate ports
	validPort := regexp.MustCompile(`^\d+:\d+$`)
	for _, port := range runOpts.Publish {
		if !validPort.MatchString(port) {
			fmt.Println("Invalid publish arg: " + port)
			os.Exit(1)
		}
	}
	// Validate links
	validLink := regexp.MustCompile(`[^:]+:[^:]+`)
	for _, link := range runOpts.Links {
		if !validLink.MatchString(link) {
			fmt.Println("Invalid link: " + link)
			os.Exit(1)
		}
	}
	return runOpts
}
<commit_msg>Use prompt lib for confirmation<commit_after>package main

import (
	"bitbucket.org/kardianos/osext"
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/fgrehm/devstep-cli/devstep"
	"github.com/segmentio/go-prompt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

var (
	project devstep.Project
	client  devstep.DockerClient
)

var commands = []cli.Command{
	hackCmd,
	buildCmd,
	bootstrapCmd,
	infoCmd,
	runCmd,
	binstubsCmd,
	cleanCmd,
	pristineCmd,
}

func main() {
	app := cli.NewApp()
	app.Name = "devstep"
	app.Author = "Fábio Rehm"
	app.Email = "fgrehm@gmail.com"
	app.Usage = "development environments made easy"
	app.Version = "0.1.0"
	app.EnableBashCompletion = true
	app.Flags = []cli.Flag{
		cli.StringFlag{Name: "log-level, l", Usage: "log level", EnvVar: "DEVSTEP_LOG"},
	}
	app.Before = func(c *cli.Context) error {
		return devstep.SetLogLevel(c.GlobalString("log-level"))
	}
	app.Commands = commands

	app.RunAndExitOnError()
}

var dockerRunFlags = []cli.Flag{
	cli.StringFlag{Name: "w, working_dir", Usage: "Working directory inside the container"},
	cli.StringSliceFlag{Name: "p, publish", Value: &cli.StringSlice{}, Usage: "Publish a container's port to the host (hostPort:containerPort)"},
	cli.StringSliceFlag{Name: "link", Value: &cli.StringSlice{}, Usage: "Add link to another container (name:alias)"},
	cli.StringSliceFlag{Name: "e, env", Value: &cli.StringSlice{}, Usage: "Set environment variables"},
	cli.BoolFlag{Name: "privileged", Usage: "Give extended privileges to this container"},
}

func loadConfig() *devstep.ProjectConfig {
	projectRoot, err := os.Getwd()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	homeDir := os.Getenv("HOME")
	client = devstep.NewClient("unix:///var/run/docker.sock")
	loader := devstep.NewConfigLoader(client, homeDir, projectRoot)

	config, err := loader.Load()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	pluginsToLoad, err := filepath.Glob(homeDir + "/devstep/plugins/*/plugin.js")
	if err != nil {
		fmt.Printf("Error searching for plugins under '%s'\n%s\n", homeDir, err.Error())
		os.Exit(1)
	}

	if len(pluginsToLoad) > 0 {
		runtime := devstep.NewPluginRuntime(config)
		for _, pluginPath := range pluginsToLoad {
			runtime.Load(pluginPath)
		}
		runtime.Trigger("configLoaded")
	}

	if devstep.LogLevel != "" {
		config.Defaults.Env["DEVSTEP_LOG"] = devstep.LogLevel
	}

	return config
}

func newProject() devstep.Project {
	config := loadConfig()
	project, _ = devstep.NewProject(config)
	return project
}

var buildCmd = cli.Command{
	Name:  "build",
	Usage: "build a docker image for the current project",
	Action: func(c *cli.Context) {
		err := newProject().Build(client)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var bootstrapCmd = cli.Command{
	Name:  "bootstrap",
	Usage: "bootstrap an environment for the current project",
	Action: func(c *cli.Context) {
		err := newProject().Bootstrap(client)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var hackCmd = cli.Command{
	Name:  "hack",
	Usage: "start a hacking session for the current project",
	Flags: dockerRunFlags,
	BashComplete: func(c *cli.Context) {
		args := c.Args()
		if len(args) == 0 {
			fmt.Println("-p")
			fmt.Println("--publish")
			fmt.Println("--link")
			fmt.Println("-w")
			fmt.Println("--working_dir")
			fmt.Println("-e")
			fmt.Println("--env")
			fmt.Println("--privileged")
		}
	},
	Action: func(c *cli.Context) {
		project := newProject()
		runOpts := parseRunOpts(c)
		err := project.Hack(client, runOpts)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var runCmd = cli.Command{
	Name:  "run",
	Usage: "Run a one off command against the current base image",
	Flags: dockerRunFlags,
	Action: func(c *cli.Context) {
		project := newProject()
		commands := project.Config().Commands

		runOpts := parseRunOpts(c)
		runOpts.Cmd = c.Args()

		// Validate command
		if len(runOpts.Cmd) == 0 {
			fmt.Printf("No command provided to `devstep run`\n\n")
			cli.ShowCommandHelp(c, "run")
			os.Exit(1)
		}

		if cmd, ok := commands[runOpts.Cmd[0]]; ok {
			runOpts = cmd.Merge(runOpts)
			runOpts.Cmd = append(cmd.Cmd, runOpts.Cmd[1:]...)
		}

		// Prepend a `--` so that it doesn't interfere with the current init
		// process args
		runOpts.Cmd = append([]string{"--"}, runOpts.Cmd...)

		err := project.Run(client, runOpts)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var binstubsCmd = cli.Command{
	Name:  "binstubs",
	Usage: "Generate binstubs for the commands specified on devstep.yml",
	Action: func(c *cli.Context) {
		project := newProject()
		commands := project.Config().Commands

		if len(commands) == 0 {
			fmt.Println("No binstubs specified!")
			os.Exit(0)
		}

		binstubsPath := ".devstep/bin"
		os.MkdirAll("./"+binstubsPath, 0700)

		executable, _ := osext.Executable()

		for _, cmd := range commands {
			script := []byte("#!/usr/bin/env bash\neval \"" + executable + " run -- " + cmd.Name + " $@\"")
			err := ioutil.WriteFile(binstubsPath+"/"+cmd.Name, script, 0755)
			if err != nil {
				fmt.Printf("Error creating binstub '%s'\n%s\n", binstubsPath+"/"+cmd.Name, err)
				os.Exit(1)
			}
			fmt.Printf("Generated binstub for '%s' in '%s'\n", cmd.Name, binstubsPath)
		}
	},
}

var cleanCmd = cli.Command{
	Name:  "clean",
	Usage: "remove previously built images for the current environment",
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "force, f", Usage: "skip confirmation"},
	},
	Action: func(c *cli.Context) {
		if !c.Bool("force") {
			if ok := prompt.Confirm("Are you sure? [N/y] "); !ok {
				fmt.Println("Aborting")
				os.Exit(1)
			}
		}

		err := newProject().Clean(client)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	},
}

var pristineCmd = cli.Command{
	Name:  "pristine",
	Usage: "rebuild project image from scratch",
	Flags: []cli.Flag{
		cli.BoolFlag{Name: "force, f", Usage: "skip clean confirmation"},
	},
	Action: func(c *cli.Context) {
		// TODO: Figure out if this is the right way to invoke other CLI actions
		cleanCmd.Action(c)
		buildCmd.Action(c)
	},
}

var infoCmd = cli.Command{
	Name:  "info",
	Usage: "show information about the current environment",
	Action: func(c *cli.Context) {
		config := loadConfig()
		printConfig(config)
	},
}

func printConfig(config *devstep.ProjectConfig) {
	fmt.Println("==> Project info")
	fmt.Printf("Repository:   %s\n", config.RepositoryName)
	fmt.Printf("Source image: %s\n", config.SourceImage)
	fmt.Printf("Base image:   %s\n", config.BaseImage)
	fmt.Printf("Host dir:     %s\n", config.HostDir)
	fmt.Printf("Guest dir:    %s\n", config.GuestDir)
	fmt.Printf("Cache dir:    %s\n", config.CacheDir)

	if config.Defaults != nil {
		fmt.Println("\n==> Default options:")
		printDockerRunOpts(config.Defaults, "")
	}

	if config.HackOpts != nil {
		fmt.Println("\n==> Hack options:")
		printDockerRunOpts(config.HackOpts, "")
	}
	if config.Commands != nil {
		fmt.Println("\n==> Commands:")
		for _, cmd := range config.Commands {
			fmt.Printf("* %s\n", cmd.Name)
			fmt.Printf("  Cmd:        %v\n", cmd.Cmd)
			fmt.Printf("  Publish:    %v\n", cmd.Publish)
			printDockerRunOpts(&cmd.DockerRunOpts, "  ")
		}
	}
}

func printDockerRunOpts(opts *devstep.DockerRunOpts, prefix string) {
	fmt.Printf("%sPrivileged: %v\n", prefix, opts.Privileged)
	fmt.Printf("%sLinks:      %v\n", prefix, opts.Links)
	fmt.Printf("%sVolumes:    %v\n", prefix, opts.Volumes)
	fmt.Printf("%sEnv:        %v\n", prefix, opts.Env)
}

func parseRunOpts(c *cli.Context) *devstep.DockerRunOpts {
	runOpts := &devstep.DockerRunOpts{
		Publish: c.StringSlice("publish"),
		Links:   c.StringSlice("link"),
		Env:     make(map[string]string),
	}

	// Only set the privileged config if it was provided
	if c.IsSet("privileged") {
		privileged := c.Bool("privileged")
		runOpts.Privileged = &privileged
	}

	// Set working dir directly on the project object so that it get passed
	// along properly
	if workingDir := c.String("working_dir"); workingDir != "" {
		project.Config().GuestDir = workingDir
	}

	// Env vars
	for _, envVar := range c.StringSlice("env") {
		varAndValue := strings.Split(envVar, "=")
		runOpts.Env[varAndValue[0]] = varAndValue[1]
	}

	// Validate ports
	validPort := regexp.MustCompile(`^\d+:\d+$`)
	for _, port := range runOpts.Publish {
		if !validPort.MatchString(port) {
			fmt.Println("Invalid publish arg: " + port)
			os.Exit(1)
		}
	}
	// Validate links
	validLink := regexp.MustCompile(`[^:]+:[^:]+`)
	for _, link := range runOpts.Links {
		if !validLink.MatchString(link) {
			fmt.Println("Invalid link: " + link)
			os.Exit(1)
		}
	}
	return runOpts
}
<|endoftext|>


<commit_before>// Copyright 2016, RadiantBlue Technologies, 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 main

import (
	"log"

	"github.com/venicegeo/pz-gocommon/elasticsearch"
	piazza "github.com/venicegeo/pz-gocommon/gocommon"
	pzlogger "github.com/venicegeo/pz-logger/logger"
	pzuuidgen "github.com/venicegeo/pz-uuidgen/uuidgen"
	pzworkflow "github.com/venicegeo/pz-workflow/workflow"
)

func main() {

	required := []piazza.ServiceName{
		piazza.PzElasticSearch,
		piazza.PzLogger,
		piazza.PzUuidgen,
		piazza.PzKafka,
		piazza.PzServiceController,
	}

	sys, err := piazza.NewSystemConfig(piazza.PzWorkflow, required)
	if err != nil {
		log.Fatal(err)
	}

	logger, err := pzlogger.NewClient(sys)
	if err != nil {
		log.Fatal(err)
	}

	//	uuidgen, err := uuidgenPkg.NewMockUuidGenService(sys)
	uuidgen, err := pzuuidgen.NewClient(sys)

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

	eventtypesIndex, err := elasticsearch.NewIndex(sys, "eventtypes002", pzworkflow.EventTypeIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	eventsIndex, err := elasticsearch.NewIndex(sys, "events003", pzworkflow.EventIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	triggersIndex, err := elasticsearch.NewIndex(sys, "triggers002", pzworkflow.TriggerIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	alertsIndex, err := elasticsearch.NewIndex(sys, "alerts002", pzworkflow.AlertIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	cronIndex, err := elasticsearch.NewIndex(sys, "crons002", pzworkflow.CronIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	testElasticsearchIndex, err := elasticsearch.NewIndex(sys, "testElasticsearch001", "")
	if err != nil {
		log.Fatal(err)
	}

	logger.Info("pz-workflow starting...")

	workflowService := &pzworkflow.Service{}
	err = workflowService.Init(
		sys,
		logger,
		uuidgen,
		eventtypesIndex,
		eventsIndex,
		triggersIndex,
		alertsIndex,
		cronIndex,
		testElasticsearchIndex)
	if err != nil {
		log.Fatal(err)
	}

	err = workflowService.InitCron()
	if err != nil {
		log.Fatal(err)
	}

	workflowServer := &pzworkflow.Server{}
	err = workflowServer.Init(workflowService)
	if err != nil {
		log.Fatal(err)
	}

	genericServer := piazza.GenericServer{Sys: sys}

	err = genericServer.Configure(workflowServer.Routes)
	if err != nil {
		log.Fatal(err)
	}

	done, err := genericServer.Start()
	if err != nil {
		log.Fatal(err)
	}

	err = <-done
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>ES is starting to annoy me.<commit_after>// Copyright 2016, RadiantBlue Technologies, 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 main

import (
	"log"

	"github.com/venicegeo/pz-gocommon/elasticsearch"
	piazza "github.com/venicegeo/pz-gocommon/gocommon"
	pzlogger "github.com/venicegeo/pz-logger/logger"
	pzuuidgen "github.com/venicegeo/pz-uuidgen/uuidgen"
	pzworkflow "github.com/venicegeo/pz-workflow/workflow"
)

func main() {

	required := []piazza.ServiceName{
		piazza.PzElasticSearch,
		piazza.PzLogger,
		piazza.PzUuidgen,
		piazza.PzKafka,
		piazza.PzServiceController,
	}

	sys, err := piazza.NewSystemConfig(piazza.PzWorkflow, required)
	if err != nil {
		log.Fatal(err)
	}

	logger, err := pzlogger.NewClient(sys)
	if err != nil {
		log.Fatal(err)
	}

	//	uuidgen, err := uuidgenPkg.NewMockUuidGenService(sys)
	uuidgen, err := pzuuidgen.NewClient(sys)

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

	eventtypesIndex, err := elasticsearch.NewIndex(sys, "eventtypes002", pzworkflow.EventTypeIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	eventsIndex, err := elasticsearch.NewIndex(sys, "events003", pzworkflow.EventIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	triggersIndex, err := elasticsearch.NewIndex(sys, "triggers002", pzworkflow.TriggerIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	alertsIndex, err := elasticsearch.NewIndex(sys, "alerts002", pzworkflow.AlertIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	cronIndex, err := elasticsearch.NewIndex(sys, "crons002", pzworkflow.CronIndexSettings)
	if err != nil {
		log.Fatal(err)
	}
	testElasticsearchIndex, err := elasticsearch.NewIndex(sys, "testelasticsearch001", "")
	if err != nil {
		log.Fatal(err)
	}

	logger.Info("pz-workflow starting...")

	workflowService := &pzworkflow.Service{}
	err = workflowService.Init(
		sys,
		logger,
		uuidgen,
		eventtypesIndex,
		eventsIndex,
		triggersIndex,
		alertsIndex,
		cronIndex,
		testElasticsearchIndex)
	if err != nil {
		log.Fatal(err)
	}

	err = workflowService.InitCron()
	if err != nil {
		log.Fatal(err)
	}

	workflowServer := &pzworkflow.Server{}
	err = workflowServer.Init(workflowService)
	if err != nil {
		log.Fatal(err)
	}

	genericServer := piazza.GenericServer{Sys: sys}

	err = genericServer.Configure(workflowServer.Routes)
	if err != nil {
		log.Fatal(err)
	}

	done, err := genericServer.Start()
	if err != nil {
		log.Fatal(err)
	}

	err = <-done
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"

	bundle "github.com/TykTechnologies/tyk-cli/bundle"
)

// tyk-cli <module> <submodule> <command> [--options] args...

var module, submodule, command string

var bundleOutput, privKey string
var forceInsecure *bool

func init() {
	if len(os.Args) == 1 {
		fmt.Println("No module specified!")
		os.Exit(1)
	}
	if len(os.Args) == 2 {
		fmt.Println("No command specified!")
		os.Exit(1)
	}

	module = os.Args[1]
	command = os.Args[2]

	os.Args = os.Args[2:]

	flag.StringVar(&bundleOutput, "output", "", "Bundle output")
	flag.StringVar(&privKey, "key", "", "Key for bundle signature")
	forceInsecure = flag.Bool("y", false, "Skip bundle signing")

	flag.Parse()
}

// main is the entrypoint.
func main() {
	fmt.Println("tyk-cli:", flag.CommandLine, os.Args)

	fmt.Println("module =", module)
	fmt.Println("command =", command)

	var err error

	switch module {
	case "bundle":
		fmt.Println("Using bundle module.")
		err = bundle.Bundle(command, bundleOutput, privKey, forceInsecure)
	default:
		err = errors.New("Invalid module")
	}

	if err != nil {
		fmt.Println("Error:", err)
		os.Exit(1)
	}
}
<commit_msg>Display program version.<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"

	bundle "github.com/TykTechnologies/tyk-cli/bundle"
)

// tyk-cli <module> <submodule> <command> [--options] args...

var module, submodule, command string

var bundleOutput, privKey string
var forceInsecure, showVersion *bool

func init() {
	if len(os.Args) == 1 {
		fmt.Println("No module specified!")
		os.Exit(1)
	}
	if len(os.Args) == 2 {
		fmt.Println("No command specified!")
		os.Exit(1)
	}

	module = os.Args[1]
	command = os.Args[2]

	os.Args = os.Args[2:]

	flag.StringVar(&bundleOutput, "output", "", "Bundle output")
	flag.StringVar(&privKey, "key", "", "Key for bundle signature")
	forceInsecure = flag.Bool("y", false, "Skip bundle signing")
	showVersion = flag.Bool("v", false, "Show version")

	flag.Parse()
}

// main is the entrypoint.
func main() {
	fmt.Println("tyk-cli:", flag.CommandLine, os.Args)

	fmt.Println("module =", module)
	fmt.Println("command =", command)

	if *showVersion {
		fmt.Println("tyk-cli", VERSION)
		os.Exit(0)
	}

	var err error

	switch module {
	case "bundle":
		fmt.Println("Using bundle module.")
		err = bundle.Bundle(command, bundleOutput, privKey, forceInsecure)
	default:
		err = errors.New("Invalid module")
	}

	if err != nil {
		fmt.Println("Error:", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
)

var credit float64
var expense float64

// Record represents a single line the the file
type Record struct {
	record []string
}

// CREDIT ...
var CREDIT = "Credit"

// DEBT ..
var DEBT = "Debt"

// Transaction is a money movement
type Transaction struct {
	field           string
	description     string
	transactionType string
}

func (t *Transaction) new(record Record) *Transaction {

	for i := 0; i < len(record.record); i++ {
		t.field = record.record[i]
		t.description = record.record[2]

		// Expense
		if i == 3 && record.record[i] != "" {
			value := t.value()
			t.transactionType = DEBT
			expense += value
			return t
		}

		// Credit
		if i == 4 && record.record[i] != "" {
			t.transactionType = CREDIT
			value := t.value()
			credit += value
			return t
		}
	}

	return t
}

func (t *Transaction) value() float64 {
	parsedField := strings.Replace(t.field, ".", "", -1)
	parsedField = strings.Replace(parsedField, ",", ".", -1)
	retValue, _ := strconv.ParseFloat(parsedField, 64)

	return retValue
}

// documentation for csv is at http://golang.org/pkg/encoding/csv/
func main() {

	file, error := os.Open("comprovativo.csv")
	if error != nil {
		fmt.Println("Error:", error)
		return
	}
	defer file.Close()

	reader := csv.NewReader(file)
	reader.Comma = ';'
	reader.FieldsPerRecord = -1 // If FieldsPerRecord is negative, no check is made and records may have a variable number of fields.
	lineCount := 0

	var report map[string]float64
	report = make(map[string]float64)

	for {
		r, error := reader.Read()
		record := Record{r}
		if error == io.EOF {
			break
		}
		if error != nil {
			fmt.Println("Error:", error)
			lineCount++
			continue
		}

		if lineCount < 4 {
			lineCount++
			continue
		}

		if len(record.record) != 8 {
			lineCount++
			continue
		}

		t := Transaction{}
		transaction := t.new(record)
		report[transaction.description] += transaction.value()
		lineCount++
	}

	for transactionDescription, transactionValue := range report {
		fmt.Printf("%24s %8.2f \n", transactionDescription, transactionValue)
	}

	fmt.Println("Expense is ", expense)
	fmt.Println("Credit is ", credit)

	// TODO: Fetch data
	// Here, i want this data
	// Initial balance
	// Final Balance
	// Expense per 'description field'
	// Total expense
	// Total credit

}
<commit_msg>Some TODO's<commit_after>package main

import (
	"encoding/csv"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
)

var credit float64
var expense float64

// Record represents a single line the the file
type Record struct {
	record []string
}

// CREDIT ...
var CREDIT = "Credit"

// DEBT ..
var DEBT = "Debt"

// Transaction is a money movement
type Transaction struct {
	field           string
	description     string
	transactionType string
}

func (t *Transaction) new(record Record) *Transaction {

	for i := 0; i < len(record.record); i++ {
		t.field = record.record[i]
		t.description = record.record[2]

		// Expense
		if i == 3 && record.record[i] != "" {
			value := t.value()
			t.transactionType = DEBT
			expense += value
			return t
		}

		// Credit
		if i == 4 && record.record[i] != "" {
			t.transactionType = CREDIT
			value := t.value()
			credit += value
			return t
		}
	}

	return t
}

func (t *Transaction) value() float64 {
	parsedField := strings.Replace(t.field, ".", "", -1)
	parsedField = strings.Replace(parsedField, ",", ".", -1)
	retValue, _ := strconv.ParseFloat(parsedField, 64)

	return retValue
}

// documentation for csv is at http://golang.org/pkg/encoding/csv/
func main() {

	file, error := os.Open("comprovativo.csv")
	if error != nil {
		fmt.Println("Error:", error)
		return
	}
	defer file.Close()

	reader := csv.NewReader(file)
	reader.Comma = ';'
	reader.FieldsPerRecord = -1 // If FieldsPerRecord is negative, no check is made and records may have a variable number of fields.
	lineCount := 0

	var report map[string]float64
	report = make(map[string]float64)

	// TODO: Open SQlite, read the initial balance.
	// TODO: Check if the initial balance matches the one comming in the file
	// TODO: Read the records and save each one to the SQlite

	for {
		r, error := reader.Read()
		record := Record{r}
		if error == io.EOF {
			break
		}
		if error != nil {
			fmt.Println("Error:", error)
			lineCount++
			continue
		}

		if lineCount < 4 {
			lineCount++
			continue
		}

		if len(record.record) != 8 {
			lineCount++
			continue
		}

		t := Transaction{}
		transaction := t.new(record)
		report[transaction.description] += transaction.value()
		lineCount++
	}

	for transactionDescription, transactionValue := range report {
		fmt.Printf("%24s %8.2f \n", transactionDescription, transactionValue)
	}

	fmt.Println("Expense is ", expense)
	fmt.Println("Credit is ", credit)

	// TODO: Fetch data
	// Here, i want this data
	// Initial balance
	// Final Balance
	// Expense per 'description field'
	// Total expense
	// Total credit

}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "log"
    "net/http"
    "os"
    "io/ioutil"
)

func main() {
  /* Check if we are run to [re]-initialise the board */
  if (len(os.Args) >= 4) {
    // TODO make a proper class here
    key, token, boardid := os.Args[1], os.Args[2], os.Args[3]
    baseurl := "https://api.trello.com/1"

    resp, err := http.Get(baseurl + "/boards/" + boardid + "?key=" + key + "&token=" + token)
    // TODO if error
    defer resp.Body.Close()
    body, err := ioutil.ReadAll(resp.Body)
    err = err

    fmt.Println(string(body[:]))

    /* Archive all lists */

    /* Check and activate GitHub powerup */

    /* Create the new lists */

    /* Happily print the JSON */
  } else {
    port := os.Getenv("PORT")

  	if port == "" {
  		log.Fatal("$PORT must be set")
  	}

    http.HandleFunc("/issues", Issues);
    log.Fatal(http.ListenAndServe(":"+port, nil))
  }
}

func Issues(w http.ResponseWriter, r *http.Request) {
  // TODO io.LimitReader
    body, err := ioutil.ReadAll(r.Body)
    if err != nil {
        log.Fatal(err)
    }

    fmt.Fprintln(w, "")
    if err := r.Body.Close(); err != nil {
        log.Fatal(err)
    }

    log.Print(string(body[:]))
}
<commit_msg>Trello class core<commit_after>package main

import (
    "fmt"
    "log"
    "net/http"
    "os"
    "io/ioutil"
)

/* TODO separate file */
type Trello struct {
  BaseURL string
  Token string
  Key string
} /* TODO make private */

func NewTrello(key string, token string) *Trello {
  t := new(Trello)
  t.BaseURL = "https://api.trello.com/1"
  t.Token = token
  t.Key = key

  return t
}

func (this *Trello) Get(rq string) string {
  resp, err := http.Get(this.BaseURL + rq + "?key=" + this.Key + "&token=" + this.Token)
  // TODO if error
  defer resp.Body.Close()
  body, err := ioutil.ReadAll(resp.Body)
  err = err

  return string(body[:])
}

func main() {
  /* Check if we are run to [re]-initialise the board */
  if (len(os.Args) >= 4) {
    // TODO make a proper class here
    key, token, boardid := os.Args[1], os.Args[2], os.Args[3]
    trello := NewTrello(key, token)

    fmt.Println(trello.Get("/boards/" + boardid))

    /* Archive all lists */

    /* Check and activate GitHub powerup */

    /* Create the new lists */

    /* Happily print the JSON */
  } else {
    port := os.Getenv("PORT")

  	if port == "" {
  		log.Fatal("$PORT must be set")
  	}

    http.HandleFunc("/issues", Issues);
    log.Fatal(http.ListenAndServe(":"+port, nil))
  }
}

func Issues(w http.ResponseWriter, r *http.Request) {
  // TODO io.LimitReader
    body, err := ioutil.ReadAll(r.Body)
    if err != nil {
        log.Fatal(err)
    }

    fmt.Fprintln(w, "")
    if err := r.Body.Close(); err != nil {
        log.Fatal(err)
    }

    log.Print(string(body[:]))
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
)

var (
	flagVerbose  = flag.Bool("v", false, "be verbose")
	flagHTTPPort = flag.Int("hport", 80, "HTTP server port")
	flagDNSPort  = flag.Int("dport", 53, "DNS server port")
	cntMsgs      = expvar.NewInt("statsQuestions")
	cntRelayed   = expvar.NewInt("statsRelayed")
	cntBlocked   = expvar.NewInt("statsBlocked")
	cntServed    = expvar.NewInt("statsServed")
	cntErrors    = expvar.NewInt("statsErrors")
	cntRules     = expvar.NewInt("statsRules")
	pixel        = "\x47\x49\x46\x38\x39\x61\x01\x00\x01\x00\x80\x00\x00\xff\xff" +
		"\xff\x00\x00\x00\x21\xf9\x04\x01\x00\x00\x00\x00\x2c\x00\x00" +
		"\x00\x00\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b"
	exts = map[string]bool{
		"jpg":  true,
		"jpeg": true,
		"png":  true,
		"gif":  true,
	}
	proxy    *net.UDPConn
	upstream *net.UDPConn
	queries  map[int]*net.UDPAddr
	blocked  map[string]bool
	answer   []byte
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] upstream proxy list.txt\n\n"+
			"upstream - real upstream DNS address, e.g. 8.8.8.8\n"+
			"proxy    - your address, e.g. 127.0.0.1\n"+
			"list.txt - text file with addresses to kill\n\n",
			os.Args[0],
		)
		flag.PrintDefaults()
		return
	}
	flag.Parse()

	if len(flag.Args()) < 3 {
		flag.Usage()
		os.Exit(1)
	}

	upIP := net.ParseIP(flag.Arg(0))
	if upIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse upstream IP '%s'\n", flag.Arg(0))
		os.Exit(2)
	}

	proxyIP := net.ParseIP(flag.Arg(1))
	if proxyIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse proxy IP '%s'\n", flag.Arg(1))
		os.Exit(2)
	}

	answer = []byte("\x00\x01\x00\x01\xff\xff\xff\xff\x00\x04")
	answerIP := proxyIP.To4()
	if answerIP == nil {
		fmt.Fprintln(os.Stderr, "ERROR: IPv6 is not supported, sorry")
		os.Exit(3)
	}
	answer = append(answer, answerIP...)

	parseList(flag.Arg(2))

	var err error
	upAddr := &net.UDPAddr{IP: upIP, Port: 53}
	upstream, err = net.DialUDP("udp", nil, upAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer upstream.Close()

	proxyAddr := &net.UDPAddr{IP: proxyIP, Port: *flagDNSPort}
	proxy, err = net.ListenUDP("udp", proxyAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer proxy.Close()

	queries = make(map[int]*net.UDPAddr, 4096)

	go runServerHTTP(flag.Arg(1))
	go runServerLocalDNS()
	go runServerUpstreamDNS()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGUSR1)

forever:
	for {
		select {
		case s := <-sig:
			switch s {
			case syscall.SIGUSR1:
				log.Println("SIGUSR1 received, reloading rules")
			default:
				log.Println("Signal received, stopping")
				break forever
			}
		}
	}
}

func parseList(path string) {
	file, err := os.Open(path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR:", err)
		os.Exit(2)
	}
	defer file.Close()

	blocked = make(map[string]bool, 4096)
	counter := 0
	scn := bufio.NewScanner(file)
	for scn.Scan() {
		counter++
		blocked[scn.Text()+"."] = true
	}
	cntRules.Set(int64(counter))
	log.Printf("DNS: Parsed %d entries from list\n", counter)

	return
}

func runServerLocalDNS() {
	log.Println("DNS: Started local server at", proxy.LocalAddr())

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, addr, err := proxy.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		msg := make([]byte, n)
		copy(msg, buf[:n])
		cntMsgs.Add(1)
		go handleDNS(msg, addr)
	}

	panic("not reachable")
}

func runServerUpstreamDNS() {
	log.Println("DNS: Started upstream server")

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, _, err := upstream.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		id := int(uint16(buf[0])<<8 + uint16(buf[1]))
		if to, ok := queries[id]; ok {
			delete(queries, id)
			_, err := proxy.WriteTo(buf[:n], to)
			if err != nil {
				log.Println("DNS ERROR:", err)
				cntErrors.Add(1)
				continue
			}
			cntRelayed.Add(1)
			if *flagVerbose {
				log.Println("DNS: Relayed answer to query", id)
			}
		}
	}

	panic("not reachable")
}

func popPart(host string) string {
	parts := strings.Split(host, ".")
	if len(parts) < 3 {
		return ""
	}

	return strings.Join(parts[1:], ".")
}

func handleDNS(msg []byte, from *net.UDPAddr) {
	var domain bytes.Buffer
	var host string
	var block bool
	var try int

	id := int(uint16(msg[0])<<8 + uint16(msg[1]))
	if *flagVerbose {
		log.Printf("DNS: Query id %d from %s\n", id, from)
	}

	// peak query
	count := uint16(msg[5]) // question counter
	offset := uint16(12)    // point to first domain name

	if count != 1 {
		log.Printf("DNS WARN: Query id %d from %s has %d questions\n", id, from, count)
		return
	}

outer:
	for count > 0 {
	inner:
		for {
			length := int8(msg[offset])
			if length == 0 {
				break inner
			}
			offset++
			domain.WriteString(string(msg[offset : offset+uint16(length)]))
			domain.WriteString(".")
			offset += uint16(length)
		}
		host = domain.String()
		testHost := host
		try = 1

	test:
		for {
			if _, ok := blocked[testHost]; ok {
				block = true
				break outer
			}
			testHost = popPart(testHost)
			if testHost == "" {
				break test
			}
			try++
		}
		domain.Reset()

		offset += 4
		count--
	}
	// end peak query

	if block {
		// fake answer
		cntBlocked.Add(1)
		if *flagVerbose {
			log.Printf("DNS: Blocking (%d) %s\n", try, host)
		}

		msg[2] = uint8(129) // flags upper byte
		msg[3] = uint8(128) // flags lower byte
		msg[7] = uint8(1)   // answer counter

		res := append(msg, msg[12:12+1+len(host)]...)
		res = append(res, answer...)
		_, err := proxy.WriteTo(res, from)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			return
		}
		if *flagVerbose {
			log.Println("DNS: Sent fake answer")
		}
		return
		// end fake answer
	} else {
		if *flagVerbose {
			log.Println("DNS: Asking upstream")
		}
		queries[id] = from
		_, err := upstream.Write(msg)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			goto clean
		}
		return
	}

clean:
	delete(queries, id)

	return
}

func handleHTTP(w http.ResponseWriter, req *http.Request) {
	if *flagVerbose {
		log.Printf("HTTP: Request %s %s %s\n", req.Method, req.Host, req.RequestURI)
	}
	cntServed.Add(1)
	w.Header()["Content-type"] = []string{"image/gif"}
	io.WriteString(w, pixel)

	return
}

func runServerHTTP(host string) {
	addr := fmt.Sprintf("%s:%d", host, *flagHTTPPort)
	http.HandleFunc("/", handleHTTP)
	log.Println("HTTP: Started at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))

	panic("not reachable")
}

// vim: ts=4 sw=4 sts=4
<commit_msg>Cleanup<commit_after>package main

import (
	"bufio"
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
)

var (
	flagVerbose  = flag.Bool("v", false, "be verbose")
	flagHTTPPort = flag.Int("hport", 80, "HTTP server port")
	flagDNSPort  = flag.Int("dport", 53, "DNS server port")
	cntMsgs      = expvar.NewInt("statsQuestions")
	cntRelayed   = expvar.NewInt("statsRelayed")
	cntBlocked   = expvar.NewInt("statsBlocked")
	cntServed    = expvar.NewInt("statsServed")
	cntErrors    = expvar.NewInt("statsErrors")
	cntRules     = expvar.NewInt("statsRules")
	answer       = []byte("\x00\x01\x00\x01\xff\xff\xff\xff\x00\x04")
	pixel        = "\x47\x49\x46\x38\x39\x61\x01\x00\x01\x00\x80\x00\x00\xff\xff" +
		"\xff\x00\x00\x00\x21\xf9\x04\x01\x00\x00\x00\x00\x2c\x00\x00" +
		"\x00\x00\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b"
	proxy    *net.UDPConn
	upstream *net.UDPConn
	queries  map[int]*net.UDPAddr
	blocked  map[string]bool
)

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] upstream proxy list.txt\n\n"+
			"upstream - real upstream DNS address, e.g. 8.8.8.8\n"+
			"proxy    - your address, e.g. 127.0.0.1\n"+
			"list.txt - text file with addresses to kill\n\n",
			os.Args[0],
		)
		flag.PrintDefaults()
		return
	}
	flag.Parse()

	if len(flag.Args()) < 3 {
		flag.Usage()
		os.Exit(1)
	}

	upIP := net.ParseIP(flag.Arg(0))
	if upIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse upstream IP '%s'\n", flag.Arg(0))
		os.Exit(2)
	}

	proxyIP := net.ParseIP(flag.Arg(1))
	if proxyIP == nil {
		fmt.Fprintf(os.Stderr, "ERROR: Can't parse proxy IP '%s'\n", flag.Arg(1))
		os.Exit(2)
	}

	answerIP := proxyIP.To4()
	if answerIP == nil {
		fmt.Fprintln(os.Stderr, "ERROR: IPv6 is not supported, sorry")
		os.Exit(3)
	}
	answer = append(answer, answerIP...)

	parseList(flag.Arg(2))

	var err error
	upAddr := &net.UDPAddr{IP: upIP, Port: 53}
	upstream, err = net.DialUDP("udp", nil, upAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer upstream.Close()

	proxyAddr := &net.UDPAddr{IP: proxyIP, Port: *flagDNSPort}
	proxy, err = net.ListenUDP("udp", proxyAddr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s\n", err)
		os.Exit(2)
	}
	defer proxy.Close()

	queries = make(map[int]*net.UDPAddr, 4096)

	go runServerHTTP(flag.Arg(1))
	go runServerUpstreamDNS()
	go runServerLocalDNS()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGUSR1)

forever:
	for {
		select {
		case s := <-sig:
			switch s {
			case syscall.SIGUSR1:
				log.Println("SIGUSR1 received, reloading rules")
			default:
				log.Println("Signal received, stopping")
				break forever
			}
		}
	}
}

func parseList(path string) {
	file, err := os.Open(path)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR:", err)
		os.Exit(2)
	}
	defer file.Close()

	blocked = make(map[string]bool, 4096)
	counter := 0
	scn := bufio.NewScanner(file)
	for scn.Scan() {
		counter++
		blocked[scn.Text()+"."] = true
	}
	log.Printf("DNS: Parsed %d entries from list\n", counter)
	cntRules.Set(int64(counter))

	return
}

func runServerLocalDNS() {
	log.Println("DNS: Started local server at", proxy.LocalAddr())

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, addr, err := proxy.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		msg := make([]byte, n)
		copy(msg, buf[:n])
		cntMsgs.Add(1)
		go handleDNS(msg, addr)
	}
}

func runServerUpstreamDNS() {
	log.Println("DNS: Started upstream server")

	buf := make([]byte, 512)
	oobuf := make([]byte, 64)

	for {
		n, _, _, _, err := upstream.ReadMsgUDP(buf, oobuf)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			continue
		}

		id := int(uint16(buf[0])<<8 + uint16(buf[1]))
		if to, ok := queries[id]; ok {
			delete(queries, id)
			_, err := proxy.WriteTo(buf[:n], to)
			if err != nil {
				log.Println("DNS ERROR:", err)
				cntErrors.Add(1)
				continue
			}
			if *flagVerbose {
				log.Println("DNS: Relayed answer to query", id)
			}
			cntRelayed.Add(1)
		}
	}
}

func handleDNS(msg []byte, from *net.UDPAddr) {
	var domain bytes.Buffer
	var block bool

	id := int(uint16(msg[0])<<8 + uint16(msg[1]))
	if *flagVerbose {
		log.Printf("DNS: Query id %d from %s\n", id, from)
	}

	// peak query
	count := uint16(msg[5]) // question counter
	offset := uint16(12)    // point to first domain name

	if count != 1 {
		log.Printf("DNS WARN: Query id %d from %s has %d questions\n", id, from, count)
		return
	}

	for {
		length := int8(msg[offset])
		if length == 0 {
			break
		}
		offset++
		domain.WriteString(string(msg[offset : offset+uint16(length)]))
		domain.WriteString(".")
		offset += uint16(length)
	}
	host := domain.String()
	testHost := host
	parts := strings.Split(host, ".")
	try := 1
	for {
		if _, ok := blocked[testHost]; ok {
			block = true
			break
		}
		parts = parts[1:]
		if len(parts) < 3 {
			break
		}
		testHost = strings.Join(parts, ".")
		try++
	}
	domain.Reset()
	// end peak query

	if block {
		// fake answer
		if *flagVerbose {
			log.Printf("DNS: Blocking (%d) %s\n", try, host)
		}
		cntBlocked.Add(1)

		msg[2] = uint8(129) // flags upper byte
		msg[3] = uint8(128) // flags lower byte
		msg[7] = uint8(1)   // answer counter

		res := append(msg, msg[12:12+1+len(host)]...)
		res = append(res, answer...)
		_, err := proxy.WriteTo(res, from)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			return
		}
		if *flagVerbose {
			log.Println("DNS: Sent fake answer")
		}
		return
		// end fake answer
	} else {
		if *flagVerbose {
			log.Println("DNS: Asking upstream")
		}
		queries[id] = from
		_, err := upstream.Write(msg)
		if err != nil {
			log.Println("DNS ERROR:", err)
			cntErrors.Add(1)
			goto clean
		}
		return
	}

clean:
	delete(queries, id)

	return
}

func handleHTTP(w http.ResponseWriter, req *http.Request) {
	if *flagVerbose {
		log.Printf("HTTP: Request %s %s %s\n", req.Method, req.Host, req.RequestURI)
	}
	cntServed.Add(1)
	w.Header()["Content-type"] = []string{"image/gif"}
	io.WriteString(w, pixel)

	return
}

func runServerHTTP(host string) {
	addr := fmt.Sprintf("%s:%d", host, *flagHTTPPort)
	http.HandleFunc("/", handleHTTP)
	log.Println("HTTP: Started at", addr)
	log.Fatalln(http.ListenAndServe(addr, nil))

	panic("not reachable")
}

// vim: ts=4 sw=4 sts=4
<|endoftext|>