<commit_before>package main

import (
	"fmt"
	"github.com/sorcix/irc"
)

type Session struct {
	*irc.Conn
	// Set this bit on construction
	Server   string
	Port     int
	UserName string
	NickName string
	RealName string
	// Set by the internals
	messages <-chan *irc.Message
}

func (session *Session) Dial() error {
	conn, err := irc.Dial(fmt.Sprintf("%s:%d", session.Server, session.Port))
	if err != nil {
		return err
	}
	session.Conn = conn
	return session.handshake()
}

func (session *Session) handshake() (err error) {
	err = session.Encode(&irc.Message{
		Command: "NICK",
		Params:  []string{session.NickName},
	})
	if err != nil {
		return err
	}
	return session.Encode(&irc.Message{
		Command:  "USER",
		Params:   []string{session.UserName, "0", "*"},
		Trailing: session.RealName,
	})
}

func (session *Session) Privmsg(name, message string) error {
	return session.Encode(&irc.Message{
		Command:  "PRIVMSG",
		Params:   []string{name},
		Trailing: message,
	})
}

func (session *Session) Quit(reason string) error {
	if reason == "" {
		reason = "Shutting down"
	}
	defer session.Close()
	return session.Encode(&irc.Message{
		Command:  "QUIT",
		Trailing: reason,
	})
}

func (session *Session) handlePing(message *irc.Message) (err error) {
	// Ha ha ha this is so dodgy
	message.Command = "PONG"
	return session.Encode(message)
}

func (session *Session) readPump() (err error) {
	var toIgnore = [...]string{"001", "002", "003", "005", "251", "252", "254", "255", "265", "266"}
	var shouldIgnore = make(map[string]bool, len(toIgnore))
	for _, num := range toIgnore {
		shouldIgnore[num] = true
	}
	var m *irc.Message

	for {
		m, err = session.Decode()
		if err != nil {
			return err
		}
		// Ignore informative spam
		if shouldIgnore[m.Command] {
			continue
		}
		// Start of glorious message type switches
		if m.Command == "PING" {
			err = session.handlePing(m)
			if err != nil {
				log.Error("Couldn't pong: %s", err)
			}
		} else if m.Command == "004" {
			// The server sends 001 002 003 and 004 in fast succession.
			// It's probably a good idea to wait until #4 happens before doing anything
			log.Info("Connection established")
			setupNickserv(session)
		} else if m.Command == "376" || m.Command == "422" {
			// MOTD is finished, start harassing people
			// log.Info("MOTD Complete, fully connected!")
		} else if m.Command == "ERROR" {
			log.Critical("Server hung up: %s", m.Trailing)
			return nil
		} else {
			log.Info("Got message: %+v", m)
		}
	}
}
<commit_msg>Do some light cleanup<commit_after>package main

import (
	"fmt"
	"github.com/sorcix/irc"
)

type Session struct {
	*irc.Conn
	// Set this bit on construction
	Server   string
	Port     int
	UserName string
	NickName string
	RealName string
	// Set by the internals
	messages <-chan *irc.Message
}

func (session *Session) Dial() error {
	conn, err := irc.Dial(fmt.Sprintf("%s:%d", session.Server, session.Port))
	if err != nil {
		return err
	}
	session.Conn = conn
	return session.handshake()
}

func (session *Session) handshake() (err error) {
	err = session.Encode(&irc.Message{
		Command: "NICK",
		Params:  []string{session.NickName},
	})
	if err != nil {
		return err
	}
	return session.Encode(&irc.Message{
		Command:  "USER",
		Params:   []string{session.UserName, "0", "*"},
		Trailing: session.RealName,
	})
}

func (session *Session) Privmsg(name, message string) error {
	return session.Encode(&irc.Message{
		Command:  "PRIVMSG",
		Params:   []string{name},
		Trailing: message,
	})
}

func (session *Session) Quit(reason string) error {
	if reason == "" {
		reason = "Shutting down"
	}
	defer session.Close()
	return session.Encode(&irc.Message{
		Command:  "QUIT",
		Trailing: reason,
	})
}

func (session *Session) handlePing(message *irc.Message) (err error) {
	// Ha ha ha this is so dodgy
	message.Command = "PONG"
	return session.Encode(message)
}

var toIgnore = [...]string{
	"002",
	"003",
	"004",
	"005",
	"251",
	"252",
	"254",
	"255",
	"265",
	"266",
	"376",
	"422",
}

func (session *Session) readPump() (err error) {
	var shouldIgnore = make(map[string]bool, len(toIgnore))
	for _, num := range toIgnore {
		shouldIgnore[num] = true
	}
	var m *irc.Message

	for {
		m, err = session.Decode()
		if err != nil {
			return err
		}
		// Ignore informative spam
		if shouldIgnore[m.Command] {
			continue
		}
		// Start of glorious message type switches
		if m.Command == "PING" {
			log.Debug("Got ping: %s", m.Trailing)
			err = session.handlePing(m)
			if err != nil {
				log.Error("Couldn't pong: %s", err)
			}
		} else if m.Command == "001" {
			log.Info("Connection established")
			setupNickserv(session)
		} else if m.Command == "ERROR" {
			log.Critical("Server hung up: %s", m.Trailing)
			return nil
		} else {
			log.Debug("Got unhandled message: %+v", m)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"code.google.com/p/gosshold/ssh"
	"errors"
	"fmt"
	"io"
	"launchpad.net/goamz/aws"
	"launchpad.net/goamz/ec2"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
)

var ErrNoInstancesFound = errors.New("No instances found; run provisioner first")

const AfterDeployHookScript = ".moltar-after-deploy"

type Job struct {
	region                  aws.Region
	env                     string
	cluster                 string
	project                 string
	packageNames            []string
	instances               []*ec2.Instance
	instanceSshClients      map[*ec2.Instance]*ssh.ClientConn
	instanceLoggers         map[*ec2.Instance]*log.Logger
	output                  io.Writer
	logger                  *log.Logger
	installVersionRev       uint64
	shouldOutputAnsiEscapes bool
}

func getInstancesTagged(ec2client *ec2.EC2, project string, env string, cluster string, packageName string) (instances []*ec2.Instance, err error) {
	instanceFilter := ec2.NewFilter()
	instanceFilter.Add("instance-state-name", "running")
	instanceFilter.Add("tag:Project", project)
	queryEnv := env
	if env == "" {
		queryEnv = "*"
	}
	instanceFilter.Add("tag:Environment", queryEnv)
	if cluster != "" {
		instanceFilter.Add("tag:Cluster", cluster)
	}

	if packageName != "" {
		instanceFilter.Add("tag:Packages", "*|"+packageName+"|*")
	}

	instancesResp, err := ec2client.Instances(nil, instanceFilter)
	if err != nil {
		return
	}

	instances = make([]*ec2.Instance, 0, 20)
	for _, res := range instancesResp.Reservations {
		for _, inst := range res.Instances {
			newInst := inst
			instances = append(instances, &newInst)
		}
	}

	return instances, nil
}

func NewJob(awsConf AWSConf, env string, cluster string, project string, packageNames []string, output io.Writer, shouldOutputAnsiEscapes bool) (job *Job, err error) {
	e := ec2.New(awsConf.Auth, awsConf.Region)

	var searchPackageNames []string
	if len(packageNames) == 0 {
		searchPackageNames = []string{""}
	} else {
		searchPackageNames = packageNames[:]
	}

	instancesSet := map[string]*ec2.Instance{}
	for _, packageName := range searchPackageNames {
		instances, err := getInstancesTagged(e, project, env, cluster, packageName)
		if err != nil {
			return nil, err
		}

		for _, instance := range instances {
			instancesSet[instance.InstanceId] = instance
		}
	}

	instances := make([]*ec2.Instance, 0, len(instancesSet))
	for _, instance := range instancesSet {
		instances = append(instances, instance)
	}

	if len(instances) == 0 {
		return nil, ErrNoInstancesFound
	}

	logger := log.New(output, "", 0)

	return &Job{region: awsConf.Region, env: env, cluster: cluster,
		project: project, packageNames: packageNames, instances: instances,
		instanceSshClients: make(map[*ec2.Instance]*ssh.ClientConn),
		instanceLoggers:    make(map[*ec2.Instance]*log.Logger),
		output:             output, logger: logger,
		shouldOutputAnsiEscapes: shouldOutputAnsiEscapes}, nil
}

func (self *Job) Exec(cmd string) (errs []error) {
	errChan := make(chan error, len(self.instances))
	errs = make([]error, 0, len(self.instances))

	for _, instance := range self.instances {
		go func(inst ec2.Instance) {
			conn, err := self.sshClient(&inst)
			if err != nil {
				errChan <- err
				return
			}

			logger := self.instanceLogger(&inst)
			_, returnChan, err := sshRunOutLogger(conn, cmd, logger, nil)
			if err == nil {
				err = <-returnChan
			} else {
				errChan <- err
			}
			errChan <- err
		}(*instance)
	}
	startStdinRead()

	for _ = range self.instances {
		if err := <-errChan; err != nil {
			errs = append(errs, err)
		}
	}
	return
}

func (self *Job) ExecList(cmds []string) (errs []error) {
	for _, cmd := range cmds {
		fmt.Printf("\n%s\n\n", cmd)
		errs = self.Exec(cmd)
		if len(errs) > 0 {
			return
		}
	}
	return []error{}
}

func (self *Job) Deploy() (errs []error) {
	errs = self.ExecList([]string{
		"sudo apt-get update -qq",
		"sudo DEBIAN_FRONTEND=noninteractive apt-get install -qy '" +
			strings.Join(self.packageNames, "' '") + "'",
		"sudo apt-get autoremove -yq",
		"sudo apt-get clean -yq",
	})
	if len(errs) > 0 {
		return
	}

	if _, err := os.Stat(AfterDeployHookScript); err != nil {
		return
	}

	prepareExec()
	pwd, err := os.Getwd()
	if err != nil {
		return
	}
	syscall.Exec(path.Join(pwd, AfterDeployHookScript),
		[]string{AfterDeployHookScript},
		append(os.Environ(), "ENV="+self.env))

	return
}

func (self *Job) Ssh(criteria string, sshArgs []string) (err error) {
	sshPath, err := exec.LookPath("ssh")
	if err != nil {
		return
	}

	var instance *ec2.Instance
	matches := self.instances

	if criteria != "-1" {
		if criteria != "" {
			matches = make([]*ec2.Instance, 0, len(self.instances))
			for _, instance = range self.instances {
				if matchCriteria(instance, criteria) {
					instanceLogName(instance)
					matches = append(matches, instance)
				}
			}
		}

		if len(matches) == 0 {
			self.logger.Fatalf("Instance '%s' not found\n", criteria)
		} else if len(matches) > 1 {
			self.logger.Printf("Multiple matches for '%s' found:\n", criteria)
			self.printInstances(matches)
			self.logger.Fatal("")
		}
	}

	instance = matches[0]

	execArgs := []string{"ssh"}
	keyFile := self.keyFile()
	if keyFile != "" {
		execArgs = append(execArgs, "-i", keyFile)
	}

	execArgs = append(execArgs,
		fmt.Sprintf("%s@%s", self.sshUserName(instance), instance.DNSName))
	execArgs = append(execArgs, sshArgs...)

	fPrintShellCommand(self.output, "", execArgs)
	fmt.Fprintln(self.output, "")

	prepareExec()
	err = syscall.Exec(sshPath, execArgs, os.Environ())
	return
}

func (self *Job) Scp(args []string) (err error) {
	scpPath, err := exec.LookPath("scp")
	if err != nil {
		return
	}

	defaultArgs := []string{"-q"}
	keyFile := self.keyFile()
	if keyFile != "" {
		defaultArgs = append(defaultArgs, []string{
			"-i", keyFile,
		}...)
	}
	scpArgs := make([]string, len(defaultArgs)+len(args))
	copy(scpArgs, defaultArgs)
	copy(scpArgs[len(defaultArgs):], args)

	var dstIndex = -1
	for i, arg := range scpArgs {
		if arg[0] == ':' {
			dstIndex = i
			break
		}
	}
	if dstIndex == -1 {
		dstIndex = len(scpArgs)
		scpArgs = append(scpArgs, ":")
	}

	errChan := make(chan error, len(self.instances))

	for _, instance := range self.instances {
		go func(instance *ec2.Instance) {
			var err error
			args := make([]string, len(scpArgs))
			copy(args, scpArgs)

			logger := self.instanceLogger(instance)
			args[dstIndex] = fmt.Sprintf("%s@%s%s",
				self.sshUserName(instance), instance.DNSName, args[dstIndex])

			fPrintShellCommand(self.output, "scp", args)

			cmd := exec.Command(scpPath, args...)
			outPipeRead, outPipeWrite, err := os.Pipe()
			if err != nil {
				logger.Printf("error creating pipe: %s\n", err)
				errChan <- err
				return
			}
			cmd.Stdout = outPipeWrite
			cmd.Stderr = outPipeWrite

			err = cmd.Start()
			if err != nil {
				logger.Printf("error starting scp: %s\n", err)
				errChan <- err
				return
			}

			outPipeWrite.Close()
			stdoutReader := bufio.NewReader(outPipeRead)
			for {
				in, err := stdoutReader.ReadString('\n')
				if (err == io.EOF && in != "") || err == nil {
					logger.Print(in)
				}
				if err != nil {
					break
				}
			}

			err = cmd.Wait()
			outPipeRead.Close()
			if err != nil {
				logger.Printf("error running scp: %s\n", err)
			}
			errChan <- err
		}(instance)
	}

	var scpErr error
	for _ = range self.instances {
		scpErr = <-errChan
		if err == nil && scpErr != nil {
			err = errors.New("at least one scp failed")
		}
	}
	return
}

func (self *Job) List() (err error) {
	self.printInstances(self.instances)
	return nil
}

func (self *Job) Hostname(instanceName string) (err error) {
	for _, instance := range self.instances {
		if instanceLogName(instance) == instanceName {
			fmt.Fprintln(self.output, instance.DNSName)
			return nil
		}
	}
	return errors.New(instanceName + " not found")
}

/// Subtasks

func (self *Job) sshClient(i *ec2.Instance) (conn *ssh.ClientConn, err error) {
	conn = self.instanceSshClients[i]
	if conn == nil {
		conn, err = self.sshDial(i)
		if err == nil {
			self.instanceSshClients[i] = conn
		}
	}
	return
}

func (self *Job) instanceLogger(i *ec2.Instance) (logger *log.Logger) {
	logger = self.instanceLoggers[i]
	if logger == nil {
		prefix := instanceLogName(i)
		if self.shouldOutputAnsiEscapes {
			prefix = "\033[1m" + prefix + "\033[0m"
		}
		logger = log.New(self.output, prefix+" ", 0)
		self.instanceLoggers[i] = logger
	}
	return
}

func (self *Job) keyFile() (path string) {
	fileName := self.project
	if len(self.packageNames) > 0 {
		fileName += fmt.Sprintf("-%s", self.packageNames[0])
	}
	path = fmt.Sprintf(os.ExpandEnv("${HOME}/Google Drive/%s Ops/Keys/%s.pem"),
		self.project, fileName)

	if _, err := os.Stat(path); err == nil {
		return path
	} else {
		return ""
	}
}

func (self *Job) sshUserName(_ *ec2.Instance) (userName string) {
	// TODO: be more clever about this
	return "ubuntu"
}

func (self *Job) sshDial(i *ec2.Instance) (conn *ssh.ClientConn, err error) {
	conn, err = sshDial(i.DNSName+":22", self.sshUserName(i), self.keyFile())
	return
}

func (self *Job) printInstances(instances []*ec2.Instance) {
	fields := make([][]string, len(instances))
	for i, instance := range instances {
		fields[i] = []string{instance.InstanceId, instanceLogName(instance),
			instance.DNSName}
	}
	fmt.Fprint(self.output, formatTable(fields))
}

func instanceLogName(i *ec2.Instance) string {
	for _, tag := range i.Tags {
		if tag.Key == "Name" && tag.Value != "" {
			return tag.Value
		}
	}
	return i.InstanceId
}

func fPrintShellCommand(w io.Writer, n string, cmd []string) {
	if n != "" {
		fmt.Fprintf(w, "%s ", n)
	}
	for i, cmdPart := range cmd {
		// TODO: this escaping will work most of the time, but isn't that great
		if strings.ContainsAny(cmdPart, " $") {
			fmt.Fprintf(w, "'%s'", cmdPart)
		} else {
			fmt.Fprint(w, cmdPart)
		}
		if i < (len(cmd) - 1) {
			fmt.Fprint(w, " ")
		}
	}
	fmt.Fprint(w, "\n")
}

func matchCriteria(instance *ec2.Instance, criteria string) bool {
	var found bool
	for _, value := range strings.Split(criteria, "/") {
		found = false
		for _, tag := range instance.Tags {
			if strings.Contains(tag.Value, value) {
				found = true
				break
			}
		}
		if !strings.Contains(instance.InstanceId, value) && !strings.Contains(instance.PrivateDNSName, value) && !strings.Contains(instance.DNSName, value) && found == false {
			return false
		}
	}
	return true
}

func prepareExec() {
	/* There appears to be a bug with goamz where some fds are left open, and
	 * just closing them causes a crash. If we ask all fds > 2 to close on
	 * exec, all is well.
	 */
	var rlimit syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rlimit)
	if err != nil {
		panic(err)
	}
	maxFds := int(rlimit.Cur)
	for fd := 3; fd < maxFds; fd++ {
		syscall.CloseOnExec(fd)
	}
}
<commit_msg>apt-get autoremove should be noninteractive too<commit_after>package main

import (
	"bufio"
	"code.google.com/p/gosshold/ssh"
	"errors"
	"fmt"
	"io"
	"launchpad.net/goamz/aws"
	"launchpad.net/goamz/ec2"
	"log"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
)

var ErrNoInstancesFound = errors.New("No instances found; run provisioner first")

const AfterDeployHookScript = ".moltar-after-deploy"

type Job struct {
	region                  aws.Region
	env                     string
	cluster                 string
	project                 string
	packageNames            []string
	instances               []*ec2.Instance
	instanceSshClients      map[*ec2.Instance]*ssh.ClientConn
	instanceLoggers         map[*ec2.Instance]*log.Logger
	output                  io.Writer
	logger                  *log.Logger
	installVersionRev       uint64
	shouldOutputAnsiEscapes bool
}

func getInstancesTagged(ec2client *ec2.EC2, project string, env string, cluster string, packageName string) (instances []*ec2.Instance, err error) {
	instanceFilter := ec2.NewFilter()
	instanceFilter.Add("instance-state-name", "running")
	instanceFilter.Add("tag:Project", project)
	queryEnv := env
	if env == "" {
		queryEnv = "*"
	}
	instanceFilter.Add("tag:Environment", queryEnv)
	if cluster != "" {
		instanceFilter.Add("tag:Cluster", cluster)
	}

	if packageName != "" {
		instanceFilter.Add("tag:Packages", "*|"+packageName+"|*")
	}

	instancesResp, err := ec2client.Instances(nil, instanceFilter)
	if err != nil {
		return
	}

	instances = make([]*ec2.Instance, 0, 20)
	for _, res := range instancesResp.Reservations {
		for _, inst := range res.Instances {
			newInst := inst
			instances = append(instances, &newInst)
		}
	}

	return instances, nil
}

func NewJob(awsConf AWSConf, env string, cluster string, project string, packageNames []string, output io.Writer, shouldOutputAnsiEscapes bool) (job *Job, err error) {
	e := ec2.New(awsConf.Auth, awsConf.Region)

	var searchPackageNames []string
	if len(packageNames) == 0 {
		searchPackageNames = []string{""}
	} else {
		searchPackageNames = packageNames[:]
	}

	instancesSet := map[string]*ec2.Instance{}
	for _, packageName := range searchPackageNames {
		instances, err := getInstancesTagged(e, project, env, cluster, packageName)
		if err != nil {
			return nil, err
		}

		for _, instance := range instances {
			instancesSet[instance.InstanceId] = instance
		}
	}

	instances := make([]*ec2.Instance, 0, len(instancesSet))
	for _, instance := range instancesSet {
		instances = append(instances, instance)
	}

	if len(instances) == 0 {
		return nil, ErrNoInstancesFound
	}

	logger := log.New(output, "", 0)

	return &Job{region: awsConf.Region, env: env, cluster: cluster,
		project: project, packageNames: packageNames, instances: instances,
		instanceSshClients: make(map[*ec2.Instance]*ssh.ClientConn),
		instanceLoggers:    make(map[*ec2.Instance]*log.Logger),
		output:             output, logger: logger,
		shouldOutputAnsiEscapes: shouldOutputAnsiEscapes}, nil
}

func (self *Job) Exec(cmd string) (errs []error) {
	errChan := make(chan error, len(self.instances))
	errs = make([]error, 0, len(self.instances))

	for _, instance := range self.instances {
		go func(inst ec2.Instance) {
			conn, err := self.sshClient(&inst)
			if err != nil {
				errChan <- err
				return
			}

			logger := self.instanceLogger(&inst)
			_, returnChan, err := sshRunOutLogger(conn, cmd, logger, nil)
			if err == nil {
				err = <-returnChan
			} else {
				errChan <- err
			}
			errChan <- err
		}(*instance)
	}
	startStdinRead()

	for _ = range self.instances {
		if err := <-errChan; err != nil {
			errs = append(errs, err)
		}
	}
	return
}

func (self *Job) ExecList(cmds []string) (errs []error) {
	for _, cmd := range cmds {
		fmt.Printf("\n%s\n\n", cmd)
		errs = self.Exec(cmd)
		if len(errs) > 0 {
			return
		}
	}
	return []error{}
}

func (self *Job) Deploy() (errs []error) {
	errs = self.ExecList([]string{
		"sudo apt-get update -qq",
		"sudo DEBIAN_FRONTEND=noninteractive apt-get install -qy '" +
			strings.Join(self.packageNames, "' '") + "'",
		"sudo DEBIAN_FRONTEND=noninteractive apt-get autoremove -yq",
		"sudo apt-get clean -yq",
	})
	if len(errs) > 0 {
		return
	}

	if _, err := os.Stat(AfterDeployHookScript); err != nil {
		return
	}

	prepareExec()
	pwd, err := os.Getwd()
	if err != nil {
		return
	}
	syscall.Exec(path.Join(pwd, AfterDeployHookScript),
		[]string{AfterDeployHookScript},
		append(os.Environ(), "ENV="+self.env))

	return
}

func (self *Job) Ssh(criteria string, sshArgs []string) (err error) {
	sshPath, err := exec.LookPath("ssh")
	if err != nil {
		return
	}

	var instance *ec2.Instance
	matches := self.instances

	if criteria != "-1" {
		if criteria != "" {
			matches = make([]*ec2.Instance, 0, len(self.instances))
			for _, instance = range self.instances {
				if matchCriteria(instance, criteria) {
					instanceLogName(instance)
					matches = append(matches, instance)
				}
			}
		}

		if len(matches) == 0 {
			self.logger.Fatalf("Instance '%s' not found\n", criteria)
		} else if len(matches) > 1 {
			self.logger.Printf("Multiple matches for '%s' found:\n", criteria)
			self.printInstances(matches)
			self.logger.Fatal("")
		}
	}

	instance = matches[0]

	execArgs := []string{"ssh"}
	keyFile := self.keyFile()
	if keyFile != "" {
		execArgs = append(execArgs, "-i", keyFile)
	}

	execArgs = append(execArgs,
		fmt.Sprintf("%s@%s", self.sshUserName(instance), instance.DNSName))
	execArgs = append(execArgs, sshArgs...)

	fPrintShellCommand(self.output, "", execArgs)
	fmt.Fprintln(self.output, "")

	prepareExec()
	err = syscall.Exec(sshPath, execArgs, os.Environ())
	return
}

func (self *Job) Scp(args []string) (err error) {
	scpPath, err := exec.LookPath("scp")
	if err != nil {
		return
	}

	defaultArgs := []string{"-q"}
	keyFile := self.keyFile()
	if keyFile != "" {
		defaultArgs = append(defaultArgs, []string{
			"-i", keyFile,
		}...)
	}
	scpArgs := make([]string, len(defaultArgs)+len(args))
	copy(scpArgs, defaultArgs)
	copy(scpArgs[len(defaultArgs):], args)

	var dstIndex = -1
	for i, arg := range scpArgs {
		if arg[0] == ':' {
			dstIndex = i
			break
		}
	}
	if dstIndex == -1 {
		dstIndex = len(scpArgs)
		scpArgs = append(scpArgs, ":")
	}

	errChan := make(chan error, len(self.instances))

	for _, instance := range self.instances {
		go func(instance *ec2.Instance) {
			var err error
			args := make([]string, len(scpArgs))
			copy(args, scpArgs)

			logger := self.instanceLogger(instance)
			args[dstIndex] = fmt.Sprintf("%s@%s%s",
				self.sshUserName(instance), instance.DNSName, args[dstIndex])

			fPrintShellCommand(self.output, "scp", args)

			cmd := exec.Command(scpPath, args...)
			outPipeRead, outPipeWrite, err := os.Pipe()
			if err != nil {
				logger.Printf("error creating pipe: %s\n", err)
				errChan <- err
				return
			}
			cmd.Stdout = outPipeWrite
			cmd.Stderr = outPipeWrite

			err = cmd.Start()
			if err != nil {
				logger.Printf("error starting scp: %s\n", err)
				errChan <- err
				return
			}

			outPipeWrite.Close()
			stdoutReader := bufio.NewReader(outPipeRead)
			for {
				in, err := stdoutReader.ReadString('\n')
				if (err == io.EOF && in != "") || err == nil {
					logger.Print(in)
				}
				if err != nil {
					break
				}
			}

			err = cmd.Wait()
			outPipeRead.Close()
			if err != nil {
				logger.Printf("error running scp: %s\n", err)
			}
			errChan <- err
		}(instance)
	}

	var scpErr error
	for _ = range self.instances {
		scpErr = <-errChan
		if err == nil && scpErr != nil {
			err = errors.New("at least one scp failed")
		}
	}
	return
}

func (self *Job) List() (err error) {
	self.printInstances(self.instances)
	return nil
}

func (self *Job) Hostname(instanceName string) (err error) {
	for _, instance := range self.instances {
		if instanceLogName(instance) == instanceName {
			fmt.Fprintln(self.output, instance.DNSName)
			return nil
		}
	}
	return errors.New(instanceName + " not found")
}

/// Subtasks

func (self *Job) sshClient(i *ec2.Instance) (conn *ssh.ClientConn, err error) {
	conn = self.instanceSshClients[i]
	if conn == nil {
		conn, err = self.sshDial(i)
		if err == nil {
			self.instanceSshClients[i] = conn
		}
	}
	return
}

func (self *Job) instanceLogger(i *ec2.Instance) (logger *log.Logger) {
	logger = self.instanceLoggers[i]
	if logger == nil {
		prefix := instanceLogName(i)
		if self.shouldOutputAnsiEscapes {
			prefix = "\033[1m" + prefix + "\033[0m"
		}
		logger = log.New(self.output, prefix+" ", 0)
		self.instanceLoggers[i] = logger
	}
	return
}

func (self *Job) keyFile() (path string) {
	fileName := self.project
	if len(self.packageNames) > 0 {
		fileName += fmt.Sprintf("-%s", self.packageNames[0])
	}
	path = fmt.Sprintf(os.ExpandEnv("${HOME}/Google Drive/%s Ops/Keys/%s.pem"),
		self.project, fileName)

	if _, err := os.Stat(path); err == nil {
		return path
	} else {
		return ""
	}
}

func (self *Job) sshUserName(_ *ec2.Instance) (userName string) {
	// TODO: be more clever about this
	return "ubuntu"
}

func (self *Job) sshDial(i *ec2.Instance) (conn *ssh.ClientConn, err error) {
	conn, err = sshDial(i.DNSName+":22", self.sshUserName(i), self.keyFile())
	return
}

func (self *Job) printInstances(instances []*ec2.Instance) {
	fields := make([][]string, len(instances))
	for i, instance := range instances {
		fields[i] = []string{instance.InstanceId, instanceLogName(instance),
			instance.DNSName}
	}
	fmt.Fprint(self.output, formatTable(fields))
}

func instanceLogName(i *ec2.Instance) string {
	for _, tag := range i.Tags {
		if tag.Key == "Name" && tag.Value != "" {
			return tag.Value
		}
	}
	return i.InstanceId
}

func fPrintShellCommand(w io.Writer, n string, cmd []string) {
	if n != "" {
		fmt.Fprintf(w, "%s ", n)
	}
	for i, cmdPart := range cmd {
		// TODO: this escaping will work most of the time, but isn't that great
		if strings.ContainsAny(cmdPart, " $") {
			fmt.Fprintf(w, "'%s'", cmdPart)
		} else {
			fmt.Fprint(w, cmdPart)
		}
		if i < (len(cmd) - 1) {
			fmt.Fprint(w, " ")
		}
	}
	fmt.Fprint(w, "\n")
}

func matchCriteria(instance *ec2.Instance, criteria string) bool {
	var found bool
	for _, value := range strings.Split(criteria, "/") {
		found = false
		for _, tag := range instance.Tags {
			if strings.Contains(tag.Value, value) {
				found = true
				break
			}
		}
		if !strings.Contains(instance.InstanceId, value) && !strings.Contains(instance.PrivateDNSName, value) && !strings.Contains(instance.DNSName, value) && found == false {
			return false
		}
	}
	return true
}

func prepareExec() {
	/* There appears to be a bug with goamz where some fds are left open, and
	 * just closing them causes a crash. If we ask all fds > 2 to close on
	 * exec, all is well.
	 */
	var rlimit syscall.Rlimit
	err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rlimit)
	if err != nil {
		panic(err)
	}
	maxFds := int(rlimit.Cur)
	for fd := 3; fd < maxFds; fd++ {
		syscall.CloseOnExec(fd)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Marin Procureur. All rights reserved.
// Use of gjp source code is governed by a MIT license
// license that can be found in the LICENSE file.

/*
Package gjp stands for Go JobPool, and is willing to be a simple jobpool manager. It maintains
a number of queues determined at the init. No priority whatsoever, just every queues are
processing one job at a time.
*/

package gjp

import (
	"encoding/json"
	"fmt"
	"time"
)

/*
   TYPES
*/

type (
	// Job
	Job struct {
		JobRunner `json:"-"` //skip the field for json
		Id        string     `json:"id"`
		Name      string     `json:"name"`   //Public property retrievable
		Status    string        `json:"status"` //Status of the current job
		Error     *JobError  `json:"-"`
		Start     time.Time  `json:"start"`
		End       time.Time  `json:"end"`
	}
)

/*
   JOB STATUS
*/
const (
	failed     string = "failed"
	success    string = "done"
	waiting    string = "queued"
	processing string = "proceeded"
)

func newJob(id string, jobRunner JobRunner, jobName string) (job *Job, jobId string) {
	job = &Job{
		JobRunner: jobRunner,
		Name:      jobName,
		Status:    waiting,
		Id:        id,
	}

	fmt.Println("New job with Id : ", job.Id)

	jobId = job.Id

	return
}

//execute the job safely and set the status back for the reportChannel
func (j *Job) executeJob(start time.Time) {
	defer catchPanic("Job", j.Name, "failed in executeJob")
	//Set the execution time for this job

	j.Start = start

	j.setJobToProcessing()

	j.NotifyStart(j)

	j.Error = j.ExecuteJob(j)

	//Set the job status
	switch j.Error {
	case nil:
		j.setJobToSuccess()
		break
	default:
		j.setJobToError()
		break
	}

	j.End = time.Now()

	j.NotifyEnd(j)

	return
}

/*
 GETTERS & SETTERS
*/

func (j *Job) HasJobErrored() (errored bool) {
	fmt.Println("Has job", j.GetJobName(), "errored ? ", j.Error != nil)
	if j.Error != nil {
		errored = true
	} else {
		errored = false
	}
	return
}

//create an error well formated
func (j *Job) GetJobError() (errorString string) {
	errorString = j.Error.FmtError()
	return
}

func (j *Job) getJobStringId() (jobId string) {
	jobId = j.Id
	return
}

func (j *Job) jobErrored() (jobError bool, error string) {
	if j.Status == failed {
		jobError = true
		error = j.GetJobError()
	}
	return
}

func (j *Job) GetJobStatus() (jobStatus string) {
	jobStatus = j.Status
	return
}

func (j *Job) getExecutionTime() (executionTime time.Duration) {
	nullTime := time.Time{}
	if j.End == nullTime {
		executionTime = j.Start.Sub(j.End)
	} else {
		executionTime = time.Since(j.Start)
	}
	return
}

func (j *Job) GetJobInfos() (jobjson []byte, err error) {
	jobjson, err = json.Marshal(j)
	if err != nil {
		fmt.Println(err.Error())
	}
	return
}

func (j *Job) GetJobId() (jobId string) {
	jobId = j.Id
	return
}

func (j *Job) GetJobName() (jobName string) {
	jobName = j.Name
	return
}

func (j *Job) setJobToWaiting() {
	j.Status = waiting
}

func (j *Job) setJobToError() {
	j.Status = failed
}

func (j *Job) setJobToSuccess() {
	j.Status = success
}

func (j *Job) setJobToProcessing() {
	j.Status = processing
}
<commit_msg>Notify when errored<commit_after>// Copyright 2016 Marin Procureur. All rights reserved.
// Use of gjp source code is governed by a MIT license
// license that can be found in the LICENSE file.

/*
Package gjp stands for Go JobPool, and is willing to be a simple jobpool manager. It maintains
a number of queues determined at the init. No priority whatsoever, just every queues are
processing one job at a time.
*/

package gjp

import (
	"encoding/json"
	"fmt"
	"time"
)

/*
   TYPES
*/

type (
	// Job
	Job struct {
		JobRunner `json:"-"` //skip the field for json
		Id        string     `json:"id"`
		Name      string     `json:"name"`   //Public property retrievable
		Status    string        `json:"status"` //Status of the current job
		Error     *JobError  `json:"-"`
		Start     time.Time  `json:"start"`
		End       time.Time  `json:"end"`
	}
)

/*
   JOB STATUS
*/
const (
	failed     string = "failed"
	success    string = "done"
	waiting    string = "queued"
	processing string = "proceeded"
)

func newJob(id string, jobRunner JobRunner, jobName string) (job *Job, jobId string) {
	job = &Job{
		JobRunner: jobRunner,
		Name:      jobName,
		Status:    waiting,
		Id:        id,
	}

	fmt.Println("New job with Id : ", job.Id)

	jobId = job.Id

	return
}

//execute the job safely and set the status back for the reportChannel
func (j *Job) executeJob(start time.Time) {
	defer catchPanic("Job", j.Name, "failed in executeJob")
	//Set the execution time for this job

	j.Start = start

	j.setJobToProcessing()

	j.NotifyStart(j)

	j.Error = j.ExecuteJob(j)

	//Set the job status
	switch j.Error {
	case nil:
		j.setJobToSuccess()
		break
	default:
		j.setJobToError()
		break
	}

	j.End = time.Now()

	j.NotifyEnd(j)

	return
}

/*
 GETTERS & SETTERS
*/

func (j *Job) HasJobErrored() (errored bool) {
	fmt.Println("Has job", j.GetJobName(), "errored ? ", j.Error != nil)
	if j.Error != nil {
		errored = true
		j.NotifyEnd(j)
	} else {
		errored = false
	}
	return
}

//create an error well formated
func (j *Job) GetJobError() (errorString string) {
	errorString = j.Error.FmtError()
	return
}

func (j *Job) getJobStringId() (jobId string) {
	jobId = j.Id
	return
}

func (j *Job) jobErrored() (jobError bool, error string) {
	if j.Status == failed {
		jobError = true
		error = j.GetJobError()
	}
	return
}

func (j *Job) GetJobStatus() (jobStatus string) {
	jobStatus = j.Status
	return
}

func (j *Job) getExecutionTime() (executionTime time.Duration) {
	nullTime := time.Time{}
	if j.End == nullTime {
		executionTime = j.Start.Sub(j.End)
	} else {
		executionTime = time.Since(j.Start)
	}
	return
}

func (j *Job) GetJobInfos() (jobjson []byte, err error) {
	jobjson, err = json.Marshal(j)
	if err != nil {
		fmt.Println(err.Error())
	}
	return
}

func (j *Job) GetJobId() (jobId string) {
	jobId = j.Id
	return
}

func (j *Job) GetJobName() (jobName string) {
	jobName = j.Name
	return
}

func (j *Job) setJobToWaiting() {
	j.Status = waiting
}

func (j *Job) setJobToError() {
	j.Status = failed
}

func (j *Job) setJobToSuccess() {
	j.Status = success
}

func (j *Job) setJobToProcessing() {
	j.Status = processing
}
<|endoftext|>


<commit_before>// Copyright 2014 Matthew Endsley
// All rights reserved
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted providing that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
// OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
// IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package gojwe

import (
	"bytes"
	"compress/flate"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/json"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"strings"
)

// Verify and decrypt a draft-7 JWE object
func VerifyAndDecryptDraft7(jwe string, key crypto.PrivateKey) ([]byte, error) {
	parts := strings.Split(jwe, ".")
	if len(parts) != 5 {
		return nil, errors.New("Wrong number of parts")
	}

	// decode the JWE header
	var header struct {
		Alg string `json:"alg"`
		Enc string `json:"enc"`
		Zip string `json:"zip"`
	}
	data, err := safeDecode(parts[0])
	if err != nil {
		return nil, fmt.Errorf("Malformed header: %v", err)
	}
	err = json.Unmarshal(data, &header)
	if err != nil {
		return nil, fmt.Errorf("Failed to decode header: %v", err)
	}

	var encryptionKey []byte
	encryptionKeyData, err := safeDecode(parts[1])
	if err != nil {
		return nil, fmt.Errorf("Malformed encryption key: %v", err)
	}

	// decode the encryption key
	switch header.Alg {
	case "RSA-OAEP", "RSA-OAEP-256":
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected an RSA private key. Got %T", key)
		}

		var h hash.Hash
		if header.Alg == "RSA-OAEP" {
			h = sha1.New()
		} else if header.Alg == "RSA-OAEP-256" {
			h = sha256.New()
		} else {
			return nil, fmt.Errorf("Unknown RSA-OAEP keytype %s", header.Alg)
		}

		encryptionKey, err = rsa.DecryptOAEP(h, rand.Reader, rsaKey, encryptionKeyData, nil)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case "RSA1_5":
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected RSA private key. Got %T", key)
		}

		encryptionKey, err = rsa.DecryptPKCS1v15(rand.Reader, rsaKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case "A128KW", "A256KW":
		aesKey, ok := key.([]byte)
		if !ok {
			return nil, fmt.Errorf("Expected shared symmetric key ([]byte). Got %T", key)
		}

		encryptionKey, err = AesKeyUnwrap(aesKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to unwrap key: %v", err)
		}

	default:
		return nil, fmt.Errorf("Unsupported ALG keytype %s", header.Alg)
	}

	// decode IV
	iv, err := safeDecode(parts[2])
	if err != nil {
		return nil, fmt.Errorf("Malformed IV: %v", err)
	}

	// decode cipher text
	cipherText, err := safeDecode(parts[3])
	if err != nil {
		return nil, fmt.Errorf("Malformed cipher text: %v", err)
	}

	// decode authtag
	authTag, err := safeDecode(parts[4])
	if err != nil {
		return nil, fmt.Errorf("Malformed authtag: %v", err)
	}

	// decrypt and verify cipher text
	var plainText []byte

	switch header.Enc {
	case "A128CBC+HS256", "A256CBC+HS512":
		// derive keys
		var encKey, macKey []byte
		var hfunc func() hash.Hash
		if header.Enc == "A128CBC+HS256" {
			encKey, macKey = concatKDF(encryptionKey, header.Enc, 128, 256)
			hfunc = sha256.New
		} else if header.Enc == "A256CBC+HS512" {
			encKey, macKey = concatKDF(encryptionKey, header.Enc, 256, 512)
			hfunc = sha512.New
		} else {
			panic("Bad ENC logic for type: " + header.Enc)
		}

		// verify authtag
		hm := hmac.New(hfunc, macKey)
		io.WriteString(hm, parts[0])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[1])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[2])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[3])
		if !hmac.Equal(authTag, hm.Sum(nil)) {
			return nil, errors.New("Integrity check failed")
		}

		// decrpyt ciphertext (in-place)
		block, err := aes.NewCipher(encKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}

		c := cipher.NewCBCDecrypter(block, iv)
		c.CryptBlocks(cipherText, cipherText)
		plainText = cipherText

		// remove PCKS#7 padding
		padding := int(plainText[len(plainText)-1])
		plainText = plainText[:len(plainText)-padding]

	case "A128GCM", "A256GCM":
		// create the "additional data" for the GCM cipher
		additionalData := new(bytes.Buffer)
		additionalData.WriteString(parts[0])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[1])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[2])

		// create the authenticating cipher
		block, err := aes.NewCipher(encryptionKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}
		c, err := cipher.NewGCM(block)
		if err != nil {
			return nil, fmt.Errorf("Failed to create GCM cipher: %v", err)
		}

		// decrypt the cipher text (in-place)
		_, err = c.Open(cipherText[:0], iv, append(cipherText, authTag...), additionalData.Bytes())
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt: %v", err)
		}
		plainText = cipherText

	default:
		return nil, fmt.Errorf("Unsupported ENC keytype %s", header.Enc)
	}

	// need to deflate plain text?
	if header.Zip == "DEF" {
		plainText, err = ioutil.ReadAll(flate.NewReader(bytes.NewReader(plainText)))
		if err != nil {
			return nil, fmt.Errorf("Failed to inflate plain text: %v", err)
		}
	}

	return plainText, nil
}
<commit_msg>Panic on header.Alg logic error.<commit_after>// Copyright 2014 Matthew Endsley
// All rights reserved
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted providing that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
// OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
// IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package gojwe

import (
	"bytes"
	"compress/flate"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/json"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"strings"
)

// Verify and decrypt a draft-7 JWE object
func VerifyAndDecryptDraft7(jwe string, key crypto.PrivateKey) ([]byte, error) {
	parts := strings.Split(jwe, ".")
	if len(parts) != 5 {
		return nil, errors.New("Wrong number of parts")
	}

	// decode the JWE header
	var header struct {
		Alg string `json:"alg"`
		Enc string `json:"enc"`
		Zip string `json:"zip"`
	}
	data, err := safeDecode(parts[0])
	if err != nil {
		return nil, fmt.Errorf("Malformed header: %v", err)
	}
	err = json.Unmarshal(data, &header)
	if err != nil {
		return nil, fmt.Errorf("Failed to decode header: %v", err)
	}

	var encryptionKey []byte
	encryptionKeyData, err := safeDecode(parts[1])
	if err != nil {
		return nil, fmt.Errorf("Malformed encryption key: %v", err)
	}

	// decode the encryption key
	switch header.Alg {
	case "RSA-OAEP", "RSA-OAEP-256":
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected an RSA private key. Got %T", key)
		}

		var h hash.Hash
		if header.Alg == "RSA-OAEP" {
			h = sha1.New()
		} else if header.Alg == "RSA-OAEP-256" {
			h = sha256.New()
		} else {
			panic("Logic error with algorithm " + header.Alg)
		}

		encryptionKey, err = rsa.DecryptOAEP(h, rand.Reader, rsaKey, encryptionKeyData, nil)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case "RSA1_5":
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected RSA private key. Got %T", key)
		}

		encryptionKey, err = rsa.DecryptPKCS1v15(rand.Reader, rsaKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case "A128KW", "A256KW":
		aesKey, ok := key.([]byte)
		if !ok {
			return nil, fmt.Errorf("Expected shared symmetric key ([]byte). Got %T", key)
		}

		encryptionKey, err = AesKeyUnwrap(aesKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to unwrap key: %v", err)
		}

	default:
		return nil, fmt.Errorf("Unsupported ALG keytype %s", header.Alg)
	}

	// decode IV
	iv, err := safeDecode(parts[2])
	if err != nil {
		return nil, fmt.Errorf("Malformed IV: %v", err)
	}

	// decode cipher text
	cipherText, err := safeDecode(parts[3])
	if err != nil {
		return nil, fmt.Errorf("Malformed cipher text: %v", err)
	}

	// decode authtag
	authTag, err := safeDecode(parts[4])
	if err != nil {
		return nil, fmt.Errorf("Malformed authtag: %v", err)
	}

	// decrypt and verify cipher text
	var plainText []byte

	switch header.Enc {
	case "A128CBC+HS256", "A256CBC+HS512":
		// derive keys
		var encKey, macKey []byte
		var hfunc func() hash.Hash
		if header.Enc == "A128CBC+HS256" {
			encKey, macKey = concatKDF(encryptionKey, header.Enc, 128, 256)
			hfunc = sha256.New
		} else if header.Enc == "A256CBC+HS512" {
			encKey, macKey = concatKDF(encryptionKey, header.Enc, 256, 512)
			hfunc = sha512.New
		} else {
			panic("Bad ENC logic for type: " + header.Enc)
		}

		// verify authtag
		hm := hmac.New(hfunc, macKey)
		io.WriteString(hm, parts[0])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[1])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[2])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[3])
		if !hmac.Equal(authTag, hm.Sum(nil)) {
			return nil, errors.New("Integrity check failed")
		}

		// decrpyt ciphertext (in-place)
		block, err := aes.NewCipher(encKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}

		c := cipher.NewCBCDecrypter(block, iv)
		c.CryptBlocks(cipherText, cipherText)
		plainText = cipherText

		// remove PCKS#7 padding
		padding := int(plainText[len(plainText)-1])
		plainText = plainText[:len(plainText)-padding]

	case "A128GCM", "A256GCM":
		// create the "additional data" for the GCM cipher
		additionalData := new(bytes.Buffer)
		additionalData.WriteString(parts[0])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[1])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[2])

		// create the authenticating cipher
		block, err := aes.NewCipher(encryptionKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}
		c, err := cipher.NewGCM(block)
		if err != nil {
			return nil, fmt.Errorf("Failed to create GCM cipher: %v", err)
		}

		// decrypt the cipher text (in-place)
		_, err = c.Open(cipherText[:0], iv, append(cipherText, authTag...), additionalData.Bytes())
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt: %v", err)
		}
		plainText = cipherText

	default:
		return nil, fmt.Errorf("Unsupported ENC keytype %s", header.Enc)
	}

	// need to deflate plain text?
	if header.Zip == "DEF" {
		plainText, err = ioutil.ReadAll(flate.NewReader(bytes.NewReader(plainText)))
		if err != nil {
			return nil, fmt.Errorf("Failed to inflate plain text: %v", err)
		}
	}

	return plainText, nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Matthew Endsley
// All rights reserved
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted providing that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
// OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
// IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package gojwe

import (
	"bytes"
	"compress/flate"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/json"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"strings"
)

// Verify and decrypt a draft-7 JWE object
func VerifyAndDecryptDraft7(jwe string, key crypto.PrivateKey) ([]byte, error) {
	parts := strings.Split(jwe, ".")
	if len(parts) != 5 {
		return nil, errors.New("Wrong number of parts")
	}

	// decode the JWE header
	var header struct {
		Alg string `json:"alg"`
		Enc string `json:"enc"`
		Zip string `json:"zip"`
	}
	data, err := safeDecode(parts[0])
	if err != nil {
		return nil, fmt.Errorf("Malformed header: %v", err)
	}
	err = json.Unmarshal(data, &header)
	if err != nil {
		return nil, fmt.Errorf("Failed to decode header: %v", err)
	}

	var encryptionKey []byte
	encryptionKeyData, err := safeDecode(parts[1])
	if err != nil {
		return nil, fmt.Errorf("Malformed encryption key: %v", err)
	}

	// decode the encryption key
	switch header.Alg {
	case "RSA-OAEP", "RSA-OAEP-256":
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected an RSA private key. Got %T", key)
		}

		var h hash.Hash
		if header.Alg == "RSA-OAEP" {
			h = sha1.New()
		} else if header.Alg == "RSA-OAEP-256" {
			h = sha256.New()
		} else {
			panic("Logic error with algorithm " + header.Alg)
		}

		encryptionKey, err = rsa.DecryptOAEP(h, rand.Reader, rsaKey, encryptionKeyData, nil)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case "RSA1_5":
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected RSA private key. Got %T", key)
		}

		encryptionKey, err = rsa.DecryptPKCS1v15(rand.Reader, rsaKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case "A128KW", "A256KW":
		aesKey, ok := key.([]byte)
		if !ok {
			return nil, fmt.Errorf("Expected shared symmetric key ([]byte). Got %T", key)
		}

		encryptionKey, err = AesKeyUnwrap(aesKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to unwrap key: %v", err)
		}

	default:
		return nil, fmt.Errorf("Unsupported ALG keytype %s", header.Alg)
	}

	// decode IV
	iv, err := safeDecode(parts[2])
	if err != nil {
		return nil, fmt.Errorf("Malformed IV: %v", err)
	}

	// decode cipher text
	cipherText, err := safeDecode(parts[3])
	if err != nil {
		return nil, fmt.Errorf("Malformed cipher text: %v", err)
	}

	// decode authtag
	authTag, err := safeDecode(parts[4])
	if err != nil {
		return nil, fmt.Errorf("Malformed authtag: %v", err)
	}

	// decrypt and verify cipher text
	var plainText []byte

	switch header.Enc {
	case "A128CBC+HS256", "A256CBC+HS512":
		// derive keys
		var encSize, macSize int
		var hfunc func() hash.Hash
		if header.Enc == "A128CBC+HS256" {
			encSize, macSize = 128, 256
			hfunc = sha256.New
		} else if header.Enc == "A256CBC+HS512" {
			encSize, macSize = 256, 512
			hfunc = sha512.New
		} else {
			panic("Bad ENC logic for type: " + header.Enc)
		}

		encKey, macKey := concatKDF(encryptionKey, header.Enc, encSize, macSize)

		// verify authtag
		hm := hmac.New(hfunc, macKey)
		io.WriteString(hm, parts[0])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[1])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[2])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[3])
		if !hmac.Equal(authTag, hm.Sum(nil)) {
			return nil, errors.New("Integrity check failed")
		}

		// decrpyt ciphertext (in-place)
		block, err := aes.NewCipher(encKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}

		c := cipher.NewCBCDecrypter(block, iv)
		c.CryptBlocks(cipherText, cipherText)
		plainText = cipherText

		// remove PCKS#7 padding
		padding := int(plainText[len(plainText)-1])
		plainText = plainText[:len(plainText)-padding]

	case "A128GCM", "A256GCM":
		// create the "additional data" for the GCM cipher
		additionalData := new(bytes.Buffer)
		additionalData.WriteString(parts[0])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[1])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[2])

		// create the authenticating cipher
		block, err := aes.NewCipher(encryptionKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}
		c, err := cipher.NewGCM(block)
		if err != nil {
			return nil, fmt.Errorf("Failed to create GCM cipher: %v", err)
		}

		// decrypt the cipher text (in-place)
		_, err = c.Open(cipherText[:0], iv, append(cipherText, authTag...), additionalData.Bytes())
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt: %v", err)
		}
		plainText = cipherText

	default:
		return nil, fmt.Errorf("Unsupported ENC keytype %s", header.Enc)
	}

	// need to deflate plain text?
	if header.Zip == "DEF" {
		plainText, err = ioutil.ReadAll(flate.NewReader(bytes.NewReader(plainText)))
		if err != nil {
			return nil, fmt.Errorf("Failed to inflate plain text: %v", err)
		}
	}

	return plainText, nil
}
<commit_msg>Replace algorithm/encryption string literals with constants<commit_after>// Copyright 2014 Matthew Endsley
// All rights reserved
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted providing that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
// IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
// OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
// IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package gojwe

import (
	"bytes"
	"compress/flate"
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/json"
	"errors"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"strings"
)

type Algorithm string

const (
	ALG_RSA_OAEP     = Algorithm("RSA-OAEP")
	ALG_RSA_OAEP_256 = Algorithm("RSA-OAEP-256")
	ALG_RSA1_5       = Algorithm("RSA1_5")
	ALG_A128KW       = Algorithm("A128KW")
	ALG_A256KW       = Algorithm("A256KW")
)

type EncryptionMethod string

const (
	ENC_A128CBC_HS256_v7 = EncryptionMethod("A128CBC+HS256")
	ENC_A256CBC_HS512_v7 = EncryptionMethod("A256CBC+H512")
	ENC_A128GCM          = EncryptionMethod("A128GCM")
	ENC_A256GCM          = EncryptionMethod("A256GCM")
)

// Verify and decrypt a draft-7 JWE object
func VerifyAndDecryptDraft7(jwe string, key crypto.PrivateKey) ([]byte, error) {
	parts := strings.Split(jwe, ".")
	if len(parts) != 5 {
		return nil, errors.New("Wrong number of parts")
	}

	// decode the JWE header
	var header struct {
		Alg Algorithm        `json:"alg"`
		Enc EncryptionMethod `json:"enc"`
		Zip string           `json:"zip"`
	}
	data, err := safeDecode(parts[0])
	if err != nil {
		return nil, fmt.Errorf("Malformed header: %v", err)
	}
	err = json.Unmarshal(data, &header)
	if err != nil {
		return nil, fmt.Errorf("Failed to decode header: %v", err)
	}

	var encryptionKey []byte
	encryptionKeyData, err := safeDecode(parts[1])
	if err != nil {
		return nil, fmt.Errorf("Malformed encryption key: %v", err)
	}

	// decode the encryption key
	switch header.Alg {
	case ALG_RSA_OAEP, ALG_RSA_OAEP_256:
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected an RSA private key. Got %T", key)
		}

		var h hash.Hash
		if header.Alg == ALG_RSA_OAEP {
			h = sha1.New()
		} else if header.Alg == ALG_RSA_OAEP_256 {
			h = sha256.New()
		} else {
			panic("Logic error with algorithm " + header.Alg)
		}

		encryptionKey, err = rsa.DecryptOAEP(h, rand.Reader, rsaKey, encryptionKeyData, nil)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case ALG_RSA1_5:
		rsaKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("Expected RSA private key. Got %T", key)
		}

		encryptionKey, err = rsa.DecryptPKCS1v15(rand.Reader, rsaKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt encryption key: %v", err)
		}

	case ALG_A128KW, ALG_A256KW:
		aesKey, ok := key.([]byte)
		if !ok {
			return nil, fmt.Errorf("Expected shared symmetric key ([]byte). Got %T", key)
		}

		encryptionKey, err = AesKeyUnwrap(aesKey, encryptionKeyData)
		if err != nil {
			return nil, fmt.Errorf("Failed to unwrap key: %v", err)
		}

	default:
		return nil, fmt.Errorf("Unsupported ALG keytype %s", header.Alg)
	}

	// decode IV
	iv, err := safeDecode(parts[2])
	if err != nil {
		return nil, fmt.Errorf("Malformed IV: %v", err)
	}

	// decode cipher text
	cipherText, err := safeDecode(parts[3])
	if err != nil {
		return nil, fmt.Errorf("Malformed cipher text: %v", err)
	}

	// decode authtag
	authTag, err := safeDecode(parts[4])
	if err != nil {
		return nil, fmt.Errorf("Malformed authtag: %v", err)
	}

	// decrypt and verify cipher text
	var plainText []byte

	switch header.Enc {
	case ENC_A128CBC_HS256_v7, ENC_A256CBC_HS512_v7:
		// derive keys
		var encSize, macSize int
		var hfunc func() hash.Hash
		if header.Enc == ENC_A128CBC_HS256_v7 {
			encSize, macSize = 128, 256
			hfunc = sha256.New
		} else if header.Enc == ENC_A256CBC_HS512_v7 {
			encSize, macSize = 256, 512
			hfunc = sha512.New
		} else {
			panic("Bad ENC logic for type: " + header.Enc)
		}

		encKey, macKey := concatKDF(encryptionKey, string(header.Enc), encSize, macSize)

		// verify authtag
		hm := hmac.New(hfunc, macKey)
		io.WriteString(hm, parts[0])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[1])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[2])
		io.WriteString(hm, ".")
		io.WriteString(hm, parts[3])
		if !hmac.Equal(authTag, hm.Sum(nil)) {
			return nil, errors.New("Integrity check failed")
		}

		// decrpyt ciphertext (in-place)
		block, err := aes.NewCipher(encKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}

		c := cipher.NewCBCDecrypter(block, iv)
		c.CryptBlocks(cipherText, cipherText)
		plainText = cipherText

		// remove PCKS#7 padding
		padding := int(plainText[len(plainText)-1])
		plainText = plainText[:len(plainText)-padding]

	case ENC_A128GCM, ENC_A256GCM:
		// create the "additional data" for the GCM cipher
		additionalData := new(bytes.Buffer)
		additionalData.WriteString(parts[0])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[1])
		additionalData.WriteRune('.')
		additionalData.WriteString(parts[2])

		// create the authenticating cipher
		block, err := aes.NewCipher(encryptionKey)
		if err != nil {
			return nil, fmt.Errorf("Failed to create an AES block cipher: %v", err)
		}
		c, err := cipher.NewGCM(block)
		if err != nil {
			return nil, fmt.Errorf("Failed to create GCM cipher: %v", err)
		}

		// decrypt the cipher text (in-place)
		_, err = c.Open(cipherText[:0], iv, append(cipherText, authTag...), additionalData.Bytes())
		if err != nil {
			return nil, fmt.Errorf("Failed to decrypt: %v", err)
		}
		plainText = cipherText

	default:
		return nil, fmt.Errorf("Unsupported ENC keytype %s", header.Enc)
	}

	// need to deflate plain text?
	if header.Zip == "DEF" {
		plainText, err = ioutil.ReadAll(flate.NewReader(bytes.NewReader(plainText)))
		if err != nil {
			return nil, fmt.Errorf("Failed to inflate plain text: %v", err)
		}
	}

	return plainText, nil
}
<|endoftext|>


<commit_before>// Package kml provides convenince methods for creating and writing KML documents.
package kml

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"image/color"
	"io"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	Header = xml.Header
	NS     = "http://www.opengis.net/kml/2.2"
	NS_GX  = "http://www.google.com/kml/ext/2.2"
)

// A Coordinate represents a single geographical coordinate.
// Lon and Lat are in degrees, Alt is in meters.
type Coordinate struct {
	Lon, Lat, Alt float64
}

// A Vec2 represents a screen position.
type Vec2 struct {
	X, Y           float64
	XUnits, YUnits string
}

// An Element represents an abstract KML element.
type Element interface {
	xml.Marshaler
	StringXML() (string, error)
	Write(io.Writer) error
}

// A SimpleElement is an Element with a single value.
type SimpleElement struct {
	xml.StartElement
	value string
}

// A CompoundElement is an Element with children.
type CompoundElement struct {
	xml.StartElement
	id       int
	children []Element
}

// Attr returns the XML attributes.
func (v2 Vec2) Attr() []xml.Attr {
	return []xml.Attr{
		{Name: xml.Name{Local: "x"}, Value: strconv.FormatFloat(v2.X, 'f', -1, 64)},
		{Name: xml.Name{Local: "y"}, Value: strconv.FormatFloat(v2.Y, 'f', -1, 64)},
		{Name: xml.Name{Local: "xunits"}, Value: v2.XUnits},
		{Name: xml.Name{Local: "yunits"}, Value: v2.YUnits},
	}
}

// MarshalXML marshals se to e. start is ignored.
func (se *SimpleElement) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	if err := e.EncodeToken(se.StartElement); err != nil {
		return err
	}
	if err := e.EncodeToken(xml.CharData(se.value)); err != nil {
		return err
	}
	endElement := xml.EndElement{Name: se.Name}
	if err := e.EncodeToken(endElement); err != nil {
		return err
	}
	return nil
}

// StringXML returns se encoded in XML.
func (se *SimpleElement) StringXML() (string, error) {
	return stringXML(se)
}

// Write writes an XML header and se to w.
func (se *SimpleElement) Write(w io.Writer) error {
	return write(w, se)
}

// Add adds children to ce.
func (ce *CompoundElement) Add(children ...Element) *CompoundElement {
	ce.children = append(ce.children, children...)
	return ce
}

// MarshalXML marshals ce to e. start is ignored.
func (ce *CompoundElement) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	if err := e.EncodeToken(ce.StartElement); err != nil {
		return err
	}
	for _, c := range ce.children {
		if err := e.EncodeElement(c, ce.StartElement); err != nil {
			return err
		}
	}
	endElement := xml.EndElement{Name: ce.Name}
	if err := e.EncodeToken(endElement); err != nil {
		return err
	}
	return nil
}

// StringXML returns ce encoded in XML.
func (ce *CompoundElement) StringXML() (string, error) {
	return stringXML(ce)
}

// Write writes an XML header and ce to w.
func (ce *CompoundElement) Write(w io.Writer) error {
	return write(w, ce)
}

func Altitude(value int) *SimpleElement                  { return newSEInt("altitude", value) }
func AltitudeMode(value string) *SimpleElement           { return newSEString("altitudeMode", value) }
func BalloonStyle(children ...Element) *CompoundElement  { return newCE("BalloonStyle", children) }
func Begin(value time.Time) *SimpleElement               { return newSETime("begin", value) }
func BgColor(value color.Color) *SimpleElement           { return newSEColor("bgColor", value) }
func Camera(children ...Element) *CompoundElement        { return newCE("Camera", children) }
func Color(value color.Color) *SimpleElement             { return newSEColor("color", value) }
func Data(children ...Element) *CompoundElement          { return newCE("Data", children) }
func Description(value string) *SimpleElement            { return newSEString("description", value) }
func Document(children ...Element) *CompoundElement      { return newCE("Document", children) }
func East(value float64) *SimpleElement                  { return newSEFloat("east", value) }
func End(value time.Time) *SimpleElement                 { return newSETime("end", value) }
func Extrude(value bool) *SimpleElement                  { return newSEBool("extrude", value) }
func Folder(children ...Element) *CompoundElement        { return newCE("Folder", children) }
func GroundOverlay(children ...Element) *CompoundElement { return newCE("GroundOverlay", children) }
func Heading(value float64) *SimpleElement               { return newSEFloat("heading", value) }
func Href(value *url.URL) *SimpleElement                 { return newSEString("href", value.String()) }
func HotSpot(value Vec2) *SimpleElement                  { return newSEPosition("hotSpot", value) }
func Icon(children ...Element) *CompoundElement          { return newCE("Icon", children) }
func IconStyle(children ...Element) *CompoundElement     { return newCE("IconStyle", children) }
func LabelStyle(children ...Element) *CompoundElement    { return newCE("LabelStyle", children) }
func LatLonBox(children ...Element) *CompoundElement     { return newCE("LatLonBox", children) }
func Latitude(value float64) *SimpleElement              { return newSEFloat("latitude", value) }
func LineString(children ...Element) *CompoundElement    { return newCE("LineString", children) }
func LineStyle(children ...Element) *CompoundElement     { return newCE("LineStyle", children) }
func ListItemType(value string) *SimpleElement           { return newSEString("listItemType", value) }
func ListStyle(children ...Element) *CompoundElement     { return newCE("ListStyle", children) }
func Longitude(value float64) *SimpleElement             { return newSEFloat("longitude", value) }
func MultiGeometry(children ...Element) *CompoundElement { return newCE("MultiGeometry", children) }
func Name(value string) *SimpleElement                   { return newSEString("name", value) }
func North(value float64) *SimpleElement                 { return newSEFloat("north", value) }
func Open(value bool) *SimpleElement                     { return newSEBool("open", value) }
func OverlayXY(value Vec2) *SimpleElement                { return newSEPosition("overlayXY", value) }
func Placemark(children ...Element) *CompoundElement     { return newCE("Placemark", children) }
func Point(children ...Element) *CompoundElement         { return newCE("Point", children) }
func PolyStyle(children ...Element) *CompoundElement     { return newCE("PolyStyle", children) }
func Roll(value float64) *SimpleElement                  { return newSEFloat("roll", value) }
func Rotation(value float64) *SimpleElement              { return newSEFloat("rotation", value) }
func Scale(value float64) *SimpleElement                 { return newSEFloat("scale", value) }
func ScreenOverlay(children ...Element) *CompoundElement { return newCE("ScreenOverlay", children) }
func ScreenXY(value Vec2) *SimpleElement                 { return newSEPosition("screenXY", value) }
func Snippet(value string) *SimpleElement                { return newSEString("snippet", value) }
func South(value float64) *SimpleElement                 { return newSEFloat("south", value) }
func Style(children ...Element) *CompoundElement         { return newCE("Style", children) }
func Tesselate(value bool) *SimpleElement                { return newSEBool("tesselate", value) }
func Text(value string) *SimpleElement                   { return newSEString("text", value) }
func Tilt(value float64) *SimpleElement                  { return newSEFloat("tilt", value) }
func TimeSpan(children ...Element) *CompoundElement      { return newCE("TimeSpan", children) }
func Value(value string) *SimpleElement                  { return newSEString("value", value) }
func Visibility(value bool) *SimpleElement               { return newSEBool("visibility", value) }
func West(value float64) *SimpleElement                  { return newSEFloat("west", value) }
func When(value time.Time) *SimpleElement                { return newSETime("time", value) }
func Width(value float64) *SimpleElement                 { return newSEFloat("width", value) }

func Coordinates(value ...Coordinate) *SimpleElement {
	cs := make([]string, len(value))
	for i, c := range value {
		cs[i] = strconv.FormatFloat(c.Lon, 'f', -1, 64) + "," +
			strconv.FormatFloat(c.Lat, 'f', -1, 64) + "," +
			strconv.FormatFloat(c.Alt, 'f', -1, 64)
	}
	return &SimpleElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Local: "coordinates"},
		},
		value: strings.Join(cs, " "),
	}
}

func HrefMustParse(value string) *SimpleElement {
	url, err := url.Parse(value)
	if err != nil {
		panic(err)
	}
	return Href(url)
}

func KML(children ...Element) *CompoundElement {
	return &CompoundElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Space: NS, Local: "kml"},
		},
		children: children,
	}
}

func stringXML(m xml.Marshaler) (string, error) {
	b := &bytes.Buffer{}
	e := xml.NewEncoder(b)
	if err := e.Encode(m); err != nil {
		return "", err
	}
	return b.String(), nil
}

func write(w io.Writer, m xml.Marshaler) error {
	if _, err := w.Write([]byte(Header)); err != nil {
		return err
	}
	e := xml.NewEncoder(w)
	if err := e.Encode(m); err != nil {
		return err
	}
	return nil
}

func newSEBool(name string, value bool) *SimpleElement {
	var v string
	if value {
		v = "1"
	} else {
		v = "0"
	}
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        v,
	}
}

func newSEColor(name string, value color.Color) *SimpleElement {
	r, g, b, a := value.RGBA()
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        fmt.Sprintf("%02x%02x%02x%02x", a/256, b/256, g/256, r/256),
	}
}

func newSEFloat(name string, value float64) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        strconv.FormatFloat(value, 'f', -1, 64),
	}
}

func newSEInt(name string, value int) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        strconv.Itoa(value),
	}
}

func newSEPosition(name string, value Vec2) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Local: name},
			Attr: value.Attr(),
		},
	}
}

func newSEString(name string, value string) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        value,
	}
}

func newSETime(name string, value time.Time) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        value.Format(time.RFC3339),
	}
}

func newSEVoid(name string) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
	}
}

func newCE(name string, children []Element) *CompoundElement {
	return &CompoundElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Local: name},
		},
		children: children,
	}
}
<commit_msg>Sort elements<commit_after>// Package kml provides convenince methods for creating and writing KML documents.
package kml

import (
	"bytes"
	"encoding/xml"
	"fmt"
	"image/color"
	"io"
	"net/url"
	"strconv"
	"strings"
	"time"
)

const (
	Header = xml.Header
	NS     = "http://www.opengis.net/kml/2.2"
	NS_GX  = "http://www.google.com/kml/ext/2.2"
)

// A Coordinate represents a single geographical coordinate.
// Lon and Lat are in degrees, Alt is in meters.
type Coordinate struct {
	Lon, Lat, Alt float64
}

// A Vec2 represents a screen position.
type Vec2 struct {
	X, Y           float64
	XUnits, YUnits string
}

// An Element represents an abstract KML element.
type Element interface {
	xml.Marshaler
	StringXML() (string, error)
	Write(io.Writer) error
}

// A SimpleElement is an Element with a single value.
type SimpleElement struct {
	xml.StartElement
	value string
}

// A CompoundElement is an Element with children.
type CompoundElement struct {
	xml.StartElement
	id       int
	children []Element
}

// Attr returns the XML attributes.
func (v2 Vec2) Attr() []xml.Attr {
	return []xml.Attr{
		{Name: xml.Name{Local: "x"}, Value: strconv.FormatFloat(v2.X, 'f', -1, 64)},
		{Name: xml.Name{Local: "y"}, Value: strconv.FormatFloat(v2.Y, 'f', -1, 64)},
		{Name: xml.Name{Local: "xunits"}, Value: v2.XUnits},
		{Name: xml.Name{Local: "yunits"}, Value: v2.YUnits},
	}
}

// MarshalXML marshals se to e. start is ignored.
func (se *SimpleElement) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	if err := e.EncodeToken(se.StartElement); err != nil {
		return err
	}
	if err := e.EncodeToken(xml.CharData(se.value)); err != nil {
		return err
	}
	endElement := xml.EndElement{Name: se.Name}
	if err := e.EncodeToken(endElement); err != nil {
		return err
	}
	return nil
}

// StringXML returns se encoded in XML.
func (se *SimpleElement) StringXML() (string, error) {
	return stringXML(se)
}

// Write writes an XML header and se to w.
func (se *SimpleElement) Write(w io.Writer) error {
	return write(w, se)
}

// Add adds children to ce.
func (ce *CompoundElement) Add(children ...Element) *CompoundElement {
	ce.children = append(ce.children, children...)
	return ce
}

// MarshalXML marshals ce to e. start is ignored.
func (ce *CompoundElement) MarshalXML(e *xml.Encoder, start xml.StartElement) error {
	if err := e.EncodeToken(ce.StartElement); err != nil {
		return err
	}
	for _, c := range ce.children {
		if err := e.EncodeElement(c, ce.StartElement); err != nil {
			return err
		}
	}
	endElement := xml.EndElement{Name: ce.Name}
	if err := e.EncodeToken(endElement); err != nil {
		return err
	}
	return nil
}

// StringXML returns ce encoded in XML.
func (ce *CompoundElement) StringXML() (string, error) {
	return stringXML(ce)
}

// Write writes an XML header and ce to w.
func (ce *CompoundElement) Write(w io.Writer) error {
	return write(w, ce)
}

func Altitude(value int) *SimpleElement                  { return newSEInt("altitude", value) }
func AltitudeMode(value string) *SimpleElement           { return newSEString("altitudeMode", value) }
func BalloonStyle(children ...Element) *CompoundElement  { return newCE("BalloonStyle", children) }
func Begin(value time.Time) *SimpleElement               { return newSETime("begin", value) }
func BgColor(value color.Color) *SimpleElement           { return newSEColor("bgColor", value) }
func Camera(children ...Element) *CompoundElement        { return newCE("Camera", children) }
func Color(value color.Color) *SimpleElement             { return newSEColor("color", value) }
func Data(children ...Element) *CompoundElement          { return newCE("Data", children) }
func Description(value string) *SimpleElement            { return newSEString("description", value) }
func Document(children ...Element) *CompoundElement      { return newCE("Document", children) }
func East(value float64) *SimpleElement                  { return newSEFloat("east", value) }
func End(value time.Time) *SimpleElement                 { return newSETime("end", value) }
func Extrude(value bool) *SimpleElement                  { return newSEBool("extrude", value) }
func Folder(children ...Element) *CompoundElement        { return newCE("Folder", children) }
func GroundOverlay(children ...Element) *CompoundElement { return newCE("GroundOverlay", children) }
func Heading(value float64) *SimpleElement               { return newSEFloat("heading", value) }
func HotSpot(value Vec2) *SimpleElement                  { return newSEPosition("hotSpot", value) }
func Href(value *url.URL) *SimpleElement                 { return newSEString("href", value.String()) }
func Icon(children ...Element) *CompoundElement          { return newCE("Icon", children) }
func IconStyle(children ...Element) *CompoundElement     { return newCE("IconStyle", children) }
func LabelStyle(children ...Element) *CompoundElement    { return newCE("LabelStyle", children) }
func LatLonBox(children ...Element) *CompoundElement     { return newCE("LatLonBox", children) }
func Latitude(value float64) *SimpleElement              { return newSEFloat("latitude", value) }
func LineString(children ...Element) *CompoundElement    { return newCE("LineString", children) }
func LineStyle(children ...Element) *CompoundElement     { return newCE("LineStyle", children) }
func ListItemType(value string) *SimpleElement           { return newSEString("listItemType", value) }
func ListStyle(children ...Element) *CompoundElement     { return newCE("ListStyle", children) }
func Longitude(value float64) *SimpleElement             { return newSEFloat("longitude", value) }
func MultiGeometry(children ...Element) *CompoundElement { return newCE("MultiGeometry", children) }
func Name(value string) *SimpleElement                   { return newSEString("name", value) }
func North(value float64) *SimpleElement                 { return newSEFloat("north", value) }
func Open(value bool) *SimpleElement                     { return newSEBool("open", value) }
func OverlayXY(value Vec2) *SimpleElement                { return newSEPosition("overlayXY", value) }
func Placemark(children ...Element) *CompoundElement     { return newCE("Placemark", children) }
func Point(children ...Element) *CompoundElement         { return newCE("Point", children) }
func PolyStyle(children ...Element) *CompoundElement     { return newCE("PolyStyle", children) }
func Roll(value float64) *SimpleElement                  { return newSEFloat("roll", value) }
func Rotation(value float64) *SimpleElement              { return newSEFloat("rotation", value) }
func Scale(value float64) *SimpleElement                 { return newSEFloat("scale", value) }
func ScreenOverlay(children ...Element) *CompoundElement { return newCE("ScreenOverlay", children) }
func ScreenXY(value Vec2) *SimpleElement                 { return newSEPosition("screenXY", value) }
func Snippet(value string) *SimpleElement                { return newSEString("snippet", value) }
func South(value float64) *SimpleElement                 { return newSEFloat("south", value) }
func Style(children ...Element) *CompoundElement         { return newCE("Style", children) }
func Tesselate(value bool) *SimpleElement                { return newSEBool("tesselate", value) }
func Text(value string) *SimpleElement                   { return newSEString("text", value) }
func Tilt(value float64) *SimpleElement                  { return newSEFloat("tilt", value) }
func TimeSpan(children ...Element) *CompoundElement      { return newCE("TimeSpan", children) }
func Value(value string) *SimpleElement                  { return newSEString("value", value) }
func Visibility(value bool) *SimpleElement               { return newSEBool("visibility", value) }
func West(value float64) *SimpleElement                  { return newSEFloat("west", value) }
func When(value time.Time) *SimpleElement                { return newSETime("time", value) }
func Width(value float64) *SimpleElement                 { return newSEFloat("width", value) }

func Coordinates(value ...Coordinate) *SimpleElement {
	cs := make([]string, len(value))
	for i, c := range value {
		cs[i] = strconv.FormatFloat(c.Lon, 'f', -1, 64) + "," +
			strconv.FormatFloat(c.Lat, 'f', -1, 64) + "," +
			strconv.FormatFloat(c.Alt, 'f', -1, 64)
	}
	return &SimpleElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Local: "coordinates"},
		},
		value: strings.Join(cs, " "),
	}
}

func HrefMustParse(value string) *SimpleElement {
	url, err := url.Parse(value)
	if err != nil {
		panic(err)
	}
	return Href(url)
}

func KML(children ...Element) *CompoundElement {
	return &CompoundElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Space: NS, Local: "kml"},
		},
		children: children,
	}
}

func stringXML(m xml.Marshaler) (string, error) {
	b := &bytes.Buffer{}
	e := xml.NewEncoder(b)
	if err := e.Encode(m); err != nil {
		return "", err
	}
	return b.String(), nil
}

func write(w io.Writer, m xml.Marshaler) error {
	if _, err := w.Write([]byte(Header)); err != nil {
		return err
	}
	e := xml.NewEncoder(w)
	if err := e.Encode(m); err != nil {
		return err
	}
	return nil
}

func newSEBool(name string, value bool) *SimpleElement {
	var v string
	if value {
		v = "1"
	} else {
		v = "0"
	}
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        v,
	}
}

func newSEColor(name string, value color.Color) *SimpleElement {
	r, g, b, a := value.RGBA()
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        fmt.Sprintf("%02x%02x%02x%02x", a/256, b/256, g/256, r/256),
	}
}

func newSEFloat(name string, value float64) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        strconv.FormatFloat(value, 'f', -1, 64),
	}
}

func newSEInt(name string, value int) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        strconv.Itoa(value),
	}
}

func newSEPosition(name string, value Vec2) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Local: name},
			Attr: value.Attr(),
		},
	}
}

func newSEString(name string, value string) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        value,
	}
}

func newSETime(name string, value time.Time) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
		value:        value.Format(time.RFC3339),
	}
}

func newSEVoid(name string) *SimpleElement {
	return &SimpleElement{
		StartElement: xml.StartElement{Name: xml.Name{Local: name}},
	}
}

func newCE(name string, children []Element) *CompoundElement {
	return &CompoundElement{
		StartElement: xml.StartElement{
			Name: xml.Name{Local: name},
		},
		children: children,
	}
}
<|endoftext|>