<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/lib/pq"
	"log"
	"net/http"
	"os"
	"runtime"
	"strconv"
)

const version = "0.0.0"

var database string
var user string
var pass string
var port uint

func init() {
	const (
		databaseDefault = ""
		databaseUsage   = "database"
		userDefault     = ""
		userUsage       = "database user"
		passDefault     = ""
		passUsage       = "database password"
		portDefault     = 0
		portUsage       = "http serve port"
	)
	flag.StringVar(&database, "database", databaseDefault, databaseUsage)
	flag.StringVar(&database, "d", databaseDefault, databaseUsage+" (shorthand)")
	flag.StringVar(&user, "user", userDefault, userUsage)
	flag.StringVar(&user, "u", userDefault, userUsage+" (shorthand)")
	flag.StringVar(&pass, "password", passDefault, passUsage)
	flag.StringVar(&pass, "pass", passDefault, passUsage+" (shorthand)")
	flag.UintVar(&port, "port", portDefault, portUsage)
	flag.UintVar(&port, "p", portDefault, portUsage+" (shorthand)")
}

func printUsage() {
	exeName := os.Args[0]
	fmt.Println(exeName + " " + version + " usage: ")
	fmt.Println("\t" + exeName + " -d star-database -u database-user-name -pass database-user-password -p serve-port")
	fmt.Println("flags:")
	flag.PrintDefaults()
	fmt.Println("example:\n\t" + exeName + " -d hyg -u jimbob -p nascarrulez -p 8008")
}

type Star struct {
	Id                int64
	Name              string
	X                 float64
	Y                 float64
	Z                 float64
	Color             float32
	AbsoluteMagnitude float32
	Spectrum          string
}

func (star *Star) Json() string {
	return "{" +
		"\"id\": " + strconv.Itoa(int(star.Id)) + ", " +
		"\"name\": \"" + star.Name + "\", " +
		"\"x\": " + strconv.FormatFloat(star.X, 'f', 15, 32) + ", " +
		"\"y\": " + strconv.FormatFloat(star.Y, 'f', 15, 32) + ", " +
		"\"z\": " + strconv.FormatFloat(star.Z, 'f', 15, 32) + ", " +
		"\"color\": " + strconv.FormatFloat(float64(star.Color), 'f', 15, 32) + ", " +
		"\"absolute-magnitude\": " + strconv.FormatFloat(float64(star.AbsoluteMagnitude), 'f', 15, 32) + ", " +
		"\"spectrum\": \"" + star.Spectrum + "\"}"
}

type NullStar struct {
	Id                sql.NullInt64
	Name              sql.NullString
	X                 sql.NullFloat64
	Y                 sql.NullFloat64
	Z                 sql.NullFloat64
	Color             sql.NullFloat64
	AbsoluteMagnitude sql.NullFloat64
	Spectrum          sql.NullString
}

/// returns a Star, with zero values for any null values in the NullStar
func (nstar *NullStar) Star() Star {
	var star Star
	if nstar.Id.Valid {
		star.Id = nstar.Id.Int64
	}
	if nstar.Name.Valid {
		star.Name = nstar.Name.String
	}
	if nstar.X.Valid {
		star.X = nstar.X.Float64
	}
	if nstar.Y.Valid {
		star.Y = nstar.Y.Float64
	}
	if nstar.Z.Valid {
		star.Z = nstar.Z.Float64
	}
	if nstar.Color.Valid {
		star.Color = float32(nstar.Color.Float64)
	}
	if nstar.AbsoluteMagnitude.Valid {
		star.AbsoluteMagnitude = float32(nstar.AbsoluteMagnitude.Float64)
	}
	if nstar.Spectrum.Valid {
		star.Spectrum = nstar.Spectrum.String
	}
	return star
}

func dbManager(user string, pass string, getStar chan struct {
	id       int64
	callback chan Star
}) {
	fmt.Println("Opening " + "postgres://" + user + ":" + pass + "@localhost/")
	db, err := sql.Open("postgres", "postgres://"+user+":"+pass+"@localhost/"+database)
	if err != nil {
		log.Fatal(err)
	}

	sql := "select propername, x, y, z, colorindex, absmag, spectrum from hygxyz where starid = $1"
	stmt, err := db.Prepare(sql)
	if err != nil {
		log.Fatal(err)
	}

	for {
		request := <-getStar

		rows, err := stmt.Query(request.id)

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

		if !rows.Next() {
			request.callback <- Star{}
			continue
		}

		var nstar NullStar
		err = rows.Scan(&nstar.Name, &nstar.X, &nstar.Y, &nstar.Z, &nstar.Color, &nstar.AbsoluteMagnitude, &nstar.Spectrum)
		if err != nil {
			log.Fatal(err)
		}
		star := nstar.Star()
		star.Id = request.id
		request.callback <- star // return zero values for null fields

		rows.Close()
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	if database == "" || user == "" || pass == "" || port == 0 || port > 65535 {
		printUsage()
		return
	}

	getStar := make(chan struct {
		id       int64
		callback chan Star
	}, 100)

	go dbManager(user, pass, getStar)

	http.HandleFunc("/star/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Println("Serving request to " + r.URL.Path)

		getStarCallback := make(chan Star)

		staridStr := r.URL.Path[len("/star/"):]
		starid, err := strconv.Atoi(staridStr)
		if err != nil {
			fmt.Println("ignoring request for non-integral star id")
			return
		}

		getStar <- struct {
			id       int64
			callback chan Star
		}{int64(starid), getStarCallback}
		star := <-getStarCallback
		fmt.Fprintf(w, star.Json())
	})

	fmt.Println("Serving on " + strconv.Itoa(int(port)) + "...")
	http.ListenAndServe(":"+strconv.Itoa(int(port)), nil)
}
<commit_msg>fixed http content-type and content-length headers<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	_ "github.com/lib/pq"
	"log"
	"net/http"
	"os"
	"runtime"
	"strconv"
)

const version = "0.0.0"

var database string
var user string
var pass string
var port uint

func init() {
	const (
		databaseDefault = ""
		databaseUsage   = "database"
		userDefault     = ""
		userUsage       = "database user"
		passDefault     = ""
		passUsage       = "database password"
		portDefault     = 0
		portUsage       = "http serve port"
	)
	flag.StringVar(&database, "database", databaseDefault, databaseUsage)
	flag.StringVar(&database, "d", databaseDefault, databaseUsage+" (shorthand)")
	flag.StringVar(&user, "user", userDefault, userUsage)
	flag.StringVar(&user, "u", userDefault, userUsage+" (shorthand)")
	flag.StringVar(&pass, "password", passDefault, passUsage)
	flag.StringVar(&pass, "pass", passDefault, passUsage+" (shorthand)")
	flag.UintVar(&port, "port", portDefault, portUsage)
	flag.UintVar(&port, "p", portDefault, portUsage+" (shorthand)")
}

func printUsage() {
	exeName := os.Args[0]
	fmt.Println(exeName + " " + version + " usage: ")
	fmt.Println("\t" + exeName + " -d star-database -u database-user-name -pass database-user-password -p serve-port")
	fmt.Println("flags:")
	flag.PrintDefaults()
	fmt.Println("example:\n\t" + exeName + " -d hyg -u jimbob -p nascarrulez -p 8008")
}

type Star struct {
	Id                int64
	Name              string
	X                 float64
	Y                 float64
	Z                 float64
	Color             float32
	AbsoluteMagnitude float32
	Spectrum          string
}

func (star *Star) Json() string {
	return "{" +
		"\"id\": " + strconv.Itoa(int(star.Id)) + ", " +
		"\"name\": \"" + star.Name + "\", " +
		"\"x\": " + strconv.FormatFloat(star.X, 'f', 15, 32) + ", " +
		"\"y\": " + strconv.FormatFloat(star.Y, 'f', 15, 32) + ", " +
		"\"z\": " + strconv.FormatFloat(star.Z, 'f', 15, 32) + ", " +
		"\"color\": " + strconv.FormatFloat(float64(star.Color), 'f', 15, 32) + ", " +
		"\"absolute-magnitude\": " + strconv.FormatFloat(float64(star.AbsoluteMagnitude), 'f', 15, 32) + ", " +
		"\"spectrum\": \"" + star.Spectrum + "\"}"
}

type NullStar struct {
	Id                sql.NullInt64
	Name              sql.NullString
	X                 sql.NullFloat64
	Y                 sql.NullFloat64
	Z                 sql.NullFloat64
	Color             sql.NullFloat64
	AbsoluteMagnitude sql.NullFloat64
	Spectrum          sql.NullString
}

/// returns a Star, with zero values for any null values in the NullStar
func (nstar *NullStar) Star() Star {
	var star Star
	if nstar.Id.Valid {
		star.Id = nstar.Id.Int64
	}
	if nstar.Name.Valid {
		star.Name = nstar.Name.String
	}
	if nstar.X.Valid {
		star.X = nstar.X.Float64
	}
	if nstar.Y.Valid {
		star.Y = nstar.Y.Float64
	}
	if nstar.Z.Valid {
		star.Z = nstar.Z.Float64
	}
	if nstar.Color.Valid {
		star.Color = float32(nstar.Color.Float64)
	}
	if nstar.AbsoluteMagnitude.Valid {
		star.AbsoluteMagnitude = float32(nstar.AbsoluteMagnitude.Float64)
	}
	if nstar.Spectrum.Valid {
		star.Spectrum = nstar.Spectrum.String
	}
	return star
}

func dbManager(user string, pass string, getStar chan struct {
	id       int64
	callback chan Star
}) {
	fmt.Println("Opening " + "postgres://" + user + ":" + pass + "@localhost/")
	db, err := sql.Open("postgres", "postgres://"+user+":"+pass+"@localhost/"+database)
	if err != nil {
		log.Fatal(err)
	}

	sql := "select propername, x, y, z, colorindex, absmag, spectrum from hygxyz where starid = $1"
	stmt, err := db.Prepare(sql)
	if err != nil {
		log.Fatal(err)
	}

	for {
		request := <-getStar

		rows, err := stmt.Query(request.id)

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

		if !rows.Next() {
			request.callback <- Star{}
			continue
		}

		var nstar NullStar
		err = rows.Scan(&nstar.Name, &nstar.X, &nstar.Y, &nstar.Z, &nstar.Color, &nstar.AbsoluteMagnitude, &nstar.Spectrum)
		if err != nil {
			log.Fatal(err)
		}
		star := nstar.Star()
		star.Id = request.id
		request.callback <- star // return zero values for null fields

		rows.Close()
	}
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.Parse()

	if database == "" || user == "" || pass == "" || port == 0 || port > 65535 {
		printUsage()
		return
	}

	getStar := make(chan struct {
		id       int64
		callback chan Star
	}, 100)

	go dbManager(user, pass, getStar)

	http.HandleFunc("/star/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Println("Serving request to " + r.URL.Path)

		getStarCallback := make(chan Star)

		staridStr := r.URL.Path[len("/star/"):]
		starid, err := strconv.Atoi(staridStr)
		if err != nil {
			fmt.Println("ignoring request for non-integral star id")
			return
		}

		getStar <- struct {
			id       int64
			callback chan Star
		}{int64(starid), getStarCallback}
		star := <-getStarCallback
		starjson := star.Json();

		w.Header().Add("Content-Type", "application/json")
		w.Header().Add("Content-Length", strconv.Itoa(len(starjson)))

		fmt.Fprintf(w, star.Json())
	})

	fmt.Println("Serving on " + strconv.Itoa(int(port)) + "...")
	http.ListenAndServe(":"+strconv.Itoa(int(port)), nil)
}
<|endoftext|>


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

// osconfig_agent interacts with the osconfig api.
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"syscall"
	"time"

	"github.com/GoogleCloudPlatform/guest-logging-go/logger"
	"github.com/GoogleCloudPlatform/osconfig/agentconfig"
	"github.com/GoogleCloudPlatform/osconfig/agentendpoint"
	"github.com/GoogleCloudPlatform/osconfig/clog"
	"github.com/GoogleCloudPlatform/osconfig/policies"
	"github.com/GoogleCloudPlatform/osconfig/tasker"
	"github.com/GoogleCloudPlatform/osconfig/util"
	"github.com/tarm/serial"

	_ "net/http/pprof"

	_ "google.golang.org/genproto/googleapis/rpc/errdetails"
)

var (
	version string
	profile = flag.Bool("profile", false, "serve profiling data at localhost:6060/debug/pprof")
)

func init() {
	if version == "" {
		version = "manual-" + time.Now().Format(time.RFC3339)
	}
	// We do this here so the -X value doesn't need the full path.
	agentconfig.SetVersion(version)

	os.MkdirAll(filepath.Dir(agentconfig.RestartFile()), 0755)
}

type serialPort struct {
	port string
}

func (s *serialPort) Write(b []byte) (int, error) {
	c := &serial.Config{Name: s.port, Baud: 115200}
	p, err := serial.OpenPort(c)
	if err != nil {
		return 0, err
	}
	defer p.Close()

	return p.Write(b)
}

var deferredFuncs []func()

// RegisterAgent is a blocking call, the RPC itself has retry logic baked in
// with jitter and backoff up to a total of 10 minutes.
// If client creation or register agent (after retries) fail we then wait for
// 5 minutes and try again.
func registerAgent(ctx context.Context) {
	for {
		if client, err := agentendpoint.NewClient(ctx); err != nil {
			logger.Errorf(err.Error())
		} else if err := client.RegisterAgent(ctx); err != nil {
			logger.Errorf(err.Error())
		} else {
			// RegisterAgent completed successfully.
			return
		}
		time.Sleep(5 * time.Minute)
	}
}

func run(ctx context.Context) {
	// Setup logging.
	opts := logger.LogOpts{LoggerName: "OSConfigAgent", UserAgent: agentconfig.UserAgent(), DisableLocalLogging: agentconfig.DisableLocalLogging()}
	if agentconfig.Stdout() {
		opts.Writers = []io.Writer{os.Stdout}
	}
	if runtime.GOOS == "windows" {
		opts.Writers = append(opts.Writers, &serialPort{"COM1"})
	}

	// If this call to WatchConfig fails (like a metadata error) we can't continue.
	if err := agentconfig.WatchConfig(ctx); err != nil {
		logger.Init(ctx, opts)
		logger.Fatalf("Error parsing metadata, agent cannot start: %v", err.Error())
	}
	opts.Debug = agentconfig.Debug()
	clog.DebugEnabled = agentconfig.Debug()
	opts.ProjectName = agentconfig.ProjectID()

	if err := logger.Init(ctx, opts); err != nil {
		fmt.Printf("Error initializing logger: %v", err)
		os.Exit(1)
	}
	ctx = clog.WithLabels(ctx, map[string]string{"instance_name": agentconfig.Name()})

	// Remove any existing restart file.
	if err := os.Remove(agentconfig.RestartFile()); err != nil && !os.IsNotExist(err) {
		clog.Errorf(ctx, "Error removing restart signal file: %v", err)
	}

	// On shutdown if the old restart file exists, and there is nothing else in that old directory,
	// cleanup that directory ignoring all errors.
	// This ensures we only cleanup this directory if we were using it with an old version of the agent.
	deferredFuncs = append(deferredFuncs, func() {
		if runtime.GOOS == "linux" && util.Exists(agentconfig.OldRestartFile()) {
			os.Remove(agentconfig.OldRestartFile())

			files, err := ioutil.ReadDir(filepath.Dir(agentconfig.OldRestartFile()))
			if err != nil || len(files) > 0 {
				return
			}
			os.RemoveAll(filepath.Dir(agentconfig.OldRestartFile()))
		}
	})

	deferredFuncs = append(deferredFuncs, logger.Close, func() { clog.Infof(ctx, "OSConfig Agent (version %s) shutting down.", agentconfig.Version()) })

	obtainLock()

	// obtainLock adds functions to clear the lock at close.
	logger.DeferredFatalFuncs = append(logger.DeferredFatalFuncs, deferredFuncs...)

	clog.Infof(ctx, "OSConfig Agent (version %s) started.", agentconfig.Version())

	// Call RegisterAgent at least once every day, on start calling
	// of RegisterAgent is handled in the service loop.
	go func() {
		c := time.Tick(24 * time.Hour)
		for range c {
			if agentconfig.TaskNotificationEnabled() || agentconfig.GuestPoliciesEnabled() {
				registerAgent(ctx)
			}
		}
	}()

	switch action := flag.Arg(0); action {
	case "", "run", "noservice":
		runServiceLoop(ctx)
	case "inventory", "osinventory":
		client, err := agentendpoint.NewClient(ctx)
		if err != nil {
			logger.Fatalf(err.Error())
		}
		tasker.Enqueue(ctx, "Report OSInventory", func() {
			client.ReportInventory(ctx)
		})
		tasker.Close()
		return
	case "gp", "policies", "guestpolicies", "ospackage":
		policies.Run(ctx)
		tasker.Close()
		return
	case "w", "waitfortasknotification", "ospatch":
		client, err := agentendpoint.NewClient(ctx)
		if err != nil {
			logger.Fatalf(err.Error())
		}
		client.WaitForTaskNotification(ctx)
		select {
		case <-ctx.Done():
		}
	default:
		logger.Fatalf("Unknown arg %q", action)
	}
}

func runTaskLoop(ctx context.Context, c chan struct{}) {
	var taskNotificationClient *agentendpoint.Client
	var err error
	for {
		// Set debug logging settings so that customers don't need to restart the agent.
		logger.SetDebugLogging(agentconfig.Debug())
		clog.DebugEnabled = agentconfig.Debug()
		if agentconfig.TaskNotificationEnabled() && taskNotificationClient == nil {
			// Call RegisterAgent now since we just either started running or were just enabled.
			// This call is blocking until successful as we can't continue unless register agent has completed.
			registerAgent(ctx)
		}
		if agentconfig.TaskNotificationEnabled() && (taskNotificationClient == nil || taskNotificationClient.Closed()) {
			// Start WaitForTaskNotification if we need to.
			taskNotificationClient, err = agentendpoint.NewClient(ctx)
			if err != nil {
				clog.Errorf(ctx, err.Error())
			} else {
				taskNotificationClient.WaitForTaskNotification(ctx)
			}
		} else if !agentconfig.TaskNotificationEnabled() && taskNotificationClient != nil && !taskNotificationClient.Closed() {
			// Cancel WaitForTaskNotification if we need to, this will block if there is
			// an existing current task running.
			if err := taskNotificationClient.Close(); err != nil {
				clog.Errorf(ctx, err.Error())
			}
		}

		// This is just to signal WaitForTaskNotification has run if needed.
		select {
		case c <- struct{}{}:
		default:
		}

		// Wait on any metadata config change.
		if err := agentconfig.WatchConfig(ctx); err != nil {
			clog.Errorf(ctx, err.Error())
		}
		select {
		case <-ctx.Done():
			return
		default:
			continue
		}
	}
}

// Runs internal functions that need to run on an interval.
func runInternalPeriodics(ctx context.Context) {
	ticker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()
	for {
		if _, err := os.Stat(agentconfig.RestartFile()); err == nil {
			clog.Infof(ctx, "Restart required marker file exists, beginning agent shutdown, waiting for tasks to complete.")
			tasker.Close()
			clog.Infof(ctx, "All tasks completed, stopping agent.")
			for _, f := range deferredFuncs {
				f()
			}
			os.Exit(2)
		}

		select {
		case <-ticker.C:
			continue
		case <-ctx.Done():
			return
		}
	}
}

func runServiceLoop(ctx context.Context) {
	go runInternalPeriodics(ctx)

	// This is just to ensure WaitForTaskNotification runs before any other tasks.
	c := make(chan struct{})
	// Configures WaitForTaskNotification, waits for config changes with WatchConfig.
	go runTaskLoop(ctx, c)
	// Don't continue any other tasks until WaitForTaskNotification has run.
	<-c

	// Runs functions that need to run on a set interval.
	ticker := time.NewTicker(agentconfig.SvcPollInterval())
	defer ticker.Stop()
	// First inventory run will be somewhere between 3 and 5 min.
	firstInventory := time.After(time.Duration(rand.Intn(120)+180) * time.Second)
	ranFirstInventory := false
	for {
		if agentconfig.GuestPoliciesEnabled() {
			policies.Run(ctx)
		}

		if agentconfig.OSInventoryEnabled() {
			if !ranFirstInventory {
				// Only run first inventory after the set waiting period or if the main poll ticker ticks.
				// The default SvcPollInterval is 10min so under normal circumstances firstInventory will
				// always fire first.
				select {
				case <-ticker.C:
				case <-firstInventory:
				case <-ctx.Done():
					return
				}
				ranFirstInventory = true
			}

			// This should always run after ospackage.SetConfig.
			tasker.Enqueue(ctx, "Report OSInventory", func() {
				client, err := agentendpoint.NewClient(ctx)
				if err != nil {
					logger.Errorf(err.Error())
				}
				client.ReportInventory(ctx)
				client.Close()
			})
		}

		select {
		case <-ticker.C:
			continue
		case <-ctx.Done():
			return
		}
	}
}

func main() {
	flag.Parse()
	ctx, cncl := context.WithCancel(context.Background())
	ctx = clog.WithLabels(ctx, map[string]string{"agent_version": agentconfig.Version()})
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		select {
		case <-c:
			cncl()
		}
	}()

	if *profile {
		go func() {
			fmt.Println(http.ListenAndServe("localhost:6060", nil))
		}()
	}

	switch action := flag.Arg(0); action {
	// wuaupdates just runs the packages.WUAUpdates function and returns it's output
	// as JSON on stdout. This avoids memory issues with the WUA api since this is
	// called often for Windows inventory runs.
	case "wuaupdates":
		if err := wuaUpdates(flag.Arg(1)); err != nil {
			fmt.Fprint(os.Stderr, err)
			os.Exit(1)
		}
		os.Exit(0)
	case "", "run":
		runService(ctx)
	default:
		run(ctx)
	}

	for _, f := range deferredFuncs {
		f()
	}
}
<commit_msg>Close client when RegisterAgent fails. (#436)<commit_after>//  Copyright 2018 Google Inc. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

// osconfig_agent interacts with the osconfig api.
package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"syscall"
	"time"

	"github.com/GoogleCloudPlatform/guest-logging-go/logger"
	"github.com/GoogleCloudPlatform/osconfig/agentconfig"
	"github.com/GoogleCloudPlatform/osconfig/agentendpoint"
	"github.com/GoogleCloudPlatform/osconfig/clog"
	"github.com/GoogleCloudPlatform/osconfig/policies"
	"github.com/GoogleCloudPlatform/osconfig/tasker"
	"github.com/GoogleCloudPlatform/osconfig/util"
	"github.com/tarm/serial"

	_ "net/http/pprof"

	_ "google.golang.org/genproto/googleapis/rpc/errdetails"
)

var (
	version string
	profile = flag.Bool("profile", false, "serve profiling data at localhost:6060/debug/pprof")
)

func init() {
	if version == "" {
		version = "manual-" + time.Now().Format(time.RFC3339)
	}
	// We do this here so the -X value doesn't need the full path.
	agentconfig.SetVersion(version)

	os.MkdirAll(filepath.Dir(agentconfig.RestartFile()), 0755)
}

type serialPort struct {
	port string
}

func (s *serialPort) Write(b []byte) (int, error) {
	c := &serial.Config{Name: s.port, Baud: 115200}
	p, err := serial.OpenPort(c)
	if err != nil {
		return 0, err
	}
	defer p.Close()

	return p.Write(b)
}

var deferredFuncs []func()

// RegisterAgent is a blocking call, the RPC itself has retry logic baked in
// with jitter and backoff up to a total of 10 minutes.
// If client creation or register agent (after retries) fail we then wait for
// 5 minutes and try again.
func registerAgent(ctx context.Context) {
	for {
		if client, err := agentendpoint.NewClient(ctx); err != nil {
			logger.Errorf(err.Error())
		} else if err := client.RegisterAgent(ctx); err != nil {
			logger.Errorf(err.Error())
			client.Close()
		} else {
			// RegisterAgent completed successfully.
			client.Close()
			return
		}
		time.Sleep(5 * time.Minute)
	}
}

func run(ctx context.Context) {
	// Setup logging.
	opts := logger.LogOpts{LoggerName: "OSConfigAgent", UserAgent: agentconfig.UserAgent(), DisableLocalLogging: agentconfig.DisableLocalLogging()}
	if agentconfig.Stdout() {
		opts.Writers = []io.Writer{os.Stdout}
	}
	if runtime.GOOS == "windows" {
		opts.Writers = append(opts.Writers, &serialPort{"COM1"})
	}

	// If this call to WatchConfig fails (like a metadata error) we can't continue.
	if err := agentconfig.WatchConfig(ctx); err != nil {
		logger.Init(ctx, opts)
		logger.Fatalf("Error parsing metadata, agent cannot start: %v", err.Error())
	}
	opts.Debug = agentconfig.Debug()
	clog.DebugEnabled = agentconfig.Debug()
	opts.ProjectName = agentconfig.ProjectID()

	if err := logger.Init(ctx, opts); err != nil {
		fmt.Printf("Error initializing logger: %v", err)
		os.Exit(1)
	}
	ctx = clog.WithLabels(ctx, map[string]string{"instance_name": agentconfig.Name()})

	// Remove any existing restart file.
	if err := os.Remove(agentconfig.RestartFile()); err != nil && !os.IsNotExist(err) {
		clog.Errorf(ctx, "Error removing restart signal file: %v", err)
	}

	// On shutdown if the old restart file exists, and there is nothing else in that old directory,
	// cleanup that directory ignoring all errors.
	// This ensures we only cleanup this directory if we were using it with an old version of the agent.
	deferredFuncs = append(deferredFuncs, func() {
		if runtime.GOOS == "linux" && util.Exists(agentconfig.OldRestartFile()) {
			os.Remove(agentconfig.OldRestartFile())

			files, err := ioutil.ReadDir(filepath.Dir(agentconfig.OldRestartFile()))
			if err != nil || len(files) > 0 {
				return
			}
			os.RemoveAll(filepath.Dir(agentconfig.OldRestartFile()))
		}
	})

	deferredFuncs = append(deferredFuncs, logger.Close, func() { clog.Infof(ctx, "OSConfig Agent (version %s) shutting down.", agentconfig.Version()) })

	obtainLock()

	// obtainLock adds functions to clear the lock at close.
	logger.DeferredFatalFuncs = append(logger.DeferredFatalFuncs, deferredFuncs...)

	clog.Infof(ctx, "OSConfig Agent (version %s) started.", agentconfig.Version())

	// Call RegisterAgent at least once every day, on start calling
	// of RegisterAgent is handled in the service loop.
	go func() {
		c := time.Tick(24 * time.Hour)
		for range c {
			if agentconfig.TaskNotificationEnabled() || agentconfig.GuestPoliciesEnabled() {
				registerAgent(ctx)
			}
		}
	}()

	switch action := flag.Arg(0); action {
	case "", "run", "noservice":
		runServiceLoop(ctx)
	case "inventory", "osinventory":
		client, err := agentendpoint.NewClient(ctx)
		if err != nil {
			logger.Fatalf(err.Error())
		}
		tasker.Enqueue(ctx, "Report OSInventory", func() {
			client.ReportInventory(ctx)
		})
		tasker.Close()
		return
	case "gp", "policies", "guestpolicies", "ospackage":
		policies.Run(ctx)
		tasker.Close()
		return
	case "w", "waitfortasknotification", "ospatch":
		client, err := agentendpoint.NewClient(ctx)
		if err != nil {
			logger.Fatalf(err.Error())
		}
		client.WaitForTaskNotification(ctx)
		select {
		case <-ctx.Done():
		}
	default:
		logger.Fatalf("Unknown arg %q", action)
	}
}

func runTaskLoop(ctx context.Context, c chan struct{}) {
	var taskNotificationClient *agentendpoint.Client
	var err error
	for {
		// Set debug logging settings so that customers don't need to restart the agent.
		logger.SetDebugLogging(agentconfig.Debug())
		clog.DebugEnabled = agentconfig.Debug()
		if agentconfig.TaskNotificationEnabled() && taskNotificationClient == nil {
			// Call RegisterAgent now since we just either started running or were just enabled.
			// This call is blocking until successful as we can't continue unless register agent has completed.
			registerAgent(ctx)
		}
		if agentconfig.TaskNotificationEnabled() && (taskNotificationClient == nil || taskNotificationClient.Closed()) {
			// Start WaitForTaskNotification if we need to.
			taskNotificationClient, err = agentendpoint.NewClient(ctx)
			if err != nil {
				clog.Errorf(ctx, err.Error())
			} else {
				taskNotificationClient.WaitForTaskNotification(ctx)
			}
		} else if !agentconfig.TaskNotificationEnabled() && taskNotificationClient != nil && !taskNotificationClient.Closed() {
			// Cancel WaitForTaskNotification if we need to, this will block if there is
			// an existing current task running.
			if err := taskNotificationClient.Close(); err != nil {
				clog.Errorf(ctx, err.Error())
			}
		}

		// This is just to signal WaitForTaskNotification has run if needed.
		select {
		case c <- struct{}{}:
		default:
		}

		// Wait on any metadata config change.
		if err := agentconfig.WatchConfig(ctx); err != nil {
			clog.Errorf(ctx, err.Error())
		}
		select {
		case <-ctx.Done():
			return
		default:
			continue
		}
	}
}

// Runs internal functions that need to run on an interval.
func runInternalPeriodics(ctx context.Context) {
	ticker := time.NewTicker(10 * time.Minute)
	defer ticker.Stop()
	for {
		if _, err := os.Stat(agentconfig.RestartFile()); err == nil {
			clog.Infof(ctx, "Restart required marker file exists, beginning agent shutdown, waiting for tasks to complete.")
			tasker.Close()
			clog.Infof(ctx, "All tasks completed, stopping agent.")
			for _, f := range deferredFuncs {
				f()
			}
			os.Exit(2)
		}

		select {
		case <-ticker.C:
			continue
		case <-ctx.Done():
			return
		}
	}
}

func runServiceLoop(ctx context.Context) {
	go runInternalPeriodics(ctx)

	// This is just to ensure WaitForTaskNotification runs before any other tasks.
	c := make(chan struct{})
	// Configures WaitForTaskNotification, waits for config changes with WatchConfig.
	go runTaskLoop(ctx, c)
	// Don't continue any other tasks until WaitForTaskNotification has run.
	<-c

	// Runs functions that need to run on a set interval.
	ticker := time.NewTicker(agentconfig.SvcPollInterval())
	defer ticker.Stop()
	// First inventory run will be somewhere between 3 and 5 min.
	firstInventory := time.After(time.Duration(rand.Intn(120)+180) * time.Second)
	ranFirstInventory := false
	for {
		if agentconfig.GuestPoliciesEnabled() {
			policies.Run(ctx)
		}

		if agentconfig.OSInventoryEnabled() {
			if !ranFirstInventory {
				// Only run first inventory after the set waiting period or if the main poll ticker ticks.
				// The default SvcPollInterval is 10min so under normal circumstances firstInventory will
				// always fire first.
				select {
				case <-ticker.C:
				case <-firstInventory:
				case <-ctx.Done():
					return
				}
				ranFirstInventory = true
			}

			// This should always run after ospackage.SetConfig.
			tasker.Enqueue(ctx, "Report OSInventory", func() {
				client, err := agentendpoint.NewClient(ctx)
				if err != nil {
					logger.Errorf(err.Error())
				}
				client.ReportInventory(ctx)
				client.Close()
			})
		}

		select {
		case <-ticker.C:
			continue
		case <-ctx.Done():
			return
		}
	}
}

func main() {
	flag.Parse()
	ctx, cncl := context.WithCancel(context.Background())
	ctx = clog.WithLabels(ctx, map[string]string{"agent_version": agentconfig.Version()})
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGTERM, syscall.SIGINT)
	go func() {
		select {
		case <-c:
			cncl()
		}
	}()

	if *profile {
		go func() {
			fmt.Println(http.ListenAndServe("localhost:6060", nil))
		}()
	}

	switch action := flag.Arg(0); action {
	// wuaupdates just runs the packages.WUAUpdates function and returns it's output
	// as JSON on stdout. This avoids memory issues with the WUA api since this is
	// called often for Windows inventory runs.
	case "wuaupdates":
		if err := wuaUpdates(flag.Arg(1)); err != nil {
			fmt.Fprint(os.Stderr, err)
			os.Exit(1)
		}
		os.Exit(0)
	case "", "run":
		runService(ctx)
	default:
		run(ctx)
	}

	for _, f := range deferredFuncs {
		f()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"context"
	"crypto"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/xenolf/lego/acme"
	"github.com/xenolf/lego/providers/dns/cloudflare"

	"github.com/ericchiang/k8s"
	apiv1 "github.com/ericchiang/k8s/api/v1"
)

const annotationLetsEncryptCertificate string = "estafette.io/letsencrypt-certificate"
const annotationLetsEncryptCertificateHostnames string = "estafette.io/letsencrypt-certificate-hostnames"

const annotationLetsEncryptCertificateState string = "estafette.io/letsencrypt-certificate-state"

// KubeLetsEncryptCertificateState represents the state of the secret with respect to Let's Encrypt certificates
type KubeLetsEncryptCertificateState struct {
	Hostnames   string `json:"hostnames"`
	LastRenewed string `json:"lastRenewed"`
}

var (
	addr = flag.String("listen-address", ":9101", "The address to listen on for HTTP requests.")

	// seed random number
	r = rand.New(rand.NewSource(time.Now().UnixNano()))
)

func main() {

	// create cloudflare api client
	cfAPIKey := os.Getenv("CF_API_KEY")
	if cfAPIKey == "" {
		log.Fatal("CF_API_KEY is required. Please set CF_API_KEY environment variable to your Cloudflare API key.")
	}
	cfAPIEmail := os.Getenv("CF_API_EMAIL")
	if cfAPIEmail == "" {
		log.Fatal("CF_API_EMAIL is required. Please set CF_API_KEY environment variable to your Cloudflare API email.")
	}

	// create kubernetes api client
	client, err := k8s.NewInClusterClient()
	if err != nil {
		log.Fatal(err)
	}

	// watch secrets for all namespaces
	go func() {
		// loop indefinitely
		for {
			fmt.Println("Watching secrets for all namespaces...")
			watcher, err := client.CoreV1().WatchSecrets(context.Background(), k8s.AllNamespaces)
			if err != nil {
				log.Println(err)
			} else {
				// loop indefinitely, unless it errors
				for {
					event, secret, err := watcher.Next()
					if err != nil {
						log.Println(err)
						break
					}

					if *event.Type == k8s.EventAdded || *event.Type == k8s.EventModified {
						fmt.Printf("Secret %v (namespace %v) has event of type %v, processing it...\n", *secret.Metadata.Name, *secret.Metadata.Namespace, *event.Type)
						err := processSecret(client, secret)
						if err != nil {
							continue
						}
					} else {
						fmt.Printf("Secret %v (namespace %v) has event of type %v, skipping it...\n", *secret.Metadata.Name, *secret.Metadata.Namespace, *event.Type)
					}
				}
			}

			// sleep random time between 22 and 37 seconds
			sleepTime := applyJitter(30)
			fmt.Printf("Sleeping for %v seconds...\n", sleepTime)
			time.Sleep(time.Duration(sleepTime) * time.Second)
		}
	}()

	// loop indefinitely
	for {

		// get secrets for all namespaces
		fmt.Println("Listing secrets for all namespaces...")
		secrets, err := client.CoreV1().ListSecrets(context.Background(), k8s.AllNamespaces)
		if err != nil {
			log.Println(err)
		}
		fmt.Printf("Cluster has %v secrets\n", len(secrets.Items))

		// loop all secrets
		if secrets != nil && secrets.Items != nil {
			for _, secret := range secrets.Items {

				err := processSecret(client, secret)
				if err != nil {
					continue
				}
			}
		}

		// sleep random time between 225 and 375 seconds
		sleepTime := applyJitter(300)
		fmt.Printf("Sleeping for %v seconds...\n", sleepTime)
		time.Sleep(time.Duration(sleepTime) * time.Second)
	}
}

func applyJitter(input int) (output int) {

	deviation := int(0.25 * float64(input))

	return input - deviation + r.Intn(2*deviation)
}

type LetsEncryptUser struct {
	Email        string
	Registration *acme.RegistrationResource
	key          crypto.PrivateKey
}

func (u LetsEncryptUser) GetEmail() string {
	return u.Email
}
func (u LetsEncryptUser) GetRegistration() *acme.RegistrationResource {
	return u.Registration
}
func (u LetsEncryptUser) GetPrivateKey() crypto.PrivateKey {
	return u.key
}

type Account struct {
	Email        string `json:"email"`
	key          crypto.PrivateKey
	Registration *acme.RegistrationResource `json:"registration"`
}

func processSecret(kubeclient *k8s.Client, secret *apiv1.Secret) error {

	cfAPIKey := os.Getenv("CF_API_KEY")
	cfAPIEmail := os.Getenv("CF_API_EMAIL")

	if &secret != nil && &secret.Metadata != nil && &secret.Metadata.Annotations != nil {

		// get annotations or set default value
		letsEncryptCertificate, ok := secret.Metadata.Annotations[annotationLetsEncryptCertificate]
		if !ok {
			letsEncryptCertificate = "false"
		}
		letsEncryptCertificateHostnames, ok := secret.Metadata.Annotations[annotationLetsEncryptCertificateHostnames]
		if !ok {
			letsEncryptCertificateHostnames = ""
		}

		// get state stored in annotations if present or set to empty struct
		var kubeLetsEncryptCertificateState KubeLetsEncryptCertificateState
		kubeLetsEncryptCertificateStateString, ok := secret.Metadata.Annotations[annotationLetsEncryptCertificateState]
		if err := json.Unmarshal([]byte(kubeLetsEncryptCertificateStateString), &kubeLetsEncryptCertificateState); err != nil {
			// couldn't deserialize, setting to default struct
			kubeLetsEncryptCertificateState = KubeLetsEncryptCertificateState{}
		}

		// parse last renewed time from state
		lastRenewed := time.Time{}
		if kubeLetsEncryptCertificateState.LastRenewed != "" {
			var err error
			lastRenewed, err = time.Parse(time.RFC3339, kubeLetsEncryptCertificateState.LastRenewed)
			if err != nil {
				lastRenewed = time.Time{}
			}
		}

		durationSinceLastRenewed := time.Since(lastRenewed)

		// check if letsencrypt is enabled for this secret, hostnames are set and either the hostnames have changed or the certificate is older than 60 days
		if letsEncryptCertificate == "true" && len(letsEncryptCertificateHostnames) > 0 && (letsEncryptCertificateHostnames != kubeLetsEncryptCertificateState.Hostnames || durationSinceLastRenewed.Hours() > float64(60*24)) {

			fmt.Printf("Certificates in secret %v (namespace %v) are more than 60 days old or hostnames have changed (%v), renewing them with Let's Encrypt...\n", *secret.Metadata.Name, *secret.Metadata.Namespace, letsEncryptCertificateHostnames)

			// load account.json
			fmt.Printf("Loading account.json for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			fileBytes, err := ioutil.ReadFile("/account/account.json")
			if err != nil {
				log.Println(err)
				return err
			}
			var letsEncryptUser LetsEncryptUser
			err = json.Unmarshal(fileBytes, &letsEncryptUser)
			if err != nil {
				log.Println(err)
				return err
			}

			// load private key
			fmt.Printf("Loading account.key for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			privateKey, err := loadPrivateKey("/account/account.key")
			if err != nil {
				log.Println(err)
				return err
			}
			letsEncryptUser.key = privateKey

			// set dns timeout
			fmt.Printf("Setting acme dns timeout to 600 seconds for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			acme.DNSTimeout = time.Duration(600) * time.Second

			// create letsencrypt acme client
			fmt.Printf("Creating acme client for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			client, err := acme.NewClient("https://acme-v01.api.letsencrypt.org/directory", letsEncryptUser, acme.RSA2048)
			if err != nil {
				log.Println(err)
				return err
			}

			// get dns challenge
			fmt.Printf("Creating cloudflare provider for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			var provider acme.ChallengeProvider
			provider, err = cloudflare.NewDNSProviderCredentials(cfAPIEmail, cfAPIKey)
			if err != nil {
				log.Println(err)
				return err
			}

			// clean up acme challenge records in advance
			hostnames := strings.Split(letsEncryptCertificateHostnames, ",")
			for _, hostname := range hostnames {
				fmt.Printf("Cleaning up TXT record _acme-challenge.%v for secret %v (namespace %v)...\n", hostname, *secret.Metadata.Name, *secret.Metadata.Namespace)
				provider.CleanUp("_acme-challenge."+hostname, "", "123d==")
			}

			// set challenge and provider
			client.SetChallengeProvider(acme.DNS01, provider)
			client.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.TLSSNI01})

			// get certificate
			fmt.Printf("Obtaining certificate for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			var certificate acme.CertificateResource
			var failures map[string]error
			certificate, failures = client.ObtainCertificate(hostnames, true, nil, false)

			// clean up acme challenge records afterwards
			for _, hostname := range hostnames {
				fmt.Printf("Cleaning up TXT record _acme-challenge.%v for secret %v (namespace %v)...\n", hostname, *secret.Metadata.Name, *secret.Metadata.Namespace)
				provider.CleanUp("_acme-challenge."+hostname, "", "123d==")
			}

			if len(failures) > 0 {
				for k, v := range failures {
					log.Printf("[%s] Could not obtain certificates\n\t%s", k, v.Error())
				}

				return errors.New("Generating certificates has failed")
			}

			// update the secret
			kubeLetsEncryptCertificateState.Hostnames = letsEncryptCertificateHostnames
			kubeLetsEncryptCertificateState.LastRenewed = time.Now().Format(time.RFC3339)

			fmt.Printf("Updating secret %v (namespace %v) because certificates have changed...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)

			// serialize state and store it in the annotation
			kubeLetsEncryptCertificateStateByteArray, err := json.Marshal(kubeLetsEncryptCertificateState)
			if err != nil {
				log.Println(err)
				return err
			}
			secret.Metadata.Annotations[annotationLetsEncryptCertificateState] = string(kubeLetsEncryptCertificateStateByteArray)

			// store the certificates
			secret.Data["ssl.crt"] = certificate.Certificate
			secret.Data["ssl.key"] = certificate.PrivateKey
			secret.Data["ssl.pem"] = bytes.Join([][]byte{certificate.Certificate, certificate.PrivateKey}, []byte{})

			// update secret, because the data and state annotation have changed
			secret, err = kubeclient.CoreV1().UpdateSecret(context.Background(), secret)
			if err != nil {
				log.Println(err)
				return err
			}
		}
	}

	return nil
}

func loadPrivateKey(file string) (crypto.PrivateKey, error) {
	keyBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}

	keyBlock, _ := pem.Decode(keyBytes)

	switch keyBlock.Type {
	case "RSA PRIVATE KEY":
		return x509.ParsePKCS1PrivateKey(keyBlock.Bytes)
	case "EC PRIVATE KEY":
		return x509.ParseECPrivateKey(keyBlock.Bytes)
	}

	return nil, errors.New("Unknown private key type.")
}
<commit_msg>if secret data map does not exist make it first<commit_after>package main

import (
	"bytes"
	"context"
	"crypto"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"strings"
	"time"

	"github.com/xenolf/lego/acme"
	"github.com/xenolf/lego/providers/dns/cloudflare"

	"github.com/ericchiang/k8s"
	apiv1 "github.com/ericchiang/k8s/api/v1"
)

const annotationLetsEncryptCertificate string = "estafette.io/letsencrypt-certificate"
const annotationLetsEncryptCertificateHostnames string = "estafette.io/letsencrypt-certificate-hostnames"

const annotationLetsEncryptCertificateState string = "estafette.io/letsencrypt-certificate-state"

// KubeLetsEncryptCertificateState represents the state of the secret with respect to Let's Encrypt certificates
type KubeLetsEncryptCertificateState struct {
	Hostnames   string `json:"hostnames"`
	LastRenewed string `json:"lastRenewed"`
}

var (
	addr = flag.String("listen-address", ":9101", "The address to listen on for HTTP requests.")

	// seed random number
	r = rand.New(rand.NewSource(time.Now().UnixNano()))
)

func main() {

	// create cloudflare api client
	cfAPIKey := os.Getenv("CF_API_KEY")
	if cfAPIKey == "" {
		log.Fatal("CF_API_KEY is required. Please set CF_API_KEY environment variable to your Cloudflare API key.")
	}
	cfAPIEmail := os.Getenv("CF_API_EMAIL")
	if cfAPIEmail == "" {
		log.Fatal("CF_API_EMAIL is required. Please set CF_API_KEY environment variable to your Cloudflare API email.")
	}

	// create kubernetes api client
	client, err := k8s.NewInClusterClient()
	if err != nil {
		log.Fatal(err)
	}

	// watch secrets for all namespaces
	go func() {
		// loop indefinitely
		for {
			fmt.Println("Watching secrets for all namespaces...")
			watcher, err := client.CoreV1().WatchSecrets(context.Background(), k8s.AllNamespaces)
			if err != nil {
				log.Println(err)
			} else {
				// loop indefinitely, unless it errors
				for {
					event, secret, err := watcher.Next()
					if err != nil {
						log.Println(err)
						break
					}

					if *event.Type == k8s.EventAdded || *event.Type == k8s.EventModified {
						fmt.Printf("Secret %v (namespace %v) has event of type %v, processing it...\n", *secret.Metadata.Name, *secret.Metadata.Namespace, *event.Type)
						err := processSecret(client, secret)
						if err != nil {
							continue
						}
					} else {
						fmt.Printf("Secret %v (namespace %v) has event of type %v, skipping it...\n", *secret.Metadata.Name, *secret.Metadata.Namespace, *event.Type)
					}
				}
			}

			// sleep random time between 22 and 37 seconds
			sleepTime := applyJitter(30)
			fmt.Printf("Sleeping for %v seconds...\n", sleepTime)
			time.Sleep(time.Duration(sleepTime) * time.Second)
		}
	}()

	// loop indefinitely
	for {

		// get secrets for all namespaces
		fmt.Println("Listing secrets for all namespaces...")
		secrets, err := client.CoreV1().ListSecrets(context.Background(), k8s.AllNamespaces)
		if err != nil {
			log.Println(err)
		}
		fmt.Printf("Cluster has %v secrets\n", len(secrets.Items))

		// loop all secrets
		if secrets != nil && secrets.Items != nil {
			for _, secret := range secrets.Items {

				err := processSecret(client, secret)
				if err != nil {
					continue
				}
			}
		}

		// sleep random time between 225 and 375 seconds
		sleepTime := applyJitter(300)
		fmt.Printf("Sleeping for %v seconds...\n", sleepTime)
		time.Sleep(time.Duration(sleepTime) * time.Second)
	}
}

func applyJitter(input int) (output int) {

	deviation := int(0.25 * float64(input))

	return input - deviation + r.Intn(2*deviation)
}

type LetsEncryptUser struct {
	Email        string
	Registration *acme.RegistrationResource
	key          crypto.PrivateKey
}

func (u LetsEncryptUser) GetEmail() string {
	return u.Email
}
func (u LetsEncryptUser) GetRegistration() *acme.RegistrationResource {
	return u.Registration
}
func (u LetsEncryptUser) GetPrivateKey() crypto.PrivateKey {
	return u.key
}

type Account struct {
	Email        string `json:"email"`
	key          crypto.PrivateKey
	Registration *acme.RegistrationResource `json:"registration"`
}

func processSecret(kubeclient *k8s.Client, secret *apiv1.Secret) error {

	cfAPIKey := os.Getenv("CF_API_KEY")
	cfAPIEmail := os.Getenv("CF_API_EMAIL")

	if &secret != nil && &secret.Metadata != nil && &secret.Metadata.Annotations != nil {

		// get annotations or set default value
		letsEncryptCertificate, ok := secret.Metadata.Annotations[annotationLetsEncryptCertificate]
		if !ok {
			letsEncryptCertificate = "false"
		}
		letsEncryptCertificateHostnames, ok := secret.Metadata.Annotations[annotationLetsEncryptCertificateHostnames]
		if !ok {
			letsEncryptCertificateHostnames = ""
		}

		// get state stored in annotations if present or set to empty struct
		var kubeLetsEncryptCertificateState KubeLetsEncryptCertificateState
		kubeLetsEncryptCertificateStateString, ok := secret.Metadata.Annotations[annotationLetsEncryptCertificateState]
		if err := json.Unmarshal([]byte(kubeLetsEncryptCertificateStateString), &kubeLetsEncryptCertificateState); err != nil {
			// couldn't deserialize, setting to default struct
			kubeLetsEncryptCertificateState = KubeLetsEncryptCertificateState{}
		}

		// parse last renewed time from state
		lastRenewed := time.Time{}
		if kubeLetsEncryptCertificateState.LastRenewed != "" {
			var err error
			lastRenewed, err = time.Parse(time.RFC3339, kubeLetsEncryptCertificateState.LastRenewed)
			if err != nil {
				lastRenewed = time.Time{}
			}
		}

		durationSinceLastRenewed := time.Since(lastRenewed)

		// check if letsencrypt is enabled for this secret, hostnames are set and either the hostnames have changed or the certificate is older than 60 days
		if letsEncryptCertificate == "true" && len(letsEncryptCertificateHostnames) > 0 && (letsEncryptCertificateHostnames != kubeLetsEncryptCertificateState.Hostnames || durationSinceLastRenewed.Hours() > float64(60*24)) {

			fmt.Printf("Certificates in secret %v (namespace %v) are more than 60 days old or hostnames have changed (%v), renewing them with Let's Encrypt...\n", *secret.Metadata.Name, *secret.Metadata.Namespace, letsEncryptCertificateHostnames)

			// load account.json
			fmt.Printf("Loading account.json for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			fileBytes, err := ioutil.ReadFile("/account/account.json")
			if err != nil {
				log.Println(err)
				return err
			}
			var letsEncryptUser LetsEncryptUser
			err = json.Unmarshal(fileBytes, &letsEncryptUser)
			if err != nil {
				log.Println(err)
				return err
			}

			// load private key
			fmt.Printf("Loading account.key for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			privateKey, err := loadPrivateKey("/account/account.key")
			if err != nil {
				log.Println(err)
				return err
			}
			letsEncryptUser.key = privateKey

			// set dns timeout
			fmt.Printf("Setting acme dns timeout to 600 seconds for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			acme.DNSTimeout = time.Duration(600) * time.Second

			// create letsencrypt acme client
			fmt.Printf("Creating acme client for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			client, err := acme.NewClient("https://acme-v01.api.letsencrypt.org/directory", letsEncryptUser, acme.RSA2048)
			if err != nil {
				log.Println(err)
				return err
			}

			// get dns challenge
			fmt.Printf("Creating cloudflare provider for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			var provider acme.ChallengeProvider
			provider, err = cloudflare.NewDNSProviderCredentials(cfAPIEmail, cfAPIKey)
			if err != nil {
				log.Println(err)
				return err
			}

			// clean up acme challenge records in advance
			hostnames := strings.Split(letsEncryptCertificateHostnames, ",")
			for _, hostname := range hostnames {
				fmt.Printf("Cleaning up TXT record _acme-challenge.%v for secret %v (namespace %v)...\n", hostname, *secret.Metadata.Name, *secret.Metadata.Namespace)
				provider.CleanUp("_acme-challenge."+hostname, "", "123d==")
			}

			// set challenge and provider
			client.SetChallengeProvider(acme.DNS01, provider)
			client.ExcludeChallenges([]acme.Challenge{acme.HTTP01, acme.TLSSNI01})

			// get certificate
			fmt.Printf("Obtaining certificate for secret %v (namespace %v)...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)
			var certificate acme.CertificateResource
			var failures map[string]error
			certificate, failures = client.ObtainCertificate(hostnames, true, nil, false)

			// clean up acme challenge records afterwards
			for _, hostname := range hostnames {
				fmt.Printf("Cleaning up TXT record _acme-challenge.%v for secret %v (namespace %v)...\n", hostname, *secret.Metadata.Name, *secret.Metadata.Namespace)
				provider.CleanUp("_acme-challenge."+hostname, "", "123d==")
			}

			if len(failures) > 0 {
				for k, v := range failures {
					log.Printf("[%s] Could not obtain certificates\n\t%s", k, v.Error())
				}

				return errors.New("Generating certificates has failed")
			}

			// update the secret
			kubeLetsEncryptCertificateState.Hostnames = letsEncryptCertificateHostnames
			kubeLetsEncryptCertificateState.LastRenewed = time.Now().Format(time.RFC3339)

			fmt.Printf("Updating secret %v (namespace %v) because certificates have changed...\n", *secret.Metadata.Name, *secret.Metadata.Namespace)

			// serialize state and store it in the annotation
			kubeLetsEncryptCertificateStateByteArray, err := json.Marshal(kubeLetsEncryptCertificateState)
			if err != nil {
				log.Println(err)
				return err
			}
			secret.Metadata.Annotations[annotationLetsEncryptCertificateState] = string(kubeLetsEncryptCertificateStateByteArray)

			// store the certificates
			if secret.Data == nil {
				secret.Data = make(map[string][]byte)
			}

			secret.Data["ssl.crt"] = certificate.Certificate
			secret.Data["ssl.key"] = certificate.PrivateKey
			secret.Data["ssl.pem"] = bytes.Join([][]byte{certificate.Certificate, certificate.PrivateKey}, []byte{})

			// update secret, because the data and state annotation have changed
			secret, err = kubeclient.CoreV1().UpdateSecret(context.Background(), secret)
			if err != nil {
				log.Println(err)
				return err
			}
		}
	}

	return nil
}

func loadPrivateKey(file string) (crypto.PrivateKey, error) {
	keyBytes, err := ioutil.ReadFile(file)
	if err != nil {
		return nil, err
	}

	keyBlock, _ := pem.Decode(keyBytes)

	switch keyBlock.Type {
	case "RSA PRIVATE KEY":
		return x509.ParsePKCS1PrivateKey(keyBlock.Bytes)
	case "EC PRIVATE KEY":
		return x509.ParseECPrivateKey(keyBlock.Bytes)
	}

	return nil, errors.New("Unknown private key type.")
}
<|endoftext|>


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

package main

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

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

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

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

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

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

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

		// Purchase tickets according to simulated demand curve.
		nextTicketPrice := s.nextTicketPriceFunc()
		if nextHeight%stakeDiffWindowSize == 0 && nextHeight != 0 {
			demand := s.calcDemand(nextHeight, nextTicketPrice)
			demandPerWindow = int32(float64(maxTicketsPerWindow) * demand)
		}

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

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

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

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

	return nil
}

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

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

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

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

	// Generate the simulation results and open them in a browser.
	if err := generateResults(sim); err != nil {
		fmt.Println(err)
		return
	}
}
<commit_msg>Panic if ticket price func returns price below min.<commit_after>// Copyright (c) 2017 Dave Collins
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package main

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

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

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

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

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

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

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	return nil
}

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

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

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

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

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


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/ogier/pflag"
	"github.com/yuya-takeyama/argf"
)

var (
	cmdName = "ltsvp"
	version = "0.5.0"

	flagset         = pflag.NewFlagSet(cmdName, pflag.ContinueOnError)
	list            = flagset.StringP("keys", "k", "", "")
	outputDelimiter = flagset.StringP("output-delimiter", "D", "\t", "")
	remainLTSV      = flagset.BoolP("remain-ltsv", "r", false, "")
	isHelp          = flagset.BoolP("help", "h", false, "")
	isVersion       = flagset.BoolP("version", "v", false, "")
)

func printUsage() {
	fmt.Fprintf(os.Stderr, `
Usage: %s OPTION... [FILE]...
Print selected parts of LTSV from each FILE to standard output.

Options:
  -k, --keys=LIST
                 select only these keys (required)
  -D, --output-delimiter=STRING
                 use STRING to separate parts (default: \t)
  -r, --remain-ltsv
                 print selected parts as LTSV
  -h, --help
                 display this help text and exit
  -v, --version
                 output version information and exit

LIST is made up of keys separated by commas.
  host           # Select host
  host,time,ua   # Select host, time, and ua
`[1:], cmdName)
}

func printVersion() {
	fmt.Fprintln(os.Stderr, version)
}

func printErr(err interface{}) {
	fmt.Fprintf(os.Stderr, "%s: %s\n", cmdName, err)
}

func guideToHelp() {
	fmt.Fprintf(os.Stderr, "Try '%s --help' for more information.\n", cmdName)
}

func do(l *LTSVScanner) error {
	for l.Scan() {
		fmt.Println(l.Text())
	}
	return l.Err()
}

func main() {
	flagset.SetOutput(ioutil.Discard)
	if err := flagset.Parse(os.Args[1:]); err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}
	if *isHelp {
		printUsage()
		os.Exit(0)
	}
	if *isVersion {
		printVersion()
		os.Exit(0)
	}

	specifiedList := false
	flagset.Visit(func(f *pflag.Flag) {
		if f.Name == "keys" {
			specifiedList = true
		}
	})
	if !specifiedList {
		printErr("no specify LIST")
		guideToHelp()
		os.Exit(2)
	}
	keys := ParseKeysList(*list)

	r, err := argf.From(flagset.Args())
	if err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}

	l := NewLTSVScanner(keys, r)
	l.OutputDelimiter = *outputDelimiter
	l.RemainLTSV = *remainLTSV
	if err := do(l); err != nil {
		printErr(err)
		os.Exit(1)
	}
	os.Exit(0)
}
<commit_msg>s/&main.version/&main.cmdVersion/<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/ogier/pflag"
	"github.com/yuya-takeyama/argf"
)

var (
	cmdName    = "ltsvp"
	cmdVersion = "0.5.0"

	flagset         = pflag.NewFlagSet(cmdName, pflag.ContinueOnError)
	list            = flagset.StringP("keys", "k", "", "")
	outputDelimiter = flagset.StringP("output-delimiter", "D", "\t", "")
	remainLTSV      = flagset.BoolP("remain-ltsv", "r", false, "")
	isHelp          = flagset.BoolP("help", "h", false, "")
	isVersion       = flagset.BoolP("version", "v", false, "")
)

func printUsage() {
	fmt.Fprintf(os.Stderr, `
Usage: %s OPTION... [FILE]...
Print selected parts of LTSV from each FILE to standard output.

Options:
  -k, --keys=LIST
                 select only these keys (required)
  -D, --output-delimiter=STRING
                 use STRING to separate parts (default: \t)
  -r, --remain-ltsv
                 print selected parts as LTSV
  -h, --help
                 display this help text and exit
  -v, --version
                 output version information and exit

LIST is made up of keys separated by commas.
  host           # Select host
  host,time,ua   # Select host, time, and ua
`[1:], cmdName)
}

func printVersion() {
	fmt.Fprintln(os.Stderr, cmdVersion)
}

func printErr(err interface{}) {
	fmt.Fprintf(os.Stderr, "%s: %s\n", cmdName, err)
}

func guideToHelp() {
	fmt.Fprintf(os.Stderr, "Try '%s --help' for more information.\n", cmdName)
}

func do(l *LTSVScanner) error {
	for l.Scan() {
		fmt.Println(l.Text())
	}
	return l.Err()
}

func main() {
	flagset.SetOutput(ioutil.Discard)
	if err := flagset.Parse(os.Args[1:]); err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}
	if *isHelp {
		printUsage()
		os.Exit(0)
	}
	if *isVersion {
		printVersion()
		os.Exit(0)
	}

	specifiedList := false
	flagset.Visit(func(f *pflag.Flag) {
		if f.Name == "keys" {
			specifiedList = true
		}
	})
	if !specifiedList {
		printErr("no specify LIST")
		guideToHelp()
		os.Exit(2)
	}
	keys := ParseKeysList(*list)

	r, err := argf.From(flagset.Args())
	if err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}

	l := NewLTSVScanner(keys, r)
	l.OutputDelimiter = *outputDelimiter
	l.RemainLTSV = *remainLTSV
	if err := do(l); err != nil {
		printErr(err)
		os.Exit(1)
	}
	os.Exit(0)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/ChrisMcKenzie/dropship/ssh"
	"github.com/google/go-github/github"
	"github.com/julienschmidt/httprouter"
	"github.com/the-control-group/data-service-api/logger"
	"github.com/thoas/stats"
)

type (
	Payload struct {
		Deployment github.Deployment `json:"deployment"`
		Repository github.Repository `json:"repository"`
		Sender     github.User       `json:"sender"`
	}

	Server struct {
		Address  string `json:"address"`
		User     string `json:"username"`
		Password string `json:"password"`
	}

	Deployment struct {
		Command string `json:"command"`
		Servers struct {
			Provider string          `json:"provider"`
			Options  json.RawMessage `json:"options"`
		} `json:"servers"`
	}
)

var log = logger.NewLogger()
var port = "3000"

func Logger(h httprouter.Handle) httprouter.Handle {
	return func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		log.Infof("[%s] %s", r.Method, r.URL)
		h(w, r, p)
	}
}

func handleGithubDeploy(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	var err error
	payload := Payload{}
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&payload); err != nil {
		log.Error(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	deployment := Deployment{}
	if err := json.Unmarshal(payload.Deployment.Payload, &deployment); err != nil {
		log.Error(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	servers := []Server{}
	if deployment.Servers.Provider == "consul" {
		// TODO(ChrisMcKenzie): Write Consul provider
	} else if deployment.Servers.Provider == "payload" {
		servers, err = getServersFromPayload(deployment.Servers.Options)
		if err != nil {
			log.Error(err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
	}

	run(deployment.Command, servers)

	fmt.Fprintf(w, "deploying %s", payload.Repository.Name)
}

func run(c string, servers []Server) {
	for _, server := range servers {
		go func() {
			log.Debugf("Excecuting Command on %s", server.Address)
			res := ssh.Execute(
				c,
				server.Address,
				ssh.NewClientConfig(server.User, server.Password),
			)
			log.Debug(res)
		}()
	}
}

func getServersFromPayload(opt json.RawMessage) (servers []Server, err error) {
	list := struct {
		List []Server `json:"list"`
	}{}
	if err = json.Unmarshal(opt, &list); err != nil {
		return
	}

	servers = list.List

	return
}

func main() {
	router := httprouter.New()
	s := stats.New()

	router.POST("/deploy/github.com/:repo_owner/:repo_name", handleGithubDeploy)

	log.Info("Dropship listening on port " + port)
	log.Fatal(http.ListenAndServe(":"+port, s.Handler(router)))
}
<commit_msg>added consul and some error handling<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/ChrisMcKenzie/dropship/ssh"
	"github.com/google/go-github/github"
	"github.com/hashicorp/consul/api"
	"github.com/julienschmidt/httprouter"
	"github.com/the-control-group/data-service-api/logger"
	"github.com/thoas/stats"
)

type (
	Payload struct {
		Deployment github.Deployment `json:"deployment"`
		Repository github.Repository `json:"repository"`
		Sender     github.User       `json:"sender"`
	}

	Server struct {
		Address  string `json:"address"`
		User     string `json:"username"`
		Password string `json:"password"`
	}

	Deployment struct {
		Command string `json:"command"`
		Servers struct {
			Provider string          `json:"provider"`
			Options  json.RawMessage `json:"options"`
		} `json:"servers"`
	}
)

var log = logger.NewLogger()
var port = "3000"

func Logger(h httprouter.Handle) httprouter.Handle {
	return func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		log.Infof("[%s] %s", r.Method, r.URL)
		h(w, r, p)
	}
}

func handleGithubDeploy(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	headers := r.Header

	if headers.Get("X-GitHub-Event") != "deploy" {
		log.Infof("Unable to handle event type %s", headers.Get("X-GitHub-Event"))
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	var err error
	payload := Payload{}
	decoder := json.NewDecoder(r.Body)
	if err := decoder.Decode(&payload); err != nil {
		log.Error(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	deployment := Deployment{}
	if err := json.Unmarshal(payload.Deployment.Payload, &deployment); err != nil {
		log.Error(err)
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	servers := []Server{}
	if deployment.Servers.Provider == "consul" {
		getServersFromConsul(deployment.Servers.Options)
	} else if deployment.Servers.Provider == "payload" {
		servers, err = getServersFromPayload(deployment.Servers.Options)
		if err != nil {
			log.Error(err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
	}

	run(deployment.Command, servers)

	fmt.Fprintf(w, "deploying %s", payload.Repository.Name)
}

func run(c string, servers []Server) {
	for _, server := range servers {
		go func() {
			log.Debugf("Excecuting Command on %s", server.Address)
			res := ssh.Execute(
				c,
				server.Address,
				ssh.NewClientConfig(server.User, server.Password),
			)
			log.Debug(res)
		}()
	}
}

func getServersFromConsul(opt json.RawMessage) (servers []Server, err error) {
	options := struct {
		serviceName string `json:"service_name"`
	}{}
	err = json.Unmarshal(opt, &options)
	if err != nil {
		return
	}

	config := api.DefaultConfig()
	client, err := api.NewClient(config)
	if err != nil {
		return
	}

	catalog := client.Catalog()
	services, _, err := catalog.Service(options.serviceName, "", nil)
	if err != nil {
		return
	}

	for _, service := range services {
		servers = append(
			servers,
			Server{
				Address: service.Address,
			},
		)
	}

	return
}

func getServersFromPayload(opt json.RawMessage) (servers []Server, err error) {
	list := struct {
		List []Server `json:"list"`
	}{}
	if err = json.Unmarshal(opt, &list); err != nil {
		return
	}

	servers = list.List

	return
}

func main() {
	router := httprouter.New()
	s := stats.New()

	router.POST("/deploy/github.com/:repo_owner/:repo_name",
		Logger(handleGithubDeploy))

	log.Info("Dropship listening on port " + port)
	log.Fatal(http.ListenAndServe(":"+port, s.Handler(router)))
}
<|endoftext|>


<commit_before>// Package main initializes a web server.
package main

import (
	"fmt"
	"net/http"
	_ "net/http/pprof" // import for side effects

	"github.com/hack4impact/transcribe4all/web"
)

func main() {
	router := web.NewRouter()
	middlewareRouter := web.ApplyMiddleware(router)
	config, configErr := parseConfigFile("config.toml")
	if configErr == nil {
		// replace this with your actual use of config
		fmt.Printf("%+v\n", *config)

		url, err := uploadFileToBackblaze("testfile4.wav", config.AccountID, config.ApplicationKey, config.BucketName)
		if err != nil {
			// replace this with your actual use of the url of the file
			fmt.Println(url)
		}
	}

	// serve http
	http.Handle("/", middlewareRouter)
	http.Handle("/static/", http.FileServer(http.Dir(".")))
	http.ListenAndServe(":8080", nil)

}
<commit_msg>fix in sample error handling<commit_after>// Package main initializes a web server.
package main

import (
	"fmt"
	"net/http"
	_ "net/http/pprof" // import for side effects

	"github.com/hack4impact/transcribe4all/web"
)

func main() {
	router := web.NewRouter()
	middlewareRouter := web.ApplyMiddleware(router)
	config, configErr := parseConfigFile("config.toml")
	if configErr == nil {
		// replace this with your actual use of config
		fmt.Printf("%+v\n", *config)

		url, err := uploadFileToBackblaze("testfile5R.wav", config.AccountID, config.ApplicationKey, config.BucketName)
		if err != nil {
			// replace this with error handling
			fmt.Println(err)
		} else {
			// replace this with your actual use of the url of the file
			fmt.Println(url)
		}
	}

	// serve http
	http.Handle("/", middlewareRouter)
	http.Handle("/static/", http.FileServer(http.Dir(".")))
	http.ListenAndServe(":8080", nil)

}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"

	"github.com/BurntSushi/toml"
	"github.com/k0pernicus/goyave/configurationFile"
	"github.com/k0pernicus/goyave/consts"
	"github.com/k0pernicus/goyave/traces"
	"github.com/k0pernicus/goyave/utils"
	"github.com/k0pernicus/goyave/walk"
	"github.com/spf13/cobra"
)

var configurationFileStructure configurationFile.ConfigurationFile
var configurationFilePath string
var userHomeDir string

/*initialize get the configuration file existing in the system (or create it), and return
 *a pointer to his content.
 */
func initialize(configurationFileStructure *configurationFile.ConfigurationFile) {
	// Initialize all different traces structures
	traces.InitTraces(os.Stdout, os.Stderr, os.Stdout, os.Stdout)

	// Get the user home directory
	userHomeDir = utils.GetUserHomeDir()
	if len(userHomeDir) == 0 {
		log.Fatalf("Cannot get the user home dir.\n")
	}
	// Set the configuration path file
	configurationFilePath = path.Join(userHomeDir, consts.ConfigurationFileName)
	filePointer, err := os.OpenFile(configurationFilePath, os.O_RDWR|os.O_CREATE, 0755)
	if err != nil {
		log.Fatalf("Cannot open the file %s, due to error '%s'.\n", configurationFilePath, err)
	}
	defer filePointer.Close()
	var bytesArray []byte
	// Get the content of the goyave configuration file
	configurationFile.GetConfigurationFileContent(filePointer, &bytesArray)
	if _, err = toml.Decode(string(bytesArray[:]), configurationFileStructure); err != nil {
		log.Fatalln(err)
	}
}

func main() {

	var rootCmd = &cobra.Command{
		Use:   "goyave",
		Short: "Goyave is a tool to take a look at your local git repositories",
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			initialize(&configurationFileStructure)
		},
		PersistentPostRun: func(cmd *cobra.Command, args []string) {
			var outputBuffer bytes.Buffer
			if err := configurationFileStructure.Encode(&outputBuffer); err != nil {
				log.Fatalln("Cannot save the current configurationFile structure!")
			}
			if err := ioutil.WriteFile(configurationFilePath, outputBuffer.Bytes(), 0777); err != nil {
				log.Fatalln("Cannot access to your file to save the configurationFile structure!")
			}
		},
	}

	/*stateCmd is a subcommand to list the state of each local git repository
	 */
	var stateCmd = &cobra.Command{
		Use:   "state",
		Short: "Get the state of each local git repository",
		Run: func(cmd *cobra.Command, args []string) {
			for _, gitStruct := range configurationFileStructure.VisibleRepositories {
				gitStruct.Init()
				gitStruct.GitObject.Status()
			}
		},
	}

	/*crawlCmd is a subcommand to crawl your hard drive in order to get and save new git repositories
	 */
	var crawlCmd = &cobra.Command{
		Use:   "crawl",
		Short: "Crawl the hard drive in order to find new git repositories",
		Run: func(cmd *cobra.Command, args []string) {
			//		Get all git paths, and display them
			gitPaths, err := walk.RetrieveGitRepositories(userHomeDir)
			if err != nil {
				log.Fatalf("There was an error retrieving your git repositories: '%s'\n", err)
			}
			// For each git repository, check if it exists, and if not add it to the default target visibility
			for _, gitPath := range gitPaths {
				if err := configurationFileStructure.AddRepository(gitPath, configurationFileStructure.Local.DefaultTarget); err != nil {
					traces.WarningTracer.Printf("[%s] %s", gitPath, err)
				}
			}
			//For each VISIBLE repository, get some informations about his state and display it
		},
	}

	rootCmd.AddCommand(crawlCmd, stateCmd)

	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	// // Save all new changes
}
<commit_msg>Code style and documentation<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"

	"github.com/BurntSushi/toml"
	"github.com/k0pernicus/goyave/configurationFile"
	"github.com/k0pernicus/goyave/consts"
	"github.com/k0pernicus/goyave/traces"
	"github.com/k0pernicus/goyave/utils"
	"github.com/k0pernicus/goyave/walk"
	"github.com/spf13/cobra"
)

var configurationFileStructure configurationFile.ConfigurationFile
var configurationFilePath string
var userHomeDir string

/*initialize get the configuration file existing in the system (or create it), and return
 *a pointer to his content.
 */
func initialize(configurationFileStructure *configurationFile.ConfigurationFile) {
	// Initialize all different traces structures
	traces.InitTraces(os.Stdout, os.Stderr, os.Stdout, os.Stdout)

	// Get the user home directory
	userHomeDir = utils.GetUserHomeDir()
	if len(userHomeDir) == 0 {
		log.Fatalf("Cannot get the user home dir.\n")
	}
	// Set the configuration path file
	configurationFilePath = path.Join(userHomeDir, consts.ConfigurationFileName)
	filePointer, err := os.OpenFile(configurationFilePath, os.O_RDWR|os.O_CREATE, 0755)
	if err != nil {
		log.Fatalf("Cannot open the file %s, due to error '%s'.\n", configurationFilePath, err)
	}
	defer filePointer.Close()
	var bytesArray []byte
	// Get the content of the goyave configuration file
	configurationFile.GetConfigurationFileContent(filePointer, &bytesArray)
	if _, err = toml.Decode(string(bytesArray[:]), configurationFileStructure); err != nil {
		log.Fatalln(err)
	}
}

func main() {

	/*rootCmd defines the global app, and some actions to run before and after the command running
	 */
	var rootCmd = &cobra.Command{
		Use:   "goyave",
		Short: "Goyave is a tool to take a look at your local git repositories",
		// Initialize the structure
		PersistentPreRun: func(cmd *cobra.Command, args []string) {
			initialize(&configurationFileStructure)
		},
		// Save the current configuration file structure, in the configuration file
		PersistentPostRun: func(cmd *cobra.Command, args []string) {
			var outputBuffer bytes.Buffer
			if err := configurationFileStructure.Encode(&outputBuffer); err != nil {
				log.Fatalln("Cannot save the current configurationFile structure!")
			}
			if err := ioutil.WriteFile(configurationFilePath, outputBuffer.Bytes(), 0777); err != nil {
				log.Fatalln("Cannot access to your file to save the configurationFile structure!")
			}
		},
	}

	/*crawlCmd is a subcommand to crawl your hard drive in order to get and save new git repositories
	 */
	var crawlCmd = &cobra.Command{
		Use:   "crawl",
		Short: "Crawl the hard drive in order to find git repositories",
		Run: func(cmd *cobra.Command, args []string) {
			//		Get all git paths, and display them
			gitPaths, err := walk.RetrieveGitRepositories(userHomeDir)
			if err != nil {
				log.Fatalf("There was an error retrieving your git repositories: '%s'\n", err)
			}
			// For each git repository, check if it exists, and if not add it to the default target visibility
			for _, gitPath := range gitPaths {
				if err := configurationFileStructure.AddRepository(gitPath, configurationFileStructure.Local.DefaultTarget); err != nil {
					traces.WarningTracer.Printf("[%s] %s", gitPath, err)
				}
			}
			//For each VISIBLE repository, get some informations about his state and display it
		},
	}

	/*stateCmd is a subcommand to list the state of each local git repository
	 */
	var stateCmd = &cobra.Command{
		Use:   "state",
		Short: "Get the state of each local git repository",
		Run: func(cmd *cobra.Command, args []string) {
			for _, gitStruct := range configurationFileStructure.VisibleRepositories {
				gitStruct.Init()
				gitStruct.GitObject.Status()
			}
		},
	}

	rootCmd.AddCommand(crawlCmd, stateCmd)

	if err := rootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}
<|endoftext|>