<commit_before>package logstash

import (
	"encoding/json"
	"errors"
	"log"
	"net"
	"regexp"
	"strings"
	"os"

	"github.com/gliderlabs/logspout/router"
)

func init() {
	router.AdapterFactories.Register(NewAdapter, "logstash")
}

var regexps = []*regexp.Regexp{
	regexp.MustCompile(`^\s`),
	regexp.MustCompile(`line \d+, in .+`),
}

// Adapter is an adapter that streams UDP JSON to Logstash.
type Adapter struct {
	conn  net.Conn
	route *router.Route
}

// NewAdapter creates an Adapter with UDP as the default transport.
func NewAdapter(route *router.Route) (router.LogAdapter, error) {
	transport, found := router.AdapterTransports.Lookup(route.AdapterTransport("udp"))
	if !found {
		return nil, errors.New("unable to find adapter: " + route.Adapter)
	}

	conn, err := transport.Dial(route.Address, route.Options)
	if err != nil {
		return nil, err
	}

	return &Adapter{
		route: route,
		conn:  conn,
	}, nil
}

// MergeMessages merges an array of Message into a string
func MergeMessages(messages []Message) string {
	var strs = make([]string, 0)

	for _, x := range messages {
		strs = append(strs, x.Message)
	}

	return strings.Join(strs, "\n")
}

// GetTags decides if a message array should be tagged multiline.
func GetTags (messages []Message) []string {
	var tags = make([]string, 0)

	if len(messages) > 1 {
		tags = append(tags, "multiline")
	} else {
		tags = append(tags, "")
	}

	return tags
}

// IsMultiline is a function that determines if a string should be in the queue map.
func IsMultiline(message string) bool {
	for _, expression := range regexps {
		if expression.Match([]byte(message)) == true {
			return true;
		}
	}

	return false;
}

// GetHostname gets the HOSTNAME variable or the container's hostname.
func GetHostname() string {
	hostname := os.Getenv("HOSTNAME")

	if hostname == "" {
		log.Println("logstash: Defaulting to container hostname.")
		hostname, err := os.Hostname()
		if err != nil {
			log.Println("logstash_hostname:", err)
		}
		return hostname
	}
	return hostname
}

// Stream implements the router.LogAdapter interface.
func (a *Adapter) Stream(logstream chan *router.Message) {
	queue := make(map[string][]Message)

	hostname := GetHostname()

	for m := range logstream {
		rawMessage := Message{
			Message:  m.Data,
		}
		finalMessage := Message{}

		_, existing := queue[m.Container.ID];

		// Create an empty slice if there is no queue slice.
		if !existing {
			queue[m.Container.ID] = []Message{}
		}

		queue[m.Container.ID] = append(queue[m.Container.ID], rawMessage)

		if IsMultiline(m.Data) {
			continue
		} else {
			if len(queue[m.Container.ID]) == 1 {
				continue
			} else {
				// remove trailing slash from container name
				containerName := strings.TrimLeft(m.Container.Name, "/")

				finalMessage = Message{
					Message: MergeMessages(queue[m.Container.ID]),
					Name: containerName,
					ID: m.Container.ID,
					Image: m.Container.Config.Image,
					Hostname: m.Container.Config.Hostname,
					Stream: m.Source,
					Tags: GetTags(queue[m.Container.ID]),
					Host: hostname,
				}
				queue[m.Container.ID] = []Message{}
			}
		}

		// Mashal the message into JSON.
		js, err := json.Marshal(finalMessage)
		if err != nil {
			log.Println("logstash_marshal:", err)
			continue
		}

		// Write the message to the Logstash server.
		_, err = a.conn.Write(js)
		if err != nil {
			log.Println("logstash_write:", err)
			continue
		}
	}
}

// Message is a simple JSON input to Logstash.
type Message struct {
	Message  string   `json:"message"`
	Name     string   `json:"container_name"`
	ID       string   `json:"container_id"`
	Image    string   `json:"image_name"`
	Hostname string   `json:"container_hostname"`
	Host     string   `json:"host"`
	Stream   string   `json:"stream"`
	Tags     []string `json:"tags"`
}
<commit_msg>Clean up the code by updating messages at the end.<commit_after>package logstash

import (
	"encoding/json"
	"errors"
	"log"
	"net"
	"regexp"
	"strings"
	"os"

	"github.com/gliderlabs/logspout/router"
)

func init() {
	router.AdapterFactories.Register(NewAdapter, "logstash")
}

var regexps = []*regexp.Regexp{
	regexp.MustCompile(`^\s`),
	regexp.MustCompile(`line \d+, in .+`),
}

// Adapter is an adapter that streams UDP JSON to Logstash.
type Adapter struct {
	conn  net.Conn
	route *router.Route
}

// NewAdapter creates an Adapter with UDP as the default transport.
func NewAdapter(route *router.Route) (router.LogAdapter, error) {
	transport, found := router.AdapterTransports.Lookup(route.AdapterTransport("udp"))
	if !found {
		return nil, errors.New("unable to find adapter: " + route.Adapter)
	}

	conn, err := transport.Dial(route.Address, route.Options)
	if err != nil {
		return nil, err
	}

	return &Adapter{
		route: route,
		conn:  conn,
	}, nil
}

// MergeMessages merges an array of Message into a string
func MergeMessages(messages []Message) string {
	var strs = make([]string, 0)

	for _, x := range messages {
		strs = append(strs, x.Message)
	}

	return strings.Join(strs, "\n")
}

// GetTags decides if a message array should be tagged multiline.
func GetTags (messages []Message) []string {
	var tags = make([]string, 0)

	if len(messages) > 1 {
		tags = append(tags, "multiline")
	} else {
		tags = append(tags, "")
	}

	return tags
}

// IsMultiline is a function that determines if a string should be in the queue map.
func IsMultiline(message string) bool {
	for _, expression := range regexps {
		if expression.Match([]byte(message)) == true {
			return true;
		}
	}

	return false;
}

// GetHostname gets the HOSTNAME variable or the container's hostname.
func GetHostname() string {
	hostname := os.Getenv("HOSTNAME")

	if hostname == "" {
		log.Println("logstash: Defaulting to container hostname.")
		hostname, err := os.Hostname()
		if err != nil {
			log.Println("logstash_hostname:", err)
		}
		return hostname
	}
	return hostname
}

// Stream implements the router.LogAdapter interface.
func (a *Adapter) Stream(logstream chan *router.Message) {
	queue := make(map[string][]Message)

	hostname := GetHostname()

	for m := range logstream {
		rawMessage := Message{
			Message:  m.Data,
		}
		finalMessage := Message{}

		messages, existing := queue[m.Container.ID];

		// Create an empty slice if there is no queue slice.
		if !existing {
			messages = []Message{}
		}

		if IsMultiline(m.Data) {
			messages = append(messages, rawMessage)
			queue[m.Container.ID] = messages;
			continue
		} else {
			if len(queue[m.Container.ID]) == 0 {
				messages = append(messages, rawMessage)
				queue[m.Container.ID] = messages;
				continue
			} else {
				// remove trailing slash from container name
				containerName := strings.TrimLeft(m.Container.Name, "/")

				finalMessage = Message{
					Message: MergeMessages(messages),
					Name: containerName,
					ID: m.Container.ID,
					Image: m.Container.Config.Image,
					Hostname: m.Container.Config.Hostname,
					Stream: m.Source,
					Tags: GetTags(messages),
					Host: hostname,
				}

				if len(messages) == 1 && !IsMultiline(messages[0].Message) {
					messages = []Message{rawMessage}
				} else {
					messages = []Message{}
				}

				queue[m.Container.ID] = messages;
			}
		}

		// Mashal the message into JSON.
		js, err := json.Marshal(finalMessage)
		if err != nil {
			log.Println("logstash_marshal:", err)
			continue
		}

		// Write the message to the Logstash server.
		_, err = a.conn.Write(js)
		if err != nil {
			log.Println("logstash_write:", err)
			continue
		}
	}
}

// Message is a simple JSON input to Logstash.
type Message struct {
	Message  string   `json:"message"`
	Name     string   `json:"container_name"`
	ID       string   `json:"container_id"`
	Image    string   `json:"image_name"`
	Hostname string   `json:"container_hostname"`
	Host     string   `json:"host"`
	Stream   string   `json:"stream"`
	Tags     []string `json:"tags"`
}
<|endoftext|>


<commit_before>package gcs

import (
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/hashicorp/terraform/backend"
	"github.com/hashicorp/terraform/state/remote"
)

func TestStateFile(t *testing.T) {
	t.Parallel()

	cases := []struct {
		prefix           string
		defaultStateFile string
		name             string
		wantStateFile    string
		wantLockFile     string
	}{
		{"state", "", "default", "state/default.tfstate", "state/default.tflock"},
		{"state", "", "test", "state/test.tfstate", "state/test.tflock"},
		{"state", "legacy.tfstate", "default", "legacy.tfstate", "legacy.tflock"},
		{"state", "legacy.tfstate", "test", "state/test.tfstate", "state/test.tflock"},
		{"state", "legacy.state", "default", "legacy.state", "legacy.state.tflock"},
		{"state", "legacy.state", "test", "state/test.tfstate", "state/test.tflock"},
	}
	for _, c := range cases {
		b := &gcsBackend{
			prefix:           c.prefix,
			defaultStateFile: c.defaultStateFile,
		}

		if got := b.stateFile(c.name); got != c.wantStateFile {
			t.Errorf("stateFile(%q) = %q, want %q", c.name, got, c.wantStateFile)
		}

		if got := b.lockFile(c.name); got != c.wantLockFile {
			t.Errorf("lockFile(%q) = %q, want %q", c.name, got, c.wantLockFile)
		}
	}
}

func TestRemoteClient(t *testing.T) {
	t.Parallel()

	be := setupBackend(t)
	defer teardownBackend(t, be)

	ss, err := be.State(backend.DefaultStateName)
	if err != nil {
		t.Fatalf("be.State(%q) = %v", backend.DefaultStateName, err)
	}

	rs, ok := ss.(*remote.State)
	if !ok {
		t.Fatalf("be.State(): got a %T, want a *remote.State", ss)
	}

	remote.TestClient(t, rs.Client)
}

func TestRemoteLocks(t *testing.T) {
	t.Parallel()

	be := setupBackend(t)
	defer teardownBackend(t, be)

	remoteClient := func() (remote.Client, error) {
		ss, err := be.State(backend.DefaultStateName)
		if err != nil {
			return nil, err
		}

		rs, ok := ss.(*remote.State)
		if !ok {
			return nil, fmt.Errorf("be.State(): got a %T, want a *remote.State", ss)
		}

		return rs.Client, nil
	}

	c0, err := remoteClient()
	if err != nil {
		t.Fatalf("remoteClient(0) = %v", err)
	}
	c1, err := remoteClient()
	if err != nil {
		t.Fatalf("remoteClient(1) = %v", err)
	}

	remote.TestRemoteLocks(t, c0, c1)
}

func TestBackend(t *testing.T) {
	t.Parallel()

	be0 := setupBackend(t)
	defer teardownBackend(t, be0)

	be1 := setupBackend(t)

	backend.TestBackend(t, be0, be1)
}

// setupBackend returns a new GCS backend.
func setupBackend(t *testing.T) backend.Backend {
	t.Helper()

	projectID := os.Getenv("GOOGLE_PROJECT")
	if projectID == "" || os.Getenv("TF_ACC") == "" {
		t.Skip("This test creates a bucket in GCS and populates it. " +
			"Since this may incur costs, it will only run if " +
			"the TF_ACC and GOOGLE_PROJECT environment variables are set.")
	}

	config := map[string]interface{}{
		"project": projectID,
		"bucket":  strings.ToLower(t.Name()),
		"prefix":  "",
	}

	if creds := os.Getenv("GOOGLE_CREDENTIALS"); creds != "" {
		config["credentials"] = creds
		t.Logf("using credentials from %q", creds)
	} else {
		t.Log("using default credentials; set GOOGLE_CREDENTIALS for custom credentials")
	}

	return backend.TestBackendConfig(t, New(), config)
}

// teardownBackend deletes all states from be except the default state.
func teardownBackend(t *testing.T, be backend.Backend) {
	t.Helper()

	// Delete all states. The bucket must be empty before it can be deleted.
	states, err := be.States()
	if err != nil {
		t.Fatalf("be.States() = %v; manual clean-up may be required", err)
	}
	for _, st := range states {
		if st == backend.DefaultStateName {
			continue
		}
		if err := be.DeleteState(st); err != nil {
			t.Fatalf("be.DeleteState(%q) = %v; manual clean-up may be required", st, err)
		}
	}

	gcsBE, ok := be.(*gcsBackend)
	if !ok {
		t.Fatalf("be is a %T, want a *gcsBackend", be)
	}
	ctx := gcsBE.storageContext

	// Delete the default state, which DeleteState() will refuse to do.
	// It's okay if this fails, not all tests create a default state.
	if err := gcsBE.storageClient.Bucket(gcsBE.bucketName).Object("default.tfstate").Delete(ctx); err != nil {
		t.Logf("deleting \"default.tfstate\": %v; manual clean-up may be required", err)
	}

	// Delete the bucket itself.
	if err := gcsBE.storageClient.Bucket(gcsBE.bucketName).Delete(ctx); err != nil {
		t.Fatalf("deleting bucket failed: %v; manual cleanup may be required, though later test runs will happily reuse an existing bucket", err)
	}
}
<commit_msg>backend/remote-state/gcs: Include project ID in bucket names when testing.<commit_after>package gcs

import (
	"fmt"
	"os"
	"testing"

	"github.com/hashicorp/terraform/backend"
	"github.com/hashicorp/terraform/state/remote"
)

func TestStateFile(t *testing.T) {
	t.Parallel()

	cases := []struct {
		prefix           string
		defaultStateFile string
		name             string
		wantStateFile    string
		wantLockFile     string
	}{
		{"state", "", "default", "state/default.tfstate", "state/default.tflock"},
		{"state", "", "test", "state/test.tfstate", "state/test.tflock"},
		{"state", "legacy.tfstate", "default", "legacy.tfstate", "legacy.tflock"},
		{"state", "legacy.tfstate", "test", "state/test.tfstate", "state/test.tflock"},
		{"state", "legacy.state", "default", "legacy.state", "legacy.state.tflock"},
		{"state", "legacy.state", "test", "state/test.tfstate", "state/test.tflock"},
	}
	for _, c := range cases {
		b := &gcsBackend{
			prefix:           c.prefix,
			defaultStateFile: c.defaultStateFile,
		}

		if got := b.stateFile(c.name); got != c.wantStateFile {
			t.Errorf("stateFile(%q) = %q, want %q", c.name, got, c.wantStateFile)
		}

		if got := b.lockFile(c.name); got != c.wantLockFile {
			t.Errorf("lockFile(%q) = %q, want %q", c.name, got, c.wantLockFile)
		}
	}
}

func TestRemoteClient(t *testing.T) {
	t.Parallel()

	be := setupBackend(t)
	defer teardownBackend(t, be)

	ss, err := be.State(backend.DefaultStateName)
	if err != nil {
		t.Fatalf("be.State(%q) = %v", backend.DefaultStateName, err)
	}

	rs, ok := ss.(*remote.State)
	if !ok {
		t.Fatalf("be.State(): got a %T, want a *remote.State", ss)
	}

	remote.TestClient(t, rs.Client)
}

func TestRemoteLocks(t *testing.T) {
	t.Parallel()

	be := setupBackend(t)
	defer teardownBackend(t, be)

	remoteClient := func() (remote.Client, error) {
		ss, err := be.State(backend.DefaultStateName)
		if err != nil {
			return nil, err
		}

		rs, ok := ss.(*remote.State)
		if !ok {
			return nil, fmt.Errorf("be.State(): got a %T, want a *remote.State", ss)
		}

		return rs.Client, nil
	}

	c0, err := remoteClient()
	if err != nil {
		t.Fatalf("remoteClient(0) = %v", err)
	}
	c1, err := remoteClient()
	if err != nil {
		t.Fatalf("remoteClient(1) = %v", err)
	}

	remote.TestRemoteLocks(t, c0, c1)
}

func TestBackend(t *testing.T) {
	t.Parallel()

	be0 := setupBackend(t)
	defer teardownBackend(t, be0)

	be1 := setupBackend(t)

	backend.TestBackend(t, be0, be1)
}

// setupBackend returns a new GCS backend.
func setupBackend(t *testing.T) backend.Backend {
	t.Helper()

	projectID := os.Getenv("GOOGLE_PROJECT")
	if projectID == "" || os.Getenv("TF_ACC") == "" {
		t.Skip("This test creates a bucket in GCS and populates it. " +
			"Since this may incur costs, it will only run if " +
			"the TF_ACC and GOOGLE_PROJECT environment variables are set.")
	}

	config := map[string]interface{}{
		"project": projectID,
		"bucket":  projectID + "-" + t.Name(),
		"prefix":  "",
	}

	if creds := os.Getenv("GOOGLE_CREDENTIALS"); creds != "" {
		config["credentials"] = creds
		t.Logf("using credentials from %q", creds)
	} else {
		t.Log("using default credentials; set GOOGLE_CREDENTIALS for custom credentials")
	}

	return backend.TestBackendConfig(t, New(), config)
}

// teardownBackend deletes all states from be except the default state.
func teardownBackend(t *testing.T, be backend.Backend) {
	t.Helper()

	// Delete all states. The bucket must be empty before it can be deleted.
	states, err := be.States()
	if err != nil {
		t.Fatalf("be.States() = %v; manual clean-up may be required", err)
	}
	for _, st := range states {
		if st == backend.DefaultStateName {
			continue
		}
		if err := be.DeleteState(st); err != nil {
			t.Fatalf("be.DeleteState(%q) = %v; manual clean-up may be required", st, err)
		}
	}

	gcsBE, ok := be.(*gcsBackend)
	if !ok {
		t.Fatalf("be is a %T, want a *gcsBackend", be)
	}
	ctx := gcsBE.storageContext

	// Delete the default state, which DeleteState() will refuse to do.
	// It's okay if this fails, not all tests create a default state.
	if err := gcsBE.storageClient.Bucket(gcsBE.bucketName).Object("default.tfstate").Delete(ctx); err != nil {
		t.Logf("deleting \"default.tfstate\": %v; manual clean-up may be required", err)
	}

	// Delete the bucket itself.
	if err := gcsBE.storageClient.Bucket(gcsBE.bucketName).Delete(ctx); err != nil {
		t.Fatalf("deleting bucket failed: %v; manual cleanup may be required, though later test runs will happily reuse an existing bucket", err)
	}
}
<|endoftext|>


<commit_before>// +build mage

package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/magefile/mage/mg" // mg contains helpful utility functions, like Deps
	"github.com/magefile/mage/sh" // mg contains helpful utility functions, like Deps
)

const localWorkDir = "./sparta"

var header = strings.Repeat("-", 80)

var ignoreSubdirectoryPaths = []string{
	".vendor",
	".sparta",
	".vscode",
	"/resources/describe",
}

// Default target to run when none is specified
// If not set, running mage will list available targets
// var Default = Build

func mageScript(commands [][]string) error {
	for _, eachCommand := range commands {
		var commandErr error
		if len(eachCommand) <= 1 {
			commandErr = sh.Run(eachCommand[0])
		} else {
			commandErr = sh.Run(eachCommand[0], eachCommand[1:]...)
		}
		if commandErr != nil {
			return commandErr
		}
	}
	return nil
}
func mageLog(formatSpecifier string, args ...interface{}) {
	if mg.Verbose() {
		if len(args) != 0 {
			log.Printf(formatSpecifier, args...)
		} else {
			log.Printf(formatSpecifier)
		}
	}
}

func sourceFilesOfType(extension string) ([]string, error) {
	testExtension := strings.TrimPrefix(extension, ".")
	testExtension = fmt.Sprintf(".%s", testExtension)

	files := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		contains := false
		for _, eachComponent := range ignoreSubdirectoryPaths {
			contains = strings.Contains(path, eachComponent)
			if contains {
				break
			}
		}
		if !contains && (filepath.Ext(path) == testExtension) {
			files = append(files, path)
		}
		return nil
	}
	goSourceFilesErr := filepath.Walk(".", walker)
	return files, goSourceFilesErr
}

func applyToSource(fileExtension string, commandParts ...string) error {
	if len(commandParts) <= 0 {
		return errors.New("applyToSource requires a command to apply to source files")
	}
	eligibleSourceFiles, eligibleSourceFilesErr := sourceFilesOfType(fileExtension)
	if eligibleSourceFilesErr != nil {
		return eligibleSourceFilesErr
	}

	mageLog(header)
	mageLog("Applying `%s` to %d `*.%s` source files", commandParts[0], len(eligibleSourceFiles), fileExtension)
	mageLog(header)

	commandArgs := []string{}
	if len(commandParts) > 1 {
		for _, eachPart := range commandParts[1:] {
			commandArgs = append(commandArgs, eachPart)
		}
	}
	for _, eachFile := range eligibleSourceFiles {
		applyArgs := append(commandArgs, eachFile)
		applyErr := sh.Run(commandParts[0], applyArgs...)
		if applyErr != nil {
			return applyErr
		}
	}
	return nil
}

func markdownSourceApply(commandParts ...string) error {
	return applyToSource("md", commandParts...)
}
func goSourceApply(commandParts ...string) error {
	return applyToSource("go", commandParts...)
}

// GenerateBuildInfo creates the automatic buildinfo.go file so that we can
// stamp the SHA into the binaries we build...
func GenerateBuildInfo() error {
	// The first thing we need is the `git` SHA
	cmd := exec.Command("git", "rev-parse", "HEAD")
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		return err
	}
	stdOutResult := strings.TrimSpace(string(stdout.Bytes()))

	// Super = update the buildinfo data
	buildInfoTemplate := `package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: %s

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "%s"
`
	updatedInfo := fmt.Sprintf(buildInfoTemplate, time.Now().UTC(), stdOutResult)
	// Write it to the output location...
	writeErr := ioutil.WriteFile("./buildinfo.go", []byte(updatedInfo), os.ModePerm)

	if writeErr != nil {
		return writeErr
	}
	commitGenerateCommands := [][]string{
		// No dirty trees
		{"git", "diff", "--exit-code"},
		{"git", "commit", "-a", "-m", `"Autogenerated build info"`},
	}
	return mageScript(commitGenerateCommands)

}

// GenerateConstants runs the set of commands that update the embedded CONSTANTS
// for both local and AWS Lambda execution
func GenerateConstants() error {
	generateCommands := [][]string{
		// Create the embedded version
		{"go", "run", "$GOPATH/src/github.com/mjibson/esc/main.go", "-o", "./CONSTANTS.go", "-private", "-pkg", "sparta", "./resources"},
		//Create a secondary CONSTANTS_AWSBINARY.go file with empty content.
		{"go", "run", "$GOPATH/src/github.com/mjibson/esc/main.go", "-o", "./CONSTANTS_AWSBINARY.go", "-private", "-pkg", "sparta", "./resources/awsbinary/README.md"},
		//The next step will insert the
		// build tags at the head of each file so that they are mutually exclusive
		{"go", "run", "./cmd/insertTags/main.go", "./CONSTANTS", "!lambdabinary"},
		{"go", "run", "./cmd/insertTags/main.go", "./CONSTANTS_AWSBINARY", "lambdabinary"},
		{"git", "commit", "-a", "-m", "Autogenerated constants"},
	}
	return mageScript(generateCommands)
}

// InstallBuildRequirements installs or updates the dependent
// packages that aren't referenced by the source, but are needed
// to build the Sparta source
func InstallBuildRequirements() error {
	mageLog("`go get` update flags (env.GO_GET_FLAG): %s", os.Getenv("GO_GET_FLAG"))

	requirements := []string{
		"github.com/golang/dep/...",
		"honnef.co/go/tools/cmd/megacheck",
		"honnef.co/go/tools/cmd/gosimple",
		"honnef.co/go/tools/cmd/unused",
		"honnef.co/go/tools/cmd/staticcheck",
		"golang.org/x/tools/cmd/goimports",
		"github.com/fzipp/gocyclo",
		"github.com/golang/lint/golint",
		"github.com/mjibson/esc",
		"github.com/securego/gosec/cmd/gosec/...",
		"github.com/client9/misspell/cmd/misspell",
	}
	for _, eachDep := range requirements {
		cmdErr := sh.Run("go",
			"get",
			os.Getenv("GO_GET_FLAG"),
			eachDep)

		// cmdErr := exec.Command(.Run()
		if cmdErr != nil {
			return cmdErr
		}
	}
	return nil
}

// EnsureSpelling ensures that there are no misspellings in the source
func EnsureSpelling() error {
	goSpelling := func() error {
		return goSourceApply("misspell", "-error")
	}
	mdSpelling := func() error {
		return markdownSourceApply("misspell", "-error")
	}
	mg.SerialDeps(
		goSpelling,
		mdSpelling)
	return nil
}

// EnsureVet ensures that the source has been `go vet`ted
func EnsureVet() error {
	return goSourceApply("go", "tool", "vet")
}

// EnsureLint ensures that the source is `golint`ed
func EnsureLint() error {
	return goSourceApply("golint")
}

// EnsureFormatted ensures that the source code is formatted with goimports
func EnsureFormatted() error {
	return goSourceApply("goimports", "-d")
}

// EnsureStaticChecks ensures that the source code passes static code checks
func EnsureStaticChecks() error {
	// Megacheck
	megacheckErr := sh.Run("megacheck",
		"-ignore",
		"github.com/mweagle/Sparta/CONSTANTS.go:*")
	if megacheckErr != nil {
		return megacheckErr
	}
	// Gosec
	return sh.Run("gosec",
		"-exclude=G204,G505,G401",
		"./...")
}

// EnsureAllPreconditions ensures that the source passes *ALL* static `ensure*`
// precondition steps
func EnsureAllPreconditions() error {
	mg.SerialDeps(
		InstallBuildRequirements,
		EnsureVet,
		EnsureLint,
		EnsureFormatted,
		EnsureStaticChecks,
		EnsureSpelling,
	)
	return nil
}

// EnsureTravisBuildEnvironment is the command that sets up the Travis
// environment to run the build.
func EnsureTravisBuildEnvironment() error {
	mg.SerialDeps(InstallBuildRequirements)

	// Super run some commands
	travisComands := [][]string{
		[]string{"dep", "version"},
		[]string{"dep", "ensure", "-v"},
		[]string{"rsync", "-a", "--quiet", "--remove-source-files", "./vendor/", "$GOPATH/src"},
	}
	return mageScript(travisComands)
}

// Build the application
func Build() error {
	mg.Deps(EnsureAllPreconditions)
	return sh.Run("go", "build", ".")
}

// Clean the working directory
func Clean() error {
	cleanCommands := [][]string{
		[]string{"go", "clean", "."},
		[]string{"rm", "-rf", "./graph.html"},
		[]string{"rsync", "-a", "--quiet", "--remove-source-files", "./vendor/", "$GOPATH/src"},
	}
	return mageScript(cleanCommands)
}

// Describe runs the `TestDescribe` test to generate a describe HTML output
// file at graph.html
func Describe() error {
	describeCommands := [][]string{
		[]string{"rm", "-rf", "./graph.html"},
		[]string{"go", "test", "-v", "-run", "TestDescribe"},
	}
	return mageScript(describeCommands)
}

// Publish the latest source
func Publish() error {
	mg.SerialDeps(GenerateBuildInfo)

	describeCommands := [][]string{
		[]string{"echo", "Checking `git` tree status"},
		[]string{"git", "diff", "--exit-code"},
		// TODO - migrate to Go
		[]string{"./buildinfo.sh"},
		[]string{"git", "commit", "-a", "-m", "Tagging Sparta commit"},
		[]string{"git", "push", "origin"},
	}
	return mageScript(describeCommands)
}

// Test runs the Sparta tests
func Test() {
	testCommand := func() error {
		return sh.Run("go",
			"test",
			"-cover",
			"-race",
			"./...")
	}
	mg.SerialDeps(
		EnsureAllPreconditions,
		testCommand,
	)
}

// TestCover runs the test and opens up the resulting report
func TestCover() error {
	// mg.SerialDeps(
	// 	EnsureAllPreconditions,
	// )
	coverageReport := fmt.Sprintf("%s/cover.out", localWorkDir)
	testCoverCommands := [][]string{
		[]string{"go", "test", fmt.Sprintf("-coverprofile=%s", coverageReport), "."},
		[]string{"go", "tool", "cover", fmt.Sprintf("-html=%s", coverageReport)},
		[]string{"rm", coverageReport},
		[]string{"open", fmt.Sprintf("%s/cover.html", localWorkDir)},
	}
	return mageScript(testCoverCommands)
}

// TravisBuild is the task to build in the context of a Travis CI pipeline
func TravisBuild() error {
	mg.SerialDeps(EnsureTravisBuildEnvironment,
		Build,
		Test)
	return nil
}
<commit_msg>DRY clean tree command<commit_after>// +build mage

package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/magefile/mage/mg" // mg contains helpful utility functions, like Deps
	"github.com/magefile/mage/sh" // mg contains helpful utility functions, like Deps
)

const localWorkDir = "./sparta"

var header = strings.Repeat("-", 80)

var ignoreSubdirectoryPaths = []string{
	".vendor",
	".sparta",
	".vscode",
	"/resources/describe",
}

// Default target to run when none is specified
// If not set, running mage will list available targets
// var Default = Build

func mageScript(commands [][]string) error {
	for _, eachCommand := range commands {
		var commandErr error
		if len(eachCommand) <= 1 {
			commandErr = sh.Run(eachCommand[0])
		} else {
			commandErr = sh.Run(eachCommand[0], eachCommand[1:]...)
		}
		if commandErr != nil {
			return commandErr
		}
	}
	return nil
}
func mageLog(formatSpecifier string, args ...interface{}) {
	if mg.Verbose() {
		if len(args) != 0 {
			log.Printf(formatSpecifier, args...)
		} else {
			log.Printf(formatSpecifier)
		}
	}
}

func sourceFilesOfType(extension string) ([]string, error) {
	testExtension := strings.TrimPrefix(extension, ".")
	testExtension = fmt.Sprintf(".%s", testExtension)

	files := make([]string, 0)
	walker := func(path string, info os.FileInfo, err error) error {
		contains := false
		for _, eachComponent := range ignoreSubdirectoryPaths {
			contains = strings.Contains(path, eachComponent)
			if contains {
				break
			}
		}
		if !contains && (filepath.Ext(path) == testExtension) {
			files = append(files, path)
		}
		return nil
	}
	goSourceFilesErr := filepath.Walk(".", walker)
	return files, goSourceFilesErr
}

func applyToSource(fileExtension string, commandParts ...string) error {
	if len(commandParts) <= 0 {
		return errors.New("applyToSource requires a command to apply to source files")
	}
	eligibleSourceFiles, eligibleSourceFilesErr := sourceFilesOfType(fileExtension)
	if eligibleSourceFilesErr != nil {
		return eligibleSourceFilesErr
	}

	mageLog(header)
	mageLog("Applying `%s` to %d `*.%s` source files", commandParts[0], len(eligibleSourceFiles), fileExtension)
	mageLog(header)

	commandArgs := []string{}
	if len(commandParts) > 1 {
		for _, eachPart := range commandParts[1:] {
			commandArgs = append(commandArgs, eachPart)
		}
	}
	for _, eachFile := range eligibleSourceFiles {
		applyArgs := append(commandArgs, eachFile)
		applyErr := sh.Run(commandParts[0], applyArgs...)
		if applyErr != nil {
			return applyErr
		}
	}
	return nil
}

func markdownSourceApply(commandParts ...string) error {
	return applyToSource("md", commandParts...)
}
func goSourceApply(commandParts ...string) error {
	return applyToSource("go", commandParts...)
}

func EnsureCleanTree() error {
	cleanTreeScript := [][]string{
		// No dirty trees
		{"git", "diff", "--exit-code"},
	}
	return mageScript(cleanTreeScript)
}

// GenerateBuildInfo creates the automatic buildinfo.go file so that we can
// stamp the SHA into the binaries we build...
func GenerateBuildInfo() error {
	mage.SerialDeps(EnsureCleanTree)

	// The first thing we need is the `git` SHA
	cmd := exec.Command("git", "rev-parse", "HEAD")
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	if err != nil {
		return err
	}
	stdOutResult := strings.TrimSpace(string(stdout.Bytes()))

	// Super = update the buildinfo data
	buildInfoTemplate := `package sparta

// THIS FILE IS AUTOMATICALLY GENERATED
// DO NOT EDIT
// CREATED: %s

// SpartaGitHash is the commit hash of this Sparta library
const SpartaGitHash = "%s"
`
	updatedInfo := fmt.Sprintf(buildInfoTemplate, time.Now().UTC(), stdOutResult)
	// Write it to the output location...
	writeErr := ioutil.WriteFile("./buildinfo.go", []byte(updatedInfo), os.ModePerm)

	if writeErr != nil {
		return writeErr
	}
	commitGenerateCommands := [][]string{
		// No dirty trees
		{"git", "commit", "-a", "-m", `"Autogenerated build info"`},
	}
	return mageScript(commitGenerateCommands)

}

// GenerateConstants runs the set of commands that update the embedded CONSTANTS
// for both local and AWS Lambda execution
func GenerateConstants() error {
	generateCommands := [][]string{
		// Create the embedded version
		{"go", "run", "$GOPATH/src/github.com/mjibson/esc/main.go", "-o", "./CONSTANTS.go", "-private", "-pkg", "sparta", "./resources"},
		//Create a secondary CONSTANTS_AWSBINARY.go file with empty content.
		{"go", "run", "$GOPATH/src/github.com/mjibson/esc/main.go", "-o", "./CONSTANTS_AWSBINARY.go", "-private", "-pkg", "sparta", "./resources/awsbinary/README.md"},
		//The next step will insert the
		// build tags at the head of each file so that they are mutually exclusive
		{"go", "run", "./cmd/insertTags/main.go", "./CONSTANTS", "!lambdabinary"},
		{"go", "run", "./cmd/insertTags/main.go", "./CONSTANTS_AWSBINARY", "lambdabinary"},
		{"git", "commit", "-a", "-m", "Autogenerated constants"},
	}
	return mageScript(generateCommands)
}

// InstallBuildRequirements installs or updates the dependent
// packages that aren't referenced by the source, but are needed
// to build the Sparta source
func InstallBuildRequirements() error {
	mageLog("`go get` update flags (env.GO_GET_FLAG): %s", os.Getenv("GO_GET_FLAG"))

	requirements := []string{
		"github.com/golang/dep/...",
		"honnef.co/go/tools/cmd/megacheck",
		"honnef.co/go/tools/cmd/gosimple",
		"honnef.co/go/tools/cmd/unused",
		"honnef.co/go/tools/cmd/staticcheck",
		"golang.org/x/tools/cmd/goimports",
		"github.com/fzipp/gocyclo",
		"github.com/golang/lint/golint",
		"github.com/mjibson/esc",
		"github.com/securego/gosec/cmd/gosec/...",
		"github.com/client9/misspell/cmd/misspell",
	}
	for _, eachDep := range requirements {
		cmdErr := sh.Run("go",
			"get",
			os.Getenv("GO_GET_FLAG"),
			eachDep)

		// cmdErr := exec.Command(.Run()
		if cmdErr != nil {
			return cmdErr
		}
	}
	return nil
}

// EnsureSpelling ensures that there are no misspellings in the source
func EnsureSpelling() error {
	goSpelling := func() error {
		return goSourceApply("misspell", "-error")
	}
	mdSpelling := func() error {
		return markdownSourceApply("misspell", "-error")
	}
	mg.SerialDeps(
		goSpelling,
		mdSpelling)
	return nil
}

// EnsureVet ensures that the source has been `go vet`ted
func EnsureVet() error {
	return goSourceApply("go", "tool", "vet")
}

// EnsureLint ensures that the source is `golint`ed
func EnsureLint() error {
	return goSourceApply("golint")
}

// EnsureFormatted ensures that the source code is formatted with goimports
func EnsureFormatted() error {
	return goSourceApply("goimports", "-d")
}

// EnsureStaticChecks ensures that the source code passes static code checks
func EnsureStaticChecks() error {
	// Megacheck
	megacheckErr := sh.Run("megacheck",
		"-ignore",
		"github.com/mweagle/Sparta/CONSTANTS.go:*")
	if megacheckErr != nil {
		return megacheckErr
	}
	// Gosec
	return sh.Run("gosec",
		"-exclude=G204,G505,G401",
		"./...")
}

// EnsureAllPreconditions ensures that the source passes *ALL* static `ensure*`
// precondition steps
func EnsureAllPreconditions() error {
	mg.SerialDeps(
		InstallBuildRequirements,
		EnsureVet,
		EnsureLint,
		EnsureFormatted,
		EnsureStaticChecks,
		EnsureSpelling,
	)
	return nil
}

// EnsureTravisBuildEnvironment is the command that sets up the Travis
// environment to run the build.
func EnsureTravisBuildEnvironment() error {
	mg.SerialDeps(InstallBuildRequirements)

	// Super run some commands
	travisComands := [][]string{
		[]string{"dep", "version"},
		[]string{"dep", "ensure", "-v"},
		[]string{"rsync", "-a", "--quiet", "--remove-source-files", "./vendor/", "$GOPATH/src"},
	}
	return mageScript(travisComands)
}

// Build the application
func Build() error {
	mg.Deps(EnsureAllPreconditions)
	return sh.Run("go", "build", ".")
}

// Clean the working directory
func Clean() error {
	cleanCommands := [][]string{
		[]string{"go", "clean", "."},
		[]string{"rm", "-rf", "./graph.html"},
		[]string{"rsync", "-a", "--quiet", "--remove-source-files", "./vendor/", "$GOPATH/src"},
	}
	return mageScript(cleanCommands)
}

// Describe runs the `TestDescribe` test to generate a describe HTML output
// file at graph.html
func Describe() error {
	describeCommands := [][]string{
		[]string{"rm", "-rf", "./graph.html"},
		[]string{"go", "test", "-v", "-run", "TestDescribe"},
	}
	return mageScript(describeCommands)
}

// Publish the latest source
func Publish() error {
	mg.SerialDeps(GenerateBuildInfo)

	describeCommands := [][]string{
		[]string{"echo", "Checking `git` tree status"},
		[]string{"git", "diff", "--exit-code"},
		// TODO - migrate to Go
		[]string{"./buildinfo.sh"},
		[]string{"git", "commit", "-a", "-m", "Tagging Sparta commit"},
		[]string{"git", "push", "origin"},
	}
	return mageScript(describeCommands)
}

// Test runs the Sparta tests
func Test() {
	testCommand := func() error {
		return sh.Run("go",
			"test",
			"-cover",
			"-race",
			"./...")
	}
	mg.SerialDeps(
		EnsureAllPreconditions,
		testCommand,
	)
}

// TestCover runs the test and opens up the resulting report
func TestCover() error {
	// mg.SerialDeps(
	// 	EnsureAllPreconditions,
	// )
	coverageReport := fmt.Sprintf("%s/cover.out", localWorkDir)
	testCoverCommands := [][]string{
		[]string{"go", "test", fmt.Sprintf("-coverprofile=%s", coverageReport), "."},
		[]string{"go", "tool", "cover", fmt.Sprintf("-html=%s", coverageReport)},
		[]string{"rm", coverageReport},
		[]string{"open", fmt.Sprintf("%s/cover.html", localWorkDir)},
	}
	return mageScript(testCoverCommands)
}

// TravisBuild is the task to build in the context of a Travis CI pipeline
func TravisBuild() error {
	mg.SerialDeps(EnsureTravisBuildEnvironment,
		Build,
		Test)
	return nil
}
<|endoftext|>


<commit_before>// +build mage

package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/magefile/mage/mg"
	"github.com/magefile/mage/sh"
)

// Default mage target
var Default = Build

var (
	chocoPath      = path.Join("chocolatey")
	binPath        = path.Join("bin")
	chocoBinPath   = path.Join("bin", "choco")
	chocoLegalPath = path.Join(chocoBinPath, "legal")
	chocoToolsPath = path.Join(chocoBinPath, "tools")
	chocoNuspec    = path.Join(chocoBinPath, "windowsspyblocker.nuspec")
	wsbPath        = path.Join(binPath, "WindowsSpyBlocker.exe")
	wsbEnv         = map[string]string{
		"GO111MODULE": "on",
		"GOPROXY":     "https://goproxy.io,direct",
		"GOOS":        "windows",
		"GOARCH":      "386",
		"CGO_ENABLED": "0",
	}
)

// Build Run go build
func Build() error {
	mg.Deps(Clean)
	mg.Deps(Generate)

	var args []string
	args = append(args, "build", "-o", wsbPath, "-v")
	args = append(args, "-ldflags", flags())

	fmt.Println("⚙ Go build...")
	if err := sh.RunWith(wsbEnv, mg.GoCmd(), args...); err != nil {
		return err
	}

	return nil
}

// Clean Remove files generated at build-time
func Clean() error {
	if err := createDir(binPath); err != nil {
		return err
	}
	if err := cleanDir(binPath); err != nil {
		return err
	}
	return nil
}

// Download Run go mod download
func Download() error {

	fmt.Println("⚙ Go mod download...")
	if err := sh.RunWith(wsbEnv, mg.GoCmd(), "mod", "download"); err != nil {
		return err
	}

	return nil
}

// Generate Run go generate
func Generate() error {
	mg.Deps(Download)
	mg.Deps(appConf)
	mg.Deps(manifest)
	mg.Deps(versionInfo)

	fmt.Println("⚙ Go generate...")
	if err := sh.RunV(mg.GoCmd(), "generate", "-v"); err != nil {
		return err
	}

	return nil
}

// ChocoPack Run choco pack
func ChocoPack() error {
	mg.Deps(ChocoPrepare)

	fmt.Println("⚙ Chocolatey package...")
	choco, err := exec.LookPath("choco")
	if err != nil {
		return err
	}

	var args []string
	args = append(args, "pack", "--out", binPath)
	args = append(args, "--version", tag())
	args = append(args, "--acceptlicense", "--yes")
	args = append(args, chocoNuspec)

	if err := sh.RunV(choco, args...); err != nil {
		return err
	}

	return nil
}

// ChocoPush Run choco push
func ChocoPush() error {
	fmt.Println("⚙ Chocolatey push...")
	choco, err := exec.LookPath("choco")
	if err != nil {
		return err
	}

	nupkg := fmt.Sprintf("windowsspyblocker.%s.nupkg", tag())

	var args []string
	args = append(args, "push", path.Join(binPath, nupkg))
	args = append(args, "--source", "https://package.chocolatey.org")
	args = append(args, "--apikey", os.Getenv("CHOCO_API_KEY"))
	args = append(args, "--acceptlicense", "--yes")

	if err := sh.RunV(choco, args...); err != nil {
		return err
	}

	return nil
}

// ChocoPrepare Generate chocolatey files
func ChocoPrepare() error {
	fmt.Println("🔨 Generating Chocolatey files...")

	if err := createDir(chocoBinPath); err != nil {
		return err
	}
	if err := cleanDir(chocoBinPath); err != nil {
		return err
	}
	if err := copyDir(chocoPath, chocoBinPath); err != nil {
		return err
	}
	if err := createDir(chocoLegalPath); err != nil {
		return err
	}
	if err := copyFile("LICENSE", path.Join(chocoLegalPath, "LICENSE.txt")); err != nil {
		return err
	}
	if err := copyFile(wsbPath, path.Join(chocoToolsPath, "WindowsSpyBlocker.exe")); err != nil {
		return err
	}

	nuspec, err := ioutil.ReadFile(chocoNuspec)
	if err != nil {
		return err
	}
	nuspecContent := strings.Replace(string(nuspec), "<version>0.0.0</version>", fmt.Sprintf("<version>%s</version>", tag()), -1)
	err = ioutil.WriteFile(chocoNuspec, []byte(nuspecContent), 0)
	if err != nil {
		return err
	}

	return nil
}

// flags returns ldflags
func flags() string {
	//hash := hash()
	mod := mod()
	tag := tag()
	return fmt.Sprintf(`-s -w -X "%s/app/utils/config.AppVersion=%s"`, mod, tag)
}

// mod returns module name
func mod() string {
	f, err := os.Open("go.mod")
	if err == nil {
		reader := bufio.NewReader(f)
		line, _, _ := reader.ReadLine()
		return strings.Replace(string(line), "module ", "", 1)
	}
	return ""
}

// tag returns the git tag for the current branch or "" if none.
func tag() string {
	s, _ := sh.Output("bash", "-c", "git describe --abbrev=0 --tags 2> /dev/null")
	if s == "" {
		return "0.0.0"
	}
	return s
}

// hash returns the git hash for the current repo or "" if none.
func hash() string {
	hash, _ := sh.Output("git", "rev-parse", "--short", "HEAD")
	return hash
}

// appConf generates app.conf file
func appConf() error {
	fmt.Println("🔨 Generating app.conf...")

	var tpl = template.Must(template.New("").Parse(`{
  "version": "{{ .Version }}",
  "debug": false,
  "useEmbeddedData": true,
  "proxifier": {
    "logPath": "C:/Users/[username]/Documents/Proxifier/Log.txt"
  },
  "sysmon": {
    "evtxPath": "C:/WINDOWS/system32/winevt/Logs/Microsoft-Windows-Sysmon%4Operational.evtx"
  },
  "wireshark": {
    "pcapngPath": "C:/Users/[username]/Documents/Wireshark/cap.pcapng",
    "capture": {
      "interface": 1,
      "filter": "not arp and port not 53 and not icmp and not icmp6 and not broadcast"
    }
  },
  "exclude": {
    "ips": [
      "0.0.0.0",
      "127.0.0.1",
      "192.168.0.0-192.168.0.255",
      "8.8.8.8",
      "8.8.4.4",
      "255.255.255.255"
    ],
    "hosts": [
      "MyComputer",
      "localhost",
      "localhost.localdomain",
      "*.local",
      "yourISP.com",
      "*.yourISP.com",
      "wireshark.org",
      "*.wireshark.org"
    ],
    "orgs": [
      "*facebook*"
    ]
  }
}
`))

	f, err := os.Create("app.conf")
	if err != nil {
		return err
	}
	defer f.Close()

	return tpl.Execute(f, struct {
		Version string
	}{
		Version: tag(),
	})
}

// manifest generates manifest for versioninfo
func manifest() error {
	fmt.Println("🔨 Generating app.manifest...")

	file, err := os.Create("app.manifest")
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.WriteString(file, `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<assembly xmlns="urn:schemas-microsoft-com:asm.v1" manifestVersion="1.0">
  <compatibility xmlns="urn:schemas-microsoft-com:compatibility.v1">
    <application>
      <!--This Id value indicates the application supports Windows 7 functionality-->
      <supportedOS Id="{35138b9a-5d96-4fbd-8e2d-a2440225f93a}"/>
      <!--This Id value indicates the application supports Windows 8 functionality-->
      <supportedOS Id="{4a2f28e3-53b9-4441-ba9c-d69d4a4a6e38}"/>
      <!--This Id value indicates the application supports Windows 8.1 functionality-->
      <supportedOS Id="{1f676c76-80e1-4239-95bb-83d0f6d0da78}"/>
      <!--This Id value indicates the application supports Windows 10 functionality-->
      <supportedOS Id="{8e0f7a12-bfb3-4fe8-b9a5-48fd50a15a9a}"/>
    </application>
  </compatibility>
  <trustInfo xmlns="urn:schemas-microsoft-com:asm.v3">
    <security>
      <requestedPrivileges>
        <requestedExecutionLevel level="requireAdministrator" uiAccess="false"/>
      </requestedPrivileges>
    </security>
  </trustInfo>
</assembly>`)
	if err != nil {
		return err
	}

	return nil
}

// versionInfo generates versioninfo.json
func versionInfo() error {
	fmt.Println("🔨 Generating versioninfo.json...")

	var tpl = template.Must(template.New("").Parse(`{
	"FixedFileInfo":
	{
		"FileFlagsMask": "3f",
		"FileFlags ": "00",
		"FileOS": "040004",
		"FileType": "01",
		"FileSubType": "00"
	},
	"StringFileInfo":
	{
		"Comments": "",
		"CompanyName": "",
		"FileDescription": "Block spying and tracking on Windows",
		"FileVersion": "{{ .Version }}.0",
		"InternalName": "",
		"LegalCopyright": "https://{{ .Package }}",
		"LegalTrademarks": "",
		"OriginalFilename": "WindowsSpyBlocker.exe",
		"PrivateBuild": "",
		"ProductName": "WindowsSpyBlocker",
		"ProductVersion": "{{ .Version }}.0",
		"SpecialBuild": ""
	},
	"VarFileInfo":
	{
		"Translation": {
			"LangID": "0409",
			"CharsetID": "04B0"
		}
	}
}`))

	f, err := os.Create("versioninfo.json")
	if err != nil {
		return err
	}
	defer f.Close()

	return tpl.Execute(f, struct {
		Package string
		Version string
	}{
		Package: mod(),
		Version: tag(),
	})
}

func createDir(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return os.MkdirAll(path, 777)
	}
	return nil
}

func cleanDir(dir string) error {
	d, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(dir, name))
		if err != nil {
			return err
		}
	}
	return nil
}

func copyDir(src string, dst string) error {
	var err error
	var fds []os.FileInfo
	var srcinfo os.FileInfo

	if srcinfo, err = os.Stat(src); err != nil {
		return err
	}

	if err = os.MkdirAll(dst, srcinfo.Mode()); err != nil {
		return err
	}

	if fds, err = ioutil.ReadDir(src); err != nil {
		return err
	}

	for _, fd := range fds {
		srcfp := path.Join(src, fd.Name())
		dstfp := path.Join(dst, fd.Name())

		if fd.IsDir() {
			if err = copyDir(srcfp, dstfp); err != nil {
				fmt.Println(err)
			}
		} else {
			if err = copyFile(srcfp, dstfp); err != nil {
				fmt.Println(err)
			}
		}
	}

	return nil
}

func copyFile(src string, dest string) error {
	srcFile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	destFile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destFile.Close()

	_, err = io.Copy(destFile, srcFile)
	if err != nil {
		return err
	}

	err = destFile.Sync()
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Fix magefile<commit_after>// +build mage

package main

import (
	"bufio"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	"github.com/magefile/mage/mg"
	"github.com/magefile/mage/sh"
)

// Default mage target
var Default = Build

var (
	chocoPath      = path.Join("chocolatey")
	binPath        = path.Join("bin")
	chocoBinPath   = path.Join("bin", "choco")
	chocoLegalPath = path.Join(chocoBinPath, "legal")
	chocoToolsPath = path.Join(chocoBinPath, "tools")
	chocoNuspec    = path.Join(chocoBinPath, "windowsspyblocker.nuspec")
	wsbPath        = path.Join(binPath, "WindowsSpyBlocker.exe")
	wsbEnv         = map[string]string{
		"GO111MODULE": "on",
		"GOOS":        "windows",
		"GOARCH":      "386",
		"CGO_ENABLED": "0",
	}
)

// Build Run go build
func Build() error {
	mg.Deps(Clean)
	mg.Deps(Generate)

	var args []string
	args = append(args, "build", "-o", wsbPath, "-v")
	args = append(args, "-ldflags", flags())

	fmt.Println("⚙ Go build...")
	if err := sh.RunWith(wsbEnv, mg.GoCmd(), args...); err != nil {
		return err
	}

	return nil
}

// Clean Remove files generated at build-time
func Clean() error {
	if err := createDir(binPath); err != nil {
		return err
	}
	if err := cleanDir(binPath); err != nil {
		return err
	}
	return nil
}

// Download Run go mod download
func Download() error {

	fmt.Println("⚙ Go mod download...")
	if err := sh.RunWith(wsbEnv, mg.GoCmd(), "mod", "download"); err != nil {
		return err
	}

	return nil
}

// Generate Run go generate
func Generate() error {
	mg.Deps(Download)
	mg.Deps(appConf)
	mg.Deps(manifest)
	mg.Deps(versionInfo)

	fmt.Println("⚙ Go generate...")
	if err := sh.RunV(mg.GoCmd(), "generate", "-v"); err != nil {
		return err
	}

	return nil
}

// ChocoPack Run choco pack
func ChocoPack() error {
	mg.Deps(ChocoPrepare)

	fmt.Println("⚙ Chocolatey package...")
	choco, err := exec.LookPath("choco")
	if err != nil {
		return err
	}

	var args []string
	args = append(args, "pack", "--out", binPath)
	args = append(args, "--version", tag())
	args = append(args, "--acceptlicense", "--yes")
	args = append(args, chocoNuspec)

	if err := sh.RunV(choco, args...); err != nil {
		return err
	}

	return nil
}

// ChocoPush Run choco push
func ChocoPush() error {
	fmt.Println("⚙ Chocolatey push...")
	choco, err := exec.LookPath("choco")
	if err != nil {
		return err
	}

	nupkg := fmt.Sprintf("windowsspyblocker.%s.nupkg", tag())

	var args []string
	args = append(args, "push", path.Join(binPath, nupkg))
	args = append(args, "--source", "https://package.chocolatey.org")
	args = append(args, "--apikey", os.Getenv("CHOCO_API_KEY"))
	args = append(args, "--acceptlicense", "--yes")

	if err := sh.RunV(choco, args...); err != nil {
		return err
	}

	return nil
}

// ChocoPrepare Generate chocolatey files
func ChocoPrepare() error {
	fmt.Println("🔨 Generating Chocolatey files...")

	if err := createDir(chocoBinPath); err != nil {
		return err
	}
	if err := cleanDir(chocoBinPath); err != nil {
		return err
	}
	if err := copyDir(chocoPath, chocoBinPath); err != nil {
		return err
	}
	if err := createDir(chocoLegalPath); err != nil {
		return err
	}
	if err := copyFile("LICENSE", path.Join(chocoLegalPath, "LICENSE.txt")); err != nil {
		return err
	}
	if err := copyFile(wsbPath, path.Join(chocoToolsPath, "WindowsSpyBlocker.exe")); err != nil {
		return err
	}

	nuspec, err := ioutil.ReadFile(chocoNuspec)
	if err != nil {
		return err
	}
	nuspecContent := strings.Replace(string(nuspec), "<version>0.0.0</version>", fmt.Sprintf("<version>%s</version>", tag()), -1)
	err = ioutil.WriteFile(chocoNuspec, []byte(nuspecContent), 0)
	if err != nil {
		return err
	}

	return nil
}

// flags returns ldflags
func flags() string {
	//hash := hash()
	mod := mod()
	tag := tag()
	return fmt.Sprintf(`-s -w -X "%s/app/utils/config.AppVersion=%s"`, mod, tag)
}

// mod returns module name
func mod() string {
	f, err := os.Open("go.mod")
	if err == nil {
		reader := bufio.NewReader(f)
		line, _, _ := reader.ReadLine()
		return strings.Replace(string(line), "module ", "", 1)
	}
	return ""
}

// tag returns the git tag for the current branch or "" if none.
func tag() string {
	s, _ := sh.Output("bash", "-c", "git describe --abbrev=0 --tags 2> /dev/null")
	if s == "" {
		return "0.0.0"
	}
	return s
}

// hash returns the git hash for the current repo or "" if none.
func hash() string {
	hash, _ := sh.Output("git", "rev-parse", "--short", "HEAD")
	return hash
}

// appConf generates app.conf file
func appConf() error {
	fmt.Println("🔨 Generating app.conf...")

	var tpl = template.Must(template.New("").Parse(`{
  "version": "{{ .Version }}",
  "debug": false,
  "useEmbeddedData": true,
  "proxifier": {
    "logPath": "C:/Users/[username]/Documents/Proxifier/Log.txt"
  },
  "sysmon": {
    "evtxPath": "C:/WINDOWS/system32/winevt/Logs/Microsoft-Windows-Sysmon%4Operational.evtx"
  },
  "wireshark": {
    "pcapngPath": "C:/Users/[username]/Documents/Wireshark/cap.pcapng",
    "capture": {
      "interface": 1,
      "filter": "not arp and port not 53 and not icmp and not icmp6 and not broadcast"
    }
  },
  "exclude": {
    "ips": [
      "0.0.0.0",
      "127.0.0.1",
      "192.168.0.0-192.168.0.255",
      "8.8.8.8",
      "8.8.4.4",
      "255.255.255.255"
    ],
    "hosts": [
      "MyComputer",
      "localhost",
      "localhost.localdomain",
      "*.local",
      "yourISP.com",
      "*.yourISP.com",
      "wireshark.org",
      "*.wireshark.org"
    ],
    "orgs": [
      "*facebook*"
    ]
  }
}
`))

	f, err := os.Create("app.conf")
	if err != nil {
		return err
	}
	defer f.Close()

	return tpl.Execute(f, struct {
		Version string
	}{
		Version: tag(),
	})
}

// manifest generates manifest for versioninfo
func manifest() error {
	fmt.Println("🔨 Generating app.manifest...")

	file, err := os.Create("app.manifest")
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.WriteString(file, `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<assembly xmlns="urn:schemas-microsoft-com:asm.v1" manifestVersion="1.0">
  <compatibility xmlns="urn:schemas-microsoft-com:compatibility.v1">
    <application>
      <!--This Id value indicates the application supports Windows 7 functionality-->
      <supportedOS Id="{35138b9a-5d96-4fbd-8e2d-a2440225f93a}"/>
      <!--This Id value indicates the application supports Windows 8 functionality-->
      <supportedOS Id="{4a2f28e3-53b9-4441-ba9c-d69d4a4a6e38}"/>
      <!--This Id value indicates the application supports Windows 8.1 functionality-->
      <supportedOS Id="{1f676c76-80e1-4239-95bb-83d0f6d0da78}"/>
      <!--This Id value indicates the application supports Windows 10 functionality-->
      <supportedOS Id="{8e0f7a12-bfb3-4fe8-b9a5-48fd50a15a9a}"/>
    </application>
  </compatibility>
  <trustInfo xmlns="urn:schemas-microsoft-com:asm.v3">
    <security>
      <requestedPrivileges>
        <requestedExecutionLevel level="requireAdministrator" uiAccess="false"/>
      </requestedPrivileges>
    </security>
  </trustInfo>
</assembly>`)
	if err != nil {
		return err
	}

	return nil
}

// versionInfo generates versioninfo.json
func versionInfo() error {
	fmt.Println("🔨 Generating versioninfo.json...")

	var tpl = template.Must(template.New("").Parse(`{
	"FixedFileInfo":
	{
		"FileFlagsMask": "3f",
		"FileFlags ": "00",
		"FileOS": "040004",
		"FileType": "01",
		"FileSubType": "00"
	},
	"StringFileInfo":
	{
		"Comments": "",
		"CompanyName": "",
		"FileDescription": "Block spying and tracking on Windows",
		"FileVersion": "{{ .Version }}.0",
		"InternalName": "",
		"LegalCopyright": "https://{{ .Package }}",
		"LegalTrademarks": "",
		"OriginalFilename": "WindowsSpyBlocker.exe",
		"PrivateBuild": "",
		"ProductName": "WindowsSpyBlocker",
		"ProductVersion": "{{ .Version }}.0",
		"SpecialBuild": ""
	},
	"VarFileInfo":
	{
		"Translation": {
			"LangID": "0409",
			"CharsetID": "04B0"
		}
	}
}`))

	f, err := os.Create("versioninfo.json")
	if err != nil {
		return err
	}
	defer f.Close()

	return tpl.Execute(f, struct {
		Package string
		Version string
	}{
		Package: mod(),
		Version: tag(),
	})
}

func createDir(path string) error {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return os.MkdirAll(path, 777)
	}
	return nil
}

func cleanDir(dir string) error {
	d, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(dir, name))
		if err != nil {
			return err
		}
	}
	return nil
}

func copyDir(src string, dst string) error {
	var err error
	var fds []os.FileInfo
	var srcinfo os.FileInfo

	if srcinfo, err = os.Stat(src); err != nil {
		return err
	}

	if err = os.MkdirAll(dst, srcinfo.Mode()); err != nil {
		return err
	}

	if fds, err = ioutil.ReadDir(src); err != nil {
		return err
	}

	for _, fd := range fds {
		srcfp := path.Join(src, fd.Name())
		dstfp := path.Join(dst, fd.Name())

		if fd.IsDir() {
			if err = copyDir(srcfp, dstfp); err != nil {
				fmt.Println(err)
			}
		} else {
			if err = copyFile(srcfp, dstfp); err != nil {
				fmt.Println(err)
			}
		}
	}

	return nil
}

func copyFile(src string, dest string) error {
	srcFile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	destFile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destFile.Close()

	_, err = io.Copy(destFile, srcFile)
	if err != nil {
		return err
	}

	err = destFile.Sync()
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package security_groups_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/runner"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers"
)

var _ = Describe("Security Groups", func() {

	type AppResource struct {
		Metadata struct {
			Url string
		}
	}
	type AppsResponse struct {
		Resources []AppResource
	}

	type Stat struct {
		Stats struct {
			Host string
			Port int
		}
	}
	type StatsResponse map[string]Stat

	type DoraCurlResponse struct {
		Stdout     string
		Stderr     string
		ReturnCode int `json:"return_code"`
	}

	var serverAppName, securityGroupName, privateHost string
	var privatePort int

	BeforeEach(func() {
		serverAppName = generator.RandomName()
		Expect(cf.Cf("push", serverAppName, "-p", helpers.NewAssets().Dora).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

		// gather app url
		var appsResponse AppsResponse
		cfResponse := cf.Cf("curl", fmt.Sprintf("/v2/apps?q=name:%s", serverAppName)).Wait(DEFAULT_TIMEOUT).Out.Contents()
		json.Unmarshal(cfResponse, &appsResponse)
		serverAppUrl := appsResponse.Resources[0].Metadata.Url

		// gather app stats for dea ip and app port
		var statsResponse StatsResponse
		cfResponse = cf.Cf("curl", fmt.Sprintf("%s/stats", serverAppUrl)).Wait(DEFAULT_TIMEOUT).Out.Contents()
		json.Unmarshal(cfResponse, &statsResponse)

		privateHost = statsResponse["0"].Stats.Host
		privatePort = statsResponse["0"].Stats.Port
	})

	AfterEach(func() {
		Expect(cf.Cf("delete", serverAppName, "-f").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
	})

	// this test assumes the default running security groups block access to the DEAs
	// the test takes advantage of the fact that the DEA ip address and internal container ip address
	//  are discoverable via the cc api and dora's myip endpoint
	It("allows previously-blocked ip traffic after applying a security group, and re-blocks it when the group is removed", func() {
		clientAppName := generator.RandomName()
		Expect(cf.Cf("push", clientAppName, "-p", helpers.NewAssets().Dora).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		defer func() { cf.Cf("delete", clientAppName, "-f").Wait(CF_PUSH_TIMEOUT) }()

		// gather container ip
		curlResponse := helpers.CurlApp(serverAppName, "/myip")
		containerIp := strings.TrimSpace(curlResponse)

		// test app egress rules
		var doraCurlResponse DoraCurlResponse
		curlResponse = helpers.CurlApp(clientAppName, fmt.Sprintf("/curl/%s/%d", privateHost, privatePort))
		json.Unmarshal([]byte(curlResponse), &doraCurlResponse)
		Expect(doraCurlResponse.ReturnCode).ToNot(Equal(0))

		// apply security group
		rules := fmt.Sprintf(
			`[{"destination":"%s","ports":"%d","protocol":"tcp"},
        {"destination":"%s","ports":"%d","protocol":"tcp"}]`,
			privateHost, privatePort, containerIp, privatePort)

		file, _ := ioutil.TempFile(os.TempDir(), "CATS-sg-rules")
		defer os.Remove(file.Name())
		file.WriteString(rules)

		rulesPath := file.Name()
		securityGroupName = fmt.Sprintf("CATS-SG-%s", generator.RandomName())

		cf.AsUser(context.AdminUserContext(), func() {
			Expect(cf.Cf("create-security-group", securityGroupName, rulesPath).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			Expect(
				cf.Cf("bind-security-group",
					securityGroupName,
					context.RegularUserContext().Org,
					context.RegularUserContext().Space).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
		defer func() {
			cf.AsUser(context.AdminUserContext(), func() {
				Expect(cf.Cf("delete-security-group", securityGroupName, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})
		}()

		Expect(cf.Cf("restart", clientAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

		// test app egress rules
		curlResponse = helpers.CurlApp(clientAppName, fmt.Sprintf("/curl/%s/%d", privateHost, privatePort))
		json.Unmarshal([]byte(curlResponse), &doraCurlResponse)
		Expect(doraCurlResponse.ReturnCode).To(Equal(0))

		// unapply security group
		cf.AsUser(context.AdminUserContext(), func() {
			Expect(cf.Cf("unbind-security-group", securityGroupName, context.RegularUserContext().Org, context.RegularUserContext().Space).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
		Expect(cf.Cf("restart", clientAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

		// test app egress rules
		curlResponse = helpers.CurlApp(clientAppName, fmt.Sprintf("/curl/%s/%d", privateHost, privatePort))
		json.Unmarshal([]byte(curlResponse), &doraCurlResponse)
		Expect(doraCurlResponse.ReturnCode).ToNot(Equal(0))
	})

	It("allows external and denies internal traffic during staging based on default staging security rules", func() {
		buildpack := fmt.Sprintf("CATS-SGBP-%s", generator.RandomName())
		testAppName := generator.RandomName()
		privateUri := fmt.Sprintf("%s:%d", privateHost, privatePort)

		buildpackZip := fmt.Sprintf("%s%s.zip", os.TempDir(), buildpack)
		Expect(runner.Run("zip", "-r", "-q", buildpackZip, helpers.NewAssets().SecurityGroupBuildpack).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		defer os.Remove(buildpackZip)

		cf.AsUser(context.AdminUserContext(), func() {
			Expect(cf.Cf("create-buildpack", buildpack, buildpackZip, "999").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
		defer func() {
			cf.AsUser(context.AdminUserContext(), func() {
				Expect(cf.Cf("delete-buildpack", buildpack, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})
		}()

		Expect(cf.Cf("push", testAppName, "-b", buildpack, "-p", helpers.NewAssets().HelloWorld, "--no-start").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		defer func() { cf.Cf("delete", testAppName, "-f").Wait(CF_PUSH_TIMEOUT) }()

		Expect(cf.Cf("set-env", testAppName, "TESTURI", "www.google.com").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		Expect(cf.Cf("start", testAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(1))
		Eventually(func() *Session {
			appLogsSession := cf.Cf("logs", "--recent", testAppName)
			Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			return appLogsSession
		}, 5).Should(Say("CURL_EXIT=0"))

		Expect(cf.Cf("set-env", testAppName, "TESTURI", privateUri).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		Expect(cf.Cf("restart", testAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(1))
		Eventually(func() *Session {
			appLogsSession := cf.Cf("logs", "--recent", testAppName)
			Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			return appLogsSession
		}, 5).Should(Say("CURL_EXIT=[^0]"))
	})
})
<commit_msg>Fix tests<commit_after>package security_groups_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"strings"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/runner"
	"github.com/cloudfoundry/cf-acceptance-tests/helpers"
)

var _ = Describe("Security Groups", func() {

	type AppResource struct {
		Metadata struct {
			Url string
		}
	}
	type AppsResponse struct {
		Resources []AppResource
	}

	type Stat struct {
		Stats struct {
			Host string
			Port int
		}
	}
	type StatsResponse map[string]Stat

	type DoraCurlResponse struct {
		Stdout     string
		Stderr     string
		ReturnCode int `json:"return_code"`
	}

	var serverAppName, securityGroupName, privateHost string
	var privatePort int

	BeforeEach(func() {
		serverAppName = generator.RandomName()
		Expect(cf.Cf("push", serverAppName, "-p", helpers.NewAssets().Dora).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

		// gather app url
		var appsResponse AppsResponse
		cfResponse := cf.Cf("curl", fmt.Sprintf("/v2/apps?q=name:%s", serverAppName)).Wait(DEFAULT_TIMEOUT).Out.Contents()
		json.Unmarshal(cfResponse, &appsResponse)
		serverAppUrl := appsResponse.Resources[0].Metadata.Url

		// gather app stats for dea ip and app port
		var statsResponse StatsResponse
		cfResponse = cf.Cf("curl", fmt.Sprintf("%s/stats", serverAppUrl)).Wait(DEFAULT_TIMEOUT).Out.Contents()
		json.Unmarshal(cfResponse, &statsResponse)

		privateHost = statsResponse["0"].Stats.Host
		privatePort = statsResponse["0"].Stats.Port
	})

	AfterEach(func() {
		Expect(cf.Cf("delete", serverAppName, "-f").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
	})

	// this test assumes the default running security groups block access to the DEAs
	// the test takes advantage of the fact that the DEA ip address and internal container ip address
	//  are discoverable via the cc api and dora's myip endpoint
	It("allows previously-blocked ip traffic after applying a security group, and re-blocks it when the group is removed", func() {
		clientAppName := generator.RandomName()
		Expect(cf.Cf("push", clientAppName, "-p", helpers.NewAssets().Dora).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		defer func() { cf.Cf("delete", clientAppName, "-f").Wait(CF_PUSH_TIMEOUT) }()

		// gather container ip
		curlResponse := helpers.CurlApp(serverAppName, "/myip")
		containerIp := strings.TrimSpace(curlResponse)

		// test app egress rules
		var doraCurlResponse DoraCurlResponse
		curlResponse = helpers.CurlApp(clientAppName, fmt.Sprintf("/curl/%s/%d", privateHost, privatePort))
		json.Unmarshal([]byte(curlResponse), &doraCurlResponse)
		Expect(doraCurlResponse.ReturnCode).ToNot(Equal(0))

		// apply security group
		rules := fmt.Sprintf(
			`[{"destination":"%s","ports":"%d","protocol":"tcp"},
        {"destination":"%s","ports":"%d","protocol":"tcp"}]`,
			privateHost, privatePort, containerIp, privatePort)

		file, _ := ioutil.TempFile(os.TempDir(), "CATS-sg-rules")
		defer os.Remove(file.Name())
		file.WriteString(rules)

		rulesPath := file.Name()
		securityGroupName = fmt.Sprintf("CATS-SG-%s", generator.RandomName())

		cf.AsUser(context.AdminUserContext(), func() {
			Expect(cf.Cf("create-security-group", securityGroupName, rulesPath).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			Expect(
				cf.Cf("bind-security-group",
					securityGroupName,
					context.RegularUserContext().Org,
					context.RegularUserContext().Space).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
		defer func() {
			cf.AsUser(context.AdminUserContext(), func() {
				Expect(cf.Cf("delete-security-group", securityGroupName, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})
		}()

		Expect(cf.Cf("restart", clientAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

		// test app egress rules
		curlResponse = helpers.CurlApp(clientAppName, fmt.Sprintf("/curl/%s/%d", privateHost, privatePort))
		json.Unmarshal([]byte(curlResponse), &doraCurlResponse)
		Expect(doraCurlResponse.ReturnCode).To(Equal(0))

		// unapply security group
		cf.AsUser(context.AdminUserContext(), func() {
			Expect(cf.Cf("unbind-security-group", securityGroupName, context.RegularUserContext().Org, context.RegularUserContext().Space).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
		Expect(cf.Cf("restart", clientAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(0))

		// test app egress rules
		curlResponse = helpers.CurlApp(clientAppName, fmt.Sprintf("/curl/%s/%d", privateHost, privatePort))
		json.Unmarshal([]byte(curlResponse), &doraCurlResponse)
		Expect(doraCurlResponse.ReturnCode).ToNot(Equal(0))
	})

	It("allows external and denies internal traffic during staging based on default staging security rules", func() {
		buildpack := fmt.Sprintf("CATS-SGBP-%s", generator.RandomName())
		testAppName := generator.RandomName()
		privateUri := fmt.Sprintf("%s:%d", privateHost, privatePort)

		buildpackZip := path.Join(os.TempDir(), Sprintf("%s.zip", buildpack))
		Expect(runner.Run("zip", "-r", "-q", buildpackZip, helpers.NewAssets().SecurityGroupBuildpack).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		defer os.Remove(buildpackZip)

		cf.AsUser(context.AdminUserContext(), func() {
			Expect(cf.Cf("create-buildpack", buildpack, buildpackZip, "999").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		})
		defer func() {
			cf.AsUser(context.AdminUserContext(), func() {
				Expect(cf.Cf("delete-buildpack", buildpack, "-f").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			})
		}()

		Expect(cf.Cf("push", testAppName, "-b", buildpack, "-p", helpers.NewAssets().HelloWorld, "--no-start").Wait(CF_PUSH_TIMEOUT)).To(Exit(0))
		defer func() { cf.Cf("delete", testAppName, "-f").Wait(CF_PUSH_TIMEOUT) }()

		Expect(cf.Cf("set-env", testAppName, "TESTURI", "www.google.com").Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		Expect(cf.Cf("start", testAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(1))
		Eventually(func() *Session {
			appLogsSession := cf.Cf("logs", "--recent", testAppName)
			Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			return appLogsSession
		}, 5).Should(Say("CURL_EXIT=0"))

		Expect(cf.Cf("set-env", testAppName, "TESTURI", privateUri).Wait(DEFAULT_TIMEOUT)).To(Exit(0))
		Expect(cf.Cf("restart", testAppName).Wait(CF_PUSH_TIMEOUT)).To(Exit(1))
		Eventually(func() *Session {
			appLogsSession := cf.Cf("logs", "--recent", testAppName)
			Expect(appLogsSession.Wait(DEFAULT_TIMEOUT)).To(Exit(0))
			return appLogsSession
		}, 5).Should(Say("CURL_EXIT=[^0]"))
	})
})
<|endoftext|>


<commit_before>package agent

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	mqtt "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/juju/loggo"
)

var AlreadyConfigured = errors.New("Already configured")
var AlreadyUnConfigured = errors.New("Already unconfigured")

//
// Acts as a bridge between local and cloud brokers, this includes reconnecting
// and emitting status changes.
//
// Once configured and started this will attempt to connect to connect to local
// and cloud brokers, if something dies it will reconnect based on the configured
// reconnect backoff.
//
type Bridge struct {
	conf   *Config
	local  *mqtt.MqttClient
	remote *mqtt.MqttClient
	log    loggo.Logger

	localTopics []replaceTopic
	cloudTopics []replaceTopic

	cloudUrl *url.URL
	token    string

	timer       *time.Timer
	reconnectCh chan bool
	shutdownCh  chan bool

	Configured bool
	Connected  bool
	Counter    int64

	IngressCounter int64
	EgressCounter  int64

	IngressBytes int64
	EgressBytes  int64

	LastError error

	bridgeLock sync.Mutex
}

type replaceTopic struct {
	on      string
	replace string
	with    string
}

func (r *replaceTopic) updated(originalTopic string) string {
	return strings.Replace(originalTopic, r.replace, r.with, 1)
}

var localTopics = []replaceTopic{
	// location related topics (TODO: move to cloud userspace RPC)
	{on: "$location/calibration", replace: "$location", with: "$cloud/location"},
	{on: "$location/delete", replace: "$location", with: "$cloud/location"},
	{on: "$device/+/+/rssi", replace: "$device", with: "$cloud/device"},

	// module health statistics
	//{on: "$node/+/module/status", replace: "$node", with: "$cloud/node"},

	// cloud userspace RPC requests
	{on: "$ninja/services/rpc/+/+", replace: "$ninja", with: "$cloud/ninja"},
	{on: "$ninja/services/+", replace: "$ninja", with: "$cloud/ninja"},

	// temporary alternate topic to distinguish remote device replies from local-destined ones
	// used by the phone app for remote actuations
	// the alternate remote_ topic is to prevent a loopback with the below rule in the other direction
	// TODO: use a tag like $mesh-source to prevent loops (never re-proxy msgs with your source)
	{on: "$device/+/channel/+/reply", replace: "$device", with: "$cloud/remote_device"},

	// push up all local RPC methods in case the cloud is responding,
	// this is currently used for the push notification channel
	{on: "$device/+/channel/+", replace: "$device", with: "$cloud/device"},

	// push up state changes to the cloud
	{on: "$device/+/channel/+/event/state", replace: "$device", with: "$cloud/device"},
}

var cloudTopics = []replaceTopic{
	// location related topics
	{on: "$cloud/location/calibration/progress", replace: "$cloud/location", with: "$location"},
	{on: "$cloud/device/+/+/location", replace: "$cloud/device", with: "$device"},

	// cloud userspace RPC replies
	{on: "$cloud/ninja/services/rpc/+/+/reply", replace: "$cloud/ninja", with: "$ninja"},

	// see comment for $device/+/channel/+/reply above
	{on: "$cloud/remote_device/+/channel/+", replace: "$cloud/remote_device", with: "$device"},

	// allow cloud to announce devices and channels (used for phone on 3G and notification subscription channel)
	{on: "$cloud/device/+/event/announce", replace: "$cloud/device", with: "$device"},
	{on: "$cloud/device/+/channel/+/event/announce", replace: "$cloud/device", with: "$device"},

	// retrieve RPC responses from the cloud,
	// this is currently used for the push notification channel
	{on: "$cloud/device/+/channel/+/reply", replace: "$cloud/device", with: "$device"},
}

func createBridge(conf *Config) *Bridge {
	return &Bridge{conf: conf, localTopics: localTopics, cloudTopics: cloudTopics, log: loggo.GetLogger("bridge")}
}

func (b *Bridge) start(cloudUrl string, token string) (err error) {

	if b.Configured {
		b.log.Warningf("Already configured.")
		return AlreadyConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Connecting the bridge")

	b.Configured = true

	url, err := url.Parse(cloudUrl)

	if err != nil {
		return err
	}

	b.cloudUrl = url
	b.token = token

	b.reconnectCh = make(chan bool, 1)
	b.shutdownCh = make(chan bool, 1)

	if err = b.connect(); err != nil {
		b.log.Errorf("Connect failed %s", err)
		b.scheduleReconnect(err)
	}

	go b.mainBridgeLoop()

	return err
}

func (b *Bridge) stop() error {

	if !b.Configured {
		b.log.Warningf("Already unconfigured.")
		return AlreadyUnConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Disconnecting bridge")

	if b.Configured {
		// tell the worker to shutdown
		b.shutdownCh <- true

		b.Configured = false
	}

	b.resetTimer()

	b.disconnectAll()

	return nil
}

func (b *Bridge) connect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true

	return nil
}

func (b *Bridge) reconnect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true
	b.LastError = nil

	return nil
}

func (b *Bridge) subscriptions() (err error) {

	if err = b.subscribe(b.local, b.remote, b.localTopics, "local"); err != nil {
		return err
	}

	if err = b.subscribe(b.remote, b.local, b.cloudTopics, "cloud"); err != nil {
		return err
	}
	return nil

}

func (b *Bridge) disconnectAll() {
	b.log.Infof("disconnectAll")
	// we are now disconnected
	b.Connected = false
	if b.local != nil && b.local.IsConnected() {
		b.local.Disconnect(100)
	}
	if b.remote != nil && b.remote.IsConnected() {
		b.remote.Disconnect(100)
	}
}

func (b *Bridge) mainBridgeLoop() {

	for {
		select {
		case <-b.reconnectCh:
			b.log.Infof("reconnecting")
			if err := b.reconnect(); err != nil {
				b.log.Errorf("Reconnect failed %s", err)
				b.scheduleReconnect(err)
			}
		case <-b.shutdownCh:
			b.log.Infof("shutting down bridge")
			return
		}

	}

}

func (b *Bridge) buildClient(server string, token string) (*mqtt.MqttClient, error) {

	b.log.Infof("building client for %s", server)

	opts := mqtt.NewClientOptions().AddBroker(server).SetTlsConfig(&tls.Config{InsecureSkipVerify: true})

	if token != "" {
		opts.SetUsername(token)
	}

	opts.SetClientId(fmt.Sprintf("%d", time.Now().Unix()))

	opts.SetKeepAlive(15) // set a 15 second ping time for ELB

	// pretty much log the reason and quit
	opts.SetOnConnectionLost(b.onConnectionLoss)

	client := mqtt.NewClient(opts)
	_, err := client.Start()

	return client, err
}

func (b *Bridge) subscribe(src *mqtt.MqttClient, dst *mqtt.MqttClient, topics []replaceTopic, tag string) (err error) {

	var (
		handler mqtt.MessageHandler
	)

	for _, topic := range topics {

		topicFilter, _ := mqtt.NewTopicFilter(topic.on, 0)
		b.log.Infof("(%s) subscribed to %s", tag, topic.on)

		handler = b.buildHandler(topic, tag, b.remote)

		if _, err = src.StartSubscription(handler, topicFilter); err != nil {
			return err
		}

		b.log.Infof("(%s) subscribed to %+v", tag, topicFilter)
	}

	return nil
}

func (b *Bridge) unsubscribe(client *mqtt.MqttClient, topics []replaceTopic, tag string) {

	topicNames := []string{}

	for _, topic := range topics {
		topicNames = append(topicNames, topic.on)
	}

	b.log.Infof("(%s) unsubscribed to %s", tag, topicNames)
	client.EndSubscription(topicNames...)
}

func (b *Bridge) buildHandler(topic replaceTopic, tag string, dst *mqtt.MqttClient) mqtt.MessageHandler {
	return func(src *mqtt.MqttClient, msg mqtt.Message) {
		if b.log.IsDebugEnabled() {
			b.log.Debugf("(%s) topic: %s updated: %s len: %d", tag, msg.Topic(), topic.updated(msg.Topic()), len(msg.Payload()))
		}
		b.Counter++
		b.updateCounters(tag, msg)
		payload := b.updateSource(msg.Payload(), b.buildSource(tag))
		dst.PublishMessage(topic.updated(msg.Topic()), mqtt.NewMessage(payload))
	}
}

func (b *Bridge) scheduleReconnect(reason error) {
	b.LastError = reason
	b.disconnectAll()
	b.resetTimer()

	switch reason {
	case mqtt.ErrBadCredentials:
		b.log.Warningf("Reconnect failed trying again in 30s")

		b.timer = time.AfterFunc(30*time.Second, func() {
			b.reconnectCh <- true
		})

	default:
		b.log.Warningf("Reconnect failed trying again in 5s")
		// TODO add exponential backoff
		b.timer = time.AfterFunc(5*time.Second, func() {
			b.reconnectCh <- true
		})
	}

}

func (b *Bridge) resetTimer() {
	if b.timer != nil {
		b.timer.Stop()
	}
}

func (b *Bridge) onConnectionLoss(client *mqtt.MqttClient, reason error) {
	b.log.Errorf("Connection failed %s", reason)

	// we are now disconnected
	b.Connected = false

	b.scheduleReconnect(reason)

}

func (b *Bridge) IsConnected() bool {
	if b.remote == nil || b.local == nil {
		return false
	}
	return (b.remote.IsConnected() && b.local.IsConnected())
}

func (b *Bridge) buildSource(tag string) string {

	switch tag {
	case "local":
		return b.conf.SerialNo
	case "cloud":
		return "cloud-" + strings.Replace(b.cloudUrl.Host, ".", "_", -1) // encoded to look less wierd
	}

	return ""
}

func (b *Bridge) updateSource(payload []byte, source string) []byte {

	if !bytes.Contains(payload, []byte("$mesh-source")) {
		payload = bytes.Replace(payload, []byte("{"), []byte(`{"$mesh-source":"`+source+`", `), 1)
	}

	b.log.Debugf("msg %s", string(payload))

	return payload
}

func (b *Bridge) updateCounters(tag string, msg mqtt.Message) {
	switch tag {
	case "local":
		b.EgressCounter++
		b.EgressBytes += int64(len(msg.Bytes())) // message size not payload size
	case "cloud":
		b.IngressCounter++
		b.IngressBytes += int64(len(msg.Bytes())) // message size not payload size
	}

}
<commit_msg>Fix for cloud actuations not working.<commit_after>package agent

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	"net/url"
	"strings"
	"sync"
	"time"

	mqtt "git.eclipse.org/gitroot/paho/org.eclipse.paho.mqtt.golang.git"
	"github.com/juju/loggo"
)

var AlreadyConfigured = errors.New("Already configured")
var AlreadyUnConfigured = errors.New("Already unconfigured")

//
// Acts as a bridge between local and cloud brokers, this includes reconnecting
// and emitting status changes.
//
// Once configured and started this will attempt to connect to connect to local
// and cloud brokers, if something dies it will reconnect based on the configured
// reconnect backoff.
//
type Bridge struct {
	conf   *Config
	local  *mqtt.MqttClient
	remote *mqtt.MqttClient
	log    loggo.Logger

	localTopics []replaceTopic
	cloudTopics []replaceTopic

	cloudUrl *url.URL
	token    string

	timer       *time.Timer
	reconnectCh chan bool
	shutdownCh  chan bool

	Configured bool
	Connected  bool
	Counter    int64

	IngressCounter int64
	EgressCounter  int64

	IngressBytes int64
	EgressBytes  int64

	LastError error

	bridgeLock sync.Mutex
}

type replaceTopic struct {
	on      string
	replace string
	with    string
}

func (r *replaceTopic) updated(originalTopic string) string {
	return strings.Replace(originalTopic, r.replace, r.with, 1)
}

var localTopics = []replaceTopic{
	// location related topics (TODO: move to cloud userspace RPC)
	{on: "$location/calibration", replace: "$location", with: "$cloud/location"},
	{on: "$location/delete", replace: "$location", with: "$cloud/location"},
	{on: "$device/+/+/rssi", replace: "$device", with: "$cloud/device"},

	// module health statistics
	//{on: "$node/+/module/status", replace: "$node", with: "$cloud/node"},

	// cloud userspace RPC requests
	{on: "$ninja/services/rpc/+/+", replace: "$ninja", with: "$cloud/ninja"},
	{on: "$ninja/services/+", replace: "$ninja", with: "$cloud/ninja"},

	// temporary alternate topic to distinguish remote device replies from local-destined ones
	// used by the phone app for remote actuations
	// the alternate remote_ topic is to prevent a loopback with the below rule in the other direction
	// TODO: use a tag like $mesh-source to prevent loops (never re-proxy msgs with your source)
	{on: "$device/+/channel/+/reply", replace: "$device", with: "$cloud/remote_device"},

	// push up all local RPC methods in case the cloud is responding,
	// this is currently used for the push notification channel
	{on: "$device/+/channel/+", replace: "$device", with: "$cloud/device"},

	// push up state changes to the cloud
	{on: "$device/+/channel/+/event/state", replace: "$device", with: "$cloud/device"},
}

var cloudTopics = []replaceTopic{
	// location related topics
	{on: "$cloud/location/calibration/progress", replace: "$cloud/location", with: "$location"},
	{on: "$cloud/device/+/+/location", replace: "$cloud/device", with: "$device"},

	// cloud userspace RPC replies
	{on: "$cloud/ninja/services/rpc/+/+/reply", replace: "$cloud/ninja", with: "$ninja"},

	// see comment for $device/+/channel/+/reply above
	{on: "$cloud/remote_device/+/channel/+", replace: "$cloud/remote_device", with: "$device"},

	// allow cloud to announce devices and channels (used for phone on 3G and notification subscription channel)
	{on: "$cloud/device/+/event/announce", replace: "$cloud/device", with: "$device"},
	{on: "$cloud/device/+/channel/+/event/announce", replace: "$cloud/device", with: "$device"},

	// retrieve RPC responses from the cloud,
	// this is currently used for the push notification channel
	{on: "$cloud/device/+/channel/+/reply", replace: "$cloud/device", with: "$device"},
}

func createBridge(conf *Config) *Bridge {
	return &Bridge{conf: conf, localTopics: localTopics, cloudTopics: cloudTopics, log: loggo.GetLogger("bridge")}
}

func (b *Bridge) start(cloudUrl string, token string) (err error) {

	if b.Configured {
		b.log.Warningf("Already configured.")
		return AlreadyConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Connecting the bridge")

	b.Configured = true

	url, err := url.Parse(cloudUrl)

	if err != nil {
		return err
	}

	b.cloudUrl = url
	b.token = token

	b.reconnectCh = make(chan bool, 1)
	b.shutdownCh = make(chan bool, 1)

	if err = b.connect(); err != nil {
		b.log.Errorf("Connect failed %s", err)
		b.scheduleReconnect(err)
	}

	go b.mainBridgeLoop()

	return err
}

func (b *Bridge) stop() error {

	if !b.Configured {
		b.log.Warningf("Already unconfigured.")
		return AlreadyUnConfigured
	}

	defer b.bridgeLock.Unlock()

	b.bridgeLock.Lock()

	b.log.Infof("Disconnecting bridge")

	if b.Configured {
		// tell the worker to shutdown
		b.shutdownCh <- true

		b.Configured = false
	}

	b.resetTimer()

	b.disconnectAll()

	return nil
}

func (b *Bridge) connect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true

	return nil
}

func (b *Bridge) reconnect() (err error) {

	if b.local, err = b.buildClient(b.conf.LocalUrl, ""); err != nil {
		b.Connected = false
		return err
	}

	if b.remote, err = b.buildClient(b.cloudUrl.String(), b.token); err != nil {
		b.Connected = false
		return err
	}

	if err = b.subscriptions(); err != nil {
		return err
	}

	// we are now connected
	b.Connected = true
	b.LastError = nil

	return nil
}

func (b *Bridge) subscriptions() (err error) {

	if err = b.subscribe(b.local, b.remote, b.localTopics, "local"); err != nil {
		return err
	}

	if err = b.subscribe(b.remote, b.local, b.cloudTopics, "cloud"); err != nil {
		return err
	}
	return nil

}

func (b *Bridge) disconnectAll() {
	b.log.Infof("disconnectAll")
	// we are now disconnected
	b.Connected = false
	if b.local != nil && b.local.IsConnected() {
		b.local.Disconnect(100)
	}
	if b.remote != nil && b.remote.IsConnected() {
		b.remote.Disconnect(100)
	}
}

func (b *Bridge) mainBridgeLoop() {

	for {
		select {
		case <-b.reconnectCh:
			b.log.Infof("reconnecting")
			if err := b.reconnect(); err != nil {
				b.log.Errorf("Reconnect failed %s", err)
				b.scheduleReconnect(err)
			}
		case <-b.shutdownCh:
			b.log.Infof("shutting down bridge")
			return
		}

	}

}

func (b *Bridge) buildClient(server string, token string) (*mqtt.MqttClient, error) {

	b.log.Infof("building client for %s", server)

	opts := mqtt.NewClientOptions().AddBroker(server).SetTlsConfig(&tls.Config{InsecureSkipVerify: true})

	if token != "" {
		opts.SetUsername(token)
	}

	opts.SetClientId(fmt.Sprintf("%d", time.Now().Unix()))

	opts.SetKeepAlive(15) // set a 15 second ping time for ELB

	// pretty much log the reason and quit
	opts.SetOnConnectionLost(b.onConnectionLoss)

	client := mqtt.NewClient(opts)
	_, err := client.Start()

	return client, err
}

func (b *Bridge) subscribe(src *mqtt.MqttClient, dst *mqtt.MqttClient, topics []replaceTopic, tag string) (err error) {

	for _, topic := range topics {

		topicFilter, _ := mqtt.NewTopicFilter(topic.on, 0)
		b.log.Infof("(%s) subscribed to %s", tag, topic.on)

		if receipt, err := src.StartSubscription(b.buildHandler(topic, tag, dst), topicFilter); err != nil {
			return err
		} else {
			<-receipt
			b.log.Infof("(%s) subscribed to %+v", tag, topicFilter)
		}
	}

	return nil
}

func (b *Bridge) unsubscribe(client *mqtt.MqttClient, topics []replaceTopic, tag string) {

	topicNames := []string{}

	for _, topic := range topics {
		topicNames = append(topicNames, topic.on)
	}

	b.log.Infof("(%s) unsubscribed to %s", tag, topicNames)
	client.EndSubscription(topicNames...)
}

func (b *Bridge) buildHandler(topic replaceTopic, tag string, dst *mqtt.MqttClient) mqtt.MessageHandler {
	return func(src *mqtt.MqttClient, msg mqtt.Message) {
		if b.log.IsDebugEnabled() {
			b.log.Debugf("(%s) topic: %s updated: %s len: %d", tag, msg.Topic(), topic.updated(msg.Topic()), len(msg.Payload()))
		}
		b.updateCounters(tag, msg)
		payload := b.updateSource(msg.Payload(), b.buildSource(tag))
		dst.PublishMessage(topic.updated(msg.Topic()), mqtt.NewMessage(payload))
	}
}

func (b *Bridge) scheduleReconnect(reason error) {
	b.LastError = reason
	b.disconnectAll()
	b.resetTimer()

	switch reason {
	case mqtt.ErrBadCredentials:
		b.log.Warningf("Reconnect failed trying again in 30s")

		b.timer = time.AfterFunc(30*time.Second, func() {
			b.reconnectCh <- true
		})

	default:
		b.log.Warningf("Reconnect failed trying again in 5s")
		// TODO add exponential backoff
		b.timer = time.AfterFunc(5*time.Second, func() {
			b.reconnectCh <- true
		})
	}

}

func (b *Bridge) resetTimer() {
	if b.timer != nil {
		b.timer.Stop()
	}
}

func (b *Bridge) onConnectionLoss(client *mqtt.MqttClient, reason error) {
	b.log.Errorf("Connection failed %s", reason)

	// we are now disconnected
	b.Connected = false

	b.scheduleReconnect(reason)

}

func (b *Bridge) IsConnected() bool {
	if b.remote == nil || b.local == nil {
		return false
	}
	return (b.remote.IsConnected() && b.local.IsConnected())
}

func (b *Bridge) buildSource(tag string) string {

	switch tag {
	case "local":
		return b.conf.SerialNo
	case "cloud":
		return "cloud-" + strings.Replace(b.cloudUrl.Host, ".", "_", -1) // encoded to look less wierd
	}

	return ""
}

func (b *Bridge) updateSource(payload []byte, source string) []byte {

	if !bytes.Contains(payload, []byte("$mesh-source")) {
		payload = bytes.Replace(payload, []byte("{"), []byte(`{"$mesh-source":"`+source+`", `), 1)
	}

	b.log.Debugf("msg %s", string(payload))

	return payload
}

func (b *Bridge) updateCounters(tag string, msg mqtt.Message) {
	switch tag {
	case "local":
		b.EgressCounter++
		b.EgressBytes += int64(len(msg.Bytes())) // message size not payload size
	case "cloud":
		b.IngressCounter++
		b.IngressBytes += int64(len(msg.Bytes())) // message size not payload size
	}

}
<|endoftext|>