<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/alyu/configparser"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
)

func main() {

	sourceProfile := flag.String("i", "default", "Source Profile")
	targetProfile := flag.String("t", "default", "Destination Profile")
	rotateKeys := flag.Bool("rotate-identity-keys", false, "Boolean flag to rotate keys of the source profile when fetching new credentials")
	overwrite := flag.Bool("o", false, "Boolean flag to overwrite profile if this is not set you can not have same source and target profile")
	printOut := flag.Bool("env", false, "Boolean flag to print commands to set environment variables")
	credFile := flag.String("c", filepath.Join(getCredentialPath(), ".aws", "credentials"), "Full path to credentials file")
	duration := flag.Int64("d", 28800, "Token Duration")
	flag.Parse()

	// If AWSMFA_ALWAYS_ROTATE is set to true, always rotate the key.
	forceRotate := strings.EqualFold(os.Getenv("AWSMFA_ALWAYS_ROTATE"), "true")

	if sourceProfile == targetProfile && !*overwrite {
		fmt.Println("Source equals target and will overwrite it you probably don't want to do this")
		return
	}

	//Get Current Credentials
	exists, err := checkProfileExists(credFile, sourceProfile)
	if err != nil || !exists {
		fmt.Println(err.Error())
		return
	}
	sess := CreateSession(sourceProfile)
	user, err := getUserMFA(sess)
	if err != nil {
		fmt.Println(err.Error())
		return
	}

	//Get MFA Code
	mfa, err := getMFACode()
	if err != nil {
		fmt.Println(err.Error())
		return
	}

	tempCreds := getSTSCredentials(sess, mfa, duration, user)
	if tempCreds != nil {
		writeNewProfile(credFile, targetProfile, sourceProfile, tempCreds)
	}

	if forceRotate || *rotateKeys {
		fmt.Print("Rotating Keys...\n")
		newKeys, err := rotateCredentialKeys(sess)
		if err != nil {
			fmt.Println(err.Error())
			return
		}
		writeNewKeys(credFile, sourceProfile, newKeys)
	}

	if *printOut {
		printNewProfile(tempCreds)
	}
}

// getMFACode prompts for MFA Token input
// It returns the value as a string and any error
func getMFACode() (string, error) {
	var mfa string
	fmt.Print("Enter MFA Token: ")
	reader := bufio.NewReader(os.Stdin)
	mfa, err := reader.ReadString('\n')
	if err != nil {
		return mfa, errors.New("failed to get token")
	}
	return strings.TrimSpace(mfa), nil
}

// CreateSession Creates AWS Session with specified profile
func CreateSession(profileName *string) *session.Session {
	profileNameValue := *profileName
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		Profile: profileNameValue,
	}))
	return sess
}

// getUserMFA takes a session
// It returns the users MFA Serial and any errors
func getUserMFA(sess *session.Session) (*string, error) {
	var newToken *string

	svc := iam.New(sess)

	params := &iam.GetUserInput{}
	resp, err := svc.GetUser(params)

	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		fmt.Println(err.Error())
		return newToken, errors.New("failed to Fetch User")
	}
	userName := *resp.User.UserName
	mfaparams := &iam.ListMFADevicesInput{
		MaxItems: aws.Int64(1),
		UserName: aws.String(userName),
	}
	mfaresp, err := svc.ListMFADevices(mfaparams)

	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		fmt.Println(err.Error())
		return newToken, errors.New("failed to Fetch User")
	}

	if len(mfaresp.MFADevices) == 0 {
		return nil, errors.New("unable to find a MFA device for this identity")
	}

	return mfaresp.MFADevices[0].SerialNumber, nil
}

// getCredentialPath returns the users home directory path as a string
func getCredentialPath() string {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	return usr.HomeDir
}

// writeNewProfile writes out the new profile keys in the credential file
// Returns nothing
func writeNewProfile(credFile *string, profileName *string, sourceProfile *string, sessionDetails *sts.GetSessionTokenOutput) {
	config, err := configparser.Read(*credFile)
	sourceSection, err := config.Section(*sourceProfile)
	region := sourceSection.ValueOf("region")
	section, err := config.Section(*profileName)
	if err != nil {
		section = config.NewSection(*profileName)
	}

	if region != "" {
		section.Add("region", region)
	}

	section.Add("aws_access_key_id", *sessionDetails.Credentials.AccessKeyId)
	section.Add("aws_secret_access_key", *sessionDetails.Credentials.SecretAccessKey)
	section.Add("aws_session_token", *sessionDetails.Credentials.SessionToken)
	section.Add("awsmfa_expiration", (*sessionDetails.Credentials.Expiration).String())
	err = configparser.Save(config, *credFile)
	if err != nil {
		log.Fatal(err)
	}
}

// printNewProfile prints out the commands to set the credentials as env variables
// Returns nothing
func printNewProfile(sessionDetails *sts.GetSessionTokenOutput) {
	fmt.Printf("AWS_ACCESS_KEY_ID=%s; export AWS_ACCESS_KEY_ID;", *sessionDetails.Credentials.AccessKeyId)
	fmt.Printf("AWS_SECRET_ACCESS_KEY=%s; export AWS_SECRET_ACCESS_KEY;", *sessionDetails.Credentials.SecretAccessKey)
	fmt.Printf("AWS_SESSION_TOKEN=%s; export AWS_SESSION_TOKEN;", *sessionDetails.Credentials.SessionToken)
	fmt.Printf("AWS_SECURITY_TOKEN=%s; export AWS_SECURITY_TOKEN;", *sessionDetails.Credentials.SessionToken)
}

// writeNewKeys replaces old aws keys in the credentials file
// Returns nothing
func writeNewKeys(credFile *string, profileName *string, newKeys *iam.CreateAccessKeyOutput) {
	config, err := configparser.Read(*credFile)
	sourceSection, err := config.Section(*profileName)
	region := sourceSection.ValueOf("region")
	section, err := config.Section(*profileName)
	if err != nil {
		section = config.NewSection(*profileName)
	}
	section.Add("region", region)
	section.Add("aws_access_key_id", *newKeys.AccessKey.AccessKeyId)
	section.Add("aws_secret_access_key", *newKeys.AccessKey.SecretAccessKey)
	err = configparser.Save(config, *credFile)
	if err != nil {
		log.Fatal(err)
	}
}

// checkProfileExists takes path to the credentials file and profile name to search for
// Returns bool and any errors
func checkProfileExists(credFile *string, profileName *string) (bool, error) {
	config, err := configparser.Read(*credFile)
	if err != nil {
		fmt.Println("Could not find credentials file")
		fmt.Println(err.Error())
		return false, err
	}
	section, err := config.Section(*profileName)
	if err != nil {
		fmt.Println("Could not find profile in credentials file")
		return false, nil
	}
	if !section.Exists("aws_access_key_id") {
		fmt.Println("Could not find access key in profile")
		return false, nil
	}

	return true, nil
}

// getSTSCredentials takes session, users inputted MFA token, duration, and device serial
// Returns GetSessionTokenOutput
func getSTSCredentials(sess *session.Session, tokenCode string, duration *int64, device *string) *sts.GetSessionTokenOutput {
	svc := sts.New(sess)
	params := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(*duration),
		SerialNumber:    aws.String(*device),
		TokenCode:       aws.String(tokenCode),
	}
	resp, err := svc.GetSessionToken(params)

	if err != nil {
		fmt.Println(err.Error())
		return nil
	}
	return resp
}

// rotateCredentialKeys takes session and will delete users existing key and create a new one
// Returns the new credentials
func rotateCredentialKeys(sess *session.Session) (*iam.CreateAccessKeyOutput, error) {
	svc := iam.New(sess)
	input := &iam.ListAccessKeysInput{}
	var currentAccessKey *iam.AccessKeyMetadata
	var createResult *iam.CreateAccessKeyOutput
	result, err := svc.ListAccessKeys(input)
	if err != nil {
		fmt.Println(err.Error())
		return nil, err
	}
	currentCreds, err := sess.Config.Credentials.Get()
	for _, accessKey := range result.AccessKeyMetadata {
		if *accessKey.AccessKeyId == currentCreds.AccessKeyID {
			currentAccessKey = accessKey
		}
	}
	if currentAccessKey != nil {
		deleteKeyInput := &iam.DeleteAccessKeyInput{
			AccessKeyId: currentAccessKey.AccessKeyId,
		}
		_, err := svc.DeleteAccessKey(deleteKeyInput)
		if err != nil {
			fmt.Println(err.Error())
			return nil, err
		}
		createKeyInput := &iam.CreateAccessKeyInput{}
		createResult, err = svc.CreateAccessKey(createKeyInput)
		if err != nil {
			fmt.Println(err.Error())
			return nil, err
		}
		fmt.Printf("Replacing %s with %s", currentCreds.AccessKeyID, *createResult.AccessKey.AccessKeyId)
		return createResult, nil
	}
	return nil, errors.New("unable to find a current access key for this Identity")
}
<commit_msg>Will skip if temp creds is nil<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/alyu/configparser"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
)

func main() {

	sourceProfile := flag.String("i", "default", "Source Profile")
	targetProfile := flag.String("t", "default", "Destination Profile")
	rotateKeys := flag.Bool("rotate-identity-keys", false, "Boolean flag to rotate keys of the source profile when fetching new credentials")
	overwrite := flag.Bool("o", false, "Boolean flag to overwrite profile if this is not set you can not have same source and target profile")
	printOut := flag.Bool("env", false, "Boolean flag to print commands to set environment variables")
	credFile := flag.String("c", filepath.Join(getCredentialPath(), ".aws", "credentials"), "Full path to credentials file")
	duration := flag.Int64("d", 28800, "Token Duration")
	flag.Parse()

	// If AWSMFA_ALWAYS_ROTATE is set to true, always rotate the key.
	forceRotate := strings.EqualFold(os.Getenv("AWSMFA_ALWAYS_ROTATE"), "true")

	if sourceProfile == targetProfile && !*overwrite {
		fmt.Println("Source equals target and will overwrite it you probably don't want to do this")
		return
	}

	//Get Current Credentials
	exists, err := checkProfileExists(credFile, sourceProfile)
	if err != nil || !exists {
		fmt.Println(err.Error())
		return
	}
	sess := CreateSession(sourceProfile)
	user, err := getUserMFA(sess)
	if err != nil {
		fmt.Println(err.Error())
		return
	}

	//Get MFA Code
	mfa, err := getMFACode()
	if err != nil {
		fmt.Println(err.Error())
		return
	}

	tempCreds := getSTSCredentials(sess, mfa, duration, user)
	if tempCreds != nil {
		writeNewProfile(credFile, targetProfile, sourceProfile, tempCreds)
	}

	if forceRotate || *rotateKeys {
		fmt.Print("Rotating Keys...\n")
		newKeys, err := rotateCredentialKeys(sess)
		if err != nil {
			fmt.Println(err.Error())
			return
		}
		writeNewKeys(credFile, sourceProfile, newKeys)
	}

	if *printOut {
		if tempCreds != nil {
			printNewProfile(tempCreds)
		}
	}
}

// getMFACode prompts for MFA Token input
// It returns the value as a string and any error
func getMFACode() (string, error) {
	var mfa string
	fmt.Print("Enter MFA Token: ")
	reader := bufio.NewReader(os.Stdin)
	mfa, err := reader.ReadString('\n')
	if err != nil {
		return mfa, errors.New("failed to get token")
	}
	return strings.TrimSpace(mfa), nil
}

// CreateSession Creates AWS Session with specified profile
func CreateSession(profileName *string) *session.Session {
	profileNameValue := *profileName
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		Profile: profileNameValue,
	}))
	return sess
}

// getUserMFA takes a session
// It returns the users MFA Serial and any errors
func getUserMFA(sess *session.Session) (*string, error) {
	var newToken *string

	svc := iam.New(sess)

	params := &iam.GetUserInput{}
	resp, err := svc.GetUser(params)

	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		fmt.Println(err.Error())
		return newToken, errors.New("failed to Fetch User")
	}
	userName := *resp.User.UserName
	mfaparams := &iam.ListMFADevicesInput{
		MaxItems: aws.Int64(1),
		UserName: aws.String(userName),
	}
	mfaresp, err := svc.ListMFADevices(mfaparams)

	if err != nil {
		// Print the error, cast err to awserr.Error to get the Code and
		// Message from an error.
		fmt.Println(err.Error())
		return newToken, errors.New("failed to Fetch User")
	}

	if len(mfaresp.MFADevices) == 0 {
		return nil, errors.New("unable to find a MFA device for this identity")
	}

	return mfaresp.MFADevices[0].SerialNumber, nil
}

// getCredentialPath returns the users home directory path as a string
func getCredentialPath() string {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	return usr.HomeDir
}

// writeNewProfile writes out the new profile keys in the credential file
// Returns nothing
func writeNewProfile(credFile *string, profileName *string, sourceProfile *string, sessionDetails *sts.GetSessionTokenOutput) {
	config, err := configparser.Read(*credFile)
	sourceSection, err := config.Section(*sourceProfile)
	region := sourceSection.ValueOf("region")
	section, err := config.Section(*profileName)
	if err != nil {
		section = config.NewSection(*profileName)
	}

	if region != "" {
		section.Add("region", region)
	}

	section.Add("aws_access_key_id", *sessionDetails.Credentials.AccessKeyId)
	section.Add("aws_secret_access_key", *sessionDetails.Credentials.SecretAccessKey)
	section.Add("aws_session_token", *sessionDetails.Credentials.SessionToken)
	section.Add("awsmfa_expiration", (*sessionDetails.Credentials.Expiration).String())
	err = configparser.Save(config, *credFile)
	if err != nil {
		log.Fatal(err)
	}
}

// printNewProfile prints out the commands to set the credentials as env variables
// Returns nothing
func printNewProfile(sessionDetails *sts.GetSessionTokenOutput) {
	fmt.Printf("AWS_ACCESS_KEY_ID=%s; export AWS_ACCESS_KEY_ID;", *sessionDetails.Credentials.AccessKeyId)
	fmt.Printf("AWS_SECRET_ACCESS_KEY=%s; export AWS_SECRET_ACCESS_KEY;", *sessionDetails.Credentials.SecretAccessKey)
	fmt.Printf("AWS_SESSION_TOKEN=%s; export AWS_SESSION_TOKEN;", *sessionDetails.Credentials.SessionToken)
	fmt.Printf("AWS_SECURITY_TOKEN=%s; export AWS_SECURITY_TOKEN;", *sessionDetails.Credentials.SessionToken)
}

// writeNewKeys replaces old aws keys in the credentials file
// Returns nothing
func writeNewKeys(credFile *string, profileName *string, newKeys *iam.CreateAccessKeyOutput) {
	config, err := configparser.Read(*credFile)
	sourceSection, err := config.Section(*profileName)
	region := sourceSection.ValueOf("region")
	section, err := config.Section(*profileName)
	if err != nil {
		section = config.NewSection(*profileName)
	}
	section.Add("region", region)
	section.Add("aws_access_key_id", *newKeys.AccessKey.AccessKeyId)
	section.Add("aws_secret_access_key", *newKeys.AccessKey.SecretAccessKey)
	err = configparser.Save(config, *credFile)
	if err != nil {
		log.Fatal(err)
	}
}

// checkProfileExists takes path to the credentials file and profile name to search for
// Returns bool and any errors
func checkProfileExists(credFile *string, profileName *string) (bool, error) {
	config, err := configparser.Read(*credFile)
	if err != nil {
		fmt.Println("Could not find credentials file")
		fmt.Println(err.Error())
		return false, err
	}
	section, err := config.Section(*profileName)
	if err != nil {
		fmt.Println("Could not find profile in credentials file")
		return false, nil
	}
	if !section.Exists("aws_access_key_id") {
		fmt.Println("Could not find access key in profile")
		return false, nil
	}

	return true, nil
}

// getSTSCredentials takes session, users inputted MFA token, duration, and device serial
// Returns GetSessionTokenOutput
func getSTSCredentials(sess *session.Session, tokenCode string, duration *int64, device *string) *sts.GetSessionTokenOutput {
	svc := sts.New(sess)
	params := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(*duration),
		SerialNumber:    aws.String(*device),
		TokenCode:       aws.String(tokenCode),
	}
	resp, err := svc.GetSessionToken(params)

	if err != nil {
		fmt.Println(err.Error())
		return nil
	}
	return resp
}

// rotateCredentialKeys takes session and will delete users existing key and create a new one
// Returns the new credentials
func rotateCredentialKeys(sess *session.Session) (*iam.CreateAccessKeyOutput, error) {
	svc := iam.New(sess)
	input := &iam.ListAccessKeysInput{}
	var currentAccessKey *iam.AccessKeyMetadata
	var createResult *iam.CreateAccessKeyOutput
	result, err := svc.ListAccessKeys(input)
	if err != nil {
		fmt.Println(err.Error())
		return nil, err
	}
	currentCreds, err := sess.Config.Credentials.Get()
	for _, accessKey := range result.AccessKeyMetadata {
		if *accessKey.AccessKeyId == currentCreds.AccessKeyID {
			currentAccessKey = accessKey
		}
	}
	if currentAccessKey != nil {
		deleteKeyInput := &iam.DeleteAccessKeyInput{
			AccessKeyId: currentAccessKey.AccessKeyId,
		}
		_, err := svc.DeleteAccessKey(deleteKeyInput)
		if err != nil {
			fmt.Println(err.Error())
			return nil, err
		}
		createKeyInput := &iam.CreateAccessKeyInput{}
		createResult, err = svc.CreateAccessKey(createKeyInput)
		if err != nil {
			fmt.Println(err.Error())
			return nil, err
		}
		fmt.Printf("Replacing %s with %s", currentCreds.AccessKeyID, *createResult.AccessKey.AccessKeyId)
		return createResult, nil
	}
	return nil, errors.New("unable to find a current access key for this Identity")
}
<|endoftext|>


<commit_before>// This is the main package for the `packer` application.

//go:generate go run ./scripts/generate-plugins.go
//go:generate go generate ./common/bootcommand/...
package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"syscall"
	"time"

	"github.com/hashicorp/go-uuid"
	"github.com/hashicorp/packer/command"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/packer/plugin"
	"github.com/hashicorp/packer/packer/tmp"
	"github.com/hashicorp/packer/version"
	"github.com/mattn/go-tty"
	"github.com/mitchellh/cli"
	"github.com/mitchellh/panicwrap"
	"github.com/mitchellh/prefixedio"
)

func main() {
	// Call realMain instead of doing the work here so we can use
	// `defer` statements within the function and have them work properly.
	// (defers aren't called with os.Exit)
	os.Exit(realMain())
}

// realMain is executed from main and returns the exit status to exit with.
func realMain() int {
	var wrapConfig panicwrap.WrapConfig
	// When following env variable is set, packer
	// wont panic wrap itself as it's already wrapped.
	// i.e.: when terraform runs it.
	wrapConfig.CookieKey = "PACKER_WRAP_COOKIE"
	wrapConfig.CookieValue = "49C22B1A-3A93-4C98-97FA-E07D18C787B5"

	if !panicwrap.Wrapped(&wrapConfig) {
		// Generate a UUID for this packer run and pass it to the environment.
		// GenerateUUID always returns a nil error (based on rand.Read) so we'll
		// just ignore it.
		UUID, _ := uuid.GenerateUUID()
		os.Setenv("PACKER_RUN_UUID", UUID)

		// Determine where logs should go in general (requested by the user)
		logWriter, err := logOutput()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup log output: %s", err)
			return 1
		}
		if logWriter == nil {
			logWriter = ioutil.Discard
		}

		packer.LogSecretFilter.SetOutput(logWriter)

		//packer.LogSecrets.

		// Disable logging here
		log.SetOutput(ioutil.Discard)

		// We always send logs to a temporary file that we use in case
		// there is a panic. Otherwise, we delete it.
		logTempFile, err := tmp.File("packer-log")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup logging tempfile: %s", err)
			return 1
		}
		defer os.Remove(logTempFile.Name())
		defer logTempFile.Close()

		// Tell the logger to log to this file
		os.Setenv(EnvLog, "")
		os.Setenv(EnvLogFile, "")

		// Setup the prefixed readers that send data properly to
		// stdout/stderr.
		doneCh := make(chan struct{})
		outR, outW := io.Pipe()
		go copyOutput(outR, doneCh)

		// Enable checkpoint for panic reporting
		if config, _ := loadConfig(); config != nil && !config.DisableCheckpoint {
			packer.CheckpointReporter = packer.NewCheckpointReporter(
				config.DisableCheckpointSignature,
			)
		}

		// Create the configuration for panicwrap and wrap our executable
		wrapConfig.Handler = panicHandler(logTempFile)
		wrapConfig.Writer = io.MultiWriter(logTempFile, &packer.LogSecretFilter)
		wrapConfig.Stdout = outW
		wrapConfig.DetectDuration = 500 * time.Millisecond
		wrapConfig.ForwardSignals = []os.Signal{syscall.SIGTERM}
		exitStatus, err := panicwrap.Wrap(&wrapConfig)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't start Packer: %s", err)
			return 1
		}

		// If >= 0, we're the parent, so just exit
		if exitStatus >= 0 {
			// Close the stdout writer so that our copy process can finish
			outW.Close()

			// Wait for the output copying to finish
			<-doneCh

			return exitStatus
		}

		// We're the child, so just close the tempfile we made in order to
		// save file handles since the tempfile is only used by the parent.
		logTempFile.Close()
	}

	// Call the real main
	return wrappedMain()
}

// wrappedMain is called only when we're wrapped by panicwrap and
// returns the exit status to exit with.
func wrappedMain() int {
	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	packer.LogSecretFilter.SetOutput(os.Stderr)
	log.SetOutput(&packer.LogSecretFilter)

	log.Printf("[INFO] Packer version: %s", version.FormattedVersion())
	log.Printf("Packer Target OS/Arch: %s %s", runtime.GOOS, runtime.GOARCH)
	log.Printf("Built with Go Version: %s", runtime.Version())

	inPlugin := os.Getenv(plugin.MagicCookieKey) == plugin.MagicCookieValue

	config, err := loadConfig()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error loading configuration: \n\n%s\n", err)
		return 1
	}
	log.Printf("Packer config: %+v", config)

	// Fire off the checkpoint.
	go runCheckpoint(config)
	if !config.DisableCheckpoint {
		packer.CheckpointReporter = packer.NewCheckpointReporter(
			config.DisableCheckpointSignature,
		)
	}

	cacheDir := os.Getenv("PACKER_CACHE_DIR")
	if cacheDir == "" {
		cacheDir = "packer_cache"
	}

	cacheDir, err = filepath.Abs(cacheDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error preparing cache directory: \n\n%s\n", err)
		return 1
	}

	log.Printf("Setting cache directory: %s", cacheDir)
	cache := &packer.FileCache{CacheDir: cacheDir}

	// Determine if we're in machine-readable mode by mucking around with
	// the arguments...
	args, machineReadable := extractMachineReadable(os.Args[1:])

	defer plugin.CleanupClients()

	var ui packer.Ui
	if machineReadable {
		ui = &packer.MachineReadableUi{
			Writer: os.Stdout,
		}

		// Set this so that we don't get colored output in our machine-
		// readable UI.
		if err := os.Setenv("PACKER_NO_COLOR", "1"); err != nil {
			fmt.Fprintf(os.Stderr, "Packer failed to initialize UI: %s\n", err)
			return 1
		}
	} else {
		var TTY packer.TTY
		if !inPlugin {
			var err error
			TTY, err = tty.Open()
			if err != nil {
				fmt.Fprintf(os.Stderr, "No tty available: %s\n", err)
			}
		}
		// Setup the UI if we're being machine-readable
		ui = &packer.BasicUi{
			Reader:      os.Stdin,
			Writer:      os.Stdout,
			ErrorWriter: os.Stdout,
			TTY:         TTY,
		}
	}
	// Create the CLI meta
	CommandMeta = &command.Meta{
		CoreConfig: &packer.CoreConfig{
			Components: packer.ComponentFinder{
				Builder:       config.LoadBuilder,
				Hook:          config.LoadHook,
				PostProcessor: config.LoadPostProcessor,
				Provisioner:   config.LoadProvisioner,
			},
			Version: version.Version,
		},
		Cache: cache,
		Ui:    ui,
	}

	cli := &cli.CLI{
		Args:         args,
		Autocomplete: true,
		Commands:     Commands,
		HelpFunc:     excludeHelpFunc(Commands, []string{"plugin"}),
		HelpWriter:   os.Stdout,
		Name:         "packer",
		Version:      version.Version,
	}

	exitCode, err := cli.Run()
	if !inPlugin {
		if err := packer.CheckpointReporter.Finalize(cli.Subcommand(), exitCode, err); err != nil {
			log.Printf("[WARN] (telemetry) Error finalizing report. This is safe to ignore. %s", err.Error())
		}
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing CLI: %s\n", err)
		return 1
	}

	return exitCode
}

// excludeHelpFunc filters commands we don't want to show from the list of
// commands displayed in packer's help text.
func excludeHelpFunc(commands map[string]cli.CommandFactory, exclude []string) cli.HelpFunc {
	// Make search slice into a map so we can use use the `if found` idiom
	// instead of a nested loop.
	var excludes = make(map[string]interface{}, len(exclude))
	for _, item := range exclude {
		excludes[item] = nil
	}

	// Create filtered list of commands
	helpCommands := []string{}
	for command := range commands {
		if _, found := excludes[command]; !found {
			helpCommands = append(helpCommands, command)
		}
	}

	return cli.FilteredHelpFunc(helpCommands, cli.BasicHelpFunc("packer"))
}

// extractMachineReadable checks the args for the machine readable
// flag and returns whether or not it is on. It modifies the args
// to remove this flag.
func extractMachineReadable(args []string) ([]string, bool) {
	for i, arg := range args {
		if arg == "-machine-readable" {
			// We found it. Slice it out.
			result := make([]string, len(args)-1)
			copy(result, args[:i])
			copy(result[i:], args[i+1:])
			return result, true
		}
	}

	return args, false
}

func loadConfig() (*config, error) {
	var config config
	config.PluginMinPort = 10000
	config.PluginMaxPort = 25000
	if err := config.Discover(); err != nil {
		return nil, err
	}

	configFilePath := os.Getenv("PACKER_CONFIG")
	if configFilePath != "" {
		log.Printf("'PACKER_CONFIG' set, loading config from environment.")
	} else {
		var err error
		configFilePath, err = packer.ConfigFile()

		if err != nil {
			log.Printf("Error detecting default config file path: %s", err)
		}
	}

	if configFilePath == "" {
		return &config, nil
	}

	log.Printf("Attempting to open config file: %s", configFilePath)
	f, err := os.Open(configFilePath)
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}

		log.Printf("[WARN] Config file doesn't exist: %s", configFilePath)
		return &config, nil
	}
	defer f.Close()

	if err := decodeConfig(f, &config); err != nil {
		return nil, err
	}

	return &config, nil
}

// copyOutput uses output prefixes to determine whether data on stdout
// should go to stdout or stderr. This is due to panicwrap using stderr
// as the log and error channel.
func copyOutput(r io.Reader, doneCh chan<- struct{}) {
	defer close(doneCh)

	pr, err := prefixedio.NewReader(r)
	if err != nil {
		panic(err)
	}

	stderrR, err := pr.Prefix(ErrorPrefix)
	if err != nil {
		panic(err)
	}
	stdoutR, err := pr.Prefix(OutputPrefix)
	if err != nil {
		panic(err)
	}
	defaultR, err := pr.Prefix("")
	if err != nil {
		panic(err)
	}

	var wg sync.WaitGroup
	wg.Add(3)
	go func() {
		defer wg.Done()
		io.Copy(os.Stderr, stderrR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(os.Stdout, stdoutR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(os.Stdout, defaultR)
	}()

	wg.Wait()
}

func init() {
	// Seed the random number generator
	rand.Seed(time.Now().UTC().UnixNano())
}
<commit_msg>Update main.go<commit_after>// This is the main package for the `packer` application.

//go:generate go run ./scripts/generate-plugins.go
//go:generate go generate ./common/bootcommand/...
package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"runtime"
	"sync"
	"syscall"
	"time"

	"github.com/hashicorp/go-uuid"
	"github.com/hashicorp/packer/command"
	"github.com/hashicorp/packer/packer"
	"github.com/hashicorp/packer/packer/plugin"
	"github.com/hashicorp/packer/packer/tmp"
	"github.com/hashicorp/packer/version"
	"github.com/mattn/go-tty"
	"github.com/mitchellh/cli"
	"github.com/mitchellh/panicwrap"
	"github.com/mitchellh/prefixedio"
)

func main() {
	// Call realMain instead of doing the work here so we can use
	// `defer` statements within the function and have them work properly.
	// (defers aren't called with os.Exit)
	os.Exit(realMain())
}

// realMain is executed from main and returns the exit status to exit with.
func realMain() int {
	var wrapConfig panicwrap.WrapConfig
	// When following env variable is set, packer
	// wont panic wrap itself as it's already wrapped.
	// i.e.: when terraform runs it.
	wrapConfig.CookieKey = "PACKER_WRAP_COOKIE"
	wrapConfig.CookieValue = "49C22B1A-3A93-4C98-97FA-E07D18C787B5"

	if !panicwrap.Wrapped(&wrapConfig) {
		// Generate a UUID for this packer run and pass it to the environment.
		// GenerateUUID always returns a nil error (based on rand.Read) so we'll
		// just ignore it.
		UUID, _ := uuid.GenerateUUID()
		os.Setenv("PACKER_RUN_UUID", UUID)

		// Determine where logs should go in general (requested by the user)
		logWriter, err := logOutput()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup log output: %s", err)
			return 1
		}
		if logWriter == nil {
			logWriter = ioutil.Discard
		}

		packer.LogSecretFilter.SetOutput(logWriter)

		//packer.LogSecrets.

		// Disable logging here
		log.SetOutput(ioutil.Discard)

		// We always send logs to a temporary file that we use in case
		// there is a panic. Otherwise, we delete it.
		logTempFile, err := tmp.File("packer-log")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't setup logging tempfile: %s", err)
			return 1
		}
		defer os.Remove(logTempFile.Name())
		defer logTempFile.Close()

		// Tell the logger to log to this file
		os.Setenv(EnvLog, "")
		os.Setenv(EnvLogFile, "")

		// Setup the prefixed readers that send data properly to
		// stdout/stderr.
		doneCh := make(chan struct{})
		outR, outW := io.Pipe()
		go copyOutput(outR, doneCh)

		// Enable checkpoint for panic reporting
		if config, _ := loadConfig(); config != nil && !config.DisableCheckpoint {
			packer.CheckpointReporter = packer.NewCheckpointReporter(
				config.DisableCheckpointSignature,
			)
		}

		// Create the configuration for panicwrap and wrap our executable
		wrapConfig.Handler = panicHandler(logTempFile)
		wrapConfig.Writer = io.MultiWriter(logTempFile, &packer.LogSecretFilter)
		wrapConfig.Stdout = outW
		wrapConfig.DetectDuration = 500 * time.Millisecond
		wrapConfig.ForwardSignals = []os.Signal{syscall.SIGTERM}
		exitStatus, err := panicwrap.Wrap(&wrapConfig)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't start Packer: %s", err)
			return 1
		}

		// If >= 0, we're the parent, so just exit
		if exitStatus >= 0 {
			// Close the stdout writer so that our copy process can finish
			outW.Close()

			// Wait for the output copying to finish
			<-doneCh

			return exitStatus
		}

		// We're the child, so just close the tempfile we made in order to
		// save file handles since the tempfile is only used by the parent.
		logTempFile.Close()
	}

	// Call the real main
	return wrappedMain()
}

// wrappedMain is called only when we're wrapped by panicwrap and
// returns the exit status to exit with.
func wrappedMain() int {
	// If there is no explicit number of Go threads to use, then set it
	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	packer.LogSecretFilter.SetOutput(os.Stderr)
	log.SetOutput(&packer.LogSecretFilter)

	log.Printf("[INFO] Packer version: %s", version.FormattedVersion())
	log.Printf("Packer Target OS/Arch: %s %s", runtime.GOOS, runtime.GOARCH)
	log.Printf("Built with Go Version: %s", runtime.Version())

	inPlugin := os.Getenv(plugin.MagicCookieKey) == plugin.MagicCookieValue

	config, err := loadConfig()
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error loading configuration: \n\n%s\n", err)
		return 1
	}
	log.Printf("Packer config: %+v", config)

	// Fire off the checkpoint.
	go runCheckpoint(config)
	if !config.DisableCheckpoint {
		packer.CheckpointReporter = packer.NewCheckpointReporter(
			config.DisableCheckpointSignature,
		)
	}

	cacheDir := os.Getenv("PACKER_CACHE_DIR")
	if cacheDir == "" {
		cacheDir = "packer_cache"
	}

	cacheDir, err = filepath.Abs(cacheDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error preparing cache directory: \n\n%s\n", err)
		return 1
	}

	log.Printf("Setting cache directory: %s", cacheDir)
	cache := &packer.FileCache{CacheDir: cacheDir}

	// Determine if we're in machine-readable mode by mucking around with
	// the arguments...
	args, machineReadable := extractMachineReadable(os.Args[1:])

	defer plugin.CleanupClients()

	var ui packer.Ui
	if machineReadable {
		// Setup the UI as we're being machine-readable
		ui = &packer.MachineReadableUi{
			Writer: os.Stdout,
		}

		// Set this so that we don't get colored output in our machine-
		// readable UI.
		if err := os.Setenv("PACKER_NO_COLOR", "1"); err != nil {
			fmt.Fprintf(os.Stderr, "Packer failed to initialize UI: %s\n", err)
			return 1
		}
	} else {
		var TTY packer.TTY
		if !inPlugin {
			var err error
			TTY, err = tty.Open()
			if err != nil {
				fmt.Fprintf(os.Stderr, "No tty available: %s\n", err)
			}
		}
		ui = &packer.BasicUi{
			Reader:      os.Stdin,
			Writer:      os.Stdout,
			ErrorWriter: os.Stdout,
			TTY:         TTY,
		}
	}
	// Create the CLI meta
	CommandMeta = &command.Meta{
		CoreConfig: &packer.CoreConfig{
			Components: packer.ComponentFinder{
				Builder:       config.LoadBuilder,
				Hook:          config.LoadHook,
				PostProcessor: config.LoadPostProcessor,
				Provisioner:   config.LoadProvisioner,
			},
			Version: version.Version,
		},
		Cache: cache,
		Ui:    ui,
	}

	cli := &cli.CLI{
		Args:         args,
		Autocomplete: true,
		Commands:     Commands,
		HelpFunc:     excludeHelpFunc(Commands, []string{"plugin"}),
		HelpWriter:   os.Stdout,
		Name:         "packer",
		Version:      version.Version,
	}

	exitCode, err := cli.Run()
	if !inPlugin {
		if err := packer.CheckpointReporter.Finalize(cli.Subcommand(), exitCode, err); err != nil {
			log.Printf("[WARN] (telemetry) Error finalizing report. This is safe to ignore. %s", err.Error())
		}
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing CLI: %s\n", err)
		return 1
	}

	return exitCode
}

// excludeHelpFunc filters commands we don't want to show from the list of
// commands displayed in packer's help text.
func excludeHelpFunc(commands map[string]cli.CommandFactory, exclude []string) cli.HelpFunc {
	// Make search slice into a map so we can use use the `if found` idiom
	// instead of a nested loop.
	var excludes = make(map[string]interface{}, len(exclude))
	for _, item := range exclude {
		excludes[item] = nil
	}

	// Create filtered list of commands
	helpCommands := []string{}
	for command := range commands {
		if _, found := excludes[command]; !found {
			helpCommands = append(helpCommands, command)
		}
	}

	return cli.FilteredHelpFunc(helpCommands, cli.BasicHelpFunc("packer"))
}

// extractMachineReadable checks the args for the machine readable
// flag and returns whether or not it is on. It modifies the args
// to remove this flag.
func extractMachineReadable(args []string) ([]string, bool) {
	for i, arg := range args {
		if arg == "-machine-readable" {
			// We found it. Slice it out.
			result := make([]string, len(args)-1)
			copy(result, args[:i])
			copy(result[i:], args[i+1:])
			return result, true
		}
	}

	return args, false
}

func loadConfig() (*config, error) {
	var config config
	config.PluginMinPort = 10000
	config.PluginMaxPort = 25000
	if err := config.Discover(); err != nil {
		return nil, err
	}

	configFilePath := os.Getenv("PACKER_CONFIG")
	if configFilePath != "" {
		log.Printf("'PACKER_CONFIG' set, loading config from environment.")
	} else {
		var err error
		configFilePath, err = packer.ConfigFile()

		if err != nil {
			log.Printf("Error detecting default config file path: %s", err)
		}
	}

	if configFilePath == "" {
		return &config, nil
	}

	log.Printf("Attempting to open config file: %s", configFilePath)
	f, err := os.Open(configFilePath)
	if err != nil {
		if !os.IsNotExist(err) {
			return nil, err
		}

		log.Printf("[WARN] Config file doesn't exist: %s", configFilePath)
		return &config, nil
	}
	defer f.Close()

	if err := decodeConfig(f, &config); err != nil {
		return nil, err
	}

	return &config, nil
}

// copyOutput uses output prefixes to determine whether data on stdout
// should go to stdout or stderr. This is due to panicwrap using stderr
// as the log and error channel.
func copyOutput(r io.Reader, doneCh chan<- struct{}) {
	defer close(doneCh)

	pr, err := prefixedio.NewReader(r)
	if err != nil {
		panic(err)
	}

	stderrR, err := pr.Prefix(ErrorPrefix)
	if err != nil {
		panic(err)
	}
	stdoutR, err := pr.Prefix(OutputPrefix)
	if err != nil {
		panic(err)
	}
	defaultR, err := pr.Prefix("")
	if err != nil {
		panic(err)
	}

	var wg sync.WaitGroup
	wg.Add(3)
	go func() {
		defer wg.Done()
		io.Copy(os.Stderr, stderrR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(os.Stdout, stdoutR)
	}()
	go func() {
		defer wg.Done()
		io.Copy(os.Stdout, defaultR)
	}()

	wg.Wait()
}

func init() {
	// Seed the random number generator
	rand.Seed(time.Now().UTC().UnixNano())
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	zones = &ZoneStore{
		store: make(map[string]Zone),
		m:     new(sync.RWMutex),
		seri:  make(map[string]uint64),
	}
	config     = &Configs{
		ip:		"10.5.4.59",
		port:	8888,
		auth:	"",
	}
	configPath	string
	listenOn    string
	recurseTo   string
	apiKey      string
	buildtime   string
	buildcommit string

)

type Configs struct {
	ip			string
	port		int
	auth		string
}

func main() {
	flag.StringVar(&configPath, "c", "10.5.4.59", "The configs ip")
	flag.StringVar(&listenOn, "l", "", "The IP to listen on (default = blank = ALL)")
	flag.StringVar(&recurseTo, "r", "", "Pass-through requests that we can't answer to other DNS server (address:port or empty=disabled)")
	flag.StringVar(&apiKey, "k", "", "API key for http notifications")
	flag.Parse()

	log.Println("firedns (2015) by Matis Hsiao is starting...")
	log.Printf("bult %s from commit %s", buildtime, buildcommit)
	config.ip = configPath
	Connect(config.ip,config.port,config.auth)
	prefetch(zones, true)

	server := &Server{
		host:     listenOn,
		port:     53,
		rTimeout: 5 * time.Second,
		wTimeout: 5 * time.Second,
		zones:    zones,
	}

	server.Run()

	go StartHTTP()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	for {
		select {
		case s := <-sig:
			log.Fatalf("Signal (%d) received, stopping\n", s)
		}
	}
}
<commit_msg>change command line args for ip,port,auth<commit_after>package main

import (
	"flag"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	zones = &ZoneStore{
		store: make(map[string]Zone),
		m:     new(sync.RWMutex),
		seri:  make(map[string]uint64),
	}
	config     = &Configs{
		ip:		"10.5.4.59",
		port:	8888,
		auth:	"",
	}
	ssdbip	string
	ssdbport	int
	ssdbauth	string
	listenOn    string
	recurseTo   string
	apiKey      string
	buildtime   string
	buildcommit string

)

type Configs struct {
	ip			string
	port		int
	auth		string
}

func main() {
	flag.StringVar(&ssdbip, "c", "10.5.4.59", "The ssdb ip")
	flag.IntVar(&ssdbport, "p", 8888, "The ssdb port")
	flag.StringVar(&ssdbauth, "a", "", "The ssdb auth password")
	flag.StringVar(&listenOn, "l", "", "The IP to listen on (default = blank = ALL)")
	flag.StringVar(&recurseTo, "r", "", "Pass-through requests that we can't answer to other DNS server (address:port or empty=disabled)")
	flag.StringVar(&apiKey, "k", "", "API key for http notifications")
	flag.Parse()

	log.Println("firedns (2015) by Matis Hsiao is starting...")
	log.Printf("bult %s from commit %s", buildtime, buildcommit)
	config.ip = ssdbip
	config.port = ssdbport
	config.auth	= ssdbauth
	Connect(config.ip,config.port,config.auth)
	prefetch(zones, true)

	server := &Server{
		host:     listenOn,
		port:     53,
		rTimeout: 5 * time.Second,
		wTimeout: 5 * time.Second,
		zones:    zones,
	}

	server.Run()

	go StartHTTP()

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	for {
		select {
		case s := <-sig:
			log.Fatalf("Signal (%d) received, stopping\n", s)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"

	"github.com/zenazn/goji"
	"github.com/zenazn/goji/web"

	"github.com/dmtar/pit/controllers"
	"github.com/dmtar/pit/middlewares"
)

func hello(c web.C, w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Hello, %s!", c.URLParams["name"])
}

func init() {
	goji.Use(middlewares.JSON)
}

func main() {
	defer goji.Serve()
	goji.Handle("/*", controllers.Root())
}
<commit_msg>Remove hello function from main.go<commit_after>package main

import (
	"github.com/zenazn/goji"

	"github.com/dmtar/pit/controllers"
	"github.com/dmtar/pit/middlewares"
)

func init() {
	goji.Use(middlewares.JSON)
}

func main() {
	defer goji.Serve()
	goji.Handle("/*", controllers.Root())
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/ninjasphere/app-scheduler/model"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/support"
)

var (
	info = ninja.LoadModuleInfo("./package.json")
)

//SchedulerApp describes the scheduler application.
type SchedulerApp struct {
	support.AppSupport
	service *SchedulerService
}

// Start is called after the ExportApp call is complete.
func (a *SchedulerApp) Start(m *model.Schedule) error {
	if a.service != nil {
		return fmt.Errorf("illegal state: scheduler is already running")
	}
	a.service = &SchedulerService{
		log:   a.Log,
		conn:  a.Conn,
		model: m,
		configStore: func(m *model.Schedule) {
			a.SendEvent("config", m)
		},
	}
	err := a.service.init(a.Info.ID)
	return err
}

// Stop the scheduler module.
func (a *SchedulerApp) Stop() error {
	var err error
	if a.service != nil {
		tmp := a.service
		a.service = nil
		err = tmp.scheduler.Stop()
	}
	return err
}

func main() {
	app := &SchedulerApp{}
	err := app.Init(info)
	if err != nil {
		app.Log.Fatalf("failed to initialize app: %v", err)
	}

	err = app.Export(app)
	if err != nil {
		app.Log.Fatalf("failed to export app: %v", err)
	}

	support.WaitUntilSignal()
}
<commit_msg>Listen to events generated by the user agent.<commit_after>package main

import (
	"fmt"
	"github.com/ninjasphere/app-scheduler/model"
	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/support"
)

var (
	info = ninja.LoadModuleInfo("./package.json")
)

//SchedulerApp describes the scheduler application.
type SchedulerApp struct {
	support.AppSupport
	service   *SchedulerService
	userAgent *userAgentListener
}

// Start is called after the ExportApp call is complete.
func (a *SchedulerApp) Start(m *model.Schedule) error {
	if a.service != nil {
		return fmt.Errorf("illegal state: scheduler is already running")
	}
	a.service = &SchedulerService{
		log:   a.Log,
		conn:  a.Conn,
		model: m,
		configStore: func(m *model.Schedule) {
			a.SendEvent("config", m)
		},
	}
	err := a.service.init(a.Info.ID)
	if err != nil {
		return err
	}
	a.userAgent = &userAgentListener{}
	err = a.userAgent.init(a.service)
	if err != nil {
		return err
	}
	return nil
}

// Stop the scheduler module.
func (a *SchedulerApp) Stop() error {
	var err error
	if a.service != nil {
		tmp := a.service
		a.service = nil
		err = tmp.scheduler.Stop()
	}
	return err
}

type userAgentListener struct {
	service *SchedulerService
	client  *ninja.ServiceClient
}

func (l *userAgentListener) init(service *SchedulerService) error {
	l.service = service
	l.client = l.service.conn.GetServiceClient("$device/:deviceId/channel/user-agent")
	l.client.OnEvent("schedule-task", l.OnTaskSchedule)
	l.client.OnEvent("cancel-task", l.OnTaskCancel)
	return nil
}

func (l *userAgentListener) OnTaskSchedule(m *model.Task, keys map[string]string) bool {
	_, err := l.service.Schedule(m)
	if err != nil {
		l.service.log.Errorf("failed while scheduling task received via user-agent notification from %s: %v: %v", keys["deviceId"], m, err)
	}
	return true
}

func (l *userAgentListener) OnTaskCancel(taskID string, keys map[string]string) bool {
	err := l.service.Cancel(taskID)
	if err != nil {
		l.service.log.Errorf("failed while canceling task received from user-agent notification from %s: %s: %v", keys["deviceId"], taskID, err)
	}
	return true
}

func main() {
	app := &SchedulerApp{}
	err := app.Init(info)
	if err != nil {
		app.Log.Fatalf("failed to initialize app: %v", err)
	}

	err = app.Export(app)
	if err != nil {
		app.Log.Fatalf("failed to export app: %v", err)
	}

	support.WaitUntilSignal()
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"flag"
	"golem/master"
	"os"
	"os/signal"
	"syscall"
)

func main() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go handleCtrlC(c)
	var (
		mode = flag.String("mode", "cli", "The operational mode of the server")
		commPort = flag.Int("commport", 10000, "The listening communication port")
		dataPort = flag.Int("dataport", 10001, "The port for data transfers")
		bindIP = flag.String("bind", "0.0.0.0", "The IP you want the server to bind to")
	)

	flag.Parse()

	if *mode == "master" {
		log.Println("master")
		master.StartMaster(*commPort, *dataPort, *bindIP, "master")
	} else if *mode == "overlord" {
		log.Println("overlord")
	} else if *mode == "peon" {
		log.Println("peon")
	} else if *mode == "cli" {
		log.Println(*mode)
	} else {
		log.Println("Invalid mode setting")
	}
}

func handleCtrlC(c chan os.Signal) {
	sig := <-c
	log.Println("\nReceived signal: ", sig)
	os.Remove("/tmp/golem_bus.sock")
	log.Println("Removed UNIX domain socket")
	os.Exit(0)
}
<commit_msg>Changed the command line parser to use the codegangsta library for command line options.  It's much nicer than the built-in one.<commit_after>package main

import (
	"log"
	"golem/master"
	"os"
	"os/signal"
	"syscall"
	"github.com/codegangsta/cli"
)

var (
	mode = "cli"
	commPort = 10000
	dataPort = 10001
	bindIP = "0.0.0.0"
)

func main() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go handleCtrlC(c)
	golem := cli.NewApp()
	golem.Name = "golem"
	golem.Author = "Charles Burton"
	golem.Email = "charles.d.burton@gmail.com"
	golem.Usage = "A fast remote execution engine written in Go"
	golem.Flags = []cli.Flag {
		cli.StringFlag {
			Name: "m, mode",
			Value: "cli",
			Usage: "Set the application mode to peon, master, or overlord.  Defaults to command line",
			Destination: &mode,
		},
		cli.IntFlag {
			Name: "d, dataport",
			Value: 10001,
			Usage: "Set the server port that data will transfer over. Default: 10001",
			Destination: &dataPort,
		},
		cli.IntFlag {
			Name: "c, commport",
			Value: 10000,
			Usage: "Set the port that golem will communicate with. Default: 10000",
			Destination: &commPort,
		},
		cli.StringFlag {
			Name: "b, bind",
			Value: "0.0.0.0",
			Usage: "Sets the IP golem will bind to.",
			Destination: &bindIP,
		},
	}

	golem.Commands = []cli.Command {

	}
	golem.Run(os.Args)

	if mode == "master" {
		master.StartMaster(commPort, dataPort, bindIP, mode)
	}
}

func handleCtrlC(c chan os.Signal) {
	sig := <-c
	log.Println("\nReceived signal: ", sig)
	os.Remove("/tmp/golem_bus.sock")
	log.Println("Removed UNIX domain socket")
	os.Exit(0)
}
<|endoftext|>


<commit_before>//--------------------------------------------
// main.go
//
// Initializes and sets up Network Observer
// for use.
//--------------------------------------------

package main

import (
	"NetworkObserver/configuration"
	"NetworkObserver/logger"
	"NetworkObserver/settings"
	"NetworkObserver/web"
	"net/http"
	"os"
)

// Random default value
var portNumber string = "5000"
var loc string = settings.AppLocation

func init() {
	// Check if the application directory exists, create it if it doesn't
	if _, err := os.Stat(loc); os.IsNotExist(err) {
		os.Mkdir(loc, 0600)
	}

	cl := loc + ".cookies"
	logger.WriteString("Removing the existing .cookies file")
	os.Remove(cl)

	// this is here because apparently auth's init() is called before this one is
	if _, err := os.Stat(cl); os.IsNotExist(err) {
		logger.WriteString("Creating a new .cookies file.")
		file, _ := os.Create(cl)
		file.Close()
	}

	if _, err := os.Stat(loc + ".password"); os.IsNotExist(err) {
		logger.WriteString("The .password file does not exist. Creating a new .password file.")
		file, _ := os.Create(loc + ".password")
		file.Close()
	}

	logger.WriteString("Setting the port number to " + configuration.GetPortNumber())
	portNumber = ":" + configuration.GetPortNumber()
}

func main() {
	// Base Pages
	http.HandleFunc("/", web.Root)
	http.HandleFunc("/checkLogin", web.CheckLogin)
	http.HandleFunc("/dashboard", web.Dashboard)
	http.HandleFunc("/createaccount", web.CreateAccount)
	http.HandleFunc("/account", web.HandleAccount)

	// Handlers
	http.HandleFunc("/saveConfig", web.SaveConfig)
	http.HandleFunc("/savetest", web.SaveTest)
	http.HandleFunc("/teststarted", web.TestStarted)
	http.HandleFunc("/logout", web.Logout)

	// Dashboard Pages
	http.HandleFunc("/dashboard/configure", web.Configure)
	http.HandleFunc("/dashboard/start_test", web.StartTest)
	http.HandleFunc("/dashboard/reports", web.Reports)

	// Start the Webserver
	http.ListenAndServe(portNumber, nil)
	// Enable SSL and HTTPS connections
	//http.ListenAndServeTLS(portNumber, "cert.pem", "key.pem", nil)
}
<commit_msg>fixed incorrect folder path for .password and .cookies<commit_after>//--------------------------------------------
// main.go
//
// Initializes and sets up Network Observer
// for use.
//--------------------------------------------

package main

import (
	"NetworkObserver/configuration"
	"NetworkObserver/logger"
	"NetworkObserver/settings"
	"NetworkObserver/web"
	"net/http"
	"os"
)

// Random default value
var portNumber string = "5000"
var loc string = settings.AppLocation

func init() {
	// Check if the application directory exists, create it if it doesn't
	if _, err := os.Stat(loc); os.IsNotExist(err) {
		os.Mkdir(loc, 0600)
	}

	cl := loc + "/.cookies"
	logger.WriteString("Removing the existing .cookies file")
	os.Remove(cl)

	// this is here because apparently auth's init() is called before this one is
	if _, err := os.Stat(cl); os.IsNotExist(err) {
		logger.WriteString("Creating a new .cookies file.")
		file, _ := os.Create(cl)
		file.Close()
	}

	if _, err := os.Stat(loc + "/.password"); os.IsNotExist(err) {
		logger.WriteString("The .password file does not exist. Creating a new .password file.")
		file, _ := os.Create(loc + ".password")
		file.Close()
	}

	logger.WriteString("Setting the port number to " + configuration.GetPortNumber())
	portNumber = ":" + configuration.GetPortNumber()
}

func main() {
	// Base Pages
	http.HandleFunc("/", web.Root)
	http.HandleFunc("/checkLogin", web.CheckLogin)
	http.HandleFunc("/dashboard", web.Dashboard)
	http.HandleFunc("/createaccount", web.CreateAccount)
	http.HandleFunc("/account", web.HandleAccount)

	// Handlers
	http.HandleFunc("/saveConfig", web.SaveConfig)
	http.HandleFunc("/savetest", web.SaveTest)
	http.HandleFunc("/teststarted", web.TestStarted)
	http.HandleFunc("/logout", web.Logout)

	// Dashboard Pages
	http.HandleFunc("/dashboard/configure", web.Configure)
	http.HandleFunc("/dashboard/start_test", web.StartTest)
	http.HandleFunc("/dashboard/reports", web.Reports)

	// Start the Webserver
	http.ListenAndServe(portNumber, nil)
	// Enable SSL and HTTPS connections
	//http.ListenAndServeTLS(portNumber, "cert.pem", "key.pem", nil)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"time"

	"github.com/jesselucas/executil"
)

const VERSION = "0.0.3"

var (
	wg sync.WaitGroup
)

type cmd struct {
	Path  string   // Path to command
	Name  string   // Name of command
	Args  []string // Args to pass to command
	Sleep int      // in Miliseconds
}

func main() {
	jsonConfigPathDefault := "commands.json"

	// Check for command line configuration flags
	var (
		configPathUsage = fmt.Sprint("Path to config json.")
		configPathPtr   = flag.String("configpath", jsonConfigPathDefault, configPathUsage)
		versionUsage    = "Prints current version" + " (v. " + VERSION + ")"
		versionPtr      = flag.Bool("version", false, versionUsage)
	)
	// Set up short hand flags
	flag.BoolVar(versionPtr, "v", false, versionUsage+" (shorthand)")
	flag.Parse()

	if *versionPtr {
		fmt.Println(VERSION)
		os.Exit(0)
	}

	// Setup commands fron commands.json
	cmds, cmdsErr := createCommandsFromJSON(*configPathPtr)
	if cmdsErr != nil {
		log.Fatal(cmdsErr)
	}

	// Run them in order
	for _, cmd := range cmds {
		sleep := time.Duration(cmd.Sleep) * time.Millisecond
		run(sleep, cmd.Path, cmd.Name, cmd.Args...)
	}

}

func run(sleep time.Duration, path string, command string, args ...string) {
	fmt.Println("[Installer] Running: ", command, args)
	err := executil.CmdStart(command, args...)
	if err != nil {
		fmt.Println(err)
	}

	if sleep != 0 {
		fmt.Printf("[Installer] Sleeping: %s for %f seconds \n", command, sleep.Seconds())
		time.Sleep(sleep)
	}
}

func createCommandsFromJSON(jsonPath string) ([]*cmd, error) {
	cmdFile, err := ioutil.ReadFile(jsonPath)
	if err != nil {
		log.Fatal("Error:", err)
		return nil, err
	}

	// Unmarshal
	var cmds []*cmd
	err = json.Unmarshal(cmdFile, &cmds)
	if err != nil {
		log.Fatalf("Can't unmarshal cmdFile.: %s", cmdFile)
		return nil, err
	}

	return cmds, nil
}
<commit_msg>Updating executil package and version number<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"time"

	"github.com/jesselucas/executil"
)

const Version = "0.0.4"

var (
	wg sync.WaitGroup
)

type cmd struct {
	Path  string   // Path to command
	Name  string   // Name of command
	Args  []string // Args to pass to command
	Sleep int      // in Miliseconds
}

func main() {
	jsonConfigPathDefault := "commands.json"

	// Check for command line configuration flags
	var (
		configPathUsage = fmt.Sprint("Path to config json.")
		configPathPtr   = flag.String("configpath", jsonConfigPathDefault, configPathUsage)
		versionUsage    = "Prints current version" + " (v. " + Version + ")"
		versionPtr      = flag.Bool("version", false, versionUsage)
	)
	// Set up short hand flags
	flag.BoolVar(versionPtr, "v", false, versionUsage+" (shorthand)")
	flag.Parse()

	if *versionPtr {
		fmt.Println(Version)
		os.Exit(0)
	}

	// Setup commands fron commands.json
	cmds, cmdsErr := createCommandsFromJSON(*configPathPtr)
	if cmdsErr != nil {
		log.Fatal(cmdsErr)
	}

	// Run them in order
	for _, cmd := range cmds {
		sleep := time.Duration(cmd.Sleep) * time.Millisecond
		run(sleep, cmd.Path, cmd.Name, cmd.Args...)
	}

}

func run(sleep time.Duration, path string, command string, args ...string) {
	fmt.Println("[Installer] Running: ", command, args)
	cmd := executil.Command(command, args...)
	cmd.OutputPrefix = command
	err := cmd.StartAndWait()
	if err != nil {
		fmt.Println(err)
	}

	if sleep != 0 {
		fmt.Printf("[Installer] Sleeping: %s for %f seconds \n", command, sleep.Seconds())
		time.Sleep(sleep)
	}
}

func createCommandsFromJSON(jsonPath string) ([]*cmd, error) {
	cmdFile, err := ioutil.ReadFile(jsonPath)
	if err != nil {
		log.Fatal("Error:", err)
		return nil, err
	}

	// Unmarshal
	var cmds []*cmd
	err = json.Unmarshal(cmdFile, &cmds)
	if err != nil {
		log.Fatalf("Can't unmarshal cmdFile.: %s", cmdFile)
		return nil, err
	}

	return cmds, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/base64"
	"log"
	"net/http"

	"goshot/utility"

	"errors"

	"github.com/charles-d-burton/gphoto2go"
	"github.com/gin-gonic/gin"
)

func main() {
	mdnsServer := utility.BroadcastServer()
	defer mdnsServer.Shutdown()
	camera := new(gphoto2go.Camera)
	err := camera.Init()
	if err > 0 {
		log.Println(err)
	}
	r := gin.Default()
	r.GET("/ping", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"message": "pong",
		})
	})

	r.GET("/shot", func(c *gin.Context) {
		camera.Interrupt()
		cameraFilePath, err := camera.TriggerCaptureToFile()
		if err == 0 {

			cameraFileReader := camera.FileReader(cameraFilePath.Folder, cameraFilePath.Name)
			defer cameraFileReader.Close()
			buf := new(bytes.Buffer)
			//buf.WriteString("\"")
			buf.ReadFrom(cameraFileReader)
			//buf.WriteString("\"")
			camera.DeleteFile(cameraFilePath.Folder, cameraFilePath.Name)
			encodedImage := base64.StdEncoding.EncodeToString(buf.Bytes())
			c.JSON(200, gin.H{
				"image": encodedImage,
			})
			buf.Reset()
			//c.Data(200, "image/jpeg", buf.Bytes())

		} else {
			log.Println(gphoto2go.CameraResultToString(err))
			c.Error(errors.New(gphoto2go.CameraResultToString(err)))
			c.JSON(http.StatusInternalServerError, gphoto2go.CameraResultToString(err))
			return
		}

	})
	r.Run()
}
<commit_msg>removed interrupt<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"log"
	"net/http"

	"goshot/utility"

	"errors"

	"github.com/charles-d-burton/gphoto2go"
	"github.com/gin-gonic/gin"
)

func main() {
	mdnsServer := utility.BroadcastServer()
	defer mdnsServer.Shutdown()
	camera := new(gphoto2go.Camera)
	err := camera.Init()
	if err > 0 {
		log.Println(err)
	}
	r := gin.Default()
	r.GET("/ping", func(c *gin.Context) {
		c.JSON(200, gin.H{
			"message": "pong",
		})
	})

	r.GET("/shot", func(c *gin.Context) {
		//camera.Interrupt()
		cameraFilePath, err := camera.TriggerCaptureToFile()
		if err == 0 {

			cameraFileReader := camera.FileReader(cameraFilePath.Folder, cameraFilePath.Name)
			defer cameraFileReader.Close()
			buf := new(bytes.Buffer)
			buf.ReadFrom(cameraFileReader)

			camera.DeleteFile(cameraFilePath.Folder, cameraFilePath.Name)
			encodedImage := base64.StdEncoding.EncodeToString(buf.Bytes())
			c.JSON(200, gin.H{
				"image": encodedImage,
			})
			buf.Reset()
			//c.Data(200, "image/jpeg", buf.Bytes())

		} else {
			log.Println(gphoto2go.CameraResultToString(err))
			c.Error(errors.New(gphoto2go.CameraResultToString(err)))
			c.JSON(http.StatusInternalServerError, gphoto2go.CameraResultToString(err))
			return
		}

	})
	r.Run()
}
<|endoftext|>


<commit_before>package main

//go:generate make

import (
	"flag"
	"github.com/gorilla/context"
	"log"
	"net"
	"net/http"
	"net/http/fcgi"
	"os"
	"path"
	"strconv"
)

var serveFcgi bool
var baseDir string
var tmpDir string
var port int
var smtpServer string
var smtpPort int
var smtpUsername string
var smtpPassword string
var reminderEmail string

func init() {
	flag.StringVar(&baseDir, "base", "./", "Base directory for server")
	flag.StringVar(&tmpDir, "tmp", "/tmp", "Directory to create temporary files in")
	flag.IntVar(&port, "port", 80, "Port to serve API/files on")
	flag.StringVar(&smtpServer, "smtp.server", "smtp.example.com", "SMTP server to send reminder emails from.")
	flag.IntVar(&smtpPort, "smtp.port", 587, "SMTP server port to connect to")
	flag.StringVar(&smtpUsername, "smtp.username", "moneygo", "SMTP username")
	flag.StringVar(&smtpPassword, "smtp.password", "password", "SMTP password")
	flag.StringVar(&reminderEmail, "email", "moneygo@example.com", "Email address to send reminder emails as.")
	flag.BoolVar(&serveFcgi, "fcgi", false, "Serve via fcgi rather than http.")
	flag.Parse()

	static_path := path.Join(baseDir, "static")

	// Ensure base directory is valid
	dir_err_str := "The base directory doesn't look like it contains the " +
		"'static' directory. Check to make sure you're passing the right " +
		"value to the -base argument."
	static_dir, err := os.Stat(static_path)
	if err != nil {
		log.Print(err)
		log.Fatal(dir_err_str)
	}
	if !static_dir.IsDir() {
		log.Fatal(dir_err_str)
	}
}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, path.Join(baseDir, "static/index.html"))
}

func staticHandler(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, path.Join(baseDir, r.URL.Path))
}

func main() {
	servemux := http.NewServeMux()
	servemux.HandleFunc("/", rootHandler)
	servemux.HandleFunc("/static/", staticHandler)
	servemux.HandleFunc("/session/", SessionHandler)
	servemux.HandleFunc("/user/", UserHandler)
	servemux.HandleFunc("/security/", SecurityHandler)
	servemux.HandleFunc("/account/", AccountHandler)
	servemux.HandleFunc("/transaction/", TransactionHandler)
	servemux.HandleFunc("/import/gnucash", GnucashImportHandler)

	listener, err := net.Listen("tcp", ":"+strconv.Itoa(port))
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Serving on port %d out of directory: %s", port, baseDir)
	if serveFcgi {
		fcgi.Serve(listener, context.ClearHandler(servemux))
	} else {
		http.Serve(listener, context.ClearHandler(servemux))
	}
}
<commit_msg>Update logging format to include source file/line<commit_after>package main

//go:generate make

import (
	"flag"
	"github.com/gorilla/context"
	"log"
	"net"
	"net/http"
	"net/http/fcgi"
	"os"
	"path"
	"strconv"
)

var serveFcgi bool
var baseDir string
var tmpDir string
var port int
var smtpServer string
var smtpPort int
var smtpUsername string
var smtpPassword string
var reminderEmail string

func init() {
	flag.StringVar(&baseDir, "base", "./", "Base directory for server")
	flag.StringVar(&tmpDir, "tmp", "/tmp", "Directory to create temporary files in")
	flag.IntVar(&port, "port", 80, "Port to serve API/files on")
	flag.StringVar(&smtpServer, "smtp.server", "smtp.example.com", "SMTP server to send reminder emails from.")
	flag.IntVar(&smtpPort, "smtp.port", 587, "SMTP server port to connect to")
	flag.StringVar(&smtpUsername, "smtp.username", "moneygo", "SMTP username")
	flag.StringVar(&smtpPassword, "smtp.password", "password", "SMTP password")
	flag.StringVar(&reminderEmail, "email", "moneygo@example.com", "Email address to send reminder emails as.")
	flag.BoolVar(&serveFcgi, "fcgi", false, "Serve via fcgi rather than http.")
	flag.Parse()

	static_path := path.Join(baseDir, "static")

	// Ensure base directory is valid
	dir_err_str := "The base directory doesn't look like it contains the " +
		"'static' directory. Check to make sure you're passing the right " +
		"value to the -base argument."
	static_dir, err := os.Stat(static_path)
	if err != nil {
		log.Print(err)
		log.Fatal(dir_err_str)
	}
	if !static_dir.IsDir() {
		log.Fatal(dir_err_str)
	}

	// Setup the logging flags to be printed
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)
}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, path.Join(baseDir, "static/index.html"))
}

func staticHandler(w http.ResponseWriter, r *http.Request) {
	http.ServeFile(w, r, path.Join(baseDir, r.URL.Path))
}

func main() {
	servemux := http.NewServeMux()
	servemux.HandleFunc("/", rootHandler)
	servemux.HandleFunc("/static/", staticHandler)
	servemux.HandleFunc("/session/", SessionHandler)
	servemux.HandleFunc("/user/", UserHandler)
	servemux.HandleFunc("/security/", SecurityHandler)
	servemux.HandleFunc("/account/", AccountHandler)
	servemux.HandleFunc("/transaction/", TransactionHandler)
	servemux.HandleFunc("/import/gnucash", GnucashImportHandler)

	listener, err := net.Listen("tcp", ":"+strconv.Itoa(port))
	if err != nil {
		log.Fatal(err)
	}

	log.Printf("Serving on port %d out of directory: %s", port, baseDir)
	if serveFcgi {
		fcgi.Serve(listener, context.ClearHandler(servemux))
	} else {
		http.Serve(listener, context.ClearHandler(servemux))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"text/template"
	"time"
)

var version = "master"

type request struct {
	command   string
	emails    string
	timeout   time.Duration
	transport string
	verbose   bool
}

type result struct {
	request request
	stdout  bytes.Buffer
	stderr  bytes.Buffer
	started time.Time
	stopped time.Time
	killed  bool
	code    int
}

func main() {
	wd, err := os.Getwd()
	if err != nil {
		log.Fatalln(err)
	}

	command := flag.String("c", "", `Command to run, like '-c "ls"'`)
	emails := flag.String("e", "", `Emails to send reports when the command fails or exceeds timeout, like '-e "john@example.com,doe@example.com"'`)
	timeout := flag.Duration("t", 1*time.Hour, `Timeout for the command, like "-t 2h", "-t 2m", or "-t 30s". After the timeout, the command is killed, defaults to 1 hour "-t 3600"`)
	transport := flag.String("p", "auto", `Transport to use, like "-p auto", "-p mail", "-p sendmail"`)
	verbose := flag.Bool("v", false, "Enable sending emails even if command is successful")
	flag.Parse()

	req := request{
		command:   *command,
		emails:    *emails,
		timeout:   *timeout,
		transport: *transport,
		verbose:   *verbose,
	}

	r := execCmd(wd, req)

	if r.killed || r.code != 0 || r.request.verbose {
		if r.request.emails == "" {
			fmt.Println(r.render().String())
		} else {
			r.sendEmail()
		}
	}
}

func execCmd(path string, req request) result {
	r := result{
		started: time.Now(),
		request: req,
	}
	cmd := exec.Command("sh", "-c", req.command)
	cmd.Dir = path
	cmd.Stdout = &r.stdout
	cmd.Stderr = &r.stderr
	cmd.Env = []string{fmt.Sprintf("HOME=%s", os.Getenv("HOME"))}
	if err := cmd.Start(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	timer := time.NewTimer(req.timeout)
	go func(timer *time.Timer, cmd *exec.Cmd) {
		for _ = range timer.C {
			r.killed = true
			if err := cmd.Process.Kill(); err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
		}
	}(timer, cmd)

	err := cmd.Wait()
	timer.Stop()
	if err != nil {
		// unsuccessful exit code?
		if exitError, ok := err.(*exec.ExitError); ok {
			r.code = exitError.Sys().(syscall.WaitStatus).ExitStatus()
		} else {
			fmt.Println(err)
			os.Exit(1)
		}
	}

	r.stopped = time.Now()

	return r
}

func (r *result) sendEmail() {
	emails := strings.Split(r.request.emails, ",")
	paths := make(map[string]string)

	if r.request.transport == "auto" {
		paths = map[string]string{"sendmail": "sendmail", "/usr/sbin/sendmail": "sendmail", "mail": "mail", "/usr/bin/mail": "mail"}
	} else if r.request.transport == "sendmail" {
		paths = map[string]string{"sendmail": "sendmail", "/usr/sbin/sendmail": "sendmail"}
	} else if r.request.transport == "mail" {
		paths = map[string]string{"mail": "mail", "/usr/bin/mail": "mail"}
	} else {
		fmt.Printf("Unsupported transport %s\n", r.request.transport)
		os.Exit(1)
	}

	var err error
	var transportType string
	var transportPath string
	for p, t := range paths {
		p, err = exec.LookPath(p)
		if err == nil {
			transportType = t
			transportPath = p
			break
		}
	}

	if transportType == "" {
		fmt.Printf("Unable to find a path for %s\n", r.request.transport)
		os.Exit(1)
	}

	if transportType == "mail" {
		for _, email := range emails {
			cmd := exec.Command(transportPath, "-s", r.subject(), strings.TrimSpace(email))
			cmd.Stdin = r.render()
			cmd.Env = []string{fmt.Sprintf("HOME=%s", os.Getenv("HOME"))}
			if err := cmd.Run(); err != nil {
				fmt.Printf("Could not send email to %s: %s\n", email, err)
				os.Exit(1)
			}
		}
		return
	}

	if transportType == "sendmail" {
		message := fmt.Sprintf("To: %s\r\nCc: %s\r\nSubject: %s\r\n\r\n%s", emails[0], strings.Join(emails[1:], ","), r.subject(), r.render().String())
		cmd := exec.Command(transportPath, "-t")
		cmd.Stdin = strings.NewReader(message)
		cmd.Env = []string{fmt.Sprintf("HOME=%s", os.Getenv("HOME"))}
		if err := cmd.Run(); err != nil {
			fmt.Printf("Could not send email to %s: %s\n", emails, err)
			os.Exit(1)
		}
		return
	}
}

func (r *result) subject() string {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "undefined"
	}

	if r.killed {
		return fmt.Sprintf("Cron on host %s: Timeout", hostname)
	}

	if r.code == 0 {
		return fmt.Sprintf("Cron on host %s: Command Successful", hostname)
	}

	return fmt.Sprintf("Cron on host %s: Failure", hostname)
}

func (r *result) title() string {
	var msg string

	if r.killed {
		msg = "Cron timeout detected"
	} else if r.code == 0 {
		msg = "Cron success"
	} else {
		msg = "Cron failure detected"
	}

	return msg + "\n" + strings.Repeat("=", len(msg))
}

func (r *result) duration() time.Duration {
	return r.stopped.Sub(r.started)
}

func (r *result) render() *bytes.Buffer {
	tpl := template.Must(template.New("email").Parse(`{{.Title}}

{{.Command}}

METADATA
--------

Exit Code: {{.Code}}
Start:     {{.Started}}
Stop:      {{.Stopped}}
Duration:  {{.Duration}}

ERROR OUTPUT
------------

{{.Stderr}}

STANDARD OUTPUT
---------------

{{.Stdout}}
`))

	data := struct {
		Title    string
		Command  string
		Started  time.Time
		Stopped  time.Time
		Duration time.Duration
		Code     int
		Stderr   string
		Stdout   string
	}{
		Title:    r.title(),
		Command:  r.request.command,
		Started:  r.started,
		Stopped:  r.stopped,
		Duration: r.duration(),
		Code:     r.code,
		Stderr:   r.stderr.String(),
		Stdout:   r.stdout.String(),
	}

	contents := bytes.Buffer{}
	if err := tpl.Execute(&contents, data); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	return &contents
}
<commit_msg>simplified code<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"text/template"
	"time"
)

var version = "master"

type request struct {
	command   string
	emails    string
	timeout   time.Duration
	transport string
	verbose   bool
}

type result struct {
	request request
	stdout  bytes.Buffer
	stderr  bytes.Buffer
	started time.Time
	stopped time.Time
	killed  bool
	code    int
}

func main() {
	wd, err := os.Getwd()
	if err != nil {
		log.Fatalln(err)
	}

	req := request{}

	flag.StringVar(&req.command, "c", "", `Command to run, like '-c "ls"'`)
	flag.StringVar(&req.emails, "e", "", `Emails to send reports when the command fails or exceeds timeout, like '-e "john@example.com,doe@example.com"'`)
	flag.DurationVar(&req.timeout, "t", 1*time.Hour, `Timeout for the command, like "-t 2h", "-t 2m", or "-t 30s". After the timeout, the command is killed, defaults to 1 hour "-t 3600"`)
	flag.StringVar(&req.transport, "p", "auto", `Transport to use, like "-p auto", "-p mail", "-p sendmail"`)
	flag.BoolVar(&req.verbose, "v", false, "Enable sending emails even if command is successful")
	flag.Parse()

	r := execCmd(wd, req)

	if r.killed || r.code != 0 || r.request.verbose {
		if r.request.emails == "" {
			fmt.Println(r.render().String())
		} else {
			r.sendEmail()
		}
	}
}

func execCmd(path string, req request) result {
	r := result{
		started: time.Now(),
		request: req,
	}
	cmd := exec.Command("sh", "-c", req.command)
	cmd.Dir = path
	cmd.Stdout = &r.stdout
	cmd.Stderr = &r.stderr
	cmd.Env = []string{fmt.Sprintf("HOME=%s", os.Getenv("HOME"))}
	if err := cmd.Start(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	timer := time.NewTimer(req.timeout)
	go func(timer *time.Timer, cmd *exec.Cmd) {
		for _ = range timer.C {
			r.killed = true
			if err := cmd.Process.Kill(); err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
		}
	}(timer, cmd)

	err := cmd.Wait()
	timer.Stop()
	if err != nil {
		// unsuccessful exit code?
		if exitError, ok := err.(*exec.ExitError); ok {
			r.code = exitError.Sys().(syscall.WaitStatus).ExitStatus()
		} else {
			fmt.Println(err)
			os.Exit(1)
		}
	}

	r.stopped = time.Now()

	return r
}

func (r *result) sendEmail() {
	emails := strings.Split(r.request.emails, ",")
	paths := make(map[string]string)

	if r.request.transport == "auto" {
		paths = map[string]string{"sendmail": "sendmail", "/usr/sbin/sendmail": "sendmail", "mail": "mail", "/usr/bin/mail": "mail"}
	} else if r.request.transport == "sendmail" {
		paths = map[string]string{"sendmail": "sendmail", "/usr/sbin/sendmail": "sendmail"}
	} else if r.request.transport == "mail" {
		paths = map[string]string{"mail": "mail", "/usr/bin/mail": "mail"}
	} else {
		fmt.Printf("Unsupported transport %s\n", r.request.transport)
		os.Exit(1)
	}

	var err error
	var transportType string
	var transportPath string
	for p, t := range paths {
		p, err = exec.LookPath(p)
		if err == nil {
			transportType = t
			transportPath = p
			break
		}
	}

	if transportType == "" {
		fmt.Printf("Unable to find a path for %s\n", r.request.transport)
		os.Exit(1)
	}

	if transportType == "mail" {
		for _, email := range emails {
			cmd := exec.Command(transportPath, "-s", r.subject(), strings.TrimSpace(email))
			cmd.Stdin = r.render()
			cmd.Env = []string{fmt.Sprintf("HOME=%s", os.Getenv("HOME"))}
			if err := cmd.Run(); err != nil {
				fmt.Printf("Could not send email to %s: %s\n", email, err)
				os.Exit(1)
			}
		}
		return
	}

	if transportType == "sendmail" {
		message := fmt.Sprintf("To: %s\r\nCc: %s\r\nSubject: %s\r\n\r\n%s", emails[0], strings.Join(emails[1:], ","), r.subject(), r.render().String())
		cmd := exec.Command(transportPath, "-t")
		cmd.Stdin = strings.NewReader(message)
		cmd.Env = []string{fmt.Sprintf("HOME=%s", os.Getenv("HOME"))}
		if err := cmd.Run(); err != nil {
			fmt.Printf("Could not send email to %s: %s\n", emails, err)
			os.Exit(1)
		}
		return
	}
}

func (r *result) subject() string {
	hostname, err := os.Hostname()
	if err != nil {
		hostname = "undefined"
	}

	if r.killed {
		return fmt.Sprintf("Cron on host %s: Timeout", hostname)
	}

	if r.code == 0 {
		return fmt.Sprintf("Cron on host %s: Command Successful", hostname)
	}

	return fmt.Sprintf("Cron on host %s: Failure", hostname)
}

func (r *result) title() string {
	var msg string

	if r.killed {
		msg = "Cron timeout detected"
	} else if r.code == 0 {
		msg = "Cron success"
	} else {
		msg = "Cron failure detected"
	}

	return msg + "\n" + strings.Repeat("=", len(msg))
}

func (r *result) duration() time.Duration {
	return r.stopped.Sub(r.started)
}

func (r *result) render() *bytes.Buffer {
	tpl := template.Must(template.New("email").Parse(`{{.Title}}

{{.Command}}

METADATA
--------

Exit Code: {{.Code}}
Start:     {{.Started}}
Stop:      {{.Stopped}}
Duration:  {{.Duration}}

ERROR OUTPUT
------------

{{.Stderr}}

STANDARD OUTPUT
---------------

{{.Stdout}}
`))

	data := struct {
		Title    string
		Command  string
		Started  time.Time
		Stopped  time.Time
		Duration time.Duration
		Code     int
		Stderr   string
		Stdout   string
	}{
		Title:    r.title(),
		Command:  r.request.command,
		Started:  r.started,
		Stopped:  r.stopped,
		Duration: r.duration(),
		Code:     r.code,
		Stderr:   r.stderr.String(),
		Stdout:   r.stdout.String(),
	}

	contents := bytes.Buffer{}
	if err := tpl.Execute(&contents, data); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	return &contents
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"strings"
	"time"

	"github.com/juju/ratelimit"
	"gopkg.in/mcuadros/go-syslog.v2"
)

// getIPs takes in an haproxy log line and returns the client IP and the CDN
// connecting IP which is in the captured request header. For us, the CDN
// connecting IP happens to be the 8th captured request header.
func getIPs(logLine string) (*net.IP, *net.IP) {
	if logLine == "" {
		return nil, nil
	}
	// This string parsing stuff was lifted from TPS
	var a, b int
	if a = strings.Index(logLine, "]:") + 3; a == -1 {
		return nil, nil
	}
	if b = strings.Index(logLine[a:], ":"); b == -1 {
		return nil, nil
	}
	clientIPString := logLine[a : a+b]
	clientIP := net.ParseIP(clientIPString)
	if clientIP == nil {
		return nil, nil
	}

	logLine = logLine[a+b:]
	// The first curly-braced block contains our request headers
	if a = strings.Index(logLine, "{"); a == -1 {
		return &clientIP, nil
	}
	if b = strings.Index(logLine[a:], "}"); b == -1 {
		return &clientIP, nil
	}
	bracketedHeaders := logLine[a : a+b]
	headers := strings.Split(bracketedHeaders, "|")
	if len(headers) < 7 {
		return &clientIP, nil
	}
	ipString := headers[7]
	cdnIP := net.ParseIP(ipString)
	if cdnIP == nil {
		return &clientIP, nil
	}
	return &clientIP, &cdnIP
}

func getLogTimestamp(logLine string) int64 {
	timestampStr := strings.Split(logLine, " ")[6]
	ts, _ := time.Parse("[02/Jan/2006:15:04:05.999]", timestampStr)

	return (ts.Unix())
}

type ipList struct {
	ips  []net.IP
	nets []net.IPNet
}

func readIPList(filename string) *ipList {
	var list ipList
	f, _ := os.Open(filename)
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Index(strings.TrimSpace(line), "#") == 0 {
			continue
		}
		if strings.TrimSpace(line) == "" {
			continue
		}
		if strings.Contains(line, "/") {
			_, n, err := net.ParseCIDR(line)
			if err != nil {
				fmt.Printf("Unable to parse CIDR.\nLine:\n%s\nError:\n%s\n", line, err)
			}
			if n != nil {
				list.nets = append(list.nets, *n)
			}
		} else {
			ip := net.ParseIP(line)
			if ip != nil {
				list.ips = append(list.ips, ip)
			} else {
				fmt.Printf("Unable to parse IP address in list: %s\n", line)
			}
		}

	}
	return &list
}

func (l *ipList) contains(checkIP *net.IP) bool {
	for _, ip := range l.ips {
		if ip.Equal(*checkIP) {
			return true
		}
	}
	for _, net := range l.nets {
		if net.Contains(*checkIP) {
			return true
		}
	}

	return false
}

func main() {
	var hits map[string]*ratelimit.Bucket
	hits = make(map[string]*ratelimit.Bucket)

	channel := make(syslog.LogPartsChannel)
	handler := syslog.NewChannelHandler(channel)

	server := syslog.NewServer()
	server.SetFormat(syslog.RFC3164)
	server.SetHandler(handler)
	if err := server.ListenUDP("0.0.0.0:514"); err != nil {
		fmt.Println(err)
		return
	}
	if err := server.Boot(); err != nil {
		fmt.Println(err)
		return
	}

	whitelist := readIPList("/etc/haproxy-shared/whitelist-ips")

	go func(channel syslog.LogPartsChannel) {
		for logParts := range channel {
			var line string
			var ok bool
			if line, ok = logParts["content"].(string); !ok || line == "" {
				continue
			}
			clientIP, cdnIP := getIPs(line)
			if cdnIP == nil || clientIP == nil {
				continue
			}
			if _, found := hits[cdnIP.String()]; !found {
				hits[cdnIP.String()] = ratelimit.NewBucket(time.Duration(3)*time.Minute, 180)
			}
			_, isSoonerThanMaxWait := hits[cdnIP.String()].TakeMaxDuration(1, 0)
			if !isSoonerThanMaxWait {
				if whitelist.contains(cdnIP) {
					//fmt.Println("but is whitelisted so we don't care.")
				} else {
					fmt.Printf("Over limit: %s\n", cdnIP.String())
				}
			}
		}
	}(channel)

	server.Wait()
}
<commit_msg>Setup some cli structure.<commit_after>package main

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"strings"
	"time"

	"github.com/juju/ratelimit"
	"github.com/urfave/cli"
	"gopkg.in/mcuadros/go-syslog.v2"
)

// getIPs takes in an haproxy log line and returns the client IP and the CDN
// connecting IP which is in the captured request header. For us, the CDN
// connecting IP happens to be the 8th captured request header.
func getIPs(logLine string) (*net.IP, *net.IP) {
	if logLine == "" {
		return nil, nil
	}
	// This string parsing stuff was lifted from TPS
	var a, b int
	if a = strings.Index(logLine, "]:") + 3; a == -1 {
		return nil, nil
	}
	if b = strings.Index(logLine[a:], ":"); b == -1 {
		return nil, nil
	}
	clientIPString := logLine[a : a+b]
	clientIP := net.ParseIP(clientIPString)
	if clientIP == nil {
		return nil, nil
	}

	logLine = logLine[a+b:]
	// The first curly-braced block contains our request headers
	if a = strings.Index(logLine, "{"); a == -1 {
		return &clientIP, nil
	}
	if b = strings.Index(logLine[a:], "}"); b == -1 {
		return &clientIP, nil
	}
	bracketedHeaders := logLine[a : a+b]
	headers := strings.Split(bracketedHeaders, "|")
	if len(headers) < 7 {
		return &clientIP, nil
	}
	ipString := headers[7]
	cdnIP := net.ParseIP(ipString)
	if cdnIP == nil {
		return &clientIP, nil
	}
	return &clientIP, &cdnIP
}

func getLogTimestamp(logLine string) int64 {
	timestampStr := strings.Split(logLine, " ")[6]
	ts, _ := time.Parse("[02/Jan/2006:15:04:05.999]", timestampStr)

	return (ts.Unix())
}

type ipList struct {
	ips  []net.IP
	nets []net.IPNet
}

func readIPList(filename string) (*ipList, error) {
	var list ipList
	f, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Index(strings.TrimSpace(line), "#") == 0 {
			continue
		}
		if strings.TrimSpace(line) == "" {
			continue
		}
		if strings.Contains(line, "/") {
			_, n, err := net.ParseCIDR(line)
			if err != nil {
				return nil, fmt.Errorf("Unable to parse CIDR.\nLine:\n%s\nError:\n%s\n", line, err)
			}
			if n != nil {
				list.nets = append(list.nets, *n)
			}
		} else {
			ip := net.ParseIP(line)
			if ip != nil {
				list.ips = append(list.ips, ip)
			} else {
				return nil, fmt.Errorf("Unable to parse IP address in list: %s\n", line)
			}
		}

	}
	return &list, nil
}

func (l *ipList) contains(checkIP *net.IP) bool {
	if l == nil {
		return false
	}
	for _, ip := range l.ips {
		if ip.Equal(*checkIP) {
			return true
		}
	}
	for _, net := range l.nets {
		if net.Contains(*checkIP) {
			return true
		}
	}

	return false
}

var hits = make(map[string]*ratelimit.Bucket)

func main() {
	app := cli.NewApp()
	app.Name = "fastly-ratelimit"

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "whitelist, w",
			Usage: "Read `FILE` containing addresses which should whitelisted from any bans.",
		},
		cli.StringFlag{
			Name:  "listen, l",
			Usage: "Specify listen `ADDRESS:PORT`.",
			Value: "0.0.0.0:514",
		},
	}
	app.Action = func(c *cli.Context) error {
		channel := make(syslog.LogPartsChannel)
		handler := syslog.NewChannelHandler(channel)

		server := syslog.NewServer()
		server.SetFormat(syslog.RFC3164)
		server.SetHandler(handler)
		if err := server.ListenUDP(c.GlobalString("listen")); err != nil {
			return cli.NewExitError(fmt.Sprintf("Unable to listen: %s\n", err), -1)
		}
		if err := server.Boot(); err != nil {
			return cli.NewExitError(fmt.Sprintf("Unable to start server: %s\n", err), -1)
		}

		var whitelist *ipList
		var err error
		whitelistFile := c.GlobalString("whitelist")
		if whitelistFile != "" {
			if whitelist, err = readIPList(whitelistFile); err != nil {
				return cli.NewExitError(fmt.Sprintf("Error reading whitelist file: %s", err), -1)
			}
		}

		go func(channel syslog.LogPartsChannel) {
			for logParts := range channel {
				var line string
				var ok bool
				if line, ok = logParts["content"].(string); !ok || line == "" {
					continue
				}
				clientIP, cdnIP := getIPs(line)
				if cdnIP == nil || clientIP == nil {
					continue
				}
				if _, found := hits[cdnIP.String()]; !found {
					hits[cdnIP.String()] = ratelimit.NewBucket(time.Duration(3)*time.Minute, 180)
				}
				_, isSoonerThanMaxWait := hits[cdnIP.String()].TakeMaxDuration(1, 0)
				if !isSoonerThanMaxWait {
					if whitelist.contains(cdnIP) {
						//fmt.Println("but is whitelisted so we don't care.")
					} else {
						fmt.Printf("Over limit: %s\n", cdnIP.String())
					}
				}
			}
		}(channel)

		server.Wait()

		return nil
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"path"
	"time"

	quic "github.com/lucas-clemente/quic-go"
	"github.com/nange/easypool"
	"github.com/nange/easyss/utils"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

func init() {
	log.SetFormatter(&log.JSONFormatter{TimestampFormat: "2006-01-02 15:04:05.000"})
	log.SetLevel(log.InfoLevel)
}

func PrintVersion() {
	const version = "RC1"
	fmt.Println("easyss version", version)
}

type Easyss struct {
	config    *Config
	localSess quic.Session
	sessChan  chan sessOpts
	tcpPool   easypool.Pool
}

func New(config *Config) (*Easyss, error) {
	ss := &Easyss{
		config:   config,
		sessChan: make(chan sessOpts, 10),
	}
	if !config.ServerModel {
		factory := func() (net.Conn, error) {
			return net.Dial("tcp", fmt.Sprintf("%s:%d", config.Server, config.ServerPort))
		}
		pconfig := &easypool.PoolConfig{
			InitialCap:  10,
			MaxCap:      50,
			MaxIdle:     10,
			Idletime:    5 * time.Minute,
			MaxLifetime: 15 * time.Minute,
			Factory:     factory,
		}
		tcppool, err := easypool.NewHeapPool(pconfig)
		if err != nil {
			return nil, err
		}
		ss.tcpPool = tcppool
	}

	go ss.sessManage()

	return ss, nil
}

func main() {
	var configFile string
	var printVer, debug bool
	var cmdConfig Config

	flag.BoolVar(&printVer, "version", false, "print version")
	flag.StringVar(&configFile, "c", "config.json", "specify config file")
	flag.StringVar(&cmdConfig.Server, "s", "", "server address")
	flag.StringVar(&cmdConfig.Password, "k", "", "password")
	flag.IntVar(&cmdConfig.ServerPort, "p", 0, "server port")
	flag.IntVar(&cmdConfig.Timeout, "t", 300, "timeout in seconds")
	flag.IntVar(&cmdConfig.LocalPort, "l", 0, "local socks5 proxy port")
	flag.StringVar(&cmdConfig.Method, "m", "", "encryption method, default: aes-256-gcm")
	flag.BoolVar(&cmdConfig.EnableQuic, "quic", false, "enable quic if set this value to be true")
	flag.BoolVar(&debug, "d", false, "print debug message")
	flag.BoolVar(&cmdConfig.ServerModel, "server", false, "server model")

	flag.Parse()

	if printVer {
		PrintVersion()
		os.Exit(0)
	}

	if debug {
		log.SetLevel(log.DebugLevel)
	}

	exists, err := utils.FileExists(configFile)
	if !exists || err != nil {
		log.Debugf("config file err:%v", err)

		binDir := path.Dir(os.Args[0])
		configFile = path.Join(binDir, "config.json")

		log.Debugf("config file not found, try config file %s", configFile)
	}

	config, err := ParseConfig(configFile)
	if err != nil {
		config = &cmdConfig
		if !os.IsNotExist(errors.Cause(err)) {
			log.Errorf("error reading %s: %+v", configFile, err)
			os.Exit(1)
		}
	} else {
		UpdateConfig(config, &cmdConfig)
	}

	if config.Method == "" {
		config.Method = "aes-256-gcm"
	}

	ss, err := New(config)
	if err != nil {
		log.Fatalf("init Easyss err:%+v", err)
	}
	if config.ServerModel {
		if config.ServerPort == 0 || config.Password == "" {
			log.Fatalln("server port and password should not empty")
		}

		ss.Remote()
	} else {
		if config.Password == "" || config.Server == "" || config.ServerPort == 0 {
			log.Fatalln("server address, server port and password should not empty")
		}

		pacChan := make(chan PACStatus, 1)

		p := NewPAC(config.LocalPort, pacChan)
		go p.Serve() // system pac configuration

		go ss.Local() // start local server

		t := NewTray(pacChan)
		t.Run() // system tray management
	}

}
<commit_msg>if do not enable quic feature, don't need to start session manage goroutine<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"path"
	"time"

	quic "github.com/lucas-clemente/quic-go"
	"github.com/nange/easypool"
	"github.com/nange/easyss/utils"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

func init() {
	log.SetFormatter(&log.JSONFormatter{TimestampFormat: "2006-01-02 15:04:05.000"})
	log.SetLevel(log.InfoLevel)
}

func PrintVersion() {
	const version = "RC1"
	fmt.Println("easyss version", version)
}

type Easyss struct {
	config    *Config
	localSess quic.Session
	sessChan  chan sessOpts
	tcpPool   easypool.Pool
}

func New(config *Config) (*Easyss, error) {
	ss := &Easyss{config: config}
	if !config.ServerModel {
		factory := func() (net.Conn, error) {
			return net.Dial("tcp", fmt.Sprintf("%s:%d", config.Server, config.ServerPort))
		}
		pconfig := &easypool.PoolConfig{
			InitialCap:  10,
			MaxCap:      50,
			MaxIdle:     10,
			Idletime:    5 * time.Minute,
			MaxLifetime: 15 * time.Minute,
			Factory:     factory,
		}
		tcppool, err := easypool.NewHeapPool(pconfig)
		if err != nil {
			return nil, err
		}
		ss.tcpPool = tcppool
	}
	if config.EnableQuic {
		ss.sessChan = make(chan sessOpts, 10)
		go ss.sessManage()
	}

	return ss, nil
}

func main() {
	var configFile string
	var printVer, debug bool
	var cmdConfig Config

	flag.BoolVar(&printVer, "version", false, "print version")
	flag.StringVar(&configFile, "c", "config.json", "specify config file")
	flag.StringVar(&cmdConfig.Server, "s", "", "server address")
	flag.StringVar(&cmdConfig.Password, "k", "", "password")
	flag.IntVar(&cmdConfig.ServerPort, "p", 0, "server port")
	flag.IntVar(&cmdConfig.Timeout, "t", 300, "timeout in seconds")
	flag.IntVar(&cmdConfig.LocalPort, "l", 0, "local socks5 proxy port")
	flag.StringVar(&cmdConfig.Method, "m", "", "encryption method, default: aes-256-gcm")
	flag.BoolVar(&cmdConfig.EnableQuic, "quic", false, "enable quic if set this value to be true")
	flag.BoolVar(&debug, "d", false, "print debug message")
	flag.BoolVar(&cmdConfig.ServerModel, "server", false, "server model")

	flag.Parse()

	if printVer {
		PrintVersion()
		os.Exit(0)
	}

	if debug {
		log.SetLevel(log.DebugLevel)
	}

	exists, err := utils.FileExists(configFile)
	if !exists || err != nil {
		log.Debugf("config file err:%v", err)

		binDir := path.Dir(os.Args[0])
		configFile = path.Join(binDir, "config.json")

		log.Debugf("config file not found, try config file %s", configFile)
	}

	config, err := ParseConfig(configFile)
	if err != nil {
		config = &cmdConfig
		if !os.IsNotExist(errors.Cause(err)) {
			log.Errorf("error reading %s: %+v", configFile, err)
			os.Exit(1)
		}
	} else {
		UpdateConfig(config, &cmdConfig)
	}

	if config.Method == "" {
		config.Method = "aes-256-gcm"
	}

	ss, err := New(config)
	if err != nil {
		log.Fatalf("init Easyss err:%+v", err)
	}
	if config.ServerModel {
		if config.ServerPort == 0 || config.Password == "" {
			log.Fatalln("server port and password should not empty")
		}

		ss.Remote()
	} else {
		if config.Password == "" || config.Server == "" || config.ServerPort == 0 {
			log.Fatalln("server address, server port and password should not empty")
		}

		pacChan := make(chan PACStatus, 1)

		p := NewPAC(config.LocalPort, pacChan)
		go p.Serve() // system pac configuration

		go ss.Local() // start local server

		t := NewTray(pacChan)
		t.Run() // system tray management
	}

}
<|endoftext|>