<commit_before>// Copyright The OpenTelemetry Authors
//
// 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 sapmexporter exports trace data using Splunk's SAPM protocol.
package sapmexporter

import (
	"context"

	sapmclient "github.com/signalfx/sapm-proto/client"
	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/consumer/consumererror"
	"go.opentelemetry.io/collector/consumer/pdata"
	"go.opentelemetry.io/collector/exporter/exporterhelper"
	"go.opentelemetry.io/collector/translator/trace/jaeger"
	"go.uber.org/zap"

	"github.com/open-telemetry/opentelemetry-collector-contrib/internal/common/splunk"
)

// sapmExporter is a wrapper struct of SAPM exporter
type sapmExporter struct {
	client *sapmclient.Client
	logger *zap.Logger
	config *Config
}

func (se *sapmExporter) Shutdown(context.Context) error {
	se.client.Stop()
	return nil
}

func newSAPMExporter(cfg *Config, params component.ExporterCreateParams) (sapmExporter, error) {
	err := cfg.validate()
	if err != nil {
		return sapmExporter{}, err
	}

	client, err := sapmclient.New(cfg.clientOptions()...)
	if err != nil {
		return sapmExporter{}, err
	}

	return sapmExporter{
		client: client,
		logger: params.Logger,
		config: cfg,
	}, err
}

func newSAPMTraceExporter(cfg *Config, params component.ExporterCreateParams) (component.TraceExporter, error) {
	se, err := newSAPMExporter(cfg, params)
	if err != nil {
		return nil, err
	}

	return exporterhelper.NewTraceExporter(
		cfg,
		se.pushTraceData,
		exporterhelper.WithShutdown(se.Shutdown))
}

// tracesByAccessToken takes a pdata.Traces struct and will iterate through its ResourceSpans' attributes,
// regrouping by any SFx access token label value if Config.AccessTokenPassthrough is enabled.  It will delete any
// set token label in any case to prevent serialization.
// It returns a map of newly constructed pdata.Traces keyed by access token, defaulting to empty string.
func (se *sapmExporter) tracesByAccessToken(td pdata.Traces) map[string]pdata.Traces {
	tracesByToken := make(map[string]pdata.Traces, 1)
	resourceSpans := td.ResourceSpans()
	for i := 0; i < resourceSpans.Len(); i++ {
		resourceSpan := resourceSpans.At(i)
		if resourceSpan.IsNil() {
			// Invalid trace so nothing to export
			continue
		}

		accessToken := ""
		if !resourceSpan.Resource().IsNil() {
			attrs := resourceSpan.Resource().Attributes()
			attributeValue, ok := attrs.Get(splunk.SFxAccessTokenLabel)
			if ok {
				attrs.Delete(splunk.SFxAccessTokenLabel)
				if se.config.AccessTokenPassthrough {
					accessToken = attributeValue.StringVal()
				}
			}
		}

		traceForToken, ok := tracesByToken[accessToken]
		if !ok {
			traceForToken = pdata.NewTraces()
			tracesByToken[accessToken] = traceForToken
		}

		// Append ResourceSpan to trace for this access token
		traceForTokenSize := traceForToken.ResourceSpans().Len()
		traceForToken.ResourceSpans().Resize(traceForTokenSize + 1)
		traceForToken.ResourceSpans().At(traceForTokenSize).InitEmpty()
		resourceSpan.CopyTo(traceForToken.ResourceSpans().At(traceForTokenSize))
	}

	return tracesByToken
}

// pushTraceData exports traces in SAPM proto by associated SFx access token and returns number of dropped spans
// and the last experienced error if any translation or export failed
func (se *sapmExporter) pushTraceData(ctx context.Context, td pdata.Traces) (droppedSpansCount int, err error) {
	traces := se.tracesByAccessToken(td)
	droppedSpansCount = 0
	for accessToken, trace := range traces {
		batches, translateErr := jaeger.InternalTracesToJaegerProto(trace)
		if translateErr != nil {
			droppedSpansCount += trace.SpanCount()
			err = consumererror.Permanent(translateErr)
			continue
		}

		exportErr := se.client.ExportWithAccessToken(ctx, batches, accessToken)
		if exportErr != nil {
			if sendErr, ok := exportErr.(*sapmclient.ErrSend); ok {
				if sendErr.Permanent {
					err = consumererror.Permanent(sendErr)
				}
			}
			droppedSpansCount += trace.SpanCount()
		}
	}
	return
}
<commit_msg>adding queue settings (#1390)<commit_after>// Copyright The OpenTelemetry Authors
//
// 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 sapmexporter exports trace data using Splunk's SAPM protocol.
package sapmexporter

import (
	"context"

	sapmclient "github.com/signalfx/sapm-proto/client"
	"go.opentelemetry.io/collector/component"
	"go.opentelemetry.io/collector/consumer/consumererror"
	"go.opentelemetry.io/collector/consumer/pdata"
	"go.opentelemetry.io/collector/exporter/exporterhelper"
	"go.opentelemetry.io/collector/translator/trace/jaeger"
	"go.uber.org/zap"

	"github.com/open-telemetry/opentelemetry-collector-contrib/internal/common/splunk"
)

// sapmExporter is a wrapper struct of SAPM exporter
type sapmExporter struct {
	client *sapmclient.Client
	logger *zap.Logger
	config *Config
}

func (se *sapmExporter) Shutdown(context.Context) error {
	se.client.Stop()
	return nil
}

func newSAPMExporter(cfg *Config, params component.ExporterCreateParams) (sapmExporter, error) {
	err := cfg.validate()
	if err != nil {
		return sapmExporter{}, err
	}

	client, err := sapmclient.New(cfg.clientOptions()...)
	if err != nil {
		return sapmExporter{}, err
	}

	return sapmExporter{
		client: client,
		logger: params.Logger,
		config: cfg,
	}, err
}

func newSAPMTraceExporter(cfg *Config, params component.ExporterCreateParams) (component.TraceExporter, error) {
	se, err := newSAPMExporter(cfg, params)
	if err != nil {
		return nil, err
	}

	return exporterhelper.NewTraceExporter(
		cfg,
		se.pushTraceData,
		exporterhelper.WithShutdown(se.Shutdown),
		exporterhelper.WithQueue(cfg.QueueSettings),
		exporterhelper.WithRetry(cfg.RetrySettings),
		exporterhelper.WithTimeout(cfg.TimeoutSettings),
	)
}

// tracesByAccessToken takes a pdata.Traces struct and will iterate through its ResourceSpans' attributes,
// regrouping by any SFx access token label value if Config.AccessTokenPassthrough is enabled.  It will delete any
// set token label in any case to prevent serialization.
// It returns a map of newly constructed pdata.Traces keyed by access token, defaulting to empty string.
func (se *sapmExporter) tracesByAccessToken(td pdata.Traces) map[string]pdata.Traces {
	tracesByToken := make(map[string]pdata.Traces, 1)
	resourceSpans := td.ResourceSpans()
	for i := 0; i < resourceSpans.Len(); i++ {
		resourceSpan := resourceSpans.At(i)
		if resourceSpan.IsNil() {
			// Invalid trace so nothing to export
			continue
		}

		accessToken := ""
		if !resourceSpan.Resource().IsNil() {
			attrs := resourceSpan.Resource().Attributes()
			attributeValue, ok := attrs.Get(splunk.SFxAccessTokenLabel)
			if ok {
				attrs.Delete(splunk.SFxAccessTokenLabel)
				if se.config.AccessTokenPassthrough {
					accessToken = attributeValue.StringVal()
				}
			}
		}

		traceForToken, ok := tracesByToken[accessToken]
		if !ok {
			traceForToken = pdata.NewTraces()
			tracesByToken[accessToken] = traceForToken
		}

		// Append ResourceSpan to trace for this access token
		traceForTokenSize := traceForToken.ResourceSpans().Len()
		traceForToken.ResourceSpans().Resize(traceForTokenSize + 1)
		traceForToken.ResourceSpans().At(traceForTokenSize).InitEmpty()
		resourceSpan.CopyTo(traceForToken.ResourceSpans().At(traceForTokenSize))
	}

	return tracesByToken
}

// pushTraceData exports traces in SAPM proto by associated SFx access token and returns number of dropped spans
// and the last experienced error if any translation or export failed
func (se *sapmExporter) pushTraceData(ctx context.Context, td pdata.Traces) (droppedSpansCount int, err error) {
	traces := se.tracesByAccessToken(td)
	droppedSpansCount = 0
	for accessToken, trace := range traces {
		batches, translateErr := jaeger.InternalTracesToJaegerProto(trace)
		if translateErr != nil {
			droppedSpansCount += trace.SpanCount()
			err = consumererror.Permanent(translateErr)
			continue
		}

		exportErr := se.client.ExportWithAccessToken(ctx, batches, accessToken)
		if exportErr != nil {
			if sendErr, ok := exportErr.(*sapmclient.ErrSend); ok {
				if sendErr.Permanent {
					err = consumererror.Permanent(sendErr)
				}
			}
			droppedSpansCount += trace.SpanCount()
		}
	}
	return
}
<|endoftext|>


<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
}

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
}

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
		}

		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)
			solvesByUser[row.Str(0)] = userSolves
		}

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

	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)

	var skippedUsers int

	for _, collection := range solvesByUser {

		if !collection.valid() {
			skippedUsers++
			continue
		}

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

	}

	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>Fixed a bug where we didn't count how many rows we actually processed since we skipped dupe rows.<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
}

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
}

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)
			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)

	var skippedUsers int

	for _, collection := range solvesByUser {

		if !collection.valid() {
			skippedUsers++
			continue
		}

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

	}

	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 memrecorder

import (
	"sync"
	"time"

	"github.com/docwhat/docker-gc/types"
)

// MemRecorder stores all seen image tags in memory.
type MemRecorder struct {
	imageTags map[string]time.Time
	mutex     sync.Mutex
}

// NewMemRecorder initializes a new MemRecorder for use.
func NewMemRecorder() *MemRecorder {
	var r MemRecorder
	r.imageTags = make(map[string]time.Time)
	return &r
}

// SawImageTagAt records when a tag was last seen.
//
// Duplicate or older times will be ignored.
func (r *MemRecorder) SawImageTagAt(tag string, when time.Time) {
	if old, ok := r.imageTags[tag]; ok {
		if when.Before(old) {
			return // We don't need to adjust the value.
		}
	}
	r.mutex.Lock()
	r.imageTags[tag] = when
	r.mutex.Unlock()
}

// SawImageTag records a tag beeing seen now.
func (r *MemRecorder) SawImageTag(tag string) {
	r.SawImageTagAt(tag, time.Now())
}

// Forget that you saw a tag.
func (r *MemRecorder) Forget(tag string) {
	r.mutex.Lock()
	delete(r.imageTags, tag)
	r.mutex.Unlock()
}

// Sweep runs a function on all tag and timestamp pairs.
func (r *MemRecorder) Sweep(sweeper types.SweepHandler) {
	for tag, when := range r.imageTags {
		if sweeper(tag, when) {
			r.Forget(tag)
		}
	}
}
<commit_msg>spelling error<commit_after>package memrecorder

import (
	"sync"
	"time"

	"github.com/docwhat/docker-gc/types"
)

// MemRecorder stores all seen image tags in memory.
type MemRecorder struct {
	imageTags map[string]time.Time
	mutex     sync.Mutex
}

// NewMemRecorder initializes a new MemRecorder for use.
func NewMemRecorder() *MemRecorder {
	var r MemRecorder
	r.imageTags = make(map[string]time.Time)
	return &r
}

// SawImageTagAt records when a tag was last seen.
//
// Duplicate or older times will be ignored.
func (r *MemRecorder) SawImageTagAt(tag string, when time.Time) {
	if old, ok := r.imageTags[tag]; ok {
		if when.Before(old) {
			return // We don't need to adjust the value.
		}
	}
	r.mutex.Lock()
	r.imageTags[tag] = when
	r.mutex.Unlock()
}

// SawImageTag records a tag being seen now.
func (r *MemRecorder) SawImageTag(tag string) {
	r.SawImageTagAt(tag, time.Now())
}

// Forget that you saw a tag.
func (r *MemRecorder) Forget(tag string) {
	r.mutex.Lock()
	delete(r.imageTags, tag)
	r.mutex.Unlock()
}

// Sweep runs a function on all tag and timestamp pairs.
func (r *MemRecorder) Sweep(sweeper types.SweepHandler) {
	for tag, when := range r.imageTags {
		if sweeper(tag, when) {
			r.Forget(tag)
		}
	}
}
<|endoftext|>


<commit_before>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// 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 api Golang bindings for Baruwa REST API
package api

// RelaySetting holds relay settings
type RelaySetting struct {
	ID              int     `json:"id,omitempty"`
	Address         string  `json:"address"`
	Username        string  `json:"username"`
	Enabled         bool    `json:"enabled"`
	RequireTLS      bool    `json:"require_tls"`
	Password1       string  `json:"password1,omitempty"`
	Password2       string  `json:"password2,omitempty"`
	Description     string  `json:"description"`
	LowScore        float64 `json:"low_score"`
	HighScore       float64 `json:"high_score"`
	SpamActions     int     `json:"spam_actions"`
	HighSpamActions int     `json:"highspam_actions"`
	BlockMacros     bool    `json:"block_macros"`
	RateLimit       int     `json:"ratelimit"`
}

// GetRelaySetting returns radius settings
// https://www.baruwa.com/docs/api/#retrieve-relay-settings
func (c *Client) GetRelaySetting(id int) (server *RelaySetting, err error) {
	return
}

// CreateRelaySetting creates radius settings
// https://www.baruwa.com/docs/api/#create-relay-settings
func (c *Client) CreateRelaySetting(server *RelaySetting) (err error) {
	return
}

// UpdateRelaySetting updates radius settings
// https://www.baruwa.com/docs/api/#update-relay-settings
func (c *Client) UpdateRelaySetting(server *RelaySetting) (err error) {
	return
}

// DeleteRelaySetting deletes radius settings
// https://www.baruwa.com/docs/api/#delete-relay-settings
func (c *Client) DeleteRelaySetting(id int) (err error) {
	return
}
<commit_msg>FET: Relay Settings implementation<commit_after>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// 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 api Golang bindings for Baruwa REST API
package api

import (
	"fmt"
	"net/url"

	"github.com/google/go-querystring/query"
)

// RelaySetting holds relay settings
type RelaySetting struct {
	ID              int     `json:"id,omitempty"`
	Address         string  `json:"address"`
	Username        string  `json:"username"`
	Enabled         bool    `json:"enabled"`
	RequireTLS      bool    `json:"require_tls"`
	Password1       string  `json:"password1,omitempty"`
	Password2       string  `json:"password2,omitempty"`
	Description     string  `json:"description"`
	LowScore        float64 `json:"low_score"`
	HighScore       float64 `json:"high_score"`
	SpamActions     int     `json:"spam_actions"`
	HighSpamActions int     `json:"highspam_actions"`
	BlockMacros     bool    `json:"block_macros"`
	RateLimit       int     `json:"ratelimit"`
}

// GetRelaySetting returns radius settings
// https://www.baruwa.com/docs/api/#retrieve-relay-settings
func (c *Client) GetRelaySetting(relayID int) (server *RelaySetting, err error) {
	if relayID <= 0 {
		err = fmt.Errorf("The relayID param should be > 0")
		return
	}

	err = c.get(fmt.Sprintf("relays/%d", relayID), server)

	return
}

// CreateRelaySetting creates radius settings
// https://www.baruwa.com/docs/api/#create-relay-settings
func (c *Client) CreateRelaySetting(organizationID int, server *RelaySetting) (err error) {
	var v url.Values

	if organizationID <= 0 {
		err = fmt.Errorf("The organizationID param should be > 0")
		return
	}

	if server == nil {
		err = fmt.Errorf("The server param cannot be nil")
		return
	}

	if v, err = query.Values(server); err != nil {
		return
	}

	err = c.post(fmt.Sprintf("relays/%d", organizationID), v, server)

	return
}

// UpdateRelaySetting updates radius settings
// https://www.baruwa.com/docs/api/#update-relay-settings
func (c *Client) UpdateRelaySetting(server *RelaySetting) (err error) {
	var v url.Values

	if server == nil {
		err = fmt.Errorf("The server param cannot be nil")
		return
	}

	if server.ID <= 0 {
		err = fmt.Errorf("The server.ID param should be > 0")
		return
	}

	if v, err = query.Values(server); err != nil {
		return
	}

	err = c.put(fmt.Sprintf("relays/%d", server.ID), v, server)

	return
}

// DeleteRelaySetting deletes radius settings
// https://www.baruwa.com/docs/api/#delete-relay-settings
func (c *Client) DeleteRelaySetting(server *RelaySetting) (err error) {
	var v url.Values

	if server == nil {
		err = fmt.Errorf("The server param cannot be nil")
		return
	}

	if server.ID <= 0 {
		err = fmt.Errorf("The server.ID param should be > 0")
		return
	}

	if v, err = query.Values(server); err != nil {
		return
	}

	err = c.delete(fmt.Sprintf("relays/%d", server.ID), v)

	return
}
<|endoftext|>


<commit_before>package trollmode

import (
	"errors"
	"fmt"
	"socialapi/models"
	"socialapi/workers/common/manager"

	"github.com/koding/bongo"
	"github.com/koding/logging"
	"github.com/koding/worker"
	"github.com/streadway/amqp"
)

const (
	MarkedAsTroll   = "api.account_marked_as_troll"
	UnMarkedAsTroll = "api.account_unmarked_as_troll"
)

func NewManager(controller worker.ErrHandler) *manager.Manager {
	m := manager.New()
	m.Controller(controller)
	m.HandleFunc(MarkedAsTroll, (*Controller).MarkedAsTroll)
	m.HandleFunc(UnMarkedAsTroll, (*Controller).UnMarkedAsTroll)
	return m
}

type Controller struct {
	log logging.Logger
}

func NewController(log logging.Logger) *Controller {
	return &Controller{
		log: log,
	}
}

// this worker is completely idempotent, so no need to cut the circuit
func (c *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	c.log.Error("an error occured putting message back to queue", err)
	delivery.Nack(false, true)
	return false
}

func (c *Controller) MarkedAsTroll(account *models.Account) error {
	if err := c.validateRequest(account); err != nil {
		c.log.Error("Validation failed for marking troll; skipping, err: %s ", err.Error())
		return nil
	}

	if err := c.markChannels(account, models.Safe); err != nil {
		c.log.Error("Error while processing channels, err: %s ", err.Error())
		return err
	}

	if err := c.markParticipations(account, models.Safe); err != nil {
		c.log.Error("Error while processing participations, err: %s ", err.Error())
		return err
	}

	if err := c.markMessages(account, models.Safe); err != nil {
		c.log.Error("Error while processing channels messages, err: %s ", err.Error())
		return err
	}

	if err := c.markInteractions(account, models.Safe); err != nil {
		c.log.Error("Error while processing interactions, err: %s ", err.Error())
		return err
	}

	return nil
}

func (c *Controller) UnMarkedAsTroll(account *models.Account) error {
	if err := c.validateRequest(account); err != nil {
		c.log.Error("Validation failed for un-marking troll; skipping, err: %s ", err.Error())
		return nil
	}

	if err := c.markChannels(account, models.Troll); err != nil {
		c.log.Error("Error while processing channels, err: %s ", err.Error())
		return err
	}

	if err := c.markParticipations(account, models.Troll); err != nil {
		c.log.Error("Error while processing participations, err: %s ", err.Error())
		return err
	}

	if err := c.markMessages(account, models.Troll); err != nil {
		c.log.Error("Error while processing channels messages, err: %s ", err.Error())
		return err
	}

	if err := c.markInteractions(account, models.Troll); err != nil {
		c.log.Error("Error while processing interactions, err: %s ", err.Error())
		return err
	}

	return nil
}

func (c *Controller) validateRequest(account *models.Account) error {
	if account == nil {
		return errors.New("account is not set (nil)")
	}

	if account.Id == 0 {
		return errors.New("account id is not set")
	}

	return nil
}

func (c *Controller) markChannels(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredChannels []models.Channel

	ch := models.NewChannel()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"creator_id":    account.Id,
			"type_constant": models.Channel_TYPE_PRIVATE_MESSAGE,
			"meta_bits":     currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {
		// set skip everytime here
		q.Pagination.Skip = skip
		var channels []models.Channel
		if err := ch.Some(&channels, q); err != nil {
			return err
		}

		// we processed all messages
		if len(channels) <= 0 {
			break
		}

		for i, channel := range channels {

			if currentStatus == models.Safe {
				channel.MetaBits.Mark(models.Troll)
			} else {
				channel.MetaBits.UnMark(models.Troll)
			}

			if err := channel.Update(); err != nil {
				erroredChannels = append(erroredChannels, channels[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredChannels) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredChannels))
	}

	return nil
}

func (c *Controller) markParticipations(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredChannelParticipants []models.ChannelParticipant

	cp := models.NewChannelParticipant()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": account.Id,
			"meta_bits":  currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var channelParticipants []models.ChannelParticipant
		if err := cp.Some(&channelParticipants, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(channelParticipants) <= 0 {
			break
		}

		for i, channelParticipant := range channelParticipants {

			if currentStatus == models.Safe {
				channelParticipant.MetaBits.Mark(models.Troll)
			} else {
				channelParticipant.MetaBits.UnMark(models.Troll)
			}

			if err := channelParticipant.Update(); err != nil {
				erroredChannelParticipants = append(erroredChannelParticipants, channelParticipants[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredChannelParticipants) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredChannelParticipants))
	}

	return nil
}

func (c *Controller) markMessages(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredMessages []models.ChannelMessage

	cm := models.NewChannelMessage()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": account.Id,
			// 0 means safe
			"meta_bits": currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var messages []models.ChannelMessage
		if err := cm.Some(&messages, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(messages) <= 0 {
			break
		}

		for i, message := range messages {
			// mark all message_list items as exempt
			if err := c.markMessageLists(&message, currentStatus); err != nil {
				return err
			}

			// mark all message_replies items as exempt
			if err := c.markMessageReplies(&message, currentStatus); err != nil {
				return err
			}

			if currentStatus == models.Safe {
				message.MetaBits.Mark(models.Troll)
			} else {
				message.MetaBits.UnMark(models.Troll)
			}

			// ChannelMessage update only updates body of the message
			if err := bongo.B.Update(message); err != nil {
				erroredMessages = append(erroredMessages, messages[i])
			}
		}

		// increment skip count
		skip += processCount
	}

	if len(erroredMessages) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredMessages))
	}

	return nil
}

func (c *Controller) markMessageLists(message *models.ChannelMessage, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredMessages []models.ChannelMessageList

	cml := models.NewChannelMessageList()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": message.Id,
			"meta_bits":  currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var messageList []models.ChannelMessageList
		if err := cml.Some(&messageList, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(messageList) <= 0 {
			break
		}

		for i, item := range messageList {

			if currentStatus == models.Safe {
				item.MetaBits.Mark(models.Troll)
			} else {
				item.MetaBits.UnMark(models.Troll)
			}

			if err := item.Update(); err != nil {
				erroredMessages = append(erroredMessages, messageList[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredMessages) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredMessages))
	}

	return nil
}

func (c *Controller) markMessageReplies(message *models.ChannelMessage, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredMessages []models.MessageReply

	mr := models.NewMessageReply()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"reply_id":  message.Id,
			"meta_bits": currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var messageList []models.MessageReply
		if err := mr.Some(&messageList, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(messageList) <= 0 {
			break
		}

		for i, messageReply := range messageList {
			if currentStatus == models.Safe {
				messageReply.MetaBits.Mark(models.Troll)
			} else {
				messageReply.MetaBits.UnMark(models.Troll)
			}

			if err := messageReply.Update(); err != nil {
				erroredMessages = append(erroredMessages, messageList[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredMessages) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredMessages))
	}

	return nil
}

func (c *Controller) markInteractions(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredInteractions []models.Interaction

	i := models.NewInteraction()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": account.Id,
			// 0 means safe
			"meta_bits": currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {
		// set skip everytime here
		q.Pagination.Skip = skip
		var interactions []models.Interaction
		if err := i.Some(&interactions, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(interactions) <= 0 {
			break
		}

		for i, interaction := range interactions {
			if currentStatus == models.Safe {
				interaction.MetaBits.Mark(models.Troll)
			} else {
				interaction.MetaBits.UnMark(models.Troll)
			}

			if err := interaction.Update(); err != nil {
				erroredInteractions = append(erroredInteractions, interactions[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredInteractions) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredInteractions))
	}

	return nil
}
<commit_msg>Social: remove manager<commit_after>package trollmode

import (
	"errors"
	"fmt"
	"socialapi/models"

	"github.com/koding/bongo"
	"github.com/koding/logging"
	"github.com/streadway/amqp"
)

const (
	MarkedAsTroll   = "api.account_marked_as_troll"
	UnMarkedAsTroll = "api.account_unmarked_as_troll"
)

type Controller struct {
	log logging.Logger
}

func NewController(log logging.Logger) *Controller {
	return &Controller{
		log: log,
	}
}

// this worker is completely idempotent, so no need to cut the circuit
func (c *Controller) DefaultErrHandler(delivery amqp.Delivery, err error) bool {
	c.log.Error("an error occured putting message back to queue", err)
	delivery.Nack(false, true)
	return false
}

func (c *Controller) MarkedAsTroll(account *models.Account) error {
	if err := c.validateRequest(account); err != nil {
		c.log.Error("Validation failed for marking troll; skipping, err: %s ", err.Error())
		return nil
	}

	if err := c.markChannels(account, models.Safe); err != nil {
		c.log.Error("Error while processing channels, err: %s ", err.Error())
		return err
	}

	if err := c.markParticipations(account, models.Safe); err != nil {
		c.log.Error("Error while processing participations, err: %s ", err.Error())
		return err
	}

	if err := c.markMessages(account, models.Safe); err != nil {
		c.log.Error("Error while processing channels messages, err: %s ", err.Error())
		return err
	}

	if err := c.markInteractions(account, models.Safe); err != nil {
		c.log.Error("Error while processing interactions, err: %s ", err.Error())
		return err
	}

	return nil
}

func (c *Controller) UnMarkedAsTroll(account *models.Account) error {
	if err := c.validateRequest(account); err != nil {
		c.log.Error("Validation failed for un-marking troll; skipping, err: %s ", err.Error())
		return nil
	}

	if err := c.markChannels(account, models.Troll); err != nil {
		c.log.Error("Error while processing channels, err: %s ", err.Error())
		return err
	}

	if err := c.markParticipations(account, models.Troll); err != nil {
		c.log.Error("Error while processing participations, err: %s ", err.Error())
		return err
	}

	if err := c.markMessages(account, models.Troll); err != nil {
		c.log.Error("Error while processing channels messages, err: %s ", err.Error())
		return err
	}

	if err := c.markInteractions(account, models.Troll); err != nil {
		c.log.Error("Error while processing interactions, err: %s ", err.Error())
		return err
	}

	return nil
}

func (c *Controller) validateRequest(account *models.Account) error {
	if account == nil {
		return errors.New("account is not set (nil)")
	}

	if account.Id == 0 {
		return errors.New("account id is not set")
	}

	return nil
}

func (c *Controller) markChannels(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredChannels []models.Channel

	ch := models.NewChannel()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"creator_id":    account.Id,
			"type_constant": models.Channel_TYPE_PRIVATE_MESSAGE,
			"meta_bits":     currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {
		// set skip everytime here
		q.Pagination.Skip = skip
		var channels []models.Channel
		if err := ch.Some(&channels, q); err != nil {
			return err
		}

		// we processed all messages
		if len(channels) <= 0 {
			break
		}

		for i, channel := range channels {

			if currentStatus == models.Safe {
				channel.MetaBits.Mark(models.Troll)
			} else {
				channel.MetaBits.UnMark(models.Troll)
			}

			if err := channel.Update(); err != nil {
				erroredChannels = append(erroredChannels, channels[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredChannels) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredChannels))
	}

	return nil
}

func (c *Controller) markParticipations(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredChannelParticipants []models.ChannelParticipant

	cp := models.NewChannelParticipant()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": account.Id,
			"meta_bits":  currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var channelParticipants []models.ChannelParticipant
		if err := cp.Some(&channelParticipants, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(channelParticipants) <= 0 {
			break
		}

		for i, channelParticipant := range channelParticipants {

			if currentStatus == models.Safe {
				channelParticipant.MetaBits.Mark(models.Troll)
			} else {
				channelParticipant.MetaBits.UnMark(models.Troll)
			}

			if err := channelParticipant.Update(); err != nil {
				erroredChannelParticipants = append(erroredChannelParticipants, channelParticipants[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredChannelParticipants) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredChannelParticipants))
	}

	return nil
}

func (c *Controller) markMessages(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredMessages []models.ChannelMessage

	cm := models.NewChannelMessage()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": account.Id,
			// 0 means safe
			"meta_bits": currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var messages []models.ChannelMessage
		if err := cm.Some(&messages, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(messages) <= 0 {
			break
		}

		for i, message := range messages {
			// mark all message_list items as exempt
			if err := c.markMessageLists(&message, currentStatus); err != nil {
				return err
			}

			// mark all message_replies items as exempt
			if err := c.markMessageReplies(&message, currentStatus); err != nil {
				return err
			}

			if currentStatus == models.Safe {
				message.MetaBits.Mark(models.Troll)
			} else {
				message.MetaBits.UnMark(models.Troll)
			}

			// ChannelMessage update only updates body of the message
			if err := bongo.B.Update(message); err != nil {
				erroredMessages = append(erroredMessages, messages[i])
			}
		}

		// increment skip count
		skip += processCount
	}

	if len(erroredMessages) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredMessages))
	}

	return nil
}

func (c *Controller) markMessageLists(message *models.ChannelMessage, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredMessages []models.ChannelMessageList

	cml := models.NewChannelMessageList()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"message_id": message.Id,
			"meta_bits":  currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var messageList []models.ChannelMessageList
		if err := cml.Some(&messageList, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(messageList) <= 0 {
			break
		}

		for i, item := range messageList {

			if currentStatus == models.Safe {
				item.MetaBits.Mark(models.Troll)
			} else {
				item.MetaBits.UnMark(models.Troll)
			}

			if err := item.Update(); err != nil {
				erroredMessages = append(erroredMessages, messageList[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredMessages) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredMessages))
	}

	return nil
}

func (c *Controller) markMessageReplies(message *models.ChannelMessage, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredMessages []models.MessageReply

	mr := models.NewMessageReply()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"reply_id":  message.Id,
			"meta_bits": currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {

		// set skip everytime here
		q.Pagination.Skip = skip
		var messageList []models.MessageReply
		if err := mr.Some(&messageList, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(messageList) <= 0 {
			break
		}

		for i, messageReply := range messageList {
			if currentStatus == models.Safe {
				messageReply.MetaBits.Mark(models.Troll)
			} else {
				messageReply.MetaBits.UnMark(models.Troll)
			}

			if err := messageReply.Update(); err != nil {
				erroredMessages = append(erroredMessages, messageList[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredMessages) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredMessages))
	}

	return nil
}

func (c *Controller) markInteractions(account *models.Account, currentStatus models.MetaBits) error {
	var processCount = 100
	var skip = 0
	var erroredInteractions []models.Interaction

	i := models.NewInteraction()
	q := &bongo.Query{
		Selector: map[string]interface{}{
			"account_id": account.Id,
			// 0 means safe
			"meta_bits": currentStatus,
		},
		Pagination: *bongo.NewPagination(processCount, 0),
	}

	for {
		// set skip everytime here
		q.Pagination.Skip = skip
		var interactions []models.Interaction
		if err := i.Some(&interactions, q); err != nil {
			return err
		}

		// we processed all channel participants
		if len(interactions) <= 0 {
			break
		}

		for i, interaction := range interactions {
			if currentStatus == models.Safe {
				interaction.MetaBits.Mark(models.Troll)
			} else {
				interaction.MetaBits.UnMark(models.Troll)
			}

			if err := interaction.Update(); err != nil {
				erroredInteractions = append(erroredInteractions, interactions[i])
			}
		}

		// increment skip count
		skip = processCount + skip
	}

	if len(erroredInteractions) != 0 {
		return errors.New(fmt.Sprintf("some errors: %v", erroredInteractions))
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Jesse van den Kieboom. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package configure provides a very simple gnu configure/make style configure
// script generating a simple Makefile and go file containing all the configured
// variables.
package configure

import (
	"bytes"
	"fmt"
	"github.com/jessevdk/go-flags"
	"io"
	"os"
	"path"
	"reflect"
	"regexp"
	"sort"
	"strings"
	"runtime"
)

// Options contains all the standard configure options to specify various
// directories. Use NewOptions to create an instance of this type with the
// common default values for each variable.
type Options struct {
	Prefix        string `long:"prefix" description:"install architecture-independent files in PREFIX"`
	ExecPrefix    string `long:"execprefix" description:"install architecture-dependent files in EPREFIX"`
	BinDir        string `long:"bindir" description:"user executables"`
	LibExecDir    string `long:"libexecdir" description:"program executables"`
	SysConfDir    string `long:"sysconfdir" description:"read-only single-machine data"`
	LibDir        string `long:"libdir" description:"program executables"`
	DataRootDir   string `long:"datarootdir" description:"read-only arch.-independent data root"`
	DataDir       string `long:"datadir" description:"read-only arc.-independent data"`
	ManDir        string `long:"mandir" description:"man documentation"`
}

// NewOptions creates a new Options with common default values.
func NewOptions() *Options {
	return &Options{
		Prefix:        "/usr/local",
		ExecPrefix:    "${prefix}",
		BinDir:        "${execprefix}/bin",
		LibExecDir:    "${execprefix}/libexec",
		LibDir:        "${execprefix}/lib",
		SysConfDir:    "${prefix}/etc",
		DataRootDir:   "${prefix}/share",
		DataDir:       "${datarootdir}",
		ManDir:        "${datarootdir}/man",
	}
}

// Package is the package name in which the GoConfig file will be written
var Package = "main"

// Makefile is the filename of the makefile that will be generated
var Makefile = "go.make"

// GoConfig is the filename of the go file that will be generated containing
// all the variable values.
var GoConfig = "appconfig"

// GoConfigVariable is the name of the variable inside the GoConfig file
// containing all the variable values.
var GoConfigVariable = "AppConfig"

// Target is the executable name to build. If left empty, the name is deduced
// from the directory (similar to what go does)
var Target = ""

// Version is the application version
var Version []int = []int{0, 1}

type expandStringPart struct {
	Value      string
	IsVariable bool
}

func (x *expandStringPart) expand(m map[string]*expandString) (string, []string) {
	if x.IsVariable {
		s, ok := m[x.Value]

		if !ok {
			return "", nil
		} else {
			ret := s.expand(m)
			rets := make([]string, len(s.dependencies), len(s.dependencies)+1)

			copy(rets, s.dependencies)

			return ret, append(rets, x.Value)
		}
	}

	return x.Value, nil
}

type expandString struct {
	Name  string
	Parts []expandStringPart

	dependencies []string
	value        string
	hasExpanded  bool
}

func (x *expandString) dependsOn(name string) bool {
	i := sort.SearchStrings(x.dependencies, name)

	return i < len(x.dependencies) && x.dependencies[i] == name
}

func (x *expandString) expand(m map[string]*expandString) string {
	if !x.hasExpanded {
		// Prevent infinite loop by circular dependencies
		x.hasExpanded = true
		buf := bytes.Buffer{}

		for _, v := range x.Parts {
			s, deps := v.expand(m)
			buf.WriteString(s)

			x.dependencies = append(x.dependencies, deps...)
		}

		sort.Strings(x.dependencies)
		x.value = buf.String()
	}

	return x.value
}

// Config represents the current configuration. See Configure for more
// information.
type Config struct {
	*flags.Parser

	values   map[string]interface{}
	expanded map[string]*expandString
}

func (x *Config) extract() map[string]interface{} {
	ret := make(map[string]interface{})

	for _, grp := range x.Parser.Groups {
		for longname, option := range grp.LongNames {
			ret[longname] = option.Value.Interface()
		}
	}

	return ret
}

func (x *Config) expand() map[string]*expandString {
	ret := make(map[string]*expandString)

	r, _ := regexp.Compile(`\$\{[^}]*\}`)

	for name, val := range x.values {
		es := expandString{
			Name: name,
		}

		// Find all variable references
		s := fmt.Sprintf("%v", val)

		matches := r.FindAllStringIndex(s, -1)

		for i, match := range matches {
			var prefix string

			if i == 0 {
				prefix = s[0:match[0]]
			} else {
				prefix = s[matches[i-1][1]:match[0]]
			}

			if len(prefix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: prefix, IsVariable: false})
			}

			varname := s[match[0]+2 : match[1]-1]
			es.Parts = append(es.Parts, expandStringPart{Value: varname, IsVariable: true})
		}

		if len(matches) == 0 {
			es.Parts = append(es.Parts, expandStringPart{Value: s, IsVariable: false})
		} else {
			last := matches[len(matches)-1]
			suffix := s[last[1]:]

			if len(suffix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: suffix, IsVariable: false})
			}
		}

		ret[name] = &es
	}

	for _, val := range ret {
		val.expand(ret)
	}

	return ret
}

// Configure runs the configure process with options as provided by the given
// data variable. If data is nil, the default options will be used
// (see NewOptions). Note that the data provided is simply passed to go-flags.
// For more information on flags parsing, see the documentation of go-flags.
// If GoConfig is not empty, then the go configuration will be written to the
// GoConfig file. Similarly, if Makefile is not empty, the Makefile will be
// written.
func Configure(data interface{}) (*Config, error) {
	if data == nil {
		data = NewOptions()
	}

	parser := flags.NewParser(data, flags.PrintErrors | flags.IgnoreUnknown)

	if _, err := parser.Parse(); err != nil {
		return nil, err
	}

	ret := &Config{
		Parser: parser,
	}

	ret.values = ret.extract()
	ret.expanded = ret.expand()

	if len(GoConfig) != 0 {
		filename := GoConfig

		if !strings.HasSuffix(filename, ".go") {
			filename += ".go"
		}

		f, err := os.Create(filename)

		if err != nil {
			return nil, err
		}

		ret.WriteGoConfig(f)
		f.Close()
	}

	if len(Makefile) != 0 {
		f, err := os.Create(Makefile)

		if err != nil {
			return nil, err
		}

		ret.WriteMakefile(f)
		f.Close()

		os.Chmod(Makefile, 0755)

		f, err = os.OpenFile(path.Join(path.Dir(Makefile), "Makefile"),
		                     os.O_CREATE | os.O_EXCL | os.O_WRONLY,
		                     0644)

		if err == nil {
			fmt.Fprintf(f, "include %s\n", path.Base(Makefile))
			f.Close()
		}
	}

	return ret, nil
}

// Expand expands the variable value indicated by name
func (x *Config) Expand(name string) string {
	return x.expanded[name].expand(x.expanded)
}

// WriteGoConfig writes the go configuration file containing all the variable
// values to the given writer. Note that it will write a package line if
// the Package variable is not empty. The GoConfigVariable name will
// be used as the variable name for the configuration.
func (x *Config) WriteGoConfig(writer io.Writer) {
	if len(Package) > 0 {
		fmt.Fprintf(writer, "package %v\n\n", Package)
	}

	fmt.Fprintf(writer, "var %s = struct {\n", GoConfigVariable)
	values := make([]string, 0)

	variables := make([]string, 0, len(x.values))
	optionmap := make(map[string]*flags.Option)

	// Write all options
	for _, grp := range x.Parser.Groups {
		for _, option := range grp.LongNames {
			name := option.Field.Name

			variables = append(variables, name)
			optionmap[name] = option
		}
	}

	sort.Strings(variables)

	for i, name := range variables {
		if i != 0 {
			io.WriteString(writer, "\n")
		}

		option := optionmap[name]
		val := option.Value.Interface()

		fmt.Fprintf(writer, "\t// %s\n", option.Description)
		fmt.Fprintf(writer, "\t%v %T\n", name, val)

		var value string

		if option.Value.Type().Kind() == reflect.String {
			value = fmt.Sprintf("%#v", x.Expand(option.LongName))
		} else {
			value = fmt.Sprintf("%#v", val)
		}

		values = append(values, value)
	}

	if len(variables) > 0 {
		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "\t// Application version\n")
	io.WriteString(writer, "\tVersion []int\n")
	fmt.Fprintln(writer, "}{")

	for _, v := range values {
		fmt.Fprintf(writer, "\t%v,\n", v)
	}

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ", ")
		} else {
			io.WriteString(writer, "\t[]int{")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	fmt.Fprintln(writer, "},")
	fmt.Fprintln(writer, "}")
}

// WriteMakefile writes a Makefile for the given parser to the given writer.
// The Makefile contains the common build, clean, distclean, install and
// uninstall rules.
func (x *Config) WriteMakefile(writer io.Writer) {
	// Write a very basic makefile
	io.WriteString(writer, "#!/usr/bin/make -f\n\n")

	vars := make([]*expandString, 0, len(x.expanded))

	for name, v := range x.expanded {
		inserted := false

		// Insert into vars based on dependencies
		for i, vv := range vars {
			if vv.dependsOn(name) {
				tail := make([]*expandString, len(vars)-i)
				copy(tail, vars[i:])

				if i == 0 {
					vars = append([]*expandString{v}, vars...)
				} else {
					vars = append(append(vars[0:i], v), tail...)
				}

				inserted = true
				break
			}
		}

		if !inserted {
			vars = append(vars, v)
		}
	}

	io.WriteString(writer, "# Variables\n")

	for _, v := range vars {
		fmt.Fprintf(writer, "%s ?= ", v.Name)

		for _, part := range v.Parts {
			if part.IsVariable {
				fmt.Fprintf(writer, "$(%s)", part.Value)
			} else {
				fmt.Fprintf(writer, "%s", part.Value)
			}
		}

		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "version ?= ")

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ".")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	io.WriteString(writer, "\n")
	fmt.Fprintf(writer, "major_version = %v\n", Version[0])

	if len(Version) > 1 {
		fmt.Fprintf(writer, "minor_version = %v\n", Version[1])
	}

	if len(Version) > 2 {
		fmt.Fprintf(writer, "micro_version = %v\n", Version[2])
	}

	io.WriteString(writer, "\n")

	target := Target

	if len(target) == 0 {
		pc := make([]uintptr, 3)
		n := runtime.Callers(1, pc)

		me, _ := runtime.FuncForPC(pc[0]).FileLine(pc[0])

		for i := 1; i < n; i++ {
			f := runtime.FuncForPC(pc[i])
			fname, _ := f.FileLine(pc[i])

			if fname != me {
				target = path.Base(path.Dir(fname))
				break
			}
		}
	}

	fmt.Fprintf(writer, "TARGET ?= %s\n", target)

	io.WriteString(writer, "\nSOURCES ?=")
	io.WriteString(writer, "\nSOURCES += $(wildcard *.go)")
	io.WriteString(writer, "\nSOURCES_UNIQUE = $(sort $(SOURCES))")

	io.WriteString(writer, "\n\n")

	io.WriteString(writer, "# Rules\n")
	io.WriteString(writer, "$(TARGET): $(SOURCES_UNIQUE)\n")
	io.WriteString(writer, "\tgo build -o $@\n\n")

	io.WriteString(writer, "clean:\n")
	io.WriteString(writer, "\trm -f $(TARGET)\n\n")

	io.WriteString(writer, "distclean: clean\n\n")

	io.WriteString(writer, "$(TARGET)_installdir ?= $(bindir)\n\n")

	io.WriteString(writer, "install: $(TARGET)\n")
	io.WriteString(writer, "\tmkdir -p $(DESTDIR)$($(TARGET)_installdir) && cp $(TARGET) $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, "uninstall:\n")
	io.WriteString(writer, "\trm -f $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, ".PHONY: install uninstall distclean clean")
}
<commit_msg>Update to go-flags changes<commit_after>// Copyright 2012 Jesse van den Kieboom. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package configure provides a very simple gnu configure/make style configure
// script generating a simple Makefile and go file containing all the configured
// variables.
package configure

import (
	"bytes"
	"fmt"
	"github.com/jessevdk/go-flags"
	"io"
	"os"
	"path"
	"regexp"
	"sort"
	"strings"
	"runtime"
)

// Options contains all the standard configure options to specify various
// directories. Use NewOptions to create an instance of this type with the
// common default values for each variable.
type Options struct {
	Prefix        string `long:"prefix" description:"install architecture-independent files in PREFIX"`
	ExecPrefix    string `long:"execprefix" description:"install architecture-dependent files in EPREFIX"`
	BinDir        string `long:"bindir" description:"user executables"`
	LibExecDir    string `long:"libexecdir" description:"program executables"`
	SysConfDir    string `long:"sysconfdir" description:"read-only single-machine data"`
	LibDir        string `long:"libdir" description:"program executables"`
	DataRootDir   string `long:"datarootdir" description:"read-only arch.-independent data root"`
	DataDir       string `long:"datadir" description:"read-only arc.-independent data"`
	ManDir        string `long:"mandir" description:"man documentation"`
}

// NewOptions creates a new Options with common default values.
func NewOptions() *Options {
	return &Options{
		Prefix:        "/usr/local",
		ExecPrefix:    "${prefix}",
		BinDir:        "${execprefix}/bin",
		LibExecDir:    "${execprefix}/libexec",
		LibDir:        "${execprefix}/lib",
		SysConfDir:    "${prefix}/etc",
		DataRootDir:   "${prefix}/share",
		DataDir:       "${datarootdir}",
		ManDir:        "${datarootdir}/man",
	}
}

// Package is the package name in which the GoConfig file will be written
var Package = "main"

// Makefile is the filename of the makefile that will be generated
var Makefile = "go.make"

// GoConfig is the filename of the go file that will be generated containing
// all the variable values.
var GoConfig = "appconfig"

// GoConfigVariable is the name of the variable inside the GoConfig file
// containing all the variable values.
var GoConfigVariable = "AppConfig"

// Target is the executable name to build. If left empty, the name is deduced
// from the directory (similar to what go does)
var Target = ""

// Version is the application version
var Version []int = []int{0, 1}

type expandStringPart struct {
	Value      string
	IsVariable bool
}

func (x *expandStringPart) expand(m map[string]*expandString) (string, []string) {
	if x.IsVariable {
		s, ok := m[x.Value]

		if !ok {
			return "", nil
		} else {
			ret := s.expand(m)
			rets := make([]string, len(s.dependencies), len(s.dependencies)+1)

			copy(rets, s.dependencies)

			return ret, append(rets, x.Value)
		}
	}

	return x.Value, nil
}

type expandString struct {
	Name  string
	Parts []expandStringPart

	dependencies []string
	value        string
	hasExpanded  bool
}

func (x *expandString) dependsOn(name string) bool {
	i := sort.SearchStrings(x.dependencies, name)

	return i < len(x.dependencies) && x.dependencies[i] == name
}

func (x *expandString) expand(m map[string]*expandString) string {
	if !x.hasExpanded {
		// Prevent infinite loop by circular dependencies
		x.hasExpanded = true
		buf := bytes.Buffer{}

		for _, v := range x.Parts {
			s, deps := v.expand(m)
			buf.WriteString(s)

			x.dependencies = append(x.dependencies, deps...)
		}

		sort.Strings(x.dependencies)
		x.value = buf.String()
	}

	return x.value
}

// Config represents the current configuration. See Configure for more
// information.
type Config struct {
	*flags.Parser

	values []*flags.Option
	valuesMap   map[string]*flags.Option
	expanded map[string]*expandString
}

func eachGroup(g *flags.Group, f func(g *flags.Group)) {
	f(g)

	for _, gg := range g.Groups() {
		eachGroup(gg, f)
	}
}

func (x *Config) extract() ([]*flags.Option, map[string]*flags.Option) {
	valuesmap := make(map[string]*flags.Option)
	var values []*flags.Option

	eachGroup(x.Parser.Command.Group, func(g *flags.Group) {
		for _, option := range g.Options() {
			if len(option.LongName) > 0 {
				valuesmap[option.LongName] = option
				values = append(values, option)
			}
		}
	})

	return values, valuesmap
}

func (x *Config) expand() map[string]*expandString {
	ret := make(map[string]*expandString)

	r, _ := regexp.Compile(`\$\{[^}]*\}`)

	for name, opt := range x.valuesMap {
		es := expandString{
			Name: name,
		}

		// Find all variable references
		s, ok := opt.Value().(string)

		if !ok {
			continue
		}

		matches := r.FindAllStringIndex(s, -1)

		for i, match := range matches {
			var prefix string

			if i == 0 {
				prefix = s[0:match[0]]
			} else {
				prefix = s[matches[i-1][1]:match[0]]
			}

			if len(prefix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: prefix, IsVariable: false})
			}

			varname := s[match[0]+2 : match[1]-1]
			es.Parts = append(es.Parts, expandStringPart{Value: varname, IsVariable: true})
		}

		if len(matches) == 0 {
			es.Parts = append(es.Parts, expandStringPart{Value: s, IsVariable: false})
		} else {
			last := matches[len(matches)-1]
			suffix := s[last[1]:]

			if len(suffix) != 0 {
				es.Parts = append(es.Parts, expandStringPart{Value: suffix, IsVariable: false})
			}
		}

		ret[name] = &es
	}

	for _, val := range ret {
		val.expand(ret)
	}

	return ret
}

// Configure runs the configure process with options as provided by the given
// data variable. If data is nil, the default options will be used
// (see NewOptions). Note that the data provided is simply passed to go-flags.
// For more information on flags parsing, see the documentation of go-flags.
// If GoConfig is not empty, then the go configuration will be written to the
// GoConfig file. Similarly, if Makefile is not empty, the Makefile will be
// written.
func Configure(data interface{}) (*Config, error) {
	if data == nil {
		data = NewOptions()
	}

	parser := flags.NewParser(data, flags.PrintErrors | flags.IgnoreUnknown)

	if _, err := parser.Parse(); err != nil {
		return nil, err
	}

	ret := &Config{
		Parser: parser,
	}

	ret.values, ret.valuesMap = ret.extract()
	ret.expanded = ret.expand()

	if len(GoConfig) != 0 {
		filename := GoConfig

		if !strings.HasSuffix(filename, ".go") {
			filename += ".go"
		}

		f, err := os.Create(filename)

		if err != nil {
			return nil, err
		}

		ret.WriteGoConfig(f)
		f.Close()
	}

	if len(Makefile) != 0 {
		f, err := os.Create(Makefile)

		if err != nil {
			return nil, err
		}

		ret.WriteMakefile(f)
		f.Close()

		os.Chmod(Makefile, 0755)

		f, err = os.OpenFile(path.Join(path.Dir(Makefile), "Makefile"),
		                     os.O_CREATE | os.O_EXCL | os.O_WRONLY,
		                     0644)

		if err == nil {
			fmt.Fprintf(f, "include %s\n", path.Base(Makefile))
			f.Close()
		}
	}

	return ret, nil
}

// Expand expands the variable value indicated by name
func (x *Config) Expand(name string) string {
	return x.expanded[name].expand(x.expanded)
}

// WriteGoConfig writes the go configuration file containing all the variable
// values to the given writer. Note that it will write a package line if
// the Package variable is not empty. The GoConfigVariable name will
// be used as the variable name for the configuration.
func (x *Config) WriteGoConfig(writer io.Writer) {
	if len(Package) > 0 {
		fmt.Fprintf(writer, "package %v\n\n", Package)
	}

	fmt.Fprintf(writer, "var %s = struct {\n", GoConfigVariable)
	values := make([]string, 0)

	variables := make([]string, 0, len(x.values))

	// Write all options
	for i, opt := range x.values {
		variables[i] = opt.LongName
	}

	sort.Strings(variables)

	for i, name := range variables {
		if i != 0 {
			io.WriteString(writer, "\n")
		}

		option := x.valuesMap[name]
		val := option.Value()

		fmt.Fprintf(writer, "\t// %s\n", option.Description)
		fmt.Fprintf(writer, "\t%v %T\n", name, val)

		var value string

		if _, ok := x.expanded[option.LongName]; ok {
			value = fmt.Sprintf("%#v", x.Expand(option.LongName))
		} else {
			value = fmt.Sprintf("%#v", val)
		}

		values = append(values, value)
	}

	if len(variables) > 0 {
		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "\t// Application version\n")
	io.WriteString(writer, "\tVersion []int\n")
	fmt.Fprintln(writer, "}{")

	for _, v := range values {
		fmt.Fprintf(writer, "\t%v,\n", v)
	}

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ", ")
		} else {
			io.WriteString(writer, "\t[]int{")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	fmt.Fprintln(writer, "},")
	fmt.Fprintln(writer, "}")
}

// WriteMakefile writes a Makefile for the given parser to the given writer.
// The Makefile contains the common build, clean, distclean, install and
// uninstall rules.
func (x *Config) WriteMakefile(writer io.Writer) {
	// Write a very basic makefile
	io.WriteString(writer, "#!/usr/bin/make -f\n\n")

	vars := make([]*expandString, 0, len(x.expanded))

	for name, v := range x.expanded {
		inserted := false

		// Insert into vars based on dependencies
		for i, vv := range vars {
			if vv.dependsOn(name) {
				tail := make([]*expandString, len(vars)-i)
				copy(tail, vars[i:])

				if i == 0 {
					vars = append([]*expandString{v}, vars...)
				} else {
					vars = append(append(vars[0:i], v), tail...)
				}

				inserted = true
				break
			}
		}

		if !inserted {
			vars = append(vars, v)
		}
	}

	io.WriteString(writer, "# Variables\n")

	for _, v := range vars {
		fmt.Fprintf(writer, "%s ?= ", v.Name)

		for _, part := range v.Parts {
			if part.IsVariable {
				fmt.Fprintf(writer, "$(%s)", part.Value)
			} else {
				fmt.Fprintf(writer, "%s", part.Value)
			}
		}

		io.WriteString(writer, "\n")
	}

	io.WriteString(writer, "version ?= ")

	for i, v := range Version {
		if i != 0 {
			io.WriteString(writer, ".")
		}

		fmt.Fprintf(writer, "%v", v)
	}

	io.WriteString(writer, "\n")
	fmt.Fprintf(writer, "major_version = %v\n", Version[0])

	if len(Version) > 1 {
		fmt.Fprintf(writer, "minor_version = %v\n", Version[1])
	}

	if len(Version) > 2 {
		fmt.Fprintf(writer, "micro_version = %v\n", Version[2])
	}

	io.WriteString(writer, "\n")

	target := Target

	if len(target) == 0 {
		pc := make([]uintptr, 3)
		n := runtime.Callers(1, pc)

		me, _ := runtime.FuncForPC(pc[0]).FileLine(pc[0])

		for i := 1; i < n; i++ {
			f := runtime.FuncForPC(pc[i])
			fname, _ := f.FileLine(pc[i])

			if fname != me {
				target = path.Base(path.Dir(fname))
				break
			}
		}
	}

	fmt.Fprintf(writer, "TARGET ?= %s\n", target)

	io.WriteString(writer, "\nSOURCES ?=")
	io.WriteString(writer, "\nSOURCES += $(wildcard *.go)")
	io.WriteString(writer, "\nSOURCES_UNIQUE = $(sort $(SOURCES))")

	io.WriteString(writer, "\n\n")

	io.WriteString(writer, "# Rules\n")
	io.WriteString(writer, "$(TARGET): $(SOURCES_UNIQUE)\n")
	io.WriteString(writer, "\tgo build -o $@\n\n")

	io.WriteString(writer, "clean:\n")
	io.WriteString(writer, "\trm -f $(TARGET)\n\n")

	io.WriteString(writer, "distclean: clean\n\n")

	io.WriteString(writer, "$(TARGET)_installdir ?= $(bindir)\n\n")

	io.WriteString(writer, "install: $(TARGET)\n")
	io.WriteString(writer, "\tmkdir -p $(DESTDIR)$($(TARGET)_installdir) && cp $(TARGET) $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, "uninstall:\n")
	io.WriteString(writer, "\trm -f $(DESTDIR)$($(TARGET)_installdir)/$(TARGET)\n\n")

	io.WriteString(writer, ".PHONY: install uninstall distclean clean")
}
<|endoftext|>


<commit_before>// Copyright 2016 The Serviced Authors.
// 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 auth

import (
	"crypto"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/fsnotify/fsnotify"
)

const (
	DelegateKeyFileName = "delegate.keys"
	MasterKeyFileName   = ".keys/master.keys"
)

var (
	delegateKeys HostKeys
	masterKeys   MasterKeys
	mKeyLock     sync.RWMutex
	dKeyCond     = &sync.Cond{L: &sync.Mutex{}}
)

type HostKeys struct {
	masterPublic crypto.PublicKey
	localPrivate crypto.PrivateKey
}

type MasterKeys struct {
	public  crypto.PublicKey
	private crypto.PrivateKey
}

func (h *HostKeys) Sign(message []byte) ([]byte, error) {
	signer, err := RSASigner(h.localPrivate)
	if err != nil {
		return nil, err
	}
	return signer.Sign(message)
}

func (h *HostKeys) Verify(message, signature []byte) error {
	verifier, err := RSAVerifier(h.masterPublic)
	if err != nil {
		return err
	}
	return verifier.Verify(message, signature)
}

func (m *MasterKeys) Sign(message []byte) ([]byte, error) {
	signer, err := RSASigner(m.private)
	if err != nil {
		return nil, err
	}
	return signer.Sign(message)
}

func (m *MasterKeys) Verify(message, signature []byte) error {
	verifier, err := RSAVerifier(m.public)
	if err != nil {
		return err
	}
	return verifier.Verify(message, signature)
}

// SignAsDelegate signs the given message with the private key local
// to the delegate running this process.
func SignAsDelegate(message []byte) ([]byte, error) {
	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	if delegateKeys.localPrivate == nil {
		return nil, ErrNoPrivateKey
	}
	return delegateKeys.Sign(message)
}

// SignAsMaster signs the given message with the master's private key
// will return an error if the delegate running this process is not the master
func SignAsMaster(message []byte) ([]byte, error) {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.private == nil {
		return nil, ErrNoPrivateKey
	}
	return masterKeys.Sign(message)
}

// VerifyMasterSignature verifies that a given message was signed by the master
// whose public key we have.
func VerifyMasterSignature(message, signature []byte) error {
	if err := verifyMasterSignatureAsMaster(message, signature); err != ErrNoPublicKey {
		return err
	}

	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	if delegateKeys.masterPublic == nil {
		return ErrNoPublicKey
	}
	return delegateKeys.Verify(message, signature)
}

func verifyMasterSignatureAsMaster(message, signature []byte) error {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.public == nil {
		return ErrNoPublicKey
	}
	return masterKeys.Verify(message, signature)

}

// GetMasterPublicKey() returns the public key of the master
//  It first checks to see if we have a set of master keys (i.e. we are the master)
//  It then checks the delegate keys.  If neither exists, returns ErrNoPublicKey
func GetMasterPublicKey() (crypto.PublicKey, error) {
	// check the master keys first
	if key, err := getMasterPublicKeyAsMaster(); err == nil {
		return key, err
	}

	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	if delegateKeys.masterPublic == nil {
		return nil, ErrNoPublicKey
	}
	return delegateKeys.masterPublic, nil
}

func getMasterPublicKeyAsMaster() (crypto.PublicKey, error) {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.public == nil {
		return nil, ErrNoPublicKey
	}
	return masterKeys.public, nil
}

func getMasterPrivateKey() (crypto.PrivateKey, error) {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.private == nil {
		return nil, ErrNoPrivateKey
	}
	return masterKeys.private, nil
}

// LoadKeysFromFile loads keys from a file on disk.
func LoadDelegateKeysFromFile(filename string) error {
	pub, priv, err := LoadKeyPairFromFile(filename)
	if err != nil {
		return err
	}
	updateDelegateKeys(pub, priv)
	return nil
}

// LoadMasterKeys sets the current master key pair to the one specified
func LoadMasterKeys(public crypto.PublicKey, private crypto.PrivateKey) {
	mKeyLock.Lock()
	defer mKeyLock.Unlock()
	masterKeys = MasterKeys{public, private}
}

// CreateOrLoadMasterKeys will load the master keys from disk
//  If the file does not exist, it will generate new keys and
//  write them to disk.
func CreateOrLoadMasterKeys(filename string) error {
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		if err = os.MkdirAll(path.Dir(filename), os.ModeDir|0700); err != nil {
			return err
		}

		pub, priv, err := GenerateRSAKeyPairPEM(nil)
		if err != nil {
			return err
		}
		err = DumpPEMKeyPairToFile(filename, pub, priv)
		if err != nil {
			return err
		}

	} else if err != nil {
		return err
	}

	return LoadMasterKeyFile(filename)
}

// LoadMasterKeyFile will load the master keys from disk if
//  the file exists.  If the file does not exist, it will
//  return an error
func LoadMasterKeyFile(filename string) error {
	publicKey, privateKey, err := LoadKeyPairFromFile(filename)
	if err != nil {
		return err
	}

	LoadMasterKeys(publicKey, privateKey)

	return nil
}

// DumpPEMKeyPairToFile dumps PEM-encoded public and private keys to a single file
func DumpPEMKeyPairToFile(filename string, public, private []byte) error {
	data, err := DumpRSAPEMKeyPair(public, private)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, data, 0666)
}

// LoadPEMKeyPair loads a private/public key pair from a reader over PEM-encoded data.
//  The private key is first, the public key is second.
func LoadKeyPairFromFile(filename string) (public crypto.PublicKey, private crypto.PrivateKey, err error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, nil, err
	}
	return LoadRSAKeyPairPackage(data)
}

// LoadDelegateKeysFromPEM loads the local delegate keys (master public, delegate private)
//  from PEM data passed in directly
//  Useful mostly for writing tests
func LoadDelegateKeysFromPEM(public, private []byte) error {
	pub, priv, err := LoadRSAKeyPair(public, private)
	if err != nil {
		return err
	}
	updateDelegateKeys(pub, priv)
	return nil
}

// LoadMasterKeysFromPEM loads the local master keys from PEM data passed in directly
//  Useful mostly for writing tests
func LoadMasterKeysFromPEM(public, private []byte) error {
	pub, priv, err := LoadRSAKeyPair(public, private)
	if err != nil {
		return err
	}

	mKeyLock.Lock()
	defer mKeyLock.Unlock()
	masterKeys = MasterKeys{pub, priv}
	return nil
}

// ClearKeys wipes the current state
func ClearKeys() {
	mKeyLock.Lock()
	masterKeys = MasterKeys{}
	mKeyLock.Unlock()
	updateDelegateKeys(nil, nil)
}

// WaitForDelegateKeys blocks until delegate keys are defined.
func WaitForDelegateKeys() {
	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	for delegateKeys.localPrivate == nil || delegateKeys.masterPublic == nil {
		dKeyCond.Wait()
	}
}

// WatchDelegateKeyFile watches the delegate key file on the filesystem and
// updates the internal delegate keys when changes are detected.
func WatchDelegateKeyFile(filename string, cancel chan interface{}) error {
	filename = filepath.Clean(filename)

	log := log.WithFields(logrus.Fields{
		"keyfile": filename,
	})

	loadKeys := func() {
		if err := LoadDelegateKeysFromFile(filename); err != nil {
			log.WithError(err).Warn("Unable to load delegate keys from file. Continuing to watch for changes")
		} else {
			log.Info("Loaded delegate keys from file")
		}
	}

	// Try an initial load without any file changes
	loadKeys()

	filechanges, err := NotifyOnChange(filename, fsnotify.Write|fsnotify.Create, cancel)
	if err != nil {
		return err
	}
	for _ = range filechanges {
		loadKeys()
	}
	return nil
}

func updateDelegateKeys(pub crypto.PublicKey, priv crypto.PrivateKey) {
	dKeyCond.L.Lock()
	delegateKeys = HostKeys{pub, priv}
	dKeyCond.L.Unlock()
	dKeyCond.Broadcast()
}

func delegateKeysChanged() <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		defer close(ch)
		dKeyCond.Wait()
	}()
	return ch
}
<commit_msg>Lock/Unlock in a cond.Wait scenario<commit_after>// Copyright 2016 The Serviced Authors.
// 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 auth

import (
	"crypto"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/fsnotify/fsnotify"
)

const (
	DelegateKeyFileName = "delegate.keys"
	MasterKeyFileName   = ".keys/master.keys"
)

var (
	delegateKeys HostKeys
	masterKeys   MasterKeys
	mKeyLock     sync.RWMutex
	dKeyCond     = &sync.Cond{L: &sync.Mutex{}}
)

type HostKeys struct {
	masterPublic crypto.PublicKey
	localPrivate crypto.PrivateKey
}

type MasterKeys struct {
	public  crypto.PublicKey
	private crypto.PrivateKey
}

func (h *HostKeys) Sign(message []byte) ([]byte, error) {
	signer, err := RSASigner(h.localPrivate)
	if err != nil {
		return nil, err
	}
	return signer.Sign(message)
}

func (h *HostKeys) Verify(message, signature []byte) error {
	verifier, err := RSAVerifier(h.masterPublic)
	if err != nil {
		return err
	}
	return verifier.Verify(message, signature)
}

func (m *MasterKeys) Sign(message []byte) ([]byte, error) {
	signer, err := RSASigner(m.private)
	if err != nil {
		return nil, err
	}
	return signer.Sign(message)
}

func (m *MasterKeys) Verify(message, signature []byte) error {
	verifier, err := RSAVerifier(m.public)
	if err != nil {
		return err
	}
	return verifier.Verify(message, signature)
}

// SignAsDelegate signs the given message with the private key local
// to the delegate running this process.
func SignAsDelegate(message []byte) ([]byte, error) {
	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	if delegateKeys.localPrivate == nil {
		return nil, ErrNoPrivateKey
	}
	return delegateKeys.Sign(message)
}

// SignAsMaster signs the given message with the master's private key
// will return an error if the delegate running this process is not the master
func SignAsMaster(message []byte) ([]byte, error) {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.private == nil {
		return nil, ErrNoPrivateKey
	}
	return masterKeys.Sign(message)
}

// VerifyMasterSignature verifies that a given message was signed by the master
// whose public key we have.
func VerifyMasterSignature(message, signature []byte) error {
	if err := verifyMasterSignatureAsMaster(message, signature); err != ErrNoPublicKey {
		return err
	}

	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	if delegateKeys.masterPublic == nil {
		return ErrNoPublicKey
	}
	return delegateKeys.Verify(message, signature)
}

func verifyMasterSignatureAsMaster(message, signature []byte) error {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.public == nil {
		return ErrNoPublicKey
	}
	return masterKeys.Verify(message, signature)

}

// GetMasterPublicKey() returns the public key of the master
//  It first checks to see if we have a set of master keys (i.e. we are the master)
//  It then checks the delegate keys.  If neither exists, returns ErrNoPublicKey
func GetMasterPublicKey() (crypto.PublicKey, error) {
	// check the master keys first
	if key, err := getMasterPublicKeyAsMaster(); err == nil {
		return key, err
	}

	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	if delegateKeys.masterPublic == nil {
		return nil, ErrNoPublicKey
	}
	return delegateKeys.masterPublic, nil
}

func getMasterPublicKeyAsMaster() (crypto.PublicKey, error) {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.public == nil {
		return nil, ErrNoPublicKey
	}
	return masterKeys.public, nil
}

func getMasterPrivateKey() (crypto.PrivateKey, error) {
	mKeyLock.RLock()
	defer mKeyLock.RUnlock()
	if masterKeys.private == nil {
		return nil, ErrNoPrivateKey
	}
	return masterKeys.private, nil
}

// LoadKeysFromFile loads keys from a file on disk.
func LoadDelegateKeysFromFile(filename string) error {
	pub, priv, err := LoadKeyPairFromFile(filename)
	if err != nil {
		return err
	}
	updateDelegateKeys(pub, priv)
	return nil
}

// LoadMasterKeys sets the current master key pair to the one specified
func LoadMasterKeys(public crypto.PublicKey, private crypto.PrivateKey) {
	mKeyLock.Lock()
	defer mKeyLock.Unlock()
	masterKeys = MasterKeys{public, private}
}

// CreateOrLoadMasterKeys will load the master keys from disk
//  If the file does not exist, it will generate new keys and
//  write them to disk.
func CreateOrLoadMasterKeys(filename string) error {
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		if err = os.MkdirAll(path.Dir(filename), os.ModeDir|0700); err != nil {
			return err
		}

		pub, priv, err := GenerateRSAKeyPairPEM(nil)
		if err != nil {
			return err
		}
		err = DumpPEMKeyPairToFile(filename, pub, priv)
		if err != nil {
			return err
		}

	} else if err != nil {
		return err
	}

	return LoadMasterKeyFile(filename)
}

// LoadMasterKeyFile will load the master keys from disk if
//  the file exists.  If the file does not exist, it will
//  return an error
func LoadMasterKeyFile(filename string) error {
	publicKey, privateKey, err := LoadKeyPairFromFile(filename)
	if err != nil {
		return err
	}

	LoadMasterKeys(publicKey, privateKey)

	return nil
}

// DumpPEMKeyPairToFile dumps PEM-encoded public and private keys to a single file
func DumpPEMKeyPairToFile(filename string, public, private []byte) error {
	data, err := DumpRSAPEMKeyPair(public, private)
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filename, data, 0666)
}

// LoadPEMKeyPair loads a private/public key pair from a reader over PEM-encoded data.
//  The private key is first, the public key is second.
func LoadKeyPairFromFile(filename string) (public crypto.PublicKey, private crypto.PrivateKey, err error) {
	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, nil, err
	}
	return LoadRSAKeyPairPackage(data)
}

// LoadDelegateKeysFromPEM loads the local delegate keys (master public, delegate private)
//  from PEM data passed in directly
//  Useful mostly for writing tests
func LoadDelegateKeysFromPEM(public, private []byte) error {
	pub, priv, err := LoadRSAKeyPair(public, private)
	if err != nil {
		return err
	}
	updateDelegateKeys(pub, priv)
	return nil
}

// LoadMasterKeysFromPEM loads the local master keys from PEM data passed in directly
//  Useful mostly for writing tests
func LoadMasterKeysFromPEM(public, private []byte) error {
	pub, priv, err := LoadRSAKeyPair(public, private)
	if err != nil {
		return err
	}

	mKeyLock.Lock()
	defer mKeyLock.Unlock()
	masterKeys = MasterKeys{pub, priv}
	return nil
}

// ClearKeys wipes the current state
func ClearKeys() {
	mKeyLock.Lock()
	masterKeys = MasterKeys{}
	mKeyLock.Unlock()
	updateDelegateKeys(nil, nil)
}

// WaitForDelegateKeys blocks until delegate keys are defined.
func WaitForDelegateKeys() {
	dKeyCond.L.Lock()
	defer dKeyCond.L.Unlock()
	for delegateKeys.localPrivate == nil || delegateKeys.masterPublic == nil {
		dKeyCond.Wait()
	}
}

// WatchDelegateKeyFile watches the delegate key file on the filesystem and
// updates the internal delegate keys when changes are detected.
func WatchDelegateKeyFile(filename string, cancel chan interface{}) error {
	filename = filepath.Clean(filename)

	log := log.WithFields(logrus.Fields{
		"keyfile": filename,
	})

	loadKeys := func() {
		if err := LoadDelegateKeysFromFile(filename); err != nil {
			log.WithError(err).Warn("Unable to load delegate keys from file. Continuing to watch for changes")
		} else {
			log.Info("Loaded delegate keys from file")
		}
	}

	// Try an initial load without any file changes
	loadKeys()

	filechanges, err := NotifyOnChange(filename, fsnotify.Write|fsnotify.Create, cancel)
	if err != nil {
		return err
	}
	for _ = range filechanges {
		loadKeys()
	}
	return nil
}

func updateDelegateKeys(pub crypto.PublicKey, priv crypto.PrivateKey) {
	dKeyCond.L.Lock()
	delegateKeys = HostKeys{pub, priv}
	dKeyCond.L.Unlock()
	dKeyCond.Broadcast()
}

func delegateKeysChanged() <-chan interface{} {
	ch := make(chan interface{})
	go func() {
		defer close(ch)
		dKeyCond.L.Lock()
		defer dKeyCond.L.Unlock()
		dKeyCond.Wait()
	}()
	return ch
}
<|endoftext|>