<commit_before>// Copyright 2013 Andreas Koch. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"github.com/andreaskoch/allmark2/common/config"
	"github.com/andreaskoch/allmark2/common/index"
	"github.com/andreaskoch/allmark2/common/logger/console"
	"github.com/andreaskoch/allmark2/common/logger/loglevel"
	"github.com/andreaskoch/allmark2/common/util/fsutil"
	"github.com/andreaskoch/allmark2/dataaccess"
	"github.com/andreaskoch/allmark2/dataaccess/filesystem"
	"github.com/andreaskoch/allmark2/services/conversion/markdowntohtml"
	"github.com/andreaskoch/allmark2/services/initialization"
	"github.com/andreaskoch/allmark2/services/parsing"
	"github.com/andreaskoch/allmark2/services/search"
	"github.com/andreaskoch/allmark2/ui/web/server"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const (
	CommandNameInit  = "init"
	CommandNameServe = "serve"
)

func main() {

	serve := func(repositoryPath string) {

		// logger
		logger := console.New(loglevel.Debug)

		// config
		config := config.Get(repositoryPath)

		// data access
		repository, err := filesystem.NewRepository(logger, repositoryPath)
		if err != nil {
			logger.Fatal("Unable to create a repository. Error: %s", err)
		}

		// parser
		parser, err := parsing.New(logger)
		if err != nil {
			logger.Fatal("Unable to instantiate a parser. Error: %s", err)
		}

		// converter
		converter, err := markdowntohtml.New(logger)
		if err != nil {
			logger.Fatal("Unable to instantiate a converter. Error: %s", err)
		}

		// repository (fallback) name
		repositoryName := filepath.Base(repositoryPath)

		// item index
		index := index.New(logger, repositoryName)

		// search
		itemSearch := search.NewItemSearch(logger, index)

		addRepositoryItemToIndex := func(repositoryItem *dataaccess.Item) {
			// parse item
			if item, err := parser.Parse(repositoryItem); err == nil {
				// register the item at the index
				index.Add(item)
			} else {
				logger.Warn("%s", err.Error())
			}
		}

		removeRepositoryItemFromIndex := func(repositoryItem *dataaccess.Item) {
			// remove the item from the index
			index.Remove(repositoryItem.Route())
		}

		// full index of the repository
		go func() {
			for itemEvent := range repository.InitialItems() {

				// validate event
				if itemEvent.Error != nil {
					logger.Warn("%s", itemEvent.Error)
					continue
				}

				repositoryItem := itemEvent.Item
				if repositoryItem == nil {
					logger.Warn("Repository item is empty.")
					continue
				}

				addRepositoryItemToIndex(repositoryItem)

			}

			// update the full-text search index
			itemSearch.Update()
		}()

		// scheduled reindex of the fulltext index
		go func() {
			sleepInterval := time.Minute * 3
			for {
				time.Sleep(sleepInterval)
				itemSearch.Update()
			}
		}()

		// todo: implement discard of old items and especially old change listeners
		// event handler: new items
		go func() {
			for itemEvent := range repository.NewItems() {
				logger.Debug("New Item")

				// validate event
				if itemEvent.Error != nil {
					logger.Warn("%s", itemEvent.Error)
					continue
				}

				repositoryItem := itemEvent.Item
				if repositoryItem == nil {
					logger.Warn("Repository item is empty.")
					continue
				}

				addRepositoryItemToIndex(repositoryItem)
			}
		}()

		// event handler: changed items
		go func() {
			for itemEvent := range repository.ChangedItems() {
				logger.Debug("Item changed")

				// validate event
				if itemEvent.Error != nil {
					logger.Warn("%s", itemEvent.Error)
					continue
				}

				repositoryItem := itemEvent.Item
				if repositoryItem == nil {
					logger.Warn("Repository item is empty.")
					continue
				}

				addRepositoryItemToIndex(repositoryItem)
			}
		}()

		// event handler: moved items
		go func() {
			for itemEvent := range repository.MovedItems() {
				logger.Debug("Item moved")

				// validate event
				if itemEvent.Error != nil {
					logger.Warn("%s", itemEvent.Error)
					continue
				}

				repositoryItem := itemEvent.Item
				if repositoryItem == nil {
					logger.Warn("Repository item is empty.")
					continue
				}

				removeRepositoryItemFromIndex(repositoryItem)
			}
		}()

		// server
		server, err := server.New(logger, config, index, converter, itemSearch)
		if err != nil {
			logger.Fatal("Unable to instantiate a server. Error: %s", err.Error())
		}

		if result := <-server.Start(); result != nil {
			logger.Info("%s", result)
		}
	}

	init := func(repositoryPath string) {
		if success, err := initialization.Initialize(repositoryPath); !success {
			fmt.Println(err)
		}
	}

	parseCommandLineArguments(os.Args, func(commandName, repositoryPath string) (commandWasFound bool) {
		switch strings.ToLower(commandName) {
		case CommandNameInit:
			init(repositoryPath)

		case CommandNameServe:
			serve(repositoryPath)

		default:
			return false
		}

		return true
	})
}

func parseCommandLineArguments(args []string, commandHandler func(commandName, repositoryPath string) (commandWasFound bool)) {

	// check if the mandatory amount of
	// command line parameters has been
	// supplied. If not, print usage information.
	if len(args) < 2 {
		printUsageInformation(args)
		return
	}

	// Read the repository path parameters
	var repositoryPath string
	if len(args) > 2 {

		// use supplied repository path
		repositoryPath = args[2]

		if isFile, _ := fsutil.IsFile(repositoryPath); isFile {
			repositoryPath = filepath.Dir(repositoryPath)
		}

	} else {

		// use the current directory
		repositoryPath = fsutil.GetWorkingDirectory()

	}

	// validate the supplied repository paths
	if !fsutil.PathExists(repositoryPath) {
		fmt.Fprintf(os.Stderr, "The specified repository paths %q is does not exist.", repositoryPath)
		return
	}

	// Read the command parameter and execute the command handler
	commandName := strings.ToLower(args[1])
	if commandWasFound := commandHandler(commandName, repositoryPath); !commandWasFound {
		printUsageInformation(args)
	}
}

// Print usage information
func printUsageInformation(args []string) {
	executeableName := args[0]

	fmt.Fprintf(os.Stderr, "%s - %s\n", executeableName, "A markdown web server and renderer")
	fmt.Fprintf(os.Stderr, "\nUsage:\n%s %s %s\n", executeableName, "<command>", "<repository path>")
	fmt.Fprintf(os.Stderr, "\nAvailable commands:\n")
	fmt.Fprintf(os.Stderr, "  %7s  %s\n", CommandNameInit, "Initialize the configuration")
	fmt.Fprintf(os.Stderr, "  %7s  %s\n", CommandNameServe, "Start serving the supplied repository via HTTP")
	fmt.Fprintf(os.Stderr, "\n")
	fmt.Fprintf(os.Stderr, "Fork me on GitHub %q\n", "https://github.com/andreaskoch/allmark")

	os.Exit(2)
}
<commit_msg>Cleaned up the main func<commit_after>// Copyright 2013 Andreas Koch. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"fmt"
	"github.com/andreaskoch/allmark2/common/config"
	"github.com/andreaskoch/allmark2/common/index"
	"github.com/andreaskoch/allmark2/common/logger/console"
	"github.com/andreaskoch/allmark2/common/logger/loglevel"
	"github.com/andreaskoch/allmark2/common/util/fsutil"
	"github.com/andreaskoch/allmark2/dataaccess"
	"github.com/andreaskoch/allmark2/dataaccess/filesystem"
	"github.com/andreaskoch/allmark2/services/conversion/markdowntohtml"
	"github.com/andreaskoch/allmark2/services/initialization"
	"github.com/andreaskoch/allmark2/services/parsing"
	"github.com/andreaskoch/allmark2/services/search"
	"github.com/andreaskoch/allmark2/ui/web/server"
	"os"
	"path/filepath"
	"strings"
	"time"
)

var (
	logger = console.New(loglevel.Debug)
)

const (
	CommandNameInit  = "init"
	CommandNameServe = "serve"
)

func main() {
	parseCommandLineArguments(os.Args, func(commandName, repositoryPath string) (commandWasFound bool) {
		switch strings.ToLower(commandName) {
		case CommandNameInit:
			return initialize(repositoryPath)

		case CommandNameServe:
			return serve(repositoryPath)

		default:
			return false
		}

		panic("Unreachable")
	})
}

func parseCommandLineArguments(args []string, commandHandler func(commandName, repositoryPath string) (commandWasFound bool)) {

	// check if the mandatory amount of
	// command line parameters has been
	// supplied. If not, print usage information.
	if len(args) < 2 {
		printUsageInformation(args)
		return
	}

	// Read the repository path parameters
	var repositoryPath string
	if len(args) > 2 {

		// use supplied repository path
		repositoryPath = args[2]

		if isFile, _ := fsutil.IsFile(repositoryPath); isFile {
			repositoryPath = filepath.Dir(repositoryPath)
		}

	} else {

		// use the current directory
		repositoryPath = fsutil.GetWorkingDirectory()

	}

	// validate the supplied repository paths
	if !fsutil.PathExists(repositoryPath) {
		fmt.Fprintf(os.Stderr, "The specified repository paths %q is does not exist.", repositoryPath)
		return
	}

	// Read the command parameter and execute the command handler
	commandName := strings.ToLower(args[1])
	if commandWasFound := commandHandler(commandName, repositoryPath); !commandWasFound {
		printUsageInformation(args)
	}
}

// Print usage information
func printUsageInformation(args []string) {
	executeableName := args[0]

	fmt.Fprintf(os.Stderr, "%s - %s\n", executeableName, "A markdown web server and renderer")
	fmt.Fprintf(os.Stderr, "\nUsage:\n%s %s %s\n", executeableName, "<command>", "<repository path>")
	fmt.Fprintf(os.Stderr, "\nAvailable commands:\n")
	fmt.Fprintf(os.Stderr, "  %7s  %s\n", CommandNameInit, "Initialize the configuration")
	fmt.Fprintf(os.Stderr, "  %7s  %s\n", CommandNameServe, "Start serving the supplied repository via HTTP")
	fmt.Fprintf(os.Stderr, "\n")
	fmt.Fprintf(os.Stderr, "Fork me on GitHub %q\n", "https://github.com/andreaskoch/allmark")

	os.Exit(2)
}

func serve(repositoryPath string) bool {

	// config
	config := config.Get(repositoryPath)

	// data access
	repository, err := filesystem.NewRepository(logger, repositoryPath)
	if err != nil {
		logger.Error("Unable to create a repository. Error: %s", err)
	}

	// parser
	parser, err := parsing.New(logger)
	if err != nil {
		logger.Error("Unable to instantiate a parser. Error: %s", err)
	}

	// converter
	converter, err := markdowntohtml.New(logger)
	if err != nil {
		logger.Error("Unable to instantiate a converter. Error: %s", err)
	}

	// repository (fallback) name
	repositoryName := filepath.Base(repositoryPath)

	// item index
	index := index.New(logger, repositoryName)

	// search
	itemSearch := search.NewItemSearch(logger, index)

	addRepositoryItemToIndex := func(repositoryItem *dataaccess.Item) {
		// parse item
		if item, err := parser.Parse(repositoryItem); err == nil {
			// register the item at the index
			index.Add(item)
		} else {
			logger.Warn("%s", err.Error())
		}
	}

	removeRepositoryItemFromIndex := func(repositoryItem *dataaccess.Item) {
		// remove the item from the index
		index.Remove(repositoryItem.Route())
	}

	// full index of the repository
	go func() {
		for itemEvent := range repository.InitialItems() {

			// validate event
			if itemEvent.Error != nil {
				logger.Warn("%s", itemEvent.Error)
				continue
			}

			repositoryItem := itemEvent.Item
			if repositoryItem == nil {
				logger.Warn("Repository item is empty.")
				continue
			}

			addRepositoryItemToIndex(repositoryItem)

		}

		// update the full-text search index
		itemSearch.Update()
	}()

	// scheduled reindex of the fulltext index
	go func() {
		sleepInterval := time.Minute * 3
		for {
			time.Sleep(sleepInterval)
			itemSearch.Update()
		}
	}()

	// todo: implement discard of old items and especially old change listeners
	// event handler: new items
	go func() {
		for itemEvent := range repository.NewItems() {
			logger.Debug("New Item")

			// validate event
			if itemEvent.Error != nil {
				logger.Warn("%s", itemEvent.Error)
				continue
			}

			repositoryItem := itemEvent.Item
			if repositoryItem == nil {
				logger.Warn("Repository item is empty.")
				continue
			}

			addRepositoryItemToIndex(repositoryItem)
		}
	}()

	// event handler: changed items
	go func() {
		for itemEvent := range repository.ChangedItems() {
			logger.Debug("Item changed")

			// validate event
			if itemEvent.Error != nil {
				logger.Warn("%s", itemEvent.Error)
				continue
			}

			repositoryItem := itemEvent.Item
			if repositoryItem == nil {
				logger.Warn("Repository item is empty.")
				continue
			}

			addRepositoryItemToIndex(repositoryItem)
		}
	}()

	// event handler: moved items
	go func() {
		for itemEvent := range repository.MovedItems() {
			logger.Debug("Item moved")

			// validate event
			if itemEvent.Error != nil {
				logger.Warn("%s", itemEvent.Error)
				continue
			}

			repositoryItem := itemEvent.Item
			if repositoryItem == nil {
				logger.Warn("Repository item is empty.")
				continue
			}

			removeRepositoryItemFromIndex(repositoryItem)
		}
	}()

	// server
	server, err := server.New(logger, config, index, converter, itemSearch)
	if err != nil {
		logger.Error("Unable to instantiate a server. Error: %s", err.Error())
		return false
	}

	if result := <-server.Start(); result != nil {
		logger.Error("%s", result)
		return false
	}

	return true
}

func initialize(repositoryPath string) bool {
	if success, err := initialization.Initialize(repositoryPath); !success {
		logger.Error("Error initializing folder %q. Error: %s", repositoryPath, err.Error())
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package main

import (
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/docker/docker/cli/config"
)

// will be filled in at compile time
var gitCommit = ""

const (
	version = "0.9.0-dev"
	usage   = "inspect and push manifest list images to a registry"
)

func main() {
	app := cli.NewApp()
	app.Name = os.Args[0]
	app.Version = version + " (commit: " + gitCommit + ")"
	app.Usage = usage
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug output",
		},
		cli.StringFlag{
			Name:  "username",
			Value: "",
			Usage: "registry username",
		},
		cli.StringFlag{
			Name:  "password",
			Value: "",
			Usage: "registry password",
		},
		cli.StringFlag{
			Name:  "docker-cfg",
			Value: config.Dir(),
			Usage: "Docker's cli config for auth",
		},
	}
	app.Before = func(c *cli.Context) error {
		if c.GlobalBool("debug") {
			logrus.SetLevel(logrus.DebugLevel)
		} else {
			logrus.SetLevel(logrus.WarnLevel)
		}
		return nil
	}
	// currently support inspect and pushml
	app.Commands = []cli.Command{
		inspectCmd,
		pushCmd,
	}

	if err := app.Run(os.Args); err != nil {
		logrus.Fatal(err)
	}
}
<commit_msg>Prepare v0.9.0 release<commit_after>package main

import (
	"os"

	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/docker/docker/cli/config"
)

// will be filled in at compile time
var gitCommit = ""

const (
	version = "0.9.0"
	usage   = "inspect and push manifest list images to a registry"
)

func main() {
	app := cli.NewApp()
	app.Name = os.Args[0]
	app.Version = version + " (commit: " + gitCommit + ")"
	app.Usage = usage
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "debug",
			Usage: "enable debug output",
		},
		cli.StringFlag{
			Name:  "username",
			Value: "",
			Usage: "registry username",
		},
		cli.StringFlag{
			Name:  "password",
			Value: "",
			Usage: "registry password",
		},
		cli.StringFlag{
			Name:  "docker-cfg",
			Value: config.Dir(),
			Usage: "Docker's cli config for auth",
		},
	}
	app.Before = func(c *cli.Context) error {
		if c.GlobalBool("debug") {
			logrus.SetLevel(logrus.DebugLevel)
		} else {
			logrus.SetLevel(logrus.WarnLevel)
		}
		return nil
	}
	// currently support inspect and pushml
	app.Commands = []cli.Command{
		inspectCmd,
		pushCmd,
	}

	if err := app.Run(os.Args); err != nil {
		logrus.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package ebpf

import (
	"bytes"
	"fmt"
	"math"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"github.com/newtools/ebpf/asm"

	"github.com/pkg/errors"
)

// Errors returned by the implementation
var (
	ErrNotSupported = errors.New("ebpf: not supported by kernel")
)

const (
	// Number of bytes to pad the output buffer for BPF_PROG_TEST_RUN.
	// This is currently the maximum of spare space allocated for SKB
	// and XDP programs, and equal to XDP_PACKET_HEADROOM + NET_IP_ALIGN.
	outputPad = 256 + 2
)

// DefaultVerifierLogSize is the default number of bytes allocated for the
// verifier log.
const DefaultVerifierLogSize = 64 * 1024

// ProgramSpec defines a Program
type ProgramSpec struct {
	// Name is passed to the kernel as a debug aid. Must only contain
	// alpha numeric and '_' characters and be less than 16 characters.
	Name          string
	Type          ProgType
	Instructions  asm.Instructions
	License       string
	KernelVersion uint32
}

// Program represents a Program file descriptor
type Program struct {
	fd   uint32
	name string
	abi  ProgramABI
}

var (
	detectHaveProgName sync.Once
	haveProgName       bool
)

// NewProgram creates a new Program.
//
// Loading a program for the first time will perform
// feature detection by loading small, temporary programs.
func NewProgram(spec *ProgramSpec) (*Program, error) {
	if len(spec.Instructions) == 0 {
		return nil, errors.Errorf("instructions cannot be empty")
	}

	detectHaveProgName.Do(func() {
		attr, err := convertProgramSpec(&ProgramSpec{
			Name: "feature_test",
			Type: SocketFilter,
			Instructions: asm.Instructions{
				asm.LoadImm(asm.R0, 0, asm.DWord),
				asm.Return(),
			},
			License: "MIT",
		}, true)
		if err != nil {
			return
		}

		fd, err := progLoad(attr)
		if err != nil {
			// This may be because we lack sufficient permissions, etc.
			return
		}

		syscall.Close(int(fd))
		haveProgName = true
	})

	attr, err := convertProgramSpec(spec, haveProgName)
	if err != nil {
		return nil, err
	}

	fd, err := progLoad(attr)
	if err == nil {
		prog := &Program{
			uint32(fd),
			spec.Name,
			ProgramABI{spec.Type},
		}
		runtime.SetFinalizer(prog, (*Program).Close)
		return prog, nil
	}

	// Something went wrong, re-run with the verifier enabled.
	logs := make([]byte, DefaultVerifierLogSize)
	attr.logLevel = 1
	attr.logSize = uint32(len(logs))
	attr.logBuf = newPtr(unsafe.Pointer(&logs[0]))

	_, nerr := progLoad(attr)
	if errors.Cause(nerr) == syscall.ENOSPC {
		return nil, errors.Wrap(err, "no debug since LogSize too small")
	}
	return nil, &loadError{nerr, convertCString(logs)}
}

func convertProgramSpec(spec *ProgramSpec, includeName bool) (*progLoadAttr, error) {
	buf := bytes.NewBuffer(make([]byte, 0, len(spec.Instructions)*asm.InstructionSize))
	err := spec.Instructions.Marshal(buf, nativeEndian)
	if err != nil {
		return nil, err
	}

	bytecode := buf.Bytes()
	insCount := uint32(len(bytecode) / asm.InstructionSize)
	lic := []byte(spec.License)
	attr := &progLoadAttr{
		progType:     spec.Type,
		insCount:     insCount,
		instructions: newPtr(unsafe.Pointer(&bytecode[0])),
		license:      newPtr(unsafe.Pointer(&lic[0])),
	}

	if err := checkName(spec.Name); err != nil {
		return nil, err
	}

	if includeName {
		copy(attr.progName[:bpfObjNameLen-1], spec.Name)
	}

	return attr, nil
}

func (bpf *Program) String() string {
	if bpf.name != "" {
		return fmt.Sprintf("%s(%s)#%d", bpf.abi.Type, bpf.name, bpf.fd)
	}
	return fmt.Sprintf("%s#%d", bpf.abi.Type, bpf.fd)
}

// FD gets the file descriptor value of the Program
func (bpf *Program) FD() int {
	return int(bpf.fd)
}

// Pin persists the Program past the lifetime of the process that created it
//
// This requires bpffs to be mounted above fileName. See http://cilium.readthedocs.io/en/doc-1.0/kubernetes/install/#mounting-the-bpf-fs-optional
func (bpf *Program) Pin(fileName string) error {
	return pinObject(fileName, bpf.fd)
}

// Close unloads the program from the kernel.
func (bpf *Program) Close() error {
	runtime.SetFinalizer(bpf, nil)
	return syscall.Close(int(bpf.fd))
}

// Test runs the Program in the kernel with the given input and returns the
// value returned by the eBPF program. outLen may be zero.
//
// Note: the kernel expects at least 14 bytes input for an ethernet header for
// XDP and SKB programs.
//
// This function requires at least Linux 4.12.
func (bpf *Program) Test(in []byte) (uint32, []byte, error) {
	ret, out, _, err := bpf.testRun(in, 1)
	return ret, out, err
}

// Benchmark runs the Program with the given input for a number of times
// and returns the time taken per iteration.
//
// The returned value is the return value of the last execution of
// the program.
//
// This function requires at least Linux 4.12.
func (bpf *Program) Benchmark(in []byte, repeat int) (uint32, time.Duration, error) {
	ret, _, total, err := bpf.testRun(in, repeat)
	return ret, total, err
}

var noProgTestRun bool
var detectProgTestRun sync.Once

func (bpf *Program) testRun(in []byte, repeat int) (uint32, []byte, time.Duration, error) {
	if uint(repeat) > math.MaxUint32 {
		return 0, nil, 0, fmt.Errorf("repeat is too high")
	}

	if len(in) == 0 {
		return 0, nil, 0, fmt.Errorf("missing input")
	}

	if uint(len(in)) > math.MaxUint32 {
		return 0, nil, 0, fmt.Errorf("input is too long")
	}

	detectProgTestRun.Do(func() {
		prog, err := NewProgram(&ProgramSpec{
			Type: SocketFilter,
			Instructions: asm.Instructions{
				asm.LoadImm(asm.R0, 0, asm.DWord),
				asm.Return(),
			},
			License: "MIT",
		})
		if err != nil {
			// This may be because we lack sufficient permissions, etc.
			return
		}
		defer prog.Close()

		// XDP progs require at least 14 bytes input
		in := make([]byte, 14)
		attr := progTestRunAttr{
			fd:         prog.fd,
			dataSizeIn: uint32(len(in)),
			dataIn:     newPtr(unsafe.Pointer(&in[0])),
		}
		_, err = bpfCall(_ProgTestRun, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
		noProgTestRun = errors.Cause(err) == syscall.EINVAL
	})

	if noProgTestRun {
		return 0, nil, 0, ErrNotSupported
	}

	// There is currently no way to tell the kernel about the size of the output buffer.
	// Combined with things like bpf_xdp_adjust_head() we don't really know what the final
	// size will be. Hence we allocate an output buffer which we hope will always be large
	// enough, and panic if the kernel wrote past the end of the allocation.
	// See https://marc.info/?l=linux-netdev&m=152283265832434&w=2
	out := make([]byte, len(in)+outputPad)

	attr := progTestRunAttr{
		fd:         bpf.fd,
		dataSizeIn: uint32(len(in)),
		// NB: dataSizeOut is not read by the kernel
		dataIn:  newPtr(unsafe.Pointer(&in[0])),
		dataOut: newPtr(unsafe.Pointer(&out[0])),
		repeat:  uint32(repeat),
	}

	_, err := bpfCall(_ProgTestRun, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if err != nil {
		return 0, nil, 0, errors.Wrap(err, "can't run test")
	}

	if int(attr.dataSizeOut) > cap(out) {
		// Houston, we have a problem. The program created more data than we allocated,
		// and the kernel wrote past the end of our buffer.
		panic("kernel wrote past end of output buffer")
	}
	out = out[:int(attr.dataSizeOut)]

	total := time.Duration(attr.duration) * time.Nanosecond
	return attr.retval, out, total, nil
}

// LoadPinnedProgram loads a Program from a BPF file.
//
// Requires at least Linux 4.13, use LoadPinnedProgramExplicit on
// earlier versions.
func LoadPinnedProgram(fileName string) (*Program, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	abi, err := newProgramABIFromFd(fd)
	if err != nil {
		return nil, err
	}
	return &Program{
		uint32(fd),
		filepath.Base(fileName),
		*abi,
	}, nil
}

// LoadPinnedProgramExplicit loads a program with explicit parameters.
func LoadPinnedProgramExplicit(fileName string, abi *ProgramABI) (*Program, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	return &Program{
		uint32(fd),
		filepath.Base(fileName),
		*abi,
	}, nil
}

type loadError struct {
	cause       error
	verifierLog string
}

func (le *loadError) Error() string {
	if le.verifierLog == "" {
		return fmt.Sprintf("failed to load program: %s", le.cause)
	}
	return fmt.Sprintf("failed to load program: %s: %s", le.cause, le.verifierLog)
}

func (le *loadError) Cause() error {
	return le.cause
}

func convertCString(in []byte) string {
	inLen := bytes.IndexByte(in, 0)
	return string(in[:inLen])
}

func checkName(name string) error {
	if len(name) > bpfObjNameLen-1 {
		return errors.Errorf("name '%s' is too long", name)
	}

	idx := strings.IndexFunc(name, func(char rune) bool {
		switch {
		case char >= 'A' && char <= 'Z':
			fallthrough
		case char >= 'a' && char <= 'z':
			fallthrough
		case char >= '0' && char <= '9':
			fallthrough
		case char == '_':
			return false
		default:
			return true
		}
	})

	if idx != -1 {
		return errors.Errorf("invalid character '%c' in name '%s'", name[idx], name)
	}

	return nil
}
<commit_msg>Pass output buffer size to BPF_PROG_TEST_RUN<commit_after>package ebpf

import (
	"bytes"
	"fmt"
	"math"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"
	"unsafe"

	"github.com/newtools/ebpf/asm"

	"github.com/pkg/errors"
)

// Errors returned by the implementation
var (
	ErrNotSupported = errors.New("ebpf: not supported by kernel")
)

const (
	// Number of bytes to pad the output buffer for BPF_PROG_TEST_RUN.
	// This is currently the maximum of spare space allocated for SKB
	// and XDP programs, and equal to XDP_PACKET_HEADROOM + NET_IP_ALIGN.
	outputPad = 256 + 2
)

// DefaultVerifierLogSize is the default number of bytes allocated for the
// verifier log.
const DefaultVerifierLogSize = 64 * 1024

// ProgramSpec defines a Program
type ProgramSpec struct {
	// Name is passed to the kernel as a debug aid. Must only contain
	// alpha numeric and '_' characters and be less than 16 characters.
	Name          string
	Type          ProgType
	Instructions  asm.Instructions
	License       string
	KernelVersion uint32
}

// Program represents a Program file descriptor
type Program struct {
	fd   uint32
	name string
	abi  ProgramABI
}

var (
	detectHaveProgName sync.Once
	haveProgName       bool
)

// NewProgram creates a new Program.
//
// Loading a program for the first time will perform
// feature detection by loading small, temporary programs.
func NewProgram(spec *ProgramSpec) (*Program, error) {
	if len(spec.Instructions) == 0 {
		return nil, errors.Errorf("instructions cannot be empty")
	}

	detectHaveProgName.Do(func() {
		attr, err := convertProgramSpec(&ProgramSpec{
			Name: "feature_test",
			Type: SocketFilter,
			Instructions: asm.Instructions{
				asm.LoadImm(asm.R0, 0, asm.DWord),
				asm.Return(),
			},
			License: "MIT",
		}, true)
		if err != nil {
			return
		}

		fd, err := progLoad(attr)
		if err != nil {
			// This may be because we lack sufficient permissions, etc.
			return
		}

		syscall.Close(int(fd))
		haveProgName = true
	})

	attr, err := convertProgramSpec(spec, haveProgName)
	if err != nil {
		return nil, err
	}

	fd, err := progLoad(attr)
	if err == nil {
		prog := &Program{
			uint32(fd),
			spec.Name,
			ProgramABI{spec.Type},
		}
		runtime.SetFinalizer(prog, (*Program).Close)
		return prog, nil
	}

	// Something went wrong, re-run with the verifier enabled.
	logs := make([]byte, DefaultVerifierLogSize)
	attr.logLevel = 1
	attr.logSize = uint32(len(logs))
	attr.logBuf = newPtr(unsafe.Pointer(&logs[0]))

	_, nerr := progLoad(attr)
	if errors.Cause(nerr) == syscall.ENOSPC {
		return nil, errors.Wrap(err, "no debug since LogSize too small")
	}
	return nil, &loadError{nerr, convertCString(logs)}
}

func convertProgramSpec(spec *ProgramSpec, includeName bool) (*progLoadAttr, error) {
	buf := bytes.NewBuffer(make([]byte, 0, len(spec.Instructions)*asm.InstructionSize))
	err := spec.Instructions.Marshal(buf, nativeEndian)
	if err != nil {
		return nil, err
	}

	bytecode := buf.Bytes()
	insCount := uint32(len(bytecode) / asm.InstructionSize)
	lic := []byte(spec.License)
	attr := &progLoadAttr{
		progType:     spec.Type,
		insCount:     insCount,
		instructions: newPtr(unsafe.Pointer(&bytecode[0])),
		license:      newPtr(unsafe.Pointer(&lic[0])),
	}

	if err := checkName(spec.Name); err != nil {
		return nil, err
	}

	if includeName {
		copy(attr.progName[:bpfObjNameLen-1], spec.Name)
	}

	return attr, nil
}

func (bpf *Program) String() string {
	if bpf.name != "" {
		return fmt.Sprintf("%s(%s)#%d", bpf.abi.Type, bpf.name, bpf.fd)
	}
	return fmt.Sprintf("%s#%d", bpf.abi.Type, bpf.fd)
}

// FD gets the file descriptor value of the Program
func (bpf *Program) FD() int {
	return int(bpf.fd)
}

// Pin persists the Program past the lifetime of the process that created it
//
// This requires bpffs to be mounted above fileName. See http://cilium.readthedocs.io/en/doc-1.0/kubernetes/install/#mounting-the-bpf-fs-optional
func (bpf *Program) Pin(fileName string) error {
	return pinObject(fileName, bpf.fd)
}

// Close unloads the program from the kernel.
func (bpf *Program) Close() error {
	runtime.SetFinalizer(bpf, nil)
	return syscall.Close(int(bpf.fd))
}

// Test runs the Program in the kernel with the given input and returns the
// value returned by the eBPF program. outLen may be zero.
//
// Note: the kernel expects at least 14 bytes input for an ethernet header for
// XDP and SKB programs.
//
// This function requires at least Linux 4.12.
func (bpf *Program) Test(in []byte) (uint32, []byte, error) {
	ret, out, _, err := bpf.testRun(in, 1)
	return ret, out, err
}

// Benchmark runs the Program with the given input for a number of times
// and returns the time taken per iteration.
//
// The returned value is the return value of the last execution of
// the program.
//
// This function requires at least Linux 4.12.
func (bpf *Program) Benchmark(in []byte, repeat int) (uint32, time.Duration, error) {
	ret, _, total, err := bpf.testRun(in, repeat)
	return ret, total, err
}

var noProgTestRun bool
var detectProgTestRun sync.Once

func (bpf *Program) testRun(in []byte, repeat int) (uint32, []byte, time.Duration, error) {
	if uint(repeat) > math.MaxUint32 {
		return 0, nil, 0, fmt.Errorf("repeat is too high")
	}

	if len(in) == 0 {
		return 0, nil, 0, fmt.Errorf("missing input")
	}

	if uint(len(in)) > math.MaxUint32 {
		return 0, nil, 0, fmt.Errorf("input is too long")
	}

	detectProgTestRun.Do(func() {
		prog, err := NewProgram(&ProgramSpec{
			Type: SocketFilter,
			Instructions: asm.Instructions{
				asm.LoadImm(asm.R0, 0, asm.DWord),
				asm.Return(),
			},
			License: "MIT",
		})
		if err != nil {
			// This may be because we lack sufficient permissions, etc.
			return
		}
		defer prog.Close()

		// XDP progs require at least 14 bytes input
		in := make([]byte, 14)
		attr := progTestRunAttr{
			fd:         prog.fd,
			dataSizeIn: uint32(len(in)),
			dataIn:     newPtr(unsafe.Pointer(&in[0])),
		}
		_, err = bpfCall(_ProgTestRun, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
		noProgTestRun = errors.Cause(err) == syscall.EINVAL
	})

	if noProgTestRun {
		return 0, nil, 0, ErrNotSupported
	}

	// Older kernels ignore the dataSizeOut argument when copying to user space.
	// Combined with things like bpf_xdp_adjust_head() we don't really know what the final
	// size will be. Hence we allocate an output buffer which we hope will always be large
	// enough, and panic if the kernel wrote past the end of the allocation.
	// See https://patchwork.ozlabs.org/cover/1006822/
	out := make([]byte, len(in)+outputPad)

	attr := progTestRunAttr{
		fd:          bpf.fd,
		dataSizeIn:  uint32(len(in)),
		dataSizeOut: uint32(len(out)),
		dataIn:      newPtr(unsafe.Pointer(&in[0])),
		dataOut:     newPtr(unsafe.Pointer(&out[0])),
		repeat:      uint32(repeat),
	}

	_, err := bpfCall(_ProgTestRun, unsafe.Pointer(&attr), unsafe.Sizeof(attr))
	if err != nil {
		return 0, nil, 0, errors.Wrap(err, "can't run test")
	}

	if int(attr.dataSizeOut) > cap(out) {
		// Houston, we have a problem. The program created more data than we allocated,
		// and the kernel wrote past the end of our buffer.
		panic("kernel wrote past end of output buffer")
	}
	out = out[:int(attr.dataSizeOut)]

	total := time.Duration(attr.duration) * time.Nanosecond
	return attr.retval, out, total, nil
}

// LoadPinnedProgram loads a Program from a BPF file.
//
// Requires at least Linux 4.13, use LoadPinnedProgramExplicit on
// earlier versions.
func LoadPinnedProgram(fileName string) (*Program, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	abi, err := newProgramABIFromFd(fd)
	if err != nil {
		return nil, err
	}
	return &Program{
		uint32(fd),
		filepath.Base(fileName),
		*abi,
	}, nil
}

// LoadPinnedProgramExplicit loads a program with explicit parameters.
func LoadPinnedProgramExplicit(fileName string, abi *ProgramABI) (*Program, error) {
	fd, err := getObject(fileName)
	if err != nil {
		return nil, err
	}
	return &Program{
		uint32(fd),
		filepath.Base(fileName),
		*abi,
	}, nil
}

type loadError struct {
	cause       error
	verifierLog string
}

func (le *loadError) Error() string {
	if le.verifierLog == "" {
		return fmt.Sprintf("failed to load program: %s", le.cause)
	}
	return fmt.Sprintf("failed to load program: %s: %s", le.cause, le.verifierLog)
}

func (le *loadError) Cause() error {
	return le.cause
}

func convertCString(in []byte) string {
	inLen := bytes.IndexByte(in, 0)
	return string(in[:inLen])
}

func checkName(name string) error {
	if len(name) > bpfObjNameLen-1 {
		return errors.Errorf("name '%s' is too long", name)
	}

	idx := strings.IndexFunc(name, func(char rune) bool {
		switch {
		case char >= 'A' && char <= 'Z':
			fallthrough
		case char >= 'a' && char <= 'z':
			fallthrough
		case char >= '0' && char <= '9':
			fallthrough
		case char == '_':
			return false
		default:
			return true
		}
	})

	if idx != -1 {
		return errors.Errorf("invalid character '%c' in name '%s'", name[idx], name)
	}

	return nil
}
<|endoftext|>


<commit_before>//   Copyright 2016 DigitalOcean
//
//   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.

// Command ceph_exporter provides a Prometheus exporter for a Ceph cluster.
package main

import (
	"errors"
	"net"
	"net/http"
	"os"
	"syscall"
	"time"

	"github.com/ianschenck/envflag"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/sirupsen/logrus"

	"github.com/digitalocean/ceph_exporter/ceph"
	"github.com/digitalocean/ceph_exporter/rados"
)

const (
	defaultCephClusterLabel = "ceph"
	defaultCephConfigPath   = "/etc/ceph/ceph.conf"
	defaultCephUser         = "admin"
	defaultRadosOpTimeout   = 30 * time.Second
)

var (
	errCephVersionUnsupported = errors.New("ceph version unsupported")
)

// This horrible thing is a copy of tcpKeepAliveListener, tweaked to
// specifically check if it hits EMFILE when doing an accept, and if so,
// terminate the process.
const keepAlive time.Duration = 3 * time.Minute

type emfileAwareTcpListener struct {
	*net.TCPListener
	logger *logrus.Logger
}

func (ln emfileAwareTcpListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		if oerr, ok := err.(*net.OpError); ok {
			if serr, ok := oerr.Err.(*os.SyscallError); ok && serr.Err == syscall.EMFILE {
				ln.logger.WithError(err).Fatal("running out of file descriptors")
			}
		}
		// Default return
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(keepAlive)
	return tc, nil
}

// Verify that the exporter implements the interface correctly.
var _ prometheus.Collector = &ceph.Exporter{}

func main() {
	var (
		metricsAddr    = envflag.String("TELEMETRY_ADDR", ":9128", "Host:Port for ceph_exporter's metrics endpoint")
		metricsPath    = envflag.String("TELEMETRY_PATH", "/metrics", "URL path for surfacing metrics to Prometheus")
		exporterConfig = envflag.String("EXPORTER_CONFIG", "/etc/ceph/exporter.yml", "Path to ceph_exporter config")
		rgwMode        = envflag.Int("RGW_MODE", 0, "Enable collection of stats from RGW (0:disabled 1:enabled 2:background)")

		logLevel = envflag.String("LOG_LEVEL", "info", "Logging level. One of: [trace, debug, info, warn, error, fatal, panic]")

		cephCluster        = envflag.String("CEPH_CLUSTER", defaultCephClusterLabel, "Ceph cluster name")
		cephConfig         = envflag.String("CEPH_CONFIG", defaultCephConfigPath, "Path to Ceph config file")
		cephUser           = envflag.String("CEPH_USER", defaultCephUser, "Ceph user to connect to cluster")
		cephRadosOpTimeout = envflag.Duration("CEPH_RADOS_OP_TIMEOUT", defaultRadosOpTimeout, "Ceph rados_osd_op_timeout and rados_mon_op_timeout used to contact cluster (0s means no limit)")

		tlsCertPath = envflag.String("TLS_CERT_FILE_PATH", "", "Path to certificate file for TLS")
		tlsKeyPath  = envflag.String("TLS_KEY_FILE_PATH", "", "Path to key file for TLS")
	)

	envflag.Parse()

	logger := logrus.New()
	logger.SetFormatter(&logrus.TextFormatter{
		FullTimestamp: true,
	})

	if v, err := logrus.ParseLevel(*logLevel); err != nil {
		logger.WithError(err).Warn("error setting log level")
	} else {
		logger.SetLevel(v)
	}

	clusterConfigs := ([]*ClusterConfig)(nil)

	if fileExists(*exporterConfig) {
		cfg, err := ParseConfig(*exporterConfig)
		if err != nil {
			logger.WithError(err).WithField(
				"file", *exporterConfig,
			).Fatal("error parsing ceph_exporter config file")
		}
		clusterConfigs = cfg.Cluster
	} else {
		clusterConfigs = []*ClusterConfig{
			{
				ClusterLabel: *cephCluster,
				User:         *cephUser,
				ConfigFile:   *cephConfig,
			},
		}
	}

	for _, cluster := range clusterConfigs {
		conn := rados.NewRadosConn(
			cluster.User,
			cluster.ConfigFile,
			*cephRadosOpTimeout,
			logger)

		prometheus.MustRegister(ceph.NewExporter(
			conn,
			cluster.ClusterLabel,
			cluster.ConfigFile,
			*rgwMode,
			logger))

		logger.WithField("cluster", cluster.ClusterLabel).Info("exporting cluster")
	}

	http.Handle(*metricsPath, promhttp.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Ceph Exporter</title>

Ceph Exporter

<a href='` + *metricsPath + `'>Metrics

			`))
	})

	logger.WithField("endpoint", *metricsAddr).Info("starting ceph_exporter listener")

	// Below is essentially http.ListenAndServe(), but using our custom
	// emfileAwareTcpListener that will die if we run out of file descriptors
	ln, err := net.Listen("tcp", *metricsAddr)
	if err != nil {
		logrus.WithError(err).Fatal("error creating listener")
	}

	if len(*tlsCertPath) != 0 && len(*tlsKeyPath) != 0 {
		err = http.ServeTLS(emfileAwareTcpListener{ln.(*net.TCPListener), logger}, nil, *tlsCertPath, *tlsKeyPath)
		if err != nil {
			logrus.WithError(err).Fatal("error serving TLS requests")
		}
	} else {
		err = http.Serve(emfileAwareTcpListener{ln.(*net.TCPListener), logger}, nil)
		if err != nil {
			logrus.WithError(err).Fatal("error serving requests")
		}
	}
}
<commit_msg>reload tls cert whenever it is requested<commit_after>//   Copyright 2016 DigitalOcean
//
//   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.

// Command ceph_exporter provides a Prometheus exporter for a Ceph cluster.
package main

import (
	"crypto/tls"
	"errors"
	"net"
	"net/http"
	"os"
	"syscall"
	"time"

	"github.com/ianschenck/envflag"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/sirupsen/logrus"

	"github.com/digitalocean/ceph_exporter/ceph"
	"github.com/digitalocean/ceph_exporter/rados"
)

const (
	defaultCephClusterLabel = "ceph"
	defaultCephConfigPath   = "/etc/ceph/ceph.conf"
	defaultCephUser         = "admin"
	defaultRadosOpTimeout   = 30 * time.Second
)

var (
	errCephVersionUnsupported = errors.New("ceph version unsupported")
)

// This horrible thing is a copy of tcpKeepAliveListener, tweaked to
// specifically check if it hits EMFILE when doing an accept, and if so,
// terminate the process.
const keepAlive time.Duration = 3 * time.Minute

type emfileAwareTcpListener struct {
	*net.TCPListener
	logger *logrus.Logger
}

func (ln emfileAwareTcpListener) Accept() (c net.Conn, err error) {
	tc, err := ln.AcceptTCP()
	if err != nil {
		if oerr, ok := err.(*net.OpError); ok {
			if serr, ok := oerr.Err.(*os.SyscallError); ok && serr.Err == syscall.EMFILE {
				ln.logger.WithError(err).Fatal("running out of file descriptors")
			}
		}
		// Default return
		return
	}
	tc.SetKeepAlive(true)
	tc.SetKeepAlivePeriod(keepAlive)
	return tc, nil
}

// Verify that the exporter implements the interface correctly.
var _ prometheus.Collector = &ceph.Exporter{}

func main() {
	var (
		metricsAddr    = envflag.String("TELEMETRY_ADDR", ":9128", "Host:Port for ceph_exporter's metrics endpoint")
		metricsPath    = envflag.String("TELEMETRY_PATH", "/metrics", "URL path for surfacing metrics to Prometheus")
		exporterConfig = envflag.String("EXPORTER_CONFIG", "/etc/ceph/exporter.yml", "Path to ceph_exporter config")
		rgwMode        = envflag.Int("RGW_MODE", 0, "Enable collection of stats from RGW (0:disabled 1:enabled 2:background)")

		logLevel = envflag.String("LOG_LEVEL", "info", "Logging level. One of: [trace, debug, info, warn, error, fatal, panic]")

		cephCluster        = envflag.String("CEPH_CLUSTER", defaultCephClusterLabel, "Ceph cluster name")
		cephConfig         = envflag.String("CEPH_CONFIG", defaultCephConfigPath, "Path to Ceph config file")
		cephUser           = envflag.String("CEPH_USER", defaultCephUser, "Ceph user to connect to cluster")
		cephRadosOpTimeout = envflag.Duration("CEPH_RADOS_OP_TIMEOUT", defaultRadosOpTimeout, "Ceph rados_osd_op_timeout and rados_mon_op_timeout used to contact cluster (0s means no limit)")

		tlsCertPath = envflag.String("TLS_CERT_FILE_PATH", "", "Path to certificate file for TLS")
		tlsKeyPath  = envflag.String("TLS_KEY_FILE_PATH", "", "Path to key file for TLS")
	)

	envflag.Parse()

	logger := logrus.New()
	logger.SetFormatter(&logrus.TextFormatter{
		FullTimestamp: true,
	})

	if v, err := logrus.ParseLevel(*logLevel); err != nil {
		logger.WithError(err).Warn("error setting log level")
	} else {
		logger.SetLevel(v)
	}

	clusterConfigs := ([]*ClusterConfig)(nil)

	if fileExists(*exporterConfig) {
		cfg, err := ParseConfig(*exporterConfig)
		if err != nil {
			logger.WithError(err).WithField(
				"file", *exporterConfig,
			).Fatal("error parsing ceph_exporter config file")
		}
		clusterConfigs = cfg.Cluster
	} else {
		clusterConfigs = []*ClusterConfig{
			{
				ClusterLabel: *cephCluster,
				User:         *cephUser,
				ConfigFile:   *cephConfig,
			},
		}
	}

	for _, cluster := range clusterConfigs {
		conn := rados.NewRadosConn(
			cluster.User,
			cluster.ConfigFile,
			*cephRadosOpTimeout,
			logger)

		prometheus.MustRegister(ceph.NewExporter(
			conn,
			cluster.ClusterLabel,
			cluster.ConfigFile,
			*rgwMode,
			logger))

		logger.WithField("cluster", cluster.ClusterLabel).Info("exporting cluster")
	}

	http.Handle(*metricsPath, promhttp.Handler())
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`
			<title>Ceph Exporter</title>

Ceph Exporter

<a href='` + *metricsPath + `'>Metrics

			`))
	})

	logger.WithField("endpoint", *metricsAddr).Info("starting ceph_exporter listener")

	// Below is essentially http.ListenAndServe(), but using our custom
	// emfileAwareTcpListener that will die if we run out of file descriptors
	ln, err := net.Listen("tcp", *metricsAddr)
	if err != nil {
		logrus.WithError(err).Fatal("error creating listener")
	}

	if len(*tlsCertPath) != 0 && len(*tlsKeyPath) != 0 {
		server := &http.Server{
			TLSConfig: &tls.Config{
				GetCertificate: func(info *tls.ClientHelloInfo) (*tls.Certificate, error) {
					caFiles, err := tls.LoadX509KeyPair(*tlsCertPath, *tlsKeyPath)
					if err != nil {
						return nil, err
					}

					return &caFiles, nil
				},
			},
		}

		err = server.ServeTLS(emfileAwareTcpListener{ln.(*net.TCPListener), logger}, "", "")
		if err != nil {
			logrus.WithError(err).Fatal("error serving TLS requests")
		}
	} else {
		err = http.Serve(emfileAwareTcpListener{ln.(*net.TCPListener), logger}, nil)
		if err != nil {
			logrus.WithError(err).Fatal("error serving requests")
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	sparta "github.com/mweagle/Sparta"
	"github.com/mweagle/SpartaCodePipeline/pipeline"
	"github.com/spf13/cobra"
	"gopkg.in/go-playground/validator.v9"
)

// PipelineName is the name of the stack to provision that supports the pipeline
var pipelineOptions pipeline.ProvisionOptions

func init() {
	sparta.RegisterCodePipelineEnvironment("test", map[string]string{
		"MESSAGE":     "Hello Test!",
		"ENVIRONMENT": "test",
	})
	sparta.RegisterCodePipelineEnvironment("production", map[string]string{
		"MESSAGE":     "Hello Production!",
		"ENVIRONMENT": "prod",
	})
}

// Standard AWS Lambda function
func helloSpartaWorld() (string, error) {
	messageText := os.Getenv("MESSAGE")
	if "" == messageText {
		messageText = "$MESSAGE not defined"
	}
	return messageText, nil
}

////////////////////////////////////////////////////////////////////////////////
// Add a command to provision a CI pipeline
var pipelineProvisionCommand = &cobra.Command{
	Use:   "provisionPipeline",
	Short: "Provision a CI/CD pipeline for this stack",
	RunE: func(cmd *cobra.Command, args []string) error {
		validate := validator.New()
		cliErrors := validate.Struct(&pipelineOptions)
		if cliErrors != nil {
			return cliErrors
		}
		return pipeline.Provision(&pipelineOptions)
	},
}

////////////////////////////////////////////////////////////////////////////////
// Main
func main() {
	// Register the provisionPipeline command
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.PipelineName, "pipeline", "p", "", "pipeline name")
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.GithubRepo, "repo", "r", "", "GitHub Repo URL")
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.GithubOAuthToken, "oauth", "o", "", "GitHub OAuth token")
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.S3Bucket,
		"s3Bucket",
		"s",
		"",
		"S3 Bucket to use for Lambda source")
	pipelineProvisionCommand.PersistentFlags().BoolVarP(&pipelineOptions.Noop, "noop",
		"n",
		false,
		"Dry-run behavior only (do not perform mutations)")
	sparta.CommandLineOptions.Root.AddCommand(pipelineProvisionCommand)

	// Normal execution
	lambdaFn := sparta.HandleAWSLambda(fmt.Sprintf("HelloWorld-%s", os.Getenv("ENVIRONMENT")),
		helloSpartaWorld,
		sparta.IAMRoleDefinition{})
	var lambdaFunctions []*sparta.LambdaAWSInfo
	lambdaFunctions = append(lambdaFunctions, lambdaFn)
	err := sparta.Main("SpartaCodePipeline",
		fmt.Sprintf("SpartaCodePipeline CodePipeline example"),
		lambdaFunctions,
		nil,
		nil)

	if err != nil {
		os.Exit(1)
	}
}
<commit_msg>DRY name<commit_after>package main

import (
	"fmt"
	"os"

	sparta "github.com/mweagle/Sparta"
	"github.com/mweagle/SpartaCodePipeline/pipeline"
	"github.com/spf13/cobra"
	"gopkg.in/go-playground/validator.v9"
)

// PipelineName is the name of the stack to provision that supports the pipeline
var pipelineOptions pipeline.ProvisionOptions

func init() {
	sparta.RegisterCodePipelineEnvironment("test", map[string]string{
		"MESSAGE":     "Hello Test!",
		"ENVIRONMENT": "test",
	})
	sparta.RegisterCodePipelineEnvironment("production", map[string]string{
		"MESSAGE":     "Hello Production!",
		"ENVIRONMENT": "prod",
	})
}

// Standard AWS Lambda function
func helloSpartaWorld() (string, error) {
	messageText := os.Getenv("MESSAGE")
	if "" == messageText {
		messageText = "$MESSAGE not defined"
	}
	return messageText, nil
}

////////////////////////////////////////////////////////////////////////////////
// Add a command to provision a CI pipeline
var pipelineProvisionCommand = &cobra.Command{
	Use:   "provisionPipeline",
	Short: "Provision a CI/CD pipeline for this stack",
	RunE: func(cmd *cobra.Command, args []string) error {
		validate := validator.New()
		cliErrors := validate.Struct(&pipelineOptions)
		if cliErrors != nil {
			return cliErrors
		}
		return pipeline.Provision(&pipelineOptions)
	},
}

////////////////////////////////////////////////////////////////////////////////
// Main
func main() {
	// Register the provisionPipeline command
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.PipelineName, "pipeline", "p", "", "pipeline name")
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.GithubRepo, "repo", "r", "", "GitHub Repo URL")
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.GithubOAuthToken, "oauth", "o", "", "GitHub OAuth token")
	pipelineProvisionCommand.PersistentFlags().StringVarP(&pipelineOptions.S3Bucket,
		"s3Bucket",
		"s",
		"",
		"S3 Bucket to use for Lambda source")
	pipelineProvisionCommand.PersistentFlags().BoolVarP(&pipelineOptions.Noop, "noop",
		"n",
		false,
		"Dry-run behavior only (do not perform mutations)")
	sparta.CommandLineOptions.Root.AddCommand(pipelineProvisionCommand)

	// Normal execution
	lambdaFn := sparta.HandleAWSLambda("HelloWorld",
		helloSpartaWorld,
		sparta.IAMRoleDefinition{})
	var lambdaFunctions []*sparta.LambdaAWSInfo
	lambdaFunctions = append(lambdaFunctions, lambdaFn)
	err := sparta.Main("SpartaCodePipeline",
		fmt.Sprintf("SpartaCodePipeline CodePipeline example"),
		lambdaFunctions,
		nil,
		nil)

	if err != nil {
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"net"
	"time"
	"fmt"
)

func handleConn(conn net.Conn) {
	fmt.Printf("New Connection %v\n", conn.RemoteAddr())
	packet := "HTTP/1.1 200 OK\nContent-Type: text/plain; charset=ascii\n\n"
	buff := make([]byte, len(packet))
	copy(buff[:], packet)
	_, err := conn.Write(buff)
	if err != nil {
		fmt.Printf("Lost Connection %v\n", conn.RemoteAddr())
		return
	}
	tb := []byte{74, 74}
	for {
//		t := time.Now()
		_, err = conn.Write(tb)
//		elapsed := time.Since(t)
//		fmt.Printf("Elapsed time for write: %v\n", elapsed)
		if err != nil {
			fmt.Printf("Lost Connection %v\n", conn.RemoteAddr())
			return
		}
		time.Sleep(time.Second)
	}
}

func main() {
	l, err := net.Listen("tcp", ":5566")
	if err != nil {
		return
	}
	for {
		conn, err := l.Accept()
		if err == nil {
			go handleConn(conn)
		}
	}
}
<commit_msg>small tweaks<commit_after>package main

import (
	"fmt"
	"net"
	"time"
)

func handleConn(conn net.Conn) {
	fmt.Printf("New Connection %v\n", conn.RemoteAddr())
	packet := "HTTP/1.1 200 OK\nContent-Type: text/plain; charset=ascii\nTransfer-Encoding: chunked\n\n"
	buff := make([]byte, len(packet))
	copy(buff[:], packet)
	_, err := conn.Write(buff)
	if err != nil {
		fmt.Printf("Lost Connection %v\n", conn.RemoteAddr())
		return
	}
	tb := []byte{74, 74}
	for {
		_, err = conn.Write(tb)
		if err != nil {
			fmt.Printf("Lost Connection %v\n", conn.RemoteAddr())
			return
		}
		time.Sleep(time.Second / 2)
	}
}

func main() {
	l, err := net.Listen("tcp", ":5566")
	if err != nil {
		return
	}
	for {
		conn, err := l.Accept()
		if err == nil {
			go handleConn(conn)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/aqatl/mal/mal"
	"github.com/urfave/cli"
	"os"
	"log"
	"sort"
	"strconv"
	"path/filepath"
	"github.com/skratchdot/open-golang/open"
	"fmt"
	"strings"
)

var dataDir = filepath.Join(homeDir(), ".mal")
var (
	CredentialsFile = filepath.Join(dataDir, "cred.dat")
	MalCacheFile    = filepath.Join(dataDir, "cache.xml")
	ConfigFile      = filepath.Join(dataDir, "config.json")
)

func main() {
	checkDataDir()

	app := cli.NewApp()
	app.Name = "mal"
	app.Usage = "App for managing your MAL"
	app.Version = "0.1"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "creds, prompt-credentials",
			Usage: "Prompt for username and password",
		},
		cli.BoolFlag{
			Name:  "sp, save-password",
			Usage: "save your password. Use with caution, your password can be decoded",
		},
		cli.BoolFlag{
			Name:  "r, refresh",
			Usage: "refreshes cache file",
		},
		cli.BoolFlag{
			Name:  "ver, verify",
			Usage: "verify credentials",
		},
		cli.IntFlag{
			Name:  "max",
			Usage: "visible entries threshold",
		},
		cli.StringFlag{
			Name:  "status",
			Usage: "display entries only with given status",
		},
	}

	app.Commands = []cli.Command{
		cli.Command{
			Name:      "inc",
			Aliases:   []string{"+1"},
			Category:  "Update",
			Usage:     "Increment selected entry by one",
			UsageText: "mal inc",
			Action:    incrementEntry,
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "n",
					Usage: "Specify exact episode to set the entry to",
				},
			},
		},
		cli.Command{
			Name:      "score",
			Category:  "Update",
			Usage:     "Set your rating for selected entry",
			UsageText: "mal score <0-10>",
			Action:    setEntryScore,
		},
		cli.Command{
			Name:      "status",
			Category:  "Update",
			Usage:     "Set your status for selected entry",
			UsageText: "mal status [watching|completed|onhold|dropped|plantowatch]",
			Action:    setEntryStatus,
		},
		cli.Command{
			Name:      "sel",
			Aliases:   []string{"select"},
			Category:  "Config",
			Usage:     "Select an entry",
			UsageText: "mal sel [entry ID]",
			Action:    selectEntry,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "by-title",
					Usage: "Select entry by name instead of by ID",
				},
			},
		},
		cli.Command{
			Name:     "cfg",
			Aliases:  []string{"config", "configuration"},
			Category: "Config",
			Usage:    "Change config values",
			Subcommands: cli.Commands{
				cli.Command{
					Name:      "max",
					Aliases:   []string{"visible"},
					Usage:     "Change amount of displayed entries",
					UsageText: "mal cfg max [number]",
					Action:    configChangeMax,
				},
				cli.Command{
					Name:      "status",
					Usage:     "Status value of displayed entries",
					UsageText: "mal cfg status [all|watching|completed|onhold|dropped|plantowatch]",
					Action:    configChangeStatus,
				},
				cli.Command{
					Name:      "status-auto-update",
					Usage:     "Allows entry to be automatically set to completed when number of all episodes is reached or exceeded",
					UsageText: "mal cfg status-auto-update [off|normal|after-threshold]",
					Action:    configChangeAutoUpdateMode,
				},
			},
		},
		cli.Command{
			Name:      "web",
			Aliases:   []string{"website", "open"},
			Category:  "Action",
			Usage:     "Open url associated with current entry",
			UsageText: "mal web",
			Action:    openWebsite,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "url",
					Usage: "Set url for current entry",
				},
				cli.BoolFlag{
					Name:  "clear",
					Usage: "Clear url for current entry",
				},
			},
			Subcommands: []cli.Command{
				cli.Command{
					Name:      "get-all",
					Usage:     "Print all set urls",
					UsageText: "mal web get-all",
					Action:    printWebsites,
				},
			},
		},
	}

	app.Action = cli.ActionFunc(defaultAction)

	if err := app.Run(os.Args); err != nil {
		log.Fatalf("Error: %v", err)
	}
}

func defaultAction(ctx *cli.Context) error {
	creds := loadCredentials(ctx)
	if ctx.GlobalBool("verify") && !mal.VerifyCredentials(creds) {
		return fmt.Errorf("invalid credentials")
	}

	c := mal.NewClient(creds)
	if c == nil {
		os.Exit(1)
	}

	config := LoadConfig()

	list := loadList(c, ctx).FilterByStatus(statusFlag(ctx))
	sort.Sort(mal.AnimeSortByLastUpdated(list))

	var visibleEntries int
	if visibleEntries = ctx.Int("max"); visibleEntries == 0 {
		//`Max` flag not specified, get value from config
		visibleEntries = config.MaxVisibleEntries
	}
	if visibleEntries > len(list) {
		visibleEntries = len(list)
	}
	visibleList := list[:visibleEntries]
	reverseAnimeSlice(visibleList)

	PrettyList.Execute(os.Stdout, PrettyListData{visibleList, config.SelectedID})

	if ctx.GlobalBool("save-password") {
		saveCredentials(creds)
	}
	return nil
}

func incrementEntry(ctx *cli.Context) error {
	creds := loadCredentials(ctx)
	if ctx.Bool("verify") && !mal.VerifyCredentials(creds) {
		return fmt.Errorf("invalid credentials")
	}

	c := mal.NewClient(creds)
	if c == nil {
		os.Exit(1)
	}
	cfg := LoadConfig()

	if cfg.SelectedID == 0 {
		return fmt.Errorf("no entry selected")
	}

	list := loadList(c, ctx)
	selectedEntry := list.GetByID(cfg.SelectedID)

	if selectedEntry == nil {
		return fmt.Errorf("no entry found")
	}

	if ctx.Int("n") > 0 {
		selectedEntry.WatchedEpisodes = ctx.Int("n")
	} else {
		selectedEntry.WatchedEpisodes++
	}

	statusAutoUpdate(cfg, selectedEntry)

	if c.Update(selectedEntry) {
		fmt.Println("Updated successfully")
		printEntryDetails(selectedEntry)

		cacheList(list)
	}
	return nil
}

func setEntryScore(ctx *cli.Context) error {
	c := mal.NewClient(loadCredentials(ctx))
	cfg := LoadConfig()

	list := loadList(c, ctx)
	selectedEntry := list.GetByID(cfg.SelectedID)

	score, err := strconv.Atoi(ctx.Args().First())
	if err != nil {
		return fmt.Errorf("error parsing score: %v", err)
	}
	parsedScore, err := mal.ParseScore(score)
	if err != nil {
		return err
	}

	selectedEntry.MyScore = parsedScore
	if c.Update(selectedEntry) {
		fmt.Println("Updated successfully")
		printEntryDetails(selectedEntry)

		cacheList(list)
	}
	return nil
}

func setEntryStatus(ctx *cli.Context) error {
	c := mal.NewClient(loadCredentials(ctx))
	cfg := LoadConfig()

	list := loadList(c, ctx)
	selectedEntry := list.GetByID(cfg.SelectedID)

	status := mal.ParseStatus(ctx.Args().First())
	if status == mal.All {
		return fmt.Errorf("invalid status; possible values: watching, completed, " +
			"onhold, dropped, plantowatch")
	}

	selectedEntry.MyStatus = status
	if c.Update(selectedEntry) {
		fmt.Println("Updated successfully")
		printEntryDetails(selectedEntry)

		cacheList(list)
	}
	return nil
}

func selectEntry(ctx *cli.Context) error {
	if ctx.Bool("by-title") {
		return selectByTitle(ctx)
	}

	id, err := strconv.Atoi(ctx.Args().First())
	if err != nil {
		return fmt.Errorf("error parsing id: %v", err)
	}

	cfg := LoadConfig()
	cfg.SelectedID = id
	cfg.Save()

	fmt.Println("Selected entry:")
	printEntryDetails(loadList(mal.NewClient(loadCredentials(ctx)), ctx).GetByID(id))

	return nil
}

func selectByTitle(ctx *cli.Context) error {
	title := strings.ToLower(ctx.Args().First())
	list := loadList(mal.NewClient(loadCredentials(ctx)), ctx)

	found := make(mal.AnimeList, 0)
	for _, entry := range list {
		if strings.Contains(strings.ToLower(entry.Title), title) {
			found = append(found, entry)
		}
	}

	var selectedEntry *mal.Anime

	if len(found) > 1 {
		fmt.Printf("Found more than 1 matching entry:\n")
		fmt.Printf("%3s%8s%7s\n", "No.", "ID", "Title")
		fmt.Println(strings.Repeat("=", 80))
		for i, entry := range found {
			fmt.Printf("%3d. %6d: %s\n", i, entry.ID, entry.Title)
		}

		fmt.Printf("Enter index of the selected entry: ")
		idx := 0
		_, err := fmt.Scanln(&idx)
		if err != nil || idx < 0 || idx > len(found)-1 {
			return fmt.Errorf("invalid input %v", err)
		}

		selectedEntry = found[idx]
	} else if len(found) == 0 {
		return fmt.Errorf("no matches")
	} else {
		selectedEntry = found[0]
	}

	cfg := LoadConfig()
	cfg.SelectedID = selectedEntry.ID
	cfg.Save()

	fmt.Println("Selected entry:")
	printEntryDetails(selectedEntry)

	return nil
}

func openWebsite(ctx *cli.Context) error {
	cfg := LoadConfig()

	if url := ctx.String("url"); url != "" {
		cfg.Websites[cfg.SelectedID] = url
		log.Printf("Set url %s for entry %d", url, cfg.SelectedID)
		cfg.Save()
		return nil
	}

	if ctx.Bool("clear") {
		delete(cfg.Websites, cfg.SelectedID)
		log.Printf("Cleared url for entry %d", cfg.SelectedID)
		cfg.Save()
		return nil
	}

	if url, ok := cfg.Websites[cfg.SelectedID]; ok {
		open.Start(url)
		fmt.Println("Opened website for:")
		printEntryDetails(
			loadList(mal.NewClient(loadCredentials(ctx)), ctx).
				GetByID(cfg.SelectedID),
		)
	} else {
		log.Println("Nothing to open")
	}

	return nil
}

func printWebsites(ctx *cli.Context) error {
	cfg := LoadConfig()
	list := loadList(mal.NewClient(loadCredentials(ctx)), ctx)

	for k, v := range cfg.Websites {
		url := fmt.Sprintf("\033[3%d;%dm%s\033[0m ", 3, 1, v)

		var title string
		if entry := list.GetByID(k); entry != nil {
			title = entry.Title
		}

		fmt.Printf("%6d (%s): %s\n", k, title, url)
	}

	return nil
}

func configChangeMax(ctx *cli.Context) error {
	cfg := LoadConfig()

	max, err := strconv.Atoi(ctx.Args().First())
	if err != nil || max < 0 {
		return fmt.Errorf("invalid value")
	}

	cfg.MaxVisibleEntries = max
	cfg.Save()
	return nil
}

func configChangeStatus(ctx *cli.Context) error {
	cfg := LoadConfig()

	status := mal.ParseStatus(ctx.Args().First())

	cfg.Status = status
	cfg.Save()
	return nil
}

func configChangeAutoUpdateMode(ctx *cli.Context) error {
	arg := strings.ToLower(ctx.Args().First())
	var mode StatusAutoUpdateMode

	if arg == "off" {
		mode = Off
	} else if arg == "normal" {
		mode = Normal
	} else if arg == "after-threshold" {
		mode = AfterThreshold
	} else {
		return fmt.Errorf("invalid option; possible values: off|normal|after-threshold")
	}

	cfg := LoadConfig()
	cfg.StatusAutoUpdateMode = mode
	cfg.Save()

	return nil
}
<commit_msg>Added 'details' command for printing selected entry details<commit_after>package main

import (
	"github.com/aqatl/mal/mal"
	"github.com/urfave/cli"
	"os"
	"log"
	"sort"
	"strconv"
	"path/filepath"
	"github.com/skratchdot/open-golang/open"
	"fmt"
	"strings"
)

var dataDir = filepath.Join(homeDir(), ".mal")
var (
	CredentialsFile = filepath.Join(dataDir, "cred.dat")
	MalCacheFile    = filepath.Join(dataDir, "cache.xml")
	ConfigFile      = filepath.Join(dataDir, "config.json")
)

func main() {
	checkDataDir()

	app := cli.NewApp()
	app.Name = "mal"
	app.Usage = "App for managing your MAL"
	app.Version = "0.1"

	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "creds, prompt-credentials",
			Usage: "Prompt for username and password",
		},
		cli.BoolFlag{
			Name:  "sp, save-password",
			Usage: "save your password. Use with caution, your password can be decoded",
		},
		cli.BoolFlag{
			Name:  "r, refresh",
			Usage: "refreshes cache file",
		},
		cli.BoolFlag{
			Name:  "ver, verify",
			Usage: "verify credentials",
		},
		cli.IntFlag{
			Name:  "max",
			Usage: "visible entries threshold",
		},
		cli.StringFlag{
			Name:  "status",
			Usage: "display entries only with given status",
		},
	}

	app.Commands = []cli.Command{
		cli.Command{
			Name:      "inc",
			Aliases:   []string{"+1"},
			Category:  "Update",
			Usage:     "Increment selected entry by one",
			UsageText: "mal inc",
			Action:    incrementEntry,
			Flags: []cli.Flag{
				cli.IntFlag{
					Name:  "n",
					Usage: "Specify exact episode to set the entry to",
				},
			},
		},
		cli.Command{
			Name:      "score",
			Category:  "Update",
			Usage:     "Set your rating for selected entry",
			UsageText: "mal score <0-10>",
			Action:    setEntryScore,
		},
		cli.Command{
			Name:      "status",
			Category:  "Update",
			Usage:     "Set your status for selected entry",
			UsageText: "mal status [watching|completed|onhold|dropped|plantowatch]",
			Action:    setEntryStatus,
		},
		cli.Command{
			Name:      "sel",
			Aliases:   []string{"select"},
			Category:  "Config",
			Usage:     "Select an entry",
			UsageText: "mal sel [entry ID]",
			Action:    selectEntry,
			Flags: []cli.Flag{
				cli.BoolFlag{
					Name:  "by-title",
					Usage: "Select entry by name instead of by ID",
				},
			},
		},
		cli.Command{
			Name:     "cfg",
			Aliases:  []string{"config", "configuration"},
			Category: "Config",
			Usage:    "Change config values",
			Subcommands: cli.Commands{
				cli.Command{
					Name:      "max",
					Aliases:   []string{"visible"},
					Usage:     "Change amount of displayed entries",
					UsageText: "mal cfg max [number]",
					Action:    configChangeMax,
				},
				cli.Command{
					Name:      "status",
					Usage:     "Status value of displayed entries",
					UsageText: "mal cfg status [all|watching|completed|onhold|dropped|plantowatch]",
					Action:    configChangeStatus,
				},
				cli.Command{
					Name:      "status-auto-update",
					Usage:     "Allows entry to be automatically set to completed when number of all episodes is reached or exceeded",
					UsageText: "mal cfg status-auto-update [off|normal|after-threshold]",
					Action:    configChangeAutoUpdateMode,
				},
			},
		},
		cli.Command{
			Name:      "web",
			Aliases:   []string{"website", "open"},
			Category:  "Action",
			Usage:     "Open url associated with selected entry",
			UsageText: "mal web",
			Action:    openWebsite,
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "url",
					Usage: "Set url for current entry",
				},
				cli.BoolFlag{
					Name:  "clear",
					Usage: "Clear url for current entry",
				},
			},
			Subcommands: []cli.Command{
				cli.Command{
					Name:      "get-all",
					Usage:     "Print all set urls",
					UsageText: "mal web get-all",
					Action:    printWebsites,
				},
			},
		},
		cli.Command{
			Name: "details",
			Category: "Action",
			Usage: "Print details about selected entry",
			UsageText: "mal details",
			Action: detailsCommand,
		},
	}

	app.Action = cli.ActionFunc(defaultAction)

	if err := app.Run(os.Args); err != nil {
		log.Fatalf("Error: %v", err)
	}
}

func defaultAction(ctx *cli.Context) error {
	creds := loadCredentials(ctx)
	if ctx.GlobalBool("verify") && !mal.VerifyCredentials(creds) {
		return fmt.Errorf("invalid credentials")
	}

	c := mal.NewClient(creds)
	if c == nil {
		os.Exit(1)
	}

	config := LoadConfig()

	list := loadList(c, ctx).FilterByStatus(statusFlag(ctx))
	sort.Sort(mal.AnimeSortByLastUpdated(list))

	var visibleEntries int
	if visibleEntries = ctx.Int("max"); visibleEntries == 0 {
		//`Max` flag not specified, get value from config
		visibleEntries = config.MaxVisibleEntries
	}
	if visibleEntries > len(list) {
		visibleEntries = len(list)
	}
	visibleList := list[:visibleEntries]
	reverseAnimeSlice(visibleList)

	PrettyList.Execute(os.Stdout, PrettyListData{visibleList, config.SelectedID})

	if ctx.GlobalBool("save-password") {
		saveCredentials(creds)
	}
	return nil
}

func incrementEntry(ctx *cli.Context) error {
	creds := loadCredentials(ctx)
	if ctx.Bool("verify") && !mal.VerifyCredentials(creds) {
		return fmt.Errorf("invalid credentials")
	}

	c := mal.NewClient(creds)
	if c == nil {
		os.Exit(1)
	}
	cfg := LoadConfig()

	if cfg.SelectedID == 0 {
		return fmt.Errorf("no entry selected")
	}

	list := loadList(c, ctx)
	selectedEntry := list.GetByID(cfg.SelectedID)

	if selectedEntry == nil {
		return fmt.Errorf("no entry found")
	}

	if ctx.Int("n") > 0 {
		selectedEntry.WatchedEpisodes = ctx.Int("n")
	} else {
		selectedEntry.WatchedEpisodes++
	}

	statusAutoUpdate(cfg, selectedEntry)

	if c.Update(selectedEntry) {
		fmt.Println("Updated successfully")
		printEntryDetails(selectedEntry)

		cacheList(list)
	}
	return nil
}

func setEntryScore(ctx *cli.Context) error {
	c := mal.NewClient(loadCredentials(ctx))
	cfg := LoadConfig()

	list := loadList(c, ctx)
	selectedEntry := list.GetByID(cfg.SelectedID)

	score, err := strconv.Atoi(ctx.Args().First())
	if err != nil {
		return fmt.Errorf("error parsing score: %v", err)
	}
	parsedScore, err := mal.ParseScore(score)
	if err != nil {
		return err
	}

	selectedEntry.MyScore = parsedScore
	if c.Update(selectedEntry) {
		fmt.Println("Updated successfully")
		printEntryDetails(selectedEntry)

		cacheList(list)
	}
	return nil
}

func setEntryStatus(ctx *cli.Context) error {
	c := mal.NewClient(loadCredentials(ctx))
	cfg := LoadConfig()

	list := loadList(c, ctx)
	selectedEntry := list.GetByID(cfg.SelectedID)

	status := mal.ParseStatus(ctx.Args().First())
	if status == mal.All {
		return fmt.Errorf("invalid status; possible values: watching, completed, " +
			"onhold, dropped, plantowatch")
	}

	selectedEntry.MyStatus = status
	if c.Update(selectedEntry) {
		fmt.Println("Updated successfully")
		printEntryDetails(selectedEntry)

		cacheList(list)
	}
	return nil
}

func selectEntry(ctx *cli.Context) error {
	if ctx.Bool("by-title") {
		return selectByTitle(ctx)
	}

	id, err := strconv.Atoi(ctx.Args().First())
	if err != nil {
		return fmt.Errorf("error parsing id: %v", err)
	}

	cfg := LoadConfig()
	cfg.SelectedID = id
	cfg.Save()

	fmt.Println("Selected entry:")
	printEntryDetails(loadList(mal.NewClient(loadCredentials(ctx)), ctx).GetByID(id))

	return nil
}

func selectByTitle(ctx *cli.Context) error {
	title := strings.ToLower(ctx.Args().First())
	list := loadList(mal.NewClient(loadCredentials(ctx)), ctx)

	found := make(mal.AnimeList, 0)
	for _, entry := range list {
		if strings.Contains(strings.ToLower(entry.Title), title) {
			found = append(found, entry)
		}
	}

	var selectedEntry *mal.Anime

	if len(found) > 1 {
		fmt.Printf("Found more than 1 matching entry:\n")
		fmt.Printf("%3s%8s%7s\n", "No.", "ID", "Title")
		fmt.Println(strings.Repeat("=", 80))
		for i, entry := range found {
			fmt.Printf("%3d. %6d: %s\n", i, entry.ID, entry.Title)
		}

		fmt.Printf("Enter index of the selected entry: ")
		idx := 0
		_, err := fmt.Scanln(&idx)
		if err != nil || idx < 0 || idx > len(found)-1 {
			return fmt.Errorf("invalid input %v", err)
		}

		selectedEntry = found[idx]
	} else if len(found) == 0 {
		return fmt.Errorf("no matches")
	} else {
		selectedEntry = found[0]
	}

	cfg := LoadConfig()
	cfg.SelectedID = selectedEntry.ID
	cfg.Save()

	fmt.Println("Selected entry:")
	printEntryDetails(selectedEntry)

	return nil
}

func openWebsite(ctx *cli.Context) error {
	cfg := LoadConfig()

	if url := ctx.String("url"); url != "" {
		cfg.Websites[cfg.SelectedID] = url
		log.Printf("Set url %s for entry %d", url, cfg.SelectedID)
		cfg.Save()
		return nil
	}

	if ctx.Bool("clear") {
		delete(cfg.Websites, cfg.SelectedID)
		log.Printf("Cleared url for entry %d", cfg.SelectedID)
		cfg.Save()
		return nil
	}

	if url, ok := cfg.Websites[cfg.SelectedID]; ok {
		open.Start(url)
		fmt.Println("Opened website for:")
		printEntryDetails(
			loadList(mal.NewClient(loadCredentials(ctx)), ctx).
				GetByID(cfg.SelectedID),
		)
	} else {
		log.Println("Nothing to open")
	}

	return nil
}

func printWebsites(ctx *cli.Context) error {
	cfg := LoadConfig()
	list := loadList(mal.NewClient(loadCredentials(ctx)), ctx)

	for k, v := range cfg.Websites {
		url := fmt.Sprintf("\033[3%d;%dm%s\033[0m ", 3, 1, v)

		var title string
		if entry := list.GetByID(k); entry != nil {
			title = entry.Title
		}

		fmt.Printf("%6d (%s): %s\n", k, title, url)
	}

	return nil
}

func detailsCommand(ctx *cli.Context) error {
	cfg := LoadConfig()
	list := loadList(mal.NewClient(loadCredentials(ctx)), ctx)

	if entry := list.GetByID(cfg.SelectedID); entry != nil {
		printEntryDetails(entry)
	}

	return nil
}

func configChangeMax(ctx *cli.Context) error {
	cfg := LoadConfig()

	max, err := strconv.Atoi(ctx.Args().First())
	if err != nil || max < 0 {
		return fmt.Errorf("invalid value")
	}

	cfg.MaxVisibleEntries = max
	cfg.Save()
	return nil
}

func configChangeStatus(ctx *cli.Context) error {
	cfg := LoadConfig()

	status := mal.ParseStatus(ctx.Args().First())

	cfg.Status = status
	cfg.Save()
	return nil
}

func configChangeAutoUpdateMode(ctx *cli.Context) error {
	arg := strings.ToLower(ctx.Args().First())
	var mode StatusAutoUpdateMode

	if arg == "off" {
		mode = Off
	} else if arg == "normal" {
		mode = Normal
	} else if arg == "after-threshold" {
		mode = AfterThreshold
	} else {
		return fmt.Errorf("invalid option; possible values: off|normal|after-threshold")
	}

	cfg := LoadConfig()
	cfg.StatusAutoUpdateMode = mode
	cfg.Save()

	return nil
}
<|endoftext|>


<commit_before>// Package main defines a command line interface for the sqlboiler package
package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/friendsofgo/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/volatiletech/sqlboiler/v4/boilingcore"
	"github.com/volatiletech/sqlboiler/v4/drivers"
	"github.com/volatiletech/sqlboiler/v4/importers"
)

//go:generate go-bindata -nometadata -pkg templatebin -o templatebin/bindata.go templates templates/singleton templates_test templates_test/singleton

const sqlBoilerVersion = "4.2.0"

var (
	flagConfigFile string
	cmdState       *boilingcore.State
	cmdConfig      *boilingcore.Config
)

func initConfig() {
	if len(flagConfigFile) != 0 {
		viper.SetConfigFile(flagConfigFile)
		if err := viper.ReadInConfig(); err != nil {
			fmt.Println("Can't read config:", err)
			os.Exit(1)
		}
		return
	}

	var err error
	viper.SetConfigName("sqlboiler")

	configHome := os.Getenv("XDG_CONFIG_HOME")
	homePath := os.Getenv("HOME")
	wd, err := os.Getwd()
	if err != nil {
		wd = "."
	}

	configPaths := []string{wd}
	if len(configHome) > 0 {
		configPaths = append(configPaths, filepath.Join(configHome, "sqlboiler"))
	} else {
		configPaths = append(configPaths, filepath.Join(homePath, ".config/sqlboiler"))
	}

	for _, p := range configPaths {
		viper.AddConfigPath(p)
	}

	// Ignore errors here, fallback to other validation methods.
	// Users can use environment variables if a config is not found.
	_ = viper.ReadInConfig()
}

func main() {
	// Too much happens between here and cobra's argument handling, for
	// something so simple just do it immediately.
	for _, arg := range os.Args {
		if arg == "--version" {
			fmt.Println("SQLBoiler v" + sqlBoilerVersion)
			return
		}
	}

	// Set up the cobra root command
	var rootCmd = &cobra.Command{
		Use:   "sqlboiler [flags] <driver>",
		Short: "SQL Boiler generates an ORM tailored to your database schema.",
		Long: "SQL Boiler generates a Go ORM from template files, tailored to your database schema.\n" +
			`Complete documentation is available at http://github.com/volatiletech/sqlboiler`,
		Example:       `sqlboiler psql`,
		PreRunE:       preRun,
		RunE:          run,
		PostRunE:      postRun,
		SilenceErrors: true,
		SilenceUsage:  true,
	}

	cobra.OnInitialize(initConfig)

	// Set up the cobra root command flags
	rootCmd.PersistentFlags().StringVarP(&flagConfigFile, "config", "c", "", "Filename of config file to override default lookup")
	rootCmd.PersistentFlags().StringP("output", "o", "models", "The name of the folder to output to")
	rootCmd.PersistentFlags().StringP("pkgname", "p", "models", "The name you wish to assign to your generated package")
	rootCmd.PersistentFlags().StringSliceP("templates", "", nil, "A templates directory, overrides the bindata'd template folders in sqlboiler")
	rootCmd.PersistentFlags().StringSliceP("tag", "t", nil, "Struct tags to be included on your models in addition to json, yaml, toml")
	rootCmd.PersistentFlags().StringSliceP("replace", "", nil, "Replace templates by directory: relpath/to_file.tpl:relpath/to_replacement.tpl")
	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug mode prints stack traces on error")
	rootCmd.PersistentFlags().BoolP("no-context", "", false, "Disable context.Context usage in the generated code")
	rootCmd.PersistentFlags().BoolP("no-tests", "", false, "Disable generated go test files")
	rootCmd.PersistentFlags().BoolP("no-hooks", "", false, "Disable hooks feature for your models")
	rootCmd.PersistentFlags().BoolP("no-rows-affected", "", false, "Disable rows affected in the generated API")
	rootCmd.PersistentFlags().BoolP("no-auto-timestamps", "", false, "Disable automatic timestamps for created_at/updated_at")
	rootCmd.PersistentFlags().BoolP("no-driver-templates", "", false, "Disable parsing of templates defined by the database driver")
	rootCmd.PersistentFlags().BoolP("no-back-referencing", "", false, "Disable back referencing in the loaded relationship structs")
	rootCmd.PersistentFlags().BoolP("add-global-variants", "", false, "Enable generation for global variants")
	rootCmd.PersistentFlags().BoolP("add-panic-variants", "", false, "Enable generation for panic variants")
	rootCmd.PersistentFlags().BoolP("add-soft-deletes", "", false, "Enable soft deletion by updating deleted_at timestamp")
	rootCmd.PersistentFlags().BoolP("version", "", false, "Print the version")
	rootCmd.PersistentFlags().BoolP("wipe", "", false, "Delete the output folder (rm -rf) before generation to ensure sanity")
	rootCmd.PersistentFlags().StringP("struct-tag-casing", "", "snake", "Decides the casing for go structure tag names. camel, title or snake (default snake)")
	rootCmd.PersistentFlags().StringP("relation-tag", "r", "-", "Relationship struct tag name")
	rootCmd.PersistentFlags().StringSliceP("tag-ignore", "", nil, "List of column names that should have tags values set to '-' (ignored during parsing)")

	// hide flags not recommended for use
	rootCmd.PersistentFlags().MarkHidden("replace")

	viper.BindPFlags(rootCmd.PersistentFlags())
	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	viper.AutomaticEnv()

	if err := rootCmd.Execute(); err != nil {
		if e, ok := err.(commandFailure); ok {
			fmt.Printf("Error: %v\n\n", string(e))
			rootCmd.Help()
		} else if !viper.GetBool("debug") {
			fmt.Printf("Error: %v\n", err)
		} else {
			fmt.Printf("Error: %+v\n", err)
		}

		os.Exit(1)
	}
}

type commandFailure string

func (c commandFailure) Error() string {
	return string(c)
}

func preRun(cmd *cobra.Command, args []string) error {
	var err error

	if len(args) == 0 {
		return commandFailure("must provide a driver name")
	}

	driverName := args[0]
	driverPath := args[0]

	if strings.ContainsRune(driverName, os.PathSeparator) {
		driverName = strings.Replace(filepath.Base(driverName), "sqlboiler-", "", 1)
		driverName = strings.Replace(driverName, ".exe", "", 1)
	} else {
		driverPath = "sqlboiler-" + driverPath
		if p, err := exec.LookPath(driverPath); err == nil {
			driverPath = p
		}
	}

	driverPath, err = filepath.Abs(driverPath)
	if err != nil {
		return errors.Wrap(err, "could not find absolute path to driver")
	}
	drivers.RegisterBinary(driverName, driverPath)

	cmdConfig = &boilingcore.Config{
		DriverName:        driverName,
		OutFolder:         viper.GetString("output"),
		PkgName:           viper.GetString("pkgname"),
		Debug:             viper.GetBool("debug"),
		AddGlobal:         viper.GetBool("add-global-variants"),
		AddPanic:          viper.GetBool("add-panic-variants"),
		AddSoftDeletes:    viper.GetBool("add-soft-deletes"),
		NoContext:         viper.GetBool("no-context"),
		NoTests:           viper.GetBool("no-tests"),
		NoHooks:           viper.GetBool("no-hooks"),
		NoRowsAffected:    viper.GetBool("no-rows-affected"),
		NoAutoTimestamps:  viper.GetBool("no-auto-timestamps"),
		NoDriverTemplates: viper.GetBool("no-driver-templates"),
		NoBackReferencing: viper.GetBool("no-back-referencing"),
		Wipe:              viper.GetBool("wipe"),
		StructTagCasing:   strings.ToLower(viper.GetString("struct-tag-casing")), // camel | snake | title
		TagIgnore:         viper.GetStringSlice("tag-ignore"),
		RelationTag:       viper.GetString("relation-tag"),
		TemplateDirs:      viper.GetStringSlice("templates"),
		Tags:              viper.GetStringSlice("tag"),
		Replacements:      viper.GetStringSlice("replace"),
		Aliases:           boilingcore.ConvertAliases(viper.Get("aliases")),
		TypeReplaces:      boilingcore.ConvertTypeReplace(viper.Get("types")),
		Version:           sqlBoilerVersion,
	}

	if cmdConfig.Debug {
		fmt.Fprintln(os.Stderr, "using driver:", driverPath)
	}

	// Configure the driver
	cmdConfig.DriverConfig = map[string]interface{}{
		"whitelist": viper.GetStringSlice(driverName + ".whitelist"),
		"blacklist": viper.GetStringSlice(driverName + ".blacklist"),
	}

	keys := allKeys(driverName)
	for _, key := range keys {
		prefixedKey := fmt.Sprintf("%s.%s", driverName, key)
		cmdConfig.DriverConfig[key] = viper.Get(prefixedKey)
	}

	cmdConfig.Imports = configureImports()

	cmdState, err = boilingcore.New(cmdConfig)
	return err
}

func configureImports() importers.Collection {
	imports := importers.NewDefaultImports()

	mustMap := func(m importers.Map, err error) importers.Map {
		if err != nil {
			panic("failed to change viper interface into importers.Map: " + err.Error())
		}

		return m
	}

	if viper.IsSet("imports.all.standard") {
		imports.All.Standard = viper.GetStringSlice("imports.all.standard")
	}
	if viper.IsSet("imports.all.third_party") {
		imports.All.ThirdParty = viper.GetStringSlice("imports.all.third_party")
	}
	if viper.IsSet("imports.test.standard") {
		imports.Test.Standard = viper.GetStringSlice("imports.test.standard")
	}
	if viper.IsSet("imports.test.third_party") {
		imports.Test.ThirdParty = viper.GetStringSlice("imports.test.third_party")
	}
	if viper.IsSet("imports.singleton") {
		imports.Singleton = mustMap(importers.MapFromInterface(viper.Get("imports.singleton")))
	}
	if viper.IsSet("imports.test_singleton") {
		imports.TestSingleton = mustMap(importers.MapFromInterface(viper.Get("imports.test_singleton")))
	}
	if viper.IsSet("imports.based_on_type") {
		imports.BasedOnType = mustMap(importers.MapFromInterface(viper.Get("imports.based_on_type")))
	}

	return imports
}

func run(cmd *cobra.Command, args []string) error {
	return cmdState.Run()
}

func postRun(cmd *cobra.Command, args []string) error {
	return cmdState.Cleanup()
}

func allKeys(prefix string) []string {
	keys := make(map[string]bool)

	prefix += "."

	for _, e := range os.Environ() {
		splits := strings.SplitN(e, "=", 2)
		key := strings.ReplaceAll(strings.ToLower(splits[0]), "_", ".")

		if strings.HasPrefix(key, prefix) {
			keys[strings.ReplaceAll(key, prefix, "")] = true
		}
	}

	for _, key := range viper.AllKeys() {
		if strings.HasPrefix(key, prefix) {
			keys[strings.ReplaceAll(key, prefix, "")] = true
		}
	}

	keySlice := make([]string, 0, len(keys))
	for k := range keys {
		keySlice = append(keySlice, k)
	}
	return keySlice
}
<commit_msg>Prevent blacklist and whitelist set via ENVs from being converted into string from stringSlice<commit_after>// Package main defines a command line interface for the sqlboiler package
package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/friendsofgo/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/volatiletech/sqlboiler/v4/boilingcore"
	"github.com/volatiletech/sqlboiler/v4/drivers"
	"github.com/volatiletech/sqlboiler/v4/importers"
)

//go:generate go-bindata -nometadata -pkg templatebin -o templatebin/bindata.go templates templates/singleton templates_test templates_test/singleton

const sqlBoilerVersion = "4.2.0"

var (
	flagConfigFile string
	cmdState       *boilingcore.State
	cmdConfig      *boilingcore.Config
)

func initConfig() {
	if len(flagConfigFile) != 0 {
		viper.SetConfigFile(flagConfigFile)
		if err := viper.ReadInConfig(); err != nil {
			fmt.Println("Can't read config:", err)
			os.Exit(1)
		}
		return
	}

	var err error
	viper.SetConfigName("sqlboiler")

	configHome := os.Getenv("XDG_CONFIG_HOME")
	homePath := os.Getenv("HOME")
	wd, err := os.Getwd()
	if err != nil {
		wd = "."
	}

	configPaths := []string{wd}
	if len(configHome) > 0 {
		configPaths = append(configPaths, filepath.Join(configHome, "sqlboiler"))
	} else {
		configPaths = append(configPaths, filepath.Join(homePath, ".config/sqlboiler"))
	}

	for _, p := range configPaths {
		viper.AddConfigPath(p)
	}

	// Ignore errors here, fallback to other validation methods.
	// Users can use environment variables if a config is not found.
	_ = viper.ReadInConfig()
}

func main() {
	// Too much happens between here and cobra's argument handling, for
	// something so simple just do it immediately.
	for _, arg := range os.Args {
		if arg == "--version" {
			fmt.Println("SQLBoiler v" + sqlBoilerVersion)
			return
		}
	}

	// Set up the cobra root command
	var rootCmd = &cobra.Command{
		Use:   "sqlboiler [flags] <driver>",
		Short: "SQL Boiler generates an ORM tailored to your database schema.",
		Long: "SQL Boiler generates a Go ORM from template files, tailored to your database schema.\n" +
			`Complete documentation is available at http://github.com/volatiletech/sqlboiler`,
		Example:       `sqlboiler psql`,
		PreRunE:       preRun,
		RunE:          run,
		PostRunE:      postRun,
		SilenceErrors: true,
		SilenceUsage:  true,
	}

	cobra.OnInitialize(initConfig)

	// Set up the cobra root command flags
	rootCmd.PersistentFlags().StringVarP(&flagConfigFile, "config", "c", "", "Filename of config file to override default lookup")
	rootCmd.PersistentFlags().StringP("output", "o", "models", "The name of the folder to output to")
	rootCmd.PersistentFlags().StringP("pkgname", "p", "models", "The name you wish to assign to your generated package")
	rootCmd.PersistentFlags().StringSliceP("templates", "", nil, "A templates directory, overrides the bindata'd template folders in sqlboiler")
	rootCmd.PersistentFlags().StringSliceP("tag", "t", nil, "Struct tags to be included on your models in addition to json, yaml, toml")
	rootCmd.PersistentFlags().StringSliceP("replace", "", nil, "Replace templates by directory: relpath/to_file.tpl:relpath/to_replacement.tpl")
	rootCmd.PersistentFlags().BoolP("debug", "d", false, "Debug mode prints stack traces on error")
	rootCmd.PersistentFlags().BoolP("no-context", "", false, "Disable context.Context usage in the generated code")
	rootCmd.PersistentFlags().BoolP("no-tests", "", false, "Disable generated go test files")
	rootCmd.PersistentFlags().BoolP("no-hooks", "", false, "Disable hooks feature for your models")
	rootCmd.PersistentFlags().BoolP("no-rows-affected", "", false, "Disable rows affected in the generated API")
	rootCmd.PersistentFlags().BoolP("no-auto-timestamps", "", false, "Disable automatic timestamps for created_at/updated_at")
	rootCmd.PersistentFlags().BoolP("no-driver-templates", "", false, "Disable parsing of templates defined by the database driver")
	rootCmd.PersistentFlags().BoolP("no-back-referencing", "", false, "Disable back referencing in the loaded relationship structs")
	rootCmd.PersistentFlags().BoolP("add-global-variants", "", false, "Enable generation for global variants")
	rootCmd.PersistentFlags().BoolP("add-panic-variants", "", false, "Enable generation for panic variants")
	rootCmd.PersistentFlags().BoolP("add-soft-deletes", "", false, "Enable soft deletion by updating deleted_at timestamp")
	rootCmd.PersistentFlags().BoolP("version", "", false, "Print the version")
	rootCmd.PersistentFlags().BoolP("wipe", "", false, "Delete the output folder (rm -rf) before generation to ensure sanity")
	rootCmd.PersistentFlags().StringP("struct-tag-casing", "", "snake", "Decides the casing for go structure tag names. camel, title or snake (default snake)")
	rootCmd.PersistentFlags().StringP("relation-tag", "r", "-", "Relationship struct tag name")
	rootCmd.PersistentFlags().StringSliceP("tag-ignore", "", nil, "List of column names that should have tags values set to '-' (ignored during parsing)")

	// hide flags not recommended for use
	rootCmd.PersistentFlags().MarkHidden("replace")

	viper.BindPFlags(rootCmd.PersistentFlags())
	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	viper.AutomaticEnv()

	if err := rootCmd.Execute(); err != nil {
		if e, ok := err.(commandFailure); ok {
			fmt.Printf("Error: %v\n\n", string(e))
			rootCmd.Help()
		} else if !viper.GetBool("debug") {
			fmt.Printf("Error: %v\n", err)
		} else {
			fmt.Printf("Error: %+v\n", err)
		}

		os.Exit(1)
	}
}

type commandFailure string

func (c commandFailure) Error() string {
	return string(c)
}

func preRun(cmd *cobra.Command, args []string) error {
	var err error

	if len(args) == 0 {
		return commandFailure("must provide a driver name")
	}

	driverName := args[0]
	driverPath := args[0]

	if strings.ContainsRune(driverName, os.PathSeparator) {
		driverName = strings.Replace(filepath.Base(driverName), "sqlboiler-", "", 1)
		driverName = strings.Replace(driverName, ".exe", "", 1)
	} else {
		driverPath = "sqlboiler-" + driverPath
		if p, err := exec.LookPath(driverPath); err == nil {
			driverPath = p
		}
	}

	driverPath, err = filepath.Abs(driverPath)
	if err != nil {
		return errors.Wrap(err, "could not find absolute path to driver")
	}
	drivers.RegisterBinary(driverName, driverPath)

	cmdConfig = &boilingcore.Config{
		DriverName:        driverName,
		OutFolder:         viper.GetString("output"),
		PkgName:           viper.GetString("pkgname"),
		Debug:             viper.GetBool("debug"),
		AddGlobal:         viper.GetBool("add-global-variants"),
		AddPanic:          viper.GetBool("add-panic-variants"),
		AddSoftDeletes:    viper.GetBool("add-soft-deletes"),
		NoContext:         viper.GetBool("no-context"),
		NoTests:           viper.GetBool("no-tests"),
		NoHooks:           viper.GetBool("no-hooks"),
		NoRowsAffected:    viper.GetBool("no-rows-affected"),
		NoAutoTimestamps:  viper.GetBool("no-auto-timestamps"),
		NoDriverTemplates: viper.GetBool("no-driver-templates"),
		NoBackReferencing: viper.GetBool("no-back-referencing"),
		Wipe:              viper.GetBool("wipe"),
		StructTagCasing:   strings.ToLower(viper.GetString("struct-tag-casing")), // camel | snake | title
		TagIgnore:         viper.GetStringSlice("tag-ignore"),
		RelationTag:       viper.GetString("relation-tag"),
		TemplateDirs:      viper.GetStringSlice("templates"),
		Tags:              viper.GetStringSlice("tag"),
		Replacements:      viper.GetStringSlice("replace"),
		Aliases:           boilingcore.ConvertAliases(viper.Get("aliases")),
		TypeReplaces:      boilingcore.ConvertTypeReplace(viper.Get("types")),
		Version:           sqlBoilerVersion,
	}

	if cmdConfig.Debug {
		fmt.Fprintln(os.Stderr, "using driver:", driverPath)
	}

	// Configure the driver
	cmdConfig.DriverConfig = map[string]interface{}{
		"whitelist": viper.GetStringSlice(driverName + ".whitelist"),
		"blacklist": viper.GetStringSlice(driverName + ".blacklist"),
	}

	keys := allKeys(driverName)
	for _, key := range keys {
		if key != "blacklist" && key != "whitelist" {
			prefixedKey := fmt.Sprintf("%s.%s", driverName, key)
			cmdConfig.DriverConfig[key] = viper.Get(prefixedKey)
		}
	}

	cmdConfig.Imports = configureImports()

	cmdState, err = boilingcore.New(cmdConfig)
	return err
}

func configureImports() importers.Collection {
	imports := importers.NewDefaultImports()

	mustMap := func(m importers.Map, err error) importers.Map {
		if err != nil {
			panic("failed to change viper interface into importers.Map: " + err.Error())
		}

		return m
	}

	if viper.IsSet("imports.all.standard") {
		imports.All.Standard = viper.GetStringSlice("imports.all.standard")
	}
	if viper.IsSet("imports.all.third_party") {
		imports.All.ThirdParty = viper.GetStringSlice("imports.all.third_party")
	}
	if viper.IsSet("imports.test.standard") {
		imports.Test.Standard = viper.GetStringSlice("imports.test.standard")
	}
	if viper.IsSet("imports.test.third_party") {
		imports.Test.ThirdParty = viper.GetStringSlice("imports.test.third_party")
	}
	if viper.IsSet("imports.singleton") {
		imports.Singleton = mustMap(importers.MapFromInterface(viper.Get("imports.singleton")))
	}
	if viper.IsSet("imports.test_singleton") {
		imports.TestSingleton = mustMap(importers.MapFromInterface(viper.Get("imports.test_singleton")))
	}
	if viper.IsSet("imports.based_on_type") {
		imports.BasedOnType = mustMap(importers.MapFromInterface(viper.Get("imports.based_on_type")))
	}

	return imports
}

func run(cmd *cobra.Command, args []string) error {
	return cmdState.Run()
}

func postRun(cmd *cobra.Command, args []string) error {
	return cmdState.Cleanup()
}

func allKeys(prefix string) []string {
	keys := make(map[string]bool)

	prefix += "."

	for _, e := range os.Environ() {
		splits := strings.SplitN(e, "=", 2)
		key := strings.ReplaceAll(strings.ToLower(splits[0]), "_", ".")

		if strings.HasPrefix(key, prefix) {
			keys[strings.ReplaceAll(key, prefix, "")] = true
		}
	}

	for _, key := range viper.AllKeys() {
		if strings.HasPrefix(key, prefix) {
			keys[strings.ReplaceAll(key, prefix, "")] = true
		}
	}

	keySlice := make([]string, 0, len(keys))
	for k := range keys {
		keySlice = append(keySlice, k)
	}
	return keySlice
}
<|endoftext|>