<commit_before>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/hashicorp/go-cleanhttp"

	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"

	"github.com/Cox-Automotive/alks-go"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
)

// Version number, to be injected at link time
// to set, add `-ldflags "-X main.versionNumber=1.2.3"` to the go build command
var versionNumber string
var ErrNoValidCredentialSources = errors.New(`No valid credential sources found for ALKS Provider.
Please see https://github.com/Cox-Automotive/terraform-provider-alks#authentication for more information on
providing credentials for the ALKS Provider`)

// Config stores ALKS configuration and credentials
type Config struct {
	URL           string
	AccessKey     string
	SecretKey     string
	Token         string
	CredsFilename string
	Profile       string
	AssumeRole    assumeRoleDetails
	Account       string
	Role          string
}

type assumeRoleDetails struct {
	RoleARN     string
	SessionName string
	ExternalID  string
	Policy      string
}

func getCredentials(c *Config) *credentials.Credentials {
	// Follow the  same priority as the AWS Terraform Provider
	// https://www.terraform.io/docs/providers/aws/#authentication

	// needed for the EC2MetaData service
	sess := session.Must(session.NewSession())

	providers := []credentials.Provider{
		&credentials.StaticProvider{Value: credentials.Value{
			AccessKeyID:     c.AccessKey,
			SecretAccessKey: c.SecretKey,
			SessionToken:    c.Token,
		}},
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{
			Filename: c.CredsFilename,
			Profile:  c.Profile,
		},
		&ec2rolecreds.EC2RoleProvider{
			Client: ec2metadata.New(sess),
		},
	}

	// Check for ECS container, for more details see:
	// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html
	if uri := os.Getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"); len(uri) > 0 {
		client := cleanhttp.DefaultClient()
		client.Timeout = 100 * time.Millisecond
		cfg := &aws.Config{
			HTTPClient: client,
		}

		providers = append(providers, defaults.RemoteCredProvider(*cfg, defaults.Handlers()))
	}

	return credentials.NewChainCredentials(providers)
}

func getCredentialsFromSession(c *Config) (*credentials.Credentials, error) {
	var sess *session.Session
	var err error
	if c.Profile == "" {
		sess, err = session.NewSession()
		if err != nil {
			return nil, ErrNoValidCredentialSources
		}
	} else {
		options := &session.Options{
			Config: aws.Config{
				HTTPClient: cleanhttp.DefaultClient(),
				MaxRetries: aws.Int(0),
				Region:     aws.String("us-east-1"),
			},
		}
		options.Profile = c.Profile
		options.SharedConfigState = session.SharedConfigEnable

		sess, err = session.NewSessionWithOptions(*options)
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
				return nil, ErrNoValidCredentialSources
			}
			return nil, fmt.Errorf("Error creating AWS session: %s", err)
		}
	}
	creds := sess.Config.Credentials
	_, err = sess.Config.Credentials.Get()
	if err != nil {
		return nil, ErrNoValidCredentialSources
	}
	return creds, nil
}

// Client returns a properly configured ALKS client or an appropriate error if initialization fails
func (c *Config) Client() (*alks.Client, error) {
	log.Println("[DEBUG] Validating STS credentials") // TODO: Fix typo.

	// lookup credentials
	creds := getCredentials(c)
	cp, cpErr := creds.Get()

	// validate we have credentials
	if cpErr != nil {
		if awsErr, ok := cpErr.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
			var err error
			creds, err = getCredentialsFromSession(c)
			if err != nil {
				return nil, err
			}
			cp, cpErr = creds.Get()
		}
	}
	if cpErr != nil {
		return nil, ErrNoValidCredentialSources
	}

	// create a new session to test credentails
	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String("us-east-1"),
		Credentials: creds,
	})

	// validate session
	if err != nil {
		return nil, fmt.Errorf("Error creating session from STS. (%v)", err)
	}

	var stsconn *sts.STS
	// we need to assume another role before creating an ALKS client
	if c.AssumeRole.RoleARN != "" {
		arCreds := stscreds.NewCredentials(sess, c.AssumeRole.RoleARN, func(p *stscreds.AssumeRoleProvider) {
			if c.AssumeRole.SessionName != "" {
				p.RoleSessionName = c.AssumeRole.SessionName
			}

			if c.AssumeRole.ExternalID != "" {
				p.ExternalID = &c.AssumeRole.ExternalID
			}

			if c.AssumeRole.Policy != "" {
				p.Policy = &c.AssumeRole.Policy
			}
		})

		cp, cpErr = arCreds.Get()
		if cpErr != nil {
			return nil, fmt.Errorf("The role %q cannot be assumed. Please verify the role ARN, role policies and your base AWS credentials", c.AssumeRole.RoleARN)
		}

		stsconn = sts.New(sess, &aws.Config{
			Region:      aws.String("us-east-1"),
			Credentials: arCreds,
		})
	} else {
		stsconn = sts.New(sess)
	}

	// make a basic api call to test creds are valid
	cident, serr := stsconn.GetCallerIdentity(&sts.GetCallerIdentityInput{})

	// check for valid creds
	if serr != nil {
		return nil, serr
	}

	// got good creds, create alks sts client
	client, err := alks.NewSTSClient(c.URL, cp.AccessKeyID, cp.SecretAccessKey, cp.SessionToken)

	if err != nil {
		return nil, err
	}

	// 1. Check if calling for a specific account
	if len(c.Account) > 0 && len(c.Role) > 0 {

		// 2. Generate client specified
		newClient, err := generateNewClient(c, cident, client)

		if err != nil {
			return nil, err
		}

		return newClient, nil
	}

	// check if the user is using a assume-role IAM admin session or MI.
	if isValidIAM(cident.Arn, client) != true {
		return nil, errors.New("Looks like you are not using ALKS IAM credentials. This will result in errors when creating roles. \n " +
			"Note: If using ALKS CLI to get credentials, be sure to use the '-i' flag. \n Please see https://coxautoinc.sharepoint.com/sites/service-internal-tools-team/SitePages/ALKS-Terraform-Provider---Troubleshooting.aspx for more information.")
	}

	client.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

	log.Println("[INFO] ALKS Client configured")

	return client, nil
}

func getPluginVersion() string {
	if versionNumber != "" {
		return versionNumber
	}

	return "unknown"
}

/*
	Validates ARN for assumed-role of:
		- Admin
		- IAMAdmin
		- Machine Identities.
*/
func isValidIAM(arn *string, client *alks.Client) bool {
	// Check if Admin || IAMAdmin
	if strings.Contains(*arn, "assumed-role/Admin/") || strings.Contains(*arn, "assumed-role/IAMAdmin/") || strings.Contains(*arn, "assumed-role/LabAdmin/") {
		return true
	}

	// Check if MI...
	arnParts := strings.FieldsFunc(*arn, splitBy)
	iamArn := fmt.Sprintf("arn:aws:iam::%s:role/acct-managed/%s", arnParts[3], arnParts[5])

	_, err := client.SearchRoleMachineIdentity(iamArn)
	if err != nil {
		return false
	}
	return true
}

func splitBy(r rune) bool {
	return r == ':' || r == '/'
}

func generateNewClient(c *Config, cident *sts.GetCallerIdentityOutput, client *alks.Client) (*alks.Client, error) {

	// 3. Create account string
	newAccDetail := c.Account + "/ALKS" + c.Role

	// Calling for the same account; fine.
	if strings.Contains(newAccDetail, client.AccountDetails.Account) {

		// check if the user is using a assume-role IAM admin session or MI.
		if isValidIAM(cident.Arn, client) != true {
			return nil, errors.New("Looks like you are not using ALKS IAM credentials. This will result in errors when creating roles. \n " +
				"Note: If using ALKS CLI to get credentials, be sure to use the '-i' flag. \n Please see https://coxautoinc.sharepoint.com/sites/service-internal-tools-team/SitePages/ALKS-Terraform-Provider---Troubleshooting.aspx for more information.")
		}

		client.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

		log.Println("[INFO] ALKS Client configured")

		return client, nil
	} else {

		// 4. Alright, new credentials needed - swap em out.
		client.AccountDetails.Account = newAccDetail
		client.AccountDetails.Role = c.Role

		newCreds, _ := client.CreateIamSession()
		newClient, err := alks.NewSTSClient(c.URL, newCreds.AccessKey, newCreds.SecretKey, newCreds.SessionToken)

		if err != nil {
			return nil, err
		}

		if isValidIAM(cident.Arn, newClient) != true {
			return nil, errors.New("Looks like you are not using ALKS IAM credentials. This will result in errors when creating roles. \n " +
				"Note: If using ALKS CLI to get credentials, be sure to use the '-i' flag. \n Please see https://coxautoinc.sharepoint.com/sites/service-internal-tools-team/SitePages/ALKS-Terraform-Provider---Troubleshooting.aspx for more information.")
		}

		newClient.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

		log.Println("[INFO] ALKS Client configured")

		// 5. Return this new client for provider
		return newClient, nil
	}
}
<commit_msg>Typo fixed.<commit_after>package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/hashicorp/go-cleanhttp"

	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/defaults"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"

	"github.com/Cox-Automotive/alks-go"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
)

// Version number, to be injected at link time
// to set, add `-ldflags "-X main.versionNumber=1.2.3"` to the go build command
var versionNumber string
var ErrNoValidCredentialSources = errors.New(`No valid credential sources found for ALKS Provider.
Please see https://github.com/Cox-Automotive/terraform-provider-alks#authentication for more information on
providing credentials for the ALKS Provider`)

// Config stores ALKS configuration and credentials
type Config struct {
	URL           string
	AccessKey     string
	SecretKey     string
	Token         string
	CredsFilename string
	Profile       string
	AssumeRole    assumeRoleDetails
	Account       string
	Role          string
}

type assumeRoleDetails struct {
	RoleARN     string
	SessionName string
	ExternalID  string
	Policy      string
}

func getCredentials(c *Config) *credentials.Credentials {
	// Follow the  same priority as the AWS Terraform Provider
	// https://www.terraform.io/docs/providers/aws/#authentication

	// needed for the EC2MetaData service
	sess := session.Must(session.NewSession())

	providers := []credentials.Provider{
		&credentials.StaticProvider{Value: credentials.Value{
			AccessKeyID:     c.AccessKey,
			SecretAccessKey: c.SecretKey,
			SessionToken:    c.Token,
		}},
		&credentials.EnvProvider{},
		&credentials.SharedCredentialsProvider{
			Filename: c.CredsFilename,
			Profile:  c.Profile,
		},
		&ec2rolecreds.EC2RoleProvider{
			Client: ec2metadata.New(sess),
		},
	}

	// Check for ECS container, for more details see:
	// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-iam-roles.html
	if uri := os.Getenv("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI"); len(uri) > 0 {
		client := cleanhttp.DefaultClient()
		client.Timeout = 100 * time.Millisecond
		cfg := &aws.Config{
			HTTPClient: client,
		}

		providers = append(providers, defaults.RemoteCredProvider(*cfg, defaults.Handlers()))
	}

	return credentials.NewChainCredentials(providers)
}

func getCredentialsFromSession(c *Config) (*credentials.Credentials, error) {
	var sess *session.Session
	var err error
	if c.Profile == "" {
		sess, err = session.NewSession()
		if err != nil {
			return nil, ErrNoValidCredentialSources
		}
	} else {
		options := &session.Options{
			Config: aws.Config{
				HTTPClient: cleanhttp.DefaultClient(),
				MaxRetries: aws.Int(0),
				Region:     aws.String("us-east-1"),
			},
		}
		options.Profile = c.Profile
		options.SharedConfigState = session.SharedConfigEnable

		sess, err = session.NewSessionWithOptions(*options)
		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
				return nil, ErrNoValidCredentialSources
			}
			return nil, fmt.Errorf("Error creating AWS session: %s", err)
		}
	}
	creds := sess.Config.Credentials
	_, err = sess.Config.Credentials.Get()
	if err != nil {
		return nil, ErrNoValidCredentialSources
	}
	return creds, nil
}

// Client returns a properly configured ALKS client or an appropriate error if initialization fails
func (c *Config) Client() (*alks.Client, error) {
	log.Println("[DEBUG] Validating STS credentials")

	// lookup credentials
	creds := getCredentials(c)
	cp, cpErr := creds.Get()

	// validate we have credentials
	if cpErr != nil {
		if awsErr, ok := cpErr.(awserr.Error); ok && awsErr.Code() == "NoCredentialProviders" {
			var err error
			creds, err = getCredentialsFromSession(c)
			if err != nil {
				return nil, err
			}
			cp, cpErr = creds.Get()
		}
	}
	if cpErr != nil {
		return nil, ErrNoValidCredentialSources
	}

	// create a new session to test credentails
	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String("us-east-1"),
		Credentials: creds,
	})

	// validate session
	if err != nil {
		return nil, fmt.Errorf("Error creating session from STS. (%v)", err)
	}

	var stsconn *sts.STS
	// we need to assume another role before creating an ALKS client
	if c.AssumeRole.RoleARN != "" {
		arCreds := stscreds.NewCredentials(sess, c.AssumeRole.RoleARN, func(p *stscreds.AssumeRoleProvider) {
			if c.AssumeRole.SessionName != "" {
				p.RoleSessionName = c.AssumeRole.SessionName
			}

			if c.AssumeRole.ExternalID != "" {
				p.ExternalID = &c.AssumeRole.ExternalID
			}

			if c.AssumeRole.Policy != "" {
				p.Policy = &c.AssumeRole.Policy
			}
		})

		cp, cpErr = arCreds.Get()
		if cpErr != nil {
			return nil, fmt.Errorf("The role %q cannot be assumed. Please verify the role ARN, role policies and your base AWS credentials", c.AssumeRole.RoleARN)
		}

		stsconn = sts.New(sess, &aws.Config{
			Region:      aws.String("us-east-1"),
			Credentials: arCreds,
		})
	} else {
		stsconn = sts.New(sess)
	}

	// make a basic api call to test creds are valid
	cident, serr := stsconn.GetCallerIdentity(&sts.GetCallerIdentityInput{})

	// check for valid creds
	if serr != nil {
		return nil, serr
	}

	// got good creds, create alks sts client
	client, err := alks.NewSTSClient(c.URL, cp.AccessKeyID, cp.SecretAccessKey, cp.SessionToken)

	if err != nil {
		return nil, err
	}

	// 1. Check if calling for a specific account
	if len(c.Account) > 0 && len(c.Role) > 0 {

		// 2. Generate client specified
		newClient, err := generateNewClient(c, cident, client)

		if err != nil {
			return nil, err
		}

		return newClient, nil
	}

	// check if the user is using a assume-role IAM admin session or MI.
	if isValidIAM(cident.Arn, client) != true {
		return nil, errors.New("Looks like you are not using ALKS IAM credentials. This will result in errors when creating roles. \n " +
			"Note: If using ALKS CLI to get credentials, be sure to use the '-i' flag. \n Please see https://coxautoinc.sharepoint.com/sites/service-internal-tools-team/SitePages/ALKS-Terraform-Provider---Troubleshooting.aspx for more information.")
	}

	client.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

	log.Println("[INFO] ALKS Client configured")

	return client, nil
}

func getPluginVersion() string {
	if versionNumber != "" {
		return versionNumber
	}

	return "unknown"
}

/*
	Validates ARN for assumed-role of:
		- Admin
		- IAMAdmin
		- Machine Identities.
*/
func isValidIAM(arn *string, client *alks.Client) bool {
	// Check if Admin || IAMAdmin
	if strings.Contains(*arn, "assumed-role/Admin/") || strings.Contains(*arn, "assumed-role/IAMAdmin/") || strings.Contains(*arn, "assumed-role/LabAdmin/") {
		return true
	}

	// Check if MI...
	arnParts := strings.FieldsFunc(*arn, splitBy)
	iamArn := fmt.Sprintf("arn:aws:iam::%s:role/acct-managed/%s", arnParts[3], arnParts[5])

	_, err := client.SearchRoleMachineIdentity(iamArn)
	if err != nil {
		return false
	}
	return true
}

func splitBy(r rune) bool {
	return r == ':' || r == '/'
}

func generateNewClient(c *Config, cident *sts.GetCallerIdentityOutput, client *alks.Client) (*alks.Client, error) {

	// 3. Create account string
	newAccDetail := c.Account + "/ALKS" + c.Role

	// Calling for the same account; fine.
	if strings.Contains(newAccDetail, client.AccountDetails.Account) {

		// check if the user is using a assume-role IAM admin session or MI.
		if isValidIAM(cident.Arn, client) != true {
			return nil, errors.New("Looks like you are not using ALKS IAM credentials. This will result in errors when creating roles. \n " +
				"Note: If using ALKS CLI to get credentials, be sure to use the '-i' flag. \n Please see https://coxautoinc.sharepoint.com/sites/service-internal-tools-team/SitePages/ALKS-Terraform-Provider---Troubleshooting.aspx for more information.")
		}

		client.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

		log.Println("[INFO] ALKS Client configured")

		return client, nil
	} else {

		// 4. Alright, new credentials needed - swap em out.
		client.AccountDetails.Account = newAccDetail
		client.AccountDetails.Role = c.Role

		newCreds, _ := client.CreateIamSession()
		newClient, err := alks.NewSTSClient(c.URL, newCreds.AccessKey, newCreds.SecretKey, newCreds.SessionToken)

		if err != nil {
			return nil, err
		}

		if isValidIAM(cident.Arn, newClient) != true {
			return nil, errors.New("Looks like you are not using ALKS IAM credentials. This will result in errors when creating roles. \n " +
				"Note: If using ALKS CLI to get credentials, be sure to use the '-i' flag. \n Please see https://coxautoinc.sharepoint.com/sites/service-internal-tools-team/SitePages/ALKS-Terraform-Provider---Troubleshooting.aspx for more information.")
		}

		newClient.SetUserAgent(fmt.Sprintf("alks-terraform-provider-%s", getPluginVersion()))

		log.Println("[INFO] ALKS Client configured")

		// 5. Return this new client for provider
		return newClient, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"github.com/matejkramny/gopos"
)

func main() {
	// your printer IP & Port
	// In this case, the printer has Ethernet interface and assigned a static ip
	// TM-T20II Ethernet
	printer := gopos.POS("192.168.18.249", 9100)

	gopos.ConnectRedis()
	defer (*gopos.Redis).Close()

	gopos.RedisListen(printer)

	// testing
	// pos.Print("I am NOT emphesized")
	// pos.LineFeed()

	// pos.SetEmphesized(gopos.MODE_ON)
	// pos.Print("I am emphesized")
	// pos.SetEmphesized(gopos.MODE_OFF)
	// pos.LineFeed()

	// pos.SetFont(gopos.MODE_ON)
	// pos.Print("I am in Font B")
	// pos.LineFeed()
	// pos.SetFont(gopos.MODE_OFF)

	// pos.SetFont(gopos.MODE_ON2)
	// pos.Print("I am in Font C")
	// pos.LineFeed()
	// pos.SetFont(gopos.MODE_OFF)

	// pos.SetReversePrint(gopos.MODE_ON)
	// pos.Print("I am reverse white/black thing")
	// pos.LineFeed()
	// pos.SetReversePrint(gopos.MODE_OFF)

	// pos.SetUnderline(gopos.MODE_ON)
	// pos.Print("I am in Underline mode")
	// pos.LineFeed()
	// pos.SetUnderline(gopos.MODE_OFF)

	// pos.SetUnderline(gopos.MODE_ON2)
	// pos.Print("I am in 2 Underline mode")
	// pos.LineFeed()
	// pos.SetUnderline(gopos.MODE_OFF)

	// pos.SetDoubleStrike(gopos.MODE_ON)
	// pos.Print("I am in Double strike")
	// pos.LineFeed()
	// pos.SetDoubleStrike(gopos.MODE_OFF)

	// pos.LineFeed()
	// pos.LineFeed()
	// pos.LineFeed()
	// pos.LineFeed()

	// pos.Cut()
}<commit_msg>Remove bufio import<commit_after>package main

import (
	"github.com/matejkramny/gopos"
)

func main() {
	// your printer IP & Port
	// In this case, the printer has Ethernet interface and assigned a static ip
	// TM-T20II Ethernet
	printer := gopos.POS("192.168.18.249", 9100)

	gopos.ConnectRedis()
	defer (*gopos.Redis).Close()

	gopos.RedisListen(printer)

	// testing
	// pos.Print("I am NOT emphesized")
	// pos.LineFeed()

	// pos.SetEmphesized(gopos.MODE_ON)
	// pos.Print("I am emphesized")
	// pos.SetEmphesized(gopos.MODE_OFF)
	// pos.LineFeed()

	// pos.SetFont(gopos.MODE_ON)
	// pos.Print("I am in Font B")
	// pos.LineFeed()
	// pos.SetFont(gopos.MODE_OFF)

	// pos.SetFont(gopos.MODE_ON2)
	// pos.Print("I am in Font C")
	// pos.LineFeed()
	// pos.SetFont(gopos.MODE_OFF)

	// pos.SetReversePrint(gopos.MODE_ON)
	// pos.Print("I am reverse white/black thing")
	// pos.LineFeed()
	// pos.SetReversePrint(gopos.MODE_OFF)

	// pos.SetUnderline(gopos.MODE_ON)
	// pos.Print("I am in Underline mode")
	// pos.LineFeed()
	// pos.SetUnderline(gopos.MODE_OFF)

	// pos.SetUnderline(gopos.MODE_ON2)
	// pos.Print("I am in 2 Underline mode")
	// pos.LineFeed()
	// pos.SetUnderline(gopos.MODE_OFF)

	// pos.SetDoubleStrike(gopos.MODE_ON)
	// pos.Print("I am in Double strike")
	// pos.LineFeed()
	// pos.SetDoubleStrike(gopos.MODE_OFF)

	// pos.LineFeed()
	// pos.LineFeed()
	// pos.LineFeed()
	// pos.LineFeed()

	// pos.Cut()
}<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
)

const (
	DEFAULT_SSH_COMMAND = "ssh"
	APPNAME             = "novassh"
	VERSION             = "0.1"
)

type Config struct {
	// Executable name of SSH
	SshCommand string

	// Option flags for SSH command
	SshOptions []string

	// Hostname to connect to the instance
	SshHost string

	// Username of SSH
	SshUser string

	// Command-name to be run on the instance
	SshRemoteCommand string
}

func (c *Config) ParseArgs(args []string) (exitWithHelp bool, err error) {
	sshargs := []string{}

	for _, arg := range args {
		if arg == "--novassh-debug" {
			// Enable debug
			log.SetLevel(log.DebugLevel)
			enableDebugTransport()

		} else if arg == "--novassh-command" {
			// Detects SSH command
			c.SshCommand = arg

		} else if arg == "--help" {
			sshargs = []string{}
			break

		} else {
			sshargs = append(sshargs, arg)
		}
	}

	// Set default SSH command if not set
	if c.SshCommand == "" {
		c.SshCommand = DEFAULT_SSH_COMMAND
	}

	// Any SSH args was not provided.
	if len(sshargs) == 0 {
		help()
		return true, nil
	}

	return false, c.parseSshArgs(sshargs)
}

func (c *Config) parseSshArgs(args []string) (err error) {
	nova := NewNova()
	if err := nova.Init(); err != nil {
		return err
	}

	found := false
	pos := len(args) - 1 // position of machine name in arguments
	for pos >= 0 {
		arg := args[pos]
		found, err = c.resolveMachineName(nova, arg)
		if err != nil {
			return err

		} else if found {
			break
		}
		pos--
	}

	if found {
		if pos > 0 {
			c.SshOptions = args[:pos]
		}
		if len(args) > 1 {
			c.SshRemoteCommand = strings.Join(args[pos+1:], " ")
		}
		log.Debugf("The machine is found: ipaddr=%s, args=%v command=%s", c.SshHost, c.SshOptions, c.SshRemoteCommand)
		return nil

	} else {
		return fmt.Errorf("Could not found the machine.")
	}
}

func (c *Config) resolveMachineName(nova *nova, arg string) (found bool, err error) {
	var user, instancename string

	if strings.Index(arg, "@") > 0 {
		ss := strings.Split(arg, "@")
		user = ss[0]
		instancename = ss[1]
	} else {
		instancename = arg
	}

	log.Debugf("Try to find the machine: instance-name=%s", instancename)

	machine, err := nova.Find(instancename)
	if err != nil {
		// error
		return false, err

	} else if machine != nil {
		// Found
		c.SshHost = machine.Ipaddr
		c.SshUser = user
		return true, nil

	} else {
		// Not found
		log.Debugf("No match: name=%s", instancename)
		return false, nil
	}
}

func help() {
	fmt.Fprintf(os.Stdout, `NAME:
	%s - The ssh wrapper program to connect OpenStack instance(nova) with the instance name.

USAGE:
	%s [ssh-options] user@hostname [comamnd]

VERSION:
	%s

OPTIONS:
	--novassh-debug: output some debug messages
	--help:          print this message

`, APPNAME, APPNAME, VERSION)
}
<commit_msg>--novassh-commandオプションを追加<commit_after>package main

import (
	"fmt"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
)

const (
	DEFAULT_SSH_COMMAND = "ssh"
	APPNAME             = "novassh"
	VERSION             = "0.1"
)

type Config struct {
	// Executable name of SSH
	SshCommand string

	// Option flags for SSH command
	SshOptions []string

	// Hostname to connect to the instance
	SshHost string

	// Username of SSH
	SshUser string

	// Command-name to be run on the instance
	SshRemoteCommand string
}

func (c *Config) ParseArgs(args []string) (exitWithHelp bool, err error) {
	// Environments
	if os.Getenv("NOVASSH_COMMAND") != "" {
		c.SshCommand = os.Getenv("NOVASSH_COMMAND")
	}

	// Aeguments
	i := 0
	sshargs := []string{}
	for i < len(args) {
		arg := args[i]
		if arg == "--novassh-debug" {
			// Enable debug
			log.SetLevel(log.DebugLevel)
			enableDebugTransport()

		} else if arg == "--novassh-command" {
			// Detects SSH command
			i++
			c.SshCommand = args[i]

		} else if arg == "--help" {
			sshargs = []string{}
			break

		} else {
			sshargs = append(sshargs, arg)
		}
		i++
	}

	// Set default SSH command if not set
	if c.SshCommand == "" {
		c.SshCommand = DEFAULT_SSH_COMMAND
	}

	// Any SSH args was not provided.
	if len(sshargs) == 0 {
		help()
		return true, nil
	}

	return false, c.parseSshArgs(sshargs)
}

func (c *Config) parseSshArgs(args []string) (err error) {
	nova := NewNova()
	if err := nova.Init(); err != nil {
		return err
	}

	found := false
	pos := len(args) - 1 // position of machine name in arguments
	for pos >= 0 {
		arg := args[pos]
		found, err = c.resolveMachineName(nova, arg)
		if err != nil {
			return err

		} else if found {
			break
		}
		pos--
	}

	if found {
		if pos > 0 {
			c.SshOptions = args[:pos]
		}
		if len(args) > 1 {
			c.SshRemoteCommand = strings.Join(args[pos+1:], " ")
		}
		log.Debugf("The machine is found: ipaddr=%s, args=%v command=%s", c.SshHost, c.SshOptions, c.SshRemoteCommand)
		return nil

	} else {
		return fmt.Errorf("Could not found the machine.")
	}
}

func (c *Config) resolveMachineName(nova *nova, arg string) (found bool, err error) {
	var user, instancename string

	if strings.Index(arg, "@") > 0 {
		ss := strings.Split(arg, "@")
		user = ss[0]
		instancename = ss[1]
	} else {
		instancename = arg
	}

	log.Debugf("Try to find the machine: instance-name=%s", instancename)

	machine, err := nova.Find(instancename)
	if err != nil {
		// error
		return false, err

	} else if machine != nil {
		// Found
		c.SshHost = machine.Ipaddr
		c.SshUser = user
		return true, nil

	} else {
		// Not found
		log.Debugf("No match: name=%s", instancename)
		return false, nil
	}
}

func help() {
	fmt.Fprintf(os.Stdout, `NAME:
	%s - The ssh wrapper program to connect OpenStack instance(nova) with the instance name.

USAGE:
	%s [ssh-options] user@hostname [comamnd]

VERSION:
	%s

OPTIONS:
	--novassh-command: Specify SSH command (default: "ssh").
	--novassh-debug:   output some debug messages.
	--help:            print this message.

ENVIRONMENTS:
	NOVASSH_COMMAND: Specify SSH command (default: "ssh").

`, APPNAME, APPNAME, VERSION)
}
<|endoftext|>


<commit_before>package service_config

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/fraenkel/candiedyaml"
)

const (
	ConfigPathEnvVar = "CONFIG_PATH"
	ConfigEnvVar     = "CONFIG"
)

var NoConfigError = errors.New("No Config or Config Path Specified. Please supply one of the following: -config, -configPath, CONFIG, or CONFIG_PATH")

type ServiceConfig struct {
	configFlag     string
	configPathFlag string
	flagSet        *flag.FlagSet
	defaultModel   interface{}
	helpWriter     io.Writer
}

func New() *ServiceConfig {
	return &ServiceConfig{
		helpWriter: os.Stderr,
	}
}

func (c *ServiceConfig) AddFlags(flagSet *flag.FlagSet) {
	c.flagSet = flagSet
	c.flagSet.StringVar(&c.configFlag, "config", "", "json encoded configuration string")
	c.flagSet.StringVar(&c.configPathFlag, "configPath", "", "path to configuration file with json encoded content")

	c.flagSet.SetOutput(c.helpWriter)
	c.flagSet.Usage = c.PrintUsage
}

func (c ServiceConfig) ConfigBytes() ([]byte, error) {
	if c.configFlag != "" {
		return []byte(c.configFlag), nil
	}

	if c.configPathFlag != "" {
		absolutePath, err := filepath.Abs(c.configPathFlag)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Making config file path absolute: %s", err.Error()))
		}

		bytes, err := ioutil.ReadFile(absolutePath)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Reading config file: %s", err.Error()))
		}

		return bytes, nil
	}

	config := os.Getenv(ConfigEnvVar)
	if config != "" {
		return []byte(config), nil
	}

	configPath := os.Getenv(ConfigPathEnvVar)
	if configPath != "" {
		absolutePath, err := filepath.Abs(configPath)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Making config file path absolute: %s", err.Error()))
		}

		bytes, err := ioutil.ReadFile(absolutePath)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Reading config file: %s", err.Error()))
		}

		return bytes, nil
	}

	return nil, NoConfigError
}

func (c ServiceConfig) ConfigPath() string {
	return c.configPathFlag
}

func (c ServiceConfig) Read(model interface{}) error {
	bytes, err := c.ConfigBytes()
	if err != nil {
		return err
	}

	reader := NewReader(bytes)

	if c.defaultModel != nil {
		err = reader.ReadWithDefaults(model, c.defaultModel)
	} else {
		err = reader.Read(model)
	}

	if err != nil {
		return err
	}

	return nil
}

func (c *ServiceConfig) AddDefaults(defaultModel interface{}) {
	c.defaultModel = defaultModel
}

func (c ServiceConfig) PrintUsage() {
	fmt.Fprint(c.helpWriter, "Expected usage:\n")
	c.flagSet.PrintDefaults()

	if c.defaultModel != nil {
		defaultStr, err := candiedyaml.Marshal(c.defaultModel)
		if err != nil {
			fmt.Fprintf(c.helpWriter, "Error printing defaults: %v", err)
		} else {
			fmt.Fprintf(c.helpWriter, "Default config values:\n%s", defaultStr)
		}
	}
}
<commit_msg>Fixed ordering issue for the Addflag() and AddDefault() to be called in any order<commit_after>package service_config

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/fraenkel/candiedyaml"
)

const (
	ConfigPathEnvVar = "CONFIG_PATH"
	ConfigEnvVar     = "CONFIG"
)

var NoConfigError = errors.New("No Config or Config Path Specified. Please supply one of the following: -config, -configPath, CONFIG, or CONFIG_PATH")

type ServiceConfig struct {
	configFlag     string
	configPathFlag string
	flagSet        *flag.FlagSet
	defaultModel   interface{}
	helpWriter     io.Writer
}

func New() *ServiceConfig {
	return &ServiceConfig{
		helpWriter: os.Stderr,
	}
}

func (c *ServiceConfig) AddFlags(flagSet *flag.FlagSet) {
	c.flagSet = flagSet
	c.flagSet.StringVar(&c.configFlag, "config", "", "json encoded configuration string")
	c.flagSet.StringVar(&c.configPathFlag, "configPath", "", "path to configuration file with json encoded content")

	c.flagSet.SetOutput(c.helpWriter)
	c.flagSet.Usage = func() {
		c.PrintUsage()
	}
}

func (c ServiceConfig) ConfigBytes() ([]byte, error) {
	if c.configFlag != "" {
		return []byte(c.configFlag), nil
	}

	if c.configPathFlag != "" {
		absolutePath, err := filepath.Abs(c.configPathFlag)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Making config file path absolute: %s", err.Error()))
		}

		bytes, err := ioutil.ReadFile(absolutePath)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Reading config file: %s", err.Error()))
		}

		return bytes, nil
	}

	config := os.Getenv(ConfigEnvVar)
	if config != "" {
		return []byte(config), nil
	}

	configPath := os.Getenv(ConfigPathEnvVar)
	if configPath != "" {
		absolutePath, err := filepath.Abs(configPath)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Making config file path absolute: %s", err.Error()))
		}

		bytes, err := ioutil.ReadFile(absolutePath)
		if err != nil {
			return nil, errors.New(fmt.Sprintf("Reading config file: %s", err.Error()))
		}

		return bytes, nil
	}

	return nil, NoConfigError
}

func (c ServiceConfig) ConfigPath() string {
	return c.configPathFlag
}

func (c ServiceConfig) Read(model interface{}) error {
	bytes, err := c.ConfigBytes()
	if err != nil {
		return err
	}

	reader := NewReader(bytes)

	if c.defaultModel != nil {
		err = reader.ReadWithDefaults(model, c.defaultModel)
	} else {
		err = reader.Read(model)
	}

	if err != nil {
		return err
	}

	return nil
}

func (c *ServiceConfig) AddDefaults(defaultModel interface{}) {
	c.defaultModel = defaultModel
}

func (c ServiceConfig) PrintUsage() {
	fmt.Fprint(c.helpWriter, "Expected usage:\n")
	c.flagSet.PrintDefaults()

	if c.defaultModel != nil {
		defaultStr, err := candiedyaml.Marshal(c.defaultModel)
		if err != nil {
			fmt.Fprintf(c.helpWriter, "Error printing defaults: %v", err)
		} else {
			fmt.Fprintf(c.helpWriter, "Default config values:\n%s", defaultStr)
		}
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/Azure/acs-engine/pkg/acsengine"
	"github.com/Azure/acs-engine/pkg/api"
	"github.com/Azure/acs-engine/pkg/armhelpers"
	"github.com/Azure/acs-engine/pkg/i18n"
	"github.com/Azure/acs-engine/pkg/operations"
	"github.com/leonelquinteros/gotext"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

type scaleCmd struct {
	authArgs

	// user input
	resourceGroupName    string
	deploymentDirectory  string
	newDesiredAgentCount int
	containerService     *api.ContainerService
	apiVersion           string
	location             string
	agentPoolToScale     string
	classicMode          bool

	// derived
	apiModelPath   string
	agentPool      *api.AgentPoolProfile
	client         armhelpers.ACSEngineClient
	locale         *gotext.Locale
	nameSuffix     string
	agentPoolIndex int
	masterFQDN     string
	logger         *log.Entry
}

const (
	scaleName             = "scale"
	scaleShortDescription = "scale a deployed cluster"
	scaleLongDescription  = "scale a deployed cluster"
)

// NewScaleCmd run a command to upgrade a Kubernetes cluster
func newScaleCmd() *cobra.Command {
	sc := scaleCmd{}

	scaleCmd := &cobra.Command{
		Use:   scaleName,
		Short: scaleShortDescription,
		Long:  scaleLongDescription,
		RunE: func(cmd *cobra.Command, args []string) error {
			return sc.run(cmd, args)
		},
	}

	f := scaleCmd.Flags()
	f.StringVar(&sc.location, "location", "", "location the cluster is deployed in")
	f.StringVar(&sc.resourceGroupName, "resource-group", "", "the resource group where the cluster is deployed")
	f.StringVar(&sc.deploymentDirectory, "deployment-dir", "", "the location of the output from `generate`")
	f.IntVar(&sc.newDesiredAgentCount, "new-node-count", 0, "desired number of nodes")
	f.BoolVar(&sc.classicMode, "classic-mode", false, "enable classic parameters and outputs")
	f.StringVar(&sc.agentPoolToScale, "node-pool", "", "node pool to scale")
	f.StringVar(&sc.masterFQDN, "master-FQDN", "", "FQDN for the master load balancer, Needed to scale down Kubernetes agent pools")

	addAuthFlags(&sc.authArgs, f)

	return scaleCmd
}

func (sc *scaleCmd) validate(cmd *cobra.Command, args []string) {
	log.Infoln("validating...")
	sc.logger = log.New().WithField("source", "scaling command line")
	var err error

	sc.locale, err = i18n.LoadTranslations()
	if err != nil {
		log.Fatalf("error loading translation files: %s", err.Error())
	}

	if sc.resourceGroupName == "" {
		cmd.Usage()
		log.Fatal("--resource-group must be specified")
	}

	if sc.location == "" {
		cmd.Usage()
		log.Fatal("--location must be specified")
	}

	if sc.newDesiredAgentCount == 0 {
		cmd.Usage()
		log.Fatal("--new-node-count must be specified")
	}

	if sc.client, err = sc.authArgs.getClient(); err != nil {
		log.Error("Failed to get client:", err)
	}

	if sc.deploymentDirectory == "" {
		cmd.Usage()
		log.Fatal("--deployment-dir must be specified")
	}

	_, err = sc.client.EnsureResourceGroup(sc.resourceGroupName, sc.location, nil)
	if err != nil {
		log.Fatalln(err)
	}

	// load apimodel from the deployment directory
	sc.apiModelPath = path.Join(sc.deploymentDirectory, "apimodel.json")

	if _, err = os.Stat(sc.apiModelPath); os.IsNotExist(err) {
		log.Fatalf("specified api model does not exist (%s)", sc.apiModelPath)
	}

	apiloader := &api.Apiloader{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}
	sc.containerService, sc.apiVersion, err = apiloader.LoadContainerServiceFromFile(sc.apiModelPath, true, true, nil)
	if err != nil {
		log.Fatalf("error parsing the api model: %s", err.Error())
	}

	if sc.agentPoolToScale == "" {
		agentPoolCount := len(sc.containerService.Properties.AgentPoolProfiles)
		if agentPoolCount > 1 {
			log.Fatal("--node-pool is required if more than one agent pool is defined in the container service")
		} else if agentPoolCount == 1 {
			sc.agentPool = sc.containerService.Properties.AgentPoolProfiles[0]
			sc.agentPoolIndex = 0
			sc.agentPoolToScale = sc.containerService.Properties.AgentPoolProfiles[0].Name
		} else {
			log.Fatal("No node pools found to scale")
		}
	} else {
		agentPoolIndex := -1
		for i, pool := range sc.containerService.Properties.AgentPoolProfiles {
			if pool.Name == sc.agentPoolToScale {
				agentPoolIndex = i
				sc.agentPool = pool
				sc.agentPoolIndex = i
			}
		}
		if agentPoolIndex == -1 {
			log.Fatalf("node pool %s wasn't in the deployed api model", sc.agentPoolToScale)
		}
	}

	templatePath := path.Join(sc.deploymentDirectory, "azuredeploy.json")
	contents, _ := ioutil.ReadFile(templatePath)

	var template interface{}
	json.Unmarshal(contents, &template)

	templateMap := template.(map[string]interface{})
	templateParameters := templateMap["parameters"].(map[string]interface{})

	nameSuffixParam := templateParameters["nameSuffix"].(map[string]interface{})
	sc.nameSuffix = nameSuffixParam["defaultValue"].(string)
	log.Infoln(fmt.Sprintf("Name suffix: %s", sc.nameSuffix))
}

func (sc *scaleCmd) run(cmd *cobra.Command, args []string) error {
	sc.validate(cmd, args)

	orchestratorInfo := sc.containerService.Properties.OrchestratorProfile
	var currentNodeCount, highestUsedIndex int
	indexes := make([]int, 0)
	indexToVM := make(map[int]string)
	if sc.agentPool.IsAvailabilitySets() {
		//TODO handle when there is a nextLink in the response and get more nodes
		vms, err := sc.client.ListVirtualMachines(sc.resourceGroupName)
		if err != nil {
			log.Fatalln("failed to get vms in the resource group. Error: %s", err.Error())
		}
		for _, vm := range *vms.Value {

			poolName, nameSuffix, index, err := armhelpers.K8sLinuxVMNameParts(*vm.Name)
			if err != nil || !strings.EqualFold(poolName, sc.agentPoolToScale) || !strings.EqualFold(nameSuffix, sc.nameSuffix) {
				continue
			}

			indexToVM[index] = *vm.Name
			indexes = append(indexes, index)
		}
		sortedIndexes := sort.IntSlice(indexes)
		sortedIndexes.Sort()
		indexes = []int(sortedIndexes)
		currentNodeCount = len(indexes)

		if currentNodeCount == sc.newDesiredAgentCount {
			return nil
		}
		highestUsedIndex = indexes[len(indexes)-1]

		// Scale down Scenario
		if currentNodeCount > sc.newDesiredAgentCount {
			if sc.masterFQDN == "" {
				cmd.Usage()
				log.Fatal("master-FQDN is required to scale down a kubernetes cluster's agent pool")
			}

			vmsToDelete := make([]string, 0)
			for i := currentNodeCount - 1; i >= sc.newDesiredAgentCount; i-- {
				vmsToDelete = append(vmsToDelete, indexToVM[i])
			}

			if orchestratorInfo.OrchestratorType == api.Kubernetes {
				err = sc.drainNodes(vmsToDelete)
				if err != nil {
					log.Errorf("Got error %+v, while draining the nodes to be deleted", err)
					return err
				}
			}

			errList := operations.ScaleDownVMs(sc.client, sc.logger, sc.resourceGroupName, vmsToDelete...)
			if errList != nil {
				errorMessage := ""
				for element := errList.Front(); element != nil; element = element.Next() {
					vmError, ok := element.Value.(*operations.VMScalingErrorDetails)
					if ok {
						error := fmt.Sprintf("Node '%s' failed to delete with error: '%s'", vmError.Name, vmError.Error.Error())
						errorMessage = errorMessage + error
					}
				}
				return fmt.Errorf(errorMessage)
			}

			return nil
		}
	} else {
		vmssList, err := sc.client.ListVirtualMachineScaleSets(sc.resourceGroupName)
		if err != nil {
			log.Fatalln("failed to get vmss list in the resource group. Error: %s", err.Error())
		}
		for _, vmss := range *vmssList.Value {
			poolName, nameSuffix, err := armhelpers.VmssNameParts(*vmss.Name)
			if err != nil || !strings.EqualFold(poolName, sc.agentPoolToScale) || !strings.EqualFold(nameSuffix, sc.nameSuffix) {
				continue
			}
			currentNodeCount = int(*vmss.Sku.Capacity)
			highestUsedIndex = 0
		}
	}

	ctx := acsengine.Context{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}
	templateGenerator, err := acsengine.InitializeTemplateGenerator(ctx, sc.classicMode)
	if err != nil {
		log.Fatalln("failed to initialize template generator: %s", err.Error())
	}

	sc.containerService.Properties.AgentPoolProfiles = []*api.AgentPoolProfile{sc.agentPool}

	template, parameters, _, err := templateGenerator.GenerateTemplate(sc.containerService, acsengine.DefaultGeneratorCode)
	if err != nil {
		log.Fatalf("error generating template %s: %s", sc.apiModelPath, err.Error())
		os.Exit(1)
	}

	if template, err = acsengine.PrettyPrintArmTemplate(template); err != nil {
		log.Fatalf("error pretty printing template: %s \n", err.Error())
	}

	templateJSON := make(map[string]interface{})
	parametersJSON := make(map[string]interface{})

	err = json.Unmarshal([]byte(template), &templateJSON)
	if err != nil {
		log.Fatalln(err)
	}

	err = json.Unmarshal([]byte(parameters), &parametersJSON)
	if err != nil {
		log.Fatalln(err)
	}

	transformer := acsengine.Transformer{Translator: ctx.Translator}
	// Our templates generate a range of nodes based on a count and offset, it is possible for there to be holes in the template
	// So we need to set the count in the template to get enough nodes for the range, if there are holes that number will be larger than the desired count
	countForTemplate := sc.newDesiredAgentCount
	if highestUsedIndex != 0 {
		countForTemplate += highestUsedIndex + 1 - currentNodeCount
	}
	addValue(parametersJSON, sc.agentPool.Name+"Count", countForTemplate)

	switch orchestratorInfo.OrchestratorType {
	case api.Kubernetes:
		err = transformer.NormalizeForK8sVMASScalingUp(sc.logger, templateJSON)
		if err != nil {
			log.Fatalf("error tranforming the template for scaling template %s: %s", sc.apiModelPath, err.Error())
			os.Exit(1)
		}
		if sc.agentPool.IsAvailabilitySets() {
			addValue(parametersJSON, fmt.Sprintf("%sOffset", sc.agentPool.Name), highestUsedIndex+1)
		}
		break
	case api.Swarm:
	case api.SwarmMode:
	case api.DCOS:
		if sc.agentPool.IsAvailabilitySets() {
			log.Fatalf("scaling isn't supported for orchestrator %s, with availability sets", orchestratorInfo.OrchestratorType)
			os.Exit(1)
		}
		transformer.NormalizeForVMSSScaling(sc.logger, templateJSON)
	}

	random := rand.New(rand.NewSource(time.Now().UnixNano()))
	deploymentSuffix := random.Int31()

	_, err = sc.client.DeployTemplate(
		sc.resourceGroupName,
		fmt.Sprintf("%s-%d", sc.resourceGroupName, deploymentSuffix),
		templateJSON,
		parametersJSON,
		nil)
	if err != nil {
		log.Fatalln(err)
	}

	apiloader := &api.Apiloader{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}
	var apiVersion string
	sc.containerService, apiVersion, err = apiloader.LoadContainerServiceFromFile(sc.apiModelPath, false, true, nil)
	sc.containerService.Properties.AgentPoolProfiles[sc.agentPoolIndex].Count = sc.newDesiredAgentCount

	b, e := apiloader.SerializeContainerService(sc.containerService, apiVersion)

	if e != nil {
		return e
	}

	f := acsengine.FileSaver{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}

	if e = f.SaveFile(sc.deploymentDirectory, "apimodel.json", b); e != nil {
		return e
	}

	return nil
}

type paramsMap map[string]interface{}

func addValue(m paramsMap, k string, v interface{}) {
	m[k] = paramsMap{
		"value": v,
	}
}

func (sc *scaleCmd) drainNodes(vmsToDelete []string) error {
	kubeConfig, err := acsengine.GenerateKubeConfig(sc.containerService.Properties, sc.location)
	if err != nil {
		log.Fatalf("failed to generate kube config") // TODO: cleanup
	}
	var errorMessage string
	masterURL := sc.masterFQDN
	if !strings.HasPrefix(masterURL, "https://") {
		masterURL = fmt.Sprintf("https://%s", masterURL)
	}
	numVmsToDrain := len(vmsToDelete)
	errChan := make(chan *operations.VMScalingErrorDetails, numVmsToDrain)
	defer close(errChan)
	for _, vmName := range vmsToDelete {
		go func(vmName string) {
			e := operations.SafelyDrainNode(sc.client, sc.logger,
				masterURL, kubeConfig, vmName, time.Duration(60)*time.Minute)
			if e != nil {
				log.Errorf("Failed to drain node %s, got error %s", vmName, e.Error())
				errChan <- &operations.VMScalingErrorDetails{Error: e, Name: vmName}
				return
			}
			errChan <- nil
		}(vmName)
	}

	for i := 0; i < numVmsToDrain; i++ {
		errDetails := <-errChan
		if errDetails != nil {
			error := fmt.Sprintf("Node '%s' failed to drain with error: '%s'", errDetails.Name, errDetails.Error.Error())
			errorMessage = errorMessage + error
			return fmt.Errorf(error)
		}
	}

	return nil
}
<commit_msg>adding a message when scaling command decides no action should be taken (#1723)<commit_after>package cmd

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"path"
	"sort"
	"strings"
	"time"

	"github.com/Azure/acs-engine/pkg/acsengine"
	"github.com/Azure/acs-engine/pkg/api"
	"github.com/Azure/acs-engine/pkg/armhelpers"
	"github.com/Azure/acs-engine/pkg/i18n"
	"github.com/Azure/acs-engine/pkg/operations"
	"github.com/leonelquinteros/gotext"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

type scaleCmd struct {
	authArgs

	// user input
	resourceGroupName    string
	deploymentDirectory  string
	newDesiredAgentCount int
	containerService     *api.ContainerService
	apiVersion           string
	location             string
	agentPoolToScale     string
	classicMode          bool

	// derived
	apiModelPath   string
	agentPool      *api.AgentPoolProfile
	client         armhelpers.ACSEngineClient
	locale         *gotext.Locale
	nameSuffix     string
	agentPoolIndex int
	masterFQDN     string
	logger         *log.Entry
}

const (
	scaleName             = "scale"
	scaleShortDescription = "scale a deployed cluster"
	scaleLongDescription  = "scale a deployed cluster"
)

// NewScaleCmd run a command to upgrade a Kubernetes cluster
func newScaleCmd() *cobra.Command {
	sc := scaleCmd{}

	scaleCmd := &cobra.Command{
		Use:   scaleName,
		Short: scaleShortDescription,
		Long:  scaleLongDescription,
		RunE: func(cmd *cobra.Command, args []string) error {
			return sc.run(cmd, args)
		},
	}

	f := scaleCmd.Flags()
	f.StringVar(&sc.location, "location", "", "location the cluster is deployed in")
	f.StringVar(&sc.resourceGroupName, "resource-group", "", "the resource group where the cluster is deployed")
	f.StringVar(&sc.deploymentDirectory, "deployment-dir", "", "the location of the output from `generate`")
	f.IntVar(&sc.newDesiredAgentCount, "new-node-count", 0, "desired number of nodes")
	f.BoolVar(&sc.classicMode, "classic-mode", false, "enable classic parameters and outputs")
	f.StringVar(&sc.agentPoolToScale, "node-pool", "", "node pool to scale")
	f.StringVar(&sc.masterFQDN, "master-FQDN", "", "FQDN for the master load balancer, Needed to scale down Kubernetes agent pools")

	addAuthFlags(&sc.authArgs, f)

	return scaleCmd
}

func (sc *scaleCmd) validate(cmd *cobra.Command, args []string) {
	log.Infoln("validating...")
	sc.logger = log.New().WithField("source", "scaling command line")
	var err error

	sc.locale, err = i18n.LoadTranslations()
	if err != nil {
		log.Fatalf("error loading translation files: %s", err.Error())
	}

	if sc.resourceGroupName == "" {
		cmd.Usage()
		log.Fatal("--resource-group must be specified")
	}

	if sc.location == "" {
		cmd.Usage()
		log.Fatal("--location must be specified")
	}

	if sc.newDesiredAgentCount == 0 {
		cmd.Usage()
		log.Fatal("--new-node-count must be specified")
	}

	if sc.client, err = sc.authArgs.getClient(); err != nil {
		log.Error("Failed to get client:", err)
	}

	if sc.deploymentDirectory == "" {
		cmd.Usage()
		log.Fatal("--deployment-dir must be specified")
	}

	_, err = sc.client.EnsureResourceGroup(sc.resourceGroupName, sc.location, nil)
	if err != nil {
		log.Fatalln(err)
	}

	// load apimodel from the deployment directory
	sc.apiModelPath = path.Join(sc.deploymentDirectory, "apimodel.json")

	if _, err = os.Stat(sc.apiModelPath); os.IsNotExist(err) {
		log.Fatalf("specified api model does not exist (%s)", sc.apiModelPath)
	}

	apiloader := &api.Apiloader{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}
	sc.containerService, sc.apiVersion, err = apiloader.LoadContainerServiceFromFile(sc.apiModelPath, true, true, nil)
	if err != nil {
		log.Fatalf("error parsing the api model: %s", err.Error())
	}

	if sc.agentPoolToScale == "" {
		agentPoolCount := len(sc.containerService.Properties.AgentPoolProfiles)
		if agentPoolCount > 1 {
			log.Fatal("--node-pool is required if more than one agent pool is defined in the container service")
		} else if agentPoolCount == 1 {
			sc.agentPool = sc.containerService.Properties.AgentPoolProfiles[0]
			sc.agentPoolIndex = 0
			sc.agentPoolToScale = sc.containerService.Properties.AgentPoolProfiles[0].Name
		} else {
			log.Fatal("No node pools found to scale")
		}
	} else {
		agentPoolIndex := -1
		for i, pool := range sc.containerService.Properties.AgentPoolProfiles {
			if pool.Name == sc.agentPoolToScale {
				agentPoolIndex = i
				sc.agentPool = pool
				sc.agentPoolIndex = i
			}
		}
		if agentPoolIndex == -1 {
			log.Fatalf("node pool %s wasn't in the deployed api model", sc.agentPoolToScale)
		}
	}

	templatePath := path.Join(sc.deploymentDirectory, "azuredeploy.json")
	contents, _ := ioutil.ReadFile(templatePath)

	var template interface{}
	json.Unmarshal(contents, &template)

	templateMap := template.(map[string]interface{})
	templateParameters := templateMap["parameters"].(map[string]interface{})

	nameSuffixParam := templateParameters["nameSuffix"].(map[string]interface{})
	sc.nameSuffix = nameSuffixParam["defaultValue"].(string)
	log.Infoln(fmt.Sprintf("Name suffix: %s", sc.nameSuffix))
}

func (sc *scaleCmd) run(cmd *cobra.Command, args []string) error {
	sc.validate(cmd, args)

	orchestratorInfo := sc.containerService.Properties.OrchestratorProfile
	var currentNodeCount, highestUsedIndex int
	indexes := make([]int, 0)
	indexToVM := make(map[int]string)
	if sc.agentPool.IsAvailabilitySets() {
		//TODO handle when there is a nextLink in the response and get more nodes
		vms, err := sc.client.ListVirtualMachines(sc.resourceGroupName)
		if err != nil {
			log.Fatalln("failed to get vms in the resource group. Error: %s", err.Error())
		}
		for _, vm := range *vms.Value {

			poolName, nameSuffix, index, err := armhelpers.K8sLinuxVMNameParts(*vm.Name)
			if err != nil || !strings.EqualFold(poolName, sc.agentPoolToScale) || !strings.EqualFold(nameSuffix, sc.nameSuffix) {
				continue
			}

			indexToVM[index] = *vm.Name
			indexes = append(indexes, index)
		}
		sortedIndexes := sort.IntSlice(indexes)
		sortedIndexes.Sort()
		indexes = []int(sortedIndexes)
		currentNodeCount = len(indexes)

		if currentNodeCount == sc.newDesiredAgentCount {
			log.Info("Cluster is currently at the desired agent count.")
			return nil
		}
		highestUsedIndex = indexes[len(indexes)-1]

		// Scale down Scenario
		if currentNodeCount > sc.newDesiredAgentCount {
			if sc.masterFQDN == "" {
				cmd.Usage()
				log.Fatal("master-FQDN is required to scale down a kubernetes cluster's agent pool")
			}

			vmsToDelete := make([]string, 0)
			for i := currentNodeCount - 1; i >= sc.newDesiredAgentCount; i-- {
				vmsToDelete = append(vmsToDelete, indexToVM[i])
			}

			if orchestratorInfo.OrchestratorType == api.Kubernetes {
				err = sc.drainNodes(vmsToDelete)
				if err != nil {
					log.Errorf("Got error %+v, while draining the nodes to be deleted", err)
					return err
				}
			}

			errList := operations.ScaleDownVMs(sc.client, sc.logger, sc.resourceGroupName, vmsToDelete...)
			if errList != nil {
				errorMessage := ""
				for element := errList.Front(); element != nil; element = element.Next() {
					vmError, ok := element.Value.(*operations.VMScalingErrorDetails)
					if ok {
						error := fmt.Sprintf("Node '%s' failed to delete with error: '%s'", vmError.Name, vmError.Error.Error())
						errorMessage = errorMessage + error
					}
				}
				return fmt.Errorf(errorMessage)
			}

			return nil
		}
	} else {
		vmssList, err := sc.client.ListVirtualMachineScaleSets(sc.resourceGroupName)
		if err != nil {
			log.Fatalln("failed to get vmss list in the resource group. Error: %s", err.Error())
		}
		for _, vmss := range *vmssList.Value {
			poolName, nameSuffix, err := armhelpers.VmssNameParts(*vmss.Name)
			if err != nil || !strings.EqualFold(poolName, sc.agentPoolToScale) || !strings.EqualFold(nameSuffix, sc.nameSuffix) {
				continue
			}
			currentNodeCount = int(*vmss.Sku.Capacity)
			highestUsedIndex = 0
		}
	}

	ctx := acsengine.Context{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}
	templateGenerator, err := acsengine.InitializeTemplateGenerator(ctx, sc.classicMode)
	if err != nil {
		log.Fatalln("failed to initialize template generator: %s", err.Error())
	}

	sc.containerService.Properties.AgentPoolProfiles = []*api.AgentPoolProfile{sc.agentPool}

	template, parameters, _, err := templateGenerator.GenerateTemplate(sc.containerService, acsengine.DefaultGeneratorCode)
	if err != nil {
		log.Fatalf("error generating template %s: %s", sc.apiModelPath, err.Error())
		os.Exit(1)
	}

	if template, err = acsengine.PrettyPrintArmTemplate(template); err != nil {
		log.Fatalf("error pretty printing template: %s \n", err.Error())
	}

	templateJSON := make(map[string]interface{})
	parametersJSON := make(map[string]interface{})

	err = json.Unmarshal([]byte(template), &templateJSON)
	if err != nil {
		log.Fatalln(err)
	}

	err = json.Unmarshal([]byte(parameters), &parametersJSON)
	if err != nil {
		log.Fatalln(err)
	}

	transformer := acsengine.Transformer{Translator: ctx.Translator}
	// Our templates generate a range of nodes based on a count and offset, it is possible for there to be holes in the template
	// So we need to set the count in the template to get enough nodes for the range, if there are holes that number will be larger than the desired count
	countForTemplate := sc.newDesiredAgentCount
	if highestUsedIndex != 0 {
		countForTemplate += highestUsedIndex + 1 - currentNodeCount
	}
	addValue(parametersJSON, sc.agentPool.Name+"Count", countForTemplate)

	switch orchestratorInfo.OrchestratorType {
	case api.Kubernetes:
		err = transformer.NormalizeForK8sVMASScalingUp(sc.logger, templateJSON)
		if err != nil {
			log.Fatalf("error tranforming the template for scaling template %s: %s", sc.apiModelPath, err.Error())
			os.Exit(1)
		}
		if sc.agentPool.IsAvailabilitySets() {
			addValue(parametersJSON, fmt.Sprintf("%sOffset", sc.agentPool.Name), highestUsedIndex+1)
		}
		break
	case api.Swarm:
	case api.SwarmMode:
	case api.DCOS:
		if sc.agentPool.IsAvailabilitySets() {
			log.Fatalf("scaling isn't supported for orchestrator %s, with availability sets", orchestratorInfo.OrchestratorType)
			os.Exit(1)
		}
		transformer.NormalizeForVMSSScaling(sc.logger, templateJSON)
	}

	random := rand.New(rand.NewSource(time.Now().UnixNano()))
	deploymentSuffix := random.Int31()

	_, err = sc.client.DeployTemplate(
		sc.resourceGroupName,
		fmt.Sprintf("%s-%d", sc.resourceGroupName, deploymentSuffix),
		templateJSON,
		parametersJSON,
		nil)
	if err != nil {
		log.Fatalln(err)
	}

	apiloader := &api.Apiloader{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}
	var apiVersion string
	sc.containerService, apiVersion, err = apiloader.LoadContainerServiceFromFile(sc.apiModelPath, false, true, nil)
	sc.containerService.Properties.AgentPoolProfiles[sc.agentPoolIndex].Count = sc.newDesiredAgentCount

	b, e := apiloader.SerializeContainerService(sc.containerService, apiVersion)

	if e != nil {
		return e
	}

	f := acsengine.FileSaver{
		Translator: &i18n.Translator{
			Locale: sc.locale,
		},
	}

	if e = f.SaveFile(sc.deploymentDirectory, "apimodel.json", b); e != nil {
		return e
	}

	return nil
}

type paramsMap map[string]interface{}

func addValue(m paramsMap, k string, v interface{}) {
	m[k] = paramsMap{
		"value": v,
	}
}

func (sc *scaleCmd) drainNodes(vmsToDelete []string) error {
	kubeConfig, err := acsengine.GenerateKubeConfig(sc.containerService.Properties, sc.location)
	if err != nil {
		log.Fatalf("failed to generate kube config") // TODO: cleanup
	}
	var errorMessage string
	masterURL := sc.masterFQDN
	if !strings.HasPrefix(masterURL, "https://") {
		masterURL = fmt.Sprintf("https://%s", masterURL)
	}
	numVmsToDrain := len(vmsToDelete)
	errChan := make(chan *operations.VMScalingErrorDetails, numVmsToDrain)
	defer close(errChan)
	for _, vmName := range vmsToDelete {
		go func(vmName string) {
			e := operations.SafelyDrainNode(sc.client, sc.logger,
				masterURL, kubeConfig, vmName, time.Duration(60)*time.Minute)
			if e != nil {
				log.Errorf("Failed to drain node %s, got error %s", vmName, e.Error())
				errChan <- &operations.VMScalingErrorDetails{Error: e, Name: vmName}
				return
			}
			errChan <- nil
		}(vmName)
	}

	for i := 0; i < numVmsToDrain; i++ {
		errDetails := <-errChan
		if errDetails != nil {
			error := fmt.Sprintf("Node '%s' failed to drain with error: '%s'", errDetails.Name, errDetails.Error.Error())
			errorMessage = errorMessage + error
			return fmt.Errorf(error)
		}
	}

	return nil
}
<|endoftext|>


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

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/spf13/cobra"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

const SPEC_ILLEGAL = -1
const MISSING_FILE_REFERENCE = -2
const MISSING_CONFIGURATION = -3
const ILLEGAL_JSON_CONFIGURATION = -4
const CONFIGURATION_FILE_IN_ROOT = -5
const ERROR_READING_FILE = -6
const BOOBER_ERROR = -7
const ILLEGAL_FILE = -8

const SPEC_IS_FILE = 1
const SPEC_IS_FOLDER = 2

// Cobra Flag variables
var overrideFiles []string
var overrideValues []string
var dryRun bool
var showConfig bool

// Struct to represent data to the Boober interface
type BooberInferface struct {
	Env         string                     `json:"env"`
	App         string                     `json:"app"`
	Affiliation string                     `json:"affiliation"`
	Files       map[string]json.RawMessage `json:"files"`
	Overrides   map[string]json.RawMessage `json:"overrides"`
}

// Struct to represent return data from the Boober interface
type BooberReturn struct {
	Sources json.RawMessage `json:"sources"`
	Errors  []string        `json:"errors"`
	Valid   bool            `json:"valid"`
}

// setupCmd represents the setup command
var setupCmd = &cobra.Command{
	Use:   `setup folder | file [-f file 'JSON Configuration String]'`,
	Short: "Deploys an application to OpenShift based upon local configuration files",
	Long:  `Usage: When used with a `,
	Run: func(cmd *cobra.Command, args []string) {
		executeSetup(args)
	},
}

func executeSetup(args []string) {
	validateCode := validateCommand(args)
	if validateCode < 0 {
		os.Exit(validateCode)
	}

	var absolutePath string

	absolutePath, _ = filepath.Abs(args[0])

	var envFile string      // Filename for app
	var envFolder string    // Short folder name (Env)
	var folder string       // Absolute path of folder
	var parentFolder string // Absolute path of parent

	switch validateCode {
	case SPEC_IS_FILE:
		folder = filepath.Dir(absolutePath)
		envFile = filepath.Base(absolutePath)
	case SPEC_IS_FOLDER:
		folder = absolutePath
		envFile = ""
	}
	parentFolder = filepath.Dir(folder)
	envFolder = filepath.Base(folder)

	if folder == parentFolder {
		fmt.Println("Application configuration file cannot reside in root directory")
		os.Exit(CONFIGURATION_FILE_IN_ROOT)
	}

	// Initialize JSON
	var jsonStr string

	var booberData BooberInferface
	booberData.App = strings.TrimSuffix(envFile, filepath.Ext(envFile)) //envFile
	booberData.Env = envFolder
	booberData.Affiliation = ""

	var returnMap = folder2map(folder, envFolder+"/")
	var returnMap2 = folder2map(parentFolder, "")

	booberData.Files = combineMaps(returnMap, returnMap2)
	booberData.Overrides = overrides2map(args)

	jsonByte, ok := json.Marshal(booberData)
	if !(ok == nil) {
		fmt.Println("Internal error in marshalling Boober data: " + ok.Error())
	}
	jsonStr = string(jsonByte)

	if dryRun {
		fmt.Println(string(prettyPrintJson(jsonStr)))
	} else {
		callBoober(jsonStr)
	}

}

func overrides2map(args []string) map[string]json.RawMessage {
	var returnMap = make(map[string]json.RawMessage)
	for i := 0; i < len(overrideFiles); i++ {
		returnMap[overrideFiles[i]] = json.RawMessage(args[i+1])
	}
	return returnMap
}

func folder2map(folder string, prefix string) map[string]json.RawMessage {
	var returnMap = make(map[string]json.RawMessage)
	var allFilesOK bool = true

	files, _ := ioutil.ReadDir(folder)
	var filesProcessed = 0
	for _, f := range files {
		absolutePath := filepath.Join(folder, f.Name())
		if isLegalFileFolder(absolutePath) == SPEC_IS_FILE { // Ignore folders
			matched, _ := filepath.Match("*.json", strings.ToLower(f.Name()))
			if matched {
				fileJson, err := ioutil.ReadFile(absolutePath)
				if err != nil {
					fmt.Println("Error in reading file " + absolutePath)
					os.Exit(ERROR_READING_FILE)
				}
				if isLegalJson(string(fileJson)) {
					filesProcessed++
					returnMap[prefix+f.Name()] = fileJson
				} else {
					fmt.Println("Illegal JSON in configuration file " + absolutePath)
					allFilesOK = false
				}
				filesProcessed++
			}
		}

	}
	if !allFilesOK {
		os.Exit(ILLEGAL_JSON_CONFIGURATION)
	}
	return returnMap
}

func combineMaps(map1 map[string]json.RawMessage, map2 map[string]json.RawMessage) map[string]json.RawMessage {
	var returnMap = make(map[string]json.RawMessage)

	for k, v := range map1 {
		returnMap[k] = v
	}
	for k, v := range map2 {
		returnMap[k] = v
	}
	return returnMap
}

func validateCommand(args []string) int {
	var errorString = ""
	var returnCode int

	if len(args) == 0 {
		returnCode = -1
		errorString += "Missing file/folder "
	} else {
		// Chceck argument 0 for legal file / folder
		returnCode = isLegalFileFolder(args[0])
		if returnCode < 0 {
			errorString += "Illegal file / folder: " + args[0]
			returnCode = ILLEGAL_FILE
		}

		// We have at least one argument, now there should be a correlation between the number of args
		// and the number of override (-f) flags
		if len(overrideFiles) < (len(args) - 1) {
			returnCode = MISSING_FILE_REFERENCE
			errorString += "Configuration override specified without file reference flag "
		}
		if len(overrideFiles) > (len(args) - 1) {
			returnCode = MISSING_CONFIGURATION
			errorString += "Configuration overide file reference flag specified without configuration "
		}

		// Check for legal JSON argument for each overrideFiles flag
		for i := 1; i < len(args); i++ {
			if !isLegalJson(args[i]) {
				errorString = "Illegal JSON configuration override: " + args[i] + " "
				returnCode = ILLEGAL_JSON_CONFIGURATION
			}
		}
	}

	if returnCode < 0 {
		fmt.Println(errorString)
	}
	return returnCode

}

func isLegalFileFolder(filespec string) int {
	var err error
	var absolutePath string
	var fi os.FileInfo

	absolutePath, err = filepath.Abs(filespec)
	fi, err = os.Stat(absolutePath)
	if os.IsNotExist(err) {
		return SPEC_ILLEGAL
	} else {
		switch mode := fi.Mode(); {
		case mode.IsDir():
			return SPEC_IS_FOLDER
		case mode.IsRegular():
			return SPEC_IS_FILE
		}
	}
	return SPEC_ILLEGAL
}

func isLegalJson(jsonString string) bool {
	var js map[string]interface{}
	return json.Unmarshal([]byte(jsonString), &js) == nil
}

func prettyPrintJson(jsonString string) string {
	var out bytes.Buffer
	err := json.Indent(&out, []byte(jsonString), "", "\t")
	if err != nil {
		return jsonString
	}
	return out.String()
}

func callBoober(combindedJson string) {
	//url := "http://localhost:8080/api/setupMock/env/app"
	url := "http://localhost:8080/setup"
	var jsonStr = []byte(combindedJson)

	req, err := http.NewRequest(http.MethodPut, url, bytes.NewBuffer(jsonStr))
	req.Header.Set("Content-Type", "application/json")
	if err != nil {
		log.Fatal("NewRequest: ", err)
		return
	}

	req.Header.Set("Authentication", "mydirtysecret")
	client := &http.Client{}

	resp, err := client.Do(req)
	if err != nil {
		fmt.Println("Error connecting to the Boober service on "+url+": ", err)
		os.Exit(BOOBER_ERROR)
	}

	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		fmt.Println("Error from the Boober service on " + url + ":")
		body, _ := ioutil.ReadAll(resp.Body)
		fmt.Println(string(body))
		os.Exit(BOOBER_ERROR)
	}
	body, _ := ioutil.ReadAll(resp.Body)

	// Check return for error
	var booberReturn BooberReturn
	err = json.Unmarshal(body, &booberReturn)
	if err != nil {
		fmt.Println("Error unmashalling Boober return: " + err.Error())
		os.Exit(BOOBER_ERROR)
	}

	if !(booberReturn.Valid) {
		fmt.Println("Error in configuration: ")
		for _, message := range booberReturn.Errors {
			fmt.Println("  " + message)
		}
	}

	if showConfig {
		fmt.Println(prettyPrintJson(string(booberReturn.Sources)))
	}

}

func init() {
	RootCmd.AddCommand(setupCmd)

	// File flag, supports multiple instances of the flag
	setupCmd.Flags().StringArrayVarP(&overrideFiles, "file",
		"f", overrideValues, "File to override")
	setupCmd.Flags().BoolVarP(&dryRun, "dryrun",
		"d", false,
		"Do not perform a setup, just collect and print the configuration files")
	setupCmd.Flags().BoolVarP(&showConfig, "showconfig",
		"s", false, "Send config to standard out")
}
<commit_msg>The -s flag prints the Config part of the return from Boober.<commit_after>// Copyright © 2017 Norwegian Tax Authority
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"bytes"
	"encoding/json"
	"fmt"
	"github.com/spf13/cobra"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

const SPEC_ILLEGAL = -1
const MISSING_FILE_REFERENCE = -2
const MISSING_CONFIGURATION = -3
const ILLEGAL_JSON_CONFIGURATION = -4
const CONFIGURATION_FILE_IN_ROOT = -5
const ERROR_READING_FILE = -6
const BOOBER_ERROR = -7
const ILLEGAL_FILE = -8

const SPEC_IS_FILE = 1
const SPEC_IS_FOLDER = 2

// Cobra Flag variables
var overrideFiles []string
var overrideValues []string
var dryRun bool
var showConfig bool

// Struct to represent data to the Boober interface
type BooberInferface struct {
	Env         string                     `json:"env"`
	App         string                     `json:"app"`
	Affiliation string                     `json:"affiliation"`
	Files       map[string]json.RawMessage `json:"files"`
	Overrides   map[string]json.RawMessage `json:"overrides"`
}

// Struct to represent return data from the Boober interface
type BooberReturn struct {
	Sources json.RawMessage `json:"sources"`
	Errors  []string        `json:"errors"`
	Valid   bool            `json:"valid"`
	Config  json.RawMessage `json:"config"`
}

// setupCmd represents the setup command
var setupCmd = &cobra.Command{
	Use:   `setup folder | file [-f file 'JSON Configuration String]'`,
	Short: "Deploys an application to OpenShift based upon local configuration files",
	Long:  `Usage: When used with a `,
	Run: func(cmd *cobra.Command, args []string) {
		executeSetup(args)
	},
}

func executeSetup(args []string) {
	validateCode := validateCommand(args)
	if validateCode < 0 {
		os.Exit(validateCode)
	}

	var absolutePath string

	absolutePath, _ = filepath.Abs(args[0])

	var envFile string      // Filename for app
	var envFolder string    // Short folder name (Env)
	var folder string       // Absolute path of folder
	var parentFolder string // Absolute path of parent

	switch validateCode {
	case SPEC_IS_FILE:
		folder = filepath.Dir(absolutePath)
		envFile = filepath.Base(absolutePath)
	case SPEC_IS_FOLDER:
		folder = absolutePath
		envFile = ""
	}
	parentFolder = filepath.Dir(folder)
	envFolder = filepath.Base(folder)

	if folder == parentFolder {
		fmt.Println("Application configuration file cannot reside in root directory")
		os.Exit(CONFIGURATION_FILE_IN_ROOT)
	}

	// Initialize JSON
	var jsonStr string

	var booberData BooberInferface
	booberData.App = strings.TrimSuffix(envFile, filepath.Ext(envFile)) //envFile
	booberData.Env = envFolder
	booberData.Affiliation = ""

	var returnMap = folder2map(folder, envFolder+"/")
	var returnMap2 = folder2map(parentFolder, "")

	booberData.Files = combineMaps(returnMap, returnMap2)
	booberData.Overrides = overrides2map(args)

	jsonByte, ok := json.Marshal(booberData)
	if !(ok == nil) {
		fmt.Println("Internal error in marshalling Boober data: " + ok.Error())
	}
	jsonStr = string(jsonByte)

	if dryRun {
		fmt.Println(string(prettyPrintJson(jsonStr)))
	} else {
		callBoober(jsonStr)
	}

}

func overrides2map(args []string) map[string]json.RawMessage {
	var returnMap = make(map[string]json.RawMessage)
	for i := 0; i < len(overrideFiles); i++ {
		returnMap[overrideFiles[i]] = json.RawMessage(args[i+1])
	}
	return returnMap
}

func folder2map(folder string, prefix string) map[string]json.RawMessage {
	var returnMap = make(map[string]json.RawMessage)
	var allFilesOK bool = true

	files, _ := ioutil.ReadDir(folder)
	var filesProcessed = 0
	for _, f := range files {
		absolutePath := filepath.Join(folder, f.Name())
		if isLegalFileFolder(absolutePath) == SPEC_IS_FILE { // Ignore folders
			matched, _ := filepath.Match("*.json", strings.ToLower(f.Name()))
			if matched {
				fileJson, err := ioutil.ReadFile(absolutePath)
				if err != nil {
					fmt.Println("Error in reading file " + absolutePath)
					os.Exit(ERROR_READING_FILE)
				}
				if isLegalJson(string(fileJson)) {
					filesProcessed++
					returnMap[prefix+f.Name()] = fileJson
				} else {
					fmt.Println("Illegal JSON in configuration file " + absolutePath)
					allFilesOK = false
				}
				filesProcessed++
			}
		}

	}
	if !allFilesOK {
		os.Exit(ILLEGAL_JSON_CONFIGURATION)
	}
	return returnMap
}

func combineMaps(map1 map[string]json.RawMessage, map2 map[string]json.RawMessage) map[string]json.RawMessage {
	var returnMap = make(map[string]json.RawMessage)

	for k, v := range map1 {
		returnMap[k] = v
	}
	for k, v := range map2 {
		returnMap[k] = v
	}
	return returnMap
}

func validateCommand(args []string) int {
	var errorString = ""
	var returnCode int

	if len(args) == 0 {
		returnCode = -1
		errorString += "Missing file/folder "
	} else {
		// Chceck argument 0 for legal file / folder
		returnCode = isLegalFileFolder(args[0])
		if returnCode < 0 {
			errorString += "Illegal file / folder: " + args[0]
			returnCode = ILLEGAL_FILE
		}

		// We have at least one argument, now there should be a correlation between the number of args
		// and the number of override (-f) flags
		if len(overrideFiles) < (len(args) - 1) {
			returnCode = MISSING_FILE_REFERENCE
			errorString += "Configuration override specified without file reference flag "
		}
		if len(overrideFiles) > (len(args) - 1) {
			returnCode = MISSING_CONFIGURATION
			errorString += "Configuration overide file reference flag specified without configuration "
		}

		// Check for legal JSON argument for each overrideFiles flag
		for i := 1; i < len(args); i++ {
			if !isLegalJson(args[i]) {
				errorString = "Illegal JSON configuration override: " + args[i] + " "
				returnCode = ILLEGAL_JSON_CONFIGURATION
			}
		}
	}

	if returnCode < 0 {
		fmt.Println(errorString)
	}
	return returnCode

}

func isLegalFileFolder(filespec string) int {
	var err error
	var absolutePath string
	var fi os.FileInfo

	absolutePath, err = filepath.Abs(filespec)
	fi, err = os.Stat(absolutePath)
	if os.IsNotExist(err) {
		return SPEC_ILLEGAL
	} else {
		switch mode := fi.Mode(); {
		case mode.IsDir():
			return SPEC_IS_FOLDER
		case mode.IsRegular():
			return SPEC_IS_FILE
		}
	}
	return SPEC_ILLEGAL
}

func isLegalJson(jsonString string) bool {
	var js map[string]interface{}
	return json.Unmarshal([]byte(jsonString), &js) == nil
}

func prettyPrintJson(jsonString string) string {
	var out bytes.Buffer
	err := json.Indent(&out, []byte(jsonString), "", "\t")
	if err != nil {
		return jsonString
	}
	return out.String()
}

func callBoober(combindedJson string) {
	//url := "http://localhost:8080/api/setupMock/env/app"
	url := "http://localhost:8080/setup"
	var jsonStr = []byte(combindedJson)

	req, err := http.NewRequest(http.MethodPut, url, bytes.NewBuffer(jsonStr))
	req.Header.Set("Content-Type", "application/json")
	if err != nil {
		log.Fatal("NewRequest: ", err)
		return
	}

	req.Header.Set("Authentication", "mydirtysecret")
	client := &http.Client{}

	resp, err := client.Do(req)
	if err != nil {
		fmt.Println("Error connecting to the Boober service on "+url+": ", err)
		os.Exit(BOOBER_ERROR)
	}

	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		fmt.Println("Error from the Boober service on " + url + ":")
		body, _ := ioutil.ReadAll(resp.Body)
		fmt.Println(string(body))
		os.Exit(BOOBER_ERROR)
	}
	body, _ := ioutil.ReadAll(resp.Body)

	// Check return for error
	var booberReturn BooberReturn
	err = json.Unmarshal(body, &booberReturn)
	if err != nil {
		fmt.Println("Error unmashalling Boober return: " + err.Error())
		os.Exit(BOOBER_ERROR)
	}

	if !(booberReturn.Valid) {
		fmt.Println("Error in configuration: ")
		for _, message := range booberReturn.Errors {
			fmt.Println("  " + message)
		}
	}

	if showConfig {
		fmt.Println(prettyPrintJson(string(booberReturn.Config)))
	}

}

func init() {
	RootCmd.AddCommand(setupCmd)

	// File flag, supports multiple instances of the flag
	setupCmd.Flags().StringArrayVarP(&overrideFiles, "file",
		"f", overrideValues, "File to override")
	setupCmd.Flags().BoolVarP(&dryRun, "dryrun",
		"d", false,
		"Do not perform a setup, just collect and print the configuration files")
	setupCmd.Flags().BoolVarP(&showConfig, "showconfig",
		"s", false, "Send config to standard out")
}
<|endoftext|>


<commit_before>package cmd

import (
	"os"

	"github.com/gsamokovarov/jump/cli"
	"github.com/gsamokovarov/jump/config"
	"github.com/gsamokovarov/jump/shell"
)

func shellCmd(args cli.Args, _ *config.Config) {
	hint := args.CommandName()
	if len(hint) == 0 {
		hint = os.Getenv("SHELL")
	}

	sh := shell.Guess(hint)

	cli.Outf("%s", sh.MustCompile("j"))
}

func init() {
	cli.RegisterCommand("shell", "Display a shell integration script.", shellCmd)
}
<commit_msg>Make the shortcut configurable so @luhova can replace z<commit_after>package cmd

import (
	"os"

	"github.com/gsamokovarov/jump/cli"
	"github.com/gsamokovarov/jump/config"
	"github.com/gsamokovarov/jump/shell"
)

func shellCmd(args cli.Args, _ *config.Config) {
	hint := args.CommandName()
	if len(hint) == 0 {
		hint = os.Getenv("SHELL")
	}

	sh := shell.Guess(hint)
	shortcut := args.Get("--shortcut", "j")

	cli.Outf("%s", sh.MustCompile(shortcut))
}

func init() {
	cli.RegisterCommand("shell", "Display a shell integration script.", shellCmd)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"code.google.com/p/gographviz"
	"fmt"
	"github.com/gonuts/commander"
	"github.com/gonuts/flag"
	"github.com/smira/aptly/debian"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

func graphvizEscape(s string) string {
	return fmt.Sprintf("\"%s\"", strings.Replace(s, "\"", "\\\"", 0))
}

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

	graph := gographviz.NewGraph()
	graph.SetDir(true)
	graph.SetName("aptly")

	fmt.Printf("Loading mirrors...\n")

	repoCollection := debian.NewRemoteRepoCollection(context.database)

	err = repoCollection.ForEach(func(repo *debian.RemoteRepo) error {
		err := repoCollection.LoadComplete(repo)
		if err != nil {
			return err
		}

		graph.AddNode("aptly", graphvizEscape(repo.UUID), map[string]string{
			"shape":     "Mrecord",
			"style":     "filled",
			"fillcolor": "darkgoldenrod1",
			"label": graphvizEscape(fmt.Sprintf("{Mirror %s|url: %s|dist: %s|comp: %s|arch: %s|pkgs: %d}",
				repo.Name, repo.ArchiveRoot, repo.Distribution, strings.Join(repo.Components, ", "),
				strings.Join(repo.Architectures, ", "), repo.NumPackages())),
		})
		return nil
	})

	if err != nil {
		return err
	}

	fmt.Printf("Loading snapshots...\n")

	snapshotCollection := debian.NewSnapshotCollection(context.database)

	err = snapshotCollection.ForEach(func(snapshot *debian.Snapshot) error {
		err := snapshotCollection.LoadComplete(snapshot)
		if err != nil {
			return err
		}

		description := snapshot.Description
		if snapshot.SourceKind == "repo" {
			description = "Snapshot from repo"
		}

		graph.AddNode("aptly", graphvizEscape(snapshot.UUID), map[string]string{
			"shape":     "Mrecord",
			"style":     "filled",
			"fillcolor": "cadetblue1",
			"label":     graphvizEscape(fmt.Sprintf("{Snapshot %s|%s|pkgs: %d}", snapshot.Name, description, snapshot.NumPackages())),
		})

		if snapshot.SourceKind == "repo" {
			for _, uuid := range snapshot.SourceIDs {
				graph.AddEdge(graphvizEscape(uuid), "", graphvizEscape(snapshot.UUID), "", true, nil)
			}
		} else if snapshot.SourceKind == "snapshot" {
			for _, uuid := range snapshot.SourceIDs {
				graph.AddEdge(graphvizEscape(uuid), "", graphvizEscape(snapshot.UUID), "", true, nil)
			}
		}
		return nil
	})

	if err != nil {
		return err
	}

	fmt.Printf("Loading published repos...\n")

	publishedCollection := debian.NewPublishedRepoCollection(context.database)

	publishedCollection.ForEach(func(repo *debian.PublishedRepo) error {
		graph.AddNode("aptly", graphvizEscape(repo.UUID), map[string]string{
			"shape":     "Mrecord",
			"style":     "filled",
			"fillcolor": "darkolivegreen1",
			"label":     graphvizEscape(fmt.Sprintf("{Published %s/%s|comp: %s|arch: %s}", repo.Prefix, repo.Distribution, repo.Component, strings.Join(repo.Architectures, ", "))),
		})

		graph.AddEdge(graphvizEscape(repo.SnapshotUUID), "", graphvizEscape(repo.UUID), "", true, nil)

		return nil
	})

	fmt.Printf("Generating graph...\n")

	buf := bytes.NewBufferString(graph.String())

	tempfile, err := ioutil.TempFile("", "aptly-graph")
	if err != nil {
		return err
	}
	tempfile.Close()

	command := exec.Command("dot", "-Tpng", "-o"+tempfile.Name())
	command.Stderr = os.Stderr

	stdin, err := command.StdinPipe()
	if err != nil {
		return err
	}

	err = command.Start()
	if err != nil {
		return fmt.Errorf("unable to execute dot: %s (is graphviz package installed?)", err)
	}

	_, err = io.Copy(stdin, buf)
	if err != nil {
		return err
	}

	err = stdin.Close()
	if err != nil {
		return err
	}

	err = command.Wait()
	if err != nil {
		return err
	}

	err = exec.Command("open", tempfile.Name()).Run()
	if err != nil {
		fmt.Printf("Rendered to PNG file: %s\n", tempfile.Name())
		err = nil
	}

	return err
}

func makeCmdGraph() *commander.Command {
	cmd := &commander.Command{
		Run:       aptlyGraph,
		UsageLine: "graph",
		Short:     "display graph of dependencies between aptly objects (requires graphviz)",
		Long: `
Command graph displays relationship between mirrors, snapshots and published repositories using
graphviz package to render graph as image.

ex:
  $ aptly graph
`,
		Flag: *flag.NewFlagSet("aptly-graph", flag.ExitOnError),
	}

	return cmd
}
<commit_msg>Use temp file with .png extension.<commit_after>package main

import (
	"bytes"
	"code.google.com/p/gographviz"
	"fmt"
	"github.com/gonuts/commander"
	"github.com/gonuts/flag"
	"github.com/smira/aptly/debian"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

func graphvizEscape(s string) string {
	return fmt.Sprintf("\"%s\"", strings.Replace(s, "\"", "\\\"", 0))
}

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

	graph := gographviz.NewGraph()
	graph.SetDir(true)
	graph.SetName("aptly")

	fmt.Printf("Loading mirrors...\n")

	repoCollection := debian.NewRemoteRepoCollection(context.database)

	err = repoCollection.ForEach(func(repo *debian.RemoteRepo) error {
		err := repoCollection.LoadComplete(repo)
		if err != nil {
			return err
		}

		graph.AddNode("aptly", graphvizEscape(repo.UUID), map[string]string{
			"shape":     "Mrecord",
			"style":     "filled",
			"fillcolor": "darkgoldenrod1",
			"label": graphvizEscape(fmt.Sprintf("{Mirror %s|url: %s|dist: %s|comp: %s|arch: %s|pkgs: %d}",
				repo.Name, repo.ArchiveRoot, repo.Distribution, strings.Join(repo.Components, ", "),
				strings.Join(repo.Architectures, ", "), repo.NumPackages())),
		})
		return nil
	})

	if err != nil {
		return err
	}

	fmt.Printf("Loading snapshots...\n")

	snapshotCollection := debian.NewSnapshotCollection(context.database)

	err = snapshotCollection.ForEach(func(snapshot *debian.Snapshot) error {
		err := snapshotCollection.LoadComplete(snapshot)
		if err != nil {
			return err
		}

		description := snapshot.Description
		if snapshot.SourceKind == "repo" {
			description = "Snapshot from repo"
		}

		graph.AddNode("aptly", graphvizEscape(snapshot.UUID), map[string]string{
			"shape":     "Mrecord",
			"style":     "filled",
			"fillcolor": "cadetblue1",
			"label":     graphvizEscape(fmt.Sprintf("{Snapshot %s|%s|pkgs: %d}", snapshot.Name, description, snapshot.NumPackages())),
		})

		if snapshot.SourceKind == "repo" {
			for _, uuid := range snapshot.SourceIDs {
				graph.AddEdge(graphvizEscape(uuid), "", graphvizEscape(snapshot.UUID), "", true, nil)
			}
		} else if snapshot.SourceKind == "snapshot" {
			for _, uuid := range snapshot.SourceIDs {
				graph.AddEdge(graphvizEscape(uuid), "", graphvizEscape(snapshot.UUID), "", true, nil)
			}
		}
		return nil
	})

	if err != nil {
		return err
	}

	fmt.Printf("Loading published repos...\n")

	publishedCollection := debian.NewPublishedRepoCollection(context.database)

	publishedCollection.ForEach(func(repo *debian.PublishedRepo) error {
		graph.AddNode("aptly", graphvizEscape(repo.UUID), map[string]string{
			"shape":     "Mrecord",
			"style":     "filled",
			"fillcolor": "darkolivegreen1",
			"label":     graphvizEscape(fmt.Sprintf("{Published %s/%s|comp: %s|arch: %s}", repo.Prefix, repo.Distribution, repo.Component, strings.Join(repo.Architectures, ", "))),
		})

		graph.AddEdge(graphvizEscape(repo.SnapshotUUID), "", graphvizEscape(repo.UUID), "", true, nil)

		return nil
	})

	fmt.Printf("Generating graph...\n")

	buf := bytes.NewBufferString(graph.String())

	tempfile, err := ioutil.TempFile("", "aptly-graph")
	if err != nil {
		return err
	}
	tempfile.Close()
	os.Remove(tempfile.Name())

	tempfilename := tempfile.Name() + ".png"

	command := exec.Command("dot", "-Tpng", "-o"+tempfilename)
	command.Stderr = os.Stderr

	stdin, err := command.StdinPipe()
	if err != nil {
		return err
	}

	err = command.Start()
	if err != nil {
		return fmt.Errorf("unable to execute dot: %s (is graphviz package installed?)", err)
	}

	_, err = io.Copy(stdin, buf)
	if err != nil {
		return err
	}

	err = stdin.Close()
	if err != nil {
		return err
	}

	err = command.Wait()
	if err != nil {
		return err
	}

	err = exec.Command("open", tempfilename).Run()
	if err != nil {
		fmt.Printf("Rendered to PNG file: %s\n", tempfilename)
		err = nil
	}

	return err
}

func makeCmdGraph() *commander.Command {
	cmd := &commander.Command{
		Run:       aptlyGraph,
		UsageLine: "graph",
		Short:     "display graph of dependencies between aptly objects (requires graphviz)",
		Long: `
Command graph displays relationship between mirrors, snapshots and published repositories using
graphviz package to render graph as image.

ex:
  $ aptly graph
`,
		Flag: *flag.NewFlagSet("aptly-graph", flag.ExitOnError),
	}

	return cmd
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015 Matthew Collins
 *
 * 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 thinkbot

import (
	"fmt"
	"github.com/thinkofdeath/thinkbot/command"
	"github.com/thinkofdeath/thinkbot/irc"
	"log"
	"regexp"
	"strings"
	"sync"
)

// Bot contains all the information of the currently
// running bot.
//
// Event should be continuously read from and the
// returned events handled.
type Bot struct {
	client   *irc.Client
	err      error
	username string

	Commands command.Registry

	chatHandlers []chatHandler
	chatLock     sync.RWMutex

	Events    chan Event
	writeChan chan irc.Message
	funcChan  chan func()

	channels      []string
	commandPrefix string
	modes         map[rune]struct{}

	permissionContainer PermissionContainer
}

type chatHandler struct {
	r *regexp.Regexp
	f func(b *Bot, sender User, target, message string) error
}

// NewBot creates an instance of a bot connecting to the target
// server with the specified username
//
// The init function is called before any messages/commands
// are handled to allow for setup
//
// Returns an error if it fails to connect
func NewBot(server string, port uint16, username string, init func(*Bot)) (*Bot, error) {
	c, err := irc.NewClient(server, port)
	if err != nil {
		return nil, err
	}
	b := &Bot{
		client:        c,
		Events:        make(chan Event, 100),
		writeChan:     make(chan irc.Message, 100),
		funcChan:      make(chan func(), 100),
		username:      username,
		channels:      []string{},
		commandPrefix: "+",
		modes:         map[rune]struct{}{},
		Commands: command.Registry{
			// User and target parameters
			ExtraParameters: 2,
		},
	}
	init(b)
	go b.run()
	return b, nil
}

func (b *Bot) run() {
	c := b.client
	rc := make(chan irc.Message)

	go func() {
		for {
			m, err := c.Read()
			if err != nil {
				b.err = err
				b.kill()
				close(rc)
				return
			}
			rc <- m
		}
	}()

	c.Write(irc.NewNick(b.username))
	c.Write(irc.NewUser(b.username, b.username))

	for {
		select {
		case f := <-b.funcChan:
			f()
		case m := <-b.writeChan:
			c.Write(m)
		case m, ok := <-rc:
			if !ok {
				return
			}
			switch m := m.(type) {
			case irc.Reply:
				b.handleReply(m)
			case irc.Ping:
				c.Write(irc.NewPong(m.Code()))
			case irc.Join:
				b.channels = append(b.channels, m.Channel())
			case irc.PrivateMessage:
				msg := m.Message()
				ctcp := len(msg) > 2 && msg[0] == '\x01'
				if ctcp {
					msg = msg[1 : len(msg)-1]
				}
				isCommand := strings.HasPrefix(msg, b.commandPrefix)

				if !isCommand {
					b.handleMessage(parseUser(m.Sender()), m.Target(), m.Message())
				} else {
					go b.handleCommand(
						parseUser(m.Sender()),
						m.Target(),
						m.Message()[len(b.commandPrefix):],
					)
				}
			case irc.Notice:
			// Notice
			case irc.Mode:
				// TODO Track others + channels
				if m.Target() == b.username {
					state := '#'
					for _, r := range m.Mode() {
						switch r {
						case '-':
							state = '-'
						case '+':
							state = '+'
						default:
							if state == '#' {
								panic("Invalid mode!")
							}
							if state == '+' {
								b.modes[r] = struct{}{}
							} else {
								delete(b.modes, r)
							}
						}
					}
				}
			default:
				log.Printf("Unhandled: %#v\n", m)
			}
		}
	}
}

// AddChatHandler adds a handler which is called
// whenever the passed regexp matches a message
func (b *Bot) AddChatHandler(r *regexp.Regexp, f func(b *Bot, sender User, target, message string) error) {
	b.chatLock.Lock()
	defer b.chatLock.Unlock()
	b.chatHandlers = append(b.chatHandlers, chatHandler{r, f})
}

func (b *Bot) handleMessage(sender User, target, message string) {
	b.chatLock.RLock()
	defer b.chatLock.RUnlock()
	for _, h := range b.chatHandlers {
		r, f := h.r, h.f
		if r.MatchString(message) {
			go func() {
				if err := f(b, sender, target, message); err != nil {
					b.Reply(sender, target, err.Error())
				}
			}()
			break
		}
	}
}

func (b *Bot) handleCommand(user User, target, msg string) {
	err := b.Commands.Execute(b, msg, user, target)
	if err != nil {
		b.Reply(user, target, err.Error())
	}
}

func (b *Bot) handleReply(r irc.Reply) {
	switch r.Code() {
	case irc.ReplyWelcome:
		b.Events <- Connected{}
	}
}

func (b *Bot) kill() {
	b.Events <- Stop{}
	close(b.Events)
}

// Error returns the error the bot stopped with or nil
// if no error has occurred
func (b *Bot) Error() error {
	return b.err
}

// JoinChannel attempts to join the target channel
//
// Channel names are generally prefixed with #
func (b *Bot) JoinChannel(channel string) {
	b.writeChan <- irc.NewJoin(channel)
}

// Channels returns the list of channels this bot is currently
// connected to
func (b *Bot) Channels() []string {
	ret := make(chan []string)
	b.funcChan <- func() {
		ret <- b.channels
	}
	return <-ret
}

// SendMessage sends a message to the target channel or user
func (b *Bot) SendMessage(target, message string) {
	b.writeChan <- irc.NewPrivateMessage(target, message)
}

// SendCTCP sends a CTCP message to the target channel or user
func (b *Bot) SendCTCP(target, message string) {
	b.writeChan <- irc.NewPrivateMessage(target, "\x01"+message+"\x01")
}

// AddMode sets the mode(s) on the bot
func (b *Bot) AddMode(modes ...rune) {
	b.writeChan <- irc.NewMode(b.username, "+"+string(modes))
}

// RemoveMode removes the mode(s) from the bot
func (b *Bot) RemoveMode(modes ...rune) {
	b.writeChan <- irc.NewMode(b.username, "-"+string(modes))
}

// Reply sends a message to the user in the same way
// the message was sent.
//
// If the message was sent in a channel the message
// will be sent back to that channel with the sender's
// nickname prefixed.
//
// If the message was sent in a private message then
// this will just reply normally
func (b *Bot) Reply(sender User, target, message string) {
	if target[0] == '#' {
		b.SendMessage(target, fmt.Sprintf("%s: %s", sender.Nickname, message))
	} else {
		b.SendMessage(sender.Nickname, message)
	}
}
<commit_msg>/: remove the lock from AddChatHandler<commit_after>/*
 * Copyright 2015 Matthew Collins
 *
 * 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 thinkbot

import (
	"fmt"
	"github.com/thinkofdeath/thinkbot/command"
	"github.com/thinkofdeath/thinkbot/irc"
	"log"
	"regexp"
	"strings"
)

// Bot contains all the information of the currently
// running bot.
//
// Event should be continuously read from and the
// returned events handled.
type Bot struct {
	client   *irc.Client
	err      error
	username string

	Commands command.Registry

	chatHandlers []chatHandler

	Events    chan Event
	writeChan chan irc.Message
	funcChan  chan func()

	channels      []string
	commandPrefix string
	modes         map[rune]struct{}

	permissionContainer PermissionContainer
}

type chatHandlerFunc func(b *Bot, sender User, target, message string) error

type chatHandler struct {
	r *regexp.Regexp
	f chatHandlerFunc
}

// NewBot creates an instance of a bot connecting to the target
// server with the specified username
//
// The init function is called before any messages/commands
// are handled to allow for setup
//
// Returns an error if it fails to connect
func NewBot(server string, port uint16, username string, init func(*Bot)) (*Bot, error) {
	c, err := irc.NewClient(server, port)
	if err != nil {
		return nil, err
	}
	b := &Bot{
		client:        c,
		Events:        make(chan Event, 100),
		writeChan:     make(chan irc.Message, 100),
		funcChan:      make(chan func(), 100),
		username:      username,
		channels:      []string{},
		commandPrefix: "+",
		modes:         map[rune]struct{}{},
		Commands: command.Registry{
			// User and target parameters
			ExtraParameters: 2,
		},
	}
	init(b)
	go b.run()
	return b, nil
}

func (b *Bot) run() {
	c := b.client
	rc := make(chan irc.Message)

	go func() {
		for {
			m, err := c.Read()
			if err != nil {
				b.err = err
				b.kill()
				close(rc)
				return
			}
			rc <- m
		}
	}()

	c.Write(irc.NewNick(b.username))
	c.Write(irc.NewUser(b.username, b.username))

	for {
		select {
		case f := <-b.funcChan:
			f()
		case m := <-b.writeChan:
			c.Write(m)
		case m, ok := <-rc:
			if !ok {
				return
			}
			switch m := m.(type) {
			case irc.Reply:
				b.handleReply(m)
			case irc.Ping:
				c.Write(irc.NewPong(m.Code()))
			case irc.Join:
				b.channels = append(b.channels, m.Channel())
			case irc.PrivateMessage:
				msg := m.Message()
				ctcp := len(msg) > 2 && msg[0] == '\x01'
				if ctcp {
					msg = msg[1 : len(msg)-1]
				}
				isCommand := strings.HasPrefix(msg, b.commandPrefix)

				if !isCommand {
					b.handleMessage(parseUser(m.Sender()), m.Target(), m.Message())
				} else {
					go b.handleCommand(
						parseUser(m.Sender()),
						m.Target(),
						m.Message()[len(b.commandPrefix):],
					)
				}
			case irc.Notice:
			// Notice
			case irc.Mode:
				// TODO Track others + channels
				if m.Target() == b.username {
					state := '#'
					for _, r := range m.Mode() {
						switch r {
						case '-':
							state = '-'
						case '+':
							state = '+'
						default:
							if state == '#' {
								panic("Invalid mode!")
							}
							if state == '+' {
								b.modes[r] = struct{}{}
							} else {
								delete(b.modes, r)
							}
						}
					}
				}
			default:
				log.Printf("Unhandled: %#v\n", m)
			}
		}
	}
}

// AddChatHandler adds a handler which is called
// whenever the passed regexp matches a message
func (b *Bot) AddChatHandler(r *regexp.Regexp, f chatHandlerFunc) {
	b.funcChan <- func() {
		b.chatHandlers = append(b.chatHandlers, chatHandler{r, f})
	}
}

func (b *Bot) handleMessage(sender User, target, message string) {
	for _, h := range b.chatHandlers {
		r, f := h.r, h.f
		if r.MatchString(message) {
			go func() {
				if err := f(b, sender, target, message); err != nil {
					b.Reply(sender, target, err.Error())
				}
			}()
			break
		}
	}
}

func (b *Bot) handleCommand(user User, target, msg string) {
	err := b.Commands.Execute(b, msg, user, target)
	if err != nil {
		b.Reply(user, target, err.Error())
	}
}

func (b *Bot) handleReply(r irc.Reply) {
	switch r.Code() {
	case irc.ReplyWelcome:
		b.Events <- Connected{}
	}
}

func (b *Bot) kill() {
	b.Events <- Stop{}
	close(b.Events)
}

// Error returns the error the bot stopped with or nil
// if no error has occurred
func (b *Bot) Error() error {
	return b.err
}

// JoinChannel attempts to join the target channel
//
// Channel names are generally prefixed with #
func (b *Bot) JoinChannel(channel string) {
	b.writeChan <- irc.NewJoin(channel)
}

// Channels returns the list of channels this bot is currently
// connected to
func (b *Bot) Channels() []string {
	ret := make(chan []string)
	b.funcChan <- func() {
		ret <- b.channels
	}
	return <-ret
}

// SendMessage sends a message to the target channel or user
func (b *Bot) SendMessage(target, message string) {
	b.writeChan <- irc.NewPrivateMessage(target, message)
}

// SendCTCP sends a CTCP message to the target channel or user
func (b *Bot) SendCTCP(target, message string) {
	b.writeChan <- irc.NewPrivateMessage(target, "\x01"+message+"\x01")
}

// AddMode sets the mode(s) on the bot
func (b *Bot) AddMode(modes ...rune) {
	b.writeChan <- irc.NewMode(b.username, "+"+string(modes))
}

// RemoveMode removes the mode(s) from the bot
func (b *Bot) RemoveMode(modes ...rune) {
	b.writeChan <- irc.NewMode(b.username, "-"+string(modes))
}

// Reply sends a message to the user in the same way
// the message was sent.
//
// If the message was sent in a channel the message
// will be sent back to that channel with the sender's
// nickname prefixed.
//
// If the message was sent in a private message then
// this will just reply normally
func (b *Bot) Reply(sender User, target, message string) {
	if target[0] == '#' {
		b.SendMessage(target, fmt.Sprintf("%s: %s", sender.Nickname, message))
	} else {
		b.SendMessage(sender.Nickname, message)
	}
}
<|endoftext|>