<commit_before>package archive

import (
	"fmt"
	"io"
	"path"
	"time"

	"crypto"
	"crypto/sha512"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/clearsign"
	"golang.org/x/crypto/openpgp/packet"

	"pault.ag/go/blobstore"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
	"pault.ag/go/debian/transput"
)

// Create a new Archive at the given `root` on the filesystem, with the
// openpgp.Entity `signer` (an Entity which contains an OpenPGP Private
// Key).
//
// This interface is intended to *write* Archives, not *read* them. Extra
// steps must be taken to load an Archive over the network, and attention
// must be paid when handling the Cryptographic chain of trust.
func New(path string, signer *openpgp.Entity) (*Archive, error) {
	store, err := blobstore.Load(path)
	if err != nil {
		return nil, err
	}

	return &Archive{
		store:      *store,
		signingKey: signer,
	}, nil
}

// Core Archive abstrcation. This contains helpers to write out package files,
// as well as handles creating underlying abstractions (such as Suites).
type Archive struct {
	store      blobstore.Store
	signingKey *openpgp.Entity
}

func (a Archive) Suite(name string) (*Suite, error) {
	suite := Suite{
		Name:       name,
		archive:    &a,
		components: map[string]Component{},
	}

	suite.Pool = Pool{store: a.store, suite: &suite}
	suite.features.Hashes = []string{"sha256", "sha1", "sha512"}
	suite.features.Duration = "240h"

	return &suite, nil
}

// Use the default backend to remove any unlinked files from the Blob store.
func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

func (a Archive) Link(blobs map[string]blobstore.Object) error {
	for path, obj := range blobs {
		if err := a.store.Link(obj, path); err != nil {
			return err
		}
	}
	return nil
}

func (suite Suite) newRelease() (*Release, error) {
	when := time.Now()

	var validUntil string = ""
	if suite.features.Duration != "" {
		duration, err := time.ParseDuration(suite.features.Duration)
		if err != nil {
			return nil, err
		}
		validUntil = when.Add(duration).Format(time.RFC1123Z)
	}

	release := Release{
		Suite:       suite.Name,
		Description: suite.Description,
		ValidUntil:  validUntil,
		Origin:      suite.Origin,
		Label:       suite.Label,
		Version:     suite.Version,
	}
	release.Date = when.Format(time.RFC1123Z)
	release.Architectures = []dependency.Arch{}
	release.Components = []string{}
	release.SHA256 = []control.SHA256FileHash{}
	release.SHA1 = []control.SHA1FileHash{}
	release.SHA512 = []control.SHA512FileHash{}
	release.MD5Sum = []control.MD5FileHash{}
	return &release, nil
}

//
func (a Archive) Engross(suite Suite) (map[string]blobstore.Object, error) {
	release, err := suite.newRelease()
	if err != nil {
		return nil, err
	}

	files := map[string]blobstore.Object{}
	arches := map[dependency.Arch]bool{}

	for name, component := range suite.components {
		release.Components = append(release.Components, name)

		for arch, writer := range component.packageWriters {
			arches[arch] = true

			suitePath := path.Join(name, fmt.Sprintf("binary-%s", arch),
				"Packages")

			obj, err := a.store.Commit(*writer.handle)
			if err != nil {
				return nil, err
			}

			for _, hasher := range writer.hashers {
				fileHash := control.FileHashFromHasher(suitePath, *hasher)
				release.AddHash(fileHash)
			}

			filePath := path.Join("dists", suite.Name, suitePath)
			files[filePath] = *obj
		}
	}

	for arch, _ := range arches {
		release.Architectures = append(release.Architectures, arch)
	}

	/* Now, let's do some magic */

	obj, sig, err := suite.archive.encodeSigned(release)
	if err != nil {
		return nil, err
	}

	filePath := path.Join("dists", suite.Name, "Release")
	files[filePath] = *obj
	files[fmt.Sprintf("%s.gpg", filePath)] = *sig

	obj, err = suite.archive.encodeClearsigned(release)
	if err != nil {
		return nil, err
	}

	files[path.Join("dists", suite.Name, "InRelease")] = *obj

	return files, nil
}

func (a Archive) encodeClearsigned(data interface{}) (*blobstore.Object, error) {

	if a.signingKey == nil {
		return nil, fmt.Errorf("No signing key loaded")
	}

	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	defer fd.Close()
	wc, err := clearsign.Encode(fd, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, err
	}

	encoder, err := control.NewEncoder(wc)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	if err := wc.Close(); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

func (a Archive) encodeSigned(data interface{}) (*blobstore.Object, *blobstore.Object, error) {
	/* Right, so, the trick here is that we secretly call out to encode,
	 * but tap it with a pipe into the signing code */

	if a.signingKey == nil {
		return nil, nil, fmt.Errorf("No signing key loaded")
	}

	signature, err := a.store.Create()
	if err != nil {
		return nil, nil, err
	}
	defer signature.Close()

	hash := sha512.New()

	obj, err := a.encode(data, hash)
	if err != nil {
		return nil, nil, err
	}

	sig := new(packet.Signature)
	sig.SigType = packet.SigTypeBinary
	sig.PubKeyAlgo = a.signingKey.PrivateKey.PubKeyAlgo

	sig.Hash = crypto.SHA512
	fmt.Printf("%x\n", hash.Sum(nil))

	sig.CreationTime = new(packet.Config).Now()
	sig.IssuerKeyId = &(a.signingKey.PrivateKey.KeyId)

	err = sig.Sign(hash, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, nil, err
	}

	if err := sig.Serialize(signature); err != nil {
		return nil, nil, err
	}

	sigObj, err := a.store.Commit(*signature)
	if err != nil {
		return nil, nil, err
	}

	return obj, sigObj, nil

}

func (a Archive) encode(data interface{}, tap io.Writer) (*blobstore.Object, error) {
	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	var writer io.Writer = fd
	if tap != nil {
		writer = io.MultiWriter(fd, tap)
	}

	encoder, err := control.NewEncoder(writer)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

//
type Suite struct {
	control.Paragraph

	archive *Archive

	Name        string `control:"Suite"`
	Description string
	Origin      string
	Label       string
	Version     string

	Pool       Pool                 `control:"-"`
	components map[string]Component `control"-"`

	features struct {
		Hashes   []string
		Duration string
	} `control:"-"`
}

/////////////////////////////////////////////////////////

func (s Suite) Component(name string) (*Component, error) {
	if _, ok := s.components[name]; !ok {
		comp, err := newComponent(&s)
		if err != nil {
			return nil, err
		}
		s.components[name] = *comp
		return comp, nil
	}
	el := s.components[name]
	return &el, nil
}

func newComponent(suite *Suite) (*Component, error) {
	return &Component{
		suite:          suite,
		packageWriters: map[dependency.Arch]*PackageWriter{},
	}, nil
}

type Component struct {
	suite          *Suite
	packageWriters map[dependency.Arch]*PackageWriter
	// sourceWriter *SourceWriter
}

func (c *Component) getWriter(arch dependency.Arch) (*PackageWriter, error) {
	if _, ok := c.packageWriters[arch]; !ok {
		writer, err := newPackageWriter(c.suite)
		if err != nil {
			return nil, err
		}
		c.packageWriters[arch] = writer
	}
	return c.packageWriters[arch], nil
}

func (c *Component) AddPackage(pkg Package) error {
	writer, err := c.getWriter(pkg.Architecture)
	if err != nil {
		return err
	}
	return writer.Add(pkg)
}

type packageWriters map[dependency.Arch]*PackageWriter

func newPackageWriter(suite *Suite) (*PackageWriter, error) {
	handle, err := suite.archive.store.Create()
	if err != nil {
		return nil, err
	}

	hashers := []*transput.Hasher{}
	writers := []io.Writer{handle}
	for _, algo := range suite.features.Hashes {
		hasher, err := transput.NewHasher(algo)
		if err != nil {
			handle.Close()
			return nil, err
		}
		hashers = append(hashers, hasher)
		writers = append(writers, hasher)
	}

	encoder, err := control.NewEncoder(io.MultiWriter(writers...))
	if err != nil {
		handle.Close()
		return nil, err
	}

	return &PackageWriter{
		archive: suite.archive,
		closer:  handle.Close,
		encoder: encoder,
		handle:  handle,
		hashers: hashers,
	}, nil
}

type PackageWriter struct {
	archive *Archive

	handle  *blobstore.Writer
	closer  func() error
	encoder *control.Encoder

	hashers []*transput.Hasher
}

func (p PackageWriter) Add(pkg Package) error {
	return p.encoder.Encode(pkg)
}

// vim: foldmethod=marker
<commit_msg>some docs<commit_after>package archive

import (
	"fmt"
	"io"
	"path"
	"time"

	"crypto"
	"crypto/sha512"

	"golang.org/x/crypto/openpgp"
	"golang.org/x/crypto/openpgp/clearsign"
	"golang.org/x/crypto/openpgp/packet"

	"pault.ag/go/blobstore"
	"pault.ag/go/debian/control"
	"pault.ag/go/debian/dependency"
	"pault.ag/go/debian/transput"
)

// Create a new Archive at the given `root` on the filesystem, with the
// openpgp.Entity `signer` (an Entity which contains an OpenPGP Private
// Key).
//
// This interface is intended to *write* Archives, not *read* them. Extra
// steps must be taken to load an Archive over the network, and attention
// must be paid when handling the Cryptographic chain of trust.
func New(path string, signer *openpgp.Entity) (*Archive, error) {
	store, err := blobstore.Load(path)
	if err != nil {
		return nil, err
	}

	return &Archive{
		store:      *store,
		signingKey: signer,
	}, nil
}

// Core Archive abstrcation. This contains helpers to write out package files,
// as well as handles creating underlying abstractions (such as Suites).
type Archive struct {
	store      blobstore.Store
	signingKey *openpgp.Entity
}

// Get a handle to write a given Suite from an Archive.
// The suite will be entirely blank, and attributes will not be
// read from the existing files, if any.
func (a Archive) Suite(name string) (*Suite, error) {
	suite := Suite{
		Name:       name,
		archive:    &a,
		components: map[string]Component{},
	}

	suite.Pool = Pool{store: a.store, suite: &suite}
	suite.features.Hashes = []string{"sha256", "sha1", "sha512"}
	suite.features.Duration = "240h"

	return &suite, nil
}

// Use the default backend to remove any unlinked files from the Blob store.
//
// If files you care about are not linked onto the stage, they will be removed
// by the garbage collector. Decruft only when you're sure the stage has been
// set.
func (a Archive) Decruft() error {
	return a.store.GC(blobstore.DumbGarbageCollector{})
}

// Given a list of objects, link them to the keyed paths.
func (a Archive) Link(blobs map[string]blobstore.Object) error {
	for path, obj := range blobs {
		if err := a.store.Link(obj, path); err != nil {
			return err
		}
	}
	return nil
}

// Create a new Release object from a Suite, passing off the Name, Description
// and constructing the rest of the goodies.
//
// This will be an entirely empty object, without anything read off disk.
func (suite Suite) newRelease() (*Release, error) {
	when := time.Now()

	var validUntil string = ""
	if suite.features.Duration != "" {
		duration, err := time.ParseDuration(suite.features.Duration)
		if err != nil {
			return nil, err
		}
		validUntil = when.Add(duration).Format(time.RFC1123Z)
	}

	release := Release{
		Suite:       suite.Name,
		Description: suite.Description,
		ValidUntil:  validUntil,
		Origin:      suite.Origin,
		Label:       suite.Label,
		Version:     suite.Version,
	}
	release.Date = when.Format(time.RFC1123Z)
	release.Architectures = []dependency.Arch{}
	release.Components = []string{}
	release.SHA256 = []control.SHA256FileHash{}
	release.SHA1 = []control.SHA1FileHash{}
	release.SHA512 = []control.SHA512FileHash{}
	release.MD5Sum = []control.MD5FileHash{}
	return &release, nil
}

// Engross a Suite for signing and final commit into the blobstore. This
// will return handle(s) to the signed and ready Objects, fit for passage
// to Link.
//
// This will contain all the related Packages and Release files.
func (a Archive) Engross(suite Suite) (map[string]blobstore.Object, error) {
	release, err := suite.newRelease()
	if err != nil {
		return nil, err
	}

	files := map[string]blobstore.Object{}
	arches := map[dependency.Arch]bool{}

	for name, component := range suite.components {
		release.Components = append(release.Components, name)

		for arch, writer := range component.packageWriters {
			arches[arch] = true

			suitePath := path.Join(name, fmt.Sprintf("binary-%s", arch),
				"Packages")

			obj, err := a.store.Commit(*writer.handle)
			if err != nil {
				return nil, err
			}

			for _, hasher := range writer.hashers {
				fileHash := control.FileHashFromHasher(suitePath, *hasher)
				release.AddHash(fileHash)
			}

			filePath := path.Join("dists", suite.Name, suitePath)
			files[filePath] = *obj
		}
	}

	for arch, _ := range arches {
		release.Architectures = append(release.Architectures, arch)
	}

	/* Now, let's do some magic */

	obj, sig, err := suite.archive.encodeSigned(release)
	if err != nil {
		return nil, err
	}

	filePath := path.Join("dists", suite.Name, "Release")
	files[filePath] = *obj
	files[fmt.Sprintf("%s.gpg", filePath)] = *sig

	obj, err = suite.archive.encodeClearsigned(release)
	if err != nil {
		return nil, err
	}

	files[path.Join("dists", suite.Name, "InRelease")] = *obj

	return files, nil
}

// Given a control.Marshal'able object, encode it to the blobstore, while
// also clearsigning the data.
func (a Archive) encodeClearsigned(data interface{}) (*blobstore.Object, error) {

	if a.signingKey == nil {
		return nil, fmt.Errorf("No signing key loaded")
	}

	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	defer fd.Close()
	wc, err := clearsign.Encode(fd, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, err
	}

	encoder, err := control.NewEncoder(wc)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	if err := wc.Close(); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

// Given a control.Marshal'able object, encode it to the blobstore, while
// also doing a detached OpenPGP signature. The objects returned (in order)
// are data, commited to the blobstore, the signature for that object, commited
// to the blobstore, and any error(s), finally.
func (a Archive) encodeSigned(data interface{}) (*blobstore.Object, *blobstore.Object, error) {
	/* Right, so, the trick here is that we secretly call out to encode,
	 * but tap it with a pipe into the signing code */

	if a.signingKey == nil {
		return nil, nil, fmt.Errorf("No signing key loaded")
	}

	signature, err := a.store.Create()
	if err != nil {
		return nil, nil, err
	}
	defer signature.Close()

	hash := sha512.New()

	obj, err := a.encode(data, hash)
	if err != nil {
		return nil, nil, err
	}

	sig := new(packet.Signature)
	sig.SigType = packet.SigTypeBinary
	sig.PubKeyAlgo = a.signingKey.PrivateKey.PubKeyAlgo

	sig.Hash = crypto.SHA512
	fmt.Printf("%x\n", hash.Sum(nil))

	sig.CreationTime = new(packet.Config).Now()
	sig.IssuerKeyId = &(a.signingKey.PrivateKey.KeyId)

	err = sig.Sign(hash, a.signingKey.PrivateKey, nil)
	if err != nil {
		return nil, nil, err
	}

	if err := sig.Serialize(signature); err != nil {
		return nil, nil, err
	}

	sigObj, err := a.store.Commit(*signature)
	if err != nil {
		return nil, nil, err
	}

	return obj, sigObj, nil

}

// Encode a given control.Marshal'able object into the Blobstore, and return
// a handle to its object.
//
// The optinal argument `tap` will be written to as the object gets sent into
// the blobstore. This may be useful if you wish to have a copy of the data
// going into the store.
func (a Archive) encode(data interface{}, tap io.Writer) (*blobstore.Object, error) {
	fd, err := a.store.Create()
	if err != nil {
		return nil, err
	}

	var writer io.Writer = fd
	if tap != nil {
		writer = io.MultiWriter(fd, tap)
	}

	encoder, err := control.NewEncoder(writer)
	if err != nil {
		return nil, err
	}

	if err := encoder.Encode(data); err != nil {
		return nil, err
	}

	return a.store.Commit(*fd)
}

// Abstraction to handle writing data into a Suite.
type Suite struct {
	control.Paragraph

	archive *Archive

	Name        string `control:"Suite"`
	Description string
	Origin      string
	Label       string
	Version     string

	Pool       Pool                 `control:"-"`
	components map[string]Component `control"-"`

	features struct {
		Hashes   []string
		Duration string
	} `control:"-"`
}

func (s Suite) Component(name string) (*Component, error) {
	if _, ok := s.components[name]; !ok {
		comp, err := newComponent(&s)
		if err != nil {
			return nil, err
		}
		s.components[name] = *comp
		return comp, nil
	}
	el := s.components[name]
	return &el, nil
}

func newComponent(suite *Suite) (*Component, error) {
	return &Component{
		suite:          suite,
		packageWriters: map[dependency.Arch]*PackageWriter{},
	}, nil
}

type Component struct {
	suite          *Suite
	packageWriters map[dependency.Arch]*PackageWriter
	// sourceWriter *SourceWriter
}

func (c *Component) getWriter(arch dependency.Arch) (*PackageWriter, error) {
	if _, ok := c.packageWriters[arch]; !ok {
		writer, err := newPackageWriter(c.suite)
		if err != nil {
			return nil, err
		}
		c.packageWriters[arch] = writer
	}
	return c.packageWriters[arch], nil
}

func (c *Component) AddPackage(pkg Package) error {
	writer, err := c.getWriter(pkg.Architecture)
	if err != nil {
		return err
	}
	return writer.Add(pkg)
}

type packageWriters map[dependency.Arch]*PackageWriter

func newPackageWriter(suite *Suite) (*PackageWriter, error) {
	handle, err := suite.archive.store.Create()
	if err != nil {
		return nil, err
	}

	hashers := []*transput.Hasher{}
	writers := []io.Writer{handle}
	for _, algo := range suite.features.Hashes {
		hasher, err := transput.NewHasher(algo)
		if err != nil {
			handle.Close()
			return nil, err
		}
		hashers = append(hashers, hasher)
		writers = append(writers, hasher)
	}

	encoder, err := control.NewEncoder(io.MultiWriter(writers...))
	if err != nil {
		handle.Close()
		return nil, err
	}

	return &PackageWriter{
		archive: suite.archive,
		closer:  handle.Close,
		encoder: encoder,
		handle:  handle,
		hashers: hashers,
	}, nil
}

type PackageWriter struct {
	archive *Archive

	handle  *blobstore.Writer
	closer  func() error
	encoder *control.Encoder

	hashers []*transput.Hasher
}

func (p PackageWriter) Add(pkg Package) error {
	return p.encoder.Encode(pkg)
}

// vim: foldmethod=marker
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mozilla-services/reaper/events"
	"github.com/mozilla-services/reaper/filters"
	"github.com/mozilla-services/reaper/reapable"
	log "github.com/mozilla-services/reaper/reaperlog"
)

var config *AWSConfig

type AWSConfig struct {
	Notifications events.NotificationsConfig
	HTTP          events.HTTPConfig
	Regions       []string
	WhitelistTag  string
	DefaultOwner  string
	DryRun        bool
}

func NewAWSConfig() *AWSConfig {
	return &AWSConfig{}
}

func SetAWSConfig(c *AWSConfig) {
	config = c
}

// convenience function that returns a map of instances in ASGs
func AllASGInstanceIds(as []AutoScalingGroup) map[reapable.Region]map[reapable.ID]bool {
	// maps region to id to bool
	inASG := make(map[reapable.Region]map[reapable.ID]bool)
	for _, region := range config.Regions {
		inASG[reapable.Region(region)] = make(map[reapable.ID]bool)
	}
	for _, a := range as {
		for _, instanceID := range a.Instances {
			// add the instance to the map
			inASG[a.Region][instanceID] = true
		}
	}
	return inASG
}

// AllAutoScalingGroups describes every AutoScalingGroup in the requested regions
// *AutoScalingGroups are created for every *autoscaling.AutoScalingGroup
// and are passed to a channel
func AllAutoScalingGroups() chan *AutoScalingGroup {
	ch := make(chan *AutoScalingGroup)
	wg := sync.WaitGroup{}

	for _, region := range config.Regions {
		go func(region string) {
			wg.Add(1)
			api := autoscaling.New(&aws.Config{Region: region})
			err := api.DescribeAutoScalingGroupsPages(&autoscaling.DescribeAutoScalingGroupsInput{}, func(resp *autoscaling.DescribeAutoScalingGroupsOutput, lastPage bool) bool {
				for _, asg := range resp.AutoScalingGroups {
					ch <- NewAutoScalingGroup(region, asg)
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		wg.Wait()
		close(ch)

	}()
	return ch
}

// AllInstances describes every instance in the requested regions
// *Instances are created for each *ec2.Instance
// and are passed to a channel
func AllInstances() chan *Instance {
	ch := make(chan *Instance)
	wg := sync.WaitGroup{}

	for _, region := range config.Regions {
		go func(region string) {
			wg.Add(1)
			api := ec2.New(&aws.Config{Region: region})
			// DescribeInstancesPages does autopagination
			err := api.DescribeInstancesPages(&ec2.DescribeInstancesInput{}, func(resp *ec2.DescribeInstancesOutput, lastPage bool) bool {
				for _, res := range resp.Reservations {
					for _, instance := range res.Instances {
						ch <- NewInstance(region, instance)
					}
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		wg.Wait()
		close(ch)
	}()
	return ch
}

func AllSnapshots() []filters.Filterable {
	regions := config.Regions

	// waitgroup for goroutines
	var wg sync.WaitGroup

	// channel for creating SecurityGroups
	in := make(chan *Snapshot)

	for _, region := range regions {
		wg.Add(1)

		sum := 0

		// goroutine per region to fetch all security groups
		go func(region string) {
			defer wg.Done()
			api := ec2.New(&aws.Config{Region: region})

			// TODO: nextToken paging
			input := &ec2.DescribeSnapshotsInput{}
			resp, err := api.DescribeSnapshots(input)
			if err != nil {
				// TODO: wee
			}

			for _, v := range resp.Snapshots {
				sum += 1
				in <- NewSnapshot(region, v)
			}
		}(region)
	}
	// aggregate
	var snapshots []filters.Filterable
	go func() {
		for s := range in {
			// Reapables[s.Region][s.ID] = s
			snapshots = append(snapshots, s)
		}
	}()

	// synchronous wait for all goroutines in wg to be done
	wg.Wait()

	// done with the channel
	close(in)

	log.Info("Found %d total snapshots.", len(snapshots))
	return snapshots
}
func AllVolumes() Volumes {
	regions := config.Regions

	// waitgroup for goroutines
	var wg sync.WaitGroup

	// channel for creating SecurityGroups
	in := make(chan *Volume)

	for _, region := range regions {
		wg.Add(1)

		sum := 0

		// goroutine per region to fetch all security groups
		go func(region string) {
			defer wg.Done()
			api := ec2.New(&aws.Config{Region: region})

			// TODO: nextToken paging
			input := &ec2.DescribeVolumesInput{}
			resp, err := api.DescribeVolumes(input)
			if err != nil {
				// TODO: wee
			}

			for _, v := range resp.Volumes {
				sum += 1
				in <- NewVolume(region, v)
			}

			log.Info(fmt.Sprintf("Found %d total volumes in %s", sum, region))
		}(region)
	}
	// aggregate
	var volumes Volumes
	go func() {
		for v := range in {
			// Reapables[v.Region][v.ID] = v
			volumes = append(volumes, v)
		}
	}()

	// synchronous wait for all goroutines in wg to be done
	wg.Wait()

	// done with the channel
	close(in)

	log.Info("Found %d total snapshots.", len(volumes))
	return volumes
}
func AllSecurityGroups() SecurityGroups {
	regions := config.Regions

	// waitgroup for goroutines
	var wg sync.WaitGroup

	// channel for creating SecurityGroups
	in := make(chan *SecurityGroup)

	for _, region := range regions {
		wg.Add(1)

		sum := 0

		// goroutine per region to fetch all security groups
		go func(region string) {
			defer wg.Done()
			api := ec2.New(&aws.Config{Region: region})

			// TODO: nextToken paging
			input := &ec2.DescribeSecurityGroupsInput{}
			resp, err := api.DescribeSecurityGroups(input)
			if err != nil {
				// TODO: wee
			}

			for _, sg := range resp.SecurityGroups {
				sum += 1
				in <- NewSecurityGroup(region, sg)
			}

			log.Info(fmt.Sprintf("Found %d total security groups in %s", sum, region))
		}(region)
	}
	// aggregate
	var securityGroups SecurityGroups
	go func() {
		for sg := range in {
			// Reapables[sg.Region][sg.ID] = sg
			securityGroups = append(securityGroups, sg)
		}
	}()

	// synchronous wait for all goroutines in wg to be done
	wg.Wait()

	// done with the channel
	close(in)

	log.Info("Found %d total security groups.", len(securityGroups))
	return securityGroups
}
<commit_msg>add comments<commit_after>package aws

import (
	"fmt"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mozilla-services/reaper/events"
	"github.com/mozilla-services/reaper/filters"
	"github.com/mozilla-services/reaper/reapable"
	log "github.com/mozilla-services/reaper/reaperlog"
)

var config *AWSConfig

type AWSConfig struct {
	Notifications events.NotificationsConfig
	HTTP          events.HTTPConfig
	Regions       []string
	WhitelistTag  string
	DefaultOwner  string
	DryRun        bool
}

func NewAWSConfig() *AWSConfig {
	return &AWSConfig{}
}

func SetAWSConfig(c *AWSConfig) {
	config = c
}

// convenience function that returns a map of instances in ASGs
func AllASGInstanceIds(as []AutoScalingGroup) map[reapable.Region]map[reapable.ID]bool {
	// maps region to id to bool
	inASG := make(map[reapable.Region]map[reapable.ID]bool)
	for _, region := range config.Regions {
		inASG[reapable.Region(region)] = make(map[reapable.ID]bool)
	}
	for _, a := range as {
		for _, instanceID := range a.Instances {
			// add the instance to the map
			inASG[a.Region][instanceID] = true
		}
	}
	return inASG
}

// AllAutoScalingGroups describes every AutoScalingGroup in the requested regions
// *AutoScalingGroups are created for every *autoscaling.AutoScalingGroup
// and are passed to a channel
func AllAutoScalingGroups() chan *AutoScalingGroup {
	ch := make(chan *AutoScalingGroup)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		go func(region string) {
			// add region to waitgroup
			wg.Add(1)
			api := autoscaling.New(&aws.Config{Region: region})
			err := api.DescribeAutoScalingGroupsPages(&autoscaling.DescribeAutoScalingGroupsInput{}, func(resp *autoscaling.DescribeAutoScalingGroupsOutput, lastPage bool) bool {
				for _, asg := range resp.AutoScalingGroups {
					ch <- NewAutoScalingGroup(region, asg)
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					// on the last page, finish this region
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)

	}()
	return ch
}

// AllInstances describes every instance in the requested regions
// *Instances are created for each *ec2.Instance
// and are passed to a channel
func AllInstances() chan *Instance {
	ch := make(chan *Instance)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		go func(region string) {
			// add region to waitgroup
			wg.Add(1)
			api := ec2.New(&aws.Config{Region: region})
			// DescribeInstancesPages does autopagination
			err := api.DescribeInstancesPages(&ec2.DescribeInstancesInput{}, func(resp *ec2.DescribeInstancesOutput, lastPage bool) bool {
				for _, res := range resp.Reservations {
					for _, instance := range res.Instances {
						ch <- NewInstance(region, instance)
					}
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)
	}()
	return ch
}

func AllSnapshots() []filters.Filterable {
	regions := config.Regions

	// waitgroup for goroutines
	var wg sync.WaitGroup

	// channel for creating SecurityGroups
	in := make(chan *Snapshot)

	for _, region := range regions {
		wg.Add(1)

		sum := 0

		// goroutine per region to fetch all security groups
		go func(region string) {
			defer wg.Done()
			api := ec2.New(&aws.Config{Region: region})

			// TODO: nextToken paging
			input := &ec2.DescribeSnapshotsInput{}
			resp, err := api.DescribeSnapshots(input)
			if err != nil {
				// TODO: wee
			}

			for _, v := range resp.Snapshots {
				sum += 1
				in <- NewSnapshot(region, v)
			}
		}(region)
	}
	// aggregate
	var snapshots []filters.Filterable
	go func() {
		for s := range in {
			// Reapables[s.Region][s.ID] = s
			snapshots = append(snapshots, s)
		}
	}()

	// synchronous wait for all goroutines in wg to be done
	wg.Wait()

	// done with the channel
	close(in)

	log.Info("Found %d total snapshots.", len(snapshots))
	return snapshots
}
func AllVolumes() Volumes {
	regions := config.Regions

	// waitgroup for goroutines
	var wg sync.WaitGroup

	// channel for creating SecurityGroups
	in := make(chan *Volume)

	for _, region := range regions {
		wg.Add(1)

		sum := 0

		// goroutine per region to fetch all security groups
		go func(region string) {
			defer wg.Done()
			api := ec2.New(&aws.Config{Region: region})

			// TODO: nextToken paging
			input := &ec2.DescribeVolumesInput{}
			resp, err := api.DescribeVolumes(input)
			if err != nil {
				// TODO: wee
			}

			for _, v := range resp.Volumes {
				sum += 1
				in <- NewVolume(region, v)
			}

			log.Info(fmt.Sprintf("Found %d total volumes in %s", sum, region))
		}(region)
	}
	// aggregate
	var volumes Volumes
	go func() {
		for v := range in {
			// Reapables[v.Region][v.ID] = v
			volumes = append(volumes, v)
		}
	}()

	// synchronous wait for all goroutines in wg to be done
	wg.Wait()

	// done with the channel
	close(in)

	log.Info("Found %d total snapshots.", len(volumes))
	return volumes
}
func AllSecurityGroups() SecurityGroups {
	regions := config.Regions

	// waitgroup for goroutines
	var wg sync.WaitGroup

	// channel for creating SecurityGroups
	in := make(chan *SecurityGroup)

	for _, region := range regions {
		wg.Add(1)

		sum := 0

		// goroutine per region to fetch all security groups
		go func(region string) {
			defer wg.Done()
			api := ec2.New(&aws.Config{Region: region})

			// TODO: nextToken paging
			input := &ec2.DescribeSecurityGroupsInput{}
			resp, err := api.DescribeSecurityGroups(input)
			if err != nil {
				// TODO: wee
			}

			for _, sg := range resp.SecurityGroups {
				sum += 1
				in <- NewSecurityGroup(region, sg)
			}

			log.Info(fmt.Sprintf("Found %d total security groups in %s", sum, region))
		}(region)
	}
	// aggregate
	var securityGroups SecurityGroups
	go func() {
		for sg := range in {
			// Reapables[sg.Region][sg.ID] = sg
			securityGroups = append(securityGroups, sg)
		}
	}()

	// synchronous wait for all goroutines in wg to be done
	wg.Wait()

	// done with the channel
	close(in)

	log.Info("Found %d total security groups.", len(securityGroups))
	return securityGroups
}
<|endoftext|>


<commit_before>package aws

import (
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mozilla-services/reaper/events"
	"github.com/mozilla-services/reaper/reapable"
	log "github.com/mozilla-services/reaper/reaperlog"
)

const (
	reaperTag           = "REAPER"
	reaperTagSeparator  = "|"
	reaperTagTimeFormat = "2006-01-02 03:04PM MST"
)

var config *AWSConfig

type AWSConfig struct {
	Notifications events.NotificationsConfig
	HTTP          events.HTTPConfig
	Regions       []string
	WhitelistTag  string
	DefaultOwner  string
	DryRun        bool
}

func NewAWSConfig() *AWSConfig {
	return &AWSConfig{}
}

func SetAWSConfig(c *AWSConfig) {
	config = c
}

// convenience function that returns a map of instances in ASGs
func AllASGInstanceIds(as []AutoScalingGroup) map[reapable.Region]map[reapable.ID]bool {
	// maps region to id to bool
	inASG := make(map[reapable.Region]map[reapable.ID]bool)
	for _, region := range config.Regions {
		inASG[reapable.Region(region)] = make(map[reapable.ID]bool)
	}
	for _, a := range as {
		for _, instanceID := range a.Instances {
			// add the instance to the map
			inASG[a.Region][instanceID] = true
		}
	}
	return inASG
}

func AllCloudformationStacks() chan *CloudformationStack {
	ch := make(chan *CloudformationStack)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		go func(region string) {
			// add region to waitgroup
			wg.Add(1)
			api := cloudformation.New(&aws.Config{Region: region})
			err := api.DescribeStacksPages(&cloudformation.DescribeStacksInput{}, func(resp *cloudformation.DescribeStacksOutput, lastPage bool) bool {
				for _, stack := range resp.Stacks {
					ch <- NewCloudformationStack(region, stack)
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					// on the last page, finish this region
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)

	}()
	return ch
}

// AllAutoScalingGroups describes every AutoScalingGroup in the requested regions
// *AutoScalingGroups are created for every *autoscaling.AutoScalingGroup
// and are passed to a channel
func AllAutoScalingGroups() chan *AutoScalingGroup {
	ch := make(chan *AutoScalingGroup)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		go func(region string) {
			// add region to waitgroup
			wg.Add(1)
			api := autoscaling.New(&aws.Config{Region: region})
			err := api.DescribeAutoScalingGroupsPages(&autoscaling.DescribeAutoScalingGroupsInput{}, func(resp *autoscaling.DescribeAutoScalingGroupsOutput, lastPage bool) bool {
				for _, asg := range resp.AutoScalingGroups {
					ch <- NewAutoScalingGroup(region, asg)
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					// on the last page, finish this region
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)

	}()
	return ch
}

// AllInstances describes every instance in the requested regions
// *Instances are created for each *ec2.Instance
// and are passed to a channel
func AllInstances() chan *Instance {
	ch := make(chan *Instance)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		go func(region string) {
			// add region to waitgroup
			wg.Add(1)
			api := ec2.New(&aws.Config{Region: region})
			// DescribeInstancesPages does autopagination
			err := api.DescribeInstancesPages(&ec2.DescribeInstancesInput{}, func(resp *ec2.DescribeInstancesOutput, lastPage bool) bool {
				for _, res := range resp.Reservations {
					for _, instance := range res.Instances {
						ch <- NewInstance(region, instance)
					}
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)
	}()
	return ch
}

func AllSecurityGroups() chan *SecurityGroup {
	ch := make(chan *SecurityGroup)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		go func(region string) {
			// add region to waitgroup
			wg.Add(1)
			api := ec2.New(&aws.Config{Region: region})
			resp, err := api.DescribeSecurityGroups(&ec2.DescribeSecurityGroupsInput{})
			for _, sg := range resp.SecurityGroups {
				ch <- NewSecurityGroup(region, sg)
			}
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
			wg.Done()
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)

	}()
	return ch
}
<commit_msg>remove race condition potential in All* resource channels<commit_after>package aws

import (
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/autoscaling"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mozilla-services/reaper/events"
	"github.com/mozilla-services/reaper/reapable"
	log "github.com/mozilla-services/reaper/reaperlog"
)

const (
	reaperTag           = "REAPER"
	reaperTagSeparator  = "|"
	reaperTagTimeFormat = "2006-01-02 03:04PM MST"
)

var config *AWSConfig

type AWSConfig struct {
	Notifications events.NotificationsConfig
	HTTP          events.HTTPConfig
	Regions       []string
	WhitelistTag  string
	DefaultOwner  string
	DryRun        bool
}

func NewAWSConfig() *AWSConfig {
	return &AWSConfig{}
}

func SetAWSConfig(c *AWSConfig) {
	config = c
}

// convenience function that returns a map of instances in ASGs
func AllASGInstanceIds(as []AutoScalingGroup) map[reapable.Region]map[reapable.ID]bool {
	// maps region to id to bool
	inASG := make(map[reapable.Region]map[reapable.ID]bool)
	for _, region := range config.Regions {
		inASG[reapable.Region(region)] = make(map[reapable.ID]bool)
	}
	for _, a := range as {
		for _, instanceID := range a.Instances {
			// add the instance to the map
			inASG[a.Region][instanceID] = true
		}
	}
	return inASG
}

func AllCloudformationStacks() chan *CloudformationStack {
	ch := make(chan *CloudformationStack)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		wg.Add(1)
		go func(region string) {
			// add region to waitgroup
			api := cloudformation.New(&aws.Config{Region: region})
			err := api.DescribeStacksPages(&cloudformation.DescribeStacksInput{}, func(resp *cloudformation.DescribeStacksOutput, lastPage bool) bool {
				for _, stack := range resp.Stacks {
					ch <- NewCloudformationStack(region, stack)
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					// on the last page, finish this region
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)

	}()
	return ch
}

// AllAutoScalingGroups describes every AutoScalingGroup in the requested regions
// *AutoScalingGroups are created for every *autoscaling.AutoScalingGroup
// and are passed to a channel
func AllAutoScalingGroups() chan *AutoScalingGroup {
	ch := make(chan *AutoScalingGroup)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		wg.Add(1)
		go func(region string) {
			// add region to waitgroup
			api := autoscaling.New(&aws.Config{Region: region})
			err := api.DescribeAutoScalingGroupsPages(&autoscaling.DescribeAutoScalingGroupsInput{}, func(resp *autoscaling.DescribeAutoScalingGroupsOutput, lastPage bool) bool {
				for _, asg := range resp.AutoScalingGroups {
					ch <- NewAutoScalingGroup(region, asg)
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					// on the last page, finish this region
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)

	}()
	return ch
}

// AllInstances describes every instance in the requested regions
// *Instances are created for each *ec2.Instance
// and are passed to a channel
func AllInstances() chan *Instance {
	ch := make(chan *Instance)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		wg.Add(1)
		go func(region string) {
			// add region to waitgroup
			api := ec2.New(&aws.Config{Region: region})
			// DescribeInstancesPages does autopagination
			err := api.DescribeInstancesPages(&ec2.DescribeInstancesInput{}, func(resp *ec2.DescribeInstancesOutput, lastPage bool) bool {
				for _, res := range resp.Reservations {
					for _, instance := range res.Instances {
						ch <- NewInstance(region, instance)
					}
				}
				// if we are at the last page, we should not continue
				// the return value of this func is "shouldContinue"
				if lastPage {
					wg.Done()
				}
				return true
			})
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)
	}()
	return ch
}

func AllSecurityGroups() chan *SecurityGroup {
	ch := make(chan *SecurityGroup)
	// waitgroup for all regions
	wg := sync.WaitGroup{}
	for _, region := range config.Regions {
		wg.Add(1)
		go func(region string) {
			// add region to waitgroup
			api := ec2.New(&aws.Config{Region: region})
			resp, err := api.DescribeSecurityGroups(&ec2.DescribeSecurityGroupsInput{})
			for _, sg := range resp.SecurityGroups {
				ch <- NewSecurityGroup(region, sg)
			}
			if err != nil {
				// probably should do something here...
				log.Error(err.Error())
			}
			wg.Done()
		}(region)
	}
	go func() {
		// in a separate goroutine, wait for all regions to finish
		// when they finish, close the chan
		wg.Wait()
		close(ch)

	}()
	return ch
}
<|endoftext|>


<commit_before>package axe

import (
	"fmt"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"

	"github.com/256dpi/fire/coal"
)

// Status defines the allowed statuses of a job.
type Status string

// The available job statuses.
const (
	// StatusEnqueued is used as the initial status when jobs are created.
	StatusEnqueued Status = "enqueued"

	// StatusDequeued is set when a job has been successfully dequeued.
	StatusDequeued Status = "dequeued"

	// StatusCompleted is set when a job has been successfully executed.
	StatusCompleted Status = "completed"

	// StatusFailed is set when an execution of a job failed.
	StatusFailed Status = "failed"

	// StatusCancelled is set when a job has been cancelled.
	StatusCancelled Status = "cancelled"
)

// Job is a single job managed by a queue.
type Job struct {
	coal.Base `json:"-" bson:",inline" coal:"jobs"`

	// The name of the job.
	Name string `json:"name" bson:"name"`

	// The custom label used to compute exclusiveness.
	Label string `json:"label" bson:"label"`

	// The data that has been supplied on creation.
	Data bson.Raw `json:"data" bson:"data"`

	// The current status of the job.
	Status Status `json:"status" bson:"status"`

	// The time when the job was created.
	Created time.Time `json:"created-at" bson:"created_at"`

	// The time when the job is available for execution.
	Available time.Time `json:"available-at" bson:"available_at"`

	// The time when the job was dequeue the last time.
	Started *time.Time `json:"started-at" bson:"started_at"`

	// The time when the last attempt ended (completed, failed or cancelled).
	Ended *time.Time `json:"ended-at" bson:"ended_at"`

	// The time when the job was finished (completed or cancelled).
	Finished *time.Time `json:"finished-at" bson:"finished_at"`

	// Attempts is incremented with each execution attempt.
	Attempts int `json:"attempts" bson:"attempts"`

	// The result submitted during completion.
	Result bson.M `json:"result" bson:"result"`

	// The last message submitted when the job was failed or cancelled.
	Reason string `json:"reason" bson:"reason"`
}

// AddJobIndexes will add job indexes to the specified indexer. If removeAfter
// is specified, completed and cancelled jobs are automatically removed when
// their finished timestamp falls behind the specified duration.
func AddJobIndexes(indexer *coal.Indexer, removeAfter time.Duration) {
	// add name index
	indexer.Add(&Job{}, false, 0, "Name")

	// add status index
	indexer.Add(&Job{}, false, 0, "Status")

	// add finished index
	indexer.Add(&Job{}, false, removeAfter, "Finished")
}

// Enqueue will enqueue a job using the specified name and data. If a delay
// is specified the job will not be dequeued until the specified time has passed.
// If exclusive is enabled the won't be queued if another job from this class
// is still available.
func Enqueue(store *coal.Store, session mongo.SessionContext, name, label string, model Model, delay time.Duration) (*Job, error) {
	// check name
	if name == "" {
		return nil, fmt.Errorf("missing name")
	}

	// set default model
	if model == nil {
		model = bson.M{}
	}

	// get time
	now := time.Now()

	// prepare job
	job := coal.Init(&Job{
		Name:      name,
		Label:     label,
		Status:    StatusEnqueued,
		Created:   now,
		Available: now.Add(delay),
	}).(*Job)

	// marshall model
	raw, err := bson.Marshal(model)
	if err != nil {
		return nil, err
	}

	// marshall into job
	err = bson.Unmarshal(raw, &job.Data)
	if err != nil {
		return nil, err
	}

	// check exclusiveness
	if label == "" {
		// insert job
		_, err = store.C(job).InsertOne(session, job)
		if err != nil {
			return nil, err
		}

		return job, nil
	}

	// insert job if there is no other job in an available state with the
	// provided label
	res, err := store.C(job).UpdateOne(session, bson.M{
		coal.F(&Job{}, "Name"):  name,
		coal.F(&Job{}, "Label"): label,
		coal.F(&Job{}, "Status"): bson.M{
			"$in": []Status{StatusEnqueued, StatusDequeued, StatusFailed},
		},
	}, bson.M{
		"$setOnInsert": job,
	}, options.Update().SetUpsert(true))
	if err != nil {
		return nil, err
	} else if res.UpsertedCount == 0 {
		return nil, nil
	}

	return job, nil
}

// Dequeue will dequeue the job with the specified id. The provided timeout
// will be set to allow the job to be dequeue if the process failed to set its
// status. Only jobs in the "enqueued", "dequeued" (passed timeout) or "failed"
// state are dequeued.
func Dequeue(store *coal.Store, id coal.ID, timeout time.Duration) (*Job, error) {
	// get time
	now := time.Now()

	// prepare options
	opts := options.FindOneAndUpdate().
		SetSort(coal.Sort("_id")).
		SetReturnDocument(options.After)

	// dequeue job
	var job Job
	err := store.C(&Job{}).FindOneAndUpdate(nil, bson.M{
		"_id": id,
		coal.F(&Job{}, "Status"): bson.M{
			"$in": []Status{StatusEnqueued, StatusDequeued, StatusFailed},
		},
		coal.F(&Job{}, "Available"): bson.M{
			"$lte": now,
		},
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):    StatusDequeued,
			coal.F(&Job{}, "Started"):   now,
			coal.F(&Job{}, "Available"): now.Add(timeout),
		},
		"$inc": bson.M{
			coal.F(&Job{}, "Attempts"): 1,
		},
	}, opts).Decode(&job)
	if err == mongo.ErrNoDocuments {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	return &job, nil
}

// Complete will complete the job with the specified id. Only jobs in the
// "dequeued" state can be completed.
func Complete(store *coal.Store, id coal.ID, result bson.M) error {
	// get time
	now := time.Now()

	// update job
	_, err := store.C(&Job{}).UpdateOne(nil, bson.M{
		"_id":                    id,
		coal.F(&Job{}, "Status"): StatusDequeued,
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):   StatusCompleted,
			coal.F(&Job{}, "Result"):   result,
			coal.F(&Job{}, "Ended"):    now,
			coal.F(&Job{}, "Finished"): now,
		},
	})
	if err != nil {
		return err
	}

	return nil
}

// Fail will fail the job with the specified id and the specified reason. It may
// delay the job if requested. Only jobs in the "dequeued" state can be failed.
func Fail(store *coal.Store, id coal.ID, reason string, delay time.Duration) error {
	// get time
	now := time.Now()

	// update job
	_, err := store.C(&Job{}).UpdateOne(nil, bson.M{
		coal.F(&Job{}, "Status"): StatusDequeued,
		"_id":                    id,
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):    StatusFailed,
			coal.F(&Job{}, "Reason"):    reason,
			coal.F(&Job{}, "Ended"):     now,
			coal.F(&Job{}, "Available"): now.Add(delay),
		},
	})
	if err != nil {
		return err
	}

	return nil
}

// Cancel will cancel the job with the specified id and the specified reason.
// Only jobs in the "dequeued" state can be cancelled.
func Cancel(store *coal.Store, id coal.ID, reason string) error {
	// get time
	now := time.Now()

	// update job
	_, err := store.C(&Job{}).UpdateOne(nil, bson.M{
		"_id":                    id,
		coal.F(&Job{}, "Status"): StatusDequeued,
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):   StatusCancelled,
			coal.F(&Job{}, "Reason"):   reason,
			coal.F(&Job{}, "Ended"):    now,
			coal.F(&Job{}, "Finished"): now,
		},
	})
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>cleanup<commit_after>package axe

import (
	"fmt"
	"time"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"

	"github.com/256dpi/fire/coal"
)

// Status defines the allowed statuses of a job.
type Status string

// The available job statuses.
const (
	// StatusEnqueued is used as the initial status when jobs are created.
	StatusEnqueued Status = "enqueued"

	// StatusDequeued is set when a job has been successfully dequeued.
	StatusDequeued Status = "dequeued"

	// StatusCompleted is set when a job has been successfully executed.
	StatusCompleted Status = "completed"

	// StatusFailed is set when an execution of a job failed.
	StatusFailed Status = "failed"

	// StatusCancelled is set when a job has been cancelled.
	StatusCancelled Status = "cancelled"
)

// Job is a single job managed by a queue.
type Job struct {
	coal.Base `json:"-" bson:",inline" coal:"jobs"`

	// The name of the job.
	Name string `json:"name"`

	// The custom label used to compute exclusiveness.
	Label string `json:"label"`

	// The data that has been supplied on creation.
	Data bson.Raw `json:"data"`

	// The current status of the job.
	Status Status `json:"status"`

	// The time when the job was created.
	Created time.Time `json:"created-at" bson:"created_at"`

	// The time when the job is available for execution.
	Available time.Time `json:"available-at" bson:"available_at"`

	// The time when the job was dequeue the last time.
	Started *time.Time `json:"started-at" bson:"started_at"`

	// The time when the last attempt ended (completed, failed or cancelled).
	Ended *time.Time `json:"ended-at" bson:"ended_at"`

	// The time when the job was finished (completed or cancelled).
	Finished *time.Time `json:"finished-at" bson:"finished_at"`

	// Attempts is incremented with each execution attempt.
	Attempts int `json:"attempts"`

	// The result submitted during completion.
	Result bson.M `json:"result"`

	// The last message submitted when the job was failed or cancelled.
	Reason string `json:"reason"`
}

// AddJobIndexes will add job indexes to the specified indexer. If a duration
// is specified, completed and cancelled jobs are automatically removed when
// their finished timestamp falls behind the specified duration.
func AddJobIndexes(indexer *coal.Indexer, removeAfter time.Duration) {
	// add name index
	indexer.Add(&Job{}, false, 0, "Name")

	// add status index
	indexer.Add(&Job{}, false, 0, "Status")

	// add finished index
	indexer.Add(&Job{}, false, removeAfter, "Finished")
}

// Enqueue will enqueue a job using the specified name and data. If a label is
// given, it will only enqueue the task if no other task is available with the
// same name and label. If a delay is specified the job will not be dequeued
// until the specified time has passed.
func Enqueue(store *coal.Store, session mongo.SessionContext, name, label string, model Model, delay time.Duration) (*Job, error) {
	// check name
	if name == "" {
		return nil, fmt.Errorf("missing name")
	}

	// set default model
	if model == nil {
		model = bson.M{}
	}

	// get time
	now := time.Now()

	// prepare job
	job := coal.Init(&Job{
		Name:      name,
		Label:     label,
		Status:    StatusEnqueued,
		Created:   now,
		Available: now.Add(delay),
	}).(*Job)

	// marshall model
	raw, err := bson.Marshal(model)
	if err != nil {
		return nil, err
	}

	// marshall into job
	err = bson.Unmarshal(raw, &job.Data)
	if err != nil {
		return nil, err
	}

	// check exclusiveness
	if label == "" {
		// insert job
		_, err = store.C(job).InsertOne(session, job)
		if err != nil {
			return nil, err
		}

		return job, nil
	}

	// insert job if there is no other job in an available state with the
	// provided label
	res, err := store.C(job).UpdateOne(session, bson.M{
		coal.F(&Job{}, "Name"):  name,
		coal.F(&Job{}, "Label"): label,
		coal.F(&Job{}, "Status"): bson.M{
			"$in": []Status{StatusEnqueued, StatusDequeued, StatusFailed},
		},
	}, bson.M{
		"$setOnInsert": job,
	}, options.Update().SetUpsert(true))
	if err != nil {
		return nil, err
	} else if res.UpsertedCount == 0 {
		return nil, nil
	}

	return job, nil
}

// Dequeue will dequeue the job with the specified id. The provided timeout
// will be set to allow the job to be dequeue if the process failed to set its
// status. Only jobs in the "enqueued", "dequeued" (passed timeout) or "failed"
// state are dequeued.
func Dequeue(store *coal.Store, id coal.ID, timeout time.Duration) (*Job, error) {
	// get time
	now := time.Now()

	// prepare options
	opts := options.FindOneAndUpdate().
		SetSort(coal.Sort("_id")).
		SetReturnDocument(options.After)

	// dequeue job
	var job Job
	err := store.C(&Job{}).FindOneAndUpdate(nil, bson.M{
		"_id": id,
		coal.F(&Job{}, "Status"): bson.M{
			"$in": []Status{StatusEnqueued, StatusDequeued, StatusFailed},
		},
		coal.F(&Job{}, "Available"): bson.M{
			"$lte": now,
		},
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):    StatusDequeued,
			coal.F(&Job{}, "Started"):   now,
			coal.F(&Job{}, "Available"): now.Add(timeout),
		},
		"$inc": bson.M{
			coal.F(&Job{}, "Attempts"): 1,
		},
	}, opts).Decode(&job)
	if err == mongo.ErrNoDocuments {
		return nil, nil
	} else if err != nil {
		return nil, err
	}

	return &job, nil
}

// Complete will complete the job with the specified id. Only jobs in the
// "dequeued" state can be completed.
func Complete(store *coal.Store, id coal.ID, result bson.M) error {
	// get time
	now := time.Now()

	// update job
	_, err := store.C(&Job{}).UpdateOne(nil, bson.M{
		"_id":                    id,
		coal.F(&Job{}, "Status"): StatusDequeued,
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):   StatusCompleted,
			coal.F(&Job{}, "Result"):   result,
			coal.F(&Job{}, "Ended"):    now,
			coal.F(&Job{}, "Finished"): now,
		},
	})
	if err != nil {
		return err
	}

	return nil
}

// Fail will fail the job with the specified id and the specified reason. It may
// delay the job if requested. Only jobs in the "dequeued" state can be failed.
func Fail(store *coal.Store, id coal.ID, reason string, delay time.Duration) error {
	// get time
	now := time.Now()

	// update job
	_, err := store.C(&Job{}).UpdateOne(nil, bson.M{
		coal.F(&Job{}, "Status"): StatusDequeued,
		"_id":                    id,
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):    StatusFailed,
			coal.F(&Job{}, "Reason"):    reason,
			coal.F(&Job{}, "Ended"):     now,
			coal.F(&Job{}, "Available"): now.Add(delay),
		},
	})
	if err != nil {
		return err
	}

	return nil
}

// Cancel will cancel the job with the specified id and the specified reason.
// Only jobs in the "dequeued" state can be cancelled.
func Cancel(store *coal.Store, id coal.ID, reason string) error {
	// get time
	now := time.Now()

	// update job
	_, err := store.C(&Job{}).UpdateOne(nil, bson.M{
		"_id":                    id,
		coal.F(&Job{}, "Status"): StatusDequeued,
	}, bson.M{
		"$set": bson.M{
			coal.F(&Job{}, "Status"):   StatusCancelled,
			coal.F(&Job{}, "Reason"):   reason,
			coal.F(&Job{}, "Ended"):    now,
			coal.F(&Job{}, "Finished"): now,
		},
	})
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/FederationOfFathers/dashboard/environment"
	"github.com/nlopes/slack"
	"go.uber.org/zap"
)

var isImage = regexp.MustCompile("\\.(jpe?g|gif|png)$")
var CdnPath = ""
var CdnPrefix = ""

func fileBytesQuietly(f *slack.File) ([]byte, error) {
	req, _ := http.NewRequest(
		"GET",
		f.URLPrivate,
		nil)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	rsp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer rsp.Body.Close()
	if rsp.StatusCode != 200 {
		buf, _ := httputil.DumpRequest(req, true)
		bufRsp, _ := httputil.DumpResponse(rsp, false)
		Logger.Info("Debugging File Request",
			zap.String("method", "old"),
			zap.ByteString("request", buf),
			zap.ByteString("response", bufRsp),
		)
		return nil, fmt.Errorf(rsp.Status)
	}
	if strings.Contains(rsp.Header.Get("Content-Type"), "text/html") {
		buf, _ := httputil.DumpRequest(req, true)
		bufRsp, _ := httputil.DumpResponse(rsp, false)
		Logger.Info("Debugging File Request",
			zap.String("method", "old"),
			zap.ByteString("request", buf),
			zap.ByteString("response", bufRsp),
		)
		return nil, fmt.Errorf("Expected non html content type, got %s", rsp.Header.Get("Content-Type"))
	}
	return ioutil.ReadAll(rsp.Body)
}

func fileBytesNoisy(f *slack.File) ([]byte, error) {
	for i := 0; i < 5; i++ {
		if _, _, _, err := rtm.ShareFilePublicURL(f.ID); err == nil {
			break
		} else {
			if i < 4 {
				sleepfor := time.Millisecond * time.Duration(((i+1)*(i+1))*100)
				Logger.Error(
					"Error making file public",
					zap.Error(err),
					zap.String("filename", f.Name),
					zap.String("fileid", f.ID),
					zap.Duration("sleepfor", sleepfor))
			} else {
				Logger.Error(
					"Error making file public: %s/%s: %s",
					zap.Error(err),
					zap.String("filename", f.Name),
					zap.String("fileid", f.ID))
				return nil, fmt.Errorf("Error making file public: %s/%s: %s", f.ID, f.Name, err.Error())
			}
		}

	}
	var pubSecret string
	pubParts := strings.Split(f.PermalinkPublic, "-")
	pubSecret = pubParts[len(pubParts)-1]
	req, _ := http.NewRequest(
		"GET",
		fmt.Sprintf(
			"%s?pub_secret=%s",
			f.URLPrivate,
			pubSecret,
		),
		nil)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	buf, _ := httputil.DumpRequest(req, true)
	Logger.Info("Debugging File Request", zap.ByteString("request", buf))
	rsp, err := http.DefaultClient.Do(req)
	bufRsp, _ := httputil.DumpResponse(rsp, false)
	Logger.Info("Debugging File Response", zap.ByteString("response", bufRsp))
	if err != nil {
		return nil, err
	}
	defer rsp.Body.Close()
	if rsp.StatusCode != 200 {
		return nil, fmt.Errorf(rsp.Status)
	}
	if strings.Contains(rsp.Header.Get("Content-Type"), "text/html") {
		return nil, fmt.Errorf("Expected non html content type, got %s", rsp.Header.Get("Content-Type"))
	}
	return ioutil.ReadAll(rsp.Body)
}

func fileBytes(f *slack.File) ([]byte, error) {
	for i := 0; i < 5; i++ {
		if buf, err := fileBytesQuietly(f); err == nil {
			return buf, err
		}
		if i < 4 {
			sleepfor := time.Millisecond * time.Duration(((i+1)*(i+1))*100)
			time.Sleep(sleepfor)
		}

	}
	return fileBytesNoisy(f)
}

func handleChannelUpload(m *slack.MessageEvent) bool {
	if CdnPath == "" || CdnPrefix == "" {
		return false
	}
	if !m.Msg.Upload {
		return false
	}
	// don't handle upload messages if dev, since dev and prod both manage slack
	if environment.IsDev {
		return false
	}
	file := &m.Files[0]

	// get the username by the ID from the DB
	var user = ":space_invader:"
	profile, err := data.User(m.User)
	if err != nil {
		Logger.Error("Could not get name for user",
			zap.String("userID", m.User),
			zap.String("filename", file.Name),
			zap.Error(err))
	} else {
		if profile.Profile.DisplayName != "" {
			user = profile.Profile.DisplayName
		} else if profile.Profile.RealName != "" {
			user = profile.Profile.RealName
		} else {
			user = profile.Name
		}
	}

	Logger.Info("File upload detected", zap.String("username", user), zap.String("filename", file.Name))
	if buf, err := fileBytes(file); err != nil {
		Logger.Error(
			"error downloading file",
			zap.Error(err),
			zap.String("username", user),
			zap.String("filename", file.Name))
	} else {
		path := fmt.Sprintf("%s/%s", CdnPath, time.Now().Format("2006/01/02/15"))
		if err := os.MkdirAll(path, 0755); err != nil {
			Logger.Error(
				"error making cdn path",
				zap.String("path", path),
				zap.String("username", user),
				zap.String("filename", file.Name))
			return false
		}
		part := &url.URL{Path: file.Name}
		urlPath := fmt.Sprintf("%s/%s-%s", path, file.ID, part.String())
		path = fmt.Sprintf("%s/%s-%s", path, file.ID, file.Name)
		if fp, err := os.Create(path); err != nil {
			Logger.Error(
				"error creating cdn file",
				zap.String("path", path),
				zap.String("username", user),
				zap.String("filename", file.Name))
			return false
		} else {
			if _, err := fp.Write(buf); err != nil {
				fp.Close()
				Logger.Error(
					"error writing to cdn file",
					zap.String("path", path),
					zap.String("username", user),
					zap.String("filename", file.Name))
				return false
			}
			fp.Close()
			fileURL := CdnPrefix + urlPath[len(CdnPath):]

			//delete file and message
			if err := rtm.DeleteFile(file.ID); err != nil {
				Logger.Error(fmt.Sprintf("unable to delete file - %s", err.Error()),
					zap.String("file", file.ID))
			}
			if _, _, err := rtm.DeleteMessage(m.Channel, m.Timestamp); err != nil {
				Logger.Error(fmt.Sprintf("unable to delete message - %s", err.Error()),
					zap.String("username", user),
					zap.String("filename", file.Name),
					zap.Error(err))
			}

			if isImage.MatchString(strings.ToLower(file.Name)) {
				for i := 0; i < 5; i++ {
					_, _, err := rtm.PostMessage(
						m.Channel,
						"",
						slack.PostMessageParameters{
							AsUser:      true,
							UnfurlLinks: true,
							UnfurlMedia: true,
							IconEmoji:   ":paperclip:",
							Attachments: []slack.Attachment{
								slack.Attachment{
									AuthorName: fmt.Sprintf("%s uploaded %s", user, file.Name),
									Pretext:    m.Text,
									ImageURL:   fileURL,
								},
							},
							ThreadTimestamp: m.ThreadTimestamp,
						})
					if err != nil {
						Logger.Error(
							"Failed postting cdn link back to slack",
							zap.String("username", user),
							zap.String("filename", file.Name),
							zap.String("url", fileURL))
					} else {
						break
					}
					time.Sleep(time.Second * time.Duration(i))
				}
			} else {
				rtm.SendMessage(&slack.OutgoingMessage{
					ID:              int(time.Now().UnixNano()),
					Channel:         m.Channel,
					Text:            fmt.Sprintf("%s uploaded the file *%s*\n%s", user, file.Title, fileURL),
					Type:            "message",
					ThreadTimestamp: m.ThreadTimestamp,
				})
			}
			Logger.Info("saved CDN file", zap.String("url", fileURL), zap.Int("size", len(buf)))
			return true
		}

	}
	return false
}

func handleDMUpload(m *slack.MessageEvent) bool {
	if CdnPath == "" || CdnPrefix == "" {
		return false
	}
	if !m.Msg.Upload {
		return false
	}
	if buf, err := fileBytes(&m.Files[0]); err != nil {
		Logger.Info("error downloading file", zap.Error(err))
	} else {
		path := fmt.Sprintf("%s/%s", CdnPath, time.Now().Format("2006/01/02/15"))
		if err := os.MkdirAll(path, 0755); err != nil {
			Logger.Error("error making cdn path", zap.String("path", path))
			return false
		}
		part := &url.URL{Path: m.Files[0].Name}
		urlPath := fmt.Sprintf("%s/%s-%s", path, m.Files[0].ID, part.String())
		path = fmt.Sprintf("%s/%s-%s", path, m.Files[0].ID, m.Files[0].Name)
		if fp, err := os.Create(path); err != nil {
			Logger.Error("error creating cdn file", zap.String("path", path))
			return false
		} else {
			if _, err := fp.Write(buf); err != nil {
				fp.Close()
				Logger.Error("error writing to cdn file", zap.String("path", path))
				return false
			}
			fp.Close()
			fileURL := CdnPrefix + urlPath[len(CdnPath):]
			rtm.DeleteFile(m.Files[0].ID)
			rtm.SendMessage(&slack.OutgoingMessage{
				ID:      int(time.Now().UnixNano()),
				Channel: m.Channel,
				Text:    fmt.Sprintf("Thanks for sending me the file instead of uploading it to a channel or group. You can paste the following link anywhere you want to show the file to others! ```%s```", fileURL),
				Type:    "message",
			})
			Logger.Info("saved CDN file", zap.String("url", fileURL), zap.Int("size", len(buf)))
		}

	}
	return true
}
<commit_msg>Ignore message_not_found errors, because the message is gone<commit_after>package bot

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/nlopes/slack"
	"go.uber.org/zap"
)

var isImage = regexp.MustCompile("\\.(jpe?g|gif|png)$")
var CdnPath = ""
var CdnPrefix = ""

func fileBytesQuietly(f *slack.File) ([]byte, error) {
	req, _ := http.NewRequest(
		"GET",
		f.URLPrivate,
		nil)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	rsp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer rsp.Body.Close()
	if rsp.StatusCode != 200 {
		buf, _ := httputil.DumpRequest(req, true)
		bufRsp, _ := httputil.DumpResponse(rsp, false)
		Logger.Info("Debugging File Request",
			zap.String("method", "old"),
			zap.ByteString("request", buf),
			zap.ByteString("response", bufRsp),
		)
		return nil, fmt.Errorf(rsp.Status)
	}
	if strings.Contains(rsp.Header.Get("Content-Type"), "text/html") {
		buf, _ := httputil.DumpRequest(req, true)
		bufRsp, _ := httputil.DumpResponse(rsp, false)
		Logger.Info("Debugging File Request",
			zap.String("method", "old"),
			zap.ByteString("request", buf),
			zap.ByteString("response", bufRsp),
		)
		return nil, fmt.Errorf("Expected non html content type, got %s", rsp.Header.Get("Content-Type"))
	}
	return ioutil.ReadAll(rsp.Body)
}

func fileBytesNoisy(f *slack.File) ([]byte, error) {
	for i := 0; i < 5; i++ {
		if _, _, _, err := rtm.ShareFilePublicURL(f.ID); err == nil {
			break
		} else {
			if i < 4 {
				sleepfor := time.Millisecond * time.Duration(((i+1)*(i+1))*100)
				Logger.Error(
					"Error making file public",
					zap.Error(err),
					zap.String("filename", f.Name),
					zap.String("fileid", f.ID),
					zap.Duration("sleepfor", sleepfor))
			} else {
				Logger.Error(
					"Error making file public: %s/%s: %s",
					zap.Error(err),
					zap.String("filename", f.Name),
					zap.String("fileid", f.ID))
				return nil, fmt.Errorf("Error making file public: %s/%s: %s", f.ID, f.Name, err.Error())
			}
		}

	}
	var pubSecret string
	pubParts := strings.Split(f.PermalinkPublic, "-")
	pubSecret = pubParts[len(pubParts)-1]
	req, _ := http.NewRequest(
		"GET",
		fmt.Sprintf(
			"%s?pub_secret=%s",
			f.URLPrivate,
			pubSecret,
		),
		nil)
	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
	buf, _ := httputil.DumpRequest(req, true)
	Logger.Info("Debugging File Request", zap.ByteString("request", buf))
	rsp, err := http.DefaultClient.Do(req)
	bufRsp, _ := httputil.DumpResponse(rsp, false)
	Logger.Info("Debugging File Response", zap.ByteString("response", bufRsp))
	if err != nil {
		return nil, err
	}
	defer rsp.Body.Close()
	if rsp.StatusCode != 200 {
		return nil, fmt.Errorf(rsp.Status)
	}
	if strings.Contains(rsp.Header.Get("Content-Type"), "text/html") {
		return nil, fmt.Errorf("Expected non html content type, got %s", rsp.Header.Get("Content-Type"))
	}
	return ioutil.ReadAll(rsp.Body)
}

func fileBytes(f *slack.File) ([]byte, error) {
	for i := 0; i < 5; i++ {
		if buf, err := fileBytesQuietly(f); err == nil {
			return buf, err
		}
		if i < 4 {
			sleepfor := time.Millisecond * time.Duration(((i+1)*(i+1))*100)
			time.Sleep(sleepfor)
		}

	}
	return fileBytesNoisy(f)
}

func handleChannelUpload(m *slack.MessageEvent) bool {
	if CdnPath == "" || CdnPrefix == "" {
		return false
	}
	if !m.Msg.Upload {
		return false
	}
	file := &m.Files[0]

	// get the username by the ID from the DB
	var user = ":space_invader:"
	profile, err := data.User(m.User)
	if err != nil {
		Logger.Error("Could not get name for user",
			zap.String("userID", m.User),
			zap.String("filename", file.Name),
			zap.Error(err))
	} else {
		if profile.Profile.DisplayName != "" {
			user = profile.Profile.DisplayName
		} else if profile.Profile.RealName != "" {
			user = profile.Profile.RealName
		} else {
			user = profile.Name
		}
	}

	Logger.Info("File upload detected", zap.String("username", user), zap.String("filename", file.Name))
	if buf, err := fileBytes(file); err != nil {
		Logger.Error(
			"error downloading file",
			zap.Error(err),
			zap.String("username", user),
			zap.String("filename", file.Name))
	} else {
		path := fmt.Sprintf("%s/%s", CdnPath, time.Now().Format("2006/01/02/15"))
		if err := os.MkdirAll(path, 0755); err != nil {
			Logger.Error(
				"error making cdn path",
				zap.String("path", path),
				zap.String("username", user),
				zap.String("filename", file.Name))
			return false
		}
		part := &url.URL{Path: file.Name}
		urlPath := fmt.Sprintf("%s/%s-%s", path, file.ID, part.String())
		path = fmt.Sprintf("%s/%s-%s", path, file.ID, file.Name)
		if fp, err := os.Create(path); err != nil {
			Logger.Error(
				"error creating cdn file",
				zap.String("path", path),
				zap.String("username", user),
				zap.String("filename", file.Name))
			return false
		} else {
			if _, err := fp.Write(buf); err != nil {
				fp.Close()
				Logger.Error(
					"error writing to cdn file",
					zap.String("path", path),
					zap.String("username", user),
					zap.String("filename", file.Name))
				return false
			}
			fp.Close()
			fileURL := CdnPrefix + urlPath[len(CdnPath):]

			//delete file and message
			if err := rtm.DeleteFile(file.ID); err != nil {
				Logger.Error(fmt.Sprintf("unable to delete file - %s", err.Error()),
					zap.String("file", file.ID))
			}
			if _, _, err := rtm.DeleteMessage(m.Channel, m.Timestamp); err != nil {
				if err.Error() != "message_not_found" { // message not found means the file deletion deleted the message
					Logger.Error(fmt.Sprintf("unable to delete message - %s", err.Error()),
						zap.String("username", user),
						zap.String("channel", m.Channel),
						zap.String("messageTime", m.Timestamp),
						zap.Error(err))
				}
			}

			if isImage.MatchString(strings.ToLower(file.Name)) {
				for i := 0; i < 5; i++ {
					_, _, err := rtm.PostMessage(
						m.Channel,
						"",
						slack.PostMessageParameters{
							AsUser:      true,
							UnfurlLinks: true,
							UnfurlMedia: true,
							IconEmoji:   ":paperclip:",
							Attachments: []slack.Attachment{
								slack.Attachment{
									AuthorName: fmt.Sprintf("%s uploaded %s", user, file.Name),
									Pretext:    m.Text,
									ImageURL:   fileURL,
								},
							},
							ThreadTimestamp: m.ThreadTimestamp,
						})
					if err != nil {
						Logger.Error(
							"Failed postting cdn link back to slack",
							zap.String("username", user),
							zap.String("filename", file.Name),
							zap.String("url", fileURL))
					} else {
						break
					}
					time.Sleep(time.Second * time.Duration(i))
				}
			} else {
				rtm.SendMessage(&slack.OutgoingMessage{
					ID:              int(time.Now().UnixNano()),
					Channel:         m.Channel,
					Text:            fmt.Sprintf("%s uploaded the file *%s*\n%s", user, file.Title, fileURL),
					Type:            "message",
					ThreadTimestamp: m.ThreadTimestamp,
				})
			}
			Logger.Info("saved CDN file", zap.String("url", fileURL), zap.Int("size", len(buf)))
			return true
		}

	}
	return false
}

func handleDMUpload(m *slack.MessageEvent) bool {
	if CdnPath == "" || CdnPrefix == "" {
		return false
	}
	if !m.Msg.Upload {
		return false
	}
	if buf, err := fileBytes(&m.Files[0]); err != nil {
		Logger.Info("error downloading file", zap.Error(err))
	} else {
		path := fmt.Sprintf("%s/%s", CdnPath, time.Now().Format("2006/01/02/15"))
		if err := os.MkdirAll(path, 0755); err != nil {
			Logger.Error("error making cdn path", zap.String("path", path))
			return false
		}
		part := &url.URL{Path: m.Files[0].Name}
		urlPath := fmt.Sprintf("%s/%s-%s", path, m.Files[0].ID, part.String())
		path = fmt.Sprintf("%s/%s-%s", path, m.Files[0].ID, m.Files[0].Name)
		if fp, err := os.Create(path); err != nil {
			Logger.Error("error creating cdn file", zap.String("path", path))
			return false
		} else {
			if _, err := fp.Write(buf); err != nil {
				fp.Close()
				Logger.Error("error writing to cdn file", zap.String("path", path))
				return false
			}
			fp.Close()
			fileURL := CdnPrefix + urlPath[len(CdnPath):]
			rtm.DeleteFile(m.Files[0].ID)
			rtm.SendMessage(&slack.OutgoingMessage{
				ID:      int(time.Now().UnixNano()),
				Channel: m.Channel,
				Text:    fmt.Sprintf("Thanks for sending me the file instead of uploading it to a channel or group. You can paste the following link anywhere you want to show the file to others! ```%s```", fileURL),
				Type:    "message",
			})
			Logger.Info("saved CDN file", zap.String("url", fileURL), zap.Int("size", len(buf)))
		}

	}
	return true
}
<|endoftext|>


<commit_before>package bot

import (
	"fmt"
	"log"
	"regexp"

	"strings"

	"encoding/json"

	"time"

	"strconv"

	. "github.com/PoolC/slack_bot/util"
	"github.com/marcmak/calc/calc"
	"github.com/nlopes/slack"
	"github.com/robfig/cron"
)

type Meu struct {
	*BaseBot
	rc        RedisClient
	timetable map[string]*TimeTable
	cr        *cron.Cron
}

var (
	calc_re        *regexp.Regexp = regexp.MustCompile("^계산하라 메우 (.+)")
	et_register    *regexp.Regexp = regexp.MustCompile("에타 등록 ([^ ]+)")
	party_register *regexp.Regexp = regexp.MustCompile("(?:(\\d+)월 *(\\d+)일 *)?(\\d+)시(?: *(\\d+)분)?(.+)")
)

func NewMeu(token string, stop *chan struct{}, redisClient RedisClient) *Meu {
	c := cron.New()
	c.Start()
	return &Meu{NewBot(token, stop), redisClient, map[string]*TimeTable{}, c}
}

func (bot *Meu) onConnectedEvent(e *slack.ConnectedEvent) {
	bot.BaseBot.onConnectedEvent(e)
	rescheduleParty(bot)
}

func (bot *Meu) onMessageEvent(e *slack.MessageEvent) {
	message := meuMessageProcess(bot, e)
	switch message.(type) {
	case string:
		bot.sendSimple(e, message.(string))
	}
}

func meuMessageProcess(bot *Meu, e *slack.MessageEvent) interface{} {
	switch {
	case e.Text == "메우, 멱살":
		return "사람은 일을 하고 살아야한다. 메우"
	case e.Text == "메우메우 펫탄탄":
		return `메메메 메메메메 메우메우
메메메 메우메우
펫땅펫땅펫땅펫땅 다이스키`
	}

	text := strings.TrimSpace(e.Text)
	if matched, ok := MatchRE(text, calc_re); ok {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered : %g\n", r)
				bot.replySimple(e, "에러났다 메우")
			}
		}()
		return fmt.Sprintf("%f", calc.Solve(matched[1]))
	} else {
		specialResponses(bot.getBase(), e)
	}
	if bot.IsBeginWithMention(e) {
		if matched, ok := MatchRE(text, et_register); ok {
			// 에브리타임 기록
			et_nick := matched[1]
			bot.rc.Set(fmt.Sprintf("et_nick_%s", e.User), et_nick, 0)
			bot.replySimple(e, "기억했다 메우. 시간표를 가져오겠다 메우.\n시간이 좀 걸릴거다 메우.")

			user, _ := bot.GetUserInfo(e.User)
			go getEveryTimeTable(bot, e.User, user.Name, et_nick)
		} else if strings.HasSuffix(text, "다음 시간") {
			// 에브리타임 다음 시간
			log.Printf("%q", bot.timetable)
			timetable, exists := bot.timetable[e.User]
			if !exists {
				log.Print("Get from redis")
				result := bot.rc.Get(TimeTableKeyName(e.User))
				if result == nil {
					bot.replySimple(e, "시간표 정보가 없다 메우. 등록부터 해달라 메우.")
					return nil
				}

				bytes, _ := result.Bytes()
				timetable = &TimeTable{}
				if json.Unmarshal(bytes, timetable) != nil {
					bot.replySimple(e, "저장된 시간표가 이상하다 메우. 새로 등록해달라 메우.")
					return nil
				}

				bot.timetable[e.User] = timetable
			}

			now := time.Now()
			weekDay := int(now.Weekday() - 1)
			curHour := now.Hour()*12 + now.Minute()/5
			var nextEvent *TimeTableEvent
			if weekDay < 0 || weekDay >= 5 {
				nextEvent = nil
			} else {
				log.Printf("%q", timetable)
				todayEvents := timetable.Days[weekDay].Events
				for _, event := range todayEvents {
					if event.StartTime > curHour {
						nextEvent = &event
						break
					}
				}
			}
			if nextEvent == nil {
				bot.replySimple(e, "오늘은 더 이상 수업이 없다 메우.")
			} else {
				subject := &timetable.Subjects[nextEvent.Id]

				bot.PostMessage(e.Channel, fmt.Sprintf("다음 수업은 \"%s\"다 메우.", subject.Name), slack.PostMessageParameters{
					AsUser:    false,
					IconEmoji: ":meu:",
					Username:  "시간표 알려주는 메우",
					Attachments: []slack.Attachment{
						nextEvent.toSlackAttachment(subject),
					},
				})
			}
		} else if matched, ok := MatchRE(text, party_register); ok {
			now := time.Now()
			month, err := strconv.Atoi(matched[1])
			if err != nil {
				month = int(now.Month())
			}
			day, err := strconv.Atoi(matched[2])
			if err != nil {
				day = now.Day()
			}
			hour, err := strconv.Atoi(matched[3])
			min, err := strconv.Atoi(matched[4])
			if err != nil {
				min = 0
			}
			keyword := matched[5]

			date := time.Date(now.Year(), time.Month(month), day, hour, min, 0, 0, now.Location())
			if date.Before(now) {
				bot.replySimple(e, "과거에 대해서 파티를 모집할 수 없다 메우")
				return nil
			}
			key := partyKey(&date, keyword)
			inserted := bot.rc.SetAdd(key, e.User)
			if inserted.Val() == 1 {
				bot.replySimple(e, "파티 대기에 들어갔다 메우.")
				cardinal := bot.rc.SetCard(key)
				if cardinal.Val() == 1 {
					scheduleParty(bot, &date, keyword)
				}
			} else {
				bot.replySimple(e, "이미 들어가있는 파티다 메우.")
			}
		}
	}

	return nil
}
<commit_msg>add new command<commit_after>package bot

import (
	"fmt"
	"log"
	"regexp"

	"strings"

	"encoding/json"

	"time"

	"strconv"

	. "github.com/PoolC/slack_bot/util"
	"github.com/marcmak/calc/calc"
	"github.com/nlopes/slack"
	"github.com/robfig/cron"
)

type Meu struct {
	*BaseBot
	rc        RedisClient
	timetable map[string]*TimeTable
	cr        *cron.Cron
}

var (
	calc_re        *regexp.Regexp = regexp.MustCompile("^계산하라 메우 (.+)")
	et_register    *regexp.Regexp = regexp.MustCompile("에타 등록 ([^ ]+)")
	et_call        *regexp.Regexp = regexp.MustCompile("다음 (시간|수업)")
	party_register *regexp.Regexp = regexp.MustCompile("(?:(\\d+)월 *(\\d+)일 *)?(\\d+)시(?: *(\\d+)분)?(.+)")
)

func NewMeu(token string, stop *chan struct{}, redisClient RedisClient) *Meu {
	c := cron.New()
	c.Start()
	return &Meu{NewBot(token, stop), redisClient, map[string]*TimeTable{}, c}
}

func (bot *Meu) onConnectedEvent(e *slack.ConnectedEvent) {
	bot.BaseBot.onConnectedEvent(e)
	rescheduleParty(bot)
}

func (bot *Meu) onMessageEvent(e *slack.MessageEvent) {
	message := meuMessageProcess(bot, e)
	switch message.(type) {
	case string:
		bot.sendSimple(e, message.(string))
	}
}

func meuMessageProcess(bot *Meu, e *slack.MessageEvent) interface{} {
	switch {
	case e.Text == "메우, 멱살":
		return "사람은 일을 하고 살아야한다. 메우"
	case e.Text == "메우메우 펫탄탄":
		return `메메메 메메메메 메우메우
메메메 메우메우
펫땅펫땅펫땅펫땅 다이스키`
	}

	text := strings.TrimSpace(e.Text)
	if matched, ok := MatchRE(text, calc_re); ok {
		defer func() {
			if r := recover(); r != nil {
				log.Printf("Recovered : %g\n", r)
				bot.replySimple(e, "에러났다 메우")
			}
		}()
		return fmt.Sprintf("%f", calc.Solve(matched[1]))
	} else {
		specialResponses(bot.getBase(), e)
	}
	if bot.IsBeginWithMention(e) {
		if matched, ok := MatchRE(text, et_register); ok {
			// 에브리타임 기록
			et_nick := matched[1]
			bot.rc.Set(fmt.Sprintf("et_nick_%s", e.User), et_nick, 0)
			bot.replySimple(e, "기억했다 메우. 시간표를 가져오겠다 메우.\n시간이 좀 걸릴거다 메우.")

			user, _ := bot.GetUserInfo(e.User)
			go getEveryTimeTable(bot, e.User, user.Name, et_nick)
		} else if _, ok := MatchRE(text, et_call); ok {
			// 에브리타임 다음 시간
			log.Printf("%q", bot.timetable)
			timetable, exists := bot.timetable[e.User]
			if !exists {
				log.Print("Get from redis")
				result := bot.rc.Get(TimeTableKeyName(e.User))
				if result == nil {
					bot.replySimple(e, "시간표 정보가 없다 메우. 등록부터 해달라 메우.")
					return nil
				}

				bytes, _ := result.Bytes()
				timetable = &TimeTable{}
				if json.Unmarshal(bytes, timetable) != nil {
					bot.replySimple(e, "저장된 시간표가 이상하다 메우. 새로 등록해달라 메우.")
					return nil
				}

				bot.timetable[e.User] = timetable
			}

			now := time.Now()
			weekDay := int(now.Weekday() - 1)
			curHour := now.Hour()*12 + now.Minute()/5
			var nextEvent *TimeTableEvent
			if weekDay < 0 || weekDay >= 5 {
				nextEvent = nil
			} else {
				log.Printf("%q", timetable)
				todayEvents := timetable.Days[weekDay].Events
				for _, event := range todayEvents {
					if event.StartTime > curHour {
						nextEvent = &event
						break
					}
				}
			}
			if nextEvent == nil {
				bot.replySimple(e, "오늘은 더 이상 수업이 없다 메우.")
			} else {
				subject := &timetable.Subjects[nextEvent.Id]

				bot.PostMessage(e.Channel, fmt.Sprintf("다음 수업은 \"%s\"다 메우.", subject.Name), slack.PostMessageParameters{
					AsUser:    false,
					IconEmoji: ":meu:",
					Username:  "시간표 알려주는 메우",
					Attachments: []slack.Attachment{
						nextEvent.toSlackAttachment(subject),
					},
				})
			}
		} else if matched, ok := MatchRE(text, party_register); ok {
			now := time.Now()
			month, err := strconv.Atoi(matched[1])
			if err != nil {
				month = int(now.Month())
			}
			day, err := strconv.Atoi(matched[2])
			if err != nil {
				day = now.Day()
			}
			hour, err := strconv.Atoi(matched[3])
			min, err := strconv.Atoi(matched[4])
			if err != nil {
				min = 0
			}
			keyword := matched[5]

			date := time.Date(now.Year(), time.Month(month), day, hour, min, 0, 0, now.Location())
			if date.Before(now) {
				bot.replySimple(e, "과거에 대해서 파티를 모집할 수 없다 메우")
				return nil
			}
			key := partyKey(&date, keyword)
			inserted := bot.rc.SetAdd(key, e.User)
			if inserted.Val() == 1 {
				bot.replySimple(e, "파티 대기에 들어갔다 메우.")
				cardinal := bot.rc.SetCard(key)
				if cardinal.Val() == 1 {
					scheduleParty(bot, &date, keyword)
				}
			} else {
				bot.replySimple(e, "이미 들어가있는 파티다 메우.")
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>// vim: tabstop=4 softtabstop=4 shiftwidth=4 noexpandtab tw=72
//http://stackoverflow.com/questions/8757389/reading-file-line-by-line-in-go
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"github.com/atotto/clipboard"
	"os"
	"strconv"
	"strings"
	"errors"
)

type existsFunc func(string) bool

func osStatExists(file string) bool {
	_, err := os.Stat(file)
	return err == nil
}

var ignoreList = [...]string{"/", ".", "./", "..", "../"}

//var rootListing, _ = ioutil.ReadDir("/")
//var pwdListing, _ = ioutil.ReadDir(".")

func ignored(file string) bool {
	for _, val := range ignoreList {
		if file == val {
			return true
		}
	}
	return false
}

func longestFileEndIndex(line []rune, exists existsFunc) int {
	// Possible optimisations:
	// 1. this should start at the end - the longest substring first
	// 2. it could be a good strategy to list files and try to
	//    find a common prefix - if not just stop right there and then
	//    from / if it starts with / and if not from `pwd`
	//    do file listing from / and `pwd` only once
	//    need to consider relative dirs though which is annoying

	maxIndex := 0
	for i, _ := range line {
		slice := line[0 : i+1]
		file := string(slice)
		if !ignored(file) {
			if exists(file) {
				// TODO if this is not a dir, stop here
				maxIndex = i
			}
		}
	}
	return maxIndex
}

func longestFileInLine(line string, exists existsFunc) (firstCharIndex int, lastCharIndex int) {
	for searchStartIndex, _ := range line {
		searchSpace := []rune(line[searchStartIndex:len(line)])
		lastCharIndexInSlice := longestFileEndIndex(searchSpace, exists)
		lastCharIndexInLine := lastCharIndexInSlice + searchStartIndex
		if lastCharIndexInSlice > 0 && lastCharIndexInLine > lastCharIndex {
			lastCharIndex = lastCharIndexInLine
			firstCharIndex = searchStartIndex
		}
	}

	return firstCharIndex, lastCharIndex
}

func askUser() (requestedNumbers []string, err error) {
	fmt.Println()
	fmt.Print("to clipboard: ")
	ttyFile, err := os.Open("/dev/tty")
	if err != nil {
		return nil, err
	}
	defer ttyFile.Close()
	ttyReader := bufio.NewReader(ttyFile)
	s, err := ttyReader.ReadString('\n')
	if err != nil {
		return nil, err
	}
	return strings.Fields(s), nil
}

type PrintFunction func(string, string, int, int)

type Processor interface {
	processFile(file string) error
	processEnd() error
}

type NumbersGiven struct {
	clip      *bytes.Buffer
	fileCount *int
	numbers   []string
}

type AskForNumbers struct {
	clip  *bytes.Buffer
	files []string
}

func printShortFormat(fileCount *int) PrintFunction {
	return func(file string, line string, firstCharIndex int, lastCharIndex int) {
		fmt.Println(strconv.Itoa(*fileCount), file)
	}
}

func printLongFormat(fileCount *int) PrintFunction {
	return func(file string, line string, firstCharIndex int, lastCharIndex int) {
		fmt.Print("{")
		fmt.Print(strconv.Itoa(*fileCount))
		fmt.Print("} ")
		fmt.Print(line[:firstCharIndex])
		fmt.Print("{")
		fmt.Print(file)
		fmt.Print("}")
		fmt.Print(line[lastCharIndex+1:])
	}
}

func (ng *NumbersGiven) processEnd() error {
	writeToClipboard(ng.clip)
	return nil
}

func (ng *NumbersGiven) processFile(file string) error {
	// collect any file position arguments to copy to the
	// clipboard later
	for _, v := range ng.numbers {
		n, err := strconv.Atoi(v)
		if err != nil {
			fmt.Fprintf(os.Stderr, "nw: %s is not a number\n", v)
			return err
		}
		if n == *ng.fileCount {
			ng.clip.WriteString(file)
			ng.clip.WriteString(" ")
		}
	}
	return nil
}

func (afn *AskForNumbers) processEnd() error {
	if len(afn.files) == 0 {
		fmt.Println("nw: no files names found, NUMBERWANG!")
		return nil
	}
	requestedNumbers, err := askUser()
	if err != nil {
		fmt.Fprintf(os.Stderr, "nw: failed to read input: %s\n", err)
		return err
	}

	for _, n := range requestedNumbers {
		i, err := strconv.Atoi(n)
		if err != nil {
			fmt.Fprintf(os.Stderr, "nw: %s is not a number\n", n)
			return err
		}
		if i <= 0 || i > len(afn.files) {
			fmt.Fprintf(os.Stderr, "nw: %s is not a valid choice\n", n)
			return errors.New("invalid choice")
		}
		afn.clip.WriteString(afn.files[i-1])
		afn.clip.WriteString(" ")
	}

	writeToClipboard(afn.clip)
	return nil
}

func (afn *AskForNumbers) processFile(file string) error {
	afn.files = append(afn.files, file)
	return nil
}

func writeToClipboard(buffer *bytes.Buffer) {
	clipboardOutput := buffer.String()
	if clipboardOutput != "" {
		clipboard.WriteAll(clipboardOutput)
		fmt.Printf("nw: wrote \"%s\" to clipboard\n", clipboardOutput)
	}
}

func main() {
	var fileCount int
	var clip bytes.Buffer

	short := flag.Bool("s", false, "short format, only display file names")
	flag.Parse()

	extraArgs := flag.Args()

	var processor Processor
	if len(extraArgs) == 0 {
		processor = &AskForNumbers{clip: &clip}
	} else {
		processor = &NumbersGiven{
			clip:      &clip,
			fileCount: &fileCount,
			numbers:   extraArgs,
		}
	}
	var printer PrintFunction
	if *short {
		printer = printShortFormat(&fileCount)
	} else {
		printer = printLongFormat(&fileCount)
	}

	reader := bufio.NewReader(os.Stdin)

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			break
		}

		firstCharIndex, lastCharIndex := longestFileInLine(line, osStatExists)

		if lastCharIndex > 0 {
			fileCount++
			file := line[firstCharIndex : lastCharIndex+1]

			printer(file, line, firstCharIndex, lastCharIndex)
			err := processor.processFile(file)
			if err != nil {
				os.Exit(1)
			}
		} else if (!*short) {
			fmt.Print(line)
		}
	}

	err := processor.processEnd()
	if err != nil {
		os.Exit(1)
	}
}
<commit_msg>remove pointless comment<commit_after>// vim: tabstop=4 softtabstop=4 shiftwidth=4 noexpandtab tw=72
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"github.com/atotto/clipboard"
	"os"
	"strconv"
	"strings"
	"errors"
)

type existsFunc func(string) bool

func osStatExists(file string) bool {
	_, err := os.Stat(file)
	return err == nil
}

var ignoreList = [...]string{"/", ".", "./", "..", "../"}

//var rootListing, _ = ioutil.ReadDir("/")
//var pwdListing, _ = ioutil.ReadDir(".")

func ignored(file string) bool {
	for _, val := range ignoreList {
		if file == val {
			return true
		}
	}
	return false
}

func longestFileEndIndex(line []rune, exists existsFunc) int {
	// Possible optimisations:
	// 1. this should start at the end - the longest substring first
	// 2. it could be a good strategy to list files and try to
	//    find a common prefix - if not just stop right there and then
	//    from / if it starts with / and if not from `pwd`
	//    do file listing from / and `pwd` only once
	//    need to consider relative dirs though which is annoying

	maxIndex := 0
	for i, _ := range line {
		slice := line[0 : i+1]
		file := string(slice)
		if !ignored(file) {
			if exists(file) {
				// TODO if this is not a dir, stop here
				maxIndex = i
			}
		}
	}
	return maxIndex
}

func longestFileInLine(line string, exists existsFunc) (firstCharIndex int, lastCharIndex int) {
	for searchStartIndex, _ := range line {
		searchSpace := []rune(line[searchStartIndex:len(line)])
		lastCharIndexInSlice := longestFileEndIndex(searchSpace, exists)
		lastCharIndexInLine := lastCharIndexInSlice + searchStartIndex
		if lastCharIndexInSlice > 0 && lastCharIndexInLine > lastCharIndex {
			lastCharIndex = lastCharIndexInLine
			firstCharIndex = searchStartIndex
		}
	}

	return firstCharIndex, lastCharIndex
}

func askUser() (requestedNumbers []string, err error) {
	fmt.Println()
	fmt.Print("to clipboard: ")
	ttyFile, err := os.Open("/dev/tty")
	if err != nil {
		return nil, err
	}
	defer ttyFile.Close()
	ttyReader := bufio.NewReader(ttyFile)
	s, err := ttyReader.ReadString('\n')
	if err != nil {
		return nil, err
	}
	return strings.Fields(s), nil
}

type PrintFunction func(string, string, int, int)

type Processor interface {
	processFile(file string) error
	processEnd() error
}

type NumbersGiven struct {
	clip      *bytes.Buffer
	fileCount *int
	numbers   []string
}

type AskForNumbers struct {
	clip  *bytes.Buffer
	files []string
}

func printShortFormat(fileCount *int) PrintFunction {
	return func(file string, line string, firstCharIndex int, lastCharIndex int) {
		fmt.Println(strconv.Itoa(*fileCount), file)
	}
}

func printLongFormat(fileCount *int) PrintFunction {
	return func(file string, line string, firstCharIndex int, lastCharIndex int) {
		fmt.Print("{")
		fmt.Print(strconv.Itoa(*fileCount))
		fmt.Print("} ")
		fmt.Print(line[:firstCharIndex])
		fmt.Print("{")
		fmt.Print(file)
		fmt.Print("}")
		fmt.Print(line[lastCharIndex+1:])
	}
}

func (ng *NumbersGiven) processEnd() error {
	writeToClipboard(ng.clip)
	return nil
}

func (ng *NumbersGiven) processFile(file string) error {
	// collect any file position arguments to copy to the
	// clipboard later
	for _, v := range ng.numbers {
		n, err := strconv.Atoi(v)
		if err != nil {
			fmt.Fprintf(os.Stderr, "nw: %s is not a number\n", v)
			return err
		}
		if n == *ng.fileCount {
			ng.clip.WriteString(file)
			ng.clip.WriteString(" ")
		}
	}
	return nil
}

func (afn *AskForNumbers) processEnd() error {
	if len(afn.files) == 0 {
		fmt.Println("nw: no files names found, NUMBERWANG!")
		return nil
	}
	requestedNumbers, err := askUser()
	if err != nil {
		fmt.Fprintf(os.Stderr, "nw: failed to read input: %s\n", err)
		return err
	}

	for _, n := range requestedNumbers {
		i, err := strconv.Atoi(n)
		if err != nil {
			fmt.Fprintf(os.Stderr, "nw: %s is not a number\n", n)
			return err
		}
		if i <= 0 || i > len(afn.files) {
			fmt.Fprintf(os.Stderr, "nw: %s is not a valid choice\n", n)
			return errors.New("invalid choice")
		}
		afn.clip.WriteString(afn.files[i-1])
		afn.clip.WriteString(" ")
	}

	writeToClipboard(afn.clip)
	return nil
}

func (afn *AskForNumbers) processFile(file string) error {
	afn.files = append(afn.files, file)
	return nil
}

func writeToClipboard(buffer *bytes.Buffer) {
	clipboardOutput := buffer.String()
	if clipboardOutput != "" {
		clipboard.WriteAll(clipboardOutput)
		fmt.Printf("nw: wrote \"%s\" to clipboard\n", clipboardOutput)
	}
}

func main() {
	var fileCount int
	var clip bytes.Buffer

	short := flag.Bool("s", false, "short format, only display file names")
	flag.Parse()

	extraArgs := flag.Args()

	var processor Processor
	if len(extraArgs) == 0 {
		processor = &AskForNumbers{clip: &clip}
	} else {
		processor = &NumbersGiven{
			clip:      &clip,
			fileCount: &fileCount,
			numbers:   extraArgs,
		}
	}
	var printer PrintFunction
	if *short {
		printer = printShortFormat(&fileCount)
	} else {
		printer = printLongFormat(&fileCount)
	}

	reader := bufio.NewReader(os.Stdin)

	for {
		line, err := reader.ReadString('\n')
		if err != nil {
			break
		}

		firstCharIndex, lastCharIndex := longestFileInLine(line, osStatExists)

		if lastCharIndex > 0 {
			fileCount++
			file := line[firstCharIndex : lastCharIndex+1]

			printer(file, line, firstCharIndex, lastCharIndex)
			err := processor.processFile(file)
			if err != nil {
				os.Exit(1)
			}
		} else if (!*short) {
			fmt.Print(line)
		}
	}

	err := processor.processEnd()
	if err != nil {
		os.Exit(1)
	}
}
<|endoftext|>