<commit_before>package main

import (
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/ziutek/mymysql/mysql"
	_ "github.com/ziutek/mymysql/native"
	"log"
	"os"
	"sort"
	"strconv"
)

const _DB_CONFIG_FILENAME = "db_config.SECRET.json"
const _OUTPUT_FILENAME = "output.csv"
const _QUERY_LIMIT = 100
const _PENALTY_PERCENTAGE_CUTOFF = 0.10

//How many solves a user must have to have their relative scale included.
//A low value gives you far more very low or very high scores than you shoul get.
const _MINIMUM_SOLVES = 10

var noLimitFlag bool
var printPuzzleDataFlag bool
var cullCheaterPercentageFlag float64
var minimumSolvesFlag int
var useMockData bool
var queryLimit int

func init() {
	flag.BoolVar(&noLimitFlag, "a", false, "Specify to execute the solves query with no limit.")
	flag.BoolVar(&printPuzzleDataFlag, "p", false, "Specify that you want puzzle data printed out in the output.")
	flag.Float64Var(&cullCheaterPercentageFlag, "c", _PENALTY_PERCENTAGE_CUTOFF, "What percentage of solve time must be penalty for someone to be considered a cheater.")
	flag.IntVar(&minimumSolvesFlag, "s", _MINIMUM_SOLVES, "How many solves a user must have their scores considered.")
	flag.IntVar(&queryLimit, "n", _QUERY_LIMIT, "Number of solves to fetch from the database.")
	flag.BoolVar(&useMockData, "m", false, "Use mock data (useful if you don't have a real database to test with).")
}

type dbConfig struct {
	Url               string
	Username          string
	Password          string
	DbName            string
	SolvesTable       string
	SolvesID          string
	SolvesPuzzleID    string
	SolvesTotalTime   string
	SolvesPenaltyTime string
	SolvesUser        string
	PuzzlesTable      string
	PuzzlesID         string
	PuzzlesDifficulty string
	PuzzlesName       string
	PuzzlesPuzzle     string
}

var config dbConfig

type solve struct {
	puzzleID    int
	totalTime   int
	penaltyTime int
}

type userSolvesCollection struct {
	solves     []solve
	max        int
	min        int
	idPosition map[int]int
}

type puzzle struct {
	id                     int
	userRelativeDifficulty float32
	difficultyRating       int
	name                   string
	puzzle                 string
}

type puzzles []puzzle

type byUserRelativeDifficulty struct {
	puzzles
}

func (self puzzles) Len() int {
	return len(self)
}

func (self puzzles) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func (self byUserRelativeDifficulty) Less(i, j int) bool {
	return self.puzzles[i].userRelativeDifficulty < self.puzzles[j].userRelativeDifficulty
}

type bySolveTimeDsc []solve

func (self bySolveTimeDsc) Len() int {
	return len(self)
}

func (self bySolveTimeDsc) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func (self bySolveTimeDsc) Less(i, j int) bool {
	//For the purposes of this algorithm, the "best" has to be lowest rank.
	return self[i].totalTime > self[j].totalTime
}

func (self *userSolvesCollection) addSolve(solve solve) bool {
	//Cull obviously incorrect solves.
	if solve.totalTime == 0 {
		return false
	}

	//Cull solves that leaned too heavily on hints.
	if float64(solve.penaltyTime)/float64(solve.totalTime) > cullCheaterPercentageFlag {
		return false
	}

	self.solves = append(self.solves, solve)
	if len(self.solves) == 1 {
		self.max = solve.totalTime
		self.min = solve.totalTime
	} else {
		if self.max < solve.totalTime {
			self.max = solve.totalTime
		}
		if self.min > solve.totalTime {
			self.min = solve.totalTime
		}
	}
	return true
}

//Whehter or not this should be included in calculation.
//Basically, whether the reltaiveDifficulties will all be valid.
//Normally this returns false if there is only one solve by the user, but could also
//happen when there are multiple solves but (crazily enough) they all have exactly the same solveTime.
//This DOES happen in the production dataset.
func (self *userSolvesCollection) valid() bool {
	if self.max == self.min {
		return false
	}

	if len(self.solves) < minimumSolvesFlag {
		return false
	}

	return true
}

func (self *userSolvesCollection) relativeDifficulties() map[int]float32 {
	//Returns a map of puzzle id to relative difficulty, normalized by our max and min.
	avgSolveTimes := make(map[int]float32)
	//Keep track of how many times we've seen each puzzle solved by this user so we can do correct averaging.
	avgSolveTimesCount := make(map[int]int)

	//First, collect the average solve time (in case the same user has solved more than once the same puzzle)

	for _, solve := range self.solves {
		currentAvgSolveTime := avgSolveTimes[solve.puzzleID]

		avgSolveTimes[solve.puzzleID] = (currentAvgSolveTime*float32(avgSolveTimesCount[solve.puzzleID]) + float32(solve.totalTime)) / float32(avgSolveTimesCount[solve.puzzleID]+1)

		avgSolveTimesCount[solve.puzzleID]++
	}

	//Now, relativize all of the scores.

	result := make(map[int]float32)

	for puzzleID, avgSolveTime := range avgSolveTimes {
		result[puzzleID] = (avgSolveTime - float32(self.min)) / float32(self.max-self.min)
	}

	return result
}

func main() {

	flag.Parse()

	file, err := os.Open(_DB_CONFIG_FILENAME)
	if err != nil {
		log.Fatal("Could not find the config file at ", _DB_CONFIG_FILENAME, ". You should copy the SAMPLE one to that filename and configure.")
		os.Exit(1)
	}
	defer file.Close()
	decoder := json.NewDecoder(file)
	if err := decoder.Decode(&config); err != nil {
		log.Fatal("There was an error parsing JSON from the config file: ", err)
		os.Exit(1)
	}

	difficutlyRatingsChan := make(chan map[int]puzzle)

	go getPuzzleDifficultyRatings(difficutlyRatingsChan)

	var db mysql.Conn

	if useMockData {
		db = &mockConnection{}
	} else {
		db = mysql.New("tcp", "", config.Url, config.Username, config.Password, config.DbName)
	}

	if err := db.Connect(); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	var solvesQuery string

	if noLimitFlag {
		log.Println("Running without a limit for number of solves to retrieve.")
		solvesQuery = "select %s, %s, %s, %s from %s"
	} else {
		log.Println("Running with a limit of ", queryLimit, " for number of solves to retrieve.")
		solvesQuery = "select %s, %s, %s, %s from %s limit " + strconv.Itoa(queryLimit)
	}

	res, err := db.Start(solvesQuery, config.SolvesUser, config.SolvesPuzzleID, config.SolvesTotalTime, config.SolvesPenaltyTime, config.SolvesTable)

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

	solvesByUser := make(map[string]*userSolvesCollection)

	var userSolves *userSolvesCollection
	var ok bool
	var i int
	var skippedSolves int
	var skippedDuplicateSolves int

	seenRows := make(map[string]bool)

	//First, process all user records in the DB to collect all solves by userName.
	for {

		row, _ := res.GetRow()

		if row == nil {
			break
		}

		i++

		rowHashValue := fmt.Sprintf("%v", row)

		if _, seen := seenRows[rowHashValue]; seen {
			skippedDuplicateSolves++
			continue
		} else {
			seenRows[rowHashValue] = true
		}

		userSolves, ok = solvesByUser[row.Str(0)]

		if !ok {
			userSolves = new(userSolvesCollection)
			userSolves.idPosition = make(map[int]int)
			solvesByUser[row.Str(0)] = userSolves
		}

		if !userSolves.addSolve(solve{row.Int(1), row.Int(2), row.Int(3)}) {
			skippedSolves++
		}

	}

	log.Println("Processed ", i, " solves by ", len(solvesByUser), " users.")
	log.Println("Skipped ", skippedSolves, " solves that cheated too much.")
	log.Println("Skipped ", skippedDuplicateSolves, " solves because they were duplicates of solves seen earlier.")

	//Now get the relative difficulty for each user's puzzles, and collect them.

	relativeDifficultiesByPuzzle := make(map[int][]float32)

	collectionByPuzzle := make(map[int]map[*userSolvesCollection]bool)

	var skippedUsers int

	for _, collection := range solvesByUser {

		/*
			//TODO: consider removing this logic and all of skipped users totally.
				if !collection.valid() {
					skippedUsers++
					continue
				}
		*/

		for puzzleID, relativeDifficulty := range collection.relativeDifficulties() {
			relativeDifficultiesByPuzzle[puzzleID] = append(relativeDifficultiesByPuzzle[puzzleID], relativeDifficulty)
		}

		sort.Sort(bySolveTimeDsc(collection.solves))

		for i, puzzle := range collection.solves {
			collection.idPosition[puzzle.puzzleID] = i

			collectionMap, ok := collectionByPuzzle[puzzle.puzzleID]
			if !ok {
				collectionMap = make(map[*userSolvesCollection]bool)
			}

			collectionMap[collection] = true

			collectionByPuzzle[puzzle.puzzleID] = collectionMap
		}

	}

	//Now, create the Markov Transition Matrix, according to algorithm MC4 of http://www.wisdom.weizmann.ac.il/~naor/PAPERS/rank_www10.html
	//We start by creating a stacked array of float64's that we'll pass to the matrix library.

	numPuzzles := len(collectionByPuzzle)

	matrixData := make([][]float64, numPuzzles)

	for i := range matrixData {
		matrixData[i] = make([]float64, numPuzzles)
	}

	//Now we will associate each observed puzzleID with an index that it will be associated with in the matrix.
	puzzleIndex := make([]int, numPuzzles)

	counter := 0
	for key, _ := range collectionByPuzzle {
		puzzleIndex[counter] = key
		counter++
	}

	//Now we start to build up the matrix according to the MC4 algorithm.

	for i := 0; i < numPuzzles; i++ {
		for j := 0; j < numPuzzles; j++ {
			//Convert the zero-index into the puzzle ID we're actually interested in.
			p := puzzleIndex[i]
			q := puzzleIndex[j]

			//Find the intersection of userSolveCollections that contain both p and q.
			pMap := collectionByPuzzle[p]
			qMap := collectionByPuzzle[q]

			var intersection []*userSolvesCollection
			for collection, _ := range pMap {
				if _, ok := qMap[collection]; ok {
					intersection = append(intersection, collection)
				}
			}

			//Next, calculate how many of the collections have q ranked better (lower!) than p.
			count := 0
			for _, collection := range intersection {
				if collection.idPosition[q] < collection.idPosition[p] {
					count++
				}
			}

			//Is it a majority? if so, transition. if not, leave at 0.
			if count > (len(intersection) / 2) {
				matrixData[i][j] = 1.0
			}

		}
	}

	//Wait, should I go through and normalize so that the rows all sum to 1 probablity wise?

	log.Println("Skipped ", skippedUsers, " users because they did not have enough solve times.")

	puzzles := make([]puzzle, len(relativeDifficultiesByPuzzle))

	var index int

	for puzzleID, difficulties := range relativeDifficultiesByPuzzle {
		var sum float32
		for _, difficulty := range difficulties {
			sum += difficulty
		}
		puzzles[index] = puzzle{id: puzzleID, userRelativeDifficulty: sum / float32(len(difficulties)), difficultyRating: -1}
		index++
	}

	//Sort the puzzles by relative user difficulty
	//We actually don't need the wrapper, since it will modify the underlying slice.
	sort.Sort(byUserRelativeDifficulty{puzzles})

	//Merge in the difficulty ratings from the server.
	difficultyRatings := <-difficutlyRatingsChan

	for i, puzzle := range puzzles {
		info, ok := difficultyRatings[puzzle.id]
		if ok {
			puzzle.difficultyRating = info.difficultyRating
			puzzle.name = info.name
			puzzle.puzzle = info.puzzle
		}
		//It's not a pointer so we have to copy it back.
		puzzles[i] = puzzle
	}

	//Now print the results to stdout.

	csvOut := csv.NewWriter(os.Stdout)

	for _, puzzle := range puzzles {
		temp := []string{strconv.Itoa(puzzle.id), strconv.Itoa(puzzle.difficultyRating), fmt.Sprintf("%g", puzzle.userRelativeDifficulty), puzzle.name}
		if printPuzzleDataFlag {
			temp = append(temp, puzzle.puzzle)
		}
		csvOut.Write(temp)
	}

	csvOut.Flush()

}

func getPuzzleDifficultyRatings(result chan map[int]puzzle) {

	var db mysql.Conn

	if useMockData {
		db = &mockConnection{}
	} else {
		db = mysql.New("tcp", "", config.Url, config.Username, config.Password, config.DbName)
	}

	if err := db.Connect(); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	res, err := db.Start("select %s, %s, %s, %s from %s", config.PuzzlesID, config.PuzzlesDifficulty, config.PuzzlesName, config.PuzzlesPuzzle, config.PuzzlesTable)

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

	puzzles := make(map[int]puzzle)

	for {

		row, _ := res.GetRow()

		if row == nil {
			break
		}

		puzzles[row.Int(0)] = puzzle{id: row.Int(0), difficultyRating: row.Int(1), name: row.Str(2), puzzle: row.Str(3)}
	}

	result <- puzzles

}
<commit_msg>Normalize each row so that its probabilities sum to 1.<commit_after>package main

import (
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/ziutek/mymysql/mysql"
	_ "github.com/ziutek/mymysql/native"
	"log"
	"os"
	"sort"
	"strconv"
)

const _DB_CONFIG_FILENAME = "db_config.SECRET.json"
const _OUTPUT_FILENAME = "output.csv"
const _QUERY_LIMIT = 100
const _PENALTY_PERCENTAGE_CUTOFF = 0.10

//How many solves a user must have to have their relative scale included.
//A low value gives you far more very low or very high scores than you shoul get.
const _MINIMUM_SOLVES = 10

var noLimitFlag bool
var printPuzzleDataFlag bool
var cullCheaterPercentageFlag float64
var minimumSolvesFlag int
var useMockData bool
var queryLimit int

func init() {
	flag.BoolVar(&noLimitFlag, "a", false, "Specify to execute the solves query with no limit.")
	flag.BoolVar(&printPuzzleDataFlag, "p", false, "Specify that you want puzzle data printed out in the output.")
	flag.Float64Var(&cullCheaterPercentageFlag, "c", _PENALTY_PERCENTAGE_CUTOFF, "What percentage of solve time must be penalty for someone to be considered a cheater.")
	flag.IntVar(&minimumSolvesFlag, "s", _MINIMUM_SOLVES, "How many solves a user must have their scores considered.")
	flag.IntVar(&queryLimit, "n", _QUERY_LIMIT, "Number of solves to fetch from the database.")
	flag.BoolVar(&useMockData, "m", false, "Use mock data (useful if you don't have a real database to test with).")
}

type dbConfig struct {
	Url               string
	Username          string
	Password          string
	DbName            string
	SolvesTable       string
	SolvesID          string
	SolvesPuzzleID    string
	SolvesTotalTime   string
	SolvesPenaltyTime string
	SolvesUser        string
	PuzzlesTable      string
	PuzzlesID         string
	PuzzlesDifficulty string
	PuzzlesName       string
	PuzzlesPuzzle     string
}

var config dbConfig

type solve struct {
	puzzleID    int
	totalTime   int
	penaltyTime int
}

type userSolvesCollection struct {
	solves     []solve
	max        int
	min        int
	idPosition map[int]int
}

type puzzle struct {
	id                     int
	userRelativeDifficulty float32
	difficultyRating       int
	name                   string
	puzzle                 string
}

type puzzles []puzzle

type byUserRelativeDifficulty struct {
	puzzles
}

func (self puzzles) Len() int {
	return len(self)
}

func (self puzzles) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func (self byUserRelativeDifficulty) Less(i, j int) bool {
	return self.puzzles[i].userRelativeDifficulty < self.puzzles[j].userRelativeDifficulty
}

type bySolveTimeDsc []solve

func (self bySolveTimeDsc) Len() int {
	return len(self)
}

func (self bySolveTimeDsc) Swap(i, j int) {
	self[i], self[j] = self[j], self[i]
}

func (self bySolveTimeDsc) Less(i, j int) bool {
	//For the purposes of this algorithm, the "best" has to be lowest rank.
	return self[i].totalTime > self[j].totalTime
}

func (self *userSolvesCollection) addSolve(solve solve) bool {
	//Cull obviously incorrect solves.
	if solve.totalTime == 0 {
		return false
	}

	//Cull solves that leaned too heavily on hints.
	if float64(solve.penaltyTime)/float64(solve.totalTime) > cullCheaterPercentageFlag {
		return false
	}

	self.solves = append(self.solves, solve)
	if len(self.solves) == 1 {
		self.max = solve.totalTime
		self.min = solve.totalTime
	} else {
		if self.max < solve.totalTime {
			self.max = solve.totalTime
		}
		if self.min > solve.totalTime {
			self.min = solve.totalTime
		}
	}
	return true
}

//Whehter or not this should be included in calculation.
//Basically, whether the reltaiveDifficulties will all be valid.
//Normally this returns false if there is only one solve by the user, but could also
//happen when there are multiple solves but (crazily enough) they all have exactly the same solveTime.
//This DOES happen in the production dataset.
func (self *userSolvesCollection) valid() bool {
	if self.max == self.min {
		return false
	}

	if len(self.solves) < minimumSolvesFlag {
		return false
	}

	return true
}

func (self *userSolvesCollection) relativeDifficulties() map[int]float32 {
	//Returns a map of puzzle id to relative difficulty, normalized by our max and min.
	avgSolveTimes := make(map[int]float32)
	//Keep track of how many times we've seen each puzzle solved by this user so we can do correct averaging.
	avgSolveTimesCount := make(map[int]int)

	//First, collect the average solve time (in case the same user has solved more than once the same puzzle)

	for _, solve := range self.solves {
		currentAvgSolveTime := avgSolveTimes[solve.puzzleID]

		avgSolveTimes[solve.puzzleID] = (currentAvgSolveTime*float32(avgSolveTimesCount[solve.puzzleID]) + float32(solve.totalTime)) / float32(avgSolveTimesCount[solve.puzzleID]+1)

		avgSolveTimesCount[solve.puzzleID]++
	}

	//Now, relativize all of the scores.

	result := make(map[int]float32)

	for puzzleID, avgSolveTime := range avgSolveTimes {
		result[puzzleID] = (avgSolveTime - float32(self.min)) / float32(self.max-self.min)
	}

	return result
}

func main() {

	flag.Parse()

	file, err := os.Open(_DB_CONFIG_FILENAME)
	if err != nil {
		log.Fatal("Could not find the config file at ", _DB_CONFIG_FILENAME, ". You should copy the SAMPLE one to that filename and configure.")
		os.Exit(1)
	}
	defer file.Close()
	decoder := json.NewDecoder(file)
	if err := decoder.Decode(&config); err != nil {
		log.Fatal("There was an error parsing JSON from the config file: ", err)
		os.Exit(1)
	}

	difficutlyRatingsChan := make(chan map[int]puzzle)

	go getPuzzleDifficultyRatings(difficutlyRatingsChan)

	var db mysql.Conn

	if useMockData {
		db = &mockConnection{}
	} else {
		db = mysql.New("tcp", "", config.Url, config.Username, config.Password, config.DbName)
	}

	if err := db.Connect(); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	var solvesQuery string

	if noLimitFlag {
		log.Println("Running without a limit for number of solves to retrieve.")
		solvesQuery = "select %s, %s, %s, %s from %s"
	} else {
		log.Println("Running with a limit of ", queryLimit, " for number of solves to retrieve.")
		solvesQuery = "select %s, %s, %s, %s from %s limit " + strconv.Itoa(queryLimit)
	}

	res, err := db.Start(solvesQuery, config.SolvesUser, config.SolvesPuzzleID, config.SolvesTotalTime, config.SolvesPenaltyTime, config.SolvesTable)

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

	solvesByUser := make(map[string]*userSolvesCollection)

	var userSolves *userSolvesCollection
	var ok bool
	var i int
	var skippedSolves int
	var skippedDuplicateSolves int

	seenRows := make(map[string]bool)

	//First, process all user records in the DB to collect all solves by userName.
	for {

		row, _ := res.GetRow()

		if row == nil {
			break
		}

		i++

		rowHashValue := fmt.Sprintf("%v", row)

		if _, seen := seenRows[rowHashValue]; seen {
			skippedDuplicateSolves++
			continue
		} else {
			seenRows[rowHashValue] = true
		}

		userSolves, ok = solvesByUser[row.Str(0)]

		if !ok {
			userSolves = new(userSolvesCollection)
			userSolves.idPosition = make(map[int]int)
			solvesByUser[row.Str(0)] = userSolves
		}

		if !userSolves.addSolve(solve{row.Int(1), row.Int(2), row.Int(3)}) {
			skippedSolves++
		}

	}

	log.Println("Processed ", i, " solves by ", len(solvesByUser), " users.")
	log.Println("Skipped ", skippedSolves, " solves that cheated too much.")
	log.Println("Skipped ", skippedDuplicateSolves, " solves because they were duplicates of solves seen earlier.")

	//Now get the relative difficulty for each user's puzzles, and collect them.

	relativeDifficultiesByPuzzle := make(map[int][]float32)

	collectionByPuzzle := make(map[int]map[*userSolvesCollection]bool)

	var skippedUsers int

	for _, collection := range solvesByUser {

		/*
			//TODO: consider removing this logic and all of skipped users totally.
				if !collection.valid() {
					skippedUsers++
					continue
				}
		*/

		for puzzleID, relativeDifficulty := range collection.relativeDifficulties() {
			relativeDifficultiesByPuzzle[puzzleID] = append(relativeDifficultiesByPuzzle[puzzleID], relativeDifficulty)
		}

		sort.Sort(bySolveTimeDsc(collection.solves))

		for i, puzzle := range collection.solves {
			collection.idPosition[puzzle.puzzleID] = i

			collectionMap, ok := collectionByPuzzle[puzzle.puzzleID]
			if !ok {
				collectionMap = make(map[*userSolvesCollection]bool)
			}

			collectionMap[collection] = true

			collectionByPuzzle[puzzle.puzzleID] = collectionMap
		}

	}

	//Now, create the Markov Transition Matrix, according to algorithm MC4 of http://www.wisdom.weizmann.ac.il/~naor/PAPERS/rank_www10.html
	//We start by creating a stacked array of float64's that we'll pass to the matrix library.

	numPuzzles := len(collectionByPuzzle)

	matrixData := make([][]float64, numPuzzles)

	for i := range matrixData {
		matrixData[i] = make([]float64, numPuzzles)
	}

	//Now we will associate each observed puzzleID with an index that it will be associated with in the matrix.
	puzzleIndex := make([]int, numPuzzles)

	counter := 0
	for key, _ := range collectionByPuzzle {
		puzzleIndex[counter] = key
		counter++
	}

	//Now we start to build up the matrix according to the MC4 algorithm.

	for i := 0; i < numPuzzles; i++ {
		for j := 0; j < numPuzzles; j++ {
			//Convert the zero-index into the puzzle ID we're actually interested in.
			p := puzzleIndex[i]
			q := puzzleIndex[j]

			//Find the intersection of userSolveCollections that contain both p and q.
			pMap := collectionByPuzzle[p]
			qMap := collectionByPuzzle[q]

			var intersection []*userSolvesCollection
			for collection, _ := range pMap {
				if _, ok := qMap[collection]; ok {
					intersection = append(intersection, collection)
				}
			}

			//Next, calculate how many of the collections have q ranked better (lower!) than p.
			count := 0
			for _, collection := range intersection {
				if collection.idPosition[q] < collection.idPosition[p] {
					count++
				}
			}

			//Is it a majority? if so, transition. if not, leave at 0.
			if count > (len(intersection) / 2) {
				matrixData[i][j] = 1.0
			}

		}
	}

	//Go through and normalize the probabilities in each row to sum to 1.
	for i := 0; i < numPuzzles; i++ {
		//Count the number of rows that are 1.0.
		count := 0
		for j := 0; j < numPuzzles; j++ {
			if matrixData[i][j] > 0.0 {
				count++
			}
		}
		probability := 1.0 / float64(count)
		for j := 0; j < numPuzzles; j++ {
			if matrixData[i][j] > 0.0 {
				matrixData[i][j] = probability
			}
		}
	}

	log.Println("Skipped ", skippedUsers, " users because they did not have enough solve times.")

	puzzles := make([]puzzle, len(relativeDifficultiesByPuzzle))

	var index int

	for puzzleID, difficulties := range relativeDifficultiesByPuzzle {
		var sum float32
		for _, difficulty := range difficulties {
			sum += difficulty
		}
		puzzles[index] = puzzle{id: puzzleID, userRelativeDifficulty: sum / float32(len(difficulties)), difficultyRating: -1}
		index++
	}

	//Sort the puzzles by relative user difficulty
	//We actually don't need the wrapper, since it will modify the underlying slice.
	sort.Sort(byUserRelativeDifficulty{puzzles})

	//Merge in the difficulty ratings from the server.
	difficultyRatings := <-difficutlyRatingsChan

	for i, puzzle := range puzzles {
		info, ok := difficultyRatings[puzzle.id]
		if ok {
			puzzle.difficultyRating = info.difficultyRating
			puzzle.name = info.name
			puzzle.puzzle = info.puzzle
		}
		//It's not a pointer so we have to copy it back.
		puzzles[i] = puzzle
	}

	//Now print the results to stdout.

	csvOut := csv.NewWriter(os.Stdout)

	for _, puzzle := range puzzles {
		temp := []string{strconv.Itoa(puzzle.id), strconv.Itoa(puzzle.difficultyRating), fmt.Sprintf("%g", puzzle.userRelativeDifficulty), puzzle.name}
		if printPuzzleDataFlag {
			temp = append(temp, puzzle.puzzle)
		}
		csvOut.Write(temp)
	}

	csvOut.Flush()

}

func getPuzzleDifficultyRatings(result chan map[int]puzzle) {

	var db mysql.Conn

	if useMockData {
		db = &mockConnection{}
	} else {
		db = mysql.New("tcp", "", config.Url, config.Username, config.Password, config.DbName)
	}

	if err := db.Connect(); err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	res, err := db.Start("select %s, %s, %s, %s from %s", config.PuzzlesID, config.PuzzlesDifficulty, config.PuzzlesName, config.PuzzlesPuzzle, config.PuzzlesTable)

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

	puzzles := make(map[int]puzzle)

	for {

		row, _ := res.GetRow()

		if row == nil {
			break
		}

		puzzles[row.Int(0)] = puzzle{id: row.Int(0), difficultyRating: row.Int(1), name: row.Str(2), puzzle: row.Str(3)}
	}

	result <- puzzles

}
<|endoftext|>


<commit_before>package acme

import (
	"crypto/rsa"
	"crypto/tls"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
	"testing"

	"github.com/square/go-jose"
)

func TestHTTPNonRootBind(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 128)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey)}

	solver := &httpChallenge{jws: jws}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: "localhost:4000", Token: "http1"}

	// validate error on non-root bind to 443
	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("BIND: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "Could not start HTTP server for challenge -> listen tcp :80: bind: permission denied"
		if err.Error() != expectedError {
			t.Errorf("Expected error \"%s\" but instead got \"%s\"", expectedError, err.Error())
		}
	}
}

func TestHTTPShortRSA(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 128)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey), nonces: []string{"test1", "test2"}}

	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: "http://localhost:4000", Token: "http2"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "Failed to post JWS message. -> crypto/rsa: message too long for RSA public key size"
		if err.Error() != expectedError {
			t.Errorf("Expected error %s but instead got %s", expectedError, err.Error())
		}
	}
}

func TestHTTPConnectionRefusal(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey), nonces: []string{"test1", "test2"}}

	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: "http://localhost:4000", Token: "http3"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		reg := "/Failed to post JWS message\\. -> Post http:\\/\\/localhost:4000: dial tcp 127\\.0\\.0\\.1:4000: (getsockopt: )?connection refused/g"
		test2 := "Failed to post JWS message. -> Post http://localhost:4000: dial tcp 127.0.0.1:4000: connection refused"
		r, _ := regexp.Compile(reg)
		if r.MatchString(err.Error()) && r.MatchString(test2) {
			t.Errorf("Expected \"%s\" to match %s", err.Error(), reg)
		}
	}
}

func TestHTTPUnexpectedServerState(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"simpleHttp\",\"status\":\"what\",\"uri\":\"http://some.url\",\"token\":\"http4\"}"))
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: ts.URL, Token: "http4"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "The server returned an unexpected state."
		if err.Error() != expectedError {
			t.Errorf("Expected error %s but instead got %s", expectedError, err.Error())
		}
	}
}

func TestHTTPChallengeServerUnexpectedDomain(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey)}

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "POST" {
			tr := &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			}
			client := &http.Client{Transport: tr}
			req, _ := client.Get("https://localhost:23456/.well-known/acme-challenge/" + "htto5")
			reqBytes, _ := ioutil.ReadAll(req.Body)
			if string(reqBytes) != "TEST" {
				t.Error("Expected simpleHTTP server to return string TEST on unexpected domain.")
			}
		}

		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"simpleHttp\",\"status\":\"invalid\",\"uri\":\"http://some.url\",\"token\":\"http5\"}"))
	}))

	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: ts.URL, Token: "http5"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	}
}

func TestHTTPServerError(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "HEAD" {
			w.Header().Add("Replay-Nonce", "12345")
		} else {
			w.WriteHeader(http.StatusInternalServerError)
			w.Header().Add("Replay-Nonce", "12345")
			w.Write([]byte("{\"type\":\"urn:acme:error:unauthorized\",\"detail\":\"Error creating new authz :: Syntax error\"}"))
		}
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: ts.URL, Token: "http6"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "acme: Error 500 - urn:acme:error:unauthorized - Error creating new authz :: Syntax error"
		if err.Error() != expectedError {
			t.Errorf("Expected error |%s| but instead got |%s|", expectedError, err.Error())
		}
	}
}

func TestHTTPInvalidServerState(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"simpleHttp\",\"status\":\"invalid\",\"uri\":\"http://some.url\",\"token\":\"http7\"}"))
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: ts.URL, Token: "http7"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "The server could not validate our request."
		if err.Error() != expectedError {
			t.Errorf("Expected error |%s| but instead got |%s|", expectedError, err.Error())
		}
	}
}

func TestHTTPValidServerResponse(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"simpleHttp\",\"status\":\"valid\",\"uri\":\"http://some.url\",\"token\":\"http8\"}"))
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: ts.URL, Token: "http8"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err != nil {
		t.Errorf("VALID: Expected Solve to return no error but the error was -> %v", err)
	}
}

func TestHTTPValidFull(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(nil)

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23457"}
	clientChallenge := challenge{Type: "simpleHttp", Status: "pending", URI: ts.URL, Token: "http9"}

	// Validate server on port 23456 which responds appropriately
	ts.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var request challenge
		w.Header().Add("Replay-Nonce", "12345")

		if r.Method == "HEAD" {
			return
		}

		clientJws, _ := ioutil.ReadAll(r.Body)
		j, err := jose.ParseSigned(string(clientJws))
		if err != nil {
			t.Errorf("Client sent invalid JWS to the server.\n\t%v", err)
			return
		}
		output, err := j.Verify(&privKey.(*rsa.PrivateKey).PublicKey)
		if err != nil {
			t.Errorf("Unable to verify client data -> %v", err)
		}
		json.Unmarshal(output, &request)

		transport := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
		client := &http.Client{Transport: transport}

		reqURL := "http://localhost:23457/.well-known/acme-challenge/" + clientChallenge.Token
		t.Logf("Request URL is: %s", reqURL)
		req, err := http.NewRequest("GET", reqURL, nil)
		if err != nil {
			t.Error(err)
		}
		req.Host = "127.0.0.1"
		resp, err := client.Do(req)
		if err != nil {
			t.Errorf("Expected the solver to listen on port 23457 -> %v", err)
		}
		defer resp.Body.Close()

		body, _ := ioutil.ReadAll(resp.Body)
		bodyStr := string(body)

		if resp.Header.Get("Content-Type") != "text/plain" {
			t.Errorf("Expected server to respond with content type text/plain.")
		}

		tokenRegex := regexp.MustCompile("^[\\w-]{43}$")
		parts := strings.Split(bodyStr, ".")

		if len(parts) != 2 {
			t.Errorf("Expected server token to be a composite of two strings, seperated by a dot")
		}

		if parts[0] != clientChallenge.Token {
			t.Errorf("Expected the first part of the server token to be the challenge token.")
		}

		if !tokenRegex.MatchString(parts[1]) {
			t.Errorf("Expected the second part of the server token to be a properly formatted key authorization")
		}

		valid := challenge{Type: "simpleHttp", Status: "valid", URI: ts.URL, Token: "1234567812"}
		jsonBytes, _ := json.Marshal(&valid)
		w.Write(jsonBytes)
	})

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err != nil {
		t.Errorf("VALID: Expected Solve to return no error but the error was -> %v", err)
	}
}
<commit_msg>Properly name challenge types for http-01 tests<commit_after>package acme

import (
	"crypto/rsa"
	"crypto/tls"
	"encoding/json"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"regexp"
	"strings"
	"testing"

	"github.com/square/go-jose"
)

func TestHTTPNonRootBind(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 128)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey)}

	solver := &httpChallenge{jws: jws}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: "localhost:4000", Token: "http1"}

	// validate error on non-root bind to 443
	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("BIND: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "Could not start HTTP server for challenge -> listen tcp :80: bind: permission denied"
		if err.Error() != expectedError {
			t.Errorf("Expected error \"%s\" but instead got \"%s\"", expectedError, err.Error())
		}
	}
}

func TestHTTPShortRSA(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 128)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey), nonces: []string{"test1", "test2"}}

	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: "http://localhost:4000", Token: "http2"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "Failed to post JWS message. -> crypto/rsa: message too long for RSA public key size"
		if err.Error() != expectedError {
			t.Errorf("Expected error %s but instead got %s", expectedError, err.Error())
		}
	}
}

func TestHTTPConnectionRefusal(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey), nonces: []string{"test1", "test2"}}

	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: "http://localhost:4000", Token: "http3"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		reg := "/Failed to post JWS message\\. -> Post http:\\/\\/localhost:4000: dial tcp 127\\.0\\.0\\.1:4000: (getsockopt: )?connection refused/g"
		test2 := "Failed to post JWS message. -> Post http://localhost:4000: dial tcp 127.0.0.1:4000: connection refused"
		r, _ := regexp.Compile(reg)
		if r.MatchString(err.Error()) && r.MatchString(test2) {
			t.Errorf("Expected \"%s\" to match %s", err.Error(), reg)
		}
	}
}

func TestHTTPUnexpectedServerState(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"http01\",\"status\":\"what\",\"uri\":\"http://some.url\",\"token\":\"http4\"}"))
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: ts.URL, Token: "http4"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "The server returned an unexpected state."
		if err.Error() != expectedError {
			t.Errorf("Expected error %s but instead got %s", expectedError, err.Error())
		}
	}
}

func TestHTTPChallengeServerUnexpectedDomain(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)
	jws := &jws{privKey: privKey.(*rsa.PrivateKey)}

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "POST" {
			tr := &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			}
			client := &http.Client{Transport: tr}
			req, _ := client.Get("https://localhost:23456/.well-known/acme-challenge/" + "htto5")
			reqBytes, _ := ioutil.ReadAll(req.Body)
			if string(reqBytes) != "TEST" {
				t.Error("Expected http01 server to return string TEST on unexpected domain.")
			}
		}

		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"http01\",\"status\":\"invalid\",\"uri\":\"http://some.url\",\"token\":\"http5\"}"))
	}))

	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: ts.URL, Token: "http5"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	}
}

func TestHTTPServerError(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == "HEAD" {
			w.Header().Add("Replay-Nonce", "12345")
		} else {
			w.WriteHeader(http.StatusInternalServerError)
			w.Header().Add("Replay-Nonce", "12345")
			w.Write([]byte("{\"type\":\"urn:acme:error:unauthorized\",\"detail\":\"Error creating new authz :: Syntax error\"}"))
		}
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: ts.URL, Token: "http6"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "acme: Error 500 - urn:acme:error:unauthorized - Error creating new authz :: Syntax error"
		if err.Error() != expectedError {
			t.Errorf("Expected error |%s| but instead got |%s|", expectedError, err.Error())
		}
	}
}

func TestHTTPInvalidServerState(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"http01\",\"status\":\"invalid\",\"uri\":\"http://some.url\",\"token\":\"http7\"}"))
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: ts.URL, Token: "http7"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err == nil {
		t.Error("UNEXPECTED: Expected Solve to return an error but the error was nil.")
	} else {
		expectedError := "The server could not validate our request."
		if err.Error() != expectedError {
			t.Errorf("Expected error |%s| but instead got |%s|", expectedError, err.Error())
		}
	}
}

func TestHTTPValidServerResponse(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Replay-Nonce", "12345")
		w.Write([]byte("{\"type\":\"http01\",\"status\":\"valid\",\"uri\":\"http://some.url\",\"token\":\"http8\"}"))
	}))

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23456"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: ts.URL, Token: "http8"}

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err != nil {
		t.Errorf("VALID: Expected Solve to return no error but the error was -> %v", err)
	}
}

func TestHTTPValidFull(t *testing.T) {
	privKey, _ := generatePrivateKey(rsakey, 512)

	ts := httptest.NewServer(nil)

	jws := &jws{privKey: privKey.(*rsa.PrivateKey), directoryURL: ts.URL}
	solver := &httpChallenge{jws: jws, optPort: "23457"}
	clientChallenge := challenge{Type: "http01", Status: "pending", URI: ts.URL, Token: "http9"}

	// Validate server on port 23456 which responds appropriately
	ts.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var request challenge
		w.Header().Add("Replay-Nonce", "12345")

		if r.Method == "HEAD" {
			return
		}

		clientJws, _ := ioutil.ReadAll(r.Body)
		j, err := jose.ParseSigned(string(clientJws))
		if err != nil {
			t.Errorf("Client sent invalid JWS to the server.\n\t%v", err)
			return
		}
		output, err := j.Verify(&privKey.(*rsa.PrivateKey).PublicKey)
		if err != nil {
			t.Errorf("Unable to verify client data -> %v", err)
		}
		json.Unmarshal(output, &request)

		transport := &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
		client := &http.Client{Transport: transport}

		reqURL := "http://localhost:23457/.well-known/acme-challenge/" + clientChallenge.Token
		t.Logf("Request URL is: %s", reqURL)
		req, err := http.NewRequest("GET", reqURL, nil)
		if err != nil {
			t.Error(err)
		}
		req.Host = "127.0.0.1"
		resp, err := client.Do(req)
		if err != nil {
			t.Errorf("Expected the solver to listen on port 23457 -> %v", err)
		}
		defer resp.Body.Close()

		body, _ := ioutil.ReadAll(resp.Body)
		bodyStr := string(body)

		if resp.Header.Get("Content-Type") != "text/plain" {
			t.Errorf("Expected server to respond with content type text/plain.")
		}

		tokenRegex := regexp.MustCompile("^[\\w-]{43}$")
		parts := strings.Split(bodyStr, ".")

		if len(parts) != 2 {
			t.Errorf("Expected server token to be a composite of two strings, seperated by a dot")
		}

		if parts[0] != clientChallenge.Token {
			t.Errorf("Expected the first part of the server token to be the challenge token.")
		}

		if !tokenRegex.MatchString(parts[1]) {
			t.Errorf("Expected the second part of the server token to be a properly formatted key authorization")
		}

		valid := challenge{Type: "http01", Status: "valid", URI: ts.URL, Token: "1234567812"}
		jsonBytes, _ := json.Marshal(&valid)
		w.Write(jsonBytes)
	})

	if err := solver.Solve(clientChallenge, "127.0.0.1"); err != nil {
		t.Errorf("VALID: Expected Solve to return no error but the error was -> %v", err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright IBM Corp. 2017 All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

                 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package channel

import (
	"fmt"
	"time"

	"github.com/hyperledger/fabric/common/localmsp"
	"github.com/hyperledger/fabric/protos/common"
	ab "github.com/hyperledger/fabric/protos/orderer"
	"github.com/hyperledger/fabric/protos/utils"
	"google.golang.org/grpc"
)

type deliverClientIntf interface {
	getSpecifiedBlock(num uint64) (*common.Block, error)
	getOldestBlock() (*common.Block, error)
	getNewestBlock() (*common.Block, error)
	Close() error
}

type deliverClient struct {
	conn    *grpc.ClientConn
	client  ab.AtomicBroadcast_DeliverClient
	chainID string
}

func newDeliverClient(conn *grpc.ClientConn, client ab.AtomicBroadcast_DeliverClient, chainID string) *deliverClient {
	return &deliverClient{conn: conn, client: client, chainID: chainID}
}

func seekHelper(chainID string, position *ab.SeekPosition) *common.Envelope {
	seekInfo := &ab.SeekInfo{
		Start:    position,
		Stop:     position,
		Behavior: ab.SeekInfo_BLOCK_UNTIL_READY,
	}

	//TODO- epoch and msgVersion may need to be obtained for nowfollowing usage in orderer/configupdate/configupdate.go
	msgVersion := int32(0)
	epoch := uint64(0)
	env, err := utils.CreateSignedEnvelope(common.HeaderType_CONFIG_UPDATE, chainID, localmsp.NewSigner(), seekInfo, msgVersion, epoch)
	if err != nil {
		logger.Errorf("Error signing envelope:  %s", err)
		return nil
	}
	return env
}

func (r *deliverClient) seekSpecified(blockNumber uint64) error {
	return r.client.Send(seekHelper(r.chainID, &ab.SeekPosition{Type: &ab.SeekPosition_Specified{Specified: &ab.SeekSpecified{Number: blockNumber}}}))
}

func (r *deliverClient) seekOldest() error {
	return r.client.Send(seekHelper(r.chainID, &ab.SeekPosition{Type: &ab.SeekPosition_Oldest{Oldest: &ab.SeekOldest{}}}))
}

func (r *deliverClient) seekNewest() error {
	return r.client.Send(seekHelper(r.chainID, &ab.SeekPosition{Type: &ab.SeekPosition_Newest{Newest: &ab.SeekNewest{}}}))
}

func (r *deliverClient) readBlock() (*common.Block, error) {
	msg, err := r.client.Recv()
	if err != nil {
		return nil, fmt.Errorf("Error receiving: %s", err)
	}

	switch t := msg.Type.(type) {
	case *ab.DeliverResponse_Status:
		logger.Debugf("Got status:%T ", t)
		return nil, fmt.Errorf("can't read the block")
	case *ab.DeliverResponse_Block:
		logger.Debugf("Received block:%v ", t.Block.Header.Number)
		r.client.Recv() // Flush the success message
		return t.Block, nil
	default:
		return nil, fmt.Errorf("response error")
	}
}

func (r *deliverClient) getSpecifiedBlock(num uint64) (*common.Block, error) {
	err := r.seekSpecified(num)
	if err != nil {
		logger.Errorf("Received error: %s", err)
		return nil, err
	}

	return r.readBlock()
}

func (r *deliverClient) getOldestBlock() (*common.Block, error) {
	err := r.seekOldest()
	if err != nil {
		return nil, fmt.Errorf("Received error: %s ", err)
	}

	return r.readBlock()
}

func (r *deliverClient) getNewestBlock() (*common.Block, error) {
	err := r.seekNewest()
	if err != nil {
		logger.Errorf("Received error: %s", err)
		return nil, err
	}

	return r.readBlock()
}

func (r *deliverClient) Close() error {
	return r.conn.Close()
}

func getGenesisBlock(cf *ChannelCmdFactory) (*common.Block, error) {
	timer := time.NewTimer(time.Second * time.Duration(timeout))
	defer timer.Stop()

	for {
		select {
		case <-timer.C:
			cf.DeliverClient.Close()
			return nil, fmt.Errorf("timeout waiting for channel creation")
		default:
			if block, err := cf.DeliverClient.getSpecifiedBlock(0); err != nil {
				cf.DeliverClient.Close()
				cf, err = InitCmdFactory(EndorserNotRequired, OrdererRequired)
				if err != nil {
					return nil, fmt.Errorf("failed connecting: %v", err)
				}
				time.Sleep(200 * time.Millisecond)
			} else {
				cf.DeliverClient.Close()
				return block, nil
			}
		}
	}
}
<commit_msg>[FAB-5365] Fix bad error in peer CLI Deliver<commit_after>/*
Copyright IBM Corp. 2017 All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

                 http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package channel

import (
	"fmt"
	"time"

	"github.com/hyperledger/fabric/common/localmsp"
	"github.com/hyperledger/fabric/protos/common"
	ab "github.com/hyperledger/fabric/protos/orderer"
	"github.com/hyperledger/fabric/protos/utils"
	"google.golang.org/grpc"
)

type deliverClientIntf interface {
	getSpecifiedBlock(num uint64) (*common.Block, error)
	getOldestBlock() (*common.Block, error)
	getNewestBlock() (*common.Block, error)
	Close() error
}

type deliverClient struct {
	conn    *grpc.ClientConn
	client  ab.AtomicBroadcast_DeliverClient
	chainID string
}

func newDeliverClient(conn *grpc.ClientConn, client ab.AtomicBroadcast_DeliverClient, chainID string) *deliverClient {
	return &deliverClient{conn: conn, client: client, chainID: chainID}
}

func seekHelper(chainID string, position *ab.SeekPosition) *common.Envelope {
	seekInfo := &ab.SeekInfo{
		Start:    position,
		Stop:     position,
		Behavior: ab.SeekInfo_BLOCK_UNTIL_READY,
	}

	//TODO- epoch and msgVersion may need to be obtained for nowfollowing usage in orderer/configupdate/configupdate.go
	msgVersion := int32(0)
	epoch := uint64(0)
	env, err := utils.CreateSignedEnvelope(common.HeaderType_CONFIG_UPDATE, chainID, localmsp.NewSigner(), seekInfo, msgVersion, epoch)
	if err != nil {
		logger.Errorf("Error signing envelope:  %s", err)
		return nil
	}
	return env
}

func (r *deliverClient) seekSpecified(blockNumber uint64) error {
	return r.client.Send(seekHelper(r.chainID, &ab.SeekPosition{Type: &ab.SeekPosition_Specified{Specified: &ab.SeekSpecified{Number: blockNumber}}}))
}

func (r *deliverClient) seekOldest() error {
	return r.client.Send(seekHelper(r.chainID, &ab.SeekPosition{Type: &ab.SeekPosition_Oldest{Oldest: &ab.SeekOldest{}}}))
}

func (r *deliverClient) seekNewest() error {
	return r.client.Send(seekHelper(r.chainID, &ab.SeekPosition{Type: &ab.SeekPosition_Newest{Newest: &ab.SeekNewest{}}}))
}

func (r *deliverClient) readBlock() (*common.Block, error) {
	msg, err := r.client.Recv()
	if err != nil {
		return nil, fmt.Errorf("Error receiving: %s", err)
	}

	switch t := msg.Type.(type) {
	case *ab.DeliverResponse_Status:
		logger.Debugf("Got status: %v", t)
		return nil, fmt.Errorf("can't read the block: %v", t)
	case *ab.DeliverResponse_Block:
		logger.Debugf("Received block: %v", t.Block.Header.Number)
		r.client.Recv() // Flush the success message
		return t.Block, nil
	default:
		return nil, fmt.Errorf("response error: unknown type %T", t)
	}
}

func (r *deliverClient) getSpecifiedBlock(num uint64) (*common.Block, error) {
	err := r.seekSpecified(num)
	if err != nil {
		logger.Errorf("Received error: %s", err)
		return nil, err
	}

	return r.readBlock()
}

func (r *deliverClient) getOldestBlock() (*common.Block, error) {
	err := r.seekOldest()
	if err != nil {
		return nil, fmt.Errorf("Received error: %s ", err)
	}

	return r.readBlock()
}

func (r *deliverClient) getNewestBlock() (*common.Block, error) {
	err := r.seekNewest()
	if err != nil {
		logger.Errorf("Received error: %s", err)
		return nil, err
	}

	return r.readBlock()
}

func (r *deliverClient) Close() error {
	return r.conn.Close()
}

func getGenesisBlock(cf *ChannelCmdFactory) (*common.Block, error) {
	timer := time.NewTimer(time.Second * time.Duration(timeout))
	defer timer.Stop()

	for {
		select {
		case <-timer.C:
			cf.DeliverClient.Close()
			return nil, fmt.Errorf("timeout waiting for channel creation")
		default:
			if block, err := cf.DeliverClient.getSpecifiedBlock(0); err != nil {
				cf.DeliverClient.Close()
				cf, err = InitCmdFactory(EndorserNotRequired, OrdererRequired)
				if err != nil {
					return nil, fmt.Errorf("failed connecting: %v", err)
				}
				time.Sleep(200 * time.Millisecond)
			} else {
				cf.DeliverClient.Close()
				return block, nil
			}
		}
	}
}
<|endoftext|>


<commit_before>package libvirt

import (
	"github.com/alexzorin/libvirt-go"
	"github.com/mistifyio/mistify-agent/client"
	"github.com/mistifyio/mistify-agent/log"
	"github.com/mistifyio/mistify-agent/rpc"
	"net/http"
	"syscall"
)

var StateNames = map[int]string{
	libvirt.VIR_DOMAIN_NOSTATE:     "No State",
	libvirt.VIR_DOMAIN_RUNNING:     "Running",
	libvirt.VIR_DOMAIN_BLOCKED:     "Blocked",
	libvirt.VIR_DOMAIN_PAUSED:      "Paused",
	libvirt.VIR_DOMAIN_SHUTDOWN:    "Shutdown",
	libvirt.VIR_DOMAIN_CRASHED:     "Crashed",
	libvirt.VIR_DOMAIN_PMSUSPENDED: "Suspended",
	libvirt.VIR_DOMAIN_SHUTOFF:     "Shutoff",
}

type (
	Libvirt struct {
		uri         string
		connections chan *libvirt.VirConnection
		max         int
	}

	Domain struct {
		*libvirt.VirDomain
		State int
	}
)

func NewLibvirt(uri string, max int) (*Libvirt, error) {
	lv := &Libvirt{
		uri:         uri,
		max:         max,
		connections: make(chan *libvirt.VirConnection, max),
	}

	for i := 0; i < max; i++ {
		v, err := libvirt.NewVirConnection(uri)
		if err != nil {
			return nil, err
		}
		lv.connections <- &v
	}

	return lv, nil
}

func (lv *Libvirt) RunHTTP(port uint) error {
	server, err := rpc.NewServer(int(port))
	if err != nil {
		return err
	}

	server.RegisterService(lv)
	return server.ListenAndServe()
}

func (lv *Libvirt) getConnection() *libvirt.VirConnection {
	conn := <-lv.connections
	defer func() {
		lv.connections <- conn
	}()

	return conn
}

func (lv *Libvirt) LookupDomainByName(name string) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	domain, err := conn.LookupDomainByName(name)
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func (lv *Libvirt) NewDomain(guest *client.Guest) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	xml, err := lv.DomainXML(guest)
	if err != nil {
		return nil, err
	}

	domain, err := conn.DomainDefineXML(xml)
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func GetState(domain *libvirt.VirDomain) (int, error) {
	state, err := domain.GetState()
	if err != nil {
		return libvirt.VIR_DOMAIN_NOSTATE, err
	}

	return state[0], nil
}

func (lv *Libvirt) DomainWrapper(fn func(*libvirt.VirDomain, int) error) func(*http.Request, *rpc.GuestRequest, *rpc.GuestResponse) error {
	return func(r *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
		if request.Guest == nil || request.Guest.Id == "" {
			return syscall.EINVAL
		}

		domain, err := lv.LookupDomainByName(request.Guest.Id)
		if err != nil {
			return err
		}
		defer domain.Free()

		state, err := GetState(domain)
		if err != nil {
			return err
		}

		err = fn(domain, state)
		if err != nil {
			return err
		}

		*response = rpc.GuestResponse{
			Guest: request.Guest,
		}

		state, err = GetState(domain)
		if err != nil {
			return err
		}

		response.Guest.State = StateNames[state]

		return nil
	}
}

// Guest Actions

func (lv *Libvirt) Restart(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Reboot(http, request, response)
}

func (lv *Libvirt) Poweroff(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Shutdown(http, request, response)
}

func (lv *Libvirt) Delete(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Delete %s\n", request.Guest.Id)

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	defer domain.Free()
	if err != nil {
		return err
	}

	err = domain.Destroy()
	if err != nil {
		return err
	}

	*response = rpc.GuestResponse{
		Guest: request.Guest,
	}

	return nil
}

func (lv *Libvirt) Create(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Create %s\n", request.Guest.Id)

	domain, err := lv.NewDomain(request.Guest)
	if err != nil {
		return err
	}
	defer domain.Free()

	err = domain.Create()
	if err != nil {
		return err
	}

	*response = rpc.GuestResponse{
		Guest: request.Guest,
	}

	state, err := GetState(domain)
	if err != nil {
		return err
	}

	response.Guest.State = StateNames[state]

	return nil
}

func (lv *Libvirt) Run(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Run %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		switch state {

		case libvirt.VIR_DOMAIN_RUNNING:
			// nothing to do

		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF, libvirt.VIR_DOMAIN_BLOCKED, libvirt.VIR_DOMAIN_NOSTATE:
			return domain.Create()

		case libvirt.VIR_DOMAIN_PAUSED, libvirt.VIR_DOMAIN_PMSUSPENDED:
			return domain.Resume()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Reboot(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Reboot %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		return domain.Reboot(0)
	})(http, request, response)
}

func (lv *Libvirt) Shutdown(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Shutdown %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {

		switch state {
		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF:
			// nothing to do

		default:
			return domain.Shutdown()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Status(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Status %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		// DomainWrapper gets the state already, no need to do anything here
		return nil
	})(http, request, response)
}

func (lv *Libvirt) CpuMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestMetricsResponse) error {
	var params libvirt.VirTypedParameters
	metrics := make([]*client.GuestCpuMetrics, request.Guest.Cpu)

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	if err != nil {
		return err
	}
	defer domain.Free()

	nparams, err := domain.GetCPUStats(nil, 0, 0, uint32(request.Guest.Cpu), 0)
	if err != nil {
		return err
	}

	_, err = domain.GetCPUStats(&params, nparams, 0, uint32(request.Guest.Cpu), 0)
	if err != nil {
		return err
	}

	*response = rpc.GuestMetricsResponse{
		CPU: metrics,
	}
	return nil
}

func (lv *Libvirt) DiskMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestMetricsResponse) error {
	metrics := make(map[string]*client.GuestDiskMetrics, len(request.Guest.Disks))
	var params libvirt.VirTypedParameters

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	if err != nil {
		return err
	}
	defer domain.Free()

	_, err = domain.BlockStatsFlags("", &params, 0, 0)
	if err != nil {
		return err
	}

	*response = rpc.GuestMetricsResponse{
		Disk: metrics,
	}
	return nil
}

func (lv *Libvirt) NicMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestMetricsResponse) error {
	metrics := make(map[string]*client.GuestNicMetrics, len(request.Guest.Nics))

	*response = rpc.GuestMetricsResponse{
		Nic: metrics,
	}
	return nil
}
<commit_msg>MIST-138: if domain is not running, use Undefine, not Destroy<commit_after>package libvirt

import (
	"github.com/alexzorin/libvirt-go"
	"github.com/mistifyio/mistify-agent/client"
	"github.com/mistifyio/mistify-agent/log"
	"github.com/mistifyio/mistify-agent/rpc"
	"net/http"
	"syscall"
)

var StateNames = map[int]string{
	libvirt.VIR_DOMAIN_NOSTATE:     "No State",
	libvirt.VIR_DOMAIN_RUNNING:     "Running",
	libvirt.VIR_DOMAIN_BLOCKED:     "Blocked",
	libvirt.VIR_DOMAIN_PAUSED:      "Paused",
	libvirt.VIR_DOMAIN_SHUTDOWN:    "Shutdown",
	libvirt.VIR_DOMAIN_CRASHED:     "Crashed",
	libvirt.VIR_DOMAIN_PMSUSPENDED: "Suspended",
	libvirt.VIR_DOMAIN_SHUTOFF:     "Shutoff",
}

type (
	Libvirt struct {
		uri         string
		connections chan *libvirt.VirConnection
		max         int
	}

	Domain struct {
		*libvirt.VirDomain
		State int
	}
)

func NewLibvirt(uri string, max int) (*Libvirt, error) {
	lv := &Libvirt{
		uri:         uri,
		max:         max,
		connections: make(chan *libvirt.VirConnection, max),
	}

	for i := 0; i < max; i++ {
		v, err := libvirt.NewVirConnection(uri)
		if err != nil {
			return nil, err
		}
		lv.connections <- &v
	}

	return lv, nil
}

func (lv *Libvirt) RunHTTP(port uint) error {
	server, err := rpc.NewServer(int(port))
	if err != nil {
		return err
	}

	server.RegisterService(lv)
	return server.ListenAndServe()
}

func (lv *Libvirt) getConnection() *libvirt.VirConnection {
	conn := <-lv.connections
	defer func() {
		lv.connections <- conn
	}()

	return conn
}

func (lv *Libvirt) LookupDomainByName(name string) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	domain, err := conn.LookupDomainByName(name)
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func (lv *Libvirt) NewDomain(guest *client.Guest) (*libvirt.VirDomain, error) {
	conn := lv.getConnection()

	xml, err := lv.DomainXML(guest)
	if err != nil {
		return nil, err
	}

	domain, err := conn.DomainDefineXML(xml)
	if err != nil {
		return nil, err
	}

	return &domain, err
}

func GetState(domain *libvirt.VirDomain) (int, error) {
	state, err := domain.GetState()
	if err != nil {
		return libvirt.VIR_DOMAIN_NOSTATE, err
	}

	return state[0], nil
}

func (lv *Libvirt) DomainWrapper(fn func(*libvirt.VirDomain, int) error) func(*http.Request, *rpc.GuestRequest, *rpc.GuestResponse) error {
	return func(r *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
		if request.Guest == nil || request.Guest.Id == "" {
			return syscall.EINVAL
		}

		domain, err := lv.LookupDomainByName(request.Guest.Id)
		if err != nil {
			return err
		}
		defer domain.Free()

		state, err := GetState(domain)
		if err != nil {
			return err
		}

		err = fn(domain, state)
		if err != nil {
			return err
		}

		*response = rpc.GuestResponse{
			Guest: request.Guest,
		}

		state, err = GetState(domain)
		if err != nil {
			return err
		}

		response.Guest.State = StateNames[state]

		return nil
	}
}

// Guest Actions

func (lv *Libvirt) Restart(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Reboot(http, request, response)
}

func (lv *Libvirt) Poweroff(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	return lv.Shutdown(http, request, response)
}

func (lv *Libvirt) Delete(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Delete %s\n", request.Guest.Id)

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	defer domain.Free()
	if err != nil {
		return err
	}

	state, err := GetState(domain)
	if err != nil {
		return err
	}

	switch state {
	case libvirt.VIR_DOMAIN_RUNNING:
		err = domain.Destroy()
		if err != nil {
			return err
		}

	case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF, libvirt.VIR_DOMAIN_BLOCKED, libvirt.VIR_DOMAIN_NOSTATE, libvirt.VIR_DOMAIN_PAUSED, libvirt.VIR_DOMAIN_PMSUSPENDED:
		err = domain.Undefine()
		if err != nil {
			return err
		}
	}

	*response = rpc.GuestResponse{
		Guest: request.Guest,
	}

	return nil
}

func (lv *Libvirt) Create(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Create %s\n", request.Guest.Id)

	domain, err := lv.NewDomain(request.Guest)
	if err != nil {
		return err
	}
	defer domain.Free()

	err = domain.Create()
	if err != nil {
		return err
	}

	*response = rpc.GuestResponse{
		Guest: request.Guest,
	}

	state, err := GetState(domain)
	if err != nil {
		return err
	}

	response.Guest.State = StateNames[state]

	return nil
}

func (lv *Libvirt) Run(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Run %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		switch state {

		case libvirt.VIR_DOMAIN_RUNNING:
			// nothing to do

		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF, libvirt.VIR_DOMAIN_BLOCKED, libvirt.VIR_DOMAIN_NOSTATE:
			return domain.Create()

		case libvirt.VIR_DOMAIN_PAUSED, libvirt.VIR_DOMAIN_PMSUSPENDED:
			return domain.Resume()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Reboot(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Reboot %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		return domain.Reboot(0)
	})(http, request, response)
}

func (lv *Libvirt) Shutdown(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Shutdown %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {

		switch state {
		case libvirt.VIR_DOMAIN_SHUTDOWN, libvirt.VIR_DOMAIN_SHUTOFF:
			// nothing to do

		default:
			return domain.Shutdown()
		}

		return nil
	})(http, request, response)
}

func (lv *Libvirt) Status(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestResponse) error {
	log.Info("Libvirt.Status %s\n", request.Guest.Id)

	return lv.DomainWrapper(func(domain *libvirt.VirDomain, state int) error {
		// DomainWrapper gets the state already, no need to do anything here
		return nil
	})(http, request, response)
}

func (lv *Libvirt) CpuMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestMetricsResponse) error {
	var params libvirt.VirTypedParameters
	metrics := make([]*client.GuestCpuMetrics, request.Guest.Cpu)

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	if err != nil {
		return err
	}
	defer domain.Free()

	nparams, err := domain.GetCPUStats(nil, 0, 0, uint32(request.Guest.Cpu), 0)
	if err != nil {
		return err
	}

	_, err = domain.GetCPUStats(&params, nparams, 0, uint32(request.Guest.Cpu), 0)
	if err != nil {
		return err
	}

	*response = rpc.GuestMetricsResponse{
		CPU: metrics,
	}
	return nil
}

func (lv *Libvirt) DiskMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestMetricsResponse) error {
	metrics := make(map[string]*client.GuestDiskMetrics, len(request.Guest.Disks))
	var params libvirt.VirTypedParameters

	domain, err := lv.LookupDomainByName(request.Guest.Id)
	if err != nil {
		return err
	}
	defer domain.Free()

	_, err = domain.BlockStatsFlags("", &params, 0, 0)
	if err != nil {
		return err
	}

	*response = rpc.GuestMetricsResponse{
		Disk: metrics,
	}
	return nil
}

func (lv *Libvirt) NicMetrics(http *http.Request, request *rpc.GuestRequest, response *rpc.GuestMetricsResponse) error {
	metrics := make(map[string]*client.GuestNicMetrics, len(request.Guest.Nics))

	*response = rpc.GuestMetricsResponse{
		Nic: metrics,
	}
	return nil
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"

	"github.com/hashicorp/consul/agent/consul/stream"
	"github.com/hashicorp/consul/proto/pbsubscribe"
	"github.com/hashicorp/go-memdb"
)

// ReadTxn is implemented by memdb.Txn to perform read operations.
type ReadTxn interface {
	Get(table, index string, args ...interface{}) (memdb.ResultIterator, error)
	First(table, index string, args ...interface{}) (interface{}, error)
	FirstWatch(table, index string, args ...interface{}) (<-chan struct{}, interface{}, error)
	Abort()
}

// WriteTxn is implemented by memdb.Txn to perform write operations.
type WriteTxn interface {
	ReadTxn
	Insert(table string, obj interface{}) error
	Commit() error
}

// Changes wraps a memdb.Changes to include the index at which these changes
// were made.
type Changes struct {
	// Index is the latest index at the time these changes were committed.
	Index   uint64
	Changes memdb.Changes
}

// changeTrackerDB is a thin wrapper around memdb.DB which enables TrackChanges on
// all write transactions. When the transaction is committed the changes are:
// 1. Used to update our internal usage tracking
// 2. Sent to the eventPublisher which will create and emit change events
type changeTrackerDB struct {
	db             *memdb.MemDB
	publisher      *stream.EventPublisher
	processChanges func(ReadTxn, Changes) ([]stream.Event, error)
}

// Txn exists to maintain backwards compatibility with memdb.DB.Txn. Preexisting
// code may use it to create a read-only transaction, but it will panic if called
// with write=true.
//
// Deprecated: use either ReadTxn, or WriteTxn.
func (c *changeTrackerDB) Txn(write bool) *txn {
	if write {
		panic("don't use db.Txn(true), use db.WriteTxn(idx uin64)")
	}
	return c.ReadTxn()
}

// ReadTxn returns a read-only transaction which behaves exactly the same as
// memdb.Txn
//
// TODO: this could return a regular memdb.Txn if all the state functions accepted
// the ReadTxn interface
func (c *changeTrackerDB) ReadTxn() *txn {
	return &txn{Txn: c.db.Txn(false)}
}

// WriteTxn returns a wrapped memdb.Txn suitable for writes to the state store.
// It will track changes and publish events for the changes when Commit
// is called.
//
// The idx argument must be the index of the current Raft operation. Almost
// all mutations to state should happen as part of a raft apply so the index of
// the log being applied can be passed to WriteTxn.
// The exceptional cases are transactions that are executed on an empty
// memdb.DB as part of Restore, and those executed by tests where we insert
// data directly into the DB. These cases may use WriteTxnRestore.
func (c *changeTrackerDB) WriteTxn(idx uint64) *txn {
	t := &txn{
		Txn:     c.db.Txn(true),
		Index:   idx,
		publish: c.publish,
	}
	t.Txn.TrackChanges()
	return t
}

func (c *changeTrackerDB) publish(tx ReadTxn, changes Changes) error {
	events, err := c.processChanges(tx, changes)
	if err != nil {
		return fmt.Errorf("failed generating events from changes: %v", err)
	}
	c.publisher.Publish(events)
	return nil
}

// WriteTxnRestore returns a wrapped RW transaction that should only be used in
// Restore where we need to replace the entire contents of the Store.
// WriteTxnRestore uses a zero index since the whole restore doesn't really
// occur at one index - the effect is to write many values that were previously
// written across many indexes. WriteTxnRestore also does not publish any
// change events to subscribers.
func (c *changeTrackerDB) WriteTxnRestore() *txn {
	t := &txn{
		Txn:   c.db.Txn(true),
		Index: 0,
	}

	// We enable change tracking so that usage data is correctly populated.
	t.Txn.TrackChanges()
	return t
}

// txn wraps a memdb.Txn to capture changes and send them to the EventPublisher.
//
// This can not be done with txn.Defer because the callback passed to Defer is
// invoked after commit completes, and because the callback can not return an
// error. Any errors from the callback would be lost,  which would result in a
// missing change event, even though the state store had changed.
type txn struct {
	*memdb.Txn
	// Index in raft where the write is occurring. The value is zero for a
	// read-only, or WriteTxnRestore transaction.
	// Index is stored so that it may be passed along to any subscribers as part
	// of a change event.
	Index   uint64
	publish func(tx ReadTxn, changes Changes) error
}

// Commit first pushes changes to EventPublisher, then calls Commit on the
// underlying transaction.
//
// Note that this function, unlike memdb.Txn, returns an error which must be checked
// by the caller. A non-nil error indicates that a commit failed and was not
// applied.
func (tx *txn) Commit() error {
	changes := Changes{
		Index:   tx.Index,
		Changes: tx.Txn.Changes(),
	}

	if len(changes.Changes) > 0 {
		if err := updateUsage(tx, changes); err != nil {
			return err
		}
	}

	// publish may be nil if this is a read-only or WriteTxnRestore transaction.
	// In those cases changes should also be empty, and there will be nothing
	// to publish.
	if tx.publish != nil {
		if err := tx.publish(tx.Txn, changes); err != nil {
			return err
		}
	}

	tx.Txn.Commit()
	return nil
}

var (
	topicServiceHealth        = pbsubscribe.Topic_ServiceHealth
	topicServiceHealthConnect = pbsubscribe.Topic_ServiceHealthConnect
)

func processDBChanges(tx ReadTxn, changes Changes) ([]stream.Event, error) {
	var events []stream.Event
	fns := []func(tx ReadTxn, changes Changes) ([]stream.Event, error){
		aclChangeUnsubscribeEvent,
		ServiceHealthEventsFromChanges,
		// TODO: add other table handlers here.
	}
	for _, fn := range fns {
		e, err := fn(tx, changes)
		if err != nil {
			return nil, err
		}
		events = append(events, e...)
	}
	return events, nil
}

func newSnapshotHandlers(s *Store) stream.SnapshotHandlers {
	return stream.SnapshotHandlers{
		topicServiceHealth:        serviceHealthSnapshot(s, topicServiceHealth),
		topicServiceHealthConnect: serviceHealthSnapshot(s, topicServiceHealthConnect),
	}
}
<commit_msg>state: disable streaming connect topic<commit_after>package state

import (
	"fmt"

	"github.com/hashicorp/go-memdb"

	"github.com/hashicorp/consul/agent/consul/stream"
	"github.com/hashicorp/consul/proto/pbsubscribe"
)

// ReadTxn is implemented by memdb.Txn to perform read operations.
type ReadTxn interface {
	Get(table, index string, args ...interface{}) (memdb.ResultIterator, error)
	First(table, index string, args ...interface{}) (interface{}, error)
	FirstWatch(table, index string, args ...interface{}) (<-chan struct{}, interface{}, error)
	Abort()
}

// WriteTxn is implemented by memdb.Txn to perform write operations.
type WriteTxn interface {
	ReadTxn
	Insert(table string, obj interface{}) error
	Commit() error
}

// Changes wraps a memdb.Changes to include the index at which these changes
// were made.
type Changes struct {
	// Index is the latest index at the time these changes were committed.
	Index   uint64
	Changes memdb.Changes
}

// changeTrackerDB is a thin wrapper around memdb.DB which enables TrackChanges on
// all write transactions. When the transaction is committed the changes are:
// 1. Used to update our internal usage tracking
// 2. Sent to the eventPublisher which will create and emit change events
type changeTrackerDB struct {
	db             *memdb.MemDB
	publisher      *stream.EventPublisher
	processChanges func(ReadTxn, Changes) ([]stream.Event, error)
}

// Txn exists to maintain backwards compatibility with memdb.DB.Txn. Preexisting
// code may use it to create a read-only transaction, but it will panic if called
// with write=true.
//
// Deprecated: use either ReadTxn, or WriteTxn.
func (c *changeTrackerDB) Txn(write bool) *txn {
	if write {
		panic("don't use db.Txn(true), use db.WriteTxn(idx uin64)")
	}
	return c.ReadTxn()
}

// ReadTxn returns a read-only transaction which behaves exactly the same as
// memdb.Txn
//
// TODO: this could return a regular memdb.Txn if all the state functions accepted
// the ReadTxn interface
func (c *changeTrackerDB) ReadTxn() *txn {
	return &txn{Txn: c.db.Txn(false)}
}

// WriteTxn returns a wrapped memdb.Txn suitable for writes to the state store.
// It will track changes and publish events for the changes when Commit
// is called.
//
// The idx argument must be the index of the current Raft operation. Almost
// all mutations to state should happen as part of a raft apply so the index of
// the log being applied can be passed to WriteTxn.
// The exceptional cases are transactions that are executed on an empty
// memdb.DB as part of Restore, and those executed by tests where we insert
// data directly into the DB. These cases may use WriteTxnRestore.
func (c *changeTrackerDB) WriteTxn(idx uint64) *txn {
	t := &txn{
		Txn:     c.db.Txn(true),
		Index:   idx,
		publish: c.publish,
	}
	t.Txn.TrackChanges()
	return t
}

func (c *changeTrackerDB) publish(tx ReadTxn, changes Changes) error {
	events, err := c.processChanges(tx, changes)
	if err != nil {
		return fmt.Errorf("failed generating events from changes: %v", err)
	}
	c.publisher.Publish(events)
	return nil
}

// WriteTxnRestore returns a wrapped RW transaction that should only be used in
// Restore where we need to replace the entire contents of the Store.
// WriteTxnRestore uses a zero index since the whole restore doesn't really
// occur at one index - the effect is to write many values that were previously
// written across many indexes. WriteTxnRestore also does not publish any
// change events to subscribers.
func (c *changeTrackerDB) WriteTxnRestore() *txn {
	t := &txn{
		Txn:   c.db.Txn(true),
		Index: 0,
	}

	// We enable change tracking so that usage data is correctly populated.
	t.Txn.TrackChanges()
	return t
}

// txn wraps a memdb.Txn to capture changes and send them to the EventPublisher.
//
// This can not be done with txn.Defer because the callback passed to Defer is
// invoked after commit completes, and because the callback can not return an
// error. Any errors from the callback would be lost,  which would result in a
// missing change event, even though the state store had changed.
type txn struct {
	*memdb.Txn
	// Index in raft where the write is occurring. The value is zero for a
	// read-only, or WriteTxnRestore transaction.
	// Index is stored so that it may be passed along to any subscribers as part
	// of a change event.
	Index   uint64
	publish func(tx ReadTxn, changes Changes) error
}

// Commit first pushes changes to EventPublisher, then calls Commit on the
// underlying transaction.
//
// Note that this function, unlike memdb.Txn, returns an error which must be checked
// by the caller. A non-nil error indicates that a commit failed and was not
// applied.
func (tx *txn) Commit() error {
	changes := Changes{
		Index:   tx.Index,
		Changes: tx.Txn.Changes(),
	}

	if len(changes.Changes) > 0 {
		if err := updateUsage(tx, changes); err != nil {
			return err
		}
	}

	// publish may be nil if this is a read-only or WriteTxnRestore transaction.
	// In those cases changes should also be empty, and there will be nothing
	// to publish.
	if tx.publish != nil {
		if err := tx.publish(tx.Txn, changes); err != nil {
			return err
		}
	}

	tx.Txn.Commit()
	return nil
}

var (
	topicServiceHealth        = pbsubscribe.Topic_ServiceHealth
	topicServiceHealthConnect = pbsubscribe.Topic_ServiceHealthConnect
)

func processDBChanges(tx ReadTxn, changes Changes) ([]stream.Event, error) {
	var events []stream.Event
	fns := []func(tx ReadTxn, changes Changes) ([]stream.Event, error){
		aclChangeUnsubscribeEvent,
		ServiceHealthEventsFromChanges,
		// TODO: add other table handlers here.
	}
	for _, fn := range fns {
		e, err := fn(tx, changes)
		if err != nil {
			return nil, err
		}
		events = append(events, e...)
	}
	return events, nil
}

func newSnapshotHandlers(s *Store) stream.SnapshotHandlers {
	return stream.SnapshotHandlers{
		topicServiceHealth: serviceHealthSnapshot(s, topicServiceHealth),
		// The connect topic is temporarily disabled until the correct events are
		// created for terminating gateway changes.
		//topicServiceHealthConnect: serviceHealthSnapshot(s, topicServiceHealthConnect),
	}
}
<|endoftext|>


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

package minio

import (
	"net/http"
	"time"
)

// BucketInfo container for bucket metadata.
type BucketInfo struct {
	// The name of the bucket.
	Name string `json:"name"`
	// Date the bucket was created.
	CreationDate time.Time `json:"creationDate"`
}

// ObjectInfo container for object metadata.
type ObjectInfo struct {
	// An ETag is optionally set to md5sum of an object.  In case of multipart objects,
	// ETag is of the form MD5SUM-N where MD5SUM is md5sum of all individual md5sums of
	// each parts concatenated into one string.
	ETag string `json:"etag"`

	Key          string    `json:"name"`         // Name of the object
	LastModified time.Time `json:"lastModified"` // Date and time the object was last modified.
	Size         int64     `json:"size"`         // Size in bytes of the object.
	ContentType  string    `json:"contentType"`  // A standard MIME type describing the format of the object data.

	// Collection of additional metadata on the object.
	// eg: x-amz-meta-*, content-encoding etc.
	Metadata http.Header `json:"metadata"`

	// Owner name.
	Owner struct {
		DisplayName string `json:"name"`
		ID          string `json:"id"`
	} `json:"owner"`

	// The class of storage used to store the object.
	StorageClass string `json:"storageClass"`

	// Error
	Err error `json:"-"`
}

// ObjectMultipartInfo container for multipart object metadata.
type ObjectMultipartInfo struct {
	// Date and time at which the multipart upload was initiated.
	Initiated time.Time `type:"timestamp" timestampFormat:"iso8601"`

	Initiator initiator
	Owner     owner

	// The type of storage to use for the object. Defaults to 'STANDARD'.
	StorageClass string

	// Key of the object for which the multipart upload was initiated.
	Key string

	// Size in bytes of the object.
	Size int64

	// Upload ID that identifies the multipart upload.
	UploadID string `xml:"UploadId"`

	// Error
	Err error
}
<commit_msg>Ignore Metadata field in xml encode/decode (#890)<commit_after>/*
 * Minio Go Library for Amazon S3 Compatible Cloud Storage
 * Copyright 2015-2017 Minio, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package minio

import (
	"net/http"
	"time"
)

// BucketInfo container for bucket metadata.
type BucketInfo struct {
	// The name of the bucket.
	Name string `json:"name"`
	// Date the bucket was created.
	CreationDate time.Time `json:"creationDate"`
}

// ObjectInfo container for object metadata.
type ObjectInfo struct {
	// An ETag is optionally set to md5sum of an object.  In case of multipart objects,
	// ETag is of the form MD5SUM-N where MD5SUM is md5sum of all individual md5sums of
	// each parts concatenated into one string.
	ETag string `json:"etag"`

	Key          string    `json:"name"`         // Name of the object
	LastModified time.Time `json:"lastModified"` // Date and time the object was last modified.
	Size         int64     `json:"size"`         // Size in bytes of the object.
	ContentType  string    `json:"contentType"`  // A standard MIME type describing the format of the object data.

	// Collection of additional metadata on the object.
	// eg: x-amz-meta-*, content-encoding etc.
	Metadata http.Header `json:"metadata" xml:"-"`

	// Owner name.
	Owner struct {
		DisplayName string `json:"name"`
		ID          string `json:"id"`
	} `json:"owner"`

	// The class of storage used to store the object.
	StorageClass string `json:"storageClass"`

	// Error
	Err error `json:"-"`
}

// ObjectMultipartInfo container for multipart object metadata.
type ObjectMultipartInfo struct {
	// Date and time at which the multipart upload was initiated.
	Initiated time.Time `type:"timestamp" timestampFormat:"iso8601"`

	Initiator initiator
	Owner     owner

	// The type of storage to use for the object. Defaults to 'STANDARD'.
	StorageClass string

	// Key of the object for which the multipart upload was initiated.
	Key string

	// Size in bytes of the object.
	Size int64

	// Upload ID that identifies the multipart upload.
	UploadID string `xml:"UploadId"`

	// Error
	Err error
}
<|endoftext|>


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

package api

import (
	"fmt"
)

var Contracts = &ContractsList{}

type Contract struct {
	ID         int    `json:"id"`
	Label      string `json:"label"`
	CustomerID int    `json:"customer"`
	Customer   *Company
}

type ContractsList struct {
	Contracts []Contract `json:"results"`
}

func (cs *ContractsList) apiURL() string {
	return "v1/my_contracts"
}

func (cs *ContractsList) slug() string {
	return "contracts"
}

// GetByID returns the contract with the given id
func (cs *ContractsList) GetByID(id int) *Contract {
	for _, c := range cs.Contracts {
		if c.ID == id {
			return &c
		}
	}
	return nil
}

// GetByID returns the contract with the given id
func (cs *ContractsList) GetByLabel(label string) *Contract {
	for _, c := range cs.Contracts {
		if c.PrettyLabel() == label {
			return &c
		}
	}
	return nil
}

func (cs *ContractsList) augment() error {
	for i := range cs.Contracts {
		co := &cs.Contracts[i]
		for _, customer := range companies.Companies {
			if customer.ID == co.CustomerID {
				co.Customer = &customer
				break
			}
		}
	}
	return nil
}

func (cs *ContractsList) PrettyPrint() {
	for _, c := range cs.Contracts {
		c.PrettyPrint()
	}
}

func (c *Contract) PrettyLabel() string {
	if c.Customer == nil {
		return fmt.Sprintf("%s [%d]", c.Label, c.CustomerID)
	}
	return fmt.Sprintf("%s [%s]", c.Label, c.Customer.Name)
}

func (c *Contract) PrettyPrint() {
	fmt.Println(c.PrettyLabel())
}

func init() {
	cache.register(Contracts)
}
<commit_msg>Document contracts methods<commit_after>// Copyright © 2017 Julien Pivotto <roidelapluie@inuits.eu>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package api

import (
	"fmt"
)

var Contracts = &ContractsList{}

type Contract struct {
	ID         int    `json:"id"`
	Label      string `json:"label"`
	CustomerID int    `json:"customer"`
	Customer   *Company
}

type ContractsList struct {
	Contracts []Contract `json:"results"`
}

func (cs *ContractsList) apiURL() string {
	return "v1/my_contracts"
}

func (cs *ContractsList) slug() string {
	return "contracts"
}

// GetByID returns the contract with the given id
func (cs *ContractsList) GetByID(id int) *Contract {
	for _, c := range cs.Contracts {
		if c.ID == id {
			return &c
		}
	}
	return nil
}

// GetByID returns the contract with the given id
func (cs *ContractsList) GetByLabel(label string) *Contract {
	for _, c := range cs.Contracts {
		if c.PrettyLabel() == label {
			return &c
		}
	}
	return nil
}

func (cs *ContractsList) augment() error {
	for i := range cs.Contracts {
		co := &cs.Contracts[i]
		for _, customer := range companies.Companies {
			if customer.ID == co.CustomerID {
				co.Customer = &customer
				break
			}
		}
	}
	return nil
}

// PrettyPrint prints contracts in a nice way to the console
func (cs *ContractsList) PrettyPrint() {
	for _, c := range cs.Contracts {
		c.PrettyPrint()
	}
}

// PrettyLabel returns the label of a contract that can be used
// in CLI etc.. to identify a contract. It contains the customer.
func (c *Contract) PrettyLabel() string {
	if c.Customer == nil {
		return fmt.Sprintf("%s [%d]", c.Label, c.CustomerID)
	}
	return fmt.Sprintf("%s [%s]", c.Label, c.Customer.Name)
}

// PrettyPrint prints contract in a nice way to the console
func (c *Contract) PrettyPrint() {
	fmt.Println(c.PrettyLabel())
}

func init() {
	cache.register(Contracts)
}
<|endoftext|>