<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"sort"
	"strings"

	"github.com/ryanuber/columnize"
	"github.com/tylerfowle/dtags/db"
)

var database *db.Database

// Entrypoint for dtags application. Initializes database instance and
// processed passed command to perform the expected action.
func main() {
	var err error
	if database, err = db.Init(); err != nil {
		panic(err)
	}
	defer database.Instance.Close()

	if len(os.Args) < 2 {
		showHelp()
		os.Exit(0)
	}

	switch os.Args[1] {
	case "add":
		addNewTag(os.Args[2:])
		break
	case "del":
		database.DeleteKey(strings.ToLower(os.Args[2:][0]))
		break
	case "list", "completion":
		printAllTags()
		break
	case "ls":
		printBoth()
		break
	default:
		args := os.Args[1:]
		printPath(args)
	}
}

// Show the help message to a user when not enough arguments are passed
// to the command.
func showHelp() {
	fmt.Println("usage: dtags command [params]")
	fmt.Println()
	fmt.Println("       add <tag> [path] - add a new tag at path; defaults to current directory")
	fmt.Println("       del <tag> - delete the provided tag from storage")
	fmt.Println("       list - list all tags currently stored")
	fmt.Println("       ls - list all tags currently stored and associated paths")
}

func addNewTag(args []string) {
	k := strings.ToLower(args[0])
	v := database.CurrentDirectory

	if len(args[0:]) > 1 {
		v = args[1]
	}

	if database.Exists(k) {
		fmt.Printf("Overwrite existing tag? [%s] (y/n)", k)
		if confirmation() == false {
			return
		}
	}

	database.AddKey(k, v)
}

func printAllTags() {
	for _, tag := range database.GetTags() {
		fmt.Println(tag)
	}
}

func printBoth() {
	var unformattedlist []string
	for tag, path := range database.All() {
		unformattedlist = append(unformattedlist, fmt.Sprintf("%s|%s\n", tag, path))
	}
	sort.Strings(unformattedlist)
	formattedList := columnize.SimpleFormat(unformattedlist)
	fmt.Println(formattedList)
}

func printPath(args []string) {
	cwd := database.GetValue(args[0])
	if cwd == "" {
		fmt.Printf("tag not found\n")
		os.Exit(1)
	}

	fmt.Fprint(os.Stdout, cwd)
	os.Exit(1)
}

func confirmation() bool {
	var response string

	_, err := fmt.Scanln(&response)
	if err != nil {
		log.Fatal(err)
	}
	y := []string{"y", "yes"}
	n := []string{"n", "no"}

	response = strings.TrimSpace(response)
	response = strings.ToLower(response)

	if containsString(y, response) {
		return true
	} else if containsString(n, response) {
		return false
	} else {
		fmt.Println("yes or no required:")
		return confirmation()
	}
}

func posString(slice []string, element string) int {
	for index, elem := range slice {
		if elem == element {
			return index
		}
	}
	return -1
}

// containsString returns true iff slice contains element
func containsString(slice []string, element string) bool {
	return !(posString(slice, element) == -1)
}
<commit_msg>Add more handling for commands that need more args<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"sort"
	"strings"

	"github.com/ryanuber/columnize"
	"github.com/tylerfowle/dtags/db"
)

var database *db.Database

// Entrypoint for dtags application. Initializes database instance and
// processed passed command to perform the expected action.
func main() {
	var err error
	if database, err = db.Init(); err != nil {
		panic(err)
	}
	defer database.Instance.Close()

	if len(os.Args) < 2 || ((os.Args[1] == "add" || os.Args[1] == "del") && len(os.Args) < 3) {
		showHelp()
	}

	switch os.Args[1] {
	case "add":
		addNewTag(os.Args[2:])
		break
	case "del":
		database.DeleteKey(strings.ToLower(os.Args[2:][0]))
		break
	case "list", "completion":
		printAllTags()
		break
	case "ls":
		printBoth()
		break
	default:
		printPath(os.Args[1:])
	}
}

// Show the help message to a user when not enough arguments are passed
// to the command.
func showHelp() {
	fmt.Println("usage: dtags command [params]")
	fmt.Println("       dtags <tag>")
	fmt.Println()
	fmt.Println("       add <tag> [path] - add a new tag at path; defaults to current directory")
	fmt.Println("       del <tag> - delete the provided tag from storage")
	fmt.Println("       list - list all tags currently stored")
	fmt.Println("       ls - list all tags currently stored and associated paths")

	os.Exit(64)
}

func addNewTag(args []string) {
	k := strings.ToLower(args[0])
	v := database.CurrentDirectory

	if len(args[0:]) > 1 {
		v = args[1]
	}

	if database.Exists(k) {
		fmt.Printf("Overwrite existing tag? [%s] (y/n)", k)
		if confirmation() == false {
			return
		}
	}

	database.AddKey(k, v)
}

func printAllTags() {
	for _, tag := range database.GetTags() {
		fmt.Println(tag)
	}
}

func printBoth() {
	var unformattedlist []string
	for tag, path := range database.All() {
		unformattedlist = append(unformattedlist, fmt.Sprintf("%s|%s\n", tag, path))
	}
	sort.Strings(unformattedlist)
	formattedList := columnize.SimpleFormat(unformattedlist)
	fmt.Println(formattedList)
}

func printPath(args []string) {
	cwd := database.GetValue(args[0])
	if cwd == "" {
		fmt.Printf("tag not found\n")
		os.Exit(1)
	}

	fmt.Fprint(os.Stdout, cwd)
	os.Exit(1)
}

func confirmation() bool {
	var response string

	_, err := fmt.Scanln(&response)
	if err != nil {
		log.Fatal(err)
	}
	y := []string{"y", "yes"}
	n := []string{"n", "no"}

	response = strings.TrimSpace(response)
	response = strings.ToLower(response)

	if containsString(y, response) {
		return true
	} else if containsString(n, response) {
		return false
	} else {
		fmt.Println("yes or no required:")
		return confirmation()
	}
}

func posString(slice []string, element string) int {
	for index, elem := range slice {
		if elem == element {
			return index
		}
	}
	return -1
}

// containsString returns true iff slice contains element
func containsString(slice []string, element string) bool {
	return !(posString(slice, element) == -1)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/gorilla/mux"
	"github.com/minotar/minecraft"
	"github.com/op/go-logging"
	"image"
	"net/http"
	"os"
	"strconv"
	"time"
)

const (
	DefaultSize = uint(180)
	MaxSize     = uint(300)
	MinSize     = uint(8)

	SkinCache

	ListenOn = ":9999"

	Minutes            uint = 60
	Hours                   = 60 * Minutes
	Days                    = 24 * Hours
	TimeoutActualSkin       = 2 * Days
	TimeoutFailedFetch      = 15 * Minutes

	MinotarVersion = "2.1"
)

type NotFoundHandler struct{}

func (h NotFoundHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(404)
	fmt.Fprintf(w, "404 not found")
}

func notFoundPage(w http.ResponseWriter, r *http.Request) {
	nfh := NotFoundHandler{}
	nfh.ServeHTTP(w, r)
}
func serverErrorPage(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(500)
	fmt.Fprintf(w, "500 internal server error")
}

func rationalizeSize(inp string) uint {
	out64, err := strconv.ParseUint(inp, 10, 0)
	out := uint(out64)
	if err != nil {
		return DefaultSize
	} else if out > MaxSize {
		return MaxSize
	} else if out < MinSize {
		return MinSize
	}
	return out
}

func addCacheTimeoutHeader(w http.ResponseWriter, timeout uint) {
	w.Header().Add("Cache-Control", fmt.Sprintf("max-age=%d", timeout))
}

func timeBetween(timeA time.Time, timeB time.Time) int64 {
	// millis between two timestamps

	if timeB.Before(timeA) {
		timeA, timeB = timeB, timeA
	}
	return timeB.Sub(timeA).Nanoseconds() / 1000000
}

func fetchImageProcessThen(callback func(minecraft.Skin) (image.Image, error)) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		timeReqStart := time.Now()

		vars := mux.Vars(r)

		username := vars["username"]
		size := rationalizeSize(vars["size"])
		ok := true

		var skin minecraft.Skin
		var err error

		skin = fetchSkin(username)

		timeFetch := time.Now()

		img, err := callback(skin)
		if err != nil {
			serverErrorPage(w, r)
			return
		}
		timeProcess := time.Now()

		imgResized := Resize(size, size, img)
		timeResize := time.Now()

		w.Header().Add("Content-Type", "image/png")
		w.Header().Add("X-Requested", "processed")
		var timeout uint
		if ok {
			w.Header().Add("X-Result", "ok")
			timeout = TimeoutActualSkin
		} else {
			w.Header().Add("X-Result", "failed")
			timeout = TimeoutFailedFetch
		}
		w.Header().Add("X-Timing", fmt.Sprintf("%d+%d+%d=%dms", timeBetween(timeReqStart, timeFetch), timeBetween(timeFetch, timeProcess), timeBetween(timeProcess, timeResize), timeBetween(timeReqStart, timeResize)))
		addCacheTimeoutHeader(w, timeout)
		WritePNG(w, imgResized)
	}
}
func skinPage(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	username := vars["username"]

	skin := fetchSkin(username)

	w.Header().Add("Content-Type", "image/png")
	w.Header().Add("X-Requested", "skin")
	w.Header().Add("X-Result", "ok")

	WritePNG(w, skin.Image)
}
func downloadPage(w http.ResponseWriter, r *http.Request) {
	headers := w.Header()
	headers.Add("Content-Disposition", "attachment; filename=\"skin.png\"")
	skinPage(w, r)
}

func fetchSkin(username string) minecraft.Skin {
	skin, err := minecraft.FetchSkinFromUrl(username)
	if err != nil {
		log.Error("Failed to get skin for " + username + " from Mojang (" + err.Error() + ")")
		skin, _ = minecraft.FetchSkinForChar()
	} else {
		log.Info("Fetched skin for " + username)
	}

	return skin

	/* We're not using this for now due to rate limiting restrictions
	skin, err := minecraft.GetSkin(minecraft.User{Name: username})
	if err != nil {
		// Problem with the returned image, probably means we have an incorrect username
		// Hit the accounts api
		user, err := minecraft.GetUser(username)

		if err != nil {
			// There's no account for this person, serve char
			skin, _ = minecraft.FetchSkinForChar()
		} else {
			// Get valid skin
			skin, err = minecraft.GetSkin(user)
			if err != nil {
				// Their skin somehow errored, fallback
				skin, _ = minecraft.FetchSkinForChar()
			}
		}
	}

	return skin
	*/
}

var log = logging.MustGetLogger("imgd")
var format = "[%{time:15:04:05.000000}] %{level:.4s} %{message}"

func main() {
	logBackend := logging.NewLogBackend(os.Stdout, "", 0)
	logging.SetBackend(logBackend)
	logging.SetFormatter(logging.MustStringFormatter(format))

	avatarPage := fetchImageProcessThen(func(skin minecraft.Skin) (image.Image, error) {
		return GetHead(skin)
	})
	helmPage := fetchImageProcessThen(func(skin minecraft.Skin) (image.Image, error) {
		return GetHelm(skin)
	})

	r := mux.NewRouter()
	r.NotFoundHandler = NotFoundHandler{}

	r.HandleFunc("/avatar/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", avatarPage)
	r.HandleFunc("/avatar/{username:"+minecraft.ValidUsernameRegex+"}/{size:[0-9]+}{extension:(.png)?}", avatarPage)

	r.HandleFunc("/helm/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", helmPage)
	r.HandleFunc("/helm/{username:"+minecraft.ValidUsernameRegex+"}/{size:[0-9]+}{extension:(.png)?}", helmPage)

	r.HandleFunc("/download/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", downloadPage)

	r.HandleFunc("/skin/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", skinPage)

	r.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "%s", MinotarVersion)
	})

	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Location", "https://minotar.net/")
	})

	http.Handle("/", r)
	err := http.ListenAndServe(ListenOn, nil)
	log.Critical(err.Error())
}
<commit_msg>Trying custom gc level and pprof in production<commit_after>package main

import (
	"fmt"
	"github.com/gorilla/mux"
	"github.com/minotar/minecraft"
	"github.com/op/go-logging"
	"image"
	"net/http"
	_ "net/http/pprof"
	"os"
	"runtime/debug"
	"strconv"
	"time"
)

const (
	DefaultSize = uint(180)
	MaxSize     = uint(300)
	MinSize     = uint(8)

	SkinCache

	ListenOn = ":9999"

	Minutes            uint = 60
	Hours                   = 60 * Minutes
	Days                    = 24 * Hours
	TimeoutActualSkin       = 2 * Days
	TimeoutFailedFetch      = 15 * Minutes

	MinotarVersion = "2.1"
)

type NotFoundHandler struct{}

func (h NotFoundHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(404)
	fmt.Fprintf(w, "404 not found")
}

func notFoundPage(w http.ResponseWriter, r *http.Request) {
	nfh := NotFoundHandler{}
	nfh.ServeHTTP(w, r)
}
func serverErrorPage(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(500)
	fmt.Fprintf(w, "500 internal server error")
}

func rationalizeSize(inp string) uint {
	out64, err := strconv.ParseUint(inp, 10, 0)
	out := uint(out64)
	if err != nil {
		return DefaultSize
	} else if out > MaxSize {
		return MaxSize
	} else if out < MinSize {
		return MinSize
	}
	return out
}

func addCacheTimeoutHeader(w http.ResponseWriter, timeout uint) {
	w.Header().Add("Cache-Control", fmt.Sprintf("max-age=%d", timeout))
}

func timeBetween(timeA time.Time, timeB time.Time) int64 {
	// millis between two timestamps

	if timeB.Before(timeA) {
		timeA, timeB = timeB, timeA
	}
	return timeB.Sub(timeA).Nanoseconds() / 1000000
}

func fetchImageProcessThen(callback func(minecraft.Skin) (image.Image, error)) func(w http.ResponseWriter, r *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		timeReqStart := time.Now()

		vars := mux.Vars(r)

		username := vars["username"]
		size := rationalizeSize(vars["size"])
		ok := true

		var skin minecraft.Skin
		var err error

		skin = fetchSkin(username)

		timeFetch := time.Now()

		img, err := callback(skin)
		if err != nil {
			serverErrorPage(w, r)
			return
		}
		timeProcess := time.Now()

		imgResized := Resize(size, size, img)
		timeResize := time.Now()

		w.Header().Add("Content-Type", "image/png")
		w.Header().Add("X-Requested", "processed")
		var timeout uint
		if ok {
			w.Header().Add("X-Result", "ok")
			timeout = TimeoutActualSkin
		} else {
			w.Header().Add("X-Result", "failed")
			timeout = TimeoutFailedFetch
		}
		w.Header().Add("X-Timing", fmt.Sprintf("%d+%d+%d=%dms", timeBetween(timeReqStart, timeFetch), timeBetween(timeFetch, timeProcess), timeBetween(timeProcess, timeResize), timeBetween(timeReqStart, timeResize)))
		addCacheTimeoutHeader(w, timeout)
		WritePNG(w, imgResized)
	}
}
func skinPage(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)

	username := vars["username"]

	skin := fetchSkin(username)

	w.Header().Add("Content-Type", "image/png")
	w.Header().Add("X-Requested", "skin")
	w.Header().Add("X-Result", "ok")

	WritePNG(w, skin.Image)
}
func downloadPage(w http.ResponseWriter, r *http.Request) {
	headers := w.Header()
	headers.Add("Content-Disposition", "attachment; filename=\"skin.png\"")
	skinPage(w, r)
}

func fetchSkin(username string) minecraft.Skin {
	skin, err := minecraft.FetchSkinFromUrl(username)
	if err != nil {
		log.Error("Failed to get skin for " + username + " from Mojang (" + err.Error() + ")")
		skin, _ = minecraft.FetchSkinForChar()
	} else {
		log.Info("Fetched skin for " + username)
	}

	return skin

	/* We're not using this for now due to rate limiting restrictions
	skin, err := minecraft.GetSkin(minecraft.User{Name: username})
	if err != nil {
		// Problem with the returned image, probably means we have an incorrect username
		// Hit the accounts api
		user, err := minecraft.GetUser(username)

		if err != nil {
			// There's no account for this person, serve char
			skin, _ = minecraft.FetchSkinForChar()
		} else {
			// Get valid skin
			skin, err = minecraft.GetSkin(user)
			if err != nil {
				// Their skin somehow errored, fallback
				skin, _ = minecraft.FetchSkinForChar()
			}
		}
	}

	return skin
	*/
}

var log = logging.MustGetLogger("imgd")
var format = "[%{time:15:04:05.000000}] %{level:.4s} %{message}"

func main() {
	logBackend := logging.NewLogBackend(os.Stdout, "", 0)
	logging.SetBackend(logBackend)
	logging.SetFormatter(logging.MustStringFormatter(format))

	debug.SetGCPercent(10)

	avatarPage := fetchImageProcessThen(func(skin minecraft.Skin) (image.Image, error) {
		return GetHead(skin)
	})
	helmPage := fetchImageProcessThen(func(skin minecraft.Skin) (image.Image, error) {
		return GetHelm(skin)
	})

	r := mux.NewRouter()
	r.NotFoundHandler = NotFoundHandler{}

	r.HandleFunc("/avatar/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", avatarPage)
	r.HandleFunc("/avatar/{username:"+minecraft.ValidUsernameRegex+"}/{size:[0-9]+}{extension:(.png)?}", avatarPage)

	r.HandleFunc("/helm/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", helmPage)
	r.HandleFunc("/helm/{username:"+minecraft.ValidUsernameRegex+"}/{size:[0-9]+}{extension:(.png)?}", helmPage)

	r.HandleFunc("/download/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", downloadPage)

	r.HandleFunc("/skin/{username:"+minecraft.ValidUsernameRegex+"}{extension:(.png)?}", skinPage)

	r.HandleFunc("/version", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "%s", MinotarVersion)
	})

	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Location", "https://minotar.net/")
	})

	http.Handle("/", r)
	err := http.ListenAndServe(ListenOn, nil)
	log.Critical(err.Error())
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

import (
	"fmt"
	"math/rand"
	"time"

	"github.com/go-gl/gl"
	glfw "github.com/go-gl/glfw3"
)

var (
	ship             *Ship
	bullets          []*Bullet
	asteroids        []*Asteroid
	explosions       []*Explosion
	lastBulletFired  float64 = -1
	bulletsPerSecond float64 = 5
	gameWidth        float64
	gameHeight       float64
	fieldSize        float64 = 400
	rng                      = rand.New(rand.NewSource(time.Now().UnixNano()))
)

func errorCallback(err glfw.ErrorCode, desc string) {
	fmt.Printf("%v: %v\n", err, desc)
}

func main() {
	glfw.SetErrorCallback(errorCallback)

	if !glfw.Init() {
		panic("can't init glfw!")
	}
	defer glfw.Terminate()

	monitor, err := glfw.GetPrimaryMonitor()
	if err != nil {
		panic(err)
	}

	videomode, err := monitor.GetVideoMode()
	if err != nil {
		panic(err)
	}
	if videomode.Height < 480 || videomode.Width < 640 {
		panic("unsupported resolution!")
	}

	window, err := glfw.CreateWindow(videomode.Width, videomode.Height, "Golang Asteroids!", monitor, nil)
	//window, err := glfw.CreateWindow(640, 480, "Golang Asteroids!", nil, nil)
	if err != nil {
		panic(err)
	}
	defer window.Destroy()

	window.SetKeyCallback(keyCallback)
	window.MakeContextCurrent()

	gl.Init()
	gl.ClearColor(0.0, 0.0, 0.0, 0.0)
	gl.Clear(gl.COLOR_BUFFER_BIT)

	width, height := window.GetFramebufferSize()
	ratio := float64(width) / float64(height)
	gameWidth = ratio * fieldSize
	gameHeight = fieldSize
	gl.Viewport(0, 0, width, height)
	gl.MatrixMode(gl.PROJECTION)
	gl.LoadIdentity()

	gl.Ortho(0, gameWidth, 0, gameHeight, -1.0, 1.0)
	gl.MatrixMode(gl.MODELVIEW)

	// init
	ship = NewShip(gameWidth/2, gameHeight/2, 0, 0.01)
	// create a couple of random asteroids
	for i := float64(1); i <= 7; i++ {
		CreateAsteroid(2+rng.Float64()*8, 3)
	}

	for !window.ShouldClose() {
		// game logic

		//check if objects are still alive
		if !ship.IsAlive() {
			fmt.Println("You lost!")
			window.SetShouldClose(true)
		}

		var bullets2 []*Bullet
		for _, bullet := range bullets {
			if bullet.IsAlive() {
				bullets2 = append(bullets2, bullet)
			}
		}
		bullets = bullets2

		var asteroids2 []*Asteroid
		for _, asteroid := range asteroids {
			if asteroid.IsAlive() {
				asteroids2 = append(asteroids2, asteroid)
			}
		}
		asteroids = asteroids2

		var explosions2 []*Explosion
		for _, explosion := range explosions {
			if explosion.IsAlive() {
				explosions2 = append(explosions2, explosion)
			}
		}
		explosions = explosions2

		if len(asteroids) == 0 {
			fmt.Println("You won!")
			window.SetShouldClose(true)
		}

		// update objects
		ship.Update()
		for _, bullet := range bullets {
			bullet.Update()
		}
		for _, asteroid := range asteroids {
			asteroid.Update()
		}
		for _, explosion := range explosions {
			explosion.Update()
		}

		// hit detection
		for _, asteroid := range asteroids {
			for _, bullet := range bullets {
				if IsColliding(&asteroid.Entity, &bullet.Entity) {
					asteroid.Destroy()
					bullet.Destroy()
				}
			}
			if IsColliding(&asteroid.Entity, &ship.Entity) {
				asteroid.Destroy()
				ship.Destroy()
			}
		}

		// ---------------------------------------------------------------
		// draw calls
		gl.Clear(gl.COLOR_BUFFER_BIT)
		gl.Color3d(1.0, 1.0, 1.0)

		ship.Draw()
		for _, bullet := range bullets {
			bullet.Draw()
		}
		for _, asteroid := range asteroids {
			asteroid.Draw()
		}
		for _, explosion := range explosions {
			explosion.Draw()
		}

		gl.Flush()
		window.SwapBuffers()
		glfw.PollEvents()
	}
}

func keyCallback(window *glfw.Window, key glfw.Key, scancode int, action glfw.Action, mods glfw.ModifierKey) {
	//fmt.Printf("%v, %v, %v, %v\n", key, scancode, action, mods)

	if key == glfw.KeyEscape && action == glfw.Press {
		window.SetShouldClose(true)
	}

	if key == glfw.KeyLeft {
		if action == glfw.Press {
			ship.RotateLeft(true)
		} else if action == glfw.Release {
			ship.RotateLeft(false)
		}
	} else if key == glfw.KeyRight {
		if action == glfw.Press {
			ship.RotateRight(true)
		} else if action == glfw.Release {
			ship.RotateRight(false)
		}
	}

	if key == glfw.KeyUp {
		if action == glfw.Press {
			ship.Accelerate(true)
		} else if action == glfw.Release {
			ship.Accelerate(false)
		}
	} else if key == glfw.KeyDown {
		if action == glfw.Press {
			ship.Decelerate(true)
		} else if action == glfw.Release {
			ship.Decelerate(false)
		}
	}

	if key == glfw.KeySpace && action == glfw.Press && glfw.GetTime() > lastBulletFired+(1/bulletsPerSecond) {
		bullet := ship.Shoot()
		bullets = append(bullets, bullet)
		lastBulletFired = glfw.GetTime()
	}
}
<commit_msg>Alt+Tab -> Switch between fullscreen and windowed<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

import (
	"errors"
	"fmt"
	"math/rand"
	"time"

	"github.com/go-gl/gl"
	glfw "github.com/go-gl/glfw3"
)

var (
	ship             *Ship
	bullets          []*Bullet
	asteroids        []*Asteroid
	explosions       []*Explosion
	lastBulletFired  float64 = -1
	bulletsPerSecond float64 = 5
	gameWidth        float64
	gameHeight       float64
	fieldSize        float64 = 400
	fullscreen       bool    = false
	won              bool    = false
	gameRunning      bool    = true
	altTab           bool    = false
	rng                      = rand.New(rand.NewSource(time.Now().UnixNano()))
)

func errorCallback(err glfw.ErrorCode, desc string) {
	fmt.Printf("%v: %v\n", err, desc)
}

func main() {
	glfw.SetErrorCallback(errorCallback)

	if !glfw.Init() {
		panic("can't init glfw!")
	}
	defer glfw.Terminate()

	var window *glfw.Window = initGame()
	runGameLoop(window)

	if won {
		fmt.Println("You won!")
	} else {
		fmt.Println("You lost!")
	}
}

func keyCallback(window *glfw.Window, key glfw.Key, scancode int, action glfw.Action, mods glfw.ModifierKey) {
	//fmt.Printf("%v, %v, %v, %v\n", key, scancode, action, mods)

	if key == glfw.KeyEscape && action == glfw.Press {
		gameRunning = false
	}

	if key == glfw.KeyLeft {
		if action == glfw.Press {
			ship.RotateLeft(true)
		} else if action == glfw.Release {
			ship.RotateLeft(false)
		}
	} else if key == glfw.KeyRight {
		if action == glfw.Press {
			ship.RotateRight(true)
		} else if action == glfw.Release {
			ship.RotateRight(false)
		}
	}

	if key == glfw.KeyUp {
		if action == glfw.Press {
			ship.Accelerate(true)
		} else if action == glfw.Release {
			ship.Accelerate(false)
		}
	} else if key == glfw.KeyDown {
		if action == glfw.Press {
			ship.Decelerate(true)
		} else if action == glfw.Release {
			ship.Decelerate(false)
		}
	}

	if key == glfw.KeySpace && action == glfw.Press && glfw.GetTime() > lastBulletFired+(1/bulletsPerSecond) {
		bullet := ship.Shoot()
		bullets = append(bullets, bullet)
		lastBulletFired = glfw.GetTime()
	}

	if key == glfw.KeyEnter && action == glfw.Press { //&& mods == glfw.ModAlt {
		altTab = true
	}
}

func initWindow() (window *glfw.Window, err error) {
	monitor, err := glfw.GetPrimaryMonitor()
	if err != nil {
		return nil, err
	}
	videomode, err := monitor.GetVideoMode()
	if err != nil {
		return nil, err
	}
	if videomode.Height < 480 || videomode.Width < 640 {
		return nil, errors.New("unsupported resolution!")
	}

	if fullscreen {
		glfw.WindowHint(glfw.Decorated, 0)
		window, err = glfw.CreateWindow(videomode.Width, videomode.Height, "Golang Asteroids!", nil, nil)
		if err != nil {
			return nil, err
		}
		window.SetPosition(0, 0)
	} else {
		glfw.WindowHint(glfw.Decorated, 1)
		window, err = glfw.CreateWindow(640, 480, "Golang Asteroids!", nil, nil)
		if err != nil {
			return nil, err
		}
		window.SetPosition(videomode.Width/2-320, videomode.Height/2-240)
	}

	window.SetKeyCallback(keyCallback)
	window.MakeContextCurrent()

	gl.Init()
	gl.ClearColor(0.0, 0.0, 0.0, 0.0)
	gl.Clear(gl.COLOR_BUFFER_BIT)

	width, height := window.GetFramebufferSize()
	ratio := float64(width) / float64(height)
	gameWidth = ratio * fieldSize
	gameHeight = fieldSize
	gl.Viewport(0, 0, width, height)
	gl.MatrixMode(gl.PROJECTION)
	gl.LoadIdentity()

	gl.Ortho(0, gameWidth, 0, gameHeight, -1.0, 1.0)
	gl.MatrixMode(gl.MODELVIEW)

	return window, nil
}

func initGame() *glfw.Window {
	window, err := initWindow()
	if err != nil {
		panic(err)
	}

	// init ship
	ship = NewShip(gameWidth/2, gameHeight/2, 0, 0.01)
	// create a couple of random asteroids
	for i := float64(1); i <= 7; i++ {
		CreateAsteroid(2+rng.Float64()*8, 3)
	}

	return window
}

func runGameLoop(window *glfw.Window) {
	for gameRunning {
		//check if objects are still alive
		if !ship.IsAlive() {
			won = false
			gameRunning = false
		}

		var bullets2 []*Bullet
		for _, bullet := range bullets {
			if bullet.IsAlive() {
				bullets2 = append(bullets2, bullet)
			}
		}
		bullets = bullets2

		var asteroids2 []*Asteroid
		for _, asteroid := range asteroids {
			if asteroid.IsAlive() {
				asteroids2 = append(asteroids2, asteroid)
			}
		}
		asteroids = asteroids2

		var explosions2 []*Explosion
		for _, explosion := range explosions {
			if explosion.IsAlive() {
				explosions2 = append(explosions2, explosion)
			}
		}
		explosions = explosions2

		if len(asteroids) == 0 {
			gameRunning = false
		}

		// update objects
		ship.Update()
		for _, bullet := range bullets {
			bullet.Update()
		}
		for _, asteroid := range asteroids {
			asteroid.Update()
		}
		for _, explosion := range explosions {
			explosion.Update()
		}

		// hit detection
		for _, asteroid := range asteroids {
			for _, bullet := range bullets {
				if IsColliding(&asteroid.Entity, &bullet.Entity) {
					asteroid.Destroy()
					bullet.Destroy()
				}
			}
			if IsColliding(&asteroid.Entity, &ship.Entity) {
				asteroid.Destroy()
				ship.Destroy()
			}
		}

		// ---------------------------------------------------------------
		// draw calls
		gl.Clear(gl.COLOR_BUFFER_BIT)
		gl.Color3d(1.0, 1.0, 1.0)

		ship.Draw()
		for _, bullet := range bullets {
			bullet.Draw()
		}
		for _, asteroid := range asteroids {
			asteroid.Draw()
		}
		for _, explosion := range explosions {
			explosion.Draw()
		}

		gl.Flush()
		window.SwapBuffers()
		glfw.PollEvents()

		// switch resolution
		if altTab {
			window.Destroy()

			fullscreen = !fullscreen
			var err error
			window, err = initWindow()
			if err != nil {
				panic(err)
			}

			altTab = false
		}
	}
	window.SetShouldClose(true)
}
<|endoftext|>


<commit_before>package main // import "github.com/gaia-docker/tugbot"

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/gaia-docker/tugbot-common"
	"github.com/gaia-docker/tugbot/actions"
	"github.com/gaia-docker/tugbot/container"
	"github.com/samalba/dockerclient"
)

var (
	client    container.Client
	names     []string
	wgr       sync.WaitGroup
	wgp       sync.WaitGroup
	publisher common.Publisher
)

const (
	// Release version
	Release = "v0.3.0"
)

func init() {
	log.SetLevel(log.InfoLevel)
}

func main() {
	rootCertPath := "/etc/ssl/docker"
	if os.Getenv("DOCKER_CERT_PATH") != "" {
		rootCertPath = os.Getenv("DOCKER_CERT_PATH")
	}

	app := cli.NewApp()
	app.Name = "Tugbot"
	app.Version = Release
	app.Usage = "Tugbot is a continuous testing framework for Docker based environments. Tugbot monitors changes in a runtime environment (host, os, container), runs tests (packaged into Test Containers), when event occured and collects test results."
	app.ArgsUsage = "test containers: name, list of names, or none (for all test containers)"
	app.Before = before
	app.Action = start
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host, H",
			Usage:  "daemon socket to connect to",
			Value:  "unix:///var/run/docker.sock",
			EnvVar: "DOCKER_HOST",
		},
		cli.BoolFlag{
			Name:  "tls",
			Usage: "use TLS; implied by --tlsverify",
		},
		cli.BoolFlag{
			Name:   "tlsverify",
			Usage:  "use TLS and verify the remote",
			EnvVar: "DOCKER_TLS_VERIFY",
		},
		cli.StringFlag{
			Name:  "tlscacert",
			Usage: "trust certs signed only by this CA",
			Value: fmt.Sprintf("%s/ca.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlscert",
			Usage: "client certificate for TLS authentication",
			Value: fmt.Sprintf("%s/cert.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlskey",
			Usage: "client key for TLS authentication",
			Value: fmt.Sprintf("%s/key.pem", rootCertPath),
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug mode with verbose logging",
		},
		cli.StringFlag{
			Name:   "webhooks",
			Usage:  "list of urls sperated by ';'",
			Value:  "http://result-service:8081/events",
			EnvVar: "TUGBOT_WEBHOOKS",
		},
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func before(c *cli.Context) error {
	if c.GlobalBool("debug") {
		log.SetLevel(log.DebugLevel)
	}

	// Set-up container client
	tls, err := tlsConfig(c)
	if err != nil {
		return err
	}

	client = container.NewClient(c.GlobalString("host"), tls, !c.GlobalBool("no-pull"))

	return nil
}

func start(c *cli.Context) {
	names = c.Args()
	startMonitorEvents(c)
	log.Infof("Tugbot Started. Debug: %v, Webhooks: %v", c.GlobalBool("debug"), c.GlobalBool("webhooks"))
	waitForInterrupt()
}

func startMonitorEvents(c *cli.Context) {
	client.StartMonitorEvents(runTestContainers)
	webhooks := c.GlobalString("webhooks")
	if webhooks != "" {
		publisher = common.NewPublisher(strings.Split(webhooks, ";"))
		client.StartMonitorEvents(publishEvent)
	}
}

func runTestContainers(e *dockerclient.Event, ec chan error, args ...interface{}) {
	log.Infof("Looking for test containers that should run on event: %+v", e)
	wgr.Add(1)
	if err := actions.Run(client, names, e); err != nil {
		log.Error(err)
	}
	wgr.Done()
}

func publishEvent(e *dockerclient.Event, ec chan error, args ...interface{}) {
	//log.Debugf("Publishing event: %+v", e)
	log.Infof("Publishing event: %+v", e)
	wgp.Add(1)
	publisher.Publish(e)
	wgp.Done()
}

func waitForInterrupt() {
	// Graceful shut-down on SIGINT/SIGTERM/SIGQUIT
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT)
	<-c
	log.Debug("Stop monitoring events ...")
	wgr.Wait()
	client.StopAllMonitorEvents()
	log.Debug("Graceful exit :-)")
	os.Exit(1)
}

// tlsConfig translates the command-line options into a tls.Config struct
func tlsConfig(c *cli.Context) (*tls.Config, error) {
	var tlsConfig *tls.Config
	var err error
	caCertFlag := c.GlobalString("tlscacert")
	certFlag := c.GlobalString("tlscert")
	keyFlag := c.GlobalString("tlskey")

	if c.GlobalBool("tls") || c.GlobalBool("tlsverify") {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: !c.GlobalBool("tlsverify"),
		}

		// Load CA cert
		if caCertFlag != "" {
			var caCert []byte

			if strings.HasPrefix(caCertFlag, "/") {
				caCert, err = ioutil.ReadFile(caCertFlag)
				if err != nil {
					return nil, err
				}
			} else {
				caCert = []byte(caCertFlag)
			}

			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)

			tlsConfig.RootCAs = caCertPool
		}

		// Load client certificate
		if certFlag != "" && keyFlag != "" {
			var cert tls.Certificate

			if strings.HasPrefix(certFlag, "/") && strings.HasPrefix(keyFlag, "/") {
				cert, err = tls.LoadX509KeyPair(certFlag, keyFlag)
				if err != nil {
					return nil, err
				}
			} else {
				cert, err = tls.X509KeyPair([]byte(certFlag), []byte(keyFlag))
				if err != nil {
					return nil, err
				}
			}
			tlsConfig.Certificates = []tls.Certificate{cert}
		}
	}

	return tlsConfig, nil
}
<commit_msg>starting ticker on startup; wait to publisher on shut-down;<commit_after>package main // import "github.com/gaia-docker/tugbot"

import (
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/gaia-docker/tugbot-common"
	"github.com/gaia-docker/tugbot/actions"
	"github.com/gaia-docker/tugbot/container"
	"github.com/samalba/dockerclient"

	"crypto/tls"
	"crypto/x509"
	"fmt"
	"golang.org/x/net/context"
	"io/ioutil"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"
)

var (
	client       container.Client
	names        []string
	publisher    common.Publisher
	wgr          sync.WaitGroup
	wgp          sync.WaitGroup
	wgt          sync.WaitGroup
	tickerCancel context.CancelFunc
)

const (
	// Release version
	Release = "v0.3.0"
)

func init() {
	log.SetLevel(log.InfoLevel)
}

func main() {
	rootCertPath := "/etc/ssl/docker"
	if os.Getenv("DOCKER_CERT_PATH") != "" {
		rootCertPath = os.Getenv("DOCKER_CERT_PATH")
	}

	app := cli.NewApp()
	app.Name = "Tugbot"
	app.Version = Release
	app.Usage = "Tugbot is a continuous testing framework for Docker based environments. Tugbot monitors changes in a runtime environment (host, os, container), runs tests (packaged into Test Containers), when event occured and collects test results."
	app.ArgsUsage = "test containers: name, list of names, or none (for all test containers)"
	app.Before = before
	app.Action = start
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "host, H",
			Usage:  "daemon socket to connect to",
			Value:  "unix:///var/run/docker.sock",
			EnvVar: "DOCKER_HOST",
		},
		cli.BoolFlag{
			Name:  "tls",
			Usage: "use TLS; implied by --tlsverify",
		},
		cli.BoolFlag{
			Name:   "tlsverify",
			Usage:  "use TLS and verify the remote",
			EnvVar: "DOCKER_TLS_VERIFY",
		},
		cli.StringFlag{
			Name:  "tlscacert",
			Usage: "trust certs signed only by this CA",
			Value: fmt.Sprintf("%s/ca.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlscert",
			Usage: "client certificate for TLS authentication",
			Value: fmt.Sprintf("%s/cert.pem", rootCertPath),
		},
		cli.StringFlag{
			Name:  "tlskey",
			Usage: "client key for TLS authentication",
			Value: fmt.Sprintf("%s/key.pem", rootCertPath),
		},
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug mode with verbose logging",
		},
		cli.StringFlag{
			Name:   "webhooks",
			Usage:  "list of urls sperated by ';'",
			Value:  "http://result-service:8081/events",
			EnvVar: "TUGBOT_WEBHOOKS",
		},
	}

	if err := app.Run(os.Args); err != nil {
		log.Fatal(err)
	}
}

func before(c *cli.Context) error {
	if c.GlobalBool("debug") {
		log.SetLevel(log.DebugLevel)
	}

	// Set-up container client
	tls, err := tlsConfig(c)
	if err != nil {
		return err
	}
	client = container.NewClient(c.GlobalString("host"), tls, !c.GlobalBool("no-pull"))

	return nil
}

func start(c *cli.Context) {
	names = c.Args()
	startMonitorEvents(c)
	startTicker()
	log.Infof("Tugbot Started. Debug: %v, Webhooks: %v", c.GlobalBool("debug"), c.GlobalBool("webhooks"))
	waitForInterrupt()
}

func startMonitorEvents(c *cli.Context) {
	client.StartMonitorEvents(runTestContainers)
	webhooks := c.GlobalString("webhooks")
	if webhooks != "" {
		publisher = common.NewPublisher(strings.Split(webhooks, ";"))
		client.StartMonitorEvents(publishEvent)
	}
}

func startTicker() {
	wgt.Add(1)
	var ctx context.Context
	ctx, tickerCancel = context.WithCancel(context.Background())
	go func() {
		actions.RunTickerTestContainers(ctx, client, time.Second*18)
		wgt.Done()
	}()
}
func runTestContainers(e *dockerclient.Event, ec chan error, args ...interface{}) {
	log.Debugf("Looking for test containers that should run on event: %+v", e)
	wgr.Add(1)
	if err := actions.Run(client, names, e); err != nil {
		log.Error(err)
	}
	wgr.Done()
}

func publishEvent(e *dockerclient.Event, ec chan error, args ...interface{}) {
	log.Debugf("Publishing event: %+v", e)
	wgp.Add(1)
	publisher.Publish(e)
	wgp.Done()
}

func waitForInterrupt() {
	// Graceful shut-down on SIGINT/SIGTERM/SIGQUIT
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT)
	<-c
	wgr.Wait()
	log.Info("Stoping monitor events...")
	client.StopAllMonitorEvents()
	wgp.Wait()
	log.Info("Stoping ticker...")
	tickerCancel()
	wgt.Wait()
	log.Debug("Graceful exit :-)")
	os.Exit(1)
}

// tlsConfig translates the command-line options into a tls.Config struct
func tlsConfig(c *cli.Context) (*tls.Config, error) {
	var tlsConfig *tls.Config
	var err error
	caCertFlag := c.GlobalString("tlscacert")
	certFlag := c.GlobalString("tlscert")
	keyFlag := c.GlobalString("tlskey")

	if c.GlobalBool("tls") || c.GlobalBool("tlsverify") {
		tlsConfig = &tls.Config{
			InsecureSkipVerify: !c.GlobalBool("tlsverify"),
		}

		// Load CA cert
		if caCertFlag != "" {
			var caCert []byte

			if strings.HasPrefix(caCertFlag, "/") {
				caCert, err = ioutil.ReadFile(caCertFlag)
				if err != nil {
					return nil, err
				}
			} else {
				caCert = []byte(caCertFlag)
			}

			caCertPool := x509.NewCertPool()
			caCertPool.AppendCertsFromPEM(caCert)

			tlsConfig.RootCAs = caCertPool
		}

		// Load client certificate
		if certFlag != "" && keyFlag != "" {
			var cert tls.Certificate

			if strings.HasPrefix(certFlag, "/") && strings.HasPrefix(keyFlag, "/") {
				cert, err = tls.LoadX509KeyPair(certFlag, keyFlag)
				if err != nil {
					return nil, err
				}
			} else {
				cert, err = tls.X509KeyPair([]byte(certFlag), []byte(keyFlag))
				if err != nil {
					return nil, err
				}
			}
			tlsConfig.Certificates = []tls.Certificate{cert}
		}
	}

	return tlsConfig, nil
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2017, UPMC Enterprises
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name UPMC Enterprises nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL UPMC ENTERPRISES BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
*/

package main

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecr"
	flag "github.com/spf13/pflag"
	"github.com/upmc-enterprises/registry-creds/k8sutil"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"k8s.io/client-go/pkg/api/v1"
)

const (
	dockerCfgTemplate  = `{"%s":{"username":"oauth2accesstoken","password":"%s","email":"none"}}`
	dockerJSONTemplate = `{"auths":{"%s":{"auth":"%s","email":"none"}}}`
)

var (
	flags               = flag.NewFlagSet("", flag.ContinueOnError)
	argKubecfgFile      = flags.String("kubecfg-file", "", `Location of kubecfg file for access to kubernetes master service; --kube_master_url overrides the URL part of this; if neither this nor --kube_master_url are provided, defaults to service account tokens`)
	argKubeMasterURL    = flags.String("kube-master-url", "", `URL to reach kubernetes master. Env variables in this flag will be expanded.`)
	argAWSSecretName    = flags.String("aws-secret-name", "awsecr-cred", `Default aws secret name`)
	argGCRSecretName    = flags.String("gcr-secret-name", "gcr-secret", `Default gcr secret name`)
	argDefaultNamespace = flags.String("default-namespace", "default", `Default namespace`)
	argGCRURL           = flags.String("gcr-url", "https://gcr.io", `Default GCR URL`)
	argAWSRegion        = flags.String("aws-region", "us-east-1", `Default AWS region`)
	argRefreshMinutes   = flags.Int("refresh-mins", 60, `Default time to wait before refreshing (60 minutes)`)
	argSkipKubeSystem   = flags.Bool("skip-kube-system", true, `If true, will not attempt to set ImagePullSecrets on the kube-system namespace`)
)

var (
	awsAccountID string
)

type controller struct {
	k8sutil   *k8sutil.K8sutilInterface
	ecrClient ecrInterface
	gcrClient gcrInterface
}

type ecrInterface interface {
	GetAuthorizationToken(input *ecr.GetAuthorizationTokenInput) (*ecr.GetAuthorizationTokenOutput, error)
}

type gcrInterface interface {
	DefaultTokenSource(ctx context.Context, scope ...string) (oauth2.TokenSource, error)
}

func newEcrClient() ecrInterface {
	return ecr.New(session.New(), aws.NewConfig().WithRegion(*argAWSRegion))
}

type gcrClient struct{}

func (gcr gcrClient) DefaultTokenSource(ctx context.Context, scope ...string) (oauth2.TokenSource, error) {
	return google.DefaultTokenSource(ctx, scope...)
}

func newGcrClient() gcrInterface {
	return gcrClient{}
}

func (c *controller) getGCRAuthorizationKey() (AuthToken, error) {
	ts, err := c.gcrClient.DefaultTokenSource(context.TODO(), "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		return AuthToken{}, err
	}

	token, err := ts.Token()
	if err != nil {
		return AuthToken{}, err
	}

	if !token.Valid() {
		return AuthToken{}, fmt.Errorf("token was invalid")
	}

	if token.Type() != "Bearer" {
		return AuthToken{}, fmt.Errorf(fmt.Sprintf("expected token type \"Bearer\" but got \"%s\"", token.Type()))
	}

	return AuthToken{
		AccessToken: token.AccessToken,
		Endpoint:    *argGCRURL}, nil
}

func (c *controller) getECRAuthorizationKey() (AuthToken, error) {
	params := &ecr.GetAuthorizationTokenInput{
		RegistryIds: []*string{
			aws.String(awsAccountID),
		},
	}

	resp, err := c.ecrClient.GetAuthorizationToken(params)

	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		logrus.Println(err.Error())
		return AuthToken{}, err
	}

	token := resp.AuthorizationData[0]

	return AuthToken{
		AccessToken: *token.AuthorizationToken,
		Endpoint:    *token.ProxyEndpoint}, err
}

func generateSecretObj(token string, endpoint string, isJSONCfg bool, secretName string) *v1.Secret {
	secret := &v1.Secret{
		ObjectMeta: v1.ObjectMeta{
			Name: secretName,
		},
	}
	if isJSONCfg {
		secret.Data = map[string][]byte{
			".dockerconfigjson": []byte(fmt.Sprintf(dockerJSONTemplate, endpoint, token))}
		secret.Type = "kubernetes.io/dockerconfigjson"
	} else {
		secret.Data = map[string][]byte{
			".dockercfg": []byte(fmt.Sprintf(dockerCfgTemplate, endpoint, token))}
		secret.Type = "kubernetes.io/dockercfg"
	}
	return secret
}

type AuthToken struct {
	AccessToken string
	Endpoint    string
}

type SecretGenerator struct {
	TokenGenFxn func() (AuthToken, error)
	IsJSONCfg   bool
	SecretName  string
}

func getSecretGenerators(c *controller) []SecretGenerator {
	secretGenerators := []SecretGenerator{}

	secretGenerators = append(secretGenerators, SecretGenerator{
		TokenGenFxn: c.getGCRAuthorizationKey,
		IsJSONCfg:   false,
		SecretName:  *argGCRSecretName,
	})

	secretGenerators = append(secretGenerators, SecretGenerator{
		TokenGenFxn: c.getECRAuthorizationKey,
		IsJSONCfg:   true,
		SecretName:  *argAWSSecretName,
	})

	return secretGenerators
}

func (c *controller) processNamespace(namespace *v1.Namespace, secret *v1.Secret) error {
	// Check if the secret exists for the namespace
	_, err := c.k8sutil.GetSecret(namespace.GetName(), secret.Name)

	if err != nil {
		// Secret not found, create
		err := c.k8sutil.CreateSecret(namespace.GetName(), secret)
		if err != nil {
			return fmt.Errorf("Could not create Secret! %v", err)
		}
	} else {
		// Existing secret needs updated
		err := c.k8sutil.UpdateSecret(namespace.GetName(), secret)
		if err != nil {
			return fmt.Errorf("Could not update Secret! %v", err)
		}
	}

	// Check if ServiceAccount exists
	serviceAccount, err := c.k8sutil.GetServiceAccount(namespace.GetName(), "default")
	if err != nil {
		fmt.Errorf("Could not get ServiceAccounts! %v", err)
		return err
	}

	// Update existing one if image pull secrets already exists for aws ecr token
	imagePullSecretFound := false
	for i, imagePullSecret := range serviceAccount.ImagePullSecrets {
		if imagePullSecret.Name == secret.Name {
			serviceAccount.ImagePullSecrets[i] = v1.LocalObjectReference{Name: secret.Name}
			imagePullSecretFound = true
			break
		}
	}

	// Append to list of existing service accounts if there isn't one already
	if !imagePullSecretFound {
		serviceAccount.ImagePullSecrets = append(serviceAccount.ImagePullSecrets, v1.LocalObjectReference{Name: secret.Name})
	}

	err = c.k8sutil.UpdateServiceAccount(namespace.GetName(), serviceAccount)
	if err != nil {
		return fmt.Errorf("Could not update ServiceAccount! %v", err)
	}

	return nil
}

func (c *controller) generateSecrets() []*v1.Secret {
	var secrets []*v1.Secret
	secretGenerators := getSecretGenerators(c)

	for _, secretGenerator := range secretGenerators {
		logrus.Printf("------------------ [%s] ----------------------\n", secretGenerator.SecretName)

		newToken, err := secretGenerator.TokenGenFxn()
		if err != nil {
			logrus.Printf("Error getting secret for provider %s. Skipping secret provider! [Err: %s]", secretGenerator.SecretName, err)
			continue
		}
		newSecret := generateSecretObj(newToken.AccessToken, newToken.Endpoint, secretGenerator.IsJSONCfg, secretGenerator.SecretName)
		secrets = append(secrets, newSecret)
	}
	return secrets
}

func validateParams() {
	// Allow environment variables to overwrite args
	awsAccountIDEnv := os.Getenv("awsaccount")
	awsRegionEnv := os.Getenv("awsregion")

	if len(awsRegionEnv) > 0 {
		argAWSRegion = &awsRegionEnv
	}

	if len(awsAccountIDEnv) > 0 {
		awsAccountID = awsAccountIDEnv
	}
}

func handler(c *controller, ns *v1.Namespace) error {
	log.Print("Refreshing credentials...")
	secrets := c.generateSecrets()
	for _, secret := range secrets {
		if *argSkipKubeSystem && ns.GetName() == "kube-system" {
			continue
		}

		if err := c.processNamespace(ns, secret); err != nil {
			return err
		}

		log.Printf("Finished processing secret for namespace %s, secret %s", ns.Name, secret.Name)
	}
	return nil
}

func main() {
	log.Print("Starting up...")
	flags.Parse(os.Args)

	validateParams()

	log.Print("Using AWS Account: ", awsAccountID)
	log.Printf("Using AWS Region: %s", *argAWSRegion)
	log.Print("Refresh Interval (minutes): ", *argRefreshMinutes)

	util, err := k8sutil.New(*argKubecfgFile, *argKubeMasterURL)

	if err != nil {
		logrus.Error("Could not create k8s client!!", err)
	}

	ecrClient := newEcrClient()
	gcrClient := newGcrClient()
	c := &controller{util, ecrClient, gcrClient}

	util.WatchNamespaces(time.Duration(*argRefreshMinutes)*time.Minute, func(ns *v1.Namespace) error {
		return handler(c, ns)
	})
}
<commit_msg>return the formatted error<commit_after>/*
Copyright (c) 2017, UPMC Enterprises
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name UPMC Enterprises nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL UPMC ENTERPRISES BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
*/

package main

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ecr"
	flag "github.com/spf13/pflag"
	"github.com/upmc-enterprises/registry-creds/k8sutil"
	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"k8s.io/client-go/pkg/api/v1"
)

const (
	dockerCfgTemplate  = `{"%s":{"username":"oauth2accesstoken","password":"%s","email":"none"}}`
	dockerJSONTemplate = `{"auths":{"%s":{"auth":"%s","email":"none"}}}`
)

var (
	flags               = flag.NewFlagSet("", flag.ContinueOnError)
	argKubecfgFile      = flags.String("kubecfg-file", "", `Location of kubecfg file for access to kubernetes master service; --kube_master_url overrides the URL part of this; if neither this nor --kube_master_url are provided, defaults to service account tokens`)
	argKubeMasterURL    = flags.String("kube-master-url", "", `URL to reach kubernetes master. Env variables in this flag will be expanded.`)
	argAWSSecretName    = flags.String("aws-secret-name", "awsecr-cred", `Default aws secret name`)
	argGCRSecretName    = flags.String("gcr-secret-name", "gcr-secret", `Default gcr secret name`)
	argDefaultNamespace = flags.String("default-namespace", "default", `Default namespace`)
	argGCRURL           = flags.String("gcr-url", "https://gcr.io", `Default GCR URL`)
	argAWSRegion        = flags.String("aws-region", "us-east-1", `Default AWS region`)
	argRefreshMinutes   = flags.Int("refresh-mins", 60, `Default time to wait before refreshing (60 minutes)`)
	argSkipKubeSystem   = flags.Bool("skip-kube-system", true, `If true, will not attempt to set ImagePullSecrets on the kube-system namespace`)
)

var (
	awsAccountID string
)

type controller struct {
	k8sutil   *k8sutil.K8sutilInterface
	ecrClient ecrInterface
	gcrClient gcrInterface
}

type ecrInterface interface {
	GetAuthorizationToken(input *ecr.GetAuthorizationTokenInput) (*ecr.GetAuthorizationTokenOutput, error)
}

type gcrInterface interface {
	DefaultTokenSource(ctx context.Context, scope ...string) (oauth2.TokenSource, error)
}

func newEcrClient() ecrInterface {
	return ecr.New(session.New(), aws.NewConfig().WithRegion(*argAWSRegion))
}

type gcrClient struct{}

func (gcr gcrClient) DefaultTokenSource(ctx context.Context, scope ...string) (oauth2.TokenSource, error) {
	return google.DefaultTokenSource(ctx, scope...)
}

func newGcrClient() gcrInterface {
	return gcrClient{}
}

func (c *controller) getGCRAuthorizationKey() (AuthToken, error) {
	ts, err := c.gcrClient.DefaultTokenSource(context.TODO(), "https://www.googleapis.com/auth/cloud-platform")
	if err != nil {
		return AuthToken{}, err
	}

	token, err := ts.Token()
	if err != nil {
		return AuthToken{}, err
	}

	if !token.Valid() {
		return AuthToken{}, fmt.Errorf("token was invalid")
	}

	if token.Type() != "Bearer" {
		return AuthToken{}, fmt.Errorf(fmt.Sprintf("expected token type \"Bearer\" but got \"%s\"", token.Type()))
	}

	return AuthToken{
		AccessToken: token.AccessToken,
		Endpoint:    *argGCRURL}, nil
}

func (c *controller) getECRAuthorizationKey() (AuthToken, error) {
	params := &ecr.GetAuthorizationTokenInput{
		RegistryIds: []*string{
			aws.String(awsAccountID),
		},
	}

	resp, err := c.ecrClient.GetAuthorizationToken(params)

	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		logrus.Println(err.Error())
		return AuthToken{}, err
	}

	token := resp.AuthorizationData[0]

	return AuthToken{
		AccessToken: *token.AuthorizationToken,
		Endpoint:    *token.ProxyEndpoint}, err
}

func generateSecretObj(token string, endpoint string, isJSONCfg bool, secretName string) *v1.Secret {
	secret := &v1.Secret{
		ObjectMeta: v1.ObjectMeta{
			Name: secretName,
		},
	}
	if isJSONCfg {
		secret.Data = map[string][]byte{
			".dockerconfigjson": []byte(fmt.Sprintf(dockerJSONTemplate, endpoint, token))}
		secret.Type = "kubernetes.io/dockerconfigjson"
	} else {
		secret.Data = map[string][]byte{
			".dockercfg": []byte(fmt.Sprintf(dockerCfgTemplate, endpoint, token))}
		secret.Type = "kubernetes.io/dockercfg"
	}
	return secret
}

type AuthToken struct {
	AccessToken string
	Endpoint    string
}

type SecretGenerator struct {
	TokenGenFxn func() (AuthToken, error)
	IsJSONCfg   bool
	SecretName  string
}

func getSecretGenerators(c *controller) []SecretGenerator {
	secretGenerators := []SecretGenerator{}

	secretGenerators = append(secretGenerators, SecretGenerator{
		TokenGenFxn: c.getGCRAuthorizationKey,
		IsJSONCfg:   false,
		SecretName:  *argGCRSecretName,
	})

	secretGenerators = append(secretGenerators, SecretGenerator{
		TokenGenFxn: c.getECRAuthorizationKey,
		IsJSONCfg:   true,
		SecretName:  *argAWSSecretName,
	})

	return secretGenerators
}

func (c *controller) processNamespace(namespace *v1.Namespace, secret *v1.Secret) error {
	// Check if the secret exists for the namespace
	_, err := c.k8sutil.GetSecret(namespace.GetName(), secret.Name)

	if err != nil {
		// Secret not found, create
		err := c.k8sutil.CreateSecret(namespace.GetName(), secret)
		if err != nil {
			return fmt.Errorf("Could not create Secret! %v", err)
		}
	} else {
		// Existing secret needs updated
		err := c.k8sutil.UpdateSecret(namespace.GetName(), secret)
		if err != nil {
			return fmt.Errorf("Could not update Secret! %v", err)
		}
	}

	// Check if ServiceAccount exists
	serviceAccount, err := c.k8sutil.GetServiceAccount(namespace.GetName(), "default")
	if err != nil {
		return fmt.Errorf("Could not get ServiceAccounts! %v", err)
	}

	// Update existing one if image pull secrets already exists for aws ecr token
	imagePullSecretFound := false
	for i, imagePullSecret := range serviceAccount.ImagePullSecrets {
		if imagePullSecret.Name == secret.Name {
			serviceAccount.ImagePullSecrets[i] = v1.LocalObjectReference{Name: secret.Name}
			imagePullSecretFound = true
			break
		}
	}

	// Append to list of existing service accounts if there isn't one already
	if !imagePullSecretFound {
		serviceAccount.ImagePullSecrets = append(serviceAccount.ImagePullSecrets, v1.LocalObjectReference{Name: secret.Name})
	}

	err = c.k8sutil.UpdateServiceAccount(namespace.GetName(), serviceAccount)
	if err != nil {
		return fmt.Errorf("Could not update ServiceAccount! %v", err)
	}

	return nil
}

func (c *controller) generateSecrets() []*v1.Secret {
	var secrets []*v1.Secret
	secretGenerators := getSecretGenerators(c)

	for _, secretGenerator := range secretGenerators {
		logrus.Printf("------------------ [%s] ----------------------\n", secretGenerator.SecretName)

		newToken, err := secretGenerator.TokenGenFxn()
		if err != nil {
			logrus.Printf("Error getting secret for provider %s. Skipping secret provider! [Err: %s]", secretGenerator.SecretName, err)
			continue
		}
		newSecret := generateSecretObj(newToken.AccessToken, newToken.Endpoint, secretGenerator.IsJSONCfg, secretGenerator.SecretName)
		secrets = append(secrets, newSecret)
	}
	return secrets
}

func validateParams() {
	// Allow environment variables to overwrite args
	awsAccountIDEnv := os.Getenv("awsaccount")
	awsRegionEnv := os.Getenv("awsregion")

	if len(awsRegionEnv) > 0 {
		argAWSRegion = &awsRegionEnv
	}

	if len(awsAccountIDEnv) > 0 {
		awsAccountID = awsAccountIDEnv
	}
}

func handler(c *controller, ns *v1.Namespace) error {
	log.Print("Refreshing credentials...")
	secrets := c.generateSecrets()
	for _, secret := range secrets {
		if *argSkipKubeSystem && ns.GetName() == "kube-system" {
			continue
		}

		if err := c.processNamespace(ns, secret); err != nil {
			return err
		}

		log.Printf("Finished processing secret for namespace %s, secret %s", ns.Name, secret.Name)
	}
	return nil
}

func main() {
	log.Print("Starting up...")
	flags.Parse(os.Args)

	validateParams()

	log.Print("Using AWS Account: ", awsAccountID)
	log.Printf("Using AWS Region: %s", *argAWSRegion)
	log.Print("Refresh Interval (minutes): ", *argRefreshMinutes)

	util, err := k8sutil.New(*argKubecfgFile, *argKubeMasterURL)

	if err != nil {
		logrus.Error("Could not create k8s client!!", err)
	}

	ecrClient := newEcrClient()
	gcrClient := newGcrClient()
	c := &controller{util, ecrClient, gcrClient}

	util.WatchNamespaces(time.Duration(*argRefreshMinutes)*time.Minute, func(ns *v1.Namespace) error {
		return handler(c, ns)
	})
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Dave Collins
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"encoding/csv"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"math/big"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrd/chaincfg/chainhash"
	"github.com/decred/dcrd/wire"
	"github.com/decred/dcrutil"
)

const (
	// fieldsPerRecord defines the number of fields expected in each line
	// of the input CSV data.
	fieldsPerRecord = 3
)

// convertRecord converts the passed record, which is expected to be parsed from
// a CSV file, and thus will be a slice of strings, into a struct with concrete
// types.
func convertRecord(record []string) (*simData, error) {
	headerBytes, err := hex.DecodeString(record[1])
	if err != nil {
		return nil, err
	}

	var header wire.BlockHeader
	if err := header.FromBytes(headerBytes); err != nil {
		return nil, err
	}
	var hashStrings []string
	if record[2] != "" {
		hashStrings = strings.Split(record[2], ":")
	}
	if len(hashStrings) != int(header.FreshStake) {
		return nil, fmt.Errorf("%d ticket hashes in CSV for %d new tickets",
			len(hashStrings), header.FreshStake)
	}
	ticketHashes := make([]chainhash.Hash, 0, len(hashStrings))
	for _, hashString := range hashStrings {
		hash, err := chainhash.NewHashFromStr(hashString)
		if err != nil {
			return nil, err
		}
		ticketHashes = append(ticketHashes, *hash)
	}

	return &simData{
		header:       headerBytes,
		voters:       header.Voters,
		prevValid:    dcrutil.IsFlagSet16(header.VoteBits, dcrutil.BlockValid),
		newTickets:   header.FreshStake,
		ticketHashes: ticketHashes,
		revocations:  uint16(header.Revocations),
	}, nil
}

// reportProgress periodically prints out the current simulator height to
// stdout.
func (s *simulator) reportProgress() {
	if s.tip.height%10000 == 0 && s.tip.height != 0 {
		fmt.Println()
	}
	if s.tip.height%1000 == 0 && s.tip.height != 0 {
		fmt.Printf("..%d", s.tip.height)
	}
}

// simulateFromCSV runs the simulation using input data from a CSV file.  It is
// realistically only intended to be used with data extracted from mainnet in
// order to exactly replicate its live ticket pool.
func (s *simulator) simulateFromCSV(csvPath string) error {
	// Open the simulation CSV data which is expected to be in the following
	// format:
	//
	// Block Header,Winning Ticket Hashes
	csvFile, err := os.Open(csvPath)
	if err != nil {
		return err
	}

	// Create a new simulator using input from the CSV file.
	r := csv.NewReader(csvFile)
	r.FieldsPerRecord = fieldsPerRecord
	var handledHeader bool
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		// Skip header fields if they exist.
		if !handledHeader {
			handledHeader = true
			_, err := strconv.Atoi(record[0])
			if err != nil {
				continue
			}
		}

		// Convert the CSV to concrete data.
		data, err := convertRecord(record)
		if err != nil {
			return err
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

// calcYieldDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// estimated yield purchasing a ticket would produce.
func calcYieldDemand(ticketPrice, perVoteSubsidy int64) float64 {
	// 100% demand when the yield is over 5%.
	yield := float64(perVoteSubsidy) / float64(ticketPrice)
	if yield > 0.05 {
		return 1.0
	}

	// No demand when the yield is under 2%.
	if yield < 0.02 {
		return 0.0
	}

	// The yield is in between 2% and 5%, so create a linear demand
	// accordingly.
	return (yield - 0.02) / 0.03
}

// calcVWAPDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// volume-weighted average ticket purchase of the previous ticket price windows.
func calcVWAPDemand(ticketPrice, ticketVWAP int64) float64 {
	// 100% demand when the ticket price is under 80% of the VWAP.
	eightyPercentVWAP := (ticketVWAP * 8) / 10
	if ticketPrice < eightyPercentVWAP {
		return 1.0
	}

	// No demand when the ticket price is over 120% of the VWAP.
	if ticketPrice > (ticketVWAP*12)/10 {
		return 0.0
	}

	// The ticket price is in between 80% and 120% of the VWAP, so create
	// a linear demand accordingly.
	fortyPercentVWAP := (ticketVWAP * 4) / 10
	return 1 - float64(ticketPrice-eightyPercentVWAP)/float64(fortyPercentVWAP)
}

// calcVWAP calculates and return the volume-weighted average ticket purchase
// price for up to 'StakeDiffWindows' worth of the previous ticket price
// windows.
func (s *simulator) calcPrevVWAP(prevNode *blockNode) int64 {
	windowSize := int32(s.params.StakeDiffWindowSize)
	stakeDiffWindows := int32(s.params.StakeDiffWindows)

	// Calculate the height the block just before the most recent ticket
	// price change.
	wantHeight := prevNode.height - (prevNode.height+1)%windowSize
	prevNode = s.ancestorNode(prevNode, wantHeight, nil)

	// Loop through previous required number of previous blocks and tally up
	// all of the weighted ticket purchase prices as well as the total
	// number of purchased tickets.
	numTickets, weightedVal := new(big.Int), new(big.Int)
	weightedSum, totalTickets := new(big.Int), new(big.Int)
	blocksToIterate := stakeDiffWindows * windowSize
	for i := int32(0); i < blocksToIterate && prevNode != nil; i++ {
		// weightedSum += numTickets*ticketPrice
		// totalTickets += numTickets
		numTickets.SetInt64(int64(len(prevNode.ticketsAdded)))
		weightedVal.SetInt64(prevNode.ticketPrice)
		weightedVal.Mul(weightedVal, numTickets)
		weightedSum.Add(weightedSum, weightedVal)
		totalTickets.Add(totalTickets, numTickets)
		prevNode = prevNode.parent
	}

	// Return minimum ticket price if there are were not any ticket
	// purchases at all in the entire period being examined.
	if totalTickets.Sign() == 0 {
		return s.params.MinimumStakeDiff
	}

	return new(big.Int).Div(weightedSum, totalTickets).Int64()
}

// calcDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval).
func (s *simulator) calcDemand(nextHeight int32, ticketPrice int64) float64 {
	// There is always 100% demand for minimum price tickets.
	if ticketPrice == s.params.MinimumStakeDiff {
		return 1.0
	}

	// Calculate the demand based on yield.
	ticketsPerBlock := s.params.TicketsPerBlock
	posSubsidy := s.calcPoSSubsidy(nextHeight - 1)
	perVoteSubsidy := posSubsidy / dcrutil.Amount(ticketsPerBlock)
	yieldDemand := calcYieldDemand(ticketPrice, int64(perVoteSubsidy))

	// Calculate the demand based on the volume-weighted average ticket
	// purchase price.
	currentVWAP := s.calcPrevVWAP(s.tip)
	vwapDemand := calcVWAPDemand(ticketPrice, currentVWAP)

	// The demand is the combination of the two unless there is full demand
	// based on yield and no demand based on the VWAP, in which case there
	// is 100% demand.
	demand := yieldDemand * vwapDemand
	if yieldDemand == 1.0 && vwapDemand == 0.0 {
		demand = 1.0
	}
	return demand
}

// simulate runs the simulation using a calculated demand curve which models
// how ticket purchasing would typically proceed based upon the price and the
// VWAP.
func (s *simulator) simulate(numBlocks uint64) error {
	// Shorter versions of some params for convenience.
	ticketsPerBlock := s.params.TicketsPerBlock
	stakeValidationHeight := int32(s.params.StakeValidationHeight)
	stakeDiffWindowSize := int32(s.params.StakeDiffWindowSize)
	maxNewTicketsPerBlock := int32(s.params.MaxFreshStakePerBlock)
	maxTicketsPerWindow := maxNewTicketsPerBlock * stakeDiffWindowSize

	demandPerWindow := maxTicketsPerWindow
	for i := uint64(0); i < numBlocks; i++ {
		var nextHeight int32
		if s.tip != nil {
			nextHeight = s.tip.height + 1
		}

		// Purchase tickets according to simulated demand curve.
		nextTicketPrice := s.nextTicketPriceFunc()
		if nextTicketPrice < s.params.MinimumStakeDiff {
			panic(fmt.Sprintf("Ticket price function returned a "+
				"price of %v which is under the minimum "+
				"allowed price of %v",
				dcrutil.Amount(nextTicketPrice),
				dcrutil.Amount(s.params.MinimumStakeDiff)))
		}
		if nextHeight%stakeDiffWindowSize == 0 && nextHeight != 0 {
			demand := s.calcDemand(nextHeight, nextTicketPrice)
			demandPerWindow = int32(float64(maxTicketsPerWindow) * demand)
		}

		newTickets := uint8(demandPerWindow / stakeDiffWindowSize)
		maxPossible := int64(s.spendableSupply) / nextTicketPrice
		if int64(newTickets) > maxPossible {
			newTickets = uint8(maxPossible)
		}

		// TODO(davec): Account for tickets being purchased.
		// Limit the total staked coins to 40% of the total supply
		// except for in between blocks that are 60% and 80% of the
		// total number of blocks to simulate which limit to 50% of the
		// total supply in order to simulate a sudden surge and drop the
		// amount of stake coins.
		stakedCoins := s.totalSupply - s.spendableSupply
		if uint64(nextHeight) < (numBlocks*3/5) ||
			uint64(nextHeight) > (numBlocks*4/5) {
			if newTickets > 0 && stakedCoins > (s.totalSupply*2/5) {
				newTickets = 0
			}
		} else {
			if newTickets > 0 && stakedCoins > (s.totalSupply/2) {
				newTickets = 0
			}
		}

		// Start voting once stake validation height is reached.  This
		// assumes no votes are missed and revokes all expired tickets
		// as soon as possible which isn't very realistic, but it
		// doesn't have any effect on the ticket prices, so it's good
		// enough.  It could be useful to make this more realistic for
		// other simulation purposes though.
		var numVotes uint16
		if nextHeight >= stakeValidationHeight {
			numVotes = ticketsPerBlock
		}
		data := &simData{
			newTickets:  newTickets,
			prevValid:   true,
			revocations: uint16(len(s.unrevokedTickets)),
			voters:      numVotes,
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

func main() {
	var cpuProfilePath = flag.String("cpuprofile", "",
		"Write CPU profile to the specified file")
	var csvPath = flag.String("inputcsv", "",
		"Path to simulation CSV input data -- This overrides numblocks")
	var numBlocks = flag.Uint64("numblocks", 100000, "Number of blocks to simulate")
	flag.Parse()

	// Generate a CPU profile if requested.
	if *cpuProfilePath != "" {
		f, err := os.Create(*cpuProfilePath)
		if err != nil {
			fmt.Println("Unable to create cpu profile:", err)
			return
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// *********************************************************************
	// NOTE: Set a different function to calculate the next required stake
	// difficulty (aka ticket price) here.
	// *********************************************************************
	sim := newSimulator(&chaincfg.MainNetParams)
	sim.nextTicketPriceFunc = sim.curCalcNextStakeDiff
	//sim.nextTicketPriceFunc = sim.calcNextStakeDiffProposal1
	//sim.nextTicketPriceFunc = sim.calcNextStakeDiffProposal2
	//sim.nextTicketPriceFunc = sim.calcNextStakeDiffProposal3

	startTime := time.Now()
	if *csvPath != "" {
		fmt.Printf("Running simulation from %q.\n", *csvPath)
		fmt.Printf("Height")
		if err := sim.simulateFromCSV(*csvPath); err != nil {
			fmt.Println(err)
			return
		}
	} else {
		fmt.Printf("Running simulation for %d blocks.\n", *numBlocks)
		fmt.Printf("Height")
		if err := sim.simulate(*numBlocks); err != nil {
			fmt.Println(err)
			return
		}
	}
	fmt.Println("..done")
	fmt.Println("Simulation took", time.Since(startTime))

	// Generate the simulation results and open them in a browser.
	if err := generateResults(sim); err != nil {
		fmt.Println(err)
		return
	}
}
<commit_msg>Calculate yield based on avg expected vote height.<commit_after>// Copyright (c) 2017 Dave Collins
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

import (
	"encoding/csv"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"math/big"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"github.com/decred/dcrd/chaincfg"
	"github.com/decred/dcrd/chaincfg/chainhash"
	"github.com/decred/dcrd/wire"
	"github.com/decred/dcrutil"
)

const (
	// fieldsPerRecord defines the number of fields expected in each line
	// of the input CSV data.
	fieldsPerRecord = 3
)

// convertRecord converts the passed record, which is expected to be parsed from
// a CSV file, and thus will be a slice of strings, into a struct with concrete
// types.
func convertRecord(record []string) (*simData, error) {
	headerBytes, err := hex.DecodeString(record[1])
	if err != nil {
		return nil, err
	}

	var header wire.BlockHeader
	if err := header.FromBytes(headerBytes); err != nil {
		return nil, err
	}
	var hashStrings []string
	if record[2] != "" {
		hashStrings = strings.Split(record[2], ":")
	}
	if len(hashStrings) != int(header.FreshStake) {
		return nil, fmt.Errorf("%d ticket hashes in CSV for %d new tickets",
			len(hashStrings), header.FreshStake)
	}
	ticketHashes := make([]chainhash.Hash, 0, len(hashStrings))
	for _, hashString := range hashStrings {
		hash, err := chainhash.NewHashFromStr(hashString)
		if err != nil {
			return nil, err
		}
		ticketHashes = append(ticketHashes, *hash)
	}

	return &simData{
		header:       headerBytes,
		voters:       header.Voters,
		prevValid:    dcrutil.IsFlagSet16(header.VoteBits, dcrutil.BlockValid),
		newTickets:   header.FreshStake,
		ticketHashes: ticketHashes,
		revocations:  uint16(header.Revocations),
	}, nil
}

// reportProgress periodically prints out the current simulator height to
// stdout.
func (s *simulator) reportProgress() {
	if s.tip.height%10000 == 0 && s.tip.height != 0 {
		fmt.Println()
	}
	if s.tip.height%1000 == 0 && s.tip.height != 0 {
		fmt.Printf("..%d", s.tip.height)
	}
}

// simulateFromCSV runs the simulation using input data from a CSV file.  It is
// realistically only intended to be used with data extracted from mainnet in
// order to exactly replicate its live ticket pool.
func (s *simulator) simulateFromCSV(csvPath string) error {
	// Open the simulation CSV data which is expected to be in the following
	// format:
	//
	// Block Header,Winning Ticket Hashes
	csvFile, err := os.Open(csvPath)
	if err != nil {
		return err
	}

	// Create a new simulator using input from the CSV file.
	r := csv.NewReader(csvFile)
	r.FieldsPerRecord = fieldsPerRecord
	var handledHeader bool
	for {
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		// Skip header fields if they exist.
		if !handledHeader {
			handledHeader = true
			_, err := strconv.Atoi(record[0])
			if err != nil {
				continue
			}
		}

		// Convert the CSV to concrete data.
		data, err := convertRecord(record)
		if err != nil {
			return err
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

// calcYieldDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// estimated yield purchasing a ticket would produce.
func calcYieldDemand(ticketPrice, perVoteSubsidy int64) float64 {
	// 100% demand when the yield is over 5%.
	yield := float64(perVoteSubsidy) / float64(ticketPrice)
	if yield > 0.05 {
		return 1.0
	}

	// No demand when the yield is under 2%.
	if yield < 0.02 {
		return 0.0
	}

	// The yield is in between 2% and 5%, so create a linear demand
	// accordingly.
	return (yield - 0.02) / 0.03
}

// calcVWAPDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval) based upon the
// volume-weighted average ticket purchase of the previous ticket price windows.
func calcVWAPDemand(ticketPrice, ticketVWAP int64) float64 {
	// 100% demand when the ticket price is under 80% of the VWAP.
	eightyPercentVWAP := (ticketVWAP * 8) / 10
	if ticketPrice < eightyPercentVWAP {
		return 1.0
	}

	// No demand when the ticket price is over 120% of the VWAP.
	if ticketPrice > (ticketVWAP*12)/10 {
		return 0.0
	}

	// The ticket price is in between 80% and 120% of the VWAP, so create
	// a linear demand accordingly.
	fortyPercentVWAP := (ticketVWAP * 4) / 10
	return 1 - float64(ticketPrice-eightyPercentVWAP)/float64(fortyPercentVWAP)
}

// calcVWAP calculates and return the volume-weighted average ticket purchase
// price for up to 'StakeDiffWindows' worth of the previous ticket price
// windows.
func (s *simulator) calcPrevVWAP(prevNode *blockNode) int64 {
	windowSize := int32(s.params.StakeDiffWindowSize)
	stakeDiffWindows := int32(s.params.StakeDiffWindows)

	// Calculate the height the block just before the most recent ticket
	// price change.
	wantHeight := prevNode.height - (prevNode.height+1)%windowSize
	prevNode = s.ancestorNode(prevNode, wantHeight, nil)

	// Loop through previous required number of previous blocks and tally up
	// all of the weighted ticket purchase prices as well as the total
	// number of purchased tickets.
	numTickets, weightedVal := new(big.Int), new(big.Int)
	weightedSum, totalTickets := new(big.Int), new(big.Int)
	blocksToIterate := stakeDiffWindows * windowSize
	for i := int32(0); i < blocksToIterate && prevNode != nil; i++ {
		// weightedSum += numTickets*ticketPrice
		// totalTickets += numTickets
		numTickets.SetInt64(int64(len(prevNode.ticketsAdded)))
		weightedVal.SetInt64(prevNode.ticketPrice)
		weightedVal.Mul(weightedVal, numTickets)
		weightedSum.Add(weightedSum, weightedVal)
		totalTickets.Add(totalTickets, numTickets)
		prevNode = prevNode.parent
	}

	// Return minimum ticket price if there are were not any ticket
	// purchases at all in the entire period being examined.
	if totalTickets.Sign() == 0 {
		return s.params.MinimumStakeDiff
	}

	return new(big.Int).Div(weightedSum, totalTickets).Int64()
}

// calcDemand returns a simulated demand (as a percentage of the number of
// tickets to purchase within a given stake difficulty interval).
func (s *simulator) calcDemand(nextHeight int32, ticketPrice int64) float64 {
	// Calculate the demand based on yield.
	expectedPayoutHeight := int32((time.Hour * 24) * 28 / s.params.TargetTimePerBlock)
	ticketsPerBlock := s.params.TicketsPerBlock
	posSubsidy := s.calcPoSSubsidy(nextHeight + expectedPayoutHeight - 1)
	perVoteSubsidy := posSubsidy / dcrutil.Amount(ticketsPerBlock)
	yieldDemand := calcYieldDemand(ticketPrice, int64(perVoteSubsidy))

	// Calculate the demand based on the volume-weighted average ticket
	// purchase price.
	currentVWAP := s.calcPrevVWAP(s.tip)
	vwapDemand := calcVWAPDemand(ticketPrice, currentVWAP)

	// The demand is the combination of the two unless there is full demand
	// based on yield and no demand based on the VWAP, in which case there
	// is 100% demand.
	demand := yieldDemand * vwapDemand
	if yieldDemand == 1.0 && vwapDemand == 0.0 {
		demand = 1.0
	}
	return demand
}

// simulate runs the simulation using a calculated demand curve which models
// how ticket purchasing would typically proceed based upon the price and the
// VWAP.
func (s *simulator) simulate(numBlocks uint64) error {
	// Shorter versions of some params for convenience.
	ticketsPerBlock := s.params.TicketsPerBlock
	stakeValidationHeight := int32(s.params.StakeValidationHeight)
	stakeDiffWindowSize := int32(s.params.StakeDiffWindowSize)
	maxNewTicketsPerBlock := int32(s.params.MaxFreshStakePerBlock)
	maxTicketsPerWindow := maxNewTicketsPerBlock * stakeDiffWindowSize

	demandPerWindow := maxTicketsPerWindow
	for i := uint64(0); i < numBlocks; i++ {
		var nextHeight int32
		if s.tip != nil {
			nextHeight = s.tip.height + 1
		}

		// Purchase tickets according to simulated demand curve.
		nextTicketPrice := s.nextTicketPriceFunc()
		if nextTicketPrice < s.params.MinimumStakeDiff {
			panic(fmt.Sprintf("Ticket price function returned a "+
				"price of %v which is under the minimum "+
				"allowed price of %v",
				dcrutil.Amount(nextTicketPrice),
				dcrutil.Amount(s.params.MinimumStakeDiff)))
		}
		if nextHeight%stakeDiffWindowSize == 0 && nextHeight != 0 {
			demand := s.calcDemand(nextHeight, nextTicketPrice)
			demandPerWindow = int32(float64(maxTicketsPerWindow) * demand)
		}

		newTickets := uint8(demandPerWindow / stakeDiffWindowSize)
		maxPossible := int64(s.spendableSupply) / nextTicketPrice
		if int64(newTickets) > maxPossible {
			newTickets = uint8(maxPossible)
		}

		// TODO(davec): Account for tickets being purchased.
		// Limit the total staked coins to 40% of the total supply
		// except for in between blocks that are 60% and 80% of the
		// total number of blocks to simulate which limit to 50% of the
		// total supply in order to simulate a sudden surge and drop the
		// amount of stake coins.
		stakedCoins := s.totalSupply - s.spendableSupply
		if uint64(nextHeight) < (numBlocks*3/5) ||
			uint64(nextHeight) > (numBlocks*4/5) {
			if newTickets > 0 && stakedCoins > (s.totalSupply*2/5) {
				newTickets = 0
			}
		} else {
			if newTickets > 0 && stakedCoins > (s.totalSupply/2) {
				newTickets = 0
			}
		}

		// Start voting once stake validation height is reached.  This
		// assumes no votes are missed and revokes all expired tickets
		// as soon as possible which isn't very realistic, but it
		// doesn't have any effect on the ticket prices, so it's good
		// enough.  It could be useful to make this more realistic for
		// other simulation purposes though.
		var numVotes uint16
		if nextHeight >= stakeValidationHeight {
			numVotes = ticketsPerBlock
		}
		data := &simData{
			newTickets:  newTickets,
			prevValid:   true,
			revocations: uint16(len(s.unrevokedTickets)),
			voters:      numVotes,
		}

		// Create a new node that extends the current tip using the
		// simulation data and potentially report the progress.
		s.nextNode(data)
		s.reportProgress()
	}

	return nil
}

func main() {
	var cpuProfilePath = flag.String("cpuprofile", "",
		"Write CPU profile to the specified file")
	var csvPath = flag.String("inputcsv", "",
		"Path to simulation CSV input data -- This overrides numblocks")
	var numBlocks = flag.Uint64("numblocks", 100000, "Number of blocks to simulate")
	flag.Parse()

	// Generate a CPU profile if requested.
	if *cpuProfilePath != "" {
		f, err := os.Create(*cpuProfilePath)
		if err != nil {
			fmt.Println("Unable to create cpu profile:", err)
			return
		}
		pprof.StartCPUProfile(f)
		defer f.Close()
		defer pprof.StopCPUProfile()
	}

	// *********************************************************************
	// NOTE: Set a different function to calculate the next required stake
	// difficulty (aka ticket price) here.
	// *********************************************************************
	sim := newSimulator(&chaincfg.MainNetParams)
	sim.nextTicketPriceFunc = sim.curCalcNextStakeDiff
	//sim.nextTicketPriceFunc = sim.calcNextStakeDiffProposal1
	//sim.nextTicketPriceFunc = sim.calcNextStakeDiffProposal2
	//sim.nextTicketPriceFunc = sim.calcNextStakeDiffProposal3

	startTime := time.Now()
	if *csvPath != "" {
		fmt.Printf("Running simulation from %q.\n", *csvPath)
		fmt.Printf("Height")
		if err := sim.simulateFromCSV(*csvPath); err != nil {
			fmt.Println(err)
			return
		}
	} else {
		fmt.Printf("Running simulation for %d blocks.\n", *numBlocks)
		fmt.Printf("Height")
		if err := sim.simulate(*numBlocks); err != nil {
			fmt.Println(err)
			return
		}
	}
	fmt.Println("..done")
	fmt.Println("Simulation took", time.Since(startTime))

	// Generate the simulation results and open them in a browser.
	if err := generateResults(sim); err != nil {
		fmt.Println(err)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/lluchs/mima"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Println("Usage: mimarun <filename>")
		return
	}

	filename := os.Args[1]
	file, err := os.Open(filename)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("Parsing...\n")
	program, err := mima.Parse(file)
	if err != nil {
		fmt.Println(err)
		return
	}
	// Print the program.
	fmt.Println("Pointer")
	for mark, pointer := range program.Marks {
		fmt.Println(mark, pointer)
	}

	fmt.Println("Instructions")
	for address, instruction := range program.Instructions {
		fmt.Printf("%s = %s(%s)\n", address, instruction.Op, instruction.Argument)
	}
}
<commit_msg>Implement assembling and running code<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/lluchs/mima"
)

func main() {
	if len(os.Args) != 2 {
		fmt.Println("Usage: mimarun <filename>")
		return
	}

	filename := os.Args[1]
	file, err := os.Open(filename)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Println("Parsing...\n")
	program, err := mima.Parse(file)
	if err != nil {
		fmt.Println(err)
		return
	}
	// Print the program.
	fmt.Println("Pointer")
	for mark, pointer := range program.Marks {
		fmt.Println(mark, pointer)
	}

	fmt.Println("Instructions")
	for address, instruction := range program.Instructions {
		fmt.Printf("%s = %s(%s)\n", address, instruction.Op, instruction.Argument)
	}

	fmt.Println("\n\nAssembling...\n")
	bytecode, err := program.Assemble()
	if err != nil {
		fmt.Println(err)
		return
	}
	// Print the assembled code.
	fmt.Printf("Start: 0x%06X\n", bytecode.Start)
	printMem(bytecode.Mem)

	fmt.Println("\n\nRunning...\n")
	mem, err := bytecode.Run()
	if err != nil {
		fmt.Println(err)
		return
	}
	// Print the resulting memory.
	printMem(mem)

}

// Print all memory locations which are not 0.
func printMem(mem []uint32) {
	for pos, content := range mem {
		if content != 0 {
			fmt.Printf("0x%06X: 0x%06X\n", pos, content)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"strconv"

	ldns "github.com/Liamraystanley/go-ldns"
	arg "github.com/alexflint/go-arg"
	"github.com/kataras/go-template/html"
	"github.com/kataras/iris"
)

// TODO: include items from public-dns.info?
// TODO: http://stackoverflow.com/a/31627459/1830159
// TODO: Other thoughts. Type selector,

// Config represents the configuration for the app
type Config struct {
	Debug           bool                `arg:"-d,help:enable debugging mode"`
	Host            string              `arg:"-h,help:host/ip for which to bind to"`
	Port            int                 `arg:"-p,help:port which to bind to"`
	Database        string              `arg:"help:file path to the database for dnscheck"`
	GeoDb           string              `arg:"help:GeoIP database location"`
	CustomResolvers []string            `arg:"-r,help:resolver to use to resolve query lookups"`
	Resolvers       map[string][]string `arg:"-"` // underlying resolver map, created during startup
	Concurrency     int                 `arg:"-c,help:number of records to use for resolving records"`
	Limit           int                 `arg:"-l,help:max queries per request"`
}

// setup some defaults
var conf = Config{
	Debug:           false,
	Host:            "localhost",
	Port:            3000,
	Database:        "dns.db",
	GeoDb:           "geoip.db",
	CustomResolvers: []string{},
	Resolvers:       make(map[string][]string),
	Concurrency:     10,
	Limit:           500,
}

var logger *log.Logger

func webLogRequest(ctx *iris.Context) {
	logger.Printf("http: request %d from %s for: %s", ctx.ConnRequestNum(), ctx.RemoteIP(), ctx.PathString())

	ctx.Next()
}

func handleError(ctx *iris.Context) {
	ctx.Write("An unknown error occurred")
}

func handleNotFound(ctx *iris.Context) {
	ctx.MustRender("404.html", "")
}

// getWebContext generates the web contexts for use with html templates
func getWebContext(c *iris.Context) map[string]interface{} {
	return iris.Map{
		"Messages": c.GetFlashes(),
		"Conf":     conf,
	}
}

func saveLookup(results *DNSResults) (string, error) {
	db, err := newDB()
	if err != nil {
		return "", err
	}
	defer db.Clean()

	key := genWord(5, 6)

	return key, db.SetStruct("records", key, results)
}

func getLookup(id string) (*DNSResults, error) {
	db, err := newDB()
	if err != nil {
		return nil, err
	}
	defer db.Clean()

	results := &DNSResults{}

	return results, db.GetStruct("records", id, results)
}

func genResolvers() error {
	if len(conf.CustomResolvers) == 0 {
		// assume defaults. Google DNS, OpenDNS, and local resolvers.
		localResolvers, err := ldns.ReadResolveConf()
		if err != nil {
			return err
		}

		conf.Resolvers["Local Resolvers"] = localResolvers
		conf.Resolvers["Google DNS"] = []string{"8.8.8.8", "8.8.4.4"}
		conf.Resolvers["OpenDNS"] = []string{"208.67.222.222", "208.67.220.220"}

		return nil
	}

	conf.Resolvers["Custom"] = conf.CustomResolvers

	return nil
}

func initWebserver() error {
	logger.Println("initializing webserver")

	iris.Config.Sessions.Cookie = "session"
	iris.Config.LoggerOut = os.Stdout // ioutil.Discard
	iris.Config.DisableBanner = true
	iris.Config.Gzip = true
	iris.Config.IsDevelopment = conf.Debug
	iris.StaticWeb("/static", "./static", 1)
	iris.UseTemplate(html.New(html.Config{Layout: "base.html"})).Directory("./static", ".html") //.Binary(Asset, AssetNames)
	iris.UseFunc(webLogRequest)

	// 500
	iris.OnError(iris.StatusInternalServerError, handleError)

	// 404
	iris.OnError(iris.StatusNotFound, handleNotFound)

	iris.Get("/", func(ctx *iris.Context) {
		ctx.MustRender("index.html", getWebContext(ctx))
	})("index")

	iris.Post("/", func(ctx *iris.Context) {
		input := ctx.FormValueString("hosts")
		lookupType := ctx.FormValueString("recordtype")
		resolvers := ctx.FormValueString("resolvers")

		if _, ok := conf.Resolvers[resolvers]; !ok {
			ctx.SetFlash("error", "Resolvers specified do not exist")
			ctx.SetFlash("originalHosts", input)

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		hosts, err := parseHosts(input)
		if err != nil {
			ctx.SetFlash("originalHosts", input)
			ctx.SetFlash("error", err.Error())

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		results, err := LookupAll(hosts, conf.Resolvers[resolvers], lookupType)
		if err != nil {
			ctx.SetFlash("originalHosts", input)
			ctx.SetFlash("error", err.Error())

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		id, err := saveLookup(results)
		if err != nil {
			ctx.SetFlash("originalHosts", input)
			ctx.SetFlash("error", err.Error())

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		ctx.RedirectTo("results", id)
	})

	iris.Get("/r/:key", func(ctx *iris.Context) {
		id := ctx.Param("key")

		result, err := getLookup(id)
		if err != nil {
			fmt.Println(err)

			ctx.MustRender("404.html", "")
			return
		}

		out := getWebContext(ctx)
		out["Results"] = result
		ctx.MustRender("results.html", out)
	})("results")

	iris.Get("/api/:key", func(ctx *iris.Context) {
		id := ctx.Param("key")

		result, err := getLookup(id)
		if err != nil {
			fmt.Println(err)

			ctx.JSON(iris.StatusNotFound, map[string]string{"error": "an entry with that key does not exist"})
			return
		}

		ctx.JSON(iris.StatusOK, result)
	})("api-results")

	iris.Get("/stats/:key", func(ctx *iris.Context) {
		id := ctx.Param("key")

		result, err := getLookup(id)
		if err != nil {
			fmt.Println(err)

			ctx.JSON(iris.StatusNotFound, map[string]string{"error": "an entry with that key does not exist"})
			return
		}

		stats, err := result.Stats()
		if err != nil {
			fmt.Println(err)

			ctx.JSON(iris.StatusInternalServerError, map[string]string{"error": "an unknown error occurred"})
			return
		}

		ctx.JSON(iris.StatusOK, stats)
	})

	listener, err := net.Listen("tcp", conf.Host+":"+strconv.Itoa(conf.Port))
	if err != nil {
		return err
	}

	return iris.Serve(listener)
}

func main() {
	// initialize app flags
	arg.MustParse(&conf)

	// initialize logger
	logger = log.New(os.Stdout, "", log.Lshortfile|log.LstdFlags)
	logger.Println("initializing logger")

	// initialize the database
	initDatabase()

	// initialize the resolvers
	if err := genResolvers(); err != nil {
		logger.Fatal(err)
	}

	// check for geoip updates (once a week is good 'nuff)
	GeoIPUpdateCheck(conf.GeoDb)

	// initialize webserver
	if err := initWebserver(); err != nil {
		logger.Fatal("error: ", err)
	}
}
<commit_msg>add some functions to help during template parsing and execution<commit_after>package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"strconv"
	"strings"

	ldns "github.com/Liamraystanley/go-ldns"
	arg "github.com/alexflint/go-arg"
	"github.com/kataras/go-template/html"
	"github.com/kataras/iris"
)

// TODO: include items from public-dns.info?
// TODO: http://stackoverflow.com/a/31627459/1830159
// TODO: Other thoughts. Type selector,

// Config represents the configuration for the app
type Config struct {
	Debug           bool                `arg:"-d,help:enable debugging mode"`
	Host            string              `arg:"-h,help:host/ip for which to bind to"`
	Port            int                 `arg:"-p,help:port which to bind to"`
	Database        string              `arg:"help:file path to the database for dnscheck"`
	GeoDb           string              `arg:"help:GeoIP database location"`
	CustomResolvers []string            `arg:"-r,help:resolver to use to resolve query lookups"`
	Resolvers       map[string][]string `arg:"-"` // underlying resolver map, created during startup
	Concurrency     int                 `arg:"-c,help:number of records to use for resolving records"`
	Limit           int                 `arg:"-l,help:max queries per request"`
}

// setup some defaults
var conf = Config{
	Debug:           false,
	Host:            "localhost",
	Port:            3000,
	Database:        "dns.db",
	GeoDb:           "geoip.db",
	CustomResolvers: []string{},
	Resolvers:       make(map[string][]string),
	Concurrency:     10,
	Limit:           500,
}

var logger *log.Logger

func webLogRequest(ctx *iris.Context) {
	logger.Printf("http: request %d from %s for: %s", ctx.ConnRequestNum(), ctx.RemoteIP(), ctx.PathString())

	ctx.Next()
}

func handleError(ctx *iris.Context) {
	ctx.Write("An unknown error occurred")
}

func handleNotFound(ctx *iris.Context) {
	ctx.MustRender("404.html", "")
}

// getWebContext generates the web contexts for use with html templates
func getWebContext(c *iris.Context) map[string]interface{} {
	return iris.Map{
		"Messages": c.GetFlashes(),
		"Conf":     conf,
	}
}

func saveLookup(results *DNSResults) (string, error) {
	db, err := newDB()
	if err != nil {
		return "", err
	}
	defer db.Clean()

	key := genWord(5, 6)

	return key, db.SetStruct("records", key, results)
}

func getLookup(id string) (*DNSResults, error) {
	db, err := newDB()
	if err != nil {
		return nil, err
	}
	defer db.Clean()

	results := &DNSResults{}

	return results, db.GetStruct("records", id, results)
}

func genResolvers() error {
	if len(conf.CustomResolvers) == 0 {
		// assume defaults. Google DNS, OpenDNS, and local resolvers.
		localResolvers, err := ldns.ReadResolveConf()
		if err != nil {
			return err
		}

		conf.Resolvers["Local Resolvers"] = localResolvers
		conf.Resolvers["Google DNS"] = []string{"8.8.8.8", "8.8.4.4"}
		conf.Resolvers["OpenDNS"] = []string{"208.67.222.222", "208.67.220.220"}

		return nil
	}

	conf.Resolvers["Custom"] = conf.CustomResolvers

	return nil
}

func initWebserver() error {
	logger.Println("initializing webserver")

	funcmap := make(map[string]interface{})
	funcmap["isip"] = func(ip string) bool {
		if isip := net.ParseIP(ip); isip == nil {
			return false
		}

		return true
	}
	funcmap["tolower"] = func(input string) string {
		return strings.ToLower(input)
	}
	funcmap["join"] = func(input []string) string {
		return strings.Join(input, ", ")
	}

	iris.Config.Sessions.Cookie = "session"
	iris.Config.LoggerOut = os.Stdout // ioutil.Discard
	iris.Config.DisableBanner = true
	iris.Config.Gzip = true
	iris.Config.IsDevelopment = conf.Debug
	iris.StaticWeb("/static", "./static", 1)
	iris.UseTemplate(html.New(html.Config{Layout: "base.html", Funcs: funcmap})).Directory("./static", ".html") //.Binary(Asset, AssetNames)
	iris.UseFunc(webLogRequest)

	// 500
	iris.OnError(iris.StatusInternalServerError, handleError)

	// 404
	iris.OnError(iris.StatusNotFound, handleNotFound)

	iris.Get("/", func(ctx *iris.Context) {
		ctx.MustRender("index.html", getWebContext(ctx))
	})("index")

	iris.Post("/", func(ctx *iris.Context) {
		input := ctx.FormValueString("hosts")
		lookupType := ctx.FormValueString("recordtype")
		resolvers := ctx.FormValueString("resolvers")

		if _, ok := conf.Resolvers[resolvers]; !ok {
			ctx.SetFlash("error", "Resolvers specified do not exist")
			ctx.SetFlash("originalHosts", input)

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		hosts, err := parseHosts(input)
		if err != nil {
			ctx.SetFlash("originalHosts", input)
			ctx.SetFlash("error", err.Error())

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		results, err := LookupAll(hosts, conf.Resolvers[resolvers], lookupType)
		if err != nil {
			ctx.SetFlash("originalHosts", input)
			ctx.SetFlash("error", err.Error())

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		id, err := saveLookup(results)
		if err != nil {
			ctx.SetFlash("originalHosts", input)
			ctx.SetFlash("error", err.Error())

			ctx.MustRender("index.html", getWebContext(ctx))
			return
		}

		ctx.RedirectTo("results", id)
	})

	iris.Get("/r/:key", func(ctx *iris.Context) {
		id := ctx.Param("key")

		result, err := getLookup(id)
		if err != nil {
			fmt.Println(err)

			ctx.MustRender("404.html", "")
			return
		}

		out := getWebContext(ctx)
		out["Results"] = result
		ctx.MustRender("results.html", out)
	})("results")

	iris.Get("/api/:key", func(ctx *iris.Context) {
		id := ctx.Param("key")

		result, err := getLookup(id)
		if err != nil {
			fmt.Println(err)

			ctx.JSON(iris.StatusNotFound, map[string]string{"error": "an entry with that key does not exist"})
			return
		}

		ctx.JSON(iris.StatusOK, result)
	})("api-results")

	iris.Get("/stats/:key", func(ctx *iris.Context) {
		id := ctx.Param("key")

		result, err := getLookup(id)
		if err != nil {
			fmt.Println(err)

			ctx.JSON(iris.StatusNotFound, map[string]string{"error": "an entry with that key does not exist"})
			return
		}

		stats, err := result.Stats()
		if err != nil {
			fmt.Println(err)

			ctx.JSON(iris.StatusInternalServerError, map[string]string{"error": "an unknown error occurred"})
			return
		}

		ctx.JSON(iris.StatusOK, stats)
	})

	listener, err := net.Listen("tcp", conf.Host+":"+strconv.Itoa(conf.Port))
	if err != nil {
		return err
	}

	return iris.Serve(listener)
}

func main() {
	// initialize app flags
	arg.MustParse(&conf)

	// initialize logger
	logger = log.New(os.Stdout, "", log.Lshortfile|log.LstdFlags)
	logger.Println("initializing logger")

	// initialize the database
	initDatabase()

	// initialize the resolvers
	if err := genResolvers(); err != nil {
		logger.Fatal(err)
	}

	// check for geoip updates (once a week is good 'nuff)
	GeoIPUpdateCheck(conf.GeoDb)

	// initialize webserver
	if err := initWebserver(); err != nil {
		logger.Fatal("error: ", err)
	}
}
<|endoftext|>