<commit_before>package main

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/dustin/go-humanize"
	"github.com/mbndr/logo"
	"github.com/mmmorris1975/aws-runas/lib"
	"gopkg.in/alecthomas/kingpin.v2"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

const (
	VERSION = "1.0.0-beta2"
)

var (
	listRoles    *bool
	listMfa      *bool
	showExpire   *bool
	sesCreds     *bool
	refresh      *bool
	verbose      *bool
	makeConf     *bool
	updateFlag   *bool
	profile      *string
	mfaArn       *string
	duration     *time.Duration
	roleDuration *time.Duration
	cmd          *[]string
	logLevel     = logo.WARN
	log          *logo.Logger
)

func init() {
	const (
		cmdDesc             = "Create an environment for interacting with the AWS API using an assumed role"
		durationArgDesc     = "duration of the retrieved session token"
		roleDurationArgDesc = "duration of the assume role credentials"
		listRoleArgDesc     = "list role ARNs you are able to assume"
		listMfaArgDesc      = "list the ARN of the MFA device associated with your account"
		showExpArgDesc      = "Show token expiration time"
		sesCredArgDesc      = "print eval()-able session token info, or run command using session token credentials"
		refreshArgDesc      = "force a refresh of the cached credentials"
		verboseArgDesc      = "print verbose/debug messages"
		profileArgDesc      = "name of profile, or role ARN"
		cmdArgDesc          = "command to execute using configured profile"
		mfaArnDesc          = "ARN of MFA device needed to perform Assume Role operation"
		makeConfArgDesc     = "Build an AWS extended switch-role plugin configuration for all available roles"
		updateArgDesc       = "Check for updates to aws-runas"
	)

	duration = kingpin.Flag("duration", durationArgDesc).Short('d').Duration()
	roleDuration = kingpin.Flag("role-duration", roleDurationArgDesc).Short('a').Duration()
	listRoles = kingpin.Flag("list-roles", listRoleArgDesc).Short('l').Bool()
	listMfa = kingpin.Flag("list-mfa", listMfaArgDesc).Short('m').Bool()
	showExpire = kingpin.Flag("expiration", showExpArgDesc).Short('e').Bool()
	makeConf = kingpin.Flag("make-conf", makeConfArgDesc).Short('c').Bool()
	sesCreds = kingpin.Flag("session", sesCredArgDesc).Short('s').Bool()
	refresh = kingpin.Flag("refresh", refreshArgDesc).Short('r').Bool()
	verbose = kingpin.Flag("verbose", verboseArgDesc).Short('v').Bool()
	mfaArn = kingpin.Flag("mfa-arn", mfaArnDesc).Short('M').String()
	updateFlag = kingpin.Flag("update", updateArgDesc).Short('u').Bool()

	// if AWS_PROFILE env var is NOT set, it MUST be 1st non-flag arg
	// if AWS_PROFILE env var is set, all non-flag args will be treated as cmd
	if v, ok := os.LookupEnv("AWS_PROFILE"); !ok {
		profile = kingpin.Arg("profile", profileArgDesc).String()
	} else {
		profile = aws.String(v)
	}

	cmd = CmdArg(kingpin.Arg("cmd", cmdArgDesc))

	kingpin.Version(VERSION)
	kingpin.CommandLine.VersionFlag.Short('V')
	kingpin.CommandLine.HelpFlag.Short('h')
	kingpin.CommandLine.Help = cmdDesc
}

func main() {
	// Tell kingpin to stop parsing flags once we start processing 'cmd', allows something like:
	// `aws-runas --verbose profile command -a --long_arg`
	// without needing an explicit `--` between 'profile' and 'cmd'
	kingpin.CommandLine.Interspersed(false)
	kingpin.Parse()

	if *verbose {
		logLevel = logo.DEBUG
	}
	log = logo.NewSimpleLogger(os.Stderr, logLevel, "aws-runas.main", true)

	log.Debugf("PROFILE: %s", *profile)
	sess := lib.AwsSession(*profile)

	cm, err := lib.NewAwsConfigManager(&lib.ConfigManagerOptions{LogLevel: logLevel})
	if err != nil {
		log.Fatalf("Error loading configuration: %v", err)
	}

	switch {
	case *listMfa:
		log.Debug("List MFA")
		mfa, err := lib.LookupMfa(sess)
		if err != nil {
			log.Fatalf("Error retrieving MFA info: %v", err)
		}

		for _, d := range mfa {
			fmt.Printf("%s\n", *d.SerialNumber)
		}
	case *listRoles, *makeConf:
		u := iamUser(sess)
		userName := *u.UserName

		rg := lib.NewAwsRoleGetter(sess, userName, &lib.RoleGetterOptions{LogLevel: logLevel})
		roles := rg.Roles()

		if *listRoles {
			log.Debug("List Roles")
			fmt.Printf("Available role ARNs for %s (%s)\n", userName, *u.Arn)
			for _, v := range roles {
				fmt.Printf("  %s\n", v)
			}
		}

		if *makeConf {
			log.Debug("Make Configuration Files.")
			var mfa *string
			if mfaArn != nil && len(*mfaArn) > 0 {
				// MFA arn provided by cmdline option
				mfa = mfaArn
			} else {
				m, err := lib.LookupMfa(sess)
				if err != nil {
					log.Errorf("MFA lookup failed, will not configure MFA: %v", err)
				}

				if len(m) > 0 {
					// use 1st MFA device found
					mfa = m[0].SerialNumber
				}
			}

			if err := cm.BuildConfig(roles, mfa); err != nil {
				log.Fatalf("Error building config file: %v", err)
			}
		}
	case *updateFlag:
		log.Debug("Update check")
		if err := lib.VersionCheck(VERSION); err != nil {
			log.Debugf("Error from VersionCheck(): %v", err)
		}
	default:
		p, err := awsProfile(cm, *profile)
		if err != nil {
			log.Fatalf("Error building profile: %v", err)
		}

		// Add command-line option overrides
		if duration != nil && (*duration).Nanoseconds() > 0 {
			p.SessionDuration = *duration
		}

		if roleDuration != nil && (*roleDuration).Nanoseconds() > 0 {
			p.CredDuration = *roleDuration
		}

		if mfaArn != nil && len(*mfaArn) > 0 {
			p.MfaSerial = *mfaArn
		}
		log.Debugf("RESOLVED PROFILE: %+v", p)

		opts := lib.SessionTokenProviderOptions{
			LogLevel:             logLevel,
			SessionTokenDuration: p.SessionDuration,
			MfaSerial:            p.MfaSerial,
		}
		t, err := lib.NewSessionTokenProvider(p, &opts)
		if err != nil {
			log.Fatalf("Unable to build credential provider: %v", err)
		}

		if *refresh {
			os.Remove(t.CacheFile())
		}

		if *showExpire {
			exp_t := t.ExpirationTime()
			fmt_t := exp_t.Format("2006-01-02 15:04:05")
			hmn_t := humanize.Time(exp_t)

			tense := "will expire"
			if exp_t.Before(time.Now()) {
				tense = "expired"
			}
			fmt.Fprintf(os.Stderr, "Session credentials %s on %s (%s)\n", tense, fmt_t, hmn_t)
		}

		var creds credentials.Value
		if *sesCreds || len(p.RoleArn.Resource) < 1 {
			log.Debugf("Getting SESSION TOKEN credentials")
			c := credentials.NewCredentials(t)
			creds, err = c.Get()
			if err != nil {
				log.Fatalf("Unable to get SessionToken credentials: %v", err)
			}
		} else {
			log.Debugf("Getting ASSUME ROLE credentials")
			in := assumeRoleInput(p)
			res, err := t.AssumeRole(in)
			if err != nil {
				log.Fatalf("Error doing AssumeRole: %+v", err)
			}
			c := res.Credentials
			creds = credentials.Value{
				AccessKeyID:     *c.AccessKeyId,
				SecretAccessKey: *c.SecretAccessKey,
				SessionToken:    *c.SessionToken,
				ProviderName:    "CachedCredentialsProvider",
			}
		}

		updateEnv(creds, p.Region)

		if len(*cmd) > 0 {
			c := exec.Command((*cmd)[0], (*cmd)[1:]...)
			c.Stdin = os.Stdin
			c.Stdout = os.Stdout
			c.Stderr = os.Stderr

			err := c.Run()
			if err != nil {
				log.Debug("Error running command")
				log.Fatalf("%v", err)
			}
		} else {
			printCredentials()
		}
	}
}

func assumeRoleInput(p *lib.AWSProfile) *sts.AssumeRoleInput {
	i := new(sts.AssumeRoleInput)
	if p.CredDuration.Seconds() == 0 {
		i.DurationSeconds = aws.Int64(int64(lib.ASSUME_ROLE_DEFAULT_DURATION.Seconds()))
	} else {
		i.DurationSeconds = aws.Int64(int64(p.CredDuration.Seconds()))
	}

	if len(p.RoleArn.String()) > 0 {
		i.RoleArn = aws.String(p.RoleArn.String())
	}

	if len(p.RoleSessionName) > 0 {
		i.RoleSessionName = aws.String(p.RoleSessionName)
	}

	if len(p.ExternalId) > 0 {
		i.ExternalId = aws.String(p.ExternalId)
	}

	return i
}

func updateEnv(creds credentials.Value, region string) {
	// Explicitly unset AWS_PROFILE to avoid unintended consequences
	os.Unsetenv("AWS_PROFILE")

	// Pass AWS_REGION through if it was set in our env, or found in config.
	// Ensure that called program gets the expected region
	if len(region) > 0 {
		os.Setenv("AWS_REGION", region)
	}

	os.Setenv("AWS_ACCESS_KEY_ID", creds.AccessKeyID)
	os.Setenv("AWS_SECRET_ACCESS_KEY", creds.SecretAccessKey)

	// If session token creds were returned, set them. Otherwise explicitly
	// unset them to keep the sdk from getting confused.  AFAIK, we should
	// always have SessionTokens, since our entire process revolves around them.
	// But always code defensively
	if len(creds.SessionToken) > 0 {
		os.Setenv("AWS_SESSION_TOKEN", creds.SessionToken)
		os.Setenv("AWS_SECURITY_TOKEN", creds.SessionToken)
	} else {
		os.Unsetenv("AWS_SESSION_TOKEN")
		os.Unsetenv("AWS_SECURITY_TOKEN")
	}
}

func iamUser(s *session.Session) *iam.User {
	i := iam.New(s)

	u, err := i.GetUser(new(iam.GetUserInput))
	if err != nil {
		log.Fatalf("Error getting IAM user info: %v", err)
	}

	log.Debugf("USER: %+v", u)
	return u.User
}

func awsProfile(cm lib.ConfigManager, name string) (*lib.AWSProfile, error) {
	var p *lib.AWSProfile

	a, err := arn.Parse(name)
	if err != nil {
		p, err = cm.GetProfile(aws.String(name))
		if err != nil {
			return nil, fmt.Errorf("unable to get configuration for profile '%s': %v", name, err)
		}
	} else {
		if strings.HasPrefix(a.String(), lib.IAM_ARN) {
			// Even though we were passed a role ARN, attempt profile info lookup so
			// we can capture any default configuration. (Ignore any errors). Lookup
			// default profile name first, so we don't blow up checking config files
			defProfile := session.DefaultSharedConfigProfile
			v, ok := os.LookupEnv("AWS_DEFAULT_PROFILE")
			if ok {
				defProfile = v
			}

			p, err = cm.GetProfile(aws.String(defProfile))
			if err != nil {
				return nil, err
			}
			p.RoleArn = a
		} else {
			return nil, fmt.Errorf("profile argument is not an IAM role ARN")
		}
	}

	return p, nil
}

func printCredentials() {
	format := "%s %s='%s'\n"
	exportToken := "export"
	switch runtime.GOOS {
	case "windows":
		exportToken = "set"
	}

	envVars := []string{
		"AWS_REGION",
		"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY",
		"AWS_SESSION_TOKEN", "AWS_SECURITY_TOKEN",
	}

	for _, v := range envVars {
		val, ok := os.LookupEnv(v)
		if ok {
			fmt.Printf(format, exportToken, v, val)
		}
	}
}
<commit_msg>rework logic for getting profile info when role arn is passed<commit_after>package main

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/arn"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/dustin/go-humanize"
	"github.com/mbndr/logo"
	"github.com/mmmorris1975/aws-runas/lib"
	"gopkg.in/alecthomas/kingpin.v2"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

const (
	VERSION = "1.0.0-beta2"
)

var (
	listRoles    *bool
	listMfa      *bool
	showExpire   *bool
	sesCreds     *bool
	refresh      *bool
	verbose      *bool
	makeConf     *bool
	updateFlag   *bool
	profile      *string
	mfaArn       *string
	duration     *time.Duration
	roleDuration *time.Duration
	cmd          *[]string
	logLevel     = logo.WARN
	log          *logo.Logger
)

func init() {
	const (
		cmdDesc             = "Create an environment for interacting with the AWS API using an assumed role"
		durationArgDesc     = "duration of the retrieved session token"
		roleDurationArgDesc = "duration of the assume role credentials"
		listRoleArgDesc     = "list role ARNs you are able to assume"
		listMfaArgDesc      = "list the ARN of the MFA device associated with your account"
		showExpArgDesc      = "Show token expiration time"
		sesCredArgDesc      = "print eval()-able session token info, or run command using session token credentials"
		refreshArgDesc      = "force a refresh of the cached credentials"
		verboseArgDesc      = "print verbose/debug messages"
		profileArgDesc      = "name of profile, or role ARN"
		cmdArgDesc          = "command to execute using configured profile"
		mfaArnDesc          = "ARN of MFA device needed to perform Assume Role operation"
		makeConfArgDesc     = "Build an AWS extended switch-role plugin configuration for all available roles"
		updateArgDesc       = "Check for updates to aws-runas"
	)

	duration = kingpin.Flag("duration", durationArgDesc).Short('d').Duration()
	roleDuration = kingpin.Flag("role-duration", roleDurationArgDesc).Short('a').Duration()
	listRoles = kingpin.Flag("list-roles", listRoleArgDesc).Short('l').Bool()
	listMfa = kingpin.Flag("list-mfa", listMfaArgDesc).Short('m').Bool()
	showExpire = kingpin.Flag("expiration", showExpArgDesc).Short('e').Bool()
	makeConf = kingpin.Flag("make-conf", makeConfArgDesc).Short('c').Bool()
	sesCreds = kingpin.Flag("session", sesCredArgDesc).Short('s').Bool()
	refresh = kingpin.Flag("refresh", refreshArgDesc).Short('r').Bool()
	verbose = kingpin.Flag("verbose", verboseArgDesc).Short('v').Bool()
	mfaArn = kingpin.Flag("mfa-arn", mfaArnDesc).Short('M').String()
	updateFlag = kingpin.Flag("update", updateArgDesc).Short('u').Bool()

	// if AWS_PROFILE env var is NOT set, it MUST be 1st non-flag arg
	// if AWS_PROFILE env var is set, all non-flag args will be treated as cmd
	if v, ok := os.LookupEnv("AWS_PROFILE"); !ok {
		profile = kingpin.Arg("profile", profileArgDesc).String()
	} else {
		profile = aws.String(v)
	}

	cmd = CmdArg(kingpin.Arg("cmd", cmdArgDesc))

	kingpin.Version(VERSION)
	kingpin.CommandLine.VersionFlag.Short('V')
	kingpin.CommandLine.HelpFlag.Short('h')
	kingpin.CommandLine.Help = cmdDesc
}

func main() {
	// Tell kingpin to stop parsing flags once we start processing 'cmd', allows something like:
	// `aws-runas --verbose profile command -a --long_arg`
	// without needing an explicit `--` between 'profile' and 'cmd'
	kingpin.CommandLine.Interspersed(false)
	kingpin.Parse()

	if *verbose {
		logLevel = logo.DEBUG
	}
	log = logo.NewSimpleLogger(os.Stderr, logLevel, "aws-runas.main", true)

	log.Debugf("PROFILE: %s", *profile)
	sess := lib.AwsSession(*profile)

	cm, err := lib.NewAwsConfigManager(&lib.ConfigManagerOptions{LogLevel: logLevel})
	if err != nil {
		log.Fatalf("Error loading configuration: %v", err)
	}

	switch {
	case *listMfa:
		log.Debug("List MFA")
		mfa, err := lib.LookupMfa(sess)
		if err != nil {
			log.Fatalf("Error retrieving MFA info: %v", err)
		}

		for _, d := range mfa {
			fmt.Printf("%s\n", *d.SerialNumber)
		}
	case *listRoles, *makeConf:
		u := iamUser(sess)
		userName := *u.UserName

		rg := lib.NewAwsRoleGetter(sess, userName, &lib.RoleGetterOptions{LogLevel: logLevel})
		roles := rg.Roles()

		if *listRoles {
			log.Debug("List Roles")
			fmt.Printf("Available role ARNs for %s (%s)\n", userName, *u.Arn)
			for _, v := range roles {
				fmt.Printf("  %s\n", v)
			}
		}

		if *makeConf {
			log.Debug("Make Configuration Files.")
			var mfa *string
			if mfaArn != nil && len(*mfaArn) > 0 {
				// MFA arn provided by cmdline option
				mfa = mfaArn
			} else {
				m, err := lib.LookupMfa(sess)
				if err != nil {
					log.Errorf("MFA lookup failed, will not configure MFA: %v", err)
				}

				if len(m) > 0 {
					// use 1st MFA device found
					mfa = m[0].SerialNumber
				}
			}

			if err := cm.BuildConfig(roles, mfa); err != nil {
				log.Fatalf("Error building config file: %v", err)
			}
		}
	case *updateFlag:
		log.Debug("Update check")
		if err := lib.VersionCheck(VERSION); err != nil {
			log.Debugf("Error from VersionCheck(): %v", err)
		}
	default:
		p, err := awsProfile(cm, *profile)
		if err != nil {
			log.Fatalf("Error building profile: %v", err)
		}

		// Add command-line option overrides
		if duration != nil && (*duration).Nanoseconds() > 0 {
			p.SessionDuration = *duration
		}

		if roleDuration != nil && (*roleDuration).Nanoseconds() > 0 {
			p.CredDuration = *roleDuration
		}

		if mfaArn != nil && len(*mfaArn) > 0 {
			p.MfaSerial = *mfaArn
		}
		log.Debugf("RESOLVED PROFILE: %+v", p)

		opts := lib.SessionTokenProviderOptions{
			LogLevel:             logLevel,
			SessionTokenDuration: p.SessionDuration,
			MfaSerial:            p.MfaSerial,
		}
		t, err := lib.NewSessionTokenProvider(p, &opts)
		if err != nil {
			log.Fatalf("Unable to build credential provider: %v", err)
		}

		if *refresh {
			os.Remove(t.CacheFile())
		}

		if *showExpire {
			exp_t := t.ExpirationTime()
			fmt_t := exp_t.Format("2006-01-02 15:04:05")
			hmn_t := humanize.Time(exp_t)

			tense := "will expire"
			if exp_t.Before(time.Now()) {
				tense = "expired"
			}
			fmt.Fprintf(os.Stderr, "Session credentials %s on %s (%s)\n", tense, fmt_t, hmn_t)
		}

		var creds credentials.Value
		if *sesCreds || len(p.RoleArn.Resource) < 1 {
			log.Debugf("Getting SESSION TOKEN credentials")
			c := credentials.NewCredentials(t)
			creds, err = c.Get()
			if err != nil {
				log.Fatalf("Unable to get SessionToken credentials: %v", err)
			}
		} else {
			log.Debugf("Getting ASSUME ROLE credentials")
			in := assumeRoleInput(p)
			res, err := t.AssumeRole(in)
			if err != nil {
				log.Fatalf("Error doing AssumeRole: %+v", err)
			}
			c := res.Credentials
			creds = credentials.Value{
				AccessKeyID:     *c.AccessKeyId,
				SecretAccessKey: *c.SecretAccessKey,
				SessionToken:    *c.SessionToken,
				ProviderName:    "CachedCredentialsProvider",
			}
		}

		updateEnv(creds, p.Region)

		if len(*cmd) > 0 {
			c := exec.Command((*cmd)[0], (*cmd)[1:]...)
			c.Stdin = os.Stdin
			c.Stdout = os.Stdout
			c.Stderr = os.Stderr

			err := c.Run()
			if err != nil {
				log.Debug("Error running command")
				log.Fatalf("%v", err)
			}
		} else {
			printCredentials()
		}
	}
}

func assumeRoleInput(p *lib.AWSProfile) *sts.AssumeRoleInput {
	i := new(sts.AssumeRoleInput)
	if p.CredDuration.Seconds() == 0 {
		i.DurationSeconds = aws.Int64(int64(lib.ASSUME_ROLE_DEFAULT_DURATION.Seconds()))
	} else {
		i.DurationSeconds = aws.Int64(int64(p.CredDuration.Seconds()))
	}

	if len(p.RoleArn.String()) > 0 {
		i.RoleArn = aws.String(p.RoleArn.String())
	}

	if len(p.RoleSessionName) > 0 {
		i.RoleSessionName = aws.String(p.RoleSessionName)
	}

	if len(p.ExternalId) > 0 {
		i.ExternalId = aws.String(p.ExternalId)
	}

	return i
}

func updateEnv(creds credentials.Value, region string) {
	// Explicitly unset AWS_PROFILE to avoid unintended consequences
	os.Unsetenv("AWS_PROFILE")

	// Pass AWS_REGION through if it was set in our env, or found in config.
	// Ensure that called program gets the expected region
	if len(region) > 0 {
		os.Setenv("AWS_REGION", region)
	}

	os.Setenv("AWS_ACCESS_KEY_ID", creds.AccessKeyID)
	os.Setenv("AWS_SECRET_ACCESS_KEY", creds.SecretAccessKey)

	// If session token creds were returned, set them. Otherwise explicitly
	// unset them to keep the sdk from getting confused.  AFAIK, we should
	// always have SessionTokens, since our entire process revolves around them.
	// But always code defensively
	if len(creds.SessionToken) > 0 {
		os.Setenv("AWS_SESSION_TOKEN", creds.SessionToken)
		os.Setenv("AWS_SECURITY_TOKEN", creds.SessionToken)
	} else {
		os.Unsetenv("AWS_SESSION_TOKEN")
		os.Unsetenv("AWS_SECURITY_TOKEN")
	}
}

func iamUser(s *session.Session) *iam.User {
	i := iam.New(s)

	u, err := i.GetUser(new(iam.GetUserInput))
	if err != nil {
		log.Fatalf("Error getting IAM user info: %v", err)
	}

	log.Debugf("USER: %+v", u)
	return u.User
}

func awsProfile(cm lib.ConfigManager, name string) (*lib.AWSProfile, error) {
	var p *lib.AWSProfile

	a, err := arn.Parse(name)
	if err != nil {
		p, err = cm.GetProfile(aws.String(name))
		if err != nil {
			return nil, fmt.Errorf("unable to get configuration for profile '%s': %v", name, err)
		}
	} else {
		if strings.HasPrefix(a.String(), lib.IAM_ARN) {
			// Even though we were passed a role ARN, attempt profile info lookup so
			// we can capture any default configuration. (Ignore any errors). Lookup
			// default profile name first, so we don't blow up checking config files
			defProfile := session.DefaultSharedConfigProfile
			v, ok := os.LookupEnv("AWS_DEFAULT_PROFILE")
			if ok {
				defProfile = v
			}

			// Unset AWS_PROFILE here, in case the role ARN came in via environment
			// variable, otherwise is messes up GetProfile
			os.Unsetenv("AWS_PROFILE")
			p, err = cm.GetProfile(aws.String(defProfile))
			if err != nil {
				return nil, err
			}
			p.RoleArn = a
		} else {
			return nil, fmt.Errorf("profile argument is not an IAM role ARN")
		}
	}

	return p, nil
}

func printCredentials() {
	format := "%s %s='%s'\n"
	exportToken := "export"
	switch runtime.GOOS {
	case "windows":
		exportToken = "set"
	}

	envVars := []string{
		"AWS_REGION",
		"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY",
		"AWS_SESSION_TOKEN", "AWS_SECURITY_TOKEN",
	}

	for _, v := range envVars {
		val, ok := os.LookupEnv(v)
		if ok {
			fmt.Printf(format, exportToken, v, val)
		}
	}
}
<|endoftext|>


<commit_before>package frame

var (
	V24CommonIDs = map[string]string{
		"Title":                      "TIT2",
		"Artist":                     "TPE1",
		"Album":                      "TALB",
		"Year":                       "TYER",
		"Genre":                      "TCON",
		"Attached Picture":           "APIC",
		"Unsynchronised Lyrics/Text": "USLT",
	}
)
<commit_msg>Add comment ID to V24CommonIDs<commit_after>package frame

var (
	V24CommonIDs = map[string]string{
		"Title":                      "TIT2",
		"Artist":                     "TPE1",
		"Album":                      "TALB",
		"Year":                       "TYER",
		"Genre":                      "TCON",
		"Attached Picture":           "APIC",
		"Unsynchronised Lyrics/Text": "USLT",
		"Comment":                    "COMM",
	}
)
<|endoftext|>


<commit_before>/*
Package mstore provides a means for serializing data in a format similar to raw
email or http requests. Messages compose headers and a body. The advantage to
this is Messages are both human readable and machine parsable. It is HIGHLY
recommended that each Message be stored in its own file as a typical MBOX file
uses a bare `From` line to seperate each message and mstore only uses two
blank lines.
*/
package mstore

import (
	"bytes"
	"io"
	"net/http"
	"net/mail"
)

const EOL = []byte("\r\n")

// Messages have two parts: Meta and Body. Messages very closely resemble a raw
// email message or raw HTTP request with the absence of a few distinctive
// headers.
type Message struct {
	Meta http.Header
	Body bytes.Buffer
}

// Creates a new Message
func NewMessage() *Message {
	return &Message{
		http.Header{},
		bytes.Buffer{},
	}
}

// Writes to the Message.Body can be done by Message.Write. There is no
// need to force using Message.Body.Write() on the API
func (p *Message) Write(b []byte) (int, error) {
	i, err := p.Body.Write(b)
	return i, err
}

// String provides a simple API to viewing the final payload (Headers + Body)
// as a string.
func (p *Message) String() string {
	buf := &bytes.Buffer{}
	p.WriteTo(buf)
	return buf.String()
}

// Write to w
func (p *Message) WriteTo(w io.Writer) error {
	var err error

	if err = p.Meta.WriteSubset(w, nil); err != nil {
		return err
	}

	w.Write(EOL)

	if _, err = p.Body.WriteTo(w); err != nil {
		return err
	}

	w.Write(EOL)
	return nil
}

// Read from r, parsing into Message{ Meta, Body }
func (p *Message) ReadFrom(r io.Reader) error {
	mm, err := mail.ReadMessage(r)
	if err != nil {
		return err
	}

	buf := bytes.Buffer{}
	if _, err := buf.ReadFrom(mm.Body); err != nil {
		return err
	}

	h := http.Header(mm.Header)

	p.Meta = h
	p.Body = buf
	return nil
}

// Parse from r, parsing into a new message
// func Parse(r io.Reader) *Message {
// 	m := NewMessage()
// 	m.ReadFrom(r)
// 	return m
// }
<commit_msg>tests passing<commit_after>/*
Package mstore provides a means for serializing data in a format similar to raw
email or http requests. Messages compose headers and a body. The advantage to
this is Messages are both human readable and machine parsable. It is HIGHLY
recommended that each Message be stored in its own file as a typical MBOX file
uses a bare `From` line to seperate each message and mstore only uses two
blank lines.
*/
package mstore

import (
	"bytes"
	"io"
	"net/http"
	"net/mail"
)

const EOL = "\r\n"

// Messages have two parts: Meta and Body. Messages very closely resemble a raw
// email message or raw HTTP request with the absence of a few distinctive
// headers.
type Message struct {
	Meta http.Header
	Body bytes.Buffer
}

// Creates a new Message
func NewMessage() *Message {
	return &Message{
		http.Header{},
		bytes.Buffer{},
	}
}

// Writes to the Message.Body can be done by Message.Write. There is no
// need to force using Message.Body.Write() on the API
func (p *Message) Write(b []byte) (int, error) {
	i, err := p.Body.Write(b)
	return i, err
}

// String provides a simple API to viewing the final payload (Headers + Body)
// as a string.
func (p *Message) String() string {
	buf := &bytes.Buffer{}
	p.WriteTo(buf)
	return buf.String()
}

// Write to w
func (p *Message) WriteTo(w io.Writer) error {
	var err error

	if err = p.Meta.WriteSubset(w, nil); err != nil {
		return err
	}

	w.Write([]byte(EOL))

	if _, err = p.Body.WriteTo(w); err != nil {
		return err
	}

	w.Write([]byte(EOL))
	return nil
}

// Read from r, parsing into Message{ Meta, Body }
func (p *Message) ReadFrom(r io.Reader) error {
	mm, err := mail.ReadMessage(r)
	if err != nil {
		return err
	}

	buf := bytes.Buffer{}
	if _, err := buf.ReadFrom(mm.Body); err != nil {
		return err
	}

	h := http.Header(mm.Header)

	p.Meta = h
	p.Body = buf
	return nil
}

// Parse from r, parsing into a new message
// func Parse(r io.Reader) *Message {
// 	m := NewMessage()
// 	m.ReadFrom(r)
// 	return m
// }
<|endoftext|>


<commit_before>package indexer

import (
	"bytes"
	"encoding/json"
	"io"
	"math"
	"sort"
	"strings"

	"github.com/dchest/stemmer/porter2"

	"github.com/dchest/kkr/search/indexer/tokenizer"
)

type Index struct {
	Docs       []*Document              `json:"docs"`
	Words      map[string][]interface{} `json:"words"`
	wordsToDoc map[string]map[*Document]int

	ContentWordWeight      float64 `json:"-"`
	HTMLTitleWeight        float64 `json:"-"`
	HTMLURLComponentWeight float64 `json:"-"`
}

type Document struct {
	URL   string `json:"u"`
	Title string `json:"t"`

	numWords  int `json:"-"`
	selfIndex int `json:"-"`
}

func New() *Index {
	return &Index{
		Docs:                   make([]*Document, 0),
		Words:                  make(map[string][]interface{}),
		wordsToDoc:             make(map[string]map[*Document]int), // word => doc => weight
		ContentWordWeight:      1,
		HTMLTitleWeight:        3,
		HTMLURLComponentWeight: 3,
	}
}

type docIndexAndWeight [2]int

func (n *Index) WriteJSON(w io.Writer) error {
	// Calculate minWeight (for normalization later),
	// and add word frequencies to doc.
	minWeight := math.MaxInt
	for _, m := range n.wordsToDoc {
		for doc, weight := range m {
			if weight < minWeight {
				minWeight = weight
			}
			doc.numWords += 1
		}
	}

	// Make JSON smaller:
	// Sort docs by most frequent, so that smaller
	// doc ids are taken for most frequent ones.
	sort.Slice(n.Docs, func(i, j int) bool {
		return n.Docs[i].numWords > n.Docs[j].numWords
	})

	// Add indexes to docs
	for i, d := range n.Docs {
		d.selfIndex = i
	}

	minWeight -= 1
	for word, m := range n.wordsToDoc {
		for doc, weight := range m {
			// Normalize weight.
			normWeight := weight - minWeight
			if normWeight == 1 {
				n.Words[word] = append(n.Words[word], doc.selfIndex)
			} else {
				n.Words[word] = append(n.Words[word], docIndexAndWeight{doc.selfIndex, normWeight})
			}
		}
		// Sort doc indexes for better compression.
		w := n.Words[word]
		getDocIndex := func(i int) int {
			switch v := w[i].(type) {
			case docIndexAndWeight:
				return v[0]
			case int:
				return v
			default:
				panic("wrong type in index")
			}
		}
		sort.SliceStable(w, func(i, j int) bool {
			return getDocIndex(i) < getDocIndex(j)
		})
	}
	return json.NewEncoder(w).Encode(n)
}

func (n *Index) addWord(word string, doc *Document, weight float64) {
	m := n.wordsToDoc[word]
	if m == nil {
		m = make(map[*Document]int)
		n.wordsToDoc[word] = m
	}
	m[doc] += int(weight * 100)
}

func (n *Index) newDocument(url, title string) *Document {
	doc := &Document{URL: url, Title: title}
	n.Docs = append(n.Docs, doc)
	return doc
}

func stem(word string) string {
	if strings.ContainsAny(word, "0123456789") {
		return word // don't stem words with digits
	}
	return porter2.Stemmer.Stem(word)
}

func (n *Index) addString(doc *Document, text string, wordWeight float64) {
	wordcnt := make(map[string]float64)
	tk := tokenizer.Words(text)
	for tk.Next() {
		w := normalizeWord(tk.Token())
		if len(w) < 1 || isStopWord(w) {
			continue
		}
		wordcnt[stem(w)] += wordWeight
	}
	// Scale each word weight and add it
	for w, c := range wordcnt {
		scaled := c / float64(len(wordcnt))
		n.addWord(w, doc, scaled)
	}
}

func (n *Index) AddText(url, title string, r io.Reader) error {
	var b bytes.Buffer
	if _, err := io.Copy(&b, r); err != nil {
		return err
	}
	n.addString(n.newDocument(url, title), b.String(), 1)
	return nil
}

func (n *Index) AddHTML(url string, r io.Reader) (indexed bool, err error) {
	title, content, indexable, err := parseHTML(r)
	if err != nil {
		return false, err
	}
	if !indexable {
		return false, nil
	}
	doc := n.newDocument(url, title)

	// Adjust word weight according to document level
	url = strings.TrimPrefix(url, "http://")
	url = strings.TrimPrefix(url, "https://")
	url = strings.TrimPrefix(url, "www.")
	url = strings.TrimSuffix(url, "/")

	level := float64(strings.Count(url, "/"))
	if level == 0 {
		level = 1
	}

	n.addString(doc, title, n.HTMLTitleWeight/level)
	n.addString(doc, content, 0.5+0.5*(n.ContentWordWeight/level))
	// Add URL components.
	url = strings.ReplaceAll(url, "/", " ")
	url = strings.ReplaceAll(url, "_", " ")
	url = strings.ReplaceAll(url, "-", " ")
	n.addString(doc, url, n.HTMLURLComponentWeight/level)
	return true, nil
}
<commit_msg>Make doc number sort stable<commit_after>package indexer

import (
	"bytes"
	"encoding/json"
	"io"
	"math"
	"sort"
	"strings"

	"github.com/dchest/stemmer/porter2"

	"github.com/dchest/kkr/search/indexer/tokenizer"
)

type Index struct {
	Docs       []*Document              `json:"docs"`
	Words      map[string][]interface{} `json:"words"`
	wordsToDoc map[string]map[*Document]int

	ContentWordWeight      float64 `json:"-"`
	HTMLTitleWeight        float64 `json:"-"`
	HTMLURLComponentWeight float64 `json:"-"`
}

type Document struct {
	URL   string `json:"u"`
	Title string `json:"t"`

	numWords  int `json:"-"`
	selfIndex int `json:"-"`
}

func New() *Index {
	return &Index{
		Docs:                   make([]*Document, 0),
		Words:                  make(map[string][]interface{}),
		wordsToDoc:             make(map[string]map[*Document]int), // word => doc => weight
		ContentWordWeight:      1,
		HTMLTitleWeight:        3,
		HTMLURLComponentWeight: 3,
	}
}

type docIndexAndWeight [2]int

func (n *Index) WriteJSON(w io.Writer) error {
	// Calculate minWeight (for normalization later),
	// and add word frequencies to doc.
	minWeight := math.MaxInt
	for _, m := range n.wordsToDoc {
		for doc, weight := range m {
			if weight < minWeight {
				minWeight = weight
			}
			doc.numWords += 1
		}
	}

	// Make JSON smaller:
	// Sort docs by most frequent, so that smaller
	// doc ids are taken for most frequent ones.
	sort.SliceStable(n.Docs, func(i, j int) bool {
		return n.Docs[i].numWords > n.Docs[j].numWords
	})

	// Add indexes to docs
	for i, d := range n.Docs {
		d.selfIndex = i
	}

	minWeight -= 1
	for word, m := range n.wordsToDoc {
		for doc, weight := range m {
			// Normalize weight.
			normWeight := weight - minWeight
			if normWeight == 1 {
				n.Words[word] = append(n.Words[word], doc.selfIndex)
			} else {
				n.Words[word] = append(n.Words[word], docIndexAndWeight{doc.selfIndex, normWeight})
			}
		}
		// Sort doc indexes for better compression.
		w := n.Words[word]
		getDocIndex := func(i int) int {
			switch v := w[i].(type) {
			case docIndexAndWeight:
				return v[0]
			case int:
				return v
			default:
				panic("wrong type in index")
			}
		}
		sort.SliceStable(w, func(i, j int) bool {
			return getDocIndex(i) < getDocIndex(j)
		})
	}
	return json.NewEncoder(w).Encode(n)
}

func (n *Index) addWord(word string, doc *Document, weight float64) {
	m := n.wordsToDoc[word]
	if m == nil {
		m = make(map[*Document]int)
		n.wordsToDoc[word] = m
	}
	m[doc] += int(weight * 100)
}

func (n *Index) newDocument(url, title string) *Document {
	doc := &Document{URL: url, Title: title}
	n.Docs = append(n.Docs, doc)
	return doc
}

func stem(word string) string {
	if strings.ContainsAny(word, "0123456789") {
		return word // don't stem words with digits
	}
	return porter2.Stemmer.Stem(word)
}

func (n *Index) addString(doc *Document, text string, wordWeight float64) {
	wordcnt := make(map[string]float64)
	tk := tokenizer.Words(text)
	for tk.Next() {
		w := normalizeWord(tk.Token())
		if len(w) < 1 || isStopWord(w) {
			continue
		}
		wordcnt[stem(w)] += wordWeight
	}
	// Scale each word weight and add it
	for w, c := range wordcnt {
		scaled := c / float64(len(wordcnt))
		n.addWord(w, doc, scaled)
	}
}

func (n *Index) AddText(url, title string, r io.Reader) error {
	var b bytes.Buffer
	if _, err := io.Copy(&b, r); err != nil {
		return err
	}
	n.addString(n.newDocument(url, title), b.String(), 1)
	return nil
}

func (n *Index) AddHTML(url string, r io.Reader) (indexed bool, err error) {
	title, content, indexable, err := parseHTML(r)
	if err != nil {
		return false, err
	}
	if !indexable {
		return false, nil
	}
	doc := n.newDocument(url, title)

	// Adjust word weight according to document level
	url = strings.TrimPrefix(url, "http://")
	url = strings.TrimPrefix(url, "https://")
	url = strings.TrimPrefix(url, "www.")
	url = strings.TrimSuffix(url, "/")

	level := float64(strings.Count(url, "/"))
	if level == 0 {
		level = 1
	}

	n.addString(doc, title, n.HTMLTitleWeight/level)
	n.addString(doc, content, 0.5+0.5*(n.ContentWordWeight/level))
	// Add URL components.
	url = strings.ReplaceAll(url, "/", " ")
	url = strings.ReplaceAll(url, "_", " ")
	url = strings.ReplaceAll(url, "-", " ")
	n.addString(doc, url, n.HTMLURLComponentWeight/level)
	return true, nil
}
<|endoftext|>


<commit_before>package geo

import (
	_ "database/sql"
	"fmt"
	"github.com/erikstmartin/go-testdb"
	"os"
	"strconv"
	"testing"
)

// TODO This paticular test is just one big integration for using the entire library.
//      Should seperate this out into sperate tests once I determine an effective
//      And reasonable way to test formulae and configuration handling.
// @spec: golang-geo should
//   - Should correctly return a set of [lat, lng] within a certain radius
func TestPointsWithinRadiusIntegration(t *testing.T) {
	// TODO Determine if we actually need to test SQL logic across databases.
	dbEnv := os.Getenv("DB")
	if dbEnv == "mock" {
		stubPointsWithinRadiusQueries()
	}

	s, sqlErr := HandleWithSQL()

	if sqlErr != nil {
		t.Error("ERROR: %s", sqlErr)
	}

	// SFO
	origin := &Point{37.619002, -122.37484}

	// Straight North
	bearing := 0.0

	// Make a point that is 1 meter within the desired radius
	in_point := origin.PointAtDistanceAndBearing(7.999, bearing)
	s.sqlConn.Exec(fmt.Sprintf("INSERT INTO points(lat, lng) VALUES(%f, %f);", in_point.lat, in_point.lng))

	// Make a point that is 1 meter outsied of the desired radius
	out_point := origin.PointAtDistanceAndBearing(8.001, bearing)
	s.sqlConn.Exec(fmt.Sprintf("INSERT INTO points(lat, lng) VALUES(%f, %f);", out_point.lat, out_point.lng))

	// Should only get the first point
	_, err := s.PointsWithinRadius(origin, 8)
	if err != nil {
		panic(err)
	}

	// TODO Write a test to check for expected results of PointAtDistanceAndBearing

	// Should get both the first point and second point
	_, err2 := s.PointsWithinRadius(origin, 9)
	if err2 != nil {
		panic(err2)
	}

	// TODO Write a test to check for expected results of PointAtDistanceAndBearing

	// Clear Test DB
	FlushTestDB(s)
}

// TODO Test sql configuration
// TODO Test Great Circle Distance
// TODO Test Point At Distance And Bearing

func FlushTestDB(s *SQLMapper) {
	s.sqlConn.Exec("DELETE FROM points;")
}

// Taken from: http://play.golang.org/p/cwJj8ZJUhl
func RoundFloat(x float64, prec int) float64 {
	frep := strconv.FormatFloat(x, 'g', prec, 64)
	f, _ := strconv.ParseFloat(frep, 64)
	return f
}

func stubPointsWithinRadiusQueries() {
	insideRangeQuery := "SELECT * FROM points a WHERE acos(sin(radians(37.619002)) * sin(radians(a.lat)) + cos(radians(37.619002)) * cos(radians(a.lat)) * cos(radians(a.lng) - radians(-122.374840))) * %!f(int=6371) <= 8.000000"
	testdb.StubQuery(insideRangeQuery, nil)

	outsideRangeQuery := "SELECT * FROM points a WHERE acos(sin(radians(37.619002)) * sin(radians(a.lat)) + cos(radians(37.619002)) * cos(radians(a.lat)) * cos(radians(a.lng) - radians(-122.374840))) * %!f(int=6371) <= 9.000000"
	testdb.StubQuery(outsideRangeQuery, nil)
}
<commit_msg>Attempting to use a more DB agnostic stubbed out query for testing<commit_after>package geo

import (
	_ "database/sql"
	"fmt"
	"github.com/erikstmartin/go-testdb"
	"os"
	"strconv"
	"testing"
)

// TODO This paticular test is just one big integration for using the entire library.
//      Should seperate this out into sperate tests once I determine an effective
//      And reasonable way to test formulae and configuration handling.
// @spec: golang-geo should
//   - Should correctly return a set of [lat, lng] within a certain radius
func TestPointsWithinRadiusIntegration(t *testing.T) {
	// TODO Determine if we actually need to test SQL logic across databases.
	dbEnv := os.Getenv("DB")
	if dbEnv == "mock" {
		stubPointsWithinRadiusQueries()
	}

	s, sqlErr := HandleWithSQL()

	if sqlErr != nil {
		t.Error("ERROR: %s", sqlErr)
	}

	// SFO
	origin := &Point{37.619002, -122.37484}

	// Straight North
	bearing := 0.0

	// Make a point that is 1 meter within the desired radius
	in_point := origin.PointAtDistanceAndBearing(7.999, bearing)
	s.sqlConn.Exec(fmt.Sprintf("INSERT INTO points(lat, lng) VALUES(%f, %f);", in_point.lat, in_point.lng))

	// Make a point that is 1 meter outsied of the desired radius
	out_point := origin.PointAtDistanceAndBearing(8.001, bearing)
	s.sqlConn.Exec(fmt.Sprintf("INSERT INTO points(lat, lng) VALUES(%f, %f);", out_point.lat, out_point.lng))

	// Should only get the first point
	_, err := s.PointsWithinRadius(origin, 8)
	if err != nil {
		panic(err)
	}

	// TODO Write a test to check for expected results of PointAtDistanceAndBearing

	// Should get both the first point and second point
	_, err2 := s.PointsWithinRadius(origin, 9)
	if err2 != nil {
		panic(err2)
	}

	// TODO Write a test to check for expected results of PointAtDistanceAndBearing

	// Clear Test DB
	FlushTestDB(s)
}

// TODO Test sql configuration
// TODO Test Great Circle Distance
// TODO Test Point At Distance And Bearing

func FlushTestDB(s *SQLMapper) {
	s.sqlConn.Exec("DELETE FROM points;")
}

// Taken from: http://play.golang.org/p/cwJj8ZJUhl
func RoundFloat(x float64, prec int) float64 {
	frep := strconv.FormatFloat(x, 'g', prec, 64)
	f, _ := strconv.ParseFloat(frep, 64)
	return f
}

func stubPointsWithinRadiusQueries() {
	insideRangeQuery := fmt.Sprintf("SELECT * FROM points a WHERE acos(sin(radians(37.619002)) * sin(radians(a.lat)) + cos(radians(37.619002)) * cos(radians(a.lat)) * cos(radians(a.lng) - radians(-122.374840))) * %f <= 8.000000", EARTH_RADIUS)
	testdb.StubQuery(insideRangeQuery, nil)

	outsideRangeQuery := fmt.Sprintf("SELECT * FROM points a WHERE acos(sin(radians(37.619002)) * sin(radians(a.lat)) + cos(radians(37.619002)) * cos(radians(a.lat)) * cos(radians(a.lng) - radians(-122.374840))) * %f <= 9.000000", EARTH_RADIUS)
	testdb.StubQuery(outsideRangeQuery, nil)
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2015, Ian Chan <icha024@gmail.com>.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of Redis nor the names of its contributors may be used
 *    to endorse or promote products derived from this software without
 *    specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
 * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
 * THE POSSIBILITY OF SUCH DAMAGE.
 */
package geoindex

/*
#include "geohash.h"
#include "geohash_helper.h"
#include <math.h>
#cgo LDFLAGS: -lm
*/
import "C"
import (
	"errors"
	"flag"
	"github.com/cznic/sortutil"
	"log"
	"math"
	"sort"
)

type GeoData struct {
	// Must be unique
	Id int
	// Generated automatically
	GeoHash uint64
	// User must specify these
	Latitude, Longitude float64
	Properties          *[]string
}

// GeoSlice sorted by GeoHash
type GeoSlice []*GeoData

func (s GeoSlice) Len() int           { return len(s) }
func (s GeoSlice) Less(i, j int) bool { return s[i].GeoHash < s[j].GeoHash }
func (s GeoSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

// GeoIdSlice sorted by ID
type GeoIdSlice []*GeoData

func (s GeoIdSlice) Len() int           { return len(s) }
func (s GeoIdSlice) Less(i, j int) bool { return s[i].Id < s[j].Id }
func (s GeoIdSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

var searchReady = false
// var geoStore = make(map[string]GeoSlice)
var geoHashStore GeoSlice
var geoIdStore GeoIdSlice

const MAX_STEPS C.uint8_t = 26

// Debug logger. Remember to init flag.Parse() in main!!!
var debug = flag.Bool("debugLib", false, "enable debug logging")

func Debugf(format string, args ...interface{}) {
	if *debug {
		log.Printf("DEBUG "+format, args...)
	}
}

// AddLocation data to search index.
func AddLocation(geoData *GeoData) {
	hash := geohashEncodeMax(geoData.Latitude, geoData.Longitude)
	geoData.GeoHash = hash
	geoHashStore = append(geoHashStore, geoData)
	geoIdStore = append(geoIdStore, geoData)
	searchReady = false
}

// GetLocation data for a location ID.
func GetLocation(id int) (geodata *GeoData, err error) {
	if id >= len(geoIdStore) {
		return nil, errors.New("index out of range")
	}
	return geoIdStore[id], nil
}

// SearchLocations around latitude/longitude in bounded area (km) for known location points.
func SearchLocations(latitude, longitude, bound float64) []*GeoData {
	if !searchReady {
		initSearch()
	}
	hashSteps := C.geohashEstimateStepsByRadius(C.double(bound * 1000))
	Debugf("Hash step: %v, for radius: %v km", hashSteps, bound)

	var hash C.GeoHashBits
	C.geohashEncodeWGS84(C.double(latitude), C.double(longitude), C.uint8_t(hashSteps), &hash)
	neighbours := getNeighbours(uint64(hash.bits), uint8(hashSteps))
	box := boundingBox(latitude, longitude, bound)

	locationsFound := make([]*GeoData, 0)
	geoStoreKeysLen := len(geoHashStore)
	for nIdx := range neighbours {
		neighboursUpperLimit := (neighbours[nIdx] + 1) << uint((MAX_STEPS-hashSteps)*2)
		neighbours[nIdx] = neighbours[nIdx] << uint((MAX_STEPS-hashSteps)*2)
		Debugf("Normalized Neighbours Hash: %v to %v", neighbours[nIdx], neighboursUpperLimit)
		searchIdx := sort.Search(geoStoreKeysLen, func(i int) bool { return geoHashStore[i].GeoHash >= neighbours[nIdx] })
		if searchIdx < geoStoreKeysLen { // Not found would turn index=N
			Debugf("found location?")
			// found location
			for i := searchIdx; i < geoStoreKeysLen; i++ {
				if geoHashStore[i].GeoHash < neighboursUpperLimit {
					data := geoHashStore[i]
					Debugf("filtering by lat/long: %v %v", data.Latitude, data.Longitude)
					Debugf("filtering by bounding box: %v %v %v %v", box[0], box[1], box[2], box[3])
					// filter by strict bounding box
					if ((data.Latitude >= box[0] && data.Latitude <= box[1]) || (data.Latitude <= box[0] && data.Latitude >= box[1])) &&
					((data.Longitude >= box[2] && data.Longitude <= box[3]) || (data.Longitude <= box[2] && data.Longitude >= box[3])) {
						locationsFound = append(locationsFound, data)
						Debugf("Search found location in geoHashStore: %v", geoHashStore[searchIdx])
					}
				} else {
					break
				}
			}
		}
	}
	return locationsFound
}

// Sort the list of geo so binary search can be used. Normally triggered by the first search.
func initSearch() {
	sort.Sort(geoHashStore)
	sort.Sort(geoIdStore)
	searchReady = true // might cause race cond, but assume add doesn't happen often.
}

func getNeighbours(hashBits uint64, steps uint8) []uint64 {
	var neighbours C.GeoHashNeighbors
	var hash C.GeoHashBits
	hash.bits = C.uint64_t(hashBits)
	hash.step = C.uint8_t(steps)
	C.geohashNeighbors(&hash, &neighbours)

	neighbourArr := sortutil.Uint64Slice{
		uint64(hashBits),
		uint64(neighbours.north.bits),
		uint64(neighbours.east.bits),
		uint64(neighbours.west.bits),
		uint64(neighbours.south.bits),
		uint64(neighbours.north_east.bits),
		uint64(neighbours.south_east.bits),
		uint64(neighbours.north_west.bits),
		uint64(neighbours.south_west.bits),
	}
	//	sort.Sort(neighbourArr)
	if steps <= 6 {
		sortutil.Dedupe(neighbourArr) // Can have duplicates if search range is large (>~5000km)
	}
	return neighbourArr
}

// Encode a geo hash to MAX(26) steps
func geohashEncodeMax(latitude, longitude float64) uint64 {
	var hash C.GeoHashBits
	C.geohashEncodeWGS84(C.double(latitude), C.double(longitude), MAX_STEPS, &hash)
	return uint64(hash.bits)
}

// The approximate conversions are (doesn't fully correct for the Earth's polar flattening):
// Latitude: 1 deg = 110.574 km
// Longitude: 1 deg = 111.320*cos(latitude) km
// See: http://stackoverflow.com/questions/1253499/simple-calculations-for-working-with-lat-lon-km-distance
// Returns: min/max latitude, min/max longitude.
func boundingBox(latitude, longitude, boundKm float64) []float64 {
	latDiff := boundKm / 110.574
	longDiff := boundKm / (111.320 * math.Cos(latitude))
	minLatitude := latitude - latDiff
	maxLatitude := latitude + latDiff
	minLongitude := longitude - longDiff
	maxLongitude := longitude + longDiff

	Debugf("min lat: %v", minLatitude)
	Debugf("max lat: %v", maxLatitude)
	Debugf("min long: %v", minLongitude)
	Debugf("max long: %v", maxLongitude)

	return []float64{
		minLatitude,
		maxLatitude,
		minLongitude,
		maxLongitude,
	}
}

<commit_msg>Clean up<commit_after>/*
 * Copyright (c) 2015, Ian Chan <icha024@gmail.com>.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice,
 *    this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *  * Neither the name of Redis nor the names of its contributors may be used
 *    to endorse or promote products derived from this software without
 *    specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
 * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
 * THE POSSIBILITY OF SUCH DAMAGE.
 */
package geoindex

/*
#include "geohash.h"
#include "geohash_helper.h"
#include <math.h>
#cgo LDFLAGS: -lm
*/
import "C"
import (
	"errors"
	"flag"
	"github.com/cznic/sortutil"
	"log"
	"math"
	"sort"
)

// GeoData location representation
type GeoData struct {
	// Must be unique
	ID int
	// Generated automatically
	GeoHash uint64
	// User must specify these
	Latitude, Longitude float64
	Properties          *[]string
}

// GeoHashSlice sorted by GeoHash
type geoHashSlice []*GeoData

func (s geoHashSlice) Len() int           { return len(s) }
func (s geoHashSlice) Less(i, j int) bool { return s[i].GeoHash < s[j].GeoHash }
func (s geoHashSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

// GeoIDSlice sorted by ID
type geoIDSlice []*GeoData

func (s geoIDSlice) Len() int           { return len(s) }
func (s geoIDSlice) Less(i, j int) bool { return s[i].ID < s[j].ID }
func (s geoIDSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

var searchReady = false
var geoHashStore geoHashSlice
var geoIDStore geoIDSlice

const maxSteps C.uint8_t = 26

// Debug logger. Remember to init flag.Parse() in main!!!
var debug = flag.Bool("debugLib", false, "enable debug logging")

func debugf(format string, args ...interface{}) {
	if *debug {
		log.Printf("DEBUG "+format, args...)
	}
}

// AddLocation data to search index.
func AddLocation(geoData *GeoData) {
	hash := geohashEncodeMax(geoData.Latitude, geoData.Longitude)
	geoData.GeoHash = hash
	geoHashStore = append(geoHashStore, geoData)
	geoIDStore = append(geoIDStore, geoData)
	searchReady = false
}

// GetLocation data for a location ID.
func GetLocation(id int) (geodata *GeoData, err error) {
	if id >= len(geoIDStore) {
		return nil, errors.New("index out of range")
	}
	return geoIDStore[id], nil
}

// SearchLocations around latitude/longitude in bounded area (km) for known location points.
func SearchLocations(latitude, longitude, bound float64) []*GeoData {
	if !searchReady {
		initSearch()
	}
	hashSteps := C.geohashEstimateStepsByRadius(C.double(bound * 1000))
	debugf("Hash step: %v, for radius: %v km", hashSteps, bound)

	var hash C.GeoHashBits
	C.geohashEncodeWGS84(C.double(latitude), C.double(longitude), C.uint8_t(hashSteps), &hash)
	neighbours := getNeighbours(uint64(hash.bits), uint8(hashSteps))
	box := boundingBox(latitude, longitude, bound)

	// locationsFound := make([]*GeoData, 0)
	var locationsFound []*GeoData
	geoStoreKeysLen := len(geoHashStore)
	for nIdx := range neighbours {
		neighboursUpperLimit := (neighbours[nIdx] + 1) << uint((maxSteps-hashSteps)*2)
		neighbours[nIdx] = neighbours[nIdx] << uint((maxSteps-hashSteps)*2)
		debugf("Normalized Neighbours Hash: %v to %v", neighbours[nIdx], neighboursUpperLimit)
		searchIdx := sort.Search(geoStoreKeysLen, func(i int) bool { return geoHashStore[i].GeoHash >= neighbours[nIdx] })
		if searchIdx < geoStoreKeysLen { // Not found would turn index=N
			debugf("found location?")
			// found location
			for i := searchIdx; i < geoStoreKeysLen; i++ {
				if geoHashStore[i].GeoHash < neighboursUpperLimit {
					data := geoHashStore[i]
					debugf("filtering by lat/long: %v %v", data.Latitude, data.Longitude)
					debugf("filtering by bounding box: %v %v %v %v", box[0], box[1], box[2], box[3])
					// filter by strict bounding box
					if ((data.Latitude >= box[0] && data.Latitude <= box[1]) || (data.Latitude <= box[0] && data.Latitude >= box[1])) &&
						((data.Longitude >= box[2] && data.Longitude <= box[3]) || (data.Longitude <= box[2] && data.Longitude >= box[3])) {
						locationsFound = append(locationsFound, data)
						debugf("Search found location in geoHashStore: %v", geoHashStore[searchIdx])
					}
				} else {
					break
				}
			}
		}
	}
	return locationsFound
}

// Sort the list of geo so binary search can be used. Normally triggered by the first search.
func initSearch() {
	sort.Sort(geoHashStore)
	sort.Sort(geoIDStore)
	searchReady = true // might cause race cond, but assume add doesn't happen often.
}

func getNeighbours(hashBits uint64, steps uint8) []uint64 {
	var neighbours C.GeoHashNeighbors
	var hash C.GeoHashBits
	hash.bits = C.uint64_t(hashBits)
	hash.step = C.uint8_t(steps)
	C.geohashNeighbors(&hash, &neighbours)

	neighbourArr := sortutil.Uint64Slice{
		uint64(hashBits),
		uint64(neighbours.north.bits),
		uint64(neighbours.east.bits),
		uint64(neighbours.west.bits),
		uint64(neighbours.south.bits),
		uint64(neighbours.north_east.bits),
		uint64(neighbours.south_east.bits),
		uint64(neighbours.north_west.bits),
		uint64(neighbours.south_west.bits),
	}
	//	sort.Sort(neighbourArr)
	if steps <= 6 {
		sortutil.Dedupe(neighbourArr) // Can have duplicates if search range is large (>~5000km)
	}
	return neighbourArr
}

// Encode a geo hash to MAX(26) steps
func geohashEncodeMax(latitude, longitude float64) uint64 {
	var hash C.GeoHashBits
	C.geohashEncodeWGS84(C.double(latitude), C.double(longitude), maxSteps, &hash)
	return uint64(hash.bits)
}

// The approximate conversions are (doesn't fully correct for the Earth's polar flattening):
// Latitude: 1 deg = 110.574 km
// Longitude: 1 deg = 111.320*cos(latitude) km
// See: http://stackoverflow.com/questions/1253499/simple-calculations-for-working-with-lat-lon-km-distance
// Returns: min/max latitude, min/max longitude.
func boundingBox(latitude, longitude, boundKm float64) []float64 {
	latDiff := boundKm / 110.574
	longDiff := boundKm / (111.320 * math.Cos(latitude))
	minLatitude := latitude - latDiff
	maxLatitude := latitude + latDiff
	minLongitude := longitude - longDiff
	maxLongitude := longitude + longDiff

	debugf("min lat: %v", minLatitude)
	debugf("max lat: %v", maxLatitude)
	debugf("min long: %v", minLongitude)
	debugf("max long: %v", maxLongitude)

	return []float64{
		minLatitude,
		maxLatitude,
		minLongitude,
		maxLongitude,
	}
}
<|endoftext|>


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

// +build !plan9

// Package fsnotify provides a platform-independent interface for file system notifications.
package fsnotify

import (
	"bytes"
	"errors"
	"fmt"
)

// Event represents a single file system notification.
type Event struct {
	Name string // Relative path to the file or directory.
	Op   Op     // File operation that triggered the event.
}

// Op describes a set of file operations.
type Op uint32

// These are the generalized file operations that can trigger a notification.
const (
	Create Op = 1 << iota
	Write
	Remove
	Rename
	Chmod
)

func (op Op) String() string {
	// Use a buffer for efficient string concatenation
	var buffer bytes.Buffer

	if op&Create == Create {
		buffer.WriteString("|CREATE")
	}
	if op&Remove == Remove {
		buffer.WriteString("|REMOVE")
	}
	if op&Write == Write {
		buffer.WriteString("|WRITE")
	}
	if op&Rename == Rename {
		buffer.WriteString("|RENAME")
	}
	if op&Chmod == Chmod {
		buffer.WriteString("|CHMOD")
	}
	if buffer.Len() == 0 {
		return ""
	}
	return buffer.String()[1:] // Strip leading pipe
}

// String returns a string representation of the event in the form
// "file: REMOVE|WRITE|..."
func (e Event) String() string {
	return fmt.Sprintf("%q: %s", e.Name, e.Op.String())
}

// Common errors that can be reported by a watcher
var ErrEventOverflow = errors.New("fsnotify queue overflow")
<commit_msg>Fix lint issue with error doc string<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build !plan9

// Package fsnotify provides a platform-independent interface for file system notifications.
package fsnotify

import (
	"bytes"
	"errors"
	"fmt"
)

// Event represents a single file system notification.
type Event struct {
	Name string // Relative path to the file or directory.
	Op   Op     // File operation that triggered the event.
}

// Op describes a set of file operations.
type Op uint32

// These are the generalized file operations that can trigger a notification.
const (
	Create Op = 1 << iota
	Write
	Remove
	Rename
	Chmod
)

func (op Op) String() string {
	// Use a buffer for efficient string concatenation
	var buffer bytes.Buffer

	if op&Create == Create {
		buffer.WriteString("|CREATE")
	}
	if op&Remove == Remove {
		buffer.WriteString("|REMOVE")
	}
	if op&Write == Write {
		buffer.WriteString("|WRITE")
	}
	if op&Rename == Rename {
		buffer.WriteString("|RENAME")
	}
	if op&Chmod == Chmod {
		buffer.WriteString("|CHMOD")
	}
	if buffer.Len() == 0 {
		return ""
	}
	return buffer.String()[1:] // Strip leading pipe
}

// String returns a string representation of the event in the form
// "file: REMOVE|WRITE|..."
func (e Event) String() string {
	return fmt.Sprintf("%q: %s", e.Name, e.Op.String())
}

// Common errors that can be reported by a watcher
var (
	ErrEventOverflow = errors.New("fsnotify queue overflow")
)
<|endoftext|>


<commit_before>package maps

import (
	"fmt"
	"image"
	"image/draw"
	"math"

	"github.com/ryankurte/go-mapbox/lib/base"
	"image/color"
)

// Tile is a wrapper around an image that includes positioning data
type Tile struct {
	draw.Image
	Level uint64 // Tile zoom level
	Size  uint64 // Tile size
	X, Y  uint64 // Tile X and Y postions (Web Mercurator projection)
}

const (
	SizeStandard uint64 = 256
	SizeHighDPI  uint64 = 512
)

// NewTile creates a tile with a base RGBA object
func NewTile(x, y, level, size uint64, src image.Image) Tile {
	// Convert image to RGBA
	b := src.Bounds()
	m := image.NewNRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
	draw.Draw(m, m.Bounds(), src, b.Min, draw.Src)

	return Tile{
		Image: m,
		X:     x,
		Y:     y,
		Level: level,
		Size:  size,
	}
}

// LocationToPixel translates a global location to a pixel on the tile
func (t *Tile) LocationToPixel(loc base.Location) (float64, float64, error) {
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	offsetX, offsetY := x-float64(t.X*t.Size), y-float64(t.Y*t.Size)

	if xMax := float64(t.Image.Bounds().Max.X); (offsetX < 0) || (offsetX > xMax) {
		return 0, 0, fmt.Errorf("Tile LocationToPixel error: global X offset not within tile space (x: %d max: %d)", offsetX, int(xMax))
	}
	if yMax := float64(t.Image.Bounds().Max.Y); (offsetY < 0) || (offsetY > yMax) {
		return 0, 0, fmt.Errorf("Tile LocationToPixel error: global Y offset not within tile space (y: %d max: %d)", offsetY, int(yMax))
	}

	return offsetX, offsetY, nil
}

// PixelToLocation translates a pixel location in the tile into a global location
func (t *Tile) PixelToLocation(x, y float64) (*base.Location, error) {
	if xMax := float64(t.Image.Bounds().Max.X); (x < 0) || (x > xMax) {
		return nil, fmt.Errorf("Tile LocationToPixel error: global X offset not within tile space (x: %.2f max: %d)", x, int(xMax))
	}
	if yMax := float64(t.Image.Bounds().Max.Y); (y < 0) || (y > yMax) {
		return nil, fmt.Errorf("Tile LocationToPixel error: global Y offset not within tile space (y: %.2f max: %d)", y, int(yMax))
	}

	offsetX, offsetY := x+float64(t.X*t.Size), y+float64(t.Y*t.Size)
	lat, lng := MercatorPixelToLocation(offsetX, offsetY, t.Level, t.Size)

	return &base.Location{Latitude: lat, Longitude: lng}, nil
}

// Justify sets image offsets for drawing
type Justify string

// Constant justification types
const (
	JustifyTop    Justify = "top"
	JustifyLeft   Justify = "left"
	JustifyCenter Justify = "center"
	JustifyBottom Justify = "bottom"
	JustifyRight  Justify = "right"
)

// DrawConfig configures image drawing
type DrawConfig struct {
	Vertical   Justify
	Horizontal Justify
}

// Center preconfigured centering helper
var Center = DrawConfig{JustifyCenter, JustifyCenter}

// DrawLocalXY draws the provided image at the local X/Y coordinates
func (t *Tile) DrawLocalXY(src image.Image, x, y int, config DrawConfig) error {
	dp := image.Point{}

	switch config.Horizontal {
	case JustifyLeft:
		dp.X = x
	case JustifyCenter:
		dp.X = x - src.Bounds().Dx()/2
	case JustifyRight:
		dp.X = x - src.Bounds().Dx()
	default:
		return fmt.Errorf("Unsupported horizontal justification (%s)", config.Horizontal)
	}

	switch config.Vertical {
	case JustifyTop:
		dp.Y = y
	case JustifyCenter:
		dp.Y = y - src.Bounds().Dy()/2
	case JustifyBottom:
		dp.Y = y - src.Bounds().Dy()
	default:
		return fmt.Errorf("Unsupported vertical justification (%s)", config.Horizontal)
	}

	r := image.Rectangle{dp, dp.Add(src.Bounds().Size())}
	draw.Draw(t.Image.(draw.Image), r, src, src.Bounds().Min, draw.Over)

	return nil
}

func (t *Tile) translateGlobalToLocalXY(x, y int) (int, int, error) {
	offsetX := x - int(t.X*t.Size)
	offsetY := y - int(t.Y*t.Size)

	if xMax := int(t.Image.Bounds().Max.X); (offsetX < 0) || (offsetX > xMax) {
		return 0, 0, fmt.Errorf("Tile DrawGlobalXY error: global X offset not within tile space (x: %d max: %d)", offsetX, xMax)
	}
	if yMax := int(t.Image.Bounds().Max.Y); (offsetY < 0) || (offsetY > yMax) {
		return 0, 0, fmt.Errorf("Tile DrawGlobalXY error: global Y offset not within tile space (y: %d max: %d)", offsetY, yMax)
	}

	return offsetX, offsetY, nil
}

// DrawGlobalXY draws the provided image at the global X/Y coordinates
func (t *Tile) DrawGlobalXY(src image.Image, x, y int, config DrawConfig) error {
	offsetX, offsetY, err := t.translateGlobalToLocalXY(x, y)
	if err != nil {
		return err
	}
	t.DrawLocalXY(src, offsetX, offsetY, config)
	return nil
}

// DrawLocation draws the provided image at the provided lat lng
func (t *Tile) DrawLocation(src image.Image, loc base.Location, config DrawConfig) {
	// Calculate location in pixel space
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	t.DrawGlobalXY(src, int(x), int(y), config)
}

// Interpolate function to be passed to generic line interpolator
type Interpolate func(pixel color.Color) color.Color

// InterpolateLocalXY interpolates a line between two local points and calls the interpolate function on each point
func (t *Tile) InterpolateLocalXY(x1, y1, x2, y2 int, interpolate Interpolate) {
	// This is a bit insane because the ordering between (x1, y1) and (x2, y2) must be preserved
	// So that points

	dx := int(float64(x2) - float64(x1))
	dy := int(float64(y2) - float64(y1))

	len := int(math.Sqrt(math.Pow(float64(dx), 2) + math.Pow(float64(dy), 2)))

	img := t.Image.(draw.Image)

	for i := 0; i < len; i++ {
		x := x1 + i*dx/len
		y := y1 + i*dy/len
		pixel := interpolate(img.At(x, y))
		img.Set(x, y, pixel)
	}
}

// InterpolateGlobalXY interpolates a line between two global points and calls the interpolate function on each point
func (t *Tile) InterpolateGlobalXY(x1, y1, x2, y2 int, interpolate Interpolate) error {
	offsetX1, offsetY1, err := t.translateGlobalToLocalXY(x1, y1)
	if err != nil {
		return err
	}
	offsetX2, offsetY2, err := t.translateGlobalToLocalXY(x2, y2)
	if err != nil {
		return err
	}
	t.InterpolateLocalXY(offsetX1, offsetY1, offsetX2, offsetY2, interpolate)
	return nil
}

// InterpolateLocations interpolates a line between two locations and calls the interpolate function on each point
func (t *Tile) InterpolateLocations(loc1, loc2 base.Location, interpolate Interpolate) error {
	x1, y1 := MercatorLocationToPixel(loc1.Latitude, loc1.Longitude, t.Level, t.Size)
	x2, y2 := MercatorLocationToPixel(loc2.Latitude, loc2.Longitude, t.Level, t.Size)
	return t.InterpolateGlobalXY(int(x1), int(y1), int(x2), int(y2), interpolate)
}

// DrawLine uses InterpolateLocations to draw a line between two points
func (t *Tile) DrawLine(loc1, loc2 base.Location, c color.Color) {
	t.InterpolateLocations(loc1, loc2, func(color.Color) color.Color {
		return c
	})
}

func (t *Tile) DrawPoint(loc base.Location, size uint64, c color.Color) error {
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	offsetX, offsetY, err := t.translateGlobalToLocalXY(int(x), int(y))
	if err != nil {
		return err
	}
	img := t.Image.(draw.Image)

	for x := uint64(0); x < size; x++ {
		for y := uint64(0); y < size; y++ {
			img.Set(offsetX+int(x-size/2), offsetY+int(y-size/2), c)
		}
	}

	return nil
}

func (t *Tile) GetAltitude(loc base.Location) (float64, error) {
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	offsetX, offsetY, err := t.translateGlobalToLocalXY(int(x), int(y))
	if err != nil {
		return 0.0, err
	}
	p := t.Image.At(offsetX, offsetY).(color.NRGBA)
	return PixelToHeight(p.R, p.G, p.B), nil
}

func (t *Tile) InterpolateAltitudes(loc1, loc2 base.Location) []float64 {
	altitudes := make([]float64, 0)
	t.InterpolateLocations(loc1, loc2, func(c color.Color) color.Color {
		p := c.(color.NRGBA)
		alt := PixelToHeight(p.R, p.G, p.B)
		altitudes = append(altitudes, alt)
		return c
	})
	return altitudes
}

func (t *Tile) GetHighestAltitude() float64 {
	p := t.Image.At(0, 0).(color.NRGBA)
	max := PixelToHeight(p.R, p.G, p.B)
	for y := 0; y < t.Image.Bounds().Dy(); y++ {
		for x := 0; x < t.Image.Bounds().Dx(); x++ {
			p := t.Image.At(x, y).(color.NRGBA)
			alt := PixelToHeight(p.R, p.G, p.B)
			if alt > max {
				max = alt
			}
		}
	}
	return max
}

func (t *Tile) FlattenAltitudes(maxHeight float64) Tile {
	img := image.NewNRGBA(t.Image.Bounds())

	for y := 0; y < t.Image.Bounds().Dy(); y++ {
		for x := 0; x < t.Image.Bounds().Dx(); x++ {
			p := t.Image.At(x, y).(color.NRGBA)
			alt := uint8(PixelToHeight(p.R, p.G, p.B) / maxHeight * 255)
			img.Set(x, y, color.RGBA{R: alt, G: alt, B: alt, A: 255})
		}
	}

	return NewTile(t.X, t.Y, t.Level, t.Size, img)
}
<commit_msg>Added tile clone method<commit_after>package maps

import (
	"fmt"
	"image"
	"image/draw"
	"math"

	"github.com/ryankurte/go-mapbox/lib/base"
	"image/color"
)

// Tile is a wrapper around an image that includes positioning data
type Tile struct {
	draw.Image
	Level uint64 // Tile zoom level
	Size  uint64 // Tile size
	X, Y  uint64 // Tile X and Y postions (Web Mercurator projection)
}

const (
	SizeStandard uint64 = 256
	SizeHighDPI  uint64 = 512
)

// NewTile creates a tile with a base RGBA object
func NewTile(x, y, level, size uint64, src image.Image) Tile {
	// Convert image to RGBA
	b := src.Bounds()
	m := image.NewNRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
	draw.Draw(m, m.Bounds(), src, b.Min, draw.Src)

	return Tile{
		Image: m,
		X:     x,
		Y:     y,
		Level: level,
		Size:  size,
	}
}

// Clone helper clones a tile instance with a new base image
func (t *Tile) Clone() Tile {
	n := *t
	b := t.Bounds()
	n.Image = image.NewRGBA(b)
	draw.Draw(n.Image, b, t, b.Min, draw.Src)
	return n
}

// LocationToPixel translates a global location to a pixel on the tile
func (t *Tile) LocationToPixel(loc base.Location) (float64, float64, error) {
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	offsetX, offsetY := x-float64(t.X*t.Size), y-float64(t.Y*t.Size)

	if xMax := float64(t.Image.Bounds().Max.X); (offsetX < 0) || (offsetX > xMax) {
		return 0, 0, fmt.Errorf("Tile LocationToPixel error: global X offset not within tile space (x: %d max: %d)", offsetX, int(xMax))
	}
	if yMax := float64(t.Image.Bounds().Max.Y); (offsetY < 0) || (offsetY > yMax) {
		return 0, 0, fmt.Errorf("Tile LocationToPixel error: global Y offset not within tile space (y: %d max: %d)", offsetY, int(yMax))
	}

	return offsetX, offsetY, nil
}

// PixelToLocation translates a pixel location in the tile into a global location
func (t *Tile) PixelToLocation(x, y float64) (*base.Location, error) {
	if xMax := float64(t.Image.Bounds().Max.X); (x < 0) || (x > xMax) {
		return nil, fmt.Errorf("Tile LocationToPixel error: global X offset not within tile space (x: %.2f max: %d)", x, int(xMax))
	}
	if yMax := float64(t.Image.Bounds().Max.Y); (y < 0) || (y > yMax) {
		return nil, fmt.Errorf("Tile LocationToPixel error: global Y offset not within tile space (y: %.2f max: %d)", y, int(yMax))
	}

	offsetX, offsetY := x+float64(t.X*t.Size), y+float64(t.Y*t.Size)
	lat, lng := MercatorPixelToLocation(offsetX, offsetY, t.Level, t.Size)

	return &base.Location{Latitude: lat, Longitude: lng}, nil
}

// Justify sets image offsets for drawing
type Justify string

// Constant justification types
const (
	JustifyTop    Justify = "top"
	JustifyLeft   Justify = "left"
	JustifyCenter Justify = "center"
	JustifyBottom Justify = "bottom"
	JustifyRight  Justify = "right"
)

// DrawConfig configures image drawing
type DrawConfig struct {
	Vertical   Justify
	Horizontal Justify
}

// Center preconfigured centering helper
var Center = DrawConfig{JustifyCenter, JustifyCenter}

// DrawLocalXY draws the provided image at the local X/Y coordinates
func (t *Tile) DrawLocalXY(src image.Image, x, y int, config DrawConfig) error {
	dp := image.Point{}

	switch config.Horizontal {
	case JustifyLeft:
		dp.X = x
	case JustifyCenter:
		dp.X = x - src.Bounds().Dx()/2
	case JustifyRight:
		dp.X = x - src.Bounds().Dx()
	default:
		return fmt.Errorf("Unsupported horizontal justification (%s)", config.Horizontal)
	}

	switch config.Vertical {
	case JustifyTop:
		dp.Y = y
	case JustifyCenter:
		dp.Y = y - src.Bounds().Dy()/2
	case JustifyBottom:
		dp.Y = y - src.Bounds().Dy()
	default:
		return fmt.Errorf("Unsupported vertical justification (%s)", config.Horizontal)
	}

	r := image.Rectangle{dp, dp.Add(src.Bounds().Size())}
	draw.Draw(t.Image.(draw.Image), r, src, src.Bounds().Min, draw.Over)

	return nil
}

func (t *Tile) translateGlobalToLocalXY(x, y int) (int, int, error) {
	offsetX := x - int(t.X*t.Size)
	offsetY := y - int(t.Y*t.Size)

	if xMax := int(t.Image.Bounds().Max.X); (offsetX < 0) || (offsetX > xMax) {
		return 0, 0, fmt.Errorf("Tile DrawGlobalXY error: global X offset not within tile space (x: %d max: %d)", offsetX, xMax)
	}
	if yMax := int(t.Image.Bounds().Max.Y); (offsetY < 0) || (offsetY > yMax) {
		return 0, 0, fmt.Errorf("Tile DrawGlobalXY error: global Y offset not within tile space (y: %d max: %d)", offsetY, yMax)
	}

	return offsetX, offsetY, nil
}

// DrawGlobalXY draws the provided image at the global X/Y coordinates
func (t *Tile) DrawGlobalXY(src image.Image, x, y int, config DrawConfig) error {
	offsetX, offsetY, err := t.translateGlobalToLocalXY(x, y)
	if err != nil {
		return err
	}
	t.DrawLocalXY(src, offsetX, offsetY, config)
	return nil
}

// DrawLocation draws the provided image at the provided lat lng
func (t *Tile) DrawLocation(src image.Image, loc base.Location, config DrawConfig) {
	// Calculate location in pixel space
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	t.DrawGlobalXY(src, int(x), int(y), config)
}

// Interpolate function to be passed to generic line interpolator
type Interpolate func(pixel color.Color) color.Color

// InterpolateLocalXY interpolates a line between two local points and calls the interpolate function on each point
func (t *Tile) InterpolateLocalXY(x1, y1, x2, y2 int, interpolate Interpolate) {
	// This is a bit insane because the ordering between (x1, y1) and (x2, y2) must be preserved
	// So that points

	dx := int(float64(x2) - float64(x1))
	dy := int(float64(y2) - float64(y1))

	len := int(math.Sqrt(math.Pow(float64(dx), 2) + math.Pow(float64(dy), 2)))

	img := t.Image.(draw.Image)

	for i := 0; i < len; i++ {
		x := x1 + i*dx/len
		y := y1 + i*dy/len
		pixel := interpolate(img.At(x, y))
		img.Set(x, y, pixel)
	}
}

// InterpolateGlobalXY interpolates a line between two global points and calls the interpolate function on each point
func (t *Tile) InterpolateGlobalXY(x1, y1, x2, y2 int, interpolate Interpolate) error {
	offsetX1, offsetY1, err := t.translateGlobalToLocalXY(x1, y1)
	if err != nil {
		return err
	}
	offsetX2, offsetY2, err := t.translateGlobalToLocalXY(x2, y2)
	if err != nil {
		return err
	}
	t.InterpolateLocalXY(offsetX1, offsetY1, offsetX2, offsetY2, interpolate)
	return nil
}

// InterpolateLocations interpolates a line between two locations and calls the interpolate function on each point
func (t *Tile) InterpolateLocations(loc1, loc2 base.Location, interpolate Interpolate) error {
	x1, y1 := MercatorLocationToPixel(loc1.Latitude, loc1.Longitude, t.Level, t.Size)
	x2, y2 := MercatorLocationToPixel(loc2.Latitude, loc2.Longitude, t.Level, t.Size)
	return t.InterpolateGlobalXY(int(x1), int(y1), int(x2), int(y2), interpolate)
}

// DrawLine uses InterpolateLocations to draw a line between two points
func (t *Tile) DrawLine(loc1, loc2 base.Location, c color.Color) {
	t.InterpolateLocations(loc1, loc2, func(color.Color) color.Color {
		return c
	})
}

func (t *Tile) DrawPoint(loc base.Location, size uint64, c color.Color) error {
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	offsetX, offsetY, err := t.translateGlobalToLocalXY(int(x), int(y))
	if err != nil {
		return err
	}
	img := t.Image.(draw.Image)

	for x := uint64(0); x < size; x++ {
		for y := uint64(0); y < size; y++ {
			img.Set(offsetX+int(x-size/2), offsetY+int(y-size/2), c)
		}
	}

	return nil
}

func (t *Tile) GetAltitude(loc base.Location) (float64, error) {
	x, y := MercatorLocationToPixel(loc.Latitude, loc.Longitude, t.Level, t.Size)
	offsetX, offsetY, err := t.translateGlobalToLocalXY(int(x), int(y))
	if err != nil {
		return 0.0, err
	}
	p := t.Image.At(offsetX, offsetY).(color.NRGBA)
	return PixelToHeight(p.R, p.G, p.B), nil
}

func (t *Tile) InterpolateAltitudes(loc1, loc2 base.Location) []float64 {
	altitudes := make([]float64, 0)
	t.InterpolateLocations(loc1, loc2, func(c color.Color) color.Color {
		p := c.(color.NRGBA)
		alt := PixelToHeight(p.R, p.G, p.B)
		altitudes = append(altitudes, alt)
		return c
	})
	return altitudes
}

func (t *Tile) GetHighestAltitude() float64 {
	p := t.Image.At(0, 0).(color.NRGBA)
	max := PixelToHeight(p.R, p.G, p.B)
	for y := 0; y < t.Image.Bounds().Dy(); y++ {
		for x := 0; x < t.Image.Bounds().Dx(); x++ {
			p := t.Image.At(x, y).(color.NRGBA)
			alt := PixelToHeight(p.R, p.G, p.B)
			if alt > max {
				max = alt
			}
		}
	}
	return max
}

func (t *Tile) FlattenAltitudes(maxHeight float64) Tile {
	img := image.NewNRGBA(t.Image.Bounds())

	for y := 0; y < t.Image.Bounds().Dy(); y++ {
		for x := 0; x < t.Image.Bounds().Dx(); x++ {
			p := t.Image.At(x, y).(color.NRGBA)
			alt := uint8(PixelToHeight(p.R, p.G, p.B) / maxHeight * 255)
			img.Set(x, y, color.RGBA{R: alt, G: alt, B: alt, A: 255})
		}
	}

	return NewTile(t.X, t.Y, t.Level, t.Size, img)
}
<|endoftext|>


<commit_before>// Package neptulon is a socket framework with middleware support.
package neptulon

import (
	"net/http"
	"sync"
)

// Neptulon framework entry point.
type Neptulon struct {
	debug       bool
	err         error
	listener    *Listener
	mutex       sync.Mutex
	middlewares []*func(ctx Context) (response interface{})
}

// Handle registers a new middleware to handle incoming messages.
func (n *Neptulon) Handle() {}

type handler func(w http.ResponseWriter, r *http.Request) error

func handle(handlers ...handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		for _, handler := range handlers {
			err := handler(w, r)
			if err != nil {
				w.Write([]byte(err.Error()))
				return
			}
		}
	})
}
<commit_msg>att top level neptulon req handlers<commit_after>// Package neptulon is a socket framework with middleware support.
package neptulon

import (
	"log"
	"net/http"
	"sync"
)

// Neptulon framework entry point.
type Neptulon struct {
	debug       bool
	err         error
	listener    *Listener
	mutex       sync.Mutex
	middlewares []*func(ctx Context) (response interface{})
}

// New creates and returns a new Neptulon app. This is the default TLS constructor.
// Debug mode dumps raw TCP data to stderr (log.Println() default).
func New(cert, privKey []byte, laddr string, debug bool) (*Neptulon, error) {
	l, err := Listen(cert, privKey, laddr, debug)
	if err != nil {
		return nil, err
	}

	return &Neptulon{
		debug:    debug,
		listener: l,
		mutex:    sync.Mutex{},
	}, nil
}

// Handle registers a new middleware to handle incoming messages.
func (n *Neptulon) Handle() {}

// Run starts accepting connections on the internal listener and handles connections with registered middleware.
// This function blocks and never returns, unless there was an error while accepting a new connection or the listner was closed.
func (n *Neptulon) Run() error {
	err := n.listener.Accept(handleMsg, handleDisconn)
	if err != nil && n.debug {
		log.Fatalln("Listener returned an error while closing:", err)
	}

	n.mutex.Lock()
	n.err = err
	n.mutex.Unlock()

	return err
}

// handleMsg handles incoming client messages.
func handleMsg(conn *Conn, session *Session, msg []byte) {

}

// handleDisconn handles client disconnection.
func handleDisconn(conn *Conn, session *Session) {}

// UseTLS enables TLS on a Neptulon app.
// func (n *Neptulon) UseTLS(cert, privKey []byte) {}

type handler func(w http.ResponseWriter, r *http.Request) error

func handle(handlers ...handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		for _, handler := range handlers {
			err := handler(w, r)
			if err != nil {
				w.Write([]byte(err.Error()))
				return
			}
		}
	})
}
<|endoftext|>


<commit_before>// Copyright © 2017 The Kubicorn 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 godoSdk

import (
	"os"
	"testing"
)

var (
	DoAccessToken = os.Getenv("DIGITALOCEAN_ACCESS_TOKEN")
)

func TestMain(m *testing.M) {
	m.Run()
	os.Setenv("DIGITALOCEAN_ACCESS_TOKEN", DoAccessToken)
}

func TestSdkHappy(t *testing.T) {
	os.Setenv("DIGITALOCEAN_ACCESS_TOKEN", "123")
	_, err := NewSdk()
	if err != nil {
		t.Fatalf("Unable to get Digital Ocean SDK: %v", err)
	}
}

func TestSdkSad(t *testing.T) {
	os.Setenv("DIGITALOCEAN_ACCESS_TOKEN", "")
	_, err := NewSdk()
	if err == nil {
		t.Fatalf("Able to get Amazon SDK with empty variables")
	}
}
<commit_msg>godoSdk: typo fix<commit_after>// Copyright © 2017 The Kubicorn 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 godoSdk

import (
	"os"
	"testing"
)

var (
	DoAccessToken = os.Getenv("DIGITALOCEAN_ACCESS_TOKEN")
)

func TestMain(m *testing.M) {
	m.Run()
	os.Setenv("DIGITALOCEAN_ACCESS_TOKEN", DoAccessToken)
}

func TestSdkHappy(t *testing.T) {
	os.Setenv("DIGITALOCEAN_ACCESS_TOKEN", "123")
	_, err := NewSdk()
	if err != nil {
		t.Fatalf("Unable to get Digital Ocean SDK: %v", err)
	}
}

func TestSdkSad(t *testing.T) {
	os.Setenv("DIGITALOCEAN_ACCESS_TOKEN", "")
	_, err := NewSdk()
	if err == nil {
		t.Fatalf("Able to get DigitalOcean SDK with empty variables")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/kisielk/whisper-go/whisper"
	"github.com/raintank/metrictank/api"
	"github.com/raintank/metrictank/conf"
	"github.com/raintank/metrictank/mdata/chunk"
	"github.com/raintank/metrictank/mdata/chunk/archive"
	"gopkg.in/raintank/schema.v1"
)

var (
	httpEndpoint = flag.String(
		"http-endpoint",
		"http://127.0.0.1:8080/chunks",
		"The http endpoint to send the data to",
	)
	namePrefix = flag.String(
		"name-prefix",
		"",
		"Prefix to prepend before every metric name, should include the '.' if necessary",
	)
	threads = flag.Int(
		"threads",
		10,
		"Number of workers threads to process and convert .wsp files",
	)
	writeUnfinishedChunks = flag.Bool(
		"write-unfinished-chunks",
		false,
		"Defines if chunks that have not completed their chunk span should be written",
	)
	orgId = flag.Int(
		"orgid",
		1,
		"Organization ID the data belongs to ",
	)
	insecureSSL = flag.Bool(
		"insecure-ssl",
		false,
		"Disables ssl certificate verification",
	)
	whisperDirectory = flag.String(
		"whisper-directory",
		"/opt/graphite/storage/whisper",
		"The directory that contains the whisper file structure",
	)
	httpAuth = flag.String(
		"http-auth",
		"",
		"The credentials used to authenticate in the format \"user:password\"",
	)
	dstSchemas = flag.String(
		"dst-schemas",
		"",
		"The filename of the output schemas definition file",
	)
	nameFilterPattern = flag.String(
		"name-filter",
		"",
		"A regex pattern to be applied to all metric names, only matching ones will be imported",
	)
	importUpTo = flag.Uint(
		"import-up-to",
		math.MaxUint32,
		"Only import up to the specified timestamp",
	)
	schemas    conf.Schemas
	nameFilter *regexp.Regexp
)

func main() {
	var err error
	flag.Parse()

	nameFilter = regexp.MustCompile(*nameFilterPattern)
	schemas, err = conf.ReadSchemas(*dstSchemas)
	if err != nil {
		panic(fmt.Sprintf("Error when parsing schemas file: %q", err))
	}

	fileChan := make(chan string)

	wg := &sync.WaitGroup{}
	wg.Add(*threads)
	for i := 0; i < *threads; i++ {
		go processFromChan(fileChan, wg)
	}

	getFileListIntoChan(fileChan)
	wg.Wait()
}

func processFromChan(files chan string, wg *sync.WaitGroup) {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: *insecureSSL},
	}
	client := &http.Client{Transport: tr}

	for file := range files {
		fd, err := os.Open(file)
		if err != nil {
			log.Error(fmt.Sprintf("Failed to open whisper file %q: %q\n", file, err))
			continue
		}
		w, err := whisper.OpenWhisper(fd)
		if err != nil {
			log.Error(fmt.Sprintf("Failed to open whisper file %q: %q\n", file, err))
			continue
		}

		name := getMetricName(file)
		if !nameFilter.Match([]byte(name)) {
			log.Info(fmt.Sprintf("Skipping file %s with name %s", file, name))
			continue
		}
		log.Info(fmt.Sprintf("Processing file %q", file))
		met, err := getMetrics(w, file, name)
		if err != nil {
			log.Error(fmt.Sprintf("Failed to get metric: %q", err))
			continue
		}

		b, err := met.MarshalCompressed()
		if err != nil {
			log.Error(fmt.Sprintf("Failed to encode metric: %q", err))
			continue
		}

		req, err := http.NewRequest("POST", *httpEndpoint, io.Reader(b))
		if err != nil {
			log.Fatal(fmt.Sprintf("Cannot construct request to http endpoint %q: %q", *httpEndpoint, err))
		}
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Content-Encoding", "gzip")

		if len(*httpAuth) > 0 {
			req.Header.Add("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(*httpAuth)))
		}

		resp, err := client.Do(req)
		if err != nil || resp.StatusCode != 200 {
			log.Error(fmt.Sprintf("Error sending request to http endpoint %q: %q", *httpEndpoint, err))
		}
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}
	wg.Done()
}

// generate the metric name based on the file name and given prefix
func getMetricName(file string) string {
	// remove all leading '/' from file name
	for file[0] == '/' {
		file = file[1:]
	}

	return *namePrefix + strings.Replace(strings.TrimSuffix(file, ".wsp"), "/", ".", -1)
}

// pointSorter sorts points by timestamp
type pointSorter []whisper.Point

func (a pointSorter) Len() int           { return len(a) }
func (a pointSorter) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a pointSorter) Less(i, j int) bool { return a[i].Timestamp < a[j].Timestamp }

// the whisper archives are organized like a ringbuffer. since we need to
// insert the points into the chunks in order we first need to sort them
func sortPoints(points pointSorter) pointSorter {
	sort.Sort(points)
	return points
}

func shortAggMethodString(aggMethod whisper.AggregationMethod) (string, error) {
	switch aggMethod {
	case whisper.AggregationAverage:
		return "avg", nil
	case whisper.AggregationSum:
		return "sum", nil
	case whisper.AggregationMin:
		return "min", nil
	case whisper.AggregationMax:
		return "max", nil
	case whisper.AggregationLast:
		return "lst", nil
	default:
		return "", fmt.Errorf("Unknown aggregation method %d", aggMethod)
	}
}

func getMetrics(w *whisper.Whisper, file, name string) (archive.Metric, error) {
	var res archive.Metric
	if len(w.Header.Archives) == 0 {
		return res, fmt.Errorf("Whisper file contains no archives: %q", file)
	}

	var archives []archive.Archive

	method, err := shortAggMethodString(w.Header.Metadata.AggregationMethod)
	if err != nil {
		return res, err
	}

	md := &schema.MetricData{
		Name:     name,
		Metric:   name,
		Interval: int(w.Header.Archives[0].SecondsPerPoint),
		Value:    0,
		Unit:     "unknown",
		Time:     0,
		Mtype:    "gauge",
		Tags:     []string{},
		OrgId:    *orgId,
	}
	md.SetId()
	_, schema := schemas.Match(md.Name, 0)

	points := make(map[int][]whisper.Point)
	for i := range w.Header.Archives {
		p, err := w.DumpArchive(i)
		if err != nil {
			return res, fmt.Errorf("Failed to dump archive %d from whisper file %s", i, file)
		}
		points[i] = p
	}

	conversion := newConversion(w.Header.Archives, points, method)
	for retIdx, retention := range schema.Retentions {
		convertedPoints := conversion.getPoints(retIdx, uint32(retention.SecondsPerPoint), uint32(retention.NumberOfPoints))
		for m, p := range convertedPoints {
			rowKey := getRowKey(retIdx, md.Id, m, retention.SecondsPerPoint)
			encodedChunks := encodedChunksFromPoints(p, uint32(retention.SecondsPerPoint), retention.ChunkSpan)
			archives = append(archives, archive.Archive{
				SecondsPerPoint: uint32(retention.SecondsPerPoint),
				Points:          uint32(retention.NumberOfPoints),
				Chunks:          encodedChunks,
				RowKey:          rowKey,
			})
			if int64(p[len(p)-1].Timestamp) > md.Time {
				md.Time = int64(p[len(p)-1].Timestamp)
			}
		}
	}

	res.AggregationMethod = uint32(w.Header.Metadata.AggregationMethod)
	res.MetricData = *md
	res.Archives = archives
	return res, nil
}

func getRowKey(retIdx int, id, meth string, secondsPerPoint int) string {
	if retIdx == 0 {
		return id
	} else {
		return api.AggMetricKey(
			id,
			meth,
			uint32(secondsPerPoint),
		)
	}
}

func encodedChunksFromPoints(points []whisper.Point, intervalIn, chunkSpan uint32) []chunk.IterGen {
	var point whisper.Point
	var t0, prevT0 uint32
	var c *chunk.Chunk
	var encodedChunks []chunk.IterGen

	for _, point = range points {
		// this shouldn't happen, but if it would we better catch it here because Metrictank wouldn't handle it well:
		// https://github.com/raintank/metrictank/blob/f1868cccfb92fc82cd853914af958f6d187c5f74/mdata/aggmetric.go#L378
		if point.Timestamp == 0 {
			continue
		}

		t0 = point.Timestamp - (point.Timestamp % chunkSpan)
		if prevT0 == 0 {
			c = chunk.New(t0)
			prevT0 = t0
		} else if prevT0 != t0 {
			c.Finish()

			encodedChunks = append(encodedChunks, *chunk.NewBareIterGen(c.Bytes(), c.T0, chunkSpan))

			c = chunk.New(t0)
			prevT0 = t0
		}

		err := c.Push(point.Timestamp, point.Value)
		if err != nil {
			panic(fmt.Sprintf("ERROR: Failed to push value into chunk at t0 %d: %q", t0, err))
		}
	}

	// if the last written point was also the last one of the current chunk,
	// or if writeUnfinishedChunks is on, we close the chunk and push it
	if point.Timestamp == t0+chunkSpan-intervalIn || *writeUnfinishedChunks {
		c.Finish()
		encodedChunks = append(encodedChunks, *chunk.NewBareIterGen(c.Bytes(), c.T0, chunkSpan))
	}

	return encodedChunks
}

// scan a directory and feed the list of whisper files relative to base into the given channel
func getFileListIntoChan(fileChan chan string) {
	filepath.Walk(
		*whisperDirectory,
		func(path string, info os.FileInfo, err error) error {
			if len(path) >= 4 && path[len(path)-4:] == ".wsp" {
				fileChan <- path
			}
			return nil
		},
	)

	close(fileChan)
}
<commit_msg>filter by name before pushing file into chan<commit_after>package main

import (
	"crypto/tls"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
	"sync"

	log "github.com/Sirupsen/logrus"
	"github.com/kisielk/whisper-go/whisper"
	"github.com/raintank/metrictank/api"
	"github.com/raintank/metrictank/conf"
	"github.com/raintank/metrictank/mdata/chunk"
	"github.com/raintank/metrictank/mdata/chunk/archive"
	"gopkg.in/raintank/schema.v1"
)

var (
	httpEndpoint = flag.String(
		"http-endpoint",
		"http://127.0.0.1:8080/chunks",
		"The http endpoint to send the data to",
	)
	namePrefix = flag.String(
		"name-prefix",
		"",
		"Prefix to prepend before every metric name, should include the '.' if necessary",
	)
	threads = flag.Int(
		"threads",
		10,
		"Number of workers threads to process and convert .wsp files",
	)
	writeUnfinishedChunks = flag.Bool(
		"write-unfinished-chunks",
		false,
		"Defines if chunks that have not completed their chunk span should be written",
	)
	orgId = flag.Int(
		"orgid",
		1,
		"Organization ID the data belongs to ",
	)
	insecureSSL = flag.Bool(
		"insecure-ssl",
		false,
		"Disables ssl certificate verification",
	)
	whisperDirectory = flag.String(
		"whisper-directory",
		"/opt/graphite/storage/whisper",
		"The directory that contains the whisper file structure",
	)
	httpAuth = flag.String(
		"http-auth",
		"",
		"The credentials used to authenticate in the format \"user:password\"",
	)
	dstSchemas = flag.String(
		"dst-schemas",
		"",
		"The filename of the output schemas definition file",
	)
	nameFilterPattern = flag.String(
		"name-filter",
		"",
		"A regex pattern to be applied to all metric names, only matching ones will be imported",
	)
	importUpTo = flag.Uint(
		"import-up-to",
		math.MaxUint32,
		"Only import up to the specified timestamp",
	)
	schemas    conf.Schemas
	nameFilter *regexp.Regexp
)

func main() {
	var err error
	flag.Parse()

	nameFilter = regexp.MustCompile(*nameFilterPattern)
	schemas, err = conf.ReadSchemas(*dstSchemas)
	if err != nil {
		panic(fmt.Sprintf("Error when parsing schemas file: %q", err))
	}

	fileChan := make(chan string)

	wg := &sync.WaitGroup{}
	wg.Add(*threads)
	for i := 0; i < *threads; i++ {
		go processFromChan(fileChan, wg)
	}

	getFileListIntoChan(fileChan)
	wg.Wait()
}

func processFromChan(files chan string, wg *sync.WaitGroup) {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: *insecureSSL},
	}
	client := &http.Client{Transport: tr}

	for file := range files {
		fd, err := os.Open(file)
		if err != nil {
			log.Error(fmt.Sprintf("Failed to open whisper file %q: %q\n", file, err))
			continue
		}
		w, err := whisper.OpenWhisper(fd)
		if err != nil {
			log.Error(fmt.Sprintf("Failed to open whisper file %q: %q\n", file, err))
			continue
		}

		name := getMetricName(file)
		log.Info(fmt.Sprintf("Processing file %s (%s)", file, name))
		met, err := getMetrics(w, file, name)
		if err != nil {
			log.Error(fmt.Sprintf("Failed to get metric: %q", err))
			continue
		}

		b, err := met.MarshalCompressed()
		if err != nil {
			log.Error(fmt.Sprintf("Failed to encode metric: %q", err))
			continue
		}

		req, err := http.NewRequest("POST", *httpEndpoint, io.Reader(b))
		if err != nil {
			log.Fatal(fmt.Sprintf("Cannot construct request to http endpoint %q: %q", *httpEndpoint, err))
		}
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Content-Encoding", "gzip")

		if len(*httpAuth) > 0 {
			req.Header.Add("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(*httpAuth)))
		}

		resp, err := client.Do(req)
		if err != nil || resp.StatusCode != 200 {
			log.Error(fmt.Sprintf("Error sending request to http endpoint %q: %q", *httpEndpoint, err))
		}
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}
	wg.Done()
}

// generate the metric name based on the file name and given prefix
func getMetricName(file string) string {
	// remove all leading '/' from file name
	for file[0] == '/' {
		file = file[1:]
	}

	return *namePrefix + strings.Replace(strings.TrimSuffix(file, ".wsp"), "/", ".", -1)
}

// pointSorter sorts points by timestamp
type pointSorter []whisper.Point

func (a pointSorter) Len() int           { return len(a) }
func (a pointSorter) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a pointSorter) Less(i, j int) bool { return a[i].Timestamp < a[j].Timestamp }

// the whisper archives are organized like a ringbuffer. since we need to
// insert the points into the chunks in order we first need to sort them
func sortPoints(points pointSorter) pointSorter {
	sort.Sort(points)
	return points
}

func shortAggMethodString(aggMethod whisper.AggregationMethod) (string, error) {
	switch aggMethod {
	case whisper.AggregationAverage:
		return "avg", nil
	case whisper.AggregationSum:
		return "sum", nil
	case whisper.AggregationMin:
		return "min", nil
	case whisper.AggregationMax:
		return "max", nil
	case whisper.AggregationLast:
		return "lst", nil
	default:
		return "", fmt.Errorf("Unknown aggregation method %d", aggMethod)
	}
}

func getMetrics(w *whisper.Whisper, file, name string) (archive.Metric, error) {
	var res archive.Metric
	if len(w.Header.Archives) == 0 {
		return res, fmt.Errorf("Whisper file contains no archives: %q", file)
	}

	var archives []archive.Archive

	method, err := shortAggMethodString(w.Header.Metadata.AggregationMethod)
	if err != nil {
		return res, err
	}

	md := &schema.MetricData{
		Name:     name,
		Metric:   name,
		Interval: int(w.Header.Archives[0].SecondsPerPoint),
		Value:    0,
		Unit:     "unknown",
		Time:     0,
		Mtype:    "gauge",
		Tags:     []string{},
		OrgId:    *orgId,
	}
	md.SetId()
	_, schema := schemas.Match(md.Name, 0)

	points := make(map[int][]whisper.Point)
	for i := range w.Header.Archives {
		p, err := w.DumpArchive(i)
		if err != nil {
			return res, fmt.Errorf("Failed to dump archive %d from whisper file %s", i, file)
		}
		points[i] = p
	}

	conversion := newConversion(w.Header.Archives, points, method)
	for retIdx, retention := range schema.Retentions {
		convertedPoints := conversion.getPoints(retIdx, uint32(retention.SecondsPerPoint), uint32(retention.NumberOfPoints))
		for m, p := range convertedPoints {
			rowKey := getRowKey(retIdx, md.Id, m, retention.SecondsPerPoint)
			encodedChunks := encodedChunksFromPoints(p, uint32(retention.SecondsPerPoint), retention.ChunkSpan)
			archives = append(archives, archive.Archive{
				SecondsPerPoint: uint32(retention.SecondsPerPoint),
				Points:          uint32(retention.NumberOfPoints),
				Chunks:          encodedChunks,
				RowKey:          rowKey,
			})
			if int64(p[len(p)-1].Timestamp) > md.Time {
				md.Time = int64(p[len(p)-1].Timestamp)
			}
		}
	}

	res.AggregationMethod = uint32(w.Header.Metadata.AggregationMethod)
	res.MetricData = *md
	res.Archives = archives
	return res, nil
}

func getRowKey(retIdx int, id, meth string, secondsPerPoint int) string {
	if retIdx == 0 {
		return id
	} else {
		return api.AggMetricKey(
			id,
			meth,
			uint32(secondsPerPoint),
		)
	}
}

func encodedChunksFromPoints(points []whisper.Point, intervalIn, chunkSpan uint32) []chunk.IterGen {
	var point whisper.Point
	var t0, prevT0 uint32
	var c *chunk.Chunk
	var encodedChunks []chunk.IterGen

	for _, point = range points {
		// this shouldn't happen, but if it would we better catch it here because Metrictank wouldn't handle it well:
		// https://github.com/raintank/metrictank/blob/f1868cccfb92fc82cd853914af958f6d187c5f74/mdata/aggmetric.go#L378
		if point.Timestamp == 0 {
			continue
		}

		t0 = point.Timestamp - (point.Timestamp % chunkSpan)
		if prevT0 == 0 {
			c = chunk.New(t0)
			prevT0 = t0
		} else if prevT0 != t0 {
			c.Finish()

			encodedChunks = append(encodedChunks, *chunk.NewBareIterGen(c.Bytes(), c.T0, chunkSpan))

			c = chunk.New(t0)
			prevT0 = t0
		}

		err := c.Push(point.Timestamp, point.Value)
		if err != nil {
			panic(fmt.Sprintf("ERROR: Failed to push value into chunk at t0 %d: %q", t0, err))
		}
	}

	// if the last written point was also the last one of the current chunk,
	// or if writeUnfinishedChunks is on, we close the chunk and push it
	if point.Timestamp == t0+chunkSpan-intervalIn || *writeUnfinishedChunks {
		c.Finish()
		encodedChunks = append(encodedChunks, *chunk.NewBareIterGen(c.Bytes(), c.T0, chunkSpan))
	}

	return encodedChunks
}

// scan a directory and feed the list of whisper files relative to base into the given channel
func getFileListIntoChan(fileChan chan string) {
	filepath.Walk(
		*whisperDirectory,
		func(path string, info os.FileInfo, err error) error {
			name := getMetricName(path)
			if !nameFilter.Match([]byte(getMetricName(name))) {
				log.Info(fmt.Sprintf("Skipping file %s with name %s", file, name))
				continue
			}
			if len(path) >= 4 && path[len(path)-4:] == ".wsp" {
				fileChan <- path
			}
			return nil
		},
	)

	close(fileChan)
}
<|endoftext|>