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

// Package cmd defines and implements command-line commands and flags
// used by probablyfine. Commands and flags are implemented using Cobra.
package cmd

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// webdir is the path to the source files for the Hugo website
const webdir = "/Users/mattstratton/src/devopsdays-web" // TODO: Change this to read an environment variable, and default to cwd if envar not set

var cfgFile string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "probablyfine",
	Short: "Run maintenance tasks for the devopsdays.org website",
	Long: `Command-line utilities for the devopsdays.org website
built with love by mattstratton in Go.

Complete documentation is available at https://github.com/mattstratton/probablyfine`,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	//	Run: func(cmd *cobra.Command, args []string) { },
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	// Here you will define your flags and configuration settings.
	// Cobra supports Persistent Flags, which, if defined here,
	// will be global for your application.

	// RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.probablyfine.yaml)")
	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	// RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" { // enable ability to specify config file via flag
		viper.SetConfigFile(cfgFile)
	}

	viper.SetConfigName(".probablyfine") // name of config file (without extension)
	viper.AddConfigPath("$HOME")         // adding home directory as first search path
	viper.AutomaticEnv()                 // read in environment variables that match

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Println("Using config file:", viper.ConfigFileUsed())
	}
}
<commit_msg>Add setWebdir function<commit_after>// Copyright © 2016 Matt Stratton <matt.stratton@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package cmd defines and implements command-line commands and flags
// used by probablyfine. Commands and flags are implemented using Cobra.
package cmd

import (
	"fmt"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// webdir is the path to the source files for the Hugo website
var webdir string = setWebdir()

// const webdir = "/Users/mattstratton/src/devopsdays-web"

var cfgFile string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "probablyfine",
	Short: "Run maintenance tasks for the devopsdays.org website",
	Long: `Command-line utilities for the devopsdays.org website
built with love by mattstratton in Go.

Complete documentation is available at https://github.com/mattstratton/probablyfine`,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	//	Run: func(cmd *cobra.Command, args []string) { },
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	// Here you will define your flags and configuration settings.
	// Cobra supports Persistent Flags, which, if defined here,
	// will be global for your application.

	// RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.probablyfine.yaml)")
	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	// RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" { // enable ability to specify config file via flag
		viper.SetConfigFile(cfgFile)
	}

	viper.SetConfigName(".probablyfine") // name of config file (without extension)
	viper.AddConfigPath("$HOME")         // adding home directory as first search path
	viper.AutomaticEnv()                 // read in environment variables that match

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Println("Using config file:", viper.ConfigFileUsed())
	}
}

func setWebdir() string {
	if os.Getenv("DODPATH") == "" {
		pwd, err := os.Getwd()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		return pwd
	} else {
		s := os.Getenv("DODPATH")
		s = strings.TrimSuffix(s, "/")
		s = strings.TrimSuffix(s, "\\")
		return s
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"text/template"

	"github.com/rsteube/carapace"
	"github.com/spf13/cobra"
	gitconfig "github.com/tcnksm/go-gitconfig"
	"github.com/zaquestion/lab/internal/git"
	lab "github.com/zaquestion/lab/internal/gitlab"
	"github.com/zaquestion/lab/internal/logger"
)

// Get internal lab logger instance
var log = logger.GetInstance()

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "lab",
	Short: "lab: A GitLab Command Line Interface Utility",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		if ok, err := cmd.Flags().GetBool("version"); err == nil && ok {
			versionCmd.Run(cmd, args)
			return
		}
		helpCmd.Run(cmd, args)
	},
}

func rpad(s string, padding int) string {
	template := fmt.Sprintf("%%-%ds", padding)
	return fmt.Sprintf(template, s)
}

var templateFuncs = template.FuncMap{
	"rpad": rpad,
}

const labUsageTmpl = `{{range .Commands}}{{if (and (not .Hidden) (or .IsAvailableCommand (ne .Name "help")) (and (ne .Name "clone") (ne .Name "version") (ne .Name "merge-request")))}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}`

func labUsageFormat(c *cobra.Command) string {
	t := template.New("top")
	t.Funcs(templateFuncs)
	template.Must(t.Parse(labUsageTmpl))

	var buf bytes.Buffer
	err := t.Execute(&buf, c)
	if err != nil {
		c.Println(err)
	}
	return buf.String()
}

func helpFunc(cmd *cobra.Command, args []string) {
	// When help func is called from the help command args will be
	// populated. When help is called with cmd.Help(), the args are not
	// passed through, so we pick them up ourselves here
	if len(args) == 0 {
		args = os.Args[1:]
	}
	rootCmd := cmd.Root()
	// Show help for sub/commands -- any commands that isn't "lab" or "help"
	if cmd, _, err := rootCmd.Find(args); err == nil {
		// Cobra will check parent commands for a helpFunc and we only
		// want the root command to actually use this custom help func.
		// Here we trick cobra into thinking that there is no help func
		// so it will use the default help for the subcommands
		cmd.Root().SetHelpFunc(nil)
		err2 := cmd.Help()
		if err2 != nil {
			log.Fatal(err)
		}
		return
	}
}

var helpCmd = &cobra.Command{
	Use:   "help [command [subcommand...]]",
	Short: "Show the help for lab",
	Long:  ``,
	Run:   helpFunc,
}

// Version is set with linker flags during build.
var Version string

// versionCmd represents the version command
var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Printf("%s %s\n", "lab version", Version)
	},
}

func init() {
	// NOTE: Calling SetHelpCommand like this causes helpFunc to be called
	// with correct arguments. If the default cobra help func is used no
	// arguments are passed through and subcommand help breaks.
	RootCmd.SetHelpCommand(helpCmd)
	RootCmd.SetHelpFunc(helpFunc)
	RootCmd.AddCommand(versionCmd)
	RootCmd.Flags().Bool("version", false, "Show the lab version")
	RootCmd.PersistentFlags().Bool("no-pager", false, "Do not pipe output into a pager")
	RootCmd.PersistentFlags().Bool("debug", false, "Enable debug logging level")
	RootCmd.PersistentFlags().Bool("quiet", false, "Turn off any sort of logging. Only command output is printed")

	// We need to set the logger level before any other piece of code is
	// called, thus we make sure we don't lose any debug message, but for
	// that we need to parse the args from command input.
	err := RootCmd.ParseFlags(os.Args[1:])
	// Handle the err != nil case later
	if err == nil {
		debugLogger, _ := RootCmd.Flags().GetBool("debug")
		quietLogger, _ := RootCmd.Flags().GetBool("quiet")
		if debugLogger && quietLogger {
			log.Fatal("option --debug cannot be combined with --quiet")
		}
		if debugLogger {
			log.SetLogLevel(logger.LogLevelDebug)
		} else if quietLogger {
			log.SetLogLevel(logger.LogLevelNone)
		}
	}
	carapace.Gen(RootCmd)
}

var (
	// Will be updated to upstream in Execute() if "upstream" remote exists
	defaultRemote = ""
	// Will be updated to lab.User() in Execute() if forkedFrom is "origin"
	forkRemote = ""
)

// Try to guess what should be the default remote.
func guessDefaultRemote() string {
	// Allow to force a default remote. If set, return early.
	if config := getMainConfig(); config != nil {
		defaultRemote := config.GetString("core.default_remote")
		if defaultRemote != "" {
			return defaultRemote
		}
	}

	guess := ""

	// defaultRemote should try to always point to the upstream project.
	// Since "origin" may have two different meanings depending on how the
	// user forked the project, thus make "upstream" as the most significant
	// remote.
	// In forkFromOrigin approach, "origin" remote is the one pointing to
	// the upstream project.
	_, err := gitconfig.Local("remote.origin.url")
	if err == nil {
		guess = "origin"
	}
	// In forkToUpstream approach, "upstream" remote is the one pointing to
	// the upstream project
	_, err = gitconfig.Local("remote.upstream.url")
	if err == nil {
		guess = "upstream"
	}

	// But it's still possible the user used a custom name
	if guess == "" {
		// use the remote tracked by the default branch if set
		if remote, err := gitconfig.Local("branch.main.remote"); err == nil {
			guess = remote
		} else if remote, err = gitconfig.Local("branch.master.remote"); err == nil {
			guess = remote
		} else {
			// use the first remote added to .git/config file, which, usually, is
			// the one from which the repo was clonned
			remotesStr, err := git.GetLocalRemotesFromFile()
			if err == nil {
				remotes := strings.Split(remotesStr, "\n")
				// remotes format: remote.<name>.<url|fetch>
				remoteName := strings.Split(remotes[0], ".")[1]
				guess = remoteName
			}
		}
	}

	return guess
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute(initSkipped bool) {
	// Try to gather remote information if running inside a git tree/repo.
	// Otherwise, skip it, since the info won't be used at all, also avoiding
	// misleading error/warning messages about missing remote.
	if git.InsideGitRepo() {
		defaultRemote = guessDefaultRemote()
		if defaultRemote == "" {
			log.Infoln("No default remote found")
		}

		// Check if the user fork exists
		_, err := gitconfig.Local("remote." + lab.User() + ".url")
		if err == nil {
			forkRemote = lab.User()
		} else {
			forkRemote = defaultRemote
		}
	}

	// Check if the user is calling a lab command or if we should passthrough
	// NOTE: The help command won't be found by Find, which we are counting on
	cmd, _, err := RootCmd.Find(os.Args[1:])
	if err != nil || cmd.Use == "clone" {
		// Determine if any undefined flags were passed to "clone"
		// TODO: Evaluate and support some of these flags
		if (cmd.Use == "clone" && len(os.Args) > 2) || os.Args[1] == "help" {
			// ParseFlags will err in these cases
			err = cmd.ParseFlags(os.Args[1:])
			if err == nil {
				if err := RootCmd.Execute(); err != nil {
					// Execute has already logged the error
					os.Exit(1)
				}
				return
			}
		}

		// Passthrough to git for any unrecognized commands
		log.Warnln("lab's git passthrough command support will be removed in a later release.")
		err = git.New(os.Args[1:]...).Run()
		if exiterr, ok := err.(*exec.ExitError); ok {
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				os.Exit(status.ExitStatus())
			}
		}
		if err != nil {
			log.Fatal(err)
		}
		return
	}

	// allow flags to the root cmd to be passed through.
	// Error msgs and codes generated by git are thrown away and the execution
	// flow returns to lab to generate the error.
	// TODO: remove for 1.0 when we stop wrapping git
	if cmd.Use == RootCmd.Use && len(os.Args) > 1 {
		// Pass unknown flags to git, if it also doesn't handle it, let lab
		// handle the exit msg (help) and code.
		if !initSkipped {
			log.Warnln("lab's git passthrough command support will be removed in a later release.")
			gitCmd := git.New(os.Args[1:]...)
			gitCmd.Stderr = nil
			gitCmd.Stdout = nil
			out, err := gitCmd.CombinedOutput()
			if err == nil {
				fmt.Println(string(out))
				return
			}
		}
	}

	// Set commandPrefix
	scmd, _, _ := cmd.Find(os.Args)
	setCommandPrefix(scmd)

	if err := RootCmd.Execute(); err != nil {
		// Execute has already logged the error
		os.Exit(1)
	}
}
<commit_msg>cmd/root: fix setting log level<commit_after>package cmd

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"text/template"

	"github.com/rsteube/carapace"
	"github.com/spf13/cobra"
	gitconfig "github.com/tcnksm/go-gitconfig"
	"github.com/zaquestion/lab/internal/git"
	lab "github.com/zaquestion/lab/internal/gitlab"
	"github.com/zaquestion/lab/internal/logger"
)

// Get internal lab logger instance
var log = logger.GetInstance()

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "lab",
	Short: "lab: A GitLab Command Line Interface Utility",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		if ok, err := cmd.Flags().GetBool("version"); err == nil && ok {
			versionCmd.Run(cmd, args)
			return
		}
		helpCmd.Run(cmd, args)
	},
}

func rpad(s string, padding int) string {
	template := fmt.Sprintf("%%-%ds", padding)
	return fmt.Sprintf(template, s)
}

var templateFuncs = template.FuncMap{
	"rpad": rpad,
}

const labUsageTmpl = `{{range .Commands}}{{if (and (not .Hidden) (or .IsAvailableCommand (ne .Name "help")) (and (ne .Name "clone") (ne .Name "version") (ne .Name "merge-request")))}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}`

func labUsageFormat(c *cobra.Command) string {
	t := template.New("top")
	t.Funcs(templateFuncs)
	template.Must(t.Parse(labUsageTmpl))

	var buf bytes.Buffer
	err := t.Execute(&buf, c)
	if err != nil {
		c.Println(err)
	}
	return buf.String()
}

func helpFunc(cmd *cobra.Command, args []string) {
	// When help func is called from the help command args will be
	// populated. When help is called with cmd.Help(), the args are not
	// passed through, so we pick them up ourselves here
	if len(args) == 0 {
		args = os.Args[1:]
	}
	rootCmd := cmd.Root()
	// Show help for sub/commands -- any commands that isn't "lab" or "help"
	if cmd, _, err := rootCmd.Find(args); err == nil {
		// Cobra will check parent commands for a helpFunc and we only
		// want the root command to actually use this custom help func.
		// Here we trick cobra into thinking that there is no help func
		// so it will use the default help for the subcommands
		cmd.Root().SetHelpFunc(nil)
		err2 := cmd.Help()
		if err2 != nil {
			log.Fatal(err)
		}
		return
	}
}

var helpCmd = &cobra.Command{
	Use:   "help [command [subcommand...]]",
	Short: "Show the help for lab",
	Long:  ``,
	Run:   helpFunc,
}

// Version is set with linker flags during build.
var Version string

// versionCmd represents the version command
var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		fmt.Printf("%s %s\n", "lab version", Version)
	},
}

func init() {
	// NOTE: Calling SetHelpCommand like this causes helpFunc to be called
	// with correct arguments. If the default cobra help func is used no
	// arguments are passed through and subcommand help breaks.
	RootCmd.SetHelpCommand(helpCmd)
	RootCmd.SetHelpFunc(helpFunc)
	RootCmd.AddCommand(versionCmd)
	RootCmd.Flags().Bool("version", false, "Show the lab version")
	RootCmd.PersistentFlags().Bool("no-pager", false, "Do not pipe output into a pager")
	RootCmd.PersistentFlags().Bool("debug", false, "Enable debug logging level")
	RootCmd.PersistentFlags().Bool("quiet", false, "Turn off any sort of logging. Only command output is printed")

	// We need to set the logger level before any other piece of code is
	// called, thus we make sure we don't lose any debug message, but for
	// that we need to parse the args from command input and let flag errors be
	// handled by the subcommands themselves.
	_ = RootCmd.ParseFlags(os.Args[1:])
	debugLogger, _ := RootCmd.Flags().GetBool("debug")
	quietLogger, _ := RootCmd.Flags().GetBool("quiet")
	if debugLogger && quietLogger {
		log.Fatal("option --debug cannot be combined with --quiet")
	}
	if debugLogger {
		log.SetLogLevel(logger.LogLevelDebug)
	} else if quietLogger {
		log.SetLogLevel(logger.LogLevelNone)
	}
	carapace.Gen(RootCmd)
}

var (
	// Will be updated to upstream in Execute() if "upstream" remote exists
	defaultRemote = ""
	// Will be updated to lab.User() in Execute() if forkedFrom is "origin"
	forkRemote = ""
)

// Try to guess what should be the default remote.
func guessDefaultRemote() string {
	// Allow to force a default remote. If set, return early.
	if config := getMainConfig(); config != nil {
		defaultRemote := config.GetString("core.default_remote")
		if defaultRemote != "" {
			return defaultRemote
		}
	}

	guess := ""

	// defaultRemote should try to always point to the upstream project.
	// Since "origin" may have two different meanings depending on how the
	// user forked the project, thus make "upstream" as the most significant
	// remote.
	// In forkFromOrigin approach, "origin" remote is the one pointing to
	// the upstream project.
	_, err := gitconfig.Local("remote.origin.url")
	if err == nil {
		guess = "origin"
	}
	// In forkToUpstream approach, "upstream" remote is the one pointing to
	// the upstream project
	_, err = gitconfig.Local("remote.upstream.url")
	if err == nil {
		guess = "upstream"
	}

	// But it's still possible the user used a custom name
	if guess == "" {
		// use the remote tracked by the default branch if set
		if remote, err := gitconfig.Local("branch.main.remote"); err == nil {
			guess = remote
		} else if remote, err = gitconfig.Local("branch.master.remote"); err == nil {
			guess = remote
		} else {
			// use the first remote added to .git/config file, which, usually, is
			// the one from which the repo was clonned
			remotesStr, err := git.GetLocalRemotesFromFile()
			if err == nil {
				remotes := strings.Split(remotesStr, "\n")
				// remotes format: remote.<name>.<url|fetch>
				remoteName := strings.Split(remotes[0], ".")[1]
				guess = remoteName
			}
		}
	}

	return guess
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute(initSkipped bool) {
	// Try to gather remote information if running inside a git tree/repo.
	// Otherwise, skip it, since the info won't be used at all, also avoiding
	// misleading error/warning messages about missing remote.
	if git.InsideGitRepo() {
		defaultRemote = guessDefaultRemote()
		if defaultRemote == "" {
			log.Infoln("No default remote found")
		}

		// Check if the user fork exists
		_, err := gitconfig.Local("remote." + lab.User() + ".url")
		if err == nil {
			forkRemote = lab.User()
		} else {
			forkRemote = defaultRemote
		}
	}

	// Check if the user is calling a lab command or if we should passthrough
	// NOTE: The help command won't be found by Find, which we are counting on
	cmd, _, err := RootCmd.Find(os.Args[1:])
	if err != nil || cmd.Use == "clone" {
		// Determine if any undefined flags were passed to "clone"
		// TODO: Evaluate and support some of these flags
		if (cmd.Use == "clone" && len(os.Args) > 2) || os.Args[1] == "help" {
			// ParseFlags will err in these cases
			err = cmd.ParseFlags(os.Args[1:])
			if err == nil {
				if err := RootCmd.Execute(); err != nil {
					// Execute has already logged the error
					os.Exit(1)
				}
				return
			}
		}

		// Passthrough to git for any unrecognized commands
		log.Warnln("lab's git passthrough command support will be removed in a later release.")
		err = git.New(os.Args[1:]...).Run()
		if exiterr, ok := err.(*exec.ExitError); ok {
			if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
				os.Exit(status.ExitStatus())
			}
		}
		if err != nil {
			log.Fatal(err)
		}
		return
	}

	// allow flags to the root cmd to be passed through.
	// Error msgs and codes generated by git are thrown away and the execution
	// flow returns to lab to generate the error.
	// TODO: remove for 1.0 when we stop wrapping git
	if cmd.Use == RootCmd.Use && len(os.Args) > 1 {
		// Pass unknown flags to git, if it also doesn't handle it, let lab
		// handle the exit msg (help) and code.
		if !initSkipped {
			log.Warnln("lab's git passthrough command support will be removed in a later release.")
			gitCmd := git.New(os.Args[1:]...)
			gitCmd.Stderr = nil
			gitCmd.Stdout = nil
			out, err := gitCmd.CombinedOutput()
			if err == nil {
				fmt.Println(string(out))
				return
			}
		}
	}

	// Set commandPrefix
	scmd, _, _ := cmd.Find(os.Args)
	setCommandPrefix(scmd)

	if err := RootCmd.Execute(); err != nil {
		// Execute has already logged the error
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>// Copyright © 2017 Josh Dvir <josh@dvir.uk>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"sync"
	"time"

	"strings"

	cron "github.com/robfig/cron"
	"github.com/spf13/cobra"
	elastic "gopkg.in/olivere/elastic.v5"
)

var (
	olderThanInDays int
	esURL           string
	wg              sync.WaitGroup
	ctx             context.Context
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "delete-aws-es-incidents",
	Short: "Delete ELK incidents on AWS ES 5.1",
	Long:  "",
	Run: func(cmd *cobra.Command, args []string) {
		if esURL == "" {
			println("No Elasticsearch URL present, can't continue.")
			os.Exit(0)
		}

		var wgm sync.WaitGroup
		cron := cron.New()
		cron.AddFunc("@hourly", func() { runCommand() })
		cron.Start()
		println("Cron run started...")
		wgm.Add(1)
		wgm.Wait()
	},
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	RootCmd.Flags().IntVarP(&olderThanInDays, "older-than-in-days", "d", 14, "delete incidents older then in days")
	RootCmd.Flags().StringVarP(&esURL, "es-url", "e", "", "Elasticsearch URL, eg. https://path-to-es.aws.com/")
}

func runCommand() {
	println("Starting deleting incidents run...")
	ctx = context.Background()
	client, err := elastic.NewClient(
		elastic.SetURL(esURL),
		elastic.SetSniff(false),
	)
	if err != nil {
		panic(err)
	}

	indexNames, err := client.IndexNames()
	if err != nil {
		panic(err)
	}

	for _, indexName := range indexNames {
		if strings.HasPrefix(indexName, "logstash") {
			date := strings.TrimPrefix(indexName, "logstash-")
			dateArr := strings.Split(date, ".")
			nowTime := time.Now()
			indexYear, _ := strconv.Atoi(dateArr[0])
			indexMonth, _ := strconv.Atoi(dateArr[1])
			indexDay, _ := strconv.Atoi(dateArr[2])
			incidentTime := time.Date(indexYear, time.Month(indexMonth), indexDay, 0, 0, 0, 0, nowTime.Location())
			if daysDiff(nowTime, incidentTime) > olderThanInDays {
				wg.Add(1)
				go deleteIncident(ctx, client, indexName)
			}
		}
	}

	wg.Wait()
	println("Ending deleting incidents run...")
}

func deleteIncident(ctx context.Context, client *elastic.Client, indexName string) {
	deleteIndex, err := client.DeleteIndex(indexName).Do(ctx)
	if err != nil {
		fmt.Printf("Error deleting index %s\n", indexName)
	}

	if deleteIndex.Acknowledged {
		fmt.Printf("index %s deleted.\n", indexName)
	}

	defer wg.Done()
}

func lastDayOfYear(t time.Time) time.Time {
	return time.Date(t.Year(), 12, 31, 0, 0, 0, 0, t.Location())
}

func firstDayOfNextYear(t time.Time) time.Time {
	return time.Date(t.Year()+1, 1, 1, 0, 0, 0, 0, t.Location())
}

// a - b in days
func daysDiff(a, b time.Time) (days int) {
	cur := b
	for cur.Year() < a.Year() {
		// add 1 to count the last day of the year too.
		days += lastDayOfYear(cur).YearDay() - cur.YearDay() + 1
		cur = firstDayOfNextYear(cur)
	}
	days += a.YearDay() - cur.YearDay()
	if b.AddDate(0, 0, days).After(a) {
		days--
	}
	return days
}
<commit_msg>checking we can reach ES<commit_after>// Copyright © 2017 Josh Dvir <josh@dvir.uk>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"sync"
	"time"

	"strings"

	cron "github.com/robfig/cron"
	"github.com/spf13/cobra"
	elastic "gopkg.in/olivere/elastic.v5"
)

var (
	olderThanInDays int
	esURL           string
	wg              sync.WaitGroup
	ctx             context.Context
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "delete-aws-es-incidents",
	Short: "Delete ELK incidents on AWS ES 5.1",
	Long:  "",
	Run: func(cmd *cobra.Command, args []string) {
		if esURL == "" {
			println("No Elasticsearch URL present, can't continue.")
			os.Exit(0)
		}

		client, err := elastic.NewClient(
			elastic.SetURL(esURL),
			elastic.SetSniff(false),
		)
		if err != nil {
			panic(err)
		}

		var wgm sync.WaitGroup
		cron := cron.New()
		cron.AddFunc("@hourly", func() { runCommand() })
		cron.Start()
		println("Cron run started...")
		wgm.Add(1)
		wgm.Wait()
	},
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	RootCmd.Flags().IntVarP(&olderThanInDays, "older-than-in-days", "d", 14, "delete incidents older then in days")
	RootCmd.Flags().StringVarP(&esURL, "es-url", "e", "", "Elasticsearch URL, eg. https://path-to-es.aws.com/")
}

func runCommand() {
	println("Starting deleting incidents run...")
	ctx = context.Background()
	client, err := elastic.NewClient(
		elastic.SetURL(esURL),
		elastic.SetSniff(false),
	)
	if err != nil {
		panic(err)
	}

	indexNames, err := client.IndexNames()
	if err != nil {
		panic(err)
	}

	for _, indexName := range indexNames {
		if strings.HasPrefix(indexName, "logstash") {
			date := strings.TrimPrefix(indexName, "logstash-")
			dateArr := strings.Split(date, ".")
			nowTime := time.Now()
			indexYear, _ := strconv.Atoi(dateArr[0])
			indexMonth, _ := strconv.Atoi(dateArr[1])
			indexDay, _ := strconv.Atoi(dateArr[2])
			incidentTime := time.Date(indexYear, time.Month(indexMonth), indexDay, 0, 0, 0, 0, nowTime.Location())
			if daysDiff(nowTime, incidentTime) > olderThanInDays {
				wg.Add(1)
				go deleteIncident(ctx, client, indexName)
			}
		}
	}

	wg.Wait()
	println("Ending deleting incidents run...")
}

func deleteIncident(ctx context.Context, client *elastic.Client, indexName string) {
	deleteIndex, err := client.DeleteIndex(indexName).Do(ctx)
	if err != nil {
		fmt.Printf("Error deleting index %s\n", indexName)
	}

	if deleteIndex.Acknowledged {
		fmt.Printf("index %s deleted.\n", indexName)
	}

	defer wg.Done()
}

func lastDayOfYear(t time.Time) time.Time {
	return time.Date(t.Year(), 12, 31, 0, 0, 0, 0, t.Location())
}

func firstDayOfNextYear(t time.Time) time.Time {
	return time.Date(t.Year()+1, 1, 1, 0, 0, 0, 0, t.Location())
}

// a - b in days
func daysDiff(a, b time.Time) (days int) {
	cur := b
	for cur.Year() < a.Year() {
		// add 1 to count the last day of the year too.
		days += lastDayOfYear(cur).YearDay() - cur.YearDay() + 1
		cur = firstDayOfNextYear(cur)
	}
	days += a.YearDay() - cur.YearDay()
	if b.AddDate(0, 0, days).After(a) {
		days--
	}
	return days
}
<|endoftext|>


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

package cmd

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/minio/mc/pkg/console"
	"github.com/minio/mc/pkg/probe"
)

// contentMessage container for content message structure.
type statMessage struct {
	Status            string            `json:"status"`
	Key               string            `json:"name"`
	Date              time.Time         `json:"lastModified"`
	Size              int64             `json:"size"`
	ETag              string            `json:"etag"`
	Type              string            `json:"type"`
	EncryptionHeaders map[string]string `json:"encryption,omitempty"`
	Metadata          map[string]string `json:"metadata"`
}

// String colorized string message.
func printStat(stat statMessage) {
	// Format properly for alignment based on maxKey length
	stat.Key = fmt.Sprintf("%-10s: %s", "Name", stat.Key)
	console.Println(console.Colorize("Name", stat.Key))
	console.Println(fmt.Sprintf("%-10s: %s ", "Date", stat.Date.Format(printDate)))
	console.Println(fmt.Sprintf("%-10s: %-6s ", "Size", humanize.IBytes(uint64(stat.Size))))
	if stat.ETag != "" {
		console.Println(fmt.Sprintf("%-10s: %s ", "ETag", stat.ETag))
	}
	console.Println(fmt.Sprintf("%-10s: %s ", "Type", stat.Type))

	var maxKey = 0
	for k := range stat.Metadata {
		if len(k) > maxKey {
			maxKey = len(k)
		}
	}
	if len(stat.Metadata) > 0 {
		console.Println(fmt.Sprintf("%-10s:", "Metadata"))
		for k, v := range stat.Metadata {
			console.Println(fmt.Sprintf("  %-*.*s: %s ", maxKey, maxKey, k, v))
		}
	}
	maxKey = 0
	for k := range stat.EncryptionHeaders {
		if len(k) > maxKey {
			maxKey = len(k)
		}
	}
	if len(stat.EncryptionHeaders) > 0 {
		console.Println(fmt.Sprintf("%-10s:", "Encrypted"))
		for k, v := range stat.EncryptionHeaders {
			console.Println(fmt.Sprintf("  %-*.*s: %s ", maxKey, maxKey, k, v))
		}
	}
	console.Println()
}

// JSON jsonified content message.
func (c statMessage) JSON() string {
	c.Status = "success"
	jsonMessageBytes, e := json.Marshal(c)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")

	return string(jsonMessageBytes)
}

// parseStat parses client Content container into statMessage struct.
func parseStat(targetAlias string, c *clientContent) statMessage {
	content := statMessage{}
	content.Date = c.Time.Local()
	// guess file type.
	content.Type = func() string {
		if c.Type.IsDir() {
			return "folder"
		}
		return "file"
	}()
	content.Size = c.Size
	content.Key = targetAlias + getKey(c)
	content.Metadata = c.Metadata
	content.ETag = strings.TrimPrefix(c.ETag, "\"")
	content.ETag = strings.TrimSuffix(content.ETag, "\"")

	content.EncryptionHeaders = c.EncryptionHeaders
	return content
}

// doStat - list all entities inside a folder.
func doStat(clnt Client, isRecursive bool, targetAlias, targetURL string) error {
	if !isRecursive || !isAliasURLDir(targetURL) {
		isIncomplete := false
		fetchMeta := true
		if st, err := clnt.Stat(isIncomplete, fetchMeta); err != nil {
			switch err.ToGoError().(type) {
			default:
				fatalIf(err.Trace(targetURL), "Unable to initialize `"+targetURL+"`.")
			}
		} else {
			st := parseStat(targetAlias, st)
			if !globalJSON {
				printStat(st)
			} else {
				console.Println(st.JSON())
			}
			return nil
		}
	}
	if !strings.HasSuffix(targetURL, string(clnt.GetURL().Separator)) {
		targetURL = targetURL + string(clnt.GetURL().Separator)
	}
	clnt, err := newClient(targetURL)
	fatalIf(err.Trace(targetURL), "Unable to initialize target `"+targetURL+"`.")

	prefixPath := clnt.GetURL().Path
	separator := string(clnt.GetURL().Separator)
	if !strings.HasSuffix(prefixPath, separator) {
		prefixPath = prefixPath[:strings.LastIndex(prefixPath, separator)+1]
	}
	var cErr error
	isIncomplete := false
	var prevStat *clientContent
	for content := range clnt.List(isRecursive, isIncomplete, DirNone) {
		if content.Err != nil {
			switch content.Err.ToGoError().(type) {
			// handle this specifically for filesystem related errors.
			case BrokenSymlink:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list broken link.")
				continue
			case TooManyLevelsSymlink:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list too many levels link.")
				continue
			case PathNotFound:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
				continue
			case PathInsufficientPermission:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
				continue
			case ObjectOnGlacier:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "")
				continue
			}
			errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
			cErr = exitStatus(globalErrorExitStatus) // Set the exit status.
			continue
		}
		// Convert any os specific delimiters to "/".
		contentURL := filepath.ToSlash(content.URL.Path)
		prefixPath = filepath.ToSlash(prefixPath)
		// Trim prefix path from the content path.
		contentURL = strings.TrimPrefix(contentURL, prefixPath)
		content.URL.Path = contentURL
		_, stat, err := url2StatWithMetadata(targetURL+getKey(content), true)
		if err != nil {
			continue
		}
		if stat != prevStat {
			st := parseStat(targetAlias, stat)
			if !globalJSON {
				printStat(st)
			} else {
				console.Println(st.JSON())
			}
		} else {
		}
		prevStat = stat
	}
	return cErr
}
<commit_msg>stat: fix output to show bucket in stat listing (#2340)<commit_after>/*
 * Minio Client (C) 2017 Minio, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cmd

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/minio/mc/pkg/console"
	"github.com/minio/mc/pkg/probe"
)

// contentMessage container for content message structure.
type statMessage struct {
	Status            string            `json:"status"`
	Key               string            `json:"name"`
	Date              time.Time         `json:"lastModified"`
	Size              int64             `json:"size"`
	ETag              string            `json:"etag"`
	Type              string            `json:"type"`
	EncryptionHeaders map[string]string `json:"encryption,omitempty"`
	Metadata          map[string]string `json:"metadata"`
}

// String colorized string message.
func printStat(stat statMessage) {
	// Format properly for alignment based on maxKey length
	stat.Key = fmt.Sprintf("%-10s: %s", "Name", stat.Key)
	console.Println(console.Colorize("Name", stat.Key))
	console.Println(fmt.Sprintf("%-10s: %s ", "Date", stat.Date.Format(printDate)))
	console.Println(fmt.Sprintf("%-10s: %-6s ", "Size", humanize.IBytes(uint64(stat.Size))))
	if stat.ETag != "" {
		console.Println(fmt.Sprintf("%-10s: %s ", "ETag", stat.ETag))
	}
	console.Println(fmt.Sprintf("%-10s: %s ", "Type", stat.Type))

	var maxKey = 0
	for k := range stat.Metadata {
		if len(k) > maxKey {
			maxKey = len(k)
		}
	}
	if len(stat.Metadata) > 0 {
		console.Println(fmt.Sprintf("%-10s:", "Metadata"))
		for k, v := range stat.Metadata {
			console.Println(fmt.Sprintf("  %-*.*s: %s ", maxKey, maxKey, k, v))
		}
	}
	maxKey = 0
	for k := range stat.EncryptionHeaders {
		if len(k) > maxKey {
			maxKey = len(k)
		}
	}
	if len(stat.EncryptionHeaders) > 0 {
		console.Println(fmt.Sprintf("%-10s:", "Encrypted"))
		for k, v := range stat.EncryptionHeaders {
			console.Println(fmt.Sprintf("  %-*.*s: %s ", maxKey, maxKey, k, v))
		}
	}
	console.Println()
}

// JSON jsonified content message.
func (c statMessage) JSON() string {
	c.Status = "success"
	jsonMessageBytes, e := json.Marshal(c)
	fatalIf(probe.NewError(e), "Unable to marshal into JSON.")

	return string(jsonMessageBytes)
}

// parseStat parses client Content container into statMessage struct.
func parseStat(targetAlias string, c *clientContent) statMessage {
	content := statMessage{}
	content.Date = c.Time.Local()
	// guess file type.
	content.Type = func() string {
		if c.Type.IsDir() {
			return "folder"
		}
		return "file"
	}()
	content.Size = c.Size
	content.Key = getKey(c)
	content.Metadata = c.Metadata
	content.ETag = strings.TrimPrefix(c.ETag, "\"")
	content.ETag = strings.TrimSuffix(content.ETag, "\"")

	content.EncryptionHeaders = c.EncryptionHeaders
	return content
}

// doStat - list all entities inside a folder.
func doStat(clnt Client, isRecursive bool, targetAlias, targetURL string) error {

	if !strings.HasSuffix(targetURL, string(clnt.GetURL().Separator)) {
		targetURL = targetURL + string(clnt.GetURL().Separator)
	}
	clnt, err := newClient(targetURL)
	fatalIf(err.Trace(targetURL), "Unable to initialize target `"+targetURL+"`.")

	prefixPath := clnt.GetURL().Path
	separator := string(clnt.GetURL().Separator)
	if !strings.HasSuffix(prefixPath, separator) {
		prefixPath = prefixPath[:strings.LastIndex(prefixPath, separator)+1]
	}
	var cErr error
	isIncomplete := false
	var prevStat *clientContent
	for content := range clnt.List(isRecursive, isIncomplete, DirNone) {
		if content.Err != nil {
			switch content.Err.ToGoError().(type) {
			// handle this specifically for filesystem related errors.
			case BrokenSymlink:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list broken link.")
				continue
			case TooManyLevelsSymlink:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list too many levels link.")
				continue
			case PathNotFound:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
				continue
			case PathInsufficientPermission:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
				continue
			case ObjectOnGlacier:
				errorIf(content.Err.Trace(clnt.GetURL().String()), "")
				continue
			}
			errorIf(content.Err.Trace(clnt.GetURL().String()), "Unable to list folder.")
			cErr = exitStatus(globalErrorExitStatus) // Set the exit status.
			continue
		}
		url := targetAlias + getKey(content)
		if !isAliasURLDir(targetURL) {
			url = content.URL.Path
		}
		_, stat, err := url2StatWithMetadata(url, true)
		if err != nil {
			stat = content
		}
		// Convert any os specific delimiters to "/".
		contentURL := filepath.ToSlash(stat.URL.Path)
		prefixPath = filepath.ToSlash(prefixPath)
		// Trim prefix path from the content path.
		contentURL = strings.TrimPrefix(contentURL, prefixPath)
		stat.URL.Path = contentURL
		if stat != prevStat {
			st := parseStat(targetAlias, stat)
			if !globalJSON {
				printStat(st)
			} else {
				console.Println(st.JSON())
			}
		} else {
		}
		prevStat = stat
	}
	return cErr
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/x509"
	"encoding/pem"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
)

var db = map[string]string{}
var csrLocation = "/var/csr/"
var crtLocation = "/var/crt/"
var confLocation = "/opt/pollendina/openssl-ca.cnf"

type Tuple struct{ CN, Token string }

var updates = make(chan Tuple)

const (
	PORT = ":33004"
)

var (
	Info    *log.Logger
	Warning *log.Logger
	Error   *log.Logger
)

func InitLogs(
	infoHandle io.Writer,
	warningHandle io.Writer,
	errorHandle io.Writer) {

	Info = log.New(infoHandle,
		"INFO: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Warning = log.New(warningHandle,
		"WARNING: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Error = log.New(errorHandle,
		"ERROR: ",
		log.Ldate|log.Ltime|log.Lshortfile)
}

func main() {

	flag.Parse()

	InitLogs(os.Stdout, os.Stdout, os.Stderr)

	rh := new(RegexHandler)

	authPathPattern, _ := regexp.Compile("/v1/authorize")
	signPathPattern, _ := regexp.Compile("/v1/sign/.*")

	rh.HandleFunc(authPathPattern, Authorize)
	rh.HandleFunc(signPathPattern, Sign)

	go MapWriter()

	// Placeholder for authentication / authorization middleware on authorize call.

	err := http.ListenAndServe(PORT, rh)
	if err != nil {
		Error.Println(err)
	}
}

func MapWriter() {
	for {
		select {
		case t, ok := <-updates:
			if !ok {
				Error.Println("Publisher channel closed. Stopping.")
				return
			}
			Info.Println("Setting key %s to value %s", t.Token, t.CN)
			db[t.Token] = t.CN
		}
	}
}

func Authorize(w http.ResponseWriter, req *http.Request) {
	Info.Println("Received authorize call.")
	// Parse input
	cn := req.FormValue("cn")
	token := req.FormValue("token")

	// queue for write to map
	// ...
	t := Tuple{cn, token}
	updates <- t

	Info.Println("Service: %s", cn)
	Info.Println("Token: %s", token)

	req.Body.Close()
}

func Sign(w http.ResponseWriter, req *http.Request) {
	Info.Println("Received sign call.")

	// Pull the token out of the path
	_, token := path.Split(req.URL.Path)
	Info.Println("Received signing request for token %s", token)

	if len(token) == 0 {
		Warning.Println("No token provided.")
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Get the registered CN for the provided token (or fail)
	authCn := db[token]

	if authCn == "" {
		Warning.Println("Unauthorized CN.")
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	// Upload the CSR and copy it to some known location
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		Error.Println(err.Error())
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	randoName := fmt.Sprintf("%d.csr", rand.Int63())
	csrFilename := csrLocation + randoName
	err = ioutil.WriteFile(csrFilename, body, 0777)
	if err != nil {
		Error.Println(err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	Info.Println("File uploaded.")

	// Parse the CSR
	rawCSR, _ := ioutil.ReadFile(csrFilename)
	decodedCSR, _ := pem.Decode(rawCSR)
	csr, err := x509.ParseCertificateRequest(decodedCSR.Bytes)
	if err != nil {
		Error.Println(err.Error())
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	Info.Println("Received CSR for: %s", csr.Subject.CommonName)

	// check authorization for the provided commonname
	if csr.Subject.CommonName != authCn {
		Warning.Println("Unauthorized CN %s", csr.Subject.CommonName)
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	// Build the command for exec
	// openssl ca -config openssl-ca.cnf -policy signing_policy -extensions signing_req -out servercert.pem -infiles servercert.csr
	app := "openssl"
	command := "ca"
	c_flag := "-config"
	p_flag := "-policy"
	p_value := "signing_policy"
	e_flag := "-extensions"
	e_value := "signing_req"
	o_flag := "-out"
	outputFile := crtLocation + randoName + ".crt"
	i_flag := "-infiles"
	b_flag := "-batch"

	// Sign the CSR with OpenSSL
	cmd := exec.Command(app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
	args := fmt.Sprintf("%s %s %s %s %s %s %s %s %s %s %s %s %s", app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
	fmt.Println(args)
	stdOut, err := cmd.Output()
	if err != nil {
		Error.Println("OpenSSL stdout: %s", string(stdOut))
		Error.Println("OpenSSL stderr: %s", err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	// Open the output file for reading and stream it back on the response
	outputData, err := ioutil.ReadFile(outputFile)
	w.Write(outputData)
}

type MuxRoute struct {
	pattern *regexp.Regexp
	handler http.Handler
}

type RegexHandler struct {
	rs []*MuxRoute
}

func (rh *RegexHandler) Handler(p *regexp.Regexp, h http.Handler) {
	rh.rs = append(rh.rs, &MuxRoute{p, h})
}

func (rh *RegexHandler) HandleFunc(p *regexp.Regexp, h func(http.ResponseWriter, *http.Request)) {
	rh.rs = append(rh.rs, &MuxRoute{p, http.HandlerFunc(h)})
}

func (rh *RegexHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	for _, route := range rh.rs {
		if route.pattern.MatchString(r.URL.Path) {
			route.handler.ServeHTTP(w, r)
			return
		}
	}
	Warning.Println("Route not found: %s", r.URL.Path)
	// no pattern matched; send 404 response
	http.NotFound(w, r)
}
<commit_msg>Flag port.<commit_after>package main

import (
	"crypto/x509"
	"encoding/pem"
	"flag"
        "fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
)

var db = map[string]string{}
var csrLocation = "/var/csr/"
var crtLocation = "/var/crt/"
var confLocation = "/opt/pollendina/openssl-ca.cnf"

var port = flag.Int("port", 33004, "Default port for Pollendina CA.")

type Tuple struct{ CN, Token string }

var updates = make(chan Tuple)

var (
	Info    *log.Logger
	Warning *log.Logger
	Error   *log.Logger
)

func InitLogs(
	infoHandle io.Writer,
	warningHandle io.Writer,
	errorHandle io.Writer) {

	Info = log.New(infoHandle,
		"INFO: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Warning = log.New(warningHandle,
		"WARNING: ",
		log.Ldate|log.Ltime|log.Lshortfile)

	Error = log.New(errorHandle,
		"ERROR: ",
		log.Ldate|log.Ltime|log.Lshortfile)
}

func main() {

	flag.Parse()

	InitLogs(os.Stdout, os.Stdout, os.Stderr)

	rh := new(RegexHandler)

	authPathPattern, _ := regexp.Compile("/v1/authorize")
	signPathPattern, _ := regexp.Compile("/v1/sign/.*")

	rh.HandleFunc(authPathPattern, Authorize)
	rh.HandleFunc(signPathPattern, Sign)

	go MapWriter()

	// Placeholder for authentication / authorization middleware on authorize call.

	err := http.ListenAndServe(*port, rh)
	if err != nil {
		Error.Println(err)
	}
}

func MapWriter() {
	for {
		select {
		case t, ok := <-updates:
			if !ok {
				Error.Println("Publisher channel closed. Stopping.")
				return
			}
			Info.Println("Setting key %s to value %s", t.Token, t.CN)
			db[t.Token] = t.CN
		}
	}
}

func Authorize(w http.ResponseWriter, req *http.Request) {
	Info.Println("Received authorize call.")
	// Parse input
	cn := req.FormValue("cn")
	token := req.FormValue("token")

	// queue for write to map
	// ...
	t := Tuple{cn, token}
	updates <- t

	Info.Println("Service: %s", cn)
	Info.Println("Token: %s", token)

	req.Body.Close()
}

func Sign(w http.ResponseWriter, req *http.Request) {
	Info.Println("Received sign call.")

	// Pull the token out of the path
	_, token := path.Split(req.URL.Path)
	Info.Println("Received signing request for token %s", token)

	if len(token) == 0 {
		Warning.Println("No token provided.")
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	// Get the registered CN for the provided token (or fail)
	authCn := db[token]

	if authCn == "" {
		Warning.Println("Unauthorized CN.")
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	// Upload the CSR and copy it to some known location
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		Error.Println(err.Error())
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	randoName := fmt.Sprintf("%d.csr", rand.Int63())
	csrFilename := csrLocation + randoName
	err = ioutil.WriteFile(csrFilename, body, 0777)
	if err != nil {
		Error.Println(err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	Info.Println("File uploaded.")

	// Parse the CSR
	rawCSR, _ := ioutil.ReadFile(csrFilename)
	decodedCSR, _ := pem.Decode(rawCSR)
	csr, err := x509.ParseCertificateRequest(decodedCSR.Bytes)
	if err != nil {
		Error.Println(err.Error())
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	Info.Println("Received CSR for: %s", csr.Subject.CommonName)

	// check authorization for the provided commonname
	if csr.Subject.CommonName != authCn {
		Warning.Println("Unauthorized CN %s", csr.Subject.CommonName)
		w.WriteHeader(http.StatusUnauthorized)
		return
	}

	// Build the command for exec
	// openssl ca -config openssl-ca.cnf -policy signing_policy -extensions signing_req -out servercert.pem -infiles servercert.csr
	app := "openssl"
	command := "ca"
	c_flag := "-config"
	p_flag := "-policy"
	p_value := "signing_policy"
	e_flag := "-extensions"
	e_value := "signing_req"
	o_flag := "-out"
	outputFile := crtLocation + randoName + ".crt"
	i_flag := "-infiles"
	b_flag := "-batch"

	// Sign the CSR with OpenSSL
	cmd := exec.Command(app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
	args := fmt.Sprintf("%s %s %s %s %s %s %s %s %s %s %s %s %s", app, command, b_flag, c_flag, confLocation, p_flag, p_value, e_flag, e_value, o_flag, outputFile, i_flag, csrFilename)
	fmt.Println(args)
	stdOut, err := cmd.Output()
	if err != nil {
		Error.Println("OpenSSL stdout: %s", string(stdOut))
		Error.Println("OpenSSL stderr: %s", err.Error())
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	// Open the output file for reading and stream it back on the response
	outputData, err := ioutil.ReadFile(outputFile)
	w.Write(outputData)
}

type MuxRoute struct {
	pattern *regexp.Regexp
	handler http.Handler
}

type RegexHandler struct {
	rs []*MuxRoute
}

func (rh *RegexHandler) Handler(p *regexp.Regexp, h http.Handler) {
	rh.rs = append(rh.rs, &MuxRoute{p, h})
}

func (rh *RegexHandler) HandleFunc(p *regexp.Regexp, h func(http.ResponseWriter, *http.Request)) {
	rh.rs = append(rh.rs, &MuxRoute{p, http.HandlerFunc(h)})
}

func (rh *RegexHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	for _, route := range rh.rs {
		if route.pattern.MatchString(r.URL.Path) {
			route.handler.ServeHTTP(w, r)
			return
		}
	}
	Warning.Println("Route not found: %s", r.URL.Path)
	// no pattern matched; send 404 response
	http.NotFound(w, r)
}
<|endoftext|>


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

package manager

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	dclient "github.com/fsouza/go-dockerclient"
	"github.com/golang/glog"
	"github.com/google/cadvisor/container/docker"
	"github.com/google/cadvisor/fs"
	info "github.com/google/cadvisor/info/v1"
	"github.com/google/cadvisor/utils"
	"github.com/google/cadvisor/utils/sysfs"
	"github.com/google/cadvisor/utils/sysinfo"
	version "github.com/google/cadvisor/version"
)

var cpuRegExp = regexp.MustCompile("processor\\t*: +([0-9]+)")
var coreRegExp = regexp.MustCompile("core id\\t*: +([0-9]+)")
var nodeRegExp = regexp.MustCompile("physical id\\t*: +([0-9]+)")
var CpuClockSpeedMHz = regexp.MustCompile("cpu MHz\\t*: +([0-9]+.[0-9]+)")
var memoryCapacityRegexp = regexp.MustCompile("MemTotal: *([0-9]+) kB")

var machineIdFilePath = flag.String("machine_id_file", "/etc/machine-id,/var/lib/dbus/machine-id", "Comma-separated list of files to check for machine-id. Use the first one that exists.")

func getClockSpeed(procInfo []byte) (uint64, error) {
	// First look through sys to find a max supported cpu frequency.
	const maxFreqFile = "/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq"
	if utils.FileExists(maxFreqFile) {
		val, err := ioutil.ReadFile(maxFreqFile)
		if err != nil {
			return 0, err
		}
		var maxFreq uint64
		n, err := fmt.Sscanf(string(val), "%d", &maxFreq)
		if err != nil || n != 1 {
			return 0, fmt.Errorf("could not parse frequency %q", val)
		}
		return maxFreq, nil
	}
	// Fall back to /proc/cpuinfo
	matches := CpuClockSpeedMHz.FindSubmatch(procInfo)
	if len(matches) != 2 {
		return 0, fmt.Errorf("could not detect clock speed from output: %q", string(procInfo))
	}
	speed, err := strconv.ParseFloat(string(matches[1]), 64)
	if err != nil {
		return 0, err
	}
	// Convert to kHz
	return uint64(speed * 1000), nil
}

func getMemoryCapacity(b []byte) (int64, error) {
	matches := memoryCapacityRegexp.FindSubmatch(b)
	if len(matches) != 2 {
		return -1, fmt.Errorf("failed to find memory capacity in output: %q", string(b))
	}
	m, err := strconv.ParseInt(string(matches[1]), 10, 64)
	if err != nil {
		return -1, err
	}

	// Convert to bytes.
	return m * 1024, err
}

func extractValue(s string, r *regexp.Regexp) (bool, int, error) {
	matches := r.FindSubmatch([]byte(s))
	if len(matches) == 2 {
		val, err := strconv.ParseInt(string(matches[1]), 10, 32)
		if err != nil {
			return true, -1, err
		}
		return true, int(val), nil
	}
	return false, -1, nil
}

func findNode(nodes []info.Node, id int) (bool, int) {
	for i, n := range nodes {
		if n.Id == id {
			return true, i
		}
	}
	return false, -1
}

func addNode(nodes *[]info.Node, id int) (int, error) {
	var idx int
	if id == -1 {
		// Some VMs don't fill topology data. Export single package.
		id = 0
	}

	ok, idx := findNode(*nodes, id)
	if !ok {
		// New node
		node := info.Node{Id: id}
		// Add per-node memory information.
		meminfo := fmt.Sprintf("/sys/devices/system/node/node%d/meminfo", id)
		out, err := ioutil.ReadFile(meminfo)
		// Ignore if per-node info is not available.
		if err == nil {
			m, err := getMemoryCapacity(out)
			if err != nil {
				return -1, err
			}
			node.Memory = uint64(m)
		}
		*nodes = append(*nodes, node)
		idx = len(*nodes) - 1
	}
	return idx, nil
}

func getTopology(sysFs sysfs.SysFs, cpuinfo string) ([]info.Node, int, error) {
	nodes := []info.Node{}
	numCores := 0
	lastThread := -1
	lastCore := -1
	lastNode := -1
	for _, line := range strings.Split(cpuinfo, "\n") {
		ok, val, err := extractValue(line, cpuRegExp)
		if err != nil {
			return nil, -1, fmt.Errorf("could not parse cpu info from %q: %v", line, err)
		}
		if ok {
			thread := val
			numCores++
			if lastThread != -1 {
				// New cpu section. Save last one.
				nodeIdx, err := addNode(&nodes, lastNode)
				if err != nil {
					return nil, -1, fmt.Errorf("failed to add node %d: %v", lastNode, err)
				}
				nodes[nodeIdx].AddThread(lastThread, lastCore)
				lastCore = -1
				lastNode = -1
			}
			lastThread = thread
		}
		ok, val, err = extractValue(line, coreRegExp)
		if err != nil {
			return nil, -1, fmt.Errorf("could not parse core info from %q: %v", line, err)
		}
		if ok {
			lastCore = val
		}
		ok, val, err = extractValue(line, nodeRegExp)
		if err != nil {
			return nil, -1, fmt.Errorf("could not parse node info from %q: %v", line, err)
		}
		if ok {
			lastNode = val
		}
	}
	nodeIdx, err := addNode(&nodes, lastNode)
	if err != nil {
		return nil, -1, fmt.Errorf("failed to add node %d: %v", lastNode, err)
	}
	nodes[nodeIdx].AddThread(lastThread, lastCore)
	if numCores < 1 {
		return nil, numCores, fmt.Errorf("could not detect any cores")
	}
	for idx, node := range nodes {
		caches, err := sysinfo.GetCacheInfo(sysFs, node.Cores[0].Threads[0])
		if err != nil {
			return nil, -1, fmt.Errorf("failed to get cache information for node %d: %v", node.Id, err)
		}
		numThreadsPerCore := len(node.Cores[0].Threads)
		numThreadsPerNode := len(node.Cores) * numThreadsPerCore
		for _, cache := range caches {
			c := info.Cache{
				Size:  cache.Size,
				Level: cache.Level,
				Type:  cache.Type,
			}
			if cache.Cpus == numThreadsPerNode && cache.Level > 2 {
				// Add a node-level cache.
				nodes[idx].AddNodeCache(c)
			} else if cache.Cpus == numThreadsPerCore {
				// Add to each core.
				nodes[idx].AddPerCoreCache(c)
			}
			// Ignore unknown caches.
		}
	}
	return nodes, numCores, nil
}

func getMachineID() string {
	if len(*machineIdFilePath) == 0 {
		return ""
	}
	for _, file := range strings.Split(*machineIdFilePath, ",") {
		id, err := ioutil.ReadFile(file)
		if err == nil {
			return strings.TrimSpace(string(id))
		}
	}
	glog.Infof("Couldn't collect machine-id from any of the files in %q", *machineIdFilePath)
	return ""
}

func getMachineInfo(sysFs sysfs.SysFs) (*info.MachineInfo, error) {
	cpuinfo, err := ioutil.ReadFile("/proc/cpuinfo")
	clockSpeed, err := getClockSpeed(cpuinfo)
	if err != nil {
		return nil, err
	}

	// Get the amount of usable memory from /proc/meminfo.
	out, err := ioutil.ReadFile("/proc/meminfo")
	if err != nil {
		return nil, err
	}

	memoryCapacity, err := getMemoryCapacity(out)
	if err != nil {
		return nil, err
	}

	fsInfo, err := fs.NewFsInfo()
	if err != nil {
		return nil, err
	}
	filesystems, err := fsInfo.GetGlobalFsInfo()
	if err != nil {
		return nil, err
	}

	diskMap, err := sysinfo.GetBlockDeviceInfo(sysFs)
	if err != nil {
		return nil, err
	}

	netDevices, err := sysinfo.GetNetworkDevices(sysFs)
	if err != nil {
		return nil, err
	}

	topology, numCores, err := getTopology(sysFs, string(cpuinfo))
	if err != nil {
		return nil, err
	}

	system_uuid, err := sysinfo.GetSystemUUID(sysFs)
	if err != nil {
		return nil, err
	}

	machineInfo := &info.MachineInfo{
		NumCores:       numCores,
		CpuFrequency:   clockSpeed,
		MemoryCapacity: memoryCapacity,
		DiskMap:        diskMap,
		NetworkDevices: netDevices,
		Topology:       topology,
		MachineID:      getMachineID(),
		SystemUUID:     system_uuid,
	}

	for _, fs := range filesystems {
		machineInfo.Filesystems = append(machineInfo.Filesystems, info.FsInfo{Device: fs.Device, Capacity: fs.Capacity})
	}

	return machineInfo, nil
}

func getVersionInfo() (*info.VersionInfo, error) {

	kernel_version := getKernelVersion()
	container_os := getContainerOsVersion()
	docker_version := getDockerVersion()

	return &info.VersionInfo{
		KernelVersion:      kernel_version,
		ContainerOsVersion: container_os,
		DockerVersion:      docker_version,
		CadvisorVersion:    version.VERSION,
	}, nil
}

func getContainerOsVersion() string {
	container_os := "Unknown"
	os_release, err := ioutil.ReadFile("/etc/os-release")
	if err == nil {
		// We might be running in a busybox or some hand-crafted image.
		// It's useful to know why cadvisor didn't come up.
		for _, line := range strings.Split(string(os_release), "\n") {
			parsed := strings.Split(line, "\"")
			if len(parsed) == 3 && parsed[0] == "PRETTY_NAME=" {
				container_os = parsed[1]
				break
			}
		}
	}
	return container_os
}

func getDockerVersion() string {
	docker_version := "Unknown"
	client, err := dclient.NewClient(*docker.ArgDockerEndpoint)
	if err == nil {
		version, err := client.Version()
		if err == nil {
			docker_version = version.Get("Version")
		}
	}
	return docker_version
}

func getKernelVersion() string {
	uname := &syscall.Utsname{}

	if err := syscall.Uname(uname); err != nil {
		return "Unknown"
	}

	release := make([]byte, len(uname.Release))
	i := 0
	for _, c := range uname.Release {
		release[i] = byte(c)
		i++
	}
	release = release[:bytes.IndexByte(release, 0)]

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

package manager

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"regexp"
	"strconv"
	"strings"
	"syscall"

	dclient "github.com/fsouza/go-dockerclient"
	"github.com/golang/glog"
	"github.com/google/cadvisor/container/docker"
	"github.com/google/cadvisor/fs"
	info "github.com/google/cadvisor/info/v1"
	"github.com/google/cadvisor/utils"
	"github.com/google/cadvisor/utils/sysfs"
	"github.com/google/cadvisor/utils/sysinfo"
	version "github.com/google/cadvisor/version"
)

var cpuRegExp = regexp.MustCompile("processor\\t*: +([0-9]+)")
var coreRegExp = regexp.MustCompile("core id\\t*: +([0-9]+)")
var nodeRegExp = regexp.MustCompile("physical id\\t*: +([0-9]+)")
var CpuClockSpeedMHz = regexp.MustCompile("cpu MHz\\t*: +([0-9]+.[0-9]+)")
var memoryCapacityRegexp = regexp.MustCompile("MemTotal: *([0-9]+) kB")

var machineIdFilePath = flag.String("machine_id_file", "/etc/machine-id,/var/lib/dbus/machine-id", "Comma-separated list of files to check for machine-id. Use the first one that exists.")

func getClockSpeed(procInfo []byte) (uint64, error) {
	// First look through sys to find a max supported cpu frequency.
	const maxFreqFile = "/sys/devices/system/cpu/cpu0/cpufreq/cpuinfo_max_freq"
	if utils.FileExists(maxFreqFile) {
		val, err := ioutil.ReadFile(maxFreqFile)
		if err != nil {
			return 0, err
		}
		var maxFreq uint64
		n, err := fmt.Sscanf(string(val), "%d", &maxFreq)
		if err != nil || n != 1 {
			return 0, fmt.Errorf("could not parse frequency %q", val)
		}
		return maxFreq, nil
	}
	// Fall back to /proc/cpuinfo
	matches := CpuClockSpeedMHz.FindSubmatch(procInfo)
	if len(matches) != 2 {
		return 0, fmt.Errorf("could not detect clock speed from output: %q", string(procInfo))
	}
	speed, err := strconv.ParseFloat(string(matches[1]), 64)
	if err != nil {
		return 0, err
	}
	// Convert to kHz
	return uint64(speed * 1000), nil
}

func getMemoryCapacity(b []byte) (int64, error) {
	matches := memoryCapacityRegexp.FindSubmatch(b)
	if len(matches) != 2 {
		return -1, fmt.Errorf("failed to find memory capacity in output: %q", string(b))
	}
	m, err := strconv.ParseInt(string(matches[1]), 10, 64)
	if err != nil {
		return -1, err
	}

	// Convert to bytes.
	return m * 1024, err
}

func extractValue(s string, r *regexp.Regexp) (bool, int, error) {
	matches := r.FindSubmatch([]byte(s))
	if len(matches) == 2 {
		val, err := strconv.ParseInt(string(matches[1]), 10, 32)
		if err != nil {
			return true, -1, err
		}
		return true, int(val), nil
	}
	return false, -1, nil
}

func findNode(nodes []info.Node, id int) (bool, int) {
	for i, n := range nodes {
		if n.Id == id {
			return true, i
		}
	}
	return false, -1
}

func addNode(nodes *[]info.Node, id int) (int, error) {
	var idx int
	if id == -1 {
		// Some VMs don't fill topology data. Export single package.
		id = 0
	}

	ok, idx := findNode(*nodes, id)
	if !ok {
		// New node
		node := info.Node{Id: id}
		// Add per-node memory information.
		meminfo := fmt.Sprintf("/sys/devices/system/node/node%d/meminfo", id)
		out, err := ioutil.ReadFile(meminfo)
		// Ignore if per-node info is not available.
		if err == nil {
			m, err := getMemoryCapacity(out)
			if err != nil {
				return -1, err
			}
			node.Memory = uint64(m)
		}
		*nodes = append(*nodes, node)
		idx = len(*nodes) - 1
	}
	return idx, nil
}

func getTopology(sysFs sysfs.SysFs, cpuinfo string) ([]info.Node, int, error) {
	nodes := []info.Node{}
	numCores := 0
	lastThread := -1
	lastCore := -1
	lastNode := -1
	for _, line := range strings.Split(cpuinfo, "\n") {
		ok, val, err := extractValue(line, cpuRegExp)
		if err != nil {
			return nil, -1, fmt.Errorf("could not parse cpu info from %q: %v", line, err)
		}
		if ok {
			thread := val
			numCores++
			if lastThread != -1 {
				// New cpu section. Save last one.
				nodeIdx, err := addNode(&nodes, lastNode)
				if err != nil {
					return nil, -1, fmt.Errorf("failed to add node %d: %v", lastNode, err)
				}
				nodes[nodeIdx].AddThread(lastThread, lastCore)
				lastCore = -1
				lastNode = -1
			}
			lastThread = thread
		}
		ok, val, err = extractValue(line, coreRegExp)
		if err != nil {
			return nil, -1, fmt.Errorf("could not parse core info from %q: %v", line, err)
		}
		if ok {
			lastCore = val
		}
		ok, val, err = extractValue(line, nodeRegExp)
		if err != nil {
			return nil, -1, fmt.Errorf("could not parse node info from %q: %v", line, err)
		}
		if ok {
			lastNode = val
		}
	}
	nodeIdx, err := addNode(&nodes, lastNode)
	if err != nil {
		return nil, -1, fmt.Errorf("failed to add node %d: %v", lastNode, err)
	}
	nodes[nodeIdx].AddThread(lastThread, lastCore)
	if numCores < 1 {
		return nil, numCores, fmt.Errorf("could not detect any cores")
	}
	for idx, node := range nodes {
		caches, err := sysinfo.GetCacheInfo(sysFs, node.Cores[0].Threads[0])
		if err != nil {
			return nil, -1, fmt.Errorf("failed to get cache information for node %d: %v", node.Id, err)
		}
		numThreadsPerCore := len(node.Cores[0].Threads)
		numThreadsPerNode := len(node.Cores) * numThreadsPerCore
		for _, cache := range caches {
			c := info.Cache{
				Size:  cache.Size,
				Level: cache.Level,
				Type:  cache.Type,
			}
			if cache.Cpus == numThreadsPerNode && cache.Level > 2 {
				// Add a node-level cache.
				nodes[idx].AddNodeCache(c)
			} else if cache.Cpus == numThreadsPerCore {
				// Add to each core.
				nodes[idx].AddPerCoreCache(c)
			}
			// Ignore unknown caches.
		}
	}
	return nodes, numCores, nil
}

func getMachineID() string {
	if len(*machineIdFilePath) == 0 {
		return ""
	}
	for _, file := range strings.Split(*machineIdFilePath, ",") {
		id, err := ioutil.ReadFile(file)
		if err == nil {
			return strings.TrimSpace(string(id))
		}
	}
	glog.Infof("Couldn't collect machine-id from any of the files in %q", *machineIdFilePath)
	return ""
}

func getMachineInfo(sysFs sysfs.SysFs) (*info.MachineInfo, error) {
	cpuinfo, err := ioutil.ReadFile("/proc/cpuinfo")
	clockSpeed, err := getClockSpeed(cpuinfo)
	if err != nil {
		return nil, err
	}

	// Get the amount of usable memory from /proc/meminfo.
	out, err := ioutil.ReadFile("/proc/meminfo")
	if err != nil {
		return nil, err
	}

	memoryCapacity, err := getMemoryCapacity(out)
	if err != nil {
		return nil, err
	}

	fsInfo, err := fs.NewFsInfo()
	if err != nil {
		return nil, err
	}
	filesystems, err := fsInfo.GetGlobalFsInfo()
	if err != nil {
		return nil, err
	}

	diskMap, err := sysinfo.GetBlockDeviceInfo(sysFs)
	if err != nil {
		return nil, err
	}

	netDevices, err := sysinfo.GetNetworkDevices(sysFs)
	if err != nil {
		return nil, err
	}

	topology, numCores, err := getTopology(sysFs, string(cpuinfo))
	if err != nil {
		return nil, err
	}

	systemUUID, err := sysinfo.GetSystemUUID(sysFs)
	if err != nil {
		glog.Errorf("Failed to get system UUID: %v", err)
		systemUUID = ""
	}

	machineInfo := &info.MachineInfo{
		NumCores:       numCores,
		CpuFrequency:   clockSpeed,
		MemoryCapacity: memoryCapacity,
		DiskMap:        diskMap,
		NetworkDevices: netDevices,
		Topology:       topology,
		MachineID:      getMachineID(),
		SystemUUID:     systemUUID,
	}

	for _, fs := range filesystems {
		machineInfo.Filesystems = append(machineInfo.Filesystems, info.FsInfo{Device: fs.Device, Capacity: fs.Capacity})
	}

	return machineInfo, nil
}

func getVersionInfo() (*info.VersionInfo, error) {

	kernel_version := getKernelVersion()
	container_os := getContainerOsVersion()
	docker_version := getDockerVersion()

	return &info.VersionInfo{
		KernelVersion:      kernel_version,
		ContainerOsVersion: container_os,
		DockerVersion:      docker_version,
		CadvisorVersion:    version.VERSION,
	}, nil
}

func getContainerOsVersion() string {
	container_os := "Unknown"
	os_release, err := ioutil.ReadFile("/etc/os-release")
	if err == nil {
		// We might be running in a busybox or some hand-crafted image.
		// It's useful to know why cadvisor didn't come up.
		for _, line := range strings.Split(string(os_release), "\n") {
			parsed := strings.Split(line, "\"")
			if len(parsed) == 3 && parsed[0] == "PRETTY_NAME=" {
				container_os = parsed[1]
				break
			}
		}
	}
	return container_os
}

func getDockerVersion() string {
	docker_version := "Unknown"
	client, err := dclient.NewClient(*docker.ArgDockerEndpoint)
	if err == nil {
		version, err := client.Version()
		if err == nil {
			docker_version = version.Get("Version")
		}
	}
	return docker_version
}

func getKernelVersion() string {
	uname := &syscall.Utsname{}

	if err := syscall.Uname(uname); err != nil {
		return "Unknown"
	}

	release := make([]byte, len(uname.Release))
	i := 0
	for _, c := range uname.Release {
		release[i] = byte(c)
		i++
	}
	release = release[:bytes.IndexByte(release, 0)]

	return string(release)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"path"
	"strconv"
	"strings"

	"github.com/animenotifier/arn"
	"github.com/animenotifier/kitsu"
	"github.com/fatih/color"
)

func main() {
	color.Yellow("Syncing characters with Kitsu DB")
	defer arn.Node.Close()

	kitsuCharacters := kitsu.StreamCharacters()

	for kitsuCharacter := range kitsuCharacters {
		character := &arn.Character{
			ID: kitsuCharacter.ID,
			Name: arn.CharacterName{
				Canonical: kitsuCharacter.Attributes.Name,
			},
			Image: arn.CharacterImage{
				Extension: path.Ext(kitsu.FixImageURL(kitsuCharacter.Attributes.Image.Original)),
			},
			Description: kitsuCharacter.Attributes.Description,
			Attributes:  []*arn.CharacterAttribute{},
			Mappings: []*arn.Mapping{
				&arn.Mapping{
					Service:   "kitsu/character",
					ServiceID: kitsuCharacter.ID,
				},
			},
		}

		if kitsuCharacter.Attributes.MalID != 0 {
			character.Mappings = append(character.Mappings, &arn.Mapping{
				Service:   "myanimelist/character",
				ServiceID: strconv.Itoa(kitsuCharacter.Attributes.MalID),
			})
		}

		// We use markdown, so replace
 with two line breaks.
		character.Description = strings.Replace(character.Description, "
", "\n\n", -1)

		// Parse attributes like these:
		// - Position: Club Manager
		// - Height: 162 cm (5' 4")
		// - Weight: 48 kg (106 lb)
		// - Birthday: November 24
		// - Hair color: Brown
		// - Eyes: Blue (anime), Green (manga)

		lines := strings.Split(character.Description, "\n\n")
		finalLines := make([]string, 0, len(lines))

		for _, line := range lines {
			originalLine := line

			if strings.HasPrefix(line, "(") {
				line = strings.TrimPrefix(line, "(")
				line = strings.TrimSuffix(line, ")")
			}

			line = strings.TrimSpace(line)

			colonPos := strings.Index(line, ":")

			if colonPos == -1 || colonPos < 2 || colonPos > 25 {
				finalLines = append(finalLines, originalLine)
				continue
			}

			key := line[:colonPos]
			value := line[colonPos+1:]

			value = strings.TrimSpace(value)

			if key == "source" {
				key = "Source"
			}

			character.Attributes = append(character.Attributes, &arn.CharacterAttribute{
				Name:  key,
				Value: value,
			})

			fmt.Println(color.CyanString(key), color.YellowString(value))
		}

		character.Description = strings.Join(finalLines, "\n\n")
		character.Description = strings.Trim(character.Description, "\n\n")
		character.Save()

		fmt.Printf("%s %s %s\n", color.GreenString("✔"), character.ID, character.Name)
	}

	color.Green("Finished.")
}
<commit_msg>New Kitsu character sync<commit_after>package main

import (
	"fmt"
	"path"
	"strconv"
	"strings"

	"github.com/animenotifier/arn"
	"github.com/animenotifier/kitsu"
	"github.com/fatih/color"
)

func main() {
	color.Yellow("Syncing characters with Kitsu DB")
	defer arn.Node.Close()

	kitsuCharacters := kitsu.StreamCharacters()

	for kitsuCharacter := range kitsuCharacters {
		character := &arn.Character{
			ID: kitsuCharacter.ID,
			Name: arn.CharacterName{
				Canonical: kitsuCharacter.Attributes.CanonicalName,
				English:   kitsuCharacter.Attributes.Names.En,
				Japanese:  kitsuCharacter.Attributes.Names.JaJp,
				Synonyms:  kitsuCharacter.Attributes.OtherNames,
			},
			Image: arn.CharacterImage{
				Extension: path.Ext(kitsu.FixImageURL(kitsuCharacter.Attributes.Image.Original)),
			},
			Description: kitsuCharacter.Attributes.Description,
			Attributes:  []*arn.CharacterAttribute{},
			Mappings: []*arn.Mapping{
				&arn.Mapping{
					Service:   "kitsu/character",
					ServiceID: kitsuCharacter.ID,
				},
			},
		}

		if kitsuCharacter.Attributes.MalID != 0 {
			character.Mappings = append(character.Mappings, &arn.Mapping{
				Service:   "myanimelist/character",
				ServiceID: strconv.Itoa(kitsuCharacter.Attributes.MalID),
			})
		}

		// We use markdown, so replace
 with two line breaks.
		character.Description = strings.Replace(character.Description, "
", "\n\n", -1)

		// Parse attributes like these:
		// - Position: Club Manager
		// - Height: 162 cm (5' 4")
		// - Weight: 48 kg (106 lb)
		// - Birthday: November 24
		// - Hair color: Brown
		// - Eyes: Blue (anime), Green (manga)

		lines := strings.Split(character.Description, "\n\n")
		finalLines := make([]string, 0, len(lines))

		for _, line := range lines {
			originalLine := line

			if strings.HasPrefix(line, "(") {
				line = strings.TrimPrefix(line, "(")
				line = strings.TrimSuffix(line, ")")
			}

			line = strings.TrimSpace(line)

			colonPos := strings.Index(line, ":")

			if colonPos == -1 || colonPos < 2 || colonPos > 25 {
				finalLines = append(finalLines, originalLine)
				continue
			}

			key := line[:colonPos]
			value := line[colonPos+1:]

			value = strings.TrimSpace(value)

			if key == "source" {
				key = "Source"
			}

			character.Attributes = append(character.Attributes, &arn.CharacterAttribute{
				Name:  key,
				Value: value,
			})

			fmt.Println(color.CyanString(key), color.YellowString(value))
		}

		character.Description = strings.Join(finalLines, "\n\n")
		character.Description = strings.Trim(character.Description, "\n\n")
		character.Save()

		// Save Kitsu character in Kitsu DB
		arn.Kitsu.Set("Character", kitsuCharacter.ID, kitsuCharacter)

		// Log
		fmt.Printf("%s %s %s\n", color.GreenString("✔"), character.ID, character.Name)
	}

	color.Green("Finished.")
}
<|endoftext|>


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

package command

import (
	"fmt"

	"github.com/spf13/cobra"
)

// NewAuthCommand returns the cobra command for "auth".
func NewAuthCommand() *cobra.Command {
	ac := &cobra.Command{
		Use:   "auth <enable or disable>",
		Short: "Enable or disable authentication",
	}

	ac.AddCommand(newAuthEnableCommand())
	ac.AddCommand(newAuthDisableCommand())

	return ac
}

func newAuthEnableCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "enable",
		Short: "Enables authentication",
		Run:   authEnableCommandFunc,
	}
}

// authEnableCommandFunc executes the "auth enable" command.
func authEnableCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 0 {
		ExitWithError(ExitBadArgs, fmt.Errorf("auth enable command does not accept any arguments."))
	}

	ctx, cancel := commandCtx(cmd)
	_, err := mustClientFromCmd(cmd).Auth.AuthEnable(ctx)
	cancel()
	if err != nil {
		ExitWithError(ExitError, err)
	}

	fmt.Println("Authentication Enabled")
}

func newAuthDisableCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "disable",
		Short: "Disables authentication",
		Run:   authDisableCommandFunc,
	}
}

// authDisableCommandFunc executes the "auth disable" command.
func authDisableCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 0 {
		ExitWithError(ExitBadArgs, fmt.Errorf("auth disable command does not accept any arguments."))
	}

	ctx, cancel := commandCtx(cmd)
	_, err := mustClientFromCmd(cmd).Auth.AuthDisable(ctx)
	cancel()
	if err != nil {
		ExitWithError(ExitError, err)
	}

	fmt.Println("Authentication Disabled")
}
<commit_msg>etcdctl: create root role on auth enable if it does not yet exist<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

import (
	"fmt"

	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
	"github.com/spf13/cobra"
)

// NewAuthCommand returns the cobra command for "auth".
func NewAuthCommand() *cobra.Command {
	ac := &cobra.Command{
		Use:   "auth <enable or disable>",
		Short: "Enable or disable authentication",
	}

	ac.AddCommand(newAuthEnableCommand())
	ac.AddCommand(newAuthDisableCommand())

	return ac
}

func newAuthEnableCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "enable",
		Short: "Enables authentication",
		Run:   authEnableCommandFunc,
	}
}

// authEnableCommandFunc executes the "auth enable" command.
func authEnableCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 0 {
		ExitWithError(ExitBadArgs, fmt.Errorf("auth enable command does not accept any arguments."))
	}

	ctx, cancel := commandCtx(cmd)
	cli := mustClientFromCmd(cmd)
	var err error
	for err == nil {
		if _, err = cli.AuthEnable(ctx); err == nil {
			break
		}
		if err == rpctypes.ErrRootRoleNotExist {
			if _, err = cli.RoleAdd(ctx, "root"); err != nil {
				break
			}
			if _, err = cli.UserGrantRole(ctx, "root", "root"); err != nil {
				break
			}
		}
	}
	cancel()
	if err != nil {
		ExitWithError(ExitError, err)
	}

	fmt.Println("Authentication Enabled")
}

func newAuthDisableCommand() *cobra.Command {
	return &cobra.Command{
		Use:   "disable",
		Short: "Disables authentication",
		Run:   authDisableCommandFunc,
	}
}

// authDisableCommandFunc executes the "auth disable" command.
func authDisableCommandFunc(cmd *cobra.Command, args []string) {
	if len(args) != 0 {
		ExitWithError(ExitBadArgs, fmt.Errorf("auth disable command does not accept any arguments."))
	}

	ctx, cancel := commandCtx(cmd)
	_, err := mustClientFromCmd(cmd).Auth.AuthDisable(ctx)
	cancel()
	if err != nil {
		ExitWithError(ExitError, err)
	}

	fmt.Println("Authentication Disabled")
}
<|endoftext|>


<commit_before>package thermocline_test

import (
	"testing"
	"time"

	"github.com/fortytw2/thermocline"
)

func TestDelayedTask(t *testing.T) {
	t.Parallel()

	_, reader, writer := SetupBroker(t)

	task, err := thermocline.NewDelayedTask("task infos", time.Second)
	if err != nil {
		t.Errorf("could not make delayed task, %s", err)
	}

	before := time.Now()
	writer <- task

	<-reader
	after := time.Now()

	if after.Sub(before) <= time.Second {
		t.Fatalf("task appeared before duration!")
	}
}

func TestLongDelayedTask(t *testing.T) {
	t.Parallel()

	if testing.Short() {
		t.Skip("skipping in short mode.")
	}

	_, reader, writer := SetupBroker(t)

	task, err := thermocline.NewDelayedTask("task infos", time.Second*10)
	if err != nil {
		t.Errorf("could not make delayed task, %s", err)
	}

	before := time.Now()
	writer <- task

	<-reader
	after := time.Now()

	if after.Sub(before) <= time.Second*9 {
		t.Fatalf("task appeared before duration!")
	}
}
<commit_msg>shorten _all_ tests with timers<commit_after>package thermocline_test

import (
	"testing"
	"time"

	"github.com/fortytw2/thermocline"
)

func TestDelayedTask(t *testing.T) {
	t.Parallel()

	_, reader, writer := SetupBroker(t)

	task, err := thermocline.NewDelayedTask("task infos", time.Second)
	if err != nil {
		t.Errorf("could not make delayed task, %s", err)
	}

	before := time.Now()
	writer <- task

	<-reader
	after := time.Now()

	if after.Sub(before) <= time.Second - time.Millisecond * 100 {
		t.Fatalf("task appeared before duration!")
	}
}

func TestLongDelayedTask(t *testing.T) {
	t.Parallel()

	if testing.Short() {
		t.Skip("skipping in short mode.")
	}

	_, reader, writer := SetupBroker(t)

	task, err := thermocline.NewDelayedTask("task infos", time.Second*10)
	if err != nil {
		t.Errorf("could not make delayed task, %s", err)
	}

	before := time.Now()
	writer <- task

	<-reader
	after := time.Now()

	if after.Sub(before) <= time.Second*9 {
		t.Fatalf("task appeared before duration!")
	}
}
<|endoftext|>


<commit_before>// +build unit

package rpcutils

import (
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"sync"
	"testing"
	"time"

	"github.com/control-center/serviced/commons/pool"
	. "gopkg.in/check.v1"
)

func Test(t *testing.T) { TestingT(t) }

type MySuite struct{}

var (
	_             = Suite(&MySuite{})
	rtt           *RPCTestType
	rpcClient     Client
	bareRpcClient *rpc.Client
)

type RPCTestType int

func (rtt *RPCTestType) Sleep(sleep time.Duration, reply *time.Duration) error {
	time.Sleep(sleep)
	*reply = sleep
	return nil
}

type TestArgs struct {
	A string
	B int
	C bool
	D []string
}

func (rtt *RPCTestType) NilReply(arg string, _ *struct{}) error {
	return nil
}

func (rtt *RPCTestType) Echo(arg string, reply *string) error {
	*reply = arg
	return nil
}

func (rtt *RPCTestType) StructCall(arg TestArgs, reply *TestArgs) error {
	*reply = arg
	return nil
}

func (s *MySuite) SetUpSuite(c *C) {
	rtt = new(RPCTestType)
	RegisterLocal("RPCTestType", rtt)
	rpc.Register(rtt)
	rpc.HandleHTTP()
	listener, err := net.Listen("tcp", ":32111")
	if err != nil {
		c.Errorf("listen error: %s", err)
	}
	go func() {
		for {
			conn, err := listener.Accept()
			defer conn.Close()
			if err != nil {
				c.Errorf("Error accepting connections: %s", err)
			}
			go rpc.ServeCodec(jsonrpc.NewServerCodec(conn))
		}
	}()

	rpcClient, _ = newClient("localhost:32111", 1, DiscardClientTimeout, connectRPC)
	bareRpcClient, _ = connectRPC("localhost:32111")
}

func (s *MySuite) TestConcurrentTimeout(c *C) {

	sleepTime := 100 * time.Millisecond
	client, err := newClient("localhost:32111", 1, DiscardClientTimeout, connectRPC)
	c.Assert(err, IsNil)

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		var reply time.Duration
		wg.Done()
		// Sleep, timeout after two. Shouldn't error.
		err := client.Call("RPCTestType.Sleep", sleepTime, &reply, 2*sleepTime)
		c.Assert(err, IsNil)
		c.Assert(reply, Equals, sleepTime)
	}()
	wg.Wait()
	var reply time.Duration
	// should time out wating for client
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, sleepTime/2)
	c.Assert(err, Equals, pool.ErrItemUnavailable)
}

func (s *MySuite) TestTimeout(c *C) {

	client, err := newClient("localhost:32111", 1, 10*time.Millisecond, connectRPC)

	sleepTime := 100 * time.Millisecond

	var reply time.Duration

	// Sleep for one second, timeout after two. Shouldn't error.
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, 2*time.Second)
	c.Assert(err, IsNil)
	c.Assert(reply, Equals, sleepTime)

	// Sleep, never timeout. Shouldn't error.
	sleepTime = sleepTime * 2
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, 0)
	c.Assert(err, IsNil)
	c.Assert(reply, Equals, sleepTime)

	// Sleep and timeout after half sleep. Should error.
	err = client.Call("RPCTestType.Sleep", &sleepTime, &reply, sleepTime/2)
	c.Assert(err, NotNil)
	c.Assert(err, ErrorMatches, "RPC call to RPCTestType.Sleep timed out after .+")
}

func (s *MySuite) TestLongCall(c *C) {
	client, err := newClient("localhost:32111", 1, 250*time.Millisecond, connectRPC)
	c.Assert(err, IsNil)

	startWg := sync.WaitGroup{}

	wg := sync.WaitGroup{}
	wg.Add(1)
	startWg.Add(1)
	go func() {
		var reply time.Duration
		startWg.Done()
		// Sleep for time , timeout after twice as much. Shouldn't error but underlying client will be invalidated
		sleepTime := 750 * time.Millisecond
		err := client.Call("RPCTestType.Sleep", sleepTime, &reply, 2*sleepTime)
		c.Assert(err, IsNil)
		c.Assert(reply, Equals, sleepTime)
		wg.Done()

	}()
	startWg.Wait()
	//after 250ms the previous call should have caused the the client to go stale
	time.Sleep(500 * time.Millisecond)
	var reply time.Duration
	sleepTime := 10 * time.Millisecond
	// should not time out wating for client
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, 20*time.Millisecond)
	c.Assert(err, IsNil)
	c.Assert(reply, Equals, sleepTime)
	wg.Wait() //wait for go routine to run asserts
}

func (s *MySuite) TestInvalidAddress(c *C) {
	orig := dialTimeoutSecs
	dialTimeoutSecs = 1
	defer func() {
		dialTimeoutSecs = orig
	}()
	defer func() {
		c.Assert(recover(), IsNil)
	}()
	client, _ := newClient("1.2.3.4:1234", 1, 1, connectRPC)
	// CC-1570: Client is lazy, so have to make a call to cause a panic
	client.Call("whatever", nil, nil, 1)
}
<commit_msg>adjusted sleep timeout ???<commit_after>// +build unit

package rpcutils

import (
	"net"
	"net/rpc"
	"net/rpc/jsonrpc"
	"sync"
	"testing"
	"time"

	"github.com/control-center/serviced/commons/pool"
	. "gopkg.in/check.v1"
)

func Test(t *testing.T) { TestingT(t) }

type MySuite struct{}

var (
	_             = Suite(&MySuite{})
	rtt           *RPCTestType
	rpcClient     Client
	bareRpcClient *rpc.Client
)

type RPCTestType int

func (rtt *RPCTestType) Sleep(sleep time.Duration, reply *time.Duration) error {
	time.Sleep(sleep)
	*reply = sleep
	return nil
}

type TestArgs struct {
	A string
	B int
	C bool
	D []string
}

func (rtt *RPCTestType) NilReply(arg string, _ *struct{}) error {
	return nil
}

func (rtt *RPCTestType) Echo(arg string, reply *string) error {
	*reply = arg
	return nil
}

func (rtt *RPCTestType) StructCall(arg TestArgs, reply *TestArgs) error {
	*reply = arg
	return nil
}

func (s *MySuite) SetUpSuite(c *C) {
	rtt = new(RPCTestType)
	RegisterLocal("RPCTestType", rtt)
	rpc.Register(rtt)
	rpc.HandleHTTP()
	listener, err := net.Listen("tcp", ":32111")
	if err != nil {
		c.Errorf("listen error: %s", err)
	}
	go func() {
		for {
			conn, err := listener.Accept()
			defer conn.Close()
			if err != nil {
				c.Errorf("Error accepting connections: %s", err)
			}
			go rpc.ServeCodec(jsonrpc.NewServerCodec(conn))
		}
	}()

	rpcClient, _ = newClient("localhost:32111", 1, DiscardClientTimeout, connectRPC)
	bareRpcClient, _ = connectRPC("localhost:32111")
}

func (s *MySuite) TestConcurrentTimeout(c *C) {

	sleepTime := 500 * time.Millisecond
	client, err := newClient("localhost:32111", 1, DiscardClientTimeout, connectRPC)
	c.Assert(err, IsNil)

	wg := sync.WaitGroup{}
	wg.Add(1)
	go func() {
		var reply time.Duration
		wg.Done()
		// Sleep, timeout after two. Shouldn't error.
		err := client.Call("RPCTestType.Sleep", sleepTime, &reply, 2*sleepTime)
		c.Assert(err, IsNil)
		c.Assert(reply, Equals, sleepTime)
	}()
	wg.Wait()
	var reply time.Duration
	// should time out wating for client
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, sleepTime/2)
	c.Assert(err, Equals, pool.ErrItemUnavailable)
}

func (s *MySuite) TestTimeout(c *C) {

	client, err := newClient("localhost:32111", 1, 10*time.Millisecond, connectRPC)

	sleepTime := 100 * time.Millisecond

	var reply time.Duration

	// Sleep for one second, timeout after two. Shouldn't error.
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, 2*time.Second)
	c.Assert(err, IsNil)
	c.Assert(reply, Equals, sleepTime)

	// Sleep, never timeout. Shouldn't error.
	sleepTime = sleepTime * 2
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, 0)
	c.Assert(err, IsNil)
	c.Assert(reply, Equals, sleepTime)

	// Sleep and timeout after half sleep. Should error.
	err = client.Call("RPCTestType.Sleep", &sleepTime, &reply, sleepTime/2)
	c.Assert(err, NotNil)
	c.Assert(err, ErrorMatches, "RPC call to RPCTestType.Sleep timed out after .+")
}

func (s *MySuite) TestLongCall(c *C) {
	client, err := newClient("localhost:32111", 1, 250*time.Millisecond, connectRPC)
	c.Assert(err, IsNil)

	startWg := sync.WaitGroup{}

	wg := sync.WaitGroup{}
	wg.Add(1)
	startWg.Add(1)
	go func() {
		var reply time.Duration
		startWg.Done()
		// Sleep for time , timeout after twice as much. Shouldn't error but underlying client will be invalidated
		sleepTime := 750 * time.Millisecond
		err := client.Call("RPCTestType.Sleep", sleepTime, &reply, 2*sleepTime)
		c.Assert(err, IsNil)
		c.Assert(reply, Equals, sleepTime)
		wg.Done()

	}()
	startWg.Wait()
	//after 250ms the previous call should have caused the the client to go stale
	time.Sleep(500 * time.Millisecond)
	var reply time.Duration
	sleepTime := 10 * time.Millisecond
	// should not time out wating for client
	err = client.Call("RPCTestType.Sleep", sleepTime, &reply, 20*time.Millisecond)
	c.Assert(err, IsNil)
	c.Assert(reply, Equals, sleepTime)
	wg.Wait() //wait for go routine to run asserts
}

func (s *MySuite) TestInvalidAddress(c *C) {
	orig := dialTimeoutSecs
	dialTimeoutSecs = 1
	defer func() {
		dialTimeoutSecs = orig
	}()
	defer func() {
		c.Assert(recover(), IsNil)
	}()
	client, _ := newClient("1.2.3.4:1234", 1, 1, connectRPC)
	// CC-1570: Client is lazy, so have to make a call to cause a panic
	client.Call("whatever", nil, nil, 1)
}
<|endoftext|>