<commit_before>package main

import (
	"crypto/md5"
	"fmt"
	"io"
	"mime"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudfront"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/ryanuber/go-glob"
)

type AWS struct {
	client   *s3.S3
	cfClient *cloudfront.CloudFront
	remote   []string
	local    []string
	plugin   *Plugin
}

func NewAWS(p *Plugin) AWS {
	sess := session.New(&aws.Config{
		Endpoint:         &p.Endpoint,
		DisableSSL:       aws.Bool(strings.HasPrefix(p.Endpoint, "http://")),
		Credentials:      credentials.NewStaticCredentials(p.Key, p.Secret, ""),
		S3ForcePathStyle: aws.Bool(p.PathStyle),
		Region:           aws.String(p.Region),
	})
	c := s3.New(sess)
	cf := cloudfront.New(sess)
	r := make([]string, 1, 1)
	l := make([]string, 1, 1)

	return AWS{c, cf, r, l, p}
}

func (a *AWS) Upload(local, remote string) error {
	p := a.plugin
	if local == "" {
		return nil
	}

	file, err := os.Open(local)
	if err != nil {
		return err
	}

	defer file.Close()

	var access string
	for pattern := range p.Access {
		if match := glob.Glob(pattern, local); match == true {
			access = p.Access[pattern]
			break
		}
	}

	if access == "" {
		access = "private"
	}

	fileExt := filepath.Ext(local)

	var contentType string
	for patternExt := range p.ContentType {
		if patternExt == fileExt {
			contentType = p.ContentType[patternExt]
			break
		}
	}

	if contentType == "" {
		contentType = mime.TypeByExtension(fileExt)
	}

	var contentEncoding string
	for patternExt := range p.ContentEncoding {
		if patternExt == fileExt {
			contentEncoding = p.ContentEncoding[patternExt]
			break
		}
	}

	var cacheControl string
	for pattern := range p.CacheControl {
		if match := glob.Glob(pattern, local); match == true {
			cacheControl = p.CacheControl[pattern]
			break
		}
	}

	metadata := map[string]*string{}
	for pattern := range p.Metadata {
		if match := glob.Glob(pattern, local); match == true {
			for k, v := range p.Metadata[pattern] {
				metadata[k] = aws.String(v)
			}
			break
		}
	}

	head, err := a.client.HeadObject(&s3.HeadObjectInput{
		Bucket: aws.String(p.Bucket),
		Key:    aws.String(remote),
	})
	if err != nil && err.(awserr.Error).Code() != "404" {
		if err.(awserr.Error).Code() == "404" {
			return err
		}

		debug("\"%s\" not found in bucket, uploading with Content-Type \"%s\" and permissions \"%s\"", local, contentType, access)
		var putObject = &s3.PutObjectInput{
			Bucket:      aws.String(p.Bucket),
			Key:         aws.String(remote),
			Body:        file,
			ContentType: aws.String(contentType),
			ACL:         aws.String(access),
			Metadata:    metadata,
		}

		if len(cacheControl) > 0 {
			putObject.CacheControl = aws.String(cacheControl)
		}

		if len(contentEncoding) > 0 {
			putObject.ContentEncoding = aws.String(contentEncoding)
		}

		// skip upload during dry run
		if a.plugin.DryRun {
			return nil
		}

		_, err = a.client.PutObject(putObject)
		return err
	}

	hash := md5.New()
	io.Copy(hash, file)
	sum := fmt.Sprintf("\"%x\"", hash.Sum(nil))

	if sum == *head.ETag {
		shouldCopy := false

		if head.ContentType == nil && contentType != "" {
			debug("Content-Type has changed from unset to %s", contentType)
			shouldCopy = true
		}

		if !shouldCopy && head.ContentType != nil && contentType != *head.ContentType {
			debug("Content-Type has changed from %s to %s", *head.ContentType, contentType)
			shouldCopy = true
		}

		if !shouldCopy && head.ContentEncoding == nil && contentEncoding != "" {
			debug("Content-Encoding has changed from unset to %s", contentEncoding)
			shouldCopy = true
		}

		if !shouldCopy && head.ContentEncoding != nil && contentEncoding != *head.ContentEncoding {
			debug("Content-Encoding has changed from %s to %s", *head.ContentEncoding, contentEncoding)
			shouldCopy = true
		}

		if !shouldCopy && head.CacheControl == nil && cacheControl != "" {
			debug("Cache-Control has changed from unset to %s", cacheControl)
			shouldCopy = true
		}

		if !shouldCopy && head.CacheControl != nil && cacheControl != *head.CacheControl {
			debug("Cache-Control has changed from %s to %s", *head.CacheControl, cacheControl)
			shouldCopy = true
		}

		if !shouldCopy && len(head.Metadata) != len(metadata) {
			debug("Count of metadata values has changed for %s", local)
			shouldCopy = true
		}

		if !shouldCopy && len(metadata) > 0 {
			for k, v := range metadata {
				if hv, ok := head.Metadata[k]; ok {
					if *v != *hv {
						debug("Metadata values have changed for %s", local)
						shouldCopy = true
						break
					}
				}
			}
		}

		if !shouldCopy {
			grant, err := a.client.GetObjectAcl(&s3.GetObjectAclInput{
				Bucket: aws.String(p.Bucket),
				Key:    aws.String(remote),
			})
			if err != nil {
				return err
			}

			previousAccess := "private"
			for _, g := range grant.Grants {
				gt := *g.Grantee
				if gt.URI != nil {
					if *gt.URI == "http://acs.amazonaws.com/groups/global/AllUsers" {
						if *g.Permission == "READ" {
							previousAccess = "public-read"
						} else if *g.Permission == "WRITE" {
							previousAccess = "public-read-write"
						}
					} else if *gt.URI == "http://acs.amazonaws.com/groups/global/AllUsers" {
						if *g.Permission == "READ" {
							previousAccess = "authenticated-read"
						}
					}
				}
			}

			if previousAccess != access {
				debug("Permissions for \"%s\" have changed from \"%s\" to \"%s\"", remote, previousAccess, access)
				shouldCopy = true
			}
		}

		if !shouldCopy {
			debug("Skipping \"%s\" because hashes and metadata match", local)
			return nil
		}

		debug("Updating metadata for \"%s\" Content-Type: \"%s\", ACL: \"%s\"", local, contentType, access)
		var copyObject = &s3.CopyObjectInput{
			Bucket:            aws.String(p.Bucket),
			Key:               aws.String(remote),
			CopySource:        aws.String(fmt.Sprintf("%s/%s", p.Bucket, remote)),
			ACL:               aws.String(access),
			ContentType:       aws.String(contentType),
			Metadata:          metadata,
			MetadataDirective: aws.String("REPLACE"),
		}

		if len(cacheControl) > 0 {
			copyObject.CacheControl = aws.String(cacheControl)
		}

		if len(contentEncoding) > 0 {
			copyObject.ContentEncoding = aws.String(contentEncoding)
		}

		// skip update if dry run
		if a.plugin.DryRun {
			return nil
		}

		_, err = a.client.CopyObject(copyObject)
		return err
	} else {
		_, err = file.Seek(0, 0)
		if err != nil {
			return err
		}

		debug("Uploading \"%s\" with Content-Type \"%s\" and permissions \"%s\"", local, contentType, access)
		var putObject = &s3.PutObjectInput{
			Bucket:      aws.String(p.Bucket),
			Key:         aws.String(remote),
			Body:        file,
			ContentType: aws.String(contentType),
			ACL:         aws.String(access),
			Metadata:    metadata,
		}

		if len(cacheControl) > 0 {
			putObject.CacheControl = aws.String(cacheControl)
		}

		if len(contentEncoding) > 0 {
			putObject.ContentEncoding = aws.String(contentEncoding)
		}

		// skip upload if dry run
		if a.plugin.DryRun {
			return nil
		}

		_, err = a.client.PutObject(putObject)
		return err
	}
}

func (a *AWS) Redirect(path, location string) error {
	p := a.plugin
	debug("Adding redirect from \"%s\" to \"%s\"", path, location)

	if a.plugin.DryRun {
		return nil
	}

	_, err := a.client.PutObject(&s3.PutObjectInput{
		Bucket: aws.String(p.Bucket),
		Key:    aws.String(path),
		ACL:    aws.String("public-read"),
		WebsiteRedirectLocation: aws.String(location),
	})
	return err
}

func (a *AWS) Delete(remote string) error {
	p := a.plugin
	debug("Removing remote file \"%s\"", remote)

	if a.plugin.DryRun {
		return nil
	}

	_, err := a.client.DeleteObject(&s3.DeleteObjectInput{
		Bucket: aws.String(p.Bucket),
		Key:    aws.String(remote),
	})
	return err
}

func (a *AWS) List(path string) ([]string, error) {
	p := a.plugin
	remote := make([]string, 1, 1)
	resp, err := a.client.ListObjects(&s3.ListObjectsInput{
		Bucket: aws.String(p.Bucket),
		Prefix: aws.String(path),
	})
	if err != nil {
		return remote, err
	}

	for _, item := range resp.Contents {
		remote = append(remote, *item.Key)
	}

	for *resp.IsTruncated {
		resp, err = a.client.ListObjects(&s3.ListObjectsInput{
			Bucket: aws.String(p.Bucket),
			Prefix: aws.String(path),
			Marker: aws.String(remote[len(remote)-1]),
		})

		if err != nil {
			return remote, err
		}

		for _, item := range resp.Contents {
			remote = append(remote, *item.Key)
		}
	}

	return remote, nil
}

func (a *AWS) Invalidate(invalidatePath string) error {
	p := a.plugin
	debug("Invalidating \"%s\"", invalidatePath)
	_, err := a.cfClient.CreateInvalidation(&cloudfront.CreateInvalidationInput{
		DistributionId: aws.String(p.CloudFrontDistribution),
		InvalidationBatch: &cloudfront.InvalidationBatch{
			CallerReference: aws.String(time.Now().Format(time.RFC3339Nano)),
			Paths: &cloudfront.Paths{
				Quantity: aws.Int64(1),
				Items: []*string{
					aws.String(invalidatePath),
				},
			},
		},
	})
	return err
}
<commit_msg>Endpoint for aws.Config is optional<commit_after>package main

import (
	"crypto/md5"
	"fmt"
	"io"
	"mime"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudfront"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/ryanuber/go-glob"
)

type AWS struct {
	client   *s3.S3
	cfClient *cloudfront.CloudFront
	remote   []string
	local    []string
	plugin   *Plugin
}

func NewAWS(p *Plugin) AWS {
	sessCfg := &aws.Config{
		Credentials:      credentials.NewStaticCredentials(p.Key, p.Secret, ""),
		S3ForcePathStyle: aws.Bool(p.PathStyle),
		Region:           aws.String(p.Region),
	}

	fmt.Println(p.Endpoint)
	if p.Endpoint != "" {
		sessCfg.Endpoint = &p.Endpoint
		sessCfg.DisableSSL = aws.Bool(strings.HasPrefix(p.Endpoint, "http://"))
	}
	sess := session.New(sessCfg)

	c := s3.New(sess)
	cf := cloudfront.New(sess)
	r := make([]string, 1, 1)
	l := make([]string, 1, 1)

	return AWS{c, cf, r, l, p}
}

func (a *AWS) Upload(local, remote string) error {
	p := a.plugin
	if local == "" {
		return nil
	}

	file, err := os.Open(local)
	if err != nil {
		return err
	}

	defer file.Close()

	var access string
	for pattern := range p.Access {
		if match := glob.Glob(pattern, local); match == true {
			access = p.Access[pattern]
			break
		}
	}

	if access == "" {
		access = "private"
	}

	fileExt := filepath.Ext(local)

	var contentType string
	for patternExt := range p.ContentType {
		if patternExt == fileExt {
			contentType = p.ContentType[patternExt]
			break
		}
	}

	if contentType == "" {
		contentType = mime.TypeByExtension(fileExt)
	}

	var contentEncoding string
	for patternExt := range p.ContentEncoding {
		if patternExt == fileExt {
			contentEncoding = p.ContentEncoding[patternExt]
			break
		}
	}

	var cacheControl string
	for pattern := range p.CacheControl {
		if match := glob.Glob(pattern, local); match == true {
			cacheControl = p.CacheControl[pattern]
			break
		}
	}

	metadata := map[string]*string{}
	for pattern := range p.Metadata {
		if match := glob.Glob(pattern, local); match == true {
			for k, v := range p.Metadata[pattern] {
				metadata[k] = aws.String(v)
			}
			break
		}
	}

	head, err := a.client.HeadObject(&s3.HeadObjectInput{
		Bucket: aws.String(p.Bucket),
		Key:    aws.String(remote),
	})
	if err != nil && err.(awserr.Error).Code() != "404" {
		if err.(awserr.Error).Code() == "404" {
			return err
		}

		debug("\"%s\" not found in bucket, uploading with Content-Type \"%s\" and permissions \"%s\"", local, contentType, access)
		var putObject = &s3.PutObjectInput{
			Bucket:      aws.String(p.Bucket),
			Key:         aws.String(remote),
			Body:        file,
			ContentType: aws.String(contentType),
			ACL:         aws.String(access),
			Metadata:    metadata,
		}

		if len(cacheControl) > 0 {
			putObject.CacheControl = aws.String(cacheControl)
		}

		if len(contentEncoding) > 0 {
			putObject.ContentEncoding = aws.String(contentEncoding)
		}

		// skip upload during dry run
		if a.plugin.DryRun {
			return nil
		}

		_, err = a.client.PutObject(putObject)
		return err
	}

	hash := md5.New()
	io.Copy(hash, file)
	sum := fmt.Sprintf("\"%x\"", hash.Sum(nil))

	if sum == *head.ETag {
		shouldCopy := false

		if head.ContentType == nil && contentType != "" {
			debug("Content-Type has changed from unset to %s", contentType)
			shouldCopy = true
		}

		if !shouldCopy && head.ContentType != nil && contentType != *head.ContentType {
			debug("Content-Type has changed from %s to %s", *head.ContentType, contentType)
			shouldCopy = true
		}

		if !shouldCopy && head.ContentEncoding == nil && contentEncoding != "" {
			debug("Content-Encoding has changed from unset to %s", contentEncoding)
			shouldCopy = true
		}

		if !shouldCopy && head.ContentEncoding != nil && contentEncoding != *head.ContentEncoding {
			debug("Content-Encoding has changed from %s to %s", *head.ContentEncoding, contentEncoding)
			shouldCopy = true
		}

		if !shouldCopy && head.CacheControl == nil && cacheControl != "" {
			debug("Cache-Control has changed from unset to %s", cacheControl)
			shouldCopy = true
		}

		if !shouldCopy && head.CacheControl != nil && cacheControl != *head.CacheControl {
			debug("Cache-Control has changed from %s to %s", *head.CacheControl, cacheControl)
			shouldCopy = true
		}

		if !shouldCopy && len(head.Metadata) != len(metadata) {
			debug("Count of metadata values has changed for %s", local)
			shouldCopy = true
		}

		if !shouldCopy && len(metadata) > 0 {
			for k, v := range metadata {
				if hv, ok := head.Metadata[k]; ok {
					if *v != *hv {
						debug("Metadata values have changed for %s", local)
						shouldCopy = true
						break
					}
				}
			}
		}

		if !shouldCopy {
			grant, err := a.client.GetObjectAcl(&s3.GetObjectAclInput{
				Bucket: aws.String(p.Bucket),
				Key:    aws.String(remote),
			})
			if err != nil {
				return err
			}

			previousAccess := "private"
			for _, g := range grant.Grants {
				gt := *g.Grantee
				if gt.URI != nil {
					if *gt.URI == "http://acs.amazonaws.com/groups/global/AllUsers" {
						if *g.Permission == "READ" {
							previousAccess = "public-read"
						} else if *g.Permission == "WRITE" {
							previousAccess = "public-read-write"
						}
					} else if *gt.URI == "http://acs.amazonaws.com/groups/global/AllUsers" {
						if *g.Permission == "READ" {
							previousAccess = "authenticated-read"
						}
					}
				}
			}

			if previousAccess != access {
				debug("Permissions for \"%s\" have changed from \"%s\" to \"%s\"", remote, previousAccess, access)
				shouldCopy = true
			}
		}

		if !shouldCopy {
			debug("Skipping \"%s\" because hashes and metadata match", local)
			return nil
		}

		debug("Updating metadata for \"%s\" Content-Type: \"%s\", ACL: \"%s\"", local, contentType, access)
		var copyObject = &s3.CopyObjectInput{
			Bucket:            aws.String(p.Bucket),
			Key:               aws.String(remote),
			CopySource:        aws.String(fmt.Sprintf("%s/%s", p.Bucket, remote)),
			ACL:               aws.String(access),
			ContentType:       aws.String(contentType),
			Metadata:          metadata,
			MetadataDirective: aws.String("REPLACE"),
		}

		if len(cacheControl) > 0 {
			copyObject.CacheControl = aws.String(cacheControl)
		}

		if len(contentEncoding) > 0 {
			copyObject.ContentEncoding = aws.String(contentEncoding)
		}

		// skip update if dry run
		if a.plugin.DryRun {
			return nil
		}

		_, err = a.client.CopyObject(copyObject)
		return err
	} else {
		_, err = file.Seek(0, 0)
		if err != nil {
			return err
		}

		debug("Uploading \"%s\" with Content-Type \"%s\" and permissions \"%s\"", local, contentType, access)
		var putObject = &s3.PutObjectInput{
			Bucket:      aws.String(p.Bucket),
			Key:         aws.String(remote),
			Body:        file,
			ContentType: aws.String(contentType),
			ACL:         aws.String(access),
			Metadata:    metadata,
		}

		if len(cacheControl) > 0 {
			putObject.CacheControl = aws.String(cacheControl)
		}

		if len(contentEncoding) > 0 {
			putObject.ContentEncoding = aws.String(contentEncoding)
		}

		// skip upload if dry run
		if a.plugin.DryRun {
			return nil
		}

		_, err = a.client.PutObject(putObject)
		return err
	}
}

func (a *AWS) Redirect(path, location string) error {
	p := a.plugin
	debug("Adding redirect from \"%s\" to \"%s\"", path, location)

	if a.plugin.DryRun {
		return nil
	}

	_, err := a.client.PutObject(&s3.PutObjectInput{
		Bucket: aws.String(p.Bucket),
		Key:    aws.String(path),
		ACL:    aws.String("public-read"),
		WebsiteRedirectLocation: aws.String(location),
	})
	return err
}

func (a *AWS) Delete(remote string) error {
	p := a.plugin
	debug("Removing remote file \"%s\"", remote)

	if a.plugin.DryRun {
		return nil
	}

	_, err := a.client.DeleteObject(&s3.DeleteObjectInput{
		Bucket: aws.String(p.Bucket),
		Key:    aws.String(remote),
	})
	return err
}

func (a *AWS) List(path string) ([]string, error) {
	p := a.plugin
	remote := make([]string, 1, 1)
	resp, err := a.client.ListObjects(&s3.ListObjectsInput{
		Bucket: aws.String(p.Bucket),
		Prefix: aws.String(path),
	})
	if err != nil {
		return remote, err
	}

	for _, item := range resp.Contents {
		remote = append(remote, *item.Key)
	}

	for *resp.IsTruncated {
		resp, err = a.client.ListObjects(&s3.ListObjectsInput{
			Bucket: aws.String(p.Bucket),
			Prefix: aws.String(path),
			Marker: aws.String(remote[len(remote)-1]),
		})

		if err != nil {
			return remote, err
		}

		for _, item := range resp.Contents {
			remote = append(remote, *item.Key)
		}
	}

	return remote, nil
}

func (a *AWS) Invalidate(invalidatePath string) error {
	p := a.plugin
	debug("Invalidating \"%s\"", invalidatePath)
	_, err := a.cfClient.CreateInvalidation(&cloudfront.CreateInvalidationInput{
		DistributionId: aws.String(p.CloudFrontDistribution),
		InvalidationBatch: &cloudfront.InvalidationBatch{
			CallerReference: aws.String(time.Now().Format(time.RFC3339Nano)),
			Paths: &cloudfront.Paths{
				Quantity: aws.Int64(1),
				Items: []*string{
					aws.String(invalidatePath),
				},
			},
		},
	})
	return err
}
<|endoftext|>


<commit_before>package namesys

import (
	"context"
	"strings"
	"time"

	path "github.com/ipfs/go-ipfs/path"

	ds "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore"
	routing "gx/ipfs/QmbkGVaN9W6RYJK4Ws5FvMKXKDqdRQ5snhtaa92qP6L8eU/go-libp2p-routing"
	peer "gx/ipfs/QmfMmLGoKzCHDN7cGgk64PJr4iipzidDRME8HABSJqvmhC/go-libp2p-peer"
	ci "gx/ipfs/QmfWDLQjGjVe4fr5CoztYW2DYYjRysMJrFe1RCsXLPTf46/go-libp2p-crypto"
)

// mpns (a multi-protocol NameSystem) implements generic IPFS naming.
//
// Uses several Resolvers:
// (a) IPFS routing naming: SFS-like PKI names.
// (b) dns domains: resolves using links in DNS TXT records
// (c) proquints: interprets string as the raw byte data.
//
// It can only publish to: (a) IPFS routing naming.
//
type mpns struct {
	resolvers  map[string]resolver
	publishers map[string]Publisher
}

// NewNameSystem will construct the IPFS naming system based on Routing
func NewNameSystem(r routing.ValueStore, ds ds.Datastore, cachesize int) NameSystem {
	return &mpns{
		resolvers: map[string]resolver{
			"dns":      newDNSResolver(),
			"proquint": new(ProquintResolver),
			"dht":      NewRoutingResolver(r, cachesize),
		},
		publishers: map[string]Publisher{
			"/ipns/": NewRoutingPublisher(r, ds),
		},
	}
}

const DefaultResolverCacheTTL = time.Minute

// Resolve implements Resolver.
func (ns *mpns) Resolve(ctx context.Context, name string) (path.Path, error) {
	return ns.ResolveN(ctx, name, DefaultDepthLimit)
}

// ResolveN implements Resolver.
func (ns *mpns) ResolveN(ctx context.Context, name string, depth int) (path.Path, error) {
	if strings.HasPrefix(name, "/ipfs/") {
		return path.ParsePath(name)
	}

	if !strings.HasPrefix(name, "/") {
		return path.ParsePath("/ipfs/" + name)
	}

	return resolve(ctx, ns, name, depth, "/ipns/")
}

// resolveOnce implements resolver.
func (ns *mpns) resolveOnce(ctx context.Context, name string) (path.Path, error) {
	if !strings.HasPrefix(name, "/ipns/") {
		name = "/ipns/" + name
	}
	segments := strings.SplitN(name, "/", 4)
	if len(segments) < 3 || segments[0] != "" {
		log.Warningf("Invalid name syntax for %s", name)
		return "", ErrResolveFailed
	}

	for protocol, resolver := range ns.resolvers {
		log.Debugf("Attempting to resolve %s with %s", segments[2], protocol)
		p, err := resolver.resolveOnce(ctx, segments[2])
		if err == nil {
			if len(segments) > 3 {
				return path.FromSegments("", strings.TrimRight(p.String(), "/"), segments[3])
			} else {
				return p, err
			}
		}
	}
	log.Warningf("No resolver found for %s", name)
	return "", ErrResolveFailed
}

// Publish implements Publisher
func (ns *mpns) Publish(ctx context.Context, name ci.PrivKey, value path.Path) error {
	err := ns.publishers["/ipns/"].Publish(ctx, name, value)
	if err != nil {
		return err
	}
	ns.addToDHTCache(name, value, time.Now().Add(time.Hour*24))
	return nil
}

func (ns *mpns) PublishWithEOL(ctx context.Context, name ci.PrivKey, value path.Path, eol time.Time) error {
	err := ns.publishers["/ipns/"].PublishWithEOL(ctx, name, value, eol)
	if err != nil {
		return err
	}
	ns.addToDHTCache(name, value, eol)
	return nil
}

func (ns *mpns) addToDHTCache(key ci.PrivKey, value path.Path, eol time.Time) {
	var err error
	value, err = path.ParsePath(value.String())
	if err != nil {
		log.Error("could not parse path")
		return
	}

	name, err := peer.IDFromPrivateKey(key)
	if err != nil {
		log.Error("while adding to cache, could not get peerid from private key")
		return
	}

	rr, ok := ns.resolvers["dht"].(*routingResolver)
	if !ok {
		// should never happen, purely for sanity
		log.Panicf("unexpected type %T as DHT resolver.", ns.resolvers["dht"])
	}
	rr.cache.Add(name.Pretty(), cacheEntry{
		val: value,
		eol: eol,
	})
}
<commit_msg>namesys: fix length of self resolve cache<commit_after>package namesys

import (
	"context"
	"strings"
	"time"

	path "github.com/ipfs/go-ipfs/path"

	ds "gx/ipfs/QmRWDav6mzWseLWeYfVd5fvUKiVe9xNH29YfMF438fG364/go-datastore"
	routing "gx/ipfs/QmbkGVaN9W6RYJK4Ws5FvMKXKDqdRQ5snhtaa92qP6L8eU/go-libp2p-routing"
	peer "gx/ipfs/QmfMmLGoKzCHDN7cGgk64PJr4iipzidDRME8HABSJqvmhC/go-libp2p-peer"
	ci "gx/ipfs/QmfWDLQjGjVe4fr5CoztYW2DYYjRysMJrFe1RCsXLPTf46/go-libp2p-crypto"
)

// mpns (a multi-protocol NameSystem) implements generic IPFS naming.
//
// Uses several Resolvers:
// (a) IPFS routing naming: SFS-like PKI names.
// (b) dns domains: resolves using links in DNS TXT records
// (c) proquints: interprets string as the raw byte data.
//
// It can only publish to: (a) IPFS routing naming.
//
type mpns struct {
	resolvers  map[string]resolver
	publishers map[string]Publisher
}

// NewNameSystem will construct the IPFS naming system based on Routing
func NewNameSystem(r routing.ValueStore, ds ds.Datastore, cachesize int) NameSystem {
	return &mpns{
		resolvers: map[string]resolver{
			"dns":      newDNSResolver(),
			"proquint": new(ProquintResolver),
			"dht":      NewRoutingResolver(r, cachesize),
		},
		publishers: map[string]Publisher{
			"/ipns/": NewRoutingPublisher(r, ds),
		},
	}
}

const DefaultResolverCacheTTL = time.Minute

// Resolve implements Resolver.
func (ns *mpns) Resolve(ctx context.Context, name string) (path.Path, error) {
	return ns.ResolveN(ctx, name, DefaultDepthLimit)
}

// ResolveN implements Resolver.
func (ns *mpns) ResolveN(ctx context.Context, name string, depth int) (path.Path, error) {
	if strings.HasPrefix(name, "/ipfs/") {
		return path.ParsePath(name)
	}

	if !strings.HasPrefix(name, "/") {
		return path.ParsePath("/ipfs/" + name)
	}

	return resolve(ctx, ns, name, depth, "/ipns/")
}

// resolveOnce implements resolver.
func (ns *mpns) resolveOnce(ctx context.Context, name string) (path.Path, error) {
	if !strings.HasPrefix(name, "/ipns/") {
		name = "/ipns/" + name
	}
	segments := strings.SplitN(name, "/", 4)
	if len(segments) < 3 || segments[0] != "" {
		log.Warningf("Invalid name syntax for %s", name)
		return "", ErrResolveFailed
	}

	for protocol, resolver := range ns.resolvers {
		log.Debugf("Attempting to resolve %s with %s", segments[2], protocol)
		p, err := resolver.resolveOnce(ctx, segments[2])
		if err == nil {
			if len(segments) > 3 {
				return path.FromSegments("", strings.TrimRight(p.String(), "/"), segments[3])
			} else {
				return p, err
			}
		}
	}
	log.Warningf("No resolver found for %s", name)
	return "", ErrResolveFailed
}

// Publish implements Publisher
func (ns *mpns) Publish(ctx context.Context, name ci.PrivKey, value path.Path) error {
	err := ns.publishers["/ipns/"].Publish(ctx, name, value)
	if err != nil {
		return err
	}
	ns.addToDHTCache(name, value, time.Now().Add(time.Hour*24))
	return nil
}

func (ns *mpns) PublishWithEOL(ctx context.Context, name ci.PrivKey, value path.Path, eol time.Time) error {
	err := ns.publishers["/ipns/"].PublishWithEOL(ctx, name, value, eol)
	if err != nil {
		return err
	}
	ns.addToDHTCache(name, value, eol)
	return nil
}

func (ns *mpns) addToDHTCache(key ci.PrivKey, value path.Path, eol time.Time) {
	var err error
	value, err = path.ParsePath(value.String())
	if err != nil {
		log.Error("could not parse path")
		return
	}

	name, err := peer.IDFromPrivateKey(key)
	if err != nil {
		log.Error("while adding to cache, could not get peerid from private key")
		return
	}

	rr, ok := ns.resolvers["dht"].(*routingResolver)
	if !ok {
		// should never happen, purely for sanity
		log.Panicf("unexpected type %T as DHT resolver.", ns.resolvers["dht"])
	}
	if time.Now().Add(DefaultResolverCacheTTL).Before(eol) {
		eol = time.Now().Add(DefaultResolverCacheTTL)
	}
	rr.cache.Add(name.Pretty(), cacheEntry{
		val: value,
		eol: eol,
	})
}
<|endoftext|>


<commit_before>package node

import (
	"bytes"
	"github.com/AutoRoute/l2"
)

// The layer two protocol takes a layer two device and returns the hash of the
// Public Key of all neighbors it can find.
type NeighborFinder interface {
	Find(l2.FrameReadWriter) <-chan string
}

type layer2 struct{}

func (nf layer2) Find(frw l2.FrameReadWriter) <-chan string {
	c := make(chan string)
	// Broadcast Hash
	broadcastAddr := l2.MacToBytesOrDie("ff:ff:ff:ff:ff:ff")
	localAddr := l2.MacToBytesOrDie("aa:bb:cc:dd:ee:00") // TODO: pass own mac address
	var protocol uint16 = 31337                          // TODO: add real protocol
	var p PublicKey                                      // TODO: pass public key
	publicKeyHash := []byte(p.Hash())
	initFrame := l2.NewEthFrame(broadcastAddr, localAddr, protocol, publicKeyHash)
	_ = frw.WriteFrame(initFrame) // TODO: check errors
	// Process Loop
	go func() {
		for {
			newInstanceFrame, _ := frw.ReadFrame()
			src := newInstanceFrame.Source()
			dest := newInstanceFrame.Destination()
			if newInstanceFrame.Type() != protocol {
				continue // Throw away if protocols don't match
			}
			if bytes.Equal(src, localAddr) {
				continue // Throw away if from me
			}
			if !(bytes.Equal(dest, localAddr) || bytes.Equal(dest, broadcastAddr)) {
				continue // Throw away if it wasn't to me or the broadcast address
			}
			c <- string(newInstanceFrame.Data())
			if bytes.Equal(dest, broadcastAddr) { // Respond if to broadcast addr
				var p PublicKey // TODO: pass public key
				publicKeyHash := []byte(p.Hash())
				initFrame := l2.NewEthFrame(src, localAddr, 31337, publicKeyHash) // TODO: add real protocol
				_ = frw.WriteFrame(initFrame)                                     // TODO: check errors
			}
		}
	}()
	return c
}
<commit_msg>Added print statements for testing.<commit_after>package node

import (
	"bytes"
	"fmt"
	"github.com/AutoRoute/l2"
)

// The layer two protocol takes a layer two device and returns the hash of the
// Public Key of all neighbors it can find.
type NeighborFinder interface {
	Find(l2.FrameReadWriter) <-chan string
}

type layer2 struct{}

func (nf layer2) Find(frw l2.FrameReadWriter) <-chan string {
	c := make(chan string)
	// Broadcast Hash
	broadcastAddr := l2.MacToBytesOrDie("ff:ff:ff:ff:ff:ff")
	localAddr := l2.MacToBytesOrDie("aa:bb:cc:dd:ee:00") // TODO: pass own mac address
	var protocol uint16 = 31337                          // TODO: add real protocol
	var p PublicKey                                      // TODO: pass public key
	publicKeyHash := []byte(p.Hash())
	initFrame := l2.NewEthFrame(broadcastAddr, localAddr, protocol, publicKeyHash)
	fmt.Println("Broadcasting packet.")
	_ = frw.WriteFrame(initFrame) // TODO: check errors
	fmt.Println("Broadcasted packet.")
	// Process Loop
	go func() {
		for {
			fmt.Println("Receiving packet.")
			newInstanceFrame, _ := frw.ReadFrame()
			src := newInstanceFrame.Source()
			dest := newInstanceFrame.Destination()
			fmt.Printf("Received packet from %v.\n", src)
			fmt.Printf("Received packet to %v.\n", dest)
			if newInstanceFrame.Type() != protocol {
				continue // Throw away if protocols don't match
			}
			if bytes.Equal(src, localAddr) {
				continue // Throw away if from me
			}
			if !(bytes.Equal(dest, localAddr) || bytes.Equal(dest, broadcastAddr)) {
				continue // Throw away if it wasn't to me or the broadcast address
			}
			c <- string(newInstanceFrame.Data())
			if bytes.Equal(dest, broadcastAddr) { // Respond if to broadcast addr
				var p PublicKey // TODO: pass public key
				publicKeyHash := []byte(p.Hash())
				initFrame := l2.NewEthFrame(src, localAddr, 31337, publicKeyHash) // TODO: add real protocol
				fmt.Printf("Sending response packet %v.\n", src)
				_ = frw.WriteFrame(initFrame) // TODO: check errors
				fmt.Println("Sent response packet.")
			}
		}
	}()
	return c
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nlopes/slack"
)

func main() {
	slackToken := os.Getenv("ARCHIVEBOT_SLACK_TOKEN")
	api := slack.New(slackToken)
	//api.SetDebug(true)

	channels, err := api.GetChannels(true)
	if err != nil {
		log.Printf("Error when loading channels: %s\n", err)
		return
	}

	var wg sync.WaitGroup
	wg.Add(2)

	go func(c []slack.Channel) {
		defer wg.Done()
		archiveEmptyChannels(api, c)
	}(channels)

	go func(c []slack.Channel) {
		defer wg.Done()
		archiveInactiveChannels(api, c)
	}(channels)

	wg.Wait()
}

func archiveEmptyChannels(api *slack.Slack, c []slack.Channel) {
	empty := filterEmptyChannels(api, c)
	archiveChannels(api, empty, "emptiness")
}

func archiveInactiveChannels(api *slack.Slack, c []slack.Channel) {
	inactive := filterInactiveChannels(api, c)
	archiveChannels(api, inactive, "inactivity")
}

func archiveChannels(api *slack.Slack, c []slack.Channel, reason string) {
	var wg sync.WaitGroup

	for _, channel := range c {
		fmt.Printf("Archiving #%s (%s) due to %s\n", channel.Name, channel.Id, reason)
		wg.Add(1)

		go func(c slack.Channel) {
			defer wg.Done()
			if err := api.ArchiveChannel(c.Id); err != nil {
				message := fmt.Sprintf(
					"Error archiving channel #%s (%s): %s\n", c.Name, c.Id, err)
				log.Printf(message)
				// send error message in a DM to onErrorNotify user/channel
				onErrorNotify := os.Getenv("ARCHIVEBOT_NOTIFY")
				params := slack.PostMessageParameters{}
				api.PostMessage(onErrorNotify, message, params)
			}
		}(channel)
	}

	wg.Wait()
}

func filterEmptyChannels(api *slack.Slack, c []slack.Channel) []slack.Channel {
	empty := []slack.Channel{}
	for _, channel := range c {
		if channel.NumMembers == 0 {
			empty = append(empty, channel)
		}
	}
	return empty
}

type LastChannelMessage struct {
	Channel   slack.Channel
	Timestamp int64
}

func filterInactiveChannels(api *slack.Slack, c []slack.Channel) []slack.Channel {
	inactiveDays, _ := strconv.ParseInt(os.Getenv("ARCHIVEBOT_INACTIVE_DAYS"), 10, 32)
	if inactiveDays == 0 {
		inactiveDays = 30
	}

	timeout := int64(time.Now().Unix()) - (86400 * inactiveDays)
	channels := []slack.Channel{}

	res := make(chan LastChannelMessage)
	for _, channel := range c {
		go func(channel slack.Channel) {
			timestamp, _ := lastMessageTimestamp(api, channel)
			res <- LastChannelMessage{Channel: channel, Timestamp: timestamp}
		}(channel)
	}

	for i := 0; i < len(c); i++ {
		lcm := <-res
		if lcm.Timestamp > 0 && lcm.Timestamp < timeout {
			channels = append(channels, lcm.Channel)
		}
	}

	close(res)
	return channels
}

func lastMessageTimestamp(api *slack.Slack, channel slack.Channel) (int64, error) {
	var latest string

	for {
		historyParams := slack.HistoryParameters{Count: 5}
		if latest != "" {
			historyParams.Latest = latest
		}

		history, err := api.GetChannelHistory(channel.Id, historyParams)

		if err != nil {
			return -1, err
		}

		if len(history.Messages) == 0 {
			return -1, nil
		}

		for _, msg := range history.Messages {
			latest = msg.Msg.Timestamp

			if msg.SubType != "channel_join" && msg.SubType != "channel_leave" {
				msgStamp := strings.Split(msg.Msg.Timestamp, ".")
				if timestamp, err := strconv.ParseInt(msgStamp[0], 10, 32); err == nil {
					return timestamp, nil
				}
			}
		}
	}
}
<commit_msg>Allow for ARCHIVEBOT_NOTIFY to not be set.<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nlopes/slack"
)

func main() {
	slackToken := os.Getenv("ARCHIVEBOT_SLACK_TOKEN")
	api := slack.New(slackToken)
	//api.SetDebug(true)

	channels, err := api.GetChannels(true)
	if err != nil {
		log.Printf("Error when loading channels: %s\n", err)
		return
	}

	var wg sync.WaitGroup
	wg.Add(2)

	go func(c []slack.Channel) {
		defer wg.Done()
		archiveEmptyChannels(api, c)
	}(channels)

	go func(c []slack.Channel) {
		defer wg.Done()
		archiveInactiveChannels(api, c)
	}(channels)

	wg.Wait()
}

func archiveEmptyChannels(api *slack.Slack, c []slack.Channel) {
	empty := filterEmptyChannels(api, c)
	archiveChannels(api, empty, "emptiness")
}

func archiveInactiveChannels(api *slack.Slack, c []slack.Channel) {
	inactive := filterInactiveChannels(api, c)
	archiveChannels(api, inactive, "inactivity")
}

func archiveChannels(api *slack.Slack, c []slack.Channel, reason string) {
	var wg sync.WaitGroup

	for _, channel := range c {
		fmt.Printf("Archiving #%s (%s) due to %s\n", channel.Name, channel.Id, reason)
		wg.Add(1)

		go func(c slack.Channel) {
			defer wg.Done()
			if err := api.ArchiveChannel(c.Id); err != nil {
				message := fmt.Sprintf(
					"Error archiving channel #%s (%s): %s\n", c.Name, c.Id, err)
				log.Printf(message)
				// send error message in a DM to onErrorNotify user/channel
				onErrorNotify := os.Getenv("ARCHIVEBOT_NOTIFY")
				if onErrorNotify != "" {
					params := slack.PostMessageParameters{}
					api.PostMessage(onErrorNotify, message, params)
				}
			}
		}(channel)
	}

	wg.Wait()
}

func filterEmptyChannels(api *slack.Slack, c []slack.Channel) []slack.Channel {
	empty := []slack.Channel{}
	for _, channel := range c {
		if channel.NumMembers == 0 {
			empty = append(empty, channel)
		}
	}
	return empty
}

type LastChannelMessage struct {
	Channel   slack.Channel
	Timestamp int64
}

func filterInactiveChannels(api *slack.Slack, c []slack.Channel) []slack.Channel {
	inactiveDays, _ := strconv.ParseInt(os.Getenv("ARCHIVEBOT_INACTIVE_DAYS"), 10, 32)
	if inactiveDays == 0 {
		inactiveDays = 30
	}

	timeout := int64(time.Now().Unix()) - (86400 * inactiveDays)
	channels := []slack.Channel{}

	res := make(chan LastChannelMessage)
	for _, channel := range c {
		go func(channel slack.Channel) {
			timestamp, _ := lastMessageTimestamp(api, channel)
			res <- LastChannelMessage{Channel: channel, Timestamp: timestamp}
		}(channel)
	}

	for i := 0; i < len(c); i++ {
		lcm := <-res
		if lcm.Timestamp > 0 && lcm.Timestamp < timeout {
			channels = append(channels, lcm.Channel)
		}
	}

	close(res)
	return channels
}

func lastMessageTimestamp(api *slack.Slack, channel slack.Channel) (int64, error) {
	var latest string

	for {
		historyParams := slack.HistoryParameters{Count: 5}
		if latest != "" {
			historyParams.Latest = latest
		}

		history, err := api.GetChannelHistory(channel.Id, historyParams)

		if err != nil {
			return -1, err
		}

		if len(history.Messages) == 0 {
			return -1, nil
		}

		for _, msg := range history.Messages {
			latest = msg.Msg.Timestamp

			if msg.SubType != "channel_join" && msg.SubType != "channel_leave" {
				msgStamp := strings.Split(msg.Msg.Timestamp, ".")
				if timestamp, err := strconv.ParseInt(msgStamp[0], 10, 32); err == nil {
					return timestamp, nil
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/hako/durafmt"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/tucnak/telebot"
)

const (
	commandStart = "/start"
	commandStop  = "/stop"
	commandHelp  = "/help"
	commandChats = "/chats"

	commandStatus     = "/status"
	commandAlerts     = "/alerts"
	commandSilences   = "/silences"
	commandSilenceAdd = "/silence_add"
	commandSilence    = "/silence"
	commandSilenceDel = "/silence_del"

	responseStart = "Hey, %s! I will now keep you up to date!\n" + commandHelp
	responseStop  = "Alright, %s! I won't talk to you again.\n" + commandHelp
	responseHelp  = `
I'm a Prometheus AlertManager telegram for Telegram. I will notify you about alerts.
You can also ask me about my ` + commandStatus + `, ` + commandAlerts + ` & ` + commandSilences + `

Available commands:
` + commandStart + ` - Subscribe for alerts.
` + commandStop + ` - Unsubscribe for alerts.
` + commandStatus + ` - Print the current status.
` + commandAlerts + ` - List all alerts.
` + commandSilences + ` - List all silences.
`
)

// BotChatStore is all the Bot needs to store and read
type BotChatStore interface {
	List() ([]telebot.Chat, error)
	Add(telebot.Chat) error
	Remove(telebot.Chat) error
}

// Bot runs the alertmanager telegram
type Bot struct {
	logger          log.Logger
	telegram        *telebot.Bot
	chats           BotChatStore
	addr            string
	admin           int
	alertmanager    *url.URL
	commandsCounter *prometheus.CounterVec
	webhooksCounter prometheus.Counter
}

// NewBot creates a Bot with the UserStore and telegram telegram
func NewBot(chats BotChatStore, addr string, alertmanager *url.URL, telegramToken string, telegramAdmin int, opts ...func(b *Bot)) (*Bot, error) {
	bot, err := telebot.NewBot(telegramToken)
	if err != nil {
		return nil, err
	}

	commandsCounter := prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "alertmanagerbot",
		Name:      "commands_total",
		Help:      "Number of commands received by command name",
	}, []string{"command"})
	webhooksCounter := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "alertmanagerbot",
		Name:      "webhooks_total",
		Help:      "Number of webhooks received by this bot",
	})
	prometheus.MustRegister(commandsCounter, webhooksCounter)

	b := &Bot{
		logger:          log.NewNopLogger(),
		telegram:        bot,
		chats:           chats,
		addr:            addr,
		admin:           telegramAdmin,
		alertmanager:    alertmanager,
		commandsCounter: commandsCounter,
		webhooksCounter: webhooksCounter,
	}

	for _, opt := range opts {
		opt(b)
	}

	return b, nil
}

func BotLogger(l log.Logger) func(b *Bot) {
	return func(b *Bot) {
		b.logger = l
	}
}

// RunWebserver starts a http server and listens for messages to send to the users
func (b *Bot) RunWebserver() {
	messages := make(chan string, 100)

	http.HandleFunc("/", HandleWebhook(b.logger, b.webhooksCounter, messages))
	http.Handle("/metrics", promhttp.Handler())
	http.HandleFunc("/health", handleHealth)
	http.HandleFunc("/healthz", handleHealth)

	go b.sendWebhook(messages)

	err := http.ListenAndServe(b.addr, nil)
	level.Error(b.logger).Log("err", err)
	os.Exit(1)
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
}

// sendWebhook sends messages received via webhook to all subscribed users
func (b *Bot) sendWebhook(messages <-chan string) {
	for m := range messages {
		chats, err := b.chats.List()
		if err != nil {
			level.Error(b.logger).Log("msg", "failed to get chat list from store", "err", err)
			continue
		}

		for _, chat := range chats {
			b.telegram.SendMessage(chat, m, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
		}
	}
}

// SendAdminMessage to the admin's ID with a message
func (b *Bot) SendAdminMessage(adminID int, message string) {
	b.telegram.SendMessage(telebot.User{ID: adminID}, message, nil)
}

// Run the telegram and listen to messages send to the telegram
func (b *Bot) Run() {
	commandSuffix := fmt.Sprintf("@%s", b.telegram.Identity.Username)

	commands := map[string]func(message telebot.Message){
		commandStart:    b.handleStart,
		commandStop:     b.handleStop,
		commandHelp:     b.handleHelp,
		commandChats:    b.handleChats,
		commandStatus:   b.handleStatus,
		commandAlerts:   b.handleAlerts,
		commandSilences: b.handleSilences,
	}

	// init counters with 0
	for command := range commands {
		b.commandsCounter.WithLabelValues(command).Add(0)
	}

	messages := make(chan telebot.Message, 100)
	b.telegram.Listen(messages, time.Second)

	for message := range messages {
		if message.IsService() {
			continue
		}

		if message.Sender.ID != b.admin {
			b.commandsCounter.WithLabelValues("dropped").Inc()
			level.Info(b.logger).Log(
				"msg", "dropped message from forbidden sender",
				"sender_id", message.Sender.ID,
				"sender_username", message.Sender.Username,
			)
			continue
		}

		b.telegram.SendChatAction(message.Chat, telebot.Typing)

		// Remove the command suffix from the text, /help@BotName => /help
		text := strings.Replace(message.Text, commandSuffix, "", -1)
		level.Debug(b.logger).Log("msg", "message received", "text", text)

		// Get the corresponding handler from the map by the commands text
		if handler, ok := commands[text]; ok {
			b.commandsCounter.WithLabelValues(text).Inc()
			handler(message)
		} else {
			b.commandsCounter.WithLabelValues("incomprehensible").Inc()
			b.telegram.SendMessage(
				message.Chat,
				"Sorry, I don't understand...",
				nil,
			)
		}
	}
}

func (b *Bot) handleStart(message telebot.Message) {
	if err := b.chats.Add(message.Chat); err != nil {
		level.Warn(b.logger).Log("msg", "failed to add chat to chat store", "err", err)
		b.telegram.SendMessage(message.Sender, "I can't add this chat to the subscribers list, see logs", nil)
		return
	}

	b.telegram.SendMessage(message.Chat, fmt.Sprintf(responseStart, message.Sender.FirstName), nil)
	level.Info(b.logger).Log(
		"user subscribed",
		"username", message.Sender.Username,
		"user_id", message.Sender.ID,
	)
}

func (b *Bot) handleStop(message telebot.Message) {
	if err := b.chats.Remove(message.Chat); err != nil {
		level.Warn(b.logger).Log("msg", "failed to remove chat from chat store", "err", err)
		b.telegram.SendMessage(message.Sender, "I can't remove this chat from the subscribers list, see logs", nil)
		return
	}

	b.telegram.SendMessage(message.Chat, fmt.Sprintf(responseStop, message.Sender.FirstName), nil)
	level.Info(b.logger).Log(
		"user unsubscribed",
		"username", message.Sender.Username,
		"user_id", message.Sender.ID,
	)
}

func (b *Bot) handleHelp(message telebot.Message) {
	b.telegram.SendMessage(message.Chat, responseHelp, nil)
}

func (b *Bot) handleChats(message telebot.Message) {
	chats, err := b.chats.List()
	if err != nil {
		level.Warn(b.logger).Log("msg", "failed to remove chat from chat store", "err", err)
		b.telegram.SendMessage(message.Sender, "I can't remove this chat from the subscribers list, see logs", nil)
		return
	}

	list := ""
	for _, chat := range chats {
		if chat.IsGroupChat() {
			list = list + fmt.Sprintf("@%s\n", chat.Title)
		} else {
			list = list + fmt.Sprintf("@%s\n", chat.Username)
		}
	}

	b.telegram.SendMessage(message.Chat, "Currently these chat have subscribed:\n"+list, nil)
}

func (b *Bot) handleStatus(message telebot.Message) {
	s, err := status(b.logger, b.alertmanager.String())
	if err != nil {
		b.telegram.SendMessage(message.Chat, fmt.Sprintf("failed to get status... %v", err), nil)
		return
	}

	uptime := durafmt.Parse(time.Since(s.Data.Uptime))
	uptimeBot := durafmt.Parse(time.Since(StartTime))

	b.telegram.SendMessage(
		message.Chat,
		fmt.Sprintf(
			"*AlertManager*\nVersion: %s\nUptime: %s\n*AlertManager Bot*\nVersion: %s\nUptime: %s",
			s.Data.VersionInfo.Version,
			uptime,
			Revision,
			uptimeBot,
		),
		&telebot.SendOptions{ParseMode: telebot.ModeMarkdown},
	)
}

func (b *Bot) handleAlerts(message telebot.Message) {
	alerts, err := listAlerts(b.logger, b.alertmanager.String())
	if err != nil {
		b.telegram.SendMessage(message.Chat, fmt.Sprintf("failed to list alerts... %v", err), nil)
		return
	}

	if len(alerts) == 0 {
		b.telegram.SendMessage(message.Chat, "No alerts right now! 🎉", nil)
		return
	}

	var out string
	for _, a := range alerts {
		out = out + AlertMessage(a) + "\n"
	}

	b.telegram.SendMessage(message.Chat, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
}

func (b *Bot) handleSilences(message telebot.Message) {
	silences, err := listSilences(b.logger, b.alertmanager.String())
	if err != nil {
		b.telegram.SendMessage(message.Chat, fmt.Sprintf("failed to list silences... %v", err), nil)
		return
	}

	if len(silences) == 0 {
		b.telegram.SendMessage(message.Chat, "No silences right now.", nil)
		return
	}

	var out string
	for _, silence := range silences {
		out = out + SilenceMessage(silence) + "\n"
	}

	b.telegram.SendMessage(message.Chat, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
}
<commit_msg>Split the message at their first space<commit_after>package main

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/hako/durafmt"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/tucnak/telebot"
)

const (
	commandStart = "/start"
	commandStop  = "/stop"
	commandHelp  = "/help"
	commandChats = "/chats"

	commandStatus     = "/status"
	commandAlerts     = "/alerts"
	commandSilences   = "/silences"
	commandSilenceAdd = "/silence_add"
	commandSilence    = "/silence"
	commandSilenceDel = "/silence_del"

	responseStart = "Hey, %s! I will now keep you up to date!\n" + commandHelp
	responseStop  = "Alright, %s! I won't talk to you again.\n" + commandHelp
	responseHelp  = `
I'm a Prometheus AlertManager telegram for Telegram. I will notify you about alerts.
You can also ask me about my ` + commandStatus + `, ` + commandAlerts + ` & ` + commandSilences + `

Available commands:
` + commandStart + ` - Subscribe for alerts.
` + commandStop + ` - Unsubscribe for alerts.
` + commandStatus + ` - Print the current status.
` + commandAlerts + ` - List all alerts.
` + commandSilences + ` - List all silences.
`
)

// BotChatStore is all the Bot needs to store and read
type BotChatStore interface {
	List() ([]telebot.Chat, error)
	Add(telebot.Chat) error
	Remove(telebot.Chat) error
}

// Bot runs the alertmanager telegram
type Bot struct {
	logger          log.Logger
	telegram        *telebot.Bot
	chats           BotChatStore
	addr            string
	admin           int
	alertmanager    *url.URL
	commandsCounter *prometheus.CounterVec
	webhooksCounter prometheus.Counter
}

// NewBot creates a Bot with the UserStore and telegram telegram
func NewBot(chats BotChatStore, addr string, alertmanager *url.URL, telegramToken string, telegramAdmin int, opts ...func(b *Bot)) (*Bot, error) {
	bot, err := telebot.NewBot(telegramToken)
	if err != nil {
		return nil, err
	}

	commandsCounter := prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "alertmanagerbot",
		Name:      "commands_total",
		Help:      "Number of commands received by command name",
	}, []string{"command"})
	webhooksCounter := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "alertmanagerbot",
		Name:      "webhooks_total",
		Help:      "Number of webhooks received by this bot",
	})
	prometheus.MustRegister(commandsCounter, webhooksCounter)

	b := &Bot{
		logger:          log.NewNopLogger(),
		telegram:        bot,
		chats:           chats,
		addr:            addr,
		admin:           telegramAdmin,
		alertmanager:    alertmanager,
		commandsCounter: commandsCounter,
		webhooksCounter: webhooksCounter,
	}

	for _, opt := range opts {
		opt(b)
	}

	return b, nil
}

func BotLogger(l log.Logger) func(b *Bot) {
	return func(b *Bot) {
		b.logger = l
	}
}

// RunWebserver starts a http server and listens for messages to send to the users
func (b *Bot) RunWebserver() {
	messages := make(chan string, 100)

	http.HandleFunc("/", HandleWebhook(b.logger, b.webhooksCounter, messages))
	http.Handle("/metrics", promhttp.Handler())
	http.HandleFunc("/health", handleHealth)
	http.HandleFunc("/healthz", handleHealth)

	go b.sendWebhook(messages)

	err := http.ListenAndServe(b.addr, nil)
	level.Error(b.logger).Log("err", err)
	os.Exit(1)
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
}

// sendWebhook sends messages received via webhook to all subscribed users
func (b *Bot) sendWebhook(messages <-chan string) {
	for m := range messages {
		chats, err := b.chats.List()
		if err != nil {
			level.Error(b.logger).Log("msg", "failed to get chat list from store", "err", err)
			continue
		}

		for _, chat := range chats {
			b.telegram.SendMessage(chat, m, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
		}
	}
}

// SendAdminMessage to the admin's ID with a message
func (b *Bot) SendAdminMessage(adminID int, message string) {
	b.telegram.SendMessage(telebot.User{ID: adminID}, message, nil)
}

// Run the telegram and listen to messages send to the telegram
func (b *Bot) Run() {
	commandSuffix := fmt.Sprintf("@%s", b.telegram.Identity.Username)

	commands := map[string]func(message telebot.Message){
		commandStart:    b.handleStart,
		commandStop:     b.handleStop,
		commandHelp:     b.handleHelp,
		commandChats:    b.handleChats,
		commandStatus:   b.handleStatus,
		commandAlerts:   b.handleAlerts,
		commandSilences: b.handleSilences,
	}

	// init counters with 0
	for command := range commands {
		b.commandsCounter.WithLabelValues(command).Add(0)
	}

	messages := make(chan telebot.Message, 100)
	b.telegram.Listen(messages, time.Second)

	for message := range messages {
		if message.IsService() {
			continue
		}

		if message.Sender.ID != b.admin {
			b.commandsCounter.WithLabelValues("dropped").Inc()
			level.Info(b.logger).Log(
				"msg", "dropped message from forbidden sender",
				"sender_id", message.Sender.ID,
				"sender_username", message.Sender.Username,
			)
			continue
		}

		b.telegram.SendChatAction(message.Chat, telebot.Typing)

		// Remove the command suffix from the text, /help@BotName => /help
		text := strings.Replace(message.Text, commandSuffix, "", -1)
		text = strings.Split(text, " ")[0]
		level.Debug(b.logger).Log("msg", "message received", "text", text)

		// Get the corresponding handler from the map by the commands text
		if handler, ok := commands[text]; ok {
			b.commandsCounter.WithLabelValues(text).Inc()
			handler(message)
		} else {
			b.commandsCounter.WithLabelValues("incomprehensible").Inc()
			b.telegram.SendMessage(
				message.Chat,
				"Sorry, I don't understand...",
				nil,
			)
		}
	}
}

func (b *Bot) handleStart(message telebot.Message) {
	if err := b.chats.Add(message.Chat); err != nil {
		level.Warn(b.logger).Log("msg", "failed to add chat to chat store", "err", err)
		b.telegram.SendMessage(message.Sender, "I can't add this chat to the subscribers list, see logs", nil)
		return
	}

	b.telegram.SendMessage(message.Chat, fmt.Sprintf(responseStart, message.Sender.FirstName), nil)
	level.Info(b.logger).Log(
		"user subscribed",
		"username", message.Sender.Username,
		"user_id", message.Sender.ID,
	)
}

func (b *Bot) handleStop(message telebot.Message) {
	if err := b.chats.Remove(message.Chat); err != nil {
		level.Warn(b.logger).Log("msg", "failed to remove chat from chat store", "err", err)
		b.telegram.SendMessage(message.Sender, "I can't remove this chat from the subscribers list, see logs", nil)
		return
	}

	b.telegram.SendMessage(message.Chat, fmt.Sprintf(responseStop, message.Sender.FirstName), nil)
	level.Info(b.logger).Log(
		"user unsubscribed",
		"username", message.Sender.Username,
		"user_id", message.Sender.ID,
	)
}

func (b *Bot) handleHelp(message telebot.Message) {
	b.telegram.SendMessage(message.Chat, responseHelp, nil)
}

func (b *Bot) handleChats(message telebot.Message) {
	chats, err := b.chats.List()
	if err != nil {
		level.Warn(b.logger).Log("msg", "failed to remove chat from chat store", "err", err)
		b.telegram.SendMessage(message.Sender, "I can't remove this chat from the subscribers list, see logs", nil)
		return
	}

	list := ""
	for _, chat := range chats {
		if chat.IsGroupChat() {
			list = list + fmt.Sprintf("@%s\n", chat.Title)
		} else {
			list = list + fmt.Sprintf("@%s\n", chat.Username)
		}
	}

	b.telegram.SendMessage(message.Chat, "Currently these chat have subscribed:\n"+list, nil)
}

func (b *Bot) handleStatus(message telebot.Message) {
	s, err := status(b.logger, b.alertmanager.String())
	if err != nil {
		b.telegram.SendMessage(message.Chat, fmt.Sprintf("failed to get status... %v", err), nil)
		return
	}

	uptime := durafmt.Parse(time.Since(s.Data.Uptime))
	uptimeBot := durafmt.Parse(time.Since(StartTime))

	b.telegram.SendMessage(
		message.Chat,
		fmt.Sprintf(
			"*AlertManager*\nVersion: %s\nUptime: %s\n*AlertManager Bot*\nVersion: %s\nUptime: %s",
			s.Data.VersionInfo.Version,
			uptime,
			Revision,
			uptimeBot,
		),
		&telebot.SendOptions{ParseMode: telebot.ModeMarkdown},
	)
}

func (b *Bot) handleAlerts(message telebot.Message) {
	alerts, err := listAlerts(b.logger, b.alertmanager.String())
	if err != nil {
		b.telegram.SendMessage(message.Chat, fmt.Sprintf("failed to list alerts... %v", err), nil)
		return
	}

	if len(alerts) == 0 {
		b.telegram.SendMessage(message.Chat, "No alerts right now! 🎉", nil)
		return
	}

	var out string
	for _, a := range alerts {
		out = out + AlertMessage(a) + "\n"
	}

	b.telegram.SendMessage(message.Chat, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
}

func (b *Bot) handleSilences(message telebot.Message) {
	silences, err := listSilences(b.logger, b.alertmanager.String())
	if err != nil {
		b.telegram.SendMessage(message.Chat, fmt.Sprintf("failed to list silences... %v", err), nil)
		return
	}

	if len(silences) == 0 {
		b.telegram.SendMessage(message.Chat, "No silences right now.", nil)
		return
	}

	var out string
	for _, silence := range silences {
		out = out + SilenceMessage(silence) + "\n"
	}

	b.telegram.SendMessage(message.Chat, out, &telebot.SendOptions{ParseMode: telebot.ModeMarkdown})
}
<|endoftext|>


<commit_before>package main

import (
    "github.com/Seklfreak/Robyul2/cache"
    "github.com/Seklfreak/Robyul2/helpers"
    Logger "github.com/Seklfreak/Robyul2/logger"
    "github.com/Seklfreak/Robyul2/metrics"
    "github.com/Seklfreak/Robyul2/modules"
    "github.com/Seklfreak/Robyul2/ratelimits"
    "github.com/getsentry/raven-go"
    "github.com/bwmarrin/discordgo"
    "regexp"
    "strings"
    "time"
    "github.com/Seklfreak/Robyul2/emojis"
    "fmt"
)

// BotOnReady gets called after the gateway connected
func BotOnReady(session *discordgo.Session, event *discordgo.Ready) {
    Logger.INFO.L("bot", "Connected to discord!")
    Logger.VERBOSE.L("bot", "Invite link: "+ fmt.Sprintf(
        "https://discordapp.com/oauth2/authorize?client_id=%s&scope=bot&permissions=%s",
        helpers.GetConfig().Path("discord.id").Data().(string),
        helpers.GetConfig().Path("discord.perms").Data().(string),
    ))

    // Cache the session
    cache.SetSession(session)

    // Load and init all modules
    modules.Init(session)

    // Run async worker for guild changes
    go helpers.GuildSettingsUpdater()

    // Run async game-changer
    go changeGameInterval(session)

    // Run auto-leaver for non-beta guilds
    //go autoLeaver(session)

    // Run ratelimiter
    ratelimits.Container.Init()

    go func() {
        time.Sleep(3 * time.Second)

        configName := helpers.GetConfig().Path("bot.name").Data().(string)
        configAvatar := helpers.GetConfig().Path("bot.avatar").Data().(string)

        // Change avatar if desired
        if configAvatar != "" && configAvatar != session.State.User.Avatar {
            session.UserUpdate(
                "",
                "",
                session.State.User.Username,
                configAvatar,
                "",
            )
        }

        // Change name if desired
        if configName != "" && configName != session.State.User.Username {
            session.UserUpdate(
                "",
                "",
                configName,
                session.State.User.Avatar,
                "",
            )
        }
    }()

    // Run async game-changer
    //go changeGameInterval(session)

    // Run auto-leaver for non-beta guilds
    //go autoLeaver(session)
}

func BotOnGuildMemberAdd(session *discordgo.Session, member *discordgo.GuildMemberAdd) {
    modules.CallExtendedPluginOnGuildMemberAdd(
        member.Member,
    )
}

func BotOnGuildMemberRemove(session *discordgo.Session, member *discordgo.GuildMemberRemove) {
    modules.CallExtendedPluginOnGuildMemberRemove(
        member.Member,
    )
}

func BotOnGuildBanAdd(session *discordgo.Session, user *discordgo.GuildBanAdd) {
    modules.CallExtendedPluginOnGuildBanAdd(
        user,
    )
}

func BotOnGuildBanRemove(session *discordgo.Session, user *discordgo.GuildBanRemove) {
    modules.CallExtendedPluginOnGuildBanRemove(
        user,
    )
}

// BotOnMessageCreate gets called after a new message was sent
// This will be called after *every* message on *every* server so it should die as soon as possible
// or spawn costly work inside of coroutines.
func BotOnMessageCreate(session *discordgo.Session, message *discordgo.MessageCreate) {
    // Ignore other bots and @everyone/@here
    if message.Author.Bot || message.MentionEveryone {
        return
    }

    // Get the channel
    // Ignore the event if we cannot resolve the channel
    channel, err := cache.Channel(message.ChannelID)
    if err != nil {
        go raven.CaptureError(err, map[string]string{})
        return
    }

    // We only do things in guilds.
    // Get a friend already and stop chatting with bots
    if channel.IsPrivate {
        // Track usage
        metrics.CleverbotRequests.Add(1)

        // Mark typing
        session.ChannelTyping(message.ChannelID)

        // Prepare content for editing
        msg := message.Content

        /// Remove our @mention
        msg = strings.Replace(msg, "<@"+session.State.User.ID+">", "", -1)

        // Trim message
        msg = strings.TrimSpace(msg)

        // Resolve other @mentions before sending the message
        for _, user := range message.Mentions {
            msg = strings.Replace(msg, "<@"+user.ID+">", user.Username, -1)
        }

        // Remove smileys
        msg = regexp.MustCompile(`:\w+:`).ReplaceAllString(msg, "")

        // Send to cleverbot
        helpers.CleverbotSend(session, channel.ID, msg)
        return
    }

    // Check if the message contains @mentions for us
    if strings.HasPrefix(message.Content, "<@") && len(message.Mentions) > 0 && message.Mentions[0].ID == session.State.User.ID {
        // Consume a key for this action
        e := ratelimits.Container.Drain(1, message.Author.ID)
        if e != nil {
            return
        }

        // Prepare content for editing
        msg := message.Content

        /// Remove our @mention
        msg = strings.Replace(msg, "<@"+session.State.User.ID+">", "", -1)

        // Trim message
        msg = strings.TrimSpace(msg)

        // Convert to []byte before matching
        bmsg := []byte(msg)

        // Match against common task patterns
        // Send to cleverbot if nothing matches
        switch {
        case regexp.MustCompile("(?i)^HELP.*").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            sendHelp(message)
            return

        case regexp.MustCompile("(?i)^PREFIX.*").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            prefix := helpers.GetPrefixForServer(channel.GuildID)
            if prefix == "" {
                cache.GetSession().ChannelMessageSend(
                    channel.ID,
                    helpers.GetText("bot.prefix.not-set"),
                )
            }

            cache.GetSession().ChannelMessageSend(
                channel.ID,
                helpers.GetTextF("bot.prefix.is", prefix),
            )
            return

        case regexp.MustCompile("(?i)^REFRESH CHAT SESSION$").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            helpers.RequireAdmin(message.Message, func() {
                // Refresh cleverbot session
                helpers.CleverbotRefreshSession(channel.ID)
                cache.GetSession().ChannelMessageSend(channel.ID, helpers.GetText("bot.cleverbot.refreshed"))
            })
            return

        case regexp.MustCompile("(?i)^SET PREFIX (.){1,25}$").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            helpers.RequireAdmin(message.Message, func() {
                // Extract prefix
                prefix := strings.Fields(regexp.MustCompile("(?i)^SET PREFIX\\s").ReplaceAllString(msg, ""))[0]

                // Set new prefix
                err := helpers.SetPrefixForServer(
                    channel.GuildID,
                    prefix,
                )

                if err != nil {
                    helpers.SendError(message.Message, err)
                } else {
                    cache.GetSession().ChannelMessageSend(channel.ID, helpers.GetTextF("bot.prefix.saved", prefix))
                }
            })
            return

        default:
            // Track usage
            metrics.CleverbotRequests.Add(1)

            // Mark typing
            session.ChannelTyping(message.ChannelID)

            // Resolve other @mentions before sending the message
            for _, user := range message.Mentions {
                msg = strings.Replace(msg, "<@"+user.ID+">", user.Username, -1)
            }

            // Remove smileys
            msg = regexp.MustCompile(`:\w+:`).ReplaceAllString(msg, "")

            // Send to cleverbot
            helpers.CleverbotSend(session, channel.ID, msg)
            return
        }
    }

    modules.CallExtendedPlugin(
        message.Content,
        message.Message,
    )

    // Only continue if a prefix is set
    prefix := helpers.GetPrefixForServer(channel.GuildID)
    if prefix == "" {
        return
    }

    // Check if the message is prefixed for us
    // If not exit
    if !strings.HasPrefix(message.Content, prefix) {
        return
    }

    // Check if the user is allowed to request commands
    if !ratelimits.Container.HasKeys(message.Author.ID) && !helpers.IsBotAdmin(message.Author.ID) {
        session.ChannelMessageSend(message.ChannelID, helpers.GetTextF("bot.ratelimit.hit", message.Author.ID))

        ratelimits.Container.Set(message.Author.ID, -1)
        return
    }

    // Split the message into parts
    parts := strings.Fields(message.Content)

    // Save a sanitized version of the command (no prefix)
    cmd := strings.Replace(parts[0], prefix, "", 1)

    // Check if the user calls for help
    if cmd == "h" || cmd == "help" {
        metrics.CommandsExecuted.Add(1)
        sendHelp(message)
        return
    }

    // Separate arguments from the command
    content := strings.TrimSpace(strings.Replace(message.Content, prefix+cmd, "", -1))

    // Check if a module matches said command
    modules.CallBotPlugin(cmd, content, message.Message)

    // Check if a trigger matches
    modules.CallTriggerPlugin(cmd, content, message.Message)
}

// BotOnReactionAdd gets called after a reaction is added
// This will be called after *every* reaction added on *every* server so it
// should die as soon as possible or spawn costly work inside of coroutines.
// This is currently used for the *poll* plugin.
func BotOnReactionAdd(session *discordgo.Session, reaction *discordgo.MessageReactionAdd) {
    modules.CallExtendedPluginOnReactionAdd(reaction)

    if reaction.UserID == session.State.User.ID {
        return
    }

    channel, err := session.State.Channel(reaction.ChannelID)
    if err != nil {
        return
    }
    if emojis.ToNumber(reaction.Emoji.Name) == -1 {
        //session.MessageReactionRemove(reaction.ChannelID, reaction.MessageID, reaction.Emoji.Name, reaction.UserID)
        return
    }
    if helpers.VotePollIfItsOne(channel.GuildID, reaction.MessageReaction) {
        helpers.UpdatePollMsg(channel.GuildID, reaction.MessageID)
    }

}

func BotOnReactionRemove(session *discordgo.Session, reaction *discordgo.MessageReactionRemove) {
    modules.CallExtendedPluginOnReactionRemove(reaction)
}

func sendHelp(message *discordgo.MessageCreate) {
    cache.GetSession().ChannelMessageSend(
        message.ChannelID,
        helpers.GetTextF("bot.help", message.Author.ID),
    )
}

// Changes the game interval every 10 seconds after called
func changeGameInterval(session *discordgo.Session) {
    for {
        users := make(map[string]string)
        guilds := session.State.Guilds

        for _, guild := range guilds {
            lastAfterMemberId := ""
            for {
                members, err := session.GuildMembers(guild.ID, lastAfterMemberId, 1000)
                if len(members) <= 0 {
                    break
                }
                lastAfterMemberId = members[len(members)-1].User.ID
                helpers.Relax(err)
                for _, u := range members {
                    users[u.User.ID] = u.User.Username
                }
            }
        }

        err := session.UpdateStatus(0, fmt.Sprintf("with %d people on %d servers", len(users), len(guilds)))
        if err != nil {
            raven.CaptureError(err, map[string]string{})
        }

        time.Sleep(1 * time.Hour)
    }
}
<commit_msg>[core] log commands<commit_after>package main

import (
    "github.com/Seklfreak/Robyul2/cache"
    "github.com/Seklfreak/Robyul2/helpers"
    Logger "github.com/Seklfreak/Robyul2/logger"
    "github.com/Seklfreak/Robyul2/metrics"
    "github.com/Seklfreak/Robyul2/modules"
    "github.com/Seklfreak/Robyul2/ratelimits"
    "github.com/getsentry/raven-go"
    "github.com/bwmarrin/discordgo"
    "regexp"
    "strings"
    "time"
    "github.com/Seklfreak/Robyul2/emojis"
    "fmt"
)

// BotOnReady gets called after the gateway connected
func BotOnReady(session *discordgo.Session, event *discordgo.Ready) {
    Logger.INFO.L("bot", "Connected to discord!")
    Logger.VERBOSE.L("bot", "Invite link: "+ fmt.Sprintf(
        "https://discordapp.com/oauth2/authorize?client_id=%s&scope=bot&permissions=%s",
        helpers.GetConfig().Path("discord.id").Data().(string),
        helpers.GetConfig().Path("discord.perms").Data().(string),
    ))

    // Cache the session
    cache.SetSession(session)

    // Load and init all modules
    modules.Init(session)

    // Run async worker for guild changes
    go helpers.GuildSettingsUpdater()

    // Run async game-changer
    go changeGameInterval(session)

    // Run auto-leaver for non-beta guilds
    //go autoLeaver(session)

    // Run ratelimiter
    ratelimits.Container.Init()

    go func() {
        time.Sleep(3 * time.Second)

        configName := helpers.GetConfig().Path("bot.name").Data().(string)
        configAvatar := helpers.GetConfig().Path("bot.avatar").Data().(string)

        // Change avatar if desired
        if configAvatar != "" && configAvatar != session.State.User.Avatar {
            session.UserUpdate(
                "",
                "",
                session.State.User.Username,
                configAvatar,
                "",
            )
        }

        // Change name if desired
        if configName != "" && configName != session.State.User.Username {
            session.UserUpdate(
                "",
                "",
                configName,
                session.State.User.Avatar,
                "",
            )
        }
    }()

    // Run async game-changer
    //go changeGameInterval(session)

    // Run auto-leaver for non-beta guilds
    //go autoLeaver(session)
}

func BotOnGuildMemberAdd(session *discordgo.Session, member *discordgo.GuildMemberAdd) {
    modules.CallExtendedPluginOnGuildMemberAdd(
        member.Member,
    )
}

func BotOnGuildMemberRemove(session *discordgo.Session, member *discordgo.GuildMemberRemove) {
    modules.CallExtendedPluginOnGuildMemberRemove(
        member.Member,
    )
}

func BotOnGuildBanAdd(session *discordgo.Session, user *discordgo.GuildBanAdd) {
    modules.CallExtendedPluginOnGuildBanAdd(
        user,
    )
}

func BotOnGuildBanRemove(session *discordgo.Session, user *discordgo.GuildBanRemove) {
    modules.CallExtendedPluginOnGuildBanRemove(
        user,
    )
}

// BotOnMessageCreate gets called after a new message was sent
// This will be called after *every* message on *every* server so it should die as soon as possible
// or spawn costly work inside of coroutines.
func BotOnMessageCreate(session *discordgo.Session, message *discordgo.MessageCreate) {
    // Ignore other bots and @everyone/@here
    if message.Author.Bot || message.MentionEveryone {
        return
    }

    // Get the channel
    // Ignore the event if we cannot resolve the channel
    channel, err := cache.Channel(message.ChannelID)
    if err != nil {
        go raven.CaptureError(err, map[string]string{})
        return
    }

    // We only do things in guilds.
    // Get a friend already and stop chatting with bots
    if channel.IsPrivate {
        // Track usage
        metrics.CleverbotRequests.Add(1)

        // Mark typing
        session.ChannelTyping(message.ChannelID)

        // Prepare content for editing
        msg := message.Content

        /// Remove our @mention
        msg = strings.Replace(msg, "<@"+session.State.User.ID+">", "", -1)

        // Trim message
        msg = strings.TrimSpace(msg)

        // Resolve other @mentions before sending the message
        for _, user := range message.Mentions {
            msg = strings.Replace(msg, "<@"+user.ID+">", user.Username, -1)
        }

        // Remove smileys
        msg = regexp.MustCompile(`:\w+:`).ReplaceAllString(msg, "")

        // Send to cleverbot
        helpers.CleverbotSend(session, channel.ID, msg)
        return
    }

    // Check if the message contains @mentions for us
    if strings.HasPrefix(message.Content, "<@") && len(message.Mentions) > 0 && message.Mentions[0].ID == session.State.User.ID {
        // Consume a key for this action
        e := ratelimits.Container.Drain(1, message.Author.ID)
        if e != nil {
            return
        }

        // Prepare content for editing
        msg := message.Content

        /// Remove our @mention
        msg = strings.Replace(msg, "<@"+session.State.User.ID+">", "", -1)

        // Trim message
        msg = strings.TrimSpace(msg)

        // Convert to []byte before matching
        bmsg := []byte(msg)

        // Match against common task patterns
        // Send to cleverbot if nothing matches
        switch {
        case regexp.MustCompile("(?i)^HELP.*").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            sendHelp(message)
            return

        case regexp.MustCompile("(?i)^PREFIX.*").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            prefix := helpers.GetPrefixForServer(channel.GuildID)
            if prefix == "" {
                cache.GetSession().ChannelMessageSend(
                    channel.ID,
                    helpers.GetText("bot.prefix.not-set"),
                )
            }

            cache.GetSession().ChannelMessageSend(
                channel.ID,
                helpers.GetTextF("bot.prefix.is", prefix),
            )
            return

        case regexp.MustCompile("(?i)^REFRESH CHAT SESSION$").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            helpers.RequireAdmin(message.Message, func() {
                // Refresh cleverbot session
                helpers.CleverbotRefreshSession(channel.ID)
                cache.GetSession().ChannelMessageSend(channel.ID, helpers.GetText("bot.cleverbot.refreshed"))
            })
            return

        case regexp.MustCompile("(?i)^SET PREFIX (.){1,25}$").Match(bmsg):
            metrics.CommandsExecuted.Add(1)
            helpers.RequireAdmin(message.Message, func() {
                // Extract prefix
                prefix := strings.Fields(regexp.MustCompile("(?i)^SET PREFIX\\s").ReplaceAllString(msg, ""))[0]

                // Set new prefix
                err := helpers.SetPrefixForServer(
                    channel.GuildID,
                    prefix,
                )

                if err != nil {
                    helpers.SendError(message.Message, err)
                } else {
                    cache.GetSession().ChannelMessageSend(channel.ID, helpers.GetTextF("bot.prefix.saved", prefix))
                }
            })
            return

        default:
            // Track usage
            metrics.CleverbotRequests.Add(1)

            // Mark typing
            session.ChannelTyping(message.ChannelID)

            // Resolve other @mentions before sending the message
            for _, user := range message.Mentions {
                msg = strings.Replace(msg, "<@"+user.ID+">", user.Username, -1)
            }

            // Remove smileys
            msg = regexp.MustCompile(`:\w+:`).ReplaceAllString(msg, "")

            // Send to cleverbot
            helpers.CleverbotSend(session, channel.ID, msg)
            return
        }
    }

    modules.CallExtendedPlugin(
        message.Content,
        message.Message,
    )

    // Only continue if a prefix is set
    prefix := helpers.GetPrefixForServer(channel.GuildID)
    if prefix == "" {
        return
    }

    // Check if the message is prefixed for us
    // If not exit
    if !strings.HasPrefix(message.Content, prefix) {
        return
    }

    // Check if the user is allowed to request commands
    if !ratelimits.Container.HasKeys(message.Author.ID) && !helpers.IsBotAdmin(message.Author.ID) {
        session.ChannelMessageSend(message.ChannelID, helpers.GetTextF("bot.ratelimit.hit", message.Author.ID))

        ratelimits.Container.Set(message.Author.ID, -1)
        return
    }

    // Split the message into parts
    parts := strings.Fields(message.Content)

    // Save a sanitized version of the command (no prefix)
    cmd := strings.Replace(parts[0], prefix, "", 1)

    // Check if the user calls for help
    if cmd == "h" || cmd == "help" {
        metrics.CommandsExecuted.Add(1)
        sendHelp(message)
        return
    }

    // Separate arguments from the command
    content := strings.TrimSpace(strings.Replace(message.Content, prefix+cmd, "", -1))

    // Log commands
    Logger.VERBOSE.L("bot", fmt.Sprintf("%s (#%s): %s",
        message.Author.Username, message.Author.ID, message.Content))

    // Check if a module matches said command
    modules.CallBotPlugin(cmd, content, message.Message)

    // Check if a trigger matches
    modules.CallTriggerPlugin(cmd, content, message.Message)
}

// BotOnReactionAdd gets called after a reaction is added
// This will be called after *every* reaction added on *every* server so it
// should die as soon as possible or spawn costly work inside of coroutines.
// This is currently used for the *poll* plugin.
func BotOnReactionAdd(session *discordgo.Session, reaction *discordgo.MessageReactionAdd) {
    modules.CallExtendedPluginOnReactionAdd(reaction)

    if reaction.UserID == session.State.User.ID {
        return
    }

    channel, err := session.State.Channel(reaction.ChannelID)
    if err != nil {
        return
    }
    if emojis.ToNumber(reaction.Emoji.Name) == -1 {
        //session.MessageReactionRemove(reaction.ChannelID, reaction.MessageID, reaction.Emoji.Name, reaction.UserID)
        return
    }
    if helpers.VotePollIfItsOne(channel.GuildID, reaction.MessageReaction) {
        helpers.UpdatePollMsg(channel.GuildID, reaction.MessageID)
    }

}

func BotOnReactionRemove(session *discordgo.Session, reaction *discordgo.MessageReactionRemove) {
    modules.CallExtendedPluginOnReactionRemove(reaction)
}

func sendHelp(message *discordgo.MessageCreate) {
    cache.GetSession().ChannelMessageSend(
        message.ChannelID,
        helpers.GetTextF("bot.help", message.Author.ID),
    )
}

// Changes the game interval every 10 seconds after called
func changeGameInterval(session *discordgo.Session) {
    for {
        users := make(map[string]string)
        guilds := session.State.Guilds

        for _, guild := range guilds {
            lastAfterMemberId := ""
            for {
                members, err := session.GuildMembers(guild.ID, lastAfterMemberId, 1000)
                if len(members) <= 0 {
                    break
                }
                lastAfterMemberId = members[len(members)-1].User.ID
                helpers.Relax(err)
                for _, u := range members {
                    users[u.User.ID] = u.User.Username
                }
            }
        }

        err := session.UpdateStatus(0, fmt.Sprintf("with %d people on %d servers", len(users), len(guilds)))
        if err != nil {
            raven.CaptureError(err, map[string]string{})
        }

        time.Sleep(1 * time.Hour)
    }
}
<|endoftext|>


<commit_before>package network

import (
	"errors"
	"net"
	"reflect"
	"strconv"
	"time"

	"github.com/NebulousLabs/Andromeda/encoding"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// Call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

func ReadPrefix(conn net.Conn) ([]byte, error) {
	prefix := make([]byte, 4)
	if n, err := conn.Read(prefix); err != nil || n != len(prefix) {
		return nil, errors.New("could not read length prefix")
	}
	msgLen := int(encoding.DecUint64(prefix))
	if msgLen > maxMsgLen {
		return nil, errors.New("message too long")
	}
	// read msgLen bytes
	var data []byte
	buf := make([]byte, 1024)
	for total := 0; total < msgLen; {
		n, err := conn.Read(buf)
		if err != nil {
			return nil, err
		}
		data = append(data, buf[:n]...)
		total += n
	}
	if len(data) != msgLen {
		return nil, errors.New("message length mismatch")
	}
	return data, nil
}

func WritePrefix(conn net.Conn, data []byte) (int, error) {
	encLen := encoding.EncUint64(uint64(len(data)))
	return conn.Write(append(encLen[:4], data...))
}

// SendVal returns a closure that can be used in conjuction with Call to send
// a value to a NetAddress. It prefixes the encoded data with a header,
// containing the message's type and length
func SendVal(t byte, val interface{}) func(net.Conn) error {
	encVal := encoding.Marshal(val)
	encLen := encoding.EncUint64(uint64(len(encVal)))
	msg := append([]byte{t},
		append(encLen[:4], encVal...)...)

	return func(conn net.Conn) error {
		_, err := conn.Write(msg)
		return err
	}
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

// RandomPeer selects and returns a random peer from the address book.
// TODO: probably not smart to depend on map iteration...
func (tcps *TCPServer) RandomPeer() (rand NetAddress) {
	for addr := range tcps.addressbook {
		rand = addr
		break
	}
	return
}

// Broadcast calls the specified function on each peer in the address book.
func (tcps *TCPServer) Broadcast(fn func(net.Conn) error) {
	for addr := range tcps.addressbook {
		addr.Call(fn)
	}
}

// RegisterHandler registers a message type with a message handler. The
// existing handler for that type will be overwritten.
func (tcps *TCPServer) RegisterHandler(t byte, fn func(net.Conn, []byte) error) {
	tcps.handlerMap[t] = fn
}

// RegisterRPC is for simple handlers. A simple handler decodes the message
// data and passes it to fn. fn must have the type signature:
//   func(Type) error
// i.e. a 1-adic function that returns an error.
func (tcps *TCPServer) RegisterRPC(t byte, fn interface{}) {
	// if fn not correct type, panic
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() != 1 ||
		typ.NumOut() != 1 || typ.Out(0) != reflect.TypeOf((*error)(nil)).Elem() {
		panic("registered function has wrong type signature")
	}

	// create and register function
	tcps.RegisterHandler(t, func(_ net.Conn, b []byte) error {
		// create object to decode into
		v := reflect.New(typ.In(0))
		if err := encoding.Unmarshal(b, v.Interface()); err != nil {
			return err
		}
		// call fn on object
		if err := val.Call([]reflect.Value{v.Elem()})[0].Interface(); err != nil {
			return err.(error)
		}
		return nil
	})
}

// NewTCPServer creates a TCPServer that listens on the specified port.
func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		myAddr:      NetAddress{"", port},
		addressbook: make(map[NetAddress]struct{}),
	}
	// default handlers
	tcps.handlerMap = map[byte]func(net.Conn, []byte) error{
		'H': sendHostname,
		'P': tcps.sharePeers,
		'A': tcps.addPeer,
	}

	// spawn listener
	go tcps.listen()
	return
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
// TODO: set deadlines?
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	msgType := make([]byte, 1)
	if n, err := conn.Read(msgType); err != nil || n != 1 {
		// TODO: log error
		return
	}
	msgData, err := ReadPrefix(conn)
	if err != nil {
		// TODO: log error
		return
	}
	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgType[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
		// no wait, send the error?
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func sendHostname(conn net.Conn, _ []byte) error {
	_, err := WritePrefix(conn, []byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {
			break
		}
		addrs = append(addrs, addr)
		num--
	}
	_, err := WritePrefix(conn, encoding.Marshal(addrs))
	return err
}

// addPeer adds the connecting peer to its address book
func (tcps *TCPServer) addPeer(_ net.Conn, data []byte) (err error) {
	var addr NetAddress
	if err = encoding.Unmarshal(data, &addr); err != nil {
		return
	}
	tcps.addressbook[addr] = struct{}{}
	return
}

// Ping returns whether a NetAddress is reachable. It accomplishes this by
// initiating a TCP connection and immediately closes it. This is pretty
// unsophisticated. I'll add a Pong later.
func (tcps *TCPServer) Ping(addr NetAddress) bool {
	conn, err := net.DialTimeout("tcp", addr.String(), timeout)
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// learnHostname learns the external IP of the TCPServer.
func (tcps *TCPServer) learnHostname(conn net.Conn) (err error) {
	// send hostname request
	if _, err = conn.Write([]byte{'H', 0, 0, 0, 0}); err != nil {
		return
	}
	// read response
	data, err := ReadPrefix(conn)
	if err != nil {
		return
	}
	// TODO: try to ping ourselves?
	host, _, err := net.SplitHostPort(string(data))
	if err != nil {
		return
	}
	tcps.myAddr.Host = host
	return
}

// requestPeers queries a peer for additional peers, and adds any new peers to
// the address book.
func (tcps *TCPServer) requestPeers(conn net.Conn) (err error) {
	// request 10 peers
	if _, err = conn.Write([]byte{'P', 1, 0, 0, 0, 10}); err != nil {
		return
	}
	// read response
	data, err := ReadPrefix(conn)
	if err != nil {
		return
	}
	var addrs []NetAddress
	if err = encoding.Unmarshal(data, &addrs); err != nil {
		return
	}
	// add peers
	for _, addr := range addrs {
		if addr != tcps.myAddr && tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}
	return
}

// Bootstrap discovers the external IP of the TCPServer, requests peers from
// the initial peer list, and announces itself to those peers.
func (tcps *TCPServer) Bootstrap() (err error) {
	// populate initial peer list
	for _, addr := range BootstrapPeers {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}

	// learn hostname
	for addr := range tcps.addressbook {
		if addr.Call(tcps.learnHostname) == nil {
			break
		}
	}

	// request peers
	// TODO: maybe iterate until we have enough new peers?
	tcps.Broadcast(tcps.requestPeers)

	// announce ourselves to new peers
	tcps.Broadcast(SendVal('A', tcps.myAddr))

	return
}

func (tcps *TCPServer) AddressBook() (book []NetAddress) {
	for address := range tcps.addressbook {
		book = append(book, address)
	}
	return
}
<commit_msg>revamp RPC<commit_after>package network

import (
	"errors"
	"net"
	"reflect"
	"strconv"
	"time"

	"github.com/NebulousLabs/Andromeda/encoding"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// Call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

func ReadPrefix(conn net.Conn) ([]byte, error) {
	prefix := make([]byte, 4)
	if n, err := conn.Read(prefix); err != nil || n != len(prefix) {
		return nil, errors.New("could not read length prefix")
	}
	msgLen := int(encoding.DecUint64(prefix))
	if msgLen > maxMsgLen {
		return nil, errors.New("message too long")
	}
	// read msgLen bytes
	var data []byte
	buf := make([]byte, 1024)
	for total := 0; total < msgLen; {
		n, err := conn.Read(buf)
		if err != nil {
			return nil, err
		}
		data = append(data, buf[:n]...)
		total += n
	}
	if len(data) != msgLen {
		return nil, errors.New("message length mismatch")
	}
	return data, nil
}

func ReadObject(conn net.Conn, obj interface{}) error {
	data, err := ReadPrefix(conn)
	if err != nil {
		return err
	}
	return encoding.Unmarshal(data, obj)
}

func WritePrefix(conn net.Conn, data []byte) (int, error) {
	encLen := encoding.EncUint64(uint64(len(data)))
	return conn.Write(append(encLen[:4], data...))
}

func WriteObject(conn net.Conn, obj interface{}) (int, error) {
	return WritePrefix(conn, encoding.Marshal(obj))
}

// RPC performs a Remote Procedure Call by sending the procedure name and
// encoded argument, and decoding the response into the supplied object.
// 'resp' must be a pointer. If arg is nil, no object is sent. If 'resp' is
// nil, no response is read.
func (na *NetAddress) RPC(t byte, arg, resp interface{}) error {
	return na.Call(func(conn net.Conn) error {
		conn.Write([]byte{t})
		if arg != nil {
			if _, err := WriteObject(conn, arg); err != nil {
				return err
			}
		}
		if resp != nil {
			return ReadObject(conn, resp)
		}
		return nil
	})
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

// RandomPeer selects and returns a random peer from the address book.
// TODO: probably not smart to depend on map iteration...
func (tcps *TCPServer) RandomPeer() (rand NetAddress) {
	for addr := range tcps.addressbook {
		rand = addr
		break
	}
	return
}

// Broadcast calls the specified function on each peer in the address book.
func (tcps *TCPServer) Broadcast(fn func(net.Conn) error) {
	for addr := range tcps.addressbook {
		addr.Call(fn)
	}
}

func (tcps *TCPServer) Register(t byte, fn interface{}) {
	// all handlers are function with at least one in and one error out
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() < 1 ||
		typ.NumOut() != 1 || typ.Out(0) != reflect.TypeOf((*error)(nil)).Elem() {
		panic("registered function has wrong type signature")
	}

	switch {
	// func(net.Conn, []byte) error
	case typ.NumIn() == 2 && typ.In(0) == reflect.TypeOf((*net.Conn)(nil)).Elem() && typ.In(1) == reflect.TypeOf([]byte{}):
		tcps.handlerMap[t] = fn.(func(net.Conn, []byte) error)
	// func(Type, *Type) error
	case typ.NumIn() == 2 && typ.In(0).Kind() != reflect.Ptr && typ.In(1).Kind() == reflect.Ptr:
		tcps.registerRPC(t, val, typ)
	// func(Type) error
	case typ.NumIn() == 1 && typ.In(0).Kind() != reflect.Ptr:
		tcps.registerArg(t, val, typ)
	// func(*Type) error
	case typ.NumIn() == 1 && typ.In(0).Kind() == reflect.Ptr:
		tcps.registerResp(t, val, typ)
	default:
		panic("registered function has wrong type signature")
	}
}

// registerRPC is for handlers that return a value. The input is decoded and
// passed to fn, which stores its result in a pointer argument. This argument
// is then written back to the caller. fn must have the type signature:
//     func(Type, *Type) error
func (tcps *TCPServer) registerRPC(t byte, fn reflect.Value, typ reflect.Type) {
	tcps.handlerMap[t] = func(conn net.Conn, b []byte) error {
		// create object to decode into
		arg := reflect.New(typ.In(0))
		if err := encoding.Unmarshal(b, arg.Interface()); err != nil {
			return err
		}
		// call fn on object
		resp := reflect.New(typ.In(1))
		if err := fn.Call([]reflect.Value{arg.Elem(), resp})[0].Interface(); err != nil {
			return err.(error)
		}
		// write response
		_, err := WriteObject(conn, resp.Elem().Interface())
		return err
	}
}

// registerArg is for RPCs that do not return a value.
func (tcps *TCPServer) registerArg(t byte, fn reflect.Value, typ reflect.Type) {
	tcps.handlerMap[t] = func(_ net.Conn, b []byte) error {
		// create object to decode into
		arg := reflect.New(typ.In(0))
		if err := encoding.Unmarshal(b, arg.Interface()); err != nil {
			return err
		}
		// call fn on object
		if err := fn.Call([]reflect.Value{arg.Elem()})[0].Interface(); err != nil {
			return err.(error)
		}
		return nil
	}
}

// registerResp is for RPCs that do not take a value.
func (tcps *TCPServer) registerResp(t byte, fn reflect.Value, typ reflect.Type) {
	tcps.handlerMap[t] = func(conn net.Conn, _ []byte) error {
		// create object to hold response
		resp := reflect.New(typ.In(0))
		// call fn
		if err := fn.Call([]reflect.Value{resp})[0].Interface(); err != nil {
			return err.(error)
		}
		// write response
		_, err := WriteObject(conn, resp.Elem().Interface())
		return err
	}
}

// NewTCPServer creates a TCPServer that listens on the specified port.
func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		myAddr:      NetAddress{"", port},
		addressbook: make(map[NetAddress]struct{}),
	}
	// default handlers
	tcps.handlerMap = map[byte]func(net.Conn, []byte) error{
		'H': sendHostname,
		'P': tcps.sharePeers,
		'A': tcps.addPeer,
	}

	// spawn listener
	go tcps.listen()
	return
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
// TODO: set deadlines?
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	msgType := make([]byte, 1)
	if n, err := conn.Read(msgType); err != nil || n != 1 {
		// TODO: log error
		return
	}
	msgData, err := ReadPrefix(conn)
	if err != nil {
		// TODO: log error
		return
	}
	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgType[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
		// no wait, send the error?
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func sendHostname(conn net.Conn, _ []byte) error {
	_, err := WritePrefix(conn, []byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {
			break
		}
		addrs = append(addrs, addr)
		num--
	}
	_, err := WritePrefix(conn, encoding.Marshal(addrs))
	return err
}

// addPeer adds the connecting peer to its address book
func (tcps *TCPServer) addPeer(_ net.Conn, data []byte) (err error) {
	var addr NetAddress
	if err = encoding.Unmarshal(data, &addr); err != nil {
		return
	}
	tcps.addressbook[addr] = struct{}{}
	return
}

// Ping returns whether a NetAddress is reachable. It accomplishes this by
// initiating a TCP connection and immediately closes it. This is pretty
// unsophisticated. I'll add a Pong later.
func (tcps *TCPServer) Ping(addr NetAddress) bool {
	conn, err := net.DialTimeout("tcp", addr.String(), timeout)
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// learnHostname learns the external IP of the TCPServer.
func (tcps *TCPServer) learnHostname(conn net.Conn) (err error) {
	// send hostname request
	if _, err = conn.Write([]byte{'H', 0, 0, 0, 0}); err != nil {
		return
	}
	// read response
	data, err := ReadPrefix(conn)
	if err != nil {
		return
	}
	// TODO: try to ping ourselves?
	host, _, err := net.SplitHostPort(string(data))
	if err != nil {
		return
	}
	tcps.myAddr.Host = host
	return
}

// requestPeers queries a peer for additional peers, and adds any new peers to
// the address book.
func (tcps *TCPServer) requestPeers(conn net.Conn) (err error) {
	// request 10 peers
	if _, err = conn.Write([]byte{'P', 1, 0, 0, 0, 10}); err != nil {
		return
	}
	// read response
	data, err := ReadPrefix(conn)
	if err != nil {
		return
	}
	var addrs []NetAddress
	if err = encoding.Unmarshal(data, &addrs); err != nil {
		return
	}
	// add peers
	for _, addr := range addrs {
		if addr != tcps.myAddr && tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}
	return
}

// anounce announces the TCPServer's NetAddress to other peers.
func (tcps *TCPServer) announce(conn net.Conn) error {
	conn.Write([]byte{'A'})
	_, err := WriteObject(conn, tcps.myAddr)
	return err
}

// Bootstrap discovers the external IP of the TCPServer, requests peers from
// the initial peer list, and announces itself to those peers.
func (tcps *TCPServer) Bootstrap() (err error) {
	// populate initial peer list
	for _, addr := range BootstrapPeers {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}

	// learn hostname
	for addr := range tcps.addressbook {
		if addr.Call(tcps.learnHostname) == nil {
			break
		}
	}

	// request peers
	// TODO: maybe iterate until we have enough new peers?
	tcps.Broadcast(tcps.requestPeers)

	// announce ourselves to new peers
	tcps.Broadcast(tcps.announce)

	return
}

func (tcps *TCPServer) AddressBook() (book []NetAddress) {
	for address := range tcps.addressbook {
		book = append(book, address)
	}
	return
}
<|endoftext|>