<commit_before>package command

import (
	"strings"
	"sync"

	"github.com/igungor/tlbot"
)

// A Command is an implementation of a bot command.
type Command struct {
	// Name of the command without the leading slash.
	Name string

	// Short description of the command
	ShortLine string

	// Hidden enables commands to be hidden from the /help output, such as
	// Telegram's built-in commands and easter eggs.
	Hidden bool

	// Run runs the command.
	Run func(bot *tlbot.Bot, msg *tlbot.Message)
}

const autocorrect = true

var (
	mu       sync.Mutex
	commands = make(map[string]*Command)
)

func register(cmd *Command) {
	mu.Lock()
	defer mu.Unlock()

	commands[cmd.Name] = cmd
}

// Lookup looks-up name from registered commands and returns
// corresponding Command if any.
func Lookup(cmdname string) *Command {
	mu.Lock()
	defer mu.Unlock()

	cmdname = strings.TrimSuffix(cmdname, "@ilberbot")
	cmd, ok := commands[cmdname]
	if ok {
		return cmd
	}

	if !autocorrect {
		return nil
	}

	//
	// we don't have an exact match. try to guess the input.
	//

	// don't even bother on single letter command inputs
	if len(cmdname) < 2 {
		return nil
	}
	// autocorrect based on levenshtein distance, if possible
	for k := range commands {
		if distance(cmdname, k) <= 2 {
			return commands[k]
		}
	}

	// at least we tried.
	return nil
}

// distance returns the levenshtein distance between given strings.
func distance(s1, s2 string) int {
	var cost, lastdiag, olddiag int
	s1len := len([]rune(s1))
	s2len := len([]rune(s2))

	col := make([]int, s1len+1)

	for y := 1; y <= s1len; y++ {
		col[y] = y
	}

	for x := 1; x <= s2len; x++ {
		col[0] = x
		lastdiag = x - 1
		for y := 1; y <= s1len; y++ {
			olddiag = col[y]
			cost = 0
			if s1[y-1] != s2[x-1] {
				cost = 1
			}
			col[y] = min(col[y]+1, min(col[y-1]+1, lastdiag+cost))
			lastdiag = olddiag
		}
	}
	return col[s1len]
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
<commit_msg>disable autocorrect<commit_after>package command

import (
	"strings"
	"sync"

	"github.com/igungor/tlbot"
)

// A Command is an implementation of a bot command.
type Command struct {
	// Name of the command without the leading slash.
	Name string

	// Short description of the command
	ShortLine string

	// Hidden enables commands to be hidden from the /help output, such as
	// Telegram's built-in commands and easter eggs.
	Hidden bool

	// Run runs the command.
	Run func(bot *tlbot.Bot, msg *tlbot.Message)
}

const autocorrect = false

var (
	mu       sync.Mutex
	commands = make(map[string]*Command)
)

func register(cmd *Command) {
	mu.Lock()
	defer mu.Unlock()

	commands[cmd.Name] = cmd
}

// Lookup looks-up name from registered commands and returns
// corresponding Command if any.
func Lookup(cmdname string) *Command {
	mu.Lock()
	defer mu.Unlock()

	cmdname = strings.TrimSuffix(cmdname, "@ilberbot")
	cmd, ok := commands[cmdname]
	if ok {
		return cmd
	}

	if !autocorrect {
		return nil
	}

	//
	// we don't have an exact match. try to guess the input.
	//

	// don't even bother on single letter command inputs
	if len(cmdname) < 2 {
		return nil
	}
	// autocorrect based on levenshtein distance, if possible
	for k := range commands {
		if distance(cmdname, k) <= 2 {
			return commands[k]
		}
	}

	// at least we tried.
	return nil
}

// distance returns the levenshtein distance between given strings.
func distance(s1, s2 string) int {
	var cost, lastdiag, olddiag int
	s1len := len([]rune(s1))
	s2len := len([]rune(s2))

	col := make([]int, s1len+1)

	for y := 1; y <= s1len; y++ {
		col[y] = y
	}

	for x := 1; x <= s2len; x++ {
		col[0] = x
		lastdiag = x - 1
		for y := 1; y <= s1len; y++ {
			olddiag = col[y]
			cost = 0
			if s1[y-1] != s2[x-1] {
				cost = 1
			}
			col[y] = min(col[y]+1, min(col[y-1]+1, lastdiag+cost))
			lastdiag = olddiag
		}
	}
	return col[s1len]
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
<|endoftext|>


<commit_before>// +build !windows

package commands

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/codegangsta/cli"
	"github.com/concourse/atc"
	"github.com/kr/pty"
	"github.com/mgutz/ansi"
	"github.com/pkg/term"
	"github.com/tedsuo/rata"
)

func remoteCommand(argv []string) (string, []string) {
	var path string
	var args []string

	switch len(argv) {
	case 0:
		path = "bash"
	case 1:
		path = argv[0]
	default:
		path = argv[0]
		args = argv[1:]
	}

	return path, args
}

type containerLocator interface {
	locate(containerFingerprint) url.Values
}

type stepContainerLocator struct {
	client       *http.Client
	reqGenerator *rata.RequestGenerator
}

func (locator stepContainerLocator) locate(fingerprint containerFingerprint) url.Values {
	build := getBuild(
		locator.client,
		locator.reqGenerator,
		fingerprint.jobName,
		fingerprint.buildName,
		fingerprint.pipelineName,
	)

	reqValues := url.Values{}
	reqValues["build-id"] = []string{strconv.Itoa(build.ID)}
	reqValues["name"] = []string{fingerprint.stepName}

	if fingerprint.stepType != "" {
		reqValues["type"] = []string{fingerprint.stepType}
	}

	return reqValues
}

type checkContainerLocator struct{}

func (locator checkContainerLocator) locate(fingerprint containerFingerprint) url.Values {
	reqValues := url.Values{}

	reqValues["type"] = []string{"check"}
	if fingerprint.checkName != "" {
		reqValues["name"] = []string{fingerprint.checkName}
	}
	if fingerprint.pipelineName != "" {
		reqValues["pipeline"] = []string{fingerprint.pipelineName}
	}

	return reqValues
}

type containerFingerprint struct {
	pipelineName string
	jobName      string
	buildName    string

	stepName string
	stepType string

	checkName string
}

func locateContainer(client *http.Client, reqGenerator *rata.RequestGenerator, fingerprint containerFingerprint) url.Values {
	var locator containerLocator

	if fingerprint.checkName == "" && fingerprint.stepType != "check" {
		locator = stepContainerLocator{
			client:       client,
			reqGenerator: reqGenerator,
		}
	} else {
		locator = checkContainerLocator{}
	}

	return locator.locate(fingerprint)
}

func constructRequest(reqGenerator *rata.RequestGenerator, spec atc.HijackProcessSpec, id string) *http.Request {
	payload, err := json.Marshal(spec)
	if err != nil {
		log.Fatalln("failed to marshal process spec:", err)
	}

	hijackReq, err := reqGenerator.CreateRequest(
		atc.HijackContainer,
		rata.Params{"id": id},
		bytes.NewBuffer(payload),
	)
	if err != nil {
		log.Fatalln("failed to create hijack request:", err)
	}

	if hijackReq.URL.User != nil {
		hijackReq.Header.Add("Authorization", basicAuth(hijackReq.URL.User))
		hijackReq.URL.User = nil
	}

	return hijackReq
}

func getContainerIDs(c *cli.Context) []atc.PresentedContainer {
	target := returnTarget(c.GlobalString("target"))
	insecure := c.GlobalBool("insecure")

	pipelineName := c.String("pipeline")
	jobName := c.String("job")
	buildName := c.String("build")
	stepName := c.String("step-name")
	stepType := c.String("step-type")
	check := c.String("check")

	fingerprint := containerFingerprint{
		pipelineName: pipelineName,
		jobName:      jobName,
		buildName:    buildName,
		stepName:     stepName,
		stepType:     stepType,
		checkName:    check,
	}

	reqGenerator := rata.NewRequestGenerator(target, atc.Routes)
	tlsConfig := &tls.Config{InsecureSkipVerify: insecure}
	transport := &http.Transport{
		TLSClientConfig: tlsConfig,
	}
	client := &http.Client{Transport: transport}
	reqValues := locateContainer(client, reqGenerator, fingerprint)

	listContainersReq, err := reqGenerator.CreateRequest(
		atc.ListContainers,
		rata.Params{},
		nil,
	)
	if err != nil {
		log.Fatalln("failed to create containers list request:", err)
	}
	listContainersReq.URL.RawQuery = reqValues.Encode()

	resp, err := client.Do(listContainersReq)
	if err != nil {
		log.Fatalln("failed to get containers:", err)
	}

	var listContainers atc.ListContainersReturn
	err = json.NewDecoder(resp.Body).Decode(&listContainers)
	if err != nil {
		log.Fatalln("failed to decode containers:", err)
	}

	return listContainers.Containers
}

func Hijack(c *cli.Context) {
	target := returnTarget(c.GlobalString("target"))
	insecure := c.GlobalBool("insecure")

	containers := getContainerIDs(c)

	var id string
	var selection int
	if len(containers) == 0 {
		log.Fatalln("no containers matched your search parameters! they may have expired if your build hasn't recently finished")
	} else if len(containers) > 1 {

		for i, container := range containers {
			fmt.Printf("%d. pipeline:%s, type:%s, name:%s, build_id:%d\n", i+1, container.PipelineName, container.Type, container.Name, container.BuildID)
		}

		fmt.Printf("Choose a container: ")
		_, err := fmt.Scanf("%d", &selection)
		for err != nil || selection > len(containers) || selection < 1 {
			fmt.Println("invalid selection")
			fmt.Printf("Choose a container: ")
			_, err = fmt.Scanf("%d", &selection)
		}

		id = containers[selection-1].ID
	} else {
		id = containers[0].ID
	}

	path, args := remoteCommand(c.Args())
	privileged := true

	reqGenerator := rata.NewRequestGenerator(target, atc.Routes)
	tlsConfig := &tls.Config{InsecureSkipVerify: insecure}

	var ttySpec *atc.HijackTTYSpec
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		ttySpec = &atc.HijackTTYSpec{
			WindowSize: atc.HijackWindowSize{
				Columns: cols,
				Rows:    rows,
			},
		}
	}

	spec := atc.HijackProcessSpec{
		Path: path,
		Args: args,
		Env:  []string{"TERM=" + os.Getenv("TERM")},
		User: "root",

		Privileged: privileged,
		TTY:        ttySpec,
	}

	hijackReq := constructRequest(reqGenerator, spec, id)
	hijackResult := performHijack(hijackReq, tlsConfig)
	os.Exit(hijackResult)
}

func performHijack(hijackReq *http.Request, tlsConfig *tls.Config) int {
	conn, err := dialEndpoint(hijackReq.URL, tlsConfig)
	if err != nil {
		log.Fatalln("failed to dial hijack endpoint:", err)
	}

	clientConn := httputil.NewClientConn(conn, nil)

	resp, err := clientConn.Do(hijackReq)
	if err != nil {
		log.Fatalln("failed to hijack:", err)
	}

	if resp.StatusCode != http.StatusOK {
		handleBadResponse("hijacking", resp)
	}

	return hijack(clientConn.Hijack())
}

func hijack(conn net.Conn, br *bufio.Reader) int {
	var in io.Reader

	term, err := term.Open(os.Stdin.Name())
	if err == nil {
		err = term.SetRaw()
		if err != nil {
			log.Fatalln("failed to set raw:", term)
		}

		defer term.Restore()

		in = term
	} else {
		in = os.Stdin
	}

	encoder := json.NewEncoder(conn)
	decoder := json.NewDecoder(br)

	resized := make(chan os.Signal, 10)
	signal.Notify(resized, syscall.SIGWINCH)

	go func() {
		for {
			<-resized
			// TODO json race
			sendSize(encoder)
		}
	}()

	go io.Copy(&stdinWriter{encoder}, in)

	var exitStatus int
	for {
		var output atc.HijackOutput
		err := decoder.Decode(&output)
		if err != nil {
			break
		}

		if output.ExitStatus != nil {
			exitStatus = *output.ExitStatus
		} else if len(output.Error) > 0 {
			fmt.Fprintf(os.Stderr, "%s\n", ansi.Color(output.Error, "red+b"))
			exitStatus = 255
		} else if len(output.Stdout) > 0 {
			os.Stdout.Write(output.Stdout)
		} else if len(output.Stderr) > 0 {
			os.Stderr.Write(output.Stderr)
		}
	}

	return exitStatus
}

func sendSize(enc *json.Encoder) {
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		enc.Encode(atc.HijackInput{
			TTYSpec: &atc.HijackTTYSpec{
				WindowSize: atc.HijackWindowSize{
					Columns: cols,
					Rows:    rows,
				},
			},
		})
	}
}

type stdinWriter struct {
	enc *json.Encoder
}

func (w *stdinWriter) Write(d []byte) (int, error) {
	err := w.enc.Encode(atc.HijackInput{
		Stdin: d,
	})
	if err != nil {
		return 0, err
	}

	return len(d), nil
}

func basicAuth(user *url.Userinfo) string {
	username := user.Username()
	password, _ := user.Password()
	return "Basic " + base64.StdEncoding.EncodeToString([]byte(username+":"+password))
}

var canonicalPortMap = map[string]string{
	"http":  "80",
	"https": "443",
}

func dialEndpoint(url *url.URL, tlsConfig *tls.Config) (net.Conn, error) {
	addr := canonicalAddr(url)

	if url.Scheme == "https" {
		return tls.Dial("tcp", addr, tlsConfig)
	} else {
		return net.Dial("tcp", addr)
	}
}

func canonicalAddr(url *url.URL) string {
	host, port, err := net.SplitHostPort(url.Host)
	if err != nil {
		if strings.Contains(err.Error(), "missing port in address") {
			host = url.Host
			port = canonicalPortMap[url.Scheme]
		} else {
			log.Fatalln("invalid host:", err)
		}
	}

	return net.JoinHostPort(host, port)
}
<commit_msg>use newAtcRequester because it's there<commit_after>// +build !windows

package commands

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/codegangsta/cli"
	"github.com/concourse/atc"
	"github.com/kr/pty"
	"github.com/mgutz/ansi"
	"github.com/pkg/term"
	"github.com/tedsuo/rata"
)

func remoteCommand(argv []string) (string, []string) {
	var path string
	var args []string

	switch len(argv) {
	case 0:
		path = "bash"
	case 1:
		path = argv[0]
	default:
		path = argv[0]
		args = argv[1:]
	}

	return path, args
}

type containerLocator interface {
	locate(containerFingerprint) url.Values
}

type stepContainerLocator struct {
	client       *http.Client
	reqGenerator *rata.RequestGenerator
}

func (locator stepContainerLocator) locate(fingerprint containerFingerprint) url.Values {
	build := getBuild(
		locator.client,
		locator.reqGenerator,
		fingerprint.jobName,
		fingerprint.buildName,
		fingerprint.pipelineName,
	)

	reqValues := url.Values{}
	reqValues["build-id"] = []string{strconv.Itoa(build.ID)}
	reqValues["name"] = []string{fingerprint.stepName}

	if fingerprint.stepType != "" {
		reqValues["type"] = []string{fingerprint.stepType}
	}

	return reqValues
}

type checkContainerLocator struct{}

func (locator checkContainerLocator) locate(fingerprint containerFingerprint) url.Values {
	reqValues := url.Values{}

	reqValues["type"] = []string{"check"}
	if fingerprint.checkName != "" {
		reqValues["name"] = []string{fingerprint.checkName}
	}
	if fingerprint.pipelineName != "" {
		reqValues["pipeline"] = []string{fingerprint.pipelineName}
	}

	return reqValues
}

type containerFingerprint struct {
	pipelineName string
	jobName      string
	buildName    string

	stepName string
	stepType string

	checkName string
}

func locateContainer(client *http.Client, reqGenerator *rata.RequestGenerator, fingerprint containerFingerprint) url.Values {
	var locator containerLocator

	if fingerprint.checkName == "" && fingerprint.stepType != "check" {
		locator = stepContainerLocator{
			client:       client,
			reqGenerator: reqGenerator,
		}
	} else {
		locator = checkContainerLocator{}
	}

	return locator.locate(fingerprint)
}

func constructRequest(reqGenerator *rata.RequestGenerator, spec atc.HijackProcessSpec, id string) *http.Request {
	payload, err := json.Marshal(spec)
	if err != nil {
		log.Fatalln("failed to marshal process spec:", err)
	}

	hijackReq, err := reqGenerator.CreateRequest(
		atc.HijackContainer,
		rata.Params{"id": id},
		bytes.NewBuffer(payload),
	)
	if err != nil {
		log.Fatalln("failed to create hijack request:", err)
	}

	if hijackReq.URL.User != nil {
		hijackReq.Header.Add("Authorization", basicAuth(hijackReq.URL.User))
		hijackReq.URL.User = nil
	}

	return hijackReq
}

func getContainerIDs(c *cli.Context) []atc.PresentedContainer {
	target := returnTarget(c.GlobalString("target"))
	insecure := c.GlobalBool("insecure")

	pipelineName := c.String("pipeline")
	jobName := c.String("job")
	buildName := c.String("build")
	stepName := c.String("step-name")
	stepType := c.String("step-type")
	check := c.String("check")

	fingerprint := containerFingerprint{
		pipelineName: pipelineName,
		jobName:      jobName,
		buildName:    buildName,
		stepName:     stepName,
		stepType:     stepType,
		checkName:    check,
	}

	atcRequester := newAtcRequester(target, insecure)
	reqValues := locateContainer(atcRequester.httpClient, atcRequester.RequestGenerator, fingerprint)

	listContainersReq, err := atcRequester.RequestGenerator.CreateRequest(
		atc.ListContainers,
		rata.Params{},
		nil,
	)
	if err != nil {
		log.Fatalln("failed to create containers list request:", err)
	}
	listContainersReq.URL.RawQuery = reqValues.Encode()

	resp, err := atcRequester.httpClient.Do(listContainersReq)
	if err != nil {
		log.Fatalln("failed to get containers:", err)
	}

	var listContainers atc.ListContainersReturn
	err = json.NewDecoder(resp.Body).Decode(&listContainers)
	if err != nil {
		log.Fatalln("failed to decode containers:", err)
	}

	return listContainers.Containers
}

func Hijack(c *cli.Context) {
	target := returnTarget(c.GlobalString("target"))
	insecure := c.GlobalBool("insecure")

	containers := getContainerIDs(c)

	var id string
	var selection int
	if len(containers) == 0 {
		log.Fatalln("no containers matched your search parameters! they may have expired if your build hasn't recently finished")
	} else if len(containers) > 1 {

		for i, container := range containers {
			fmt.Printf("%d. pipeline:%s, type:%s, name:%s, build_id:%d\n", i+1, container.PipelineName, container.Type, container.Name, container.BuildID)
		}

		fmt.Printf("Choose a container: ")
		_, err := fmt.Scanf("%d", &selection)
		for err != nil || selection > len(containers) || selection < 1 {
			fmt.Println("invalid selection")
			fmt.Printf("Choose a container: ")
			_, err = fmt.Scanf("%d", &selection)
		}

		id = containers[selection-1].ID
	} else {
		id = containers[0].ID
	}

	path, args := remoteCommand(c.Args())
	privileged := true

	reqGenerator := rata.NewRequestGenerator(target, atc.Routes)
	tlsConfig := &tls.Config{InsecureSkipVerify: insecure}

	var ttySpec *atc.HijackTTYSpec
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		ttySpec = &atc.HijackTTYSpec{
			WindowSize: atc.HijackWindowSize{
				Columns: cols,
				Rows:    rows,
			},
		}
	}

	spec := atc.HijackProcessSpec{
		Path: path,
		Args: args,
		Env:  []string{"TERM=" + os.Getenv("TERM")},
		User: "root",

		Privileged: privileged,
		TTY:        ttySpec,
	}

	hijackReq := constructRequest(reqGenerator, spec, id)
	hijackResult := performHijack(hijackReq, tlsConfig)
	os.Exit(hijackResult)
}

func performHijack(hijackReq *http.Request, tlsConfig *tls.Config) int {
	conn, err := dialEndpoint(hijackReq.URL, tlsConfig)
	if err != nil {
		log.Fatalln("failed to dial hijack endpoint:", err)
	}

	clientConn := httputil.NewClientConn(conn, nil)

	resp, err := clientConn.Do(hijackReq)
	if err != nil {
		log.Fatalln("failed to hijack:", err)
	}

	if resp.StatusCode != http.StatusOK {
		handleBadResponse("hijacking", resp)
	}

	return hijack(clientConn.Hijack())
}

func hijack(conn net.Conn, br *bufio.Reader) int {
	var in io.Reader

	term, err := term.Open(os.Stdin.Name())
	if err == nil {
		err = term.SetRaw()
		if err != nil {
			log.Fatalln("failed to set raw:", term)
		}

		defer term.Restore()

		in = term
	} else {
		in = os.Stdin
	}

	encoder := json.NewEncoder(conn)
	decoder := json.NewDecoder(br)

	resized := make(chan os.Signal, 10)
	signal.Notify(resized, syscall.SIGWINCH)

	go func() {
		for {
			<-resized
			// TODO json race
			sendSize(encoder)
		}
	}()

	go io.Copy(&stdinWriter{encoder}, in)

	var exitStatus int
	for {
		var output atc.HijackOutput
		err := decoder.Decode(&output)
		if err != nil {
			break
		}

		if output.ExitStatus != nil {
			exitStatus = *output.ExitStatus
		} else if len(output.Error) > 0 {
			fmt.Fprintf(os.Stderr, "%s\n", ansi.Color(output.Error, "red+b"))
			exitStatus = 255
		} else if len(output.Stdout) > 0 {
			os.Stdout.Write(output.Stdout)
		} else if len(output.Stderr) > 0 {
			os.Stderr.Write(output.Stderr)
		}
	}

	return exitStatus
}

func sendSize(enc *json.Encoder) {
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		enc.Encode(atc.HijackInput{
			TTYSpec: &atc.HijackTTYSpec{
				WindowSize: atc.HijackWindowSize{
					Columns: cols,
					Rows:    rows,
				},
			},
		})
	}
}

type stdinWriter struct {
	enc *json.Encoder
}

func (w *stdinWriter) Write(d []byte) (int, error) {
	err := w.enc.Encode(atc.HijackInput{
		Stdin: d,
	})
	if err != nil {
		return 0, err
	}

	return len(d), nil
}

func basicAuth(user *url.Userinfo) string {
	username := user.Username()
	password, _ := user.Password()
	return "Basic " + base64.StdEncoding.EncodeToString([]byte(username+":"+password))
}

var canonicalPortMap = map[string]string{
	"http":  "80",
	"https": "443",
}

func dialEndpoint(url *url.URL, tlsConfig *tls.Config) (net.Conn, error) {
	addr := canonicalAddr(url)

	if url.Scheme == "https" {
		return tls.Dial("tcp", addr, tlsConfig)
	} else {
		return net.Dial("tcp", addr)
	}
}

func canonicalAddr(url *url.URL) string {
	host, port, err := net.SplitHostPort(url.Host)
	if err != nil {
		if strings.Contains(err.Error(), "missing port in address") {
			host = url.Host
			port = canonicalPortMap[url.Scheme]
		} else {
			log.Fatalln("invalid host:", err)
		}
	}

	return net.JoinHostPort(host, port)
}
<|endoftext|>


<commit_before>package commands

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/concourse/atc"
	"github.com/concourse/fly/commands/internal/flaghelpers"
	"github.com/concourse/fly/pty"
	"github.com/concourse/fly/rc"
	"github.com/concourse/go-concourse/concourse"
	"github.com/mgutz/ansi"
	"github.com/tedsuo/rata"
	"github.com/vito/go-interact/interact"
)

type HijackCommand struct {
	Job      flaghelpers.JobFlag      `short:"j" long:"job"   value-name:"PIPELINE/JOB"   description:"Name of a job to hijack"`
	Check    flaghelpers.ResourceFlag `short:"c" long:"check" value-name:"PIPELINE/CHECK" description:"Name of a resource's checking container to hijack"`
	Build    string                   `short:"b" long:"build"                               description:"Build number within the job, or global build ID"`
	StepName string                   `short:"s" long:"step"                                description:"Name of step to hijack (e.g. build, unit, resource name)"`
}

func remoteCommand(argv []string) (string, []string) {
	var path string
	var args []string

	switch len(argv) {
	case 0:
		path = "bash"
	case 1:
		path = argv[0]
	default:
		path = argv[0]
		args = argv[1:]
	}

	return path, args
}

type containerLocator interface {
	locate(containerFingerprint) (map[string]string, error)
}

type stepContainerLocator struct {
	client concourse.Client
}

func (locator stepContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	build, err := GetBuild(
		locator.client,
		fingerprint.jobName,
		fingerprint.buildName,
		fingerprint.pipelineName,
	)
	if err != nil {
		return reqValues, err
	}

	reqValues["build-id"] = strconv.Itoa(build.ID)
	reqValues["name"] = fingerprint.stepName

	return reqValues, nil
}

type checkContainerLocator struct{}

func (locator checkContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	reqValues["type"] = "check"
	if fingerprint.checkName != "" {
		reqValues["name"] = fingerprint.checkName
	}
	if fingerprint.pipelineName != "" {
		reqValues["pipeline_name"] = fingerprint.pipelineName
	}

	return reqValues, nil
}

type containerFingerprint struct {
	pipelineName string
	jobName      string
	buildName    string

	stepName string

	checkName string
}

func locateContainer(client concourse.Client, fingerprint containerFingerprint) (map[string]string, error) {
	var locator containerLocator

	if fingerprint.checkName == "" {
		locator = stepContainerLocator{
			client: client,
		}
	} else {
		locator = checkContainerLocator{}
	}

	return locator.locate(fingerprint)
}

func constructRequest(reqGenerator *rata.RequestGenerator, spec atc.HijackProcessSpec, id string, token *rc.TargetToken) *http.Request {
	payload, err := json.Marshal(spec)
	if err != nil {
		log.Fatalln("failed to marshal process spec:", err)
	}

	hijackReq, err := reqGenerator.CreateRequest(
		atc.HijackContainer,
		rata.Params{"id": id},
		bytes.NewBuffer(payload),
	)
	if err != nil {
		log.Fatalln("failed to create hijack request:", err)
	}

	if token != nil {
		hijackReq.Header.Add("Authorization", token.Type+" "+token.Value)
	}

	return hijackReq
}

func getContainerIDs(c *HijackCommand) []atc.Container {
	var pipelineName string
	if c.Job.PipelineName != "" {
		pipelineName = c.Job.PipelineName
	} else {
		pipelineName = c.Check.PipelineName
	}

	buildName := c.Build
	stepName := c.StepName
	jobName := c.Job.JobName
	check := c.Check.ResourceName

	fingerprint := containerFingerprint{
		pipelineName: pipelineName,
		jobName:      jobName,
		buildName:    buildName,
		stepName:     stepName,
		checkName:    check,
	}

	connection, err := rc.TargetConnection(Fly.Target)
	if err != nil {
		log.Fatalln("failed to create client:", err)
	}
	client := concourse.NewClient(connection)

	reqValues, err := locateContainer(client, fingerprint)
	if err != nil {
		log.Fatalln(err)
	}

	containers, err := client.ListContainers(reqValues)
	if err != nil {
		log.Fatalln("failed to get containers:", err)
	}
	return containers
}

func (command *HijackCommand) Execute(args []string) error {
	target, err := rc.SelectTarget(Fly.Target)
	if err != nil {
		log.Fatalln(err)
		return nil
	}

	containers := getContainerIDs(command)

	var chosenContainer atc.Container
	if len(containers) == 0 {
		fmt.Fprintln(os.Stderr, "no containers matched your search parameters! they may have expired if your build hasn't recently finished")
		os.Exit(1)
	} else if len(containers) > 1 {
		var choices []interact.Choice
		for _, container := range containers {
			var infos []string
			if container.PipelineName != "" {
				infos = append(infos, fmt.Sprintf("pipeline: %s", container.PipelineName))
			}

			if container.BuildID != 0 {
				infos = append(infos, fmt.Sprintf("build id: %d", container.BuildID))
			}

			infos = append(infos, fmt.Sprintf("type: %s", container.Type))
			infos = append(infos, fmt.Sprintf("name: %s", container.Name))

			choices = append(choices, interact.Choice{
				Display: strings.Join(infos, ", "),
				Value:   container,
			})
		}

		err = interact.NewInteraction("choose a container", choices...).Resolve(&chosenContainer)
		if err == io.EOF {
			os.Exit(0)
		}

		if err != nil {
			return err
		}
	} else {
		chosenContainer = containers[0]
	}

	path, args := remoteCommand(args)
	privileged := true

	reqGenerator := rata.NewRequestGenerator(target.API, atc.Routes)
	tlsConfig := &tls.Config{InsecureSkipVerify: target.Insecure}

	var ttySpec *atc.HijackTTYSpec
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		ttySpec = &atc.HijackTTYSpec{
			WindowSize: atc.HijackWindowSize{
				Columns: cols,
				Rows:    rows,
			},
		}
	}

	envVariables := append(chosenContainer.EnvironmentVariables, "TERM="+os.Getenv("TERM"))

	spec := atc.HijackProcessSpec{
		Path: path,
		Args: args,
		Env:  envVariables,
		User: "root",
		Dir:  chosenContainer.WorkingDirectory,

		Privileged: privileged,
		TTY:        ttySpec,
	}

	hijackReq := constructRequest(reqGenerator, spec, chosenContainer.ID, target.Token)
	hijackResult := performHijack(hijackReq, tlsConfig)
	os.Exit(hijackResult)

	return nil
}

func performHijack(hijackReq *http.Request, tlsConfig *tls.Config) int {
	conn, err := dialEndpoint(hijackReq.URL, tlsConfig)
	if err != nil {
		log.Fatalln("failed to dial hijack endpoint:", err)
	}

	clientConn := httputil.NewClientConn(conn, nil)

	resp, err := clientConn.Do(hijackReq)
	if err != nil {
		log.Fatalln("failed to hijack:", err)
	}

	if resp.StatusCode != http.StatusOK {
		handleBadResponse("hijacking", resp)
	}

	return hijack(clientConn.Hijack())
}

func hijack(conn net.Conn, br *bufio.Reader) int {
	var in io.Reader

	term, err := pty.OpenRawTerm()
	if err == nil {
		defer term.Restore()

		in = term
	} else {
		in = os.Stdin
	}

	encoder := json.NewEncoder(conn)
	decoder := json.NewDecoder(br)

	resized := pty.ResizeNotifier()

	go func() {
		for {
			<-resized
			// TODO json race
			sendSize(encoder)
		}
	}()

	go io.Copy(&stdinWriter{encoder}, in)

	var exitStatus int
	for {
		var output atc.HijackOutput
		err := decoder.Decode(&output)
		if err != nil {
			break
		}

		if output.ExitStatus != nil {
			exitStatus = *output.ExitStatus
		} else if len(output.Error) > 0 {
			fmt.Fprintf(os.Stderr, "%s\n", ansi.Color(output.Error, "red+b"))
			exitStatus = 255
		} else if len(output.Stdout) > 0 {
			os.Stdout.Write(output.Stdout)
		} else if len(output.Stderr) > 0 {
			os.Stderr.Write(output.Stderr)
		}
	}

	return exitStatus
}

func sendSize(enc *json.Encoder) {
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		enc.Encode(atc.HijackInput{
			TTYSpec: &atc.HijackTTYSpec{
				WindowSize: atc.HijackWindowSize{
					Columns: cols,
					Rows:    rows,
				},
			},
		})
	}
}

type stdinWriter struct {
	enc *json.Encoder
}

func (w *stdinWriter) Write(d []byte) (int, error) {
	err := w.enc.Encode(atc.HijackInput{
		Stdin: d,
	})
	if err != nil {
		return 0, err
	}

	return len(d), nil
}

var canonicalPortMap = map[string]string{
	"http":  "80",
	"https": "443",
}

func dialEndpoint(url *url.URL, tlsConfig *tls.Config) (net.Conn, error) {
	addr := canonicalAddr(url)

	if url.Scheme == "https" {
		return tls.Dial("tcp", addr, tlsConfig)
	}

	return net.Dial("tcp", addr)
}

func canonicalAddr(url *url.URL) string {
	host, port, err := net.SplitHostPort(url.Host)
	if err != nil {
		if strings.Contains(err.Error(), "missing port in address") {
			host = url.Host
			port = canonicalPortMap[url.Scheme]
		} else {
			log.Fatalln("invalid host:", err)
		}
	}

	return net.JoinHostPort(host, port)
}
<commit_msg>BuildName->BuildNameOrID<commit_after>package commands

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/concourse/atc"
	"github.com/concourse/fly/commands/internal/flaghelpers"
	"github.com/concourse/fly/pty"
	"github.com/concourse/fly/rc"
	"github.com/concourse/go-concourse/concourse"
	"github.com/mgutz/ansi"
	"github.com/tedsuo/rata"
	"github.com/vito/go-interact/interact"
)

type HijackCommand struct {
	Job      flaghelpers.JobFlag      `short:"j" long:"job"   value-name:"PIPELINE/JOB"   description:"Name of a job to hijack"`
	Check    flaghelpers.ResourceFlag `short:"c" long:"check" value-name:"PIPELINE/CHECK" description:"Name of a resource's checking container to hijack"`
	Build    string                   `short:"b" long:"build"                               description:"Build number within the job, or global build ID"`
	StepName string                   `short:"s" long:"step"                                description:"Name of step to hijack (e.g. build, unit, resource name)"`
}

func remoteCommand(argv []string) (string, []string) {
	var path string
	var args []string

	switch len(argv) {
	case 0:
		path = "bash"
	case 1:
		path = argv[0]
	default:
		path = argv[0]
		args = argv[1:]
	}

	return path, args
}

type containerLocator interface {
	locate(containerFingerprint) (map[string]string, error)
}

type stepContainerLocator struct {
	client concourse.Client
}

func (locator stepContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	build, err := GetBuild(
		locator.client,
		fingerprint.jobName,
		fingerprint.buildNameOrID,
		fingerprint.pipelineName,
	)
	if err != nil {
		return reqValues, err
	}

	reqValues["build-id"] = strconv.Itoa(build.ID)
	reqValues["name"] = fingerprint.stepName

	return reqValues, nil
}

type checkContainerLocator struct{}

func (locator checkContainerLocator) locate(fingerprint containerFingerprint) (map[string]string, error) {
	reqValues := map[string]string{}

	reqValues["type"] = "check"
	if fingerprint.checkName != "" {
		reqValues["name"] = fingerprint.checkName
	}
	if fingerprint.pipelineName != "" {
		reqValues["pipeline_name"] = fingerprint.pipelineName
	}

	return reqValues, nil
}

type containerFingerprint struct {
	pipelineName  string
	jobName       string
	buildNameOrID string

	stepName string

	checkName string
}

func locateContainer(client concourse.Client, fingerprint containerFingerprint) (map[string]string, error) {
	var locator containerLocator

	if fingerprint.checkName == "" {
		locator = stepContainerLocator{
			client: client,
		}
	} else {
		locator = checkContainerLocator{}
	}

	return locator.locate(fingerprint)
}

func constructRequest(reqGenerator *rata.RequestGenerator, spec atc.HijackProcessSpec, id string, token *rc.TargetToken) *http.Request {
	payload, err := json.Marshal(spec)
	if err != nil {
		log.Fatalln("failed to marshal process spec:", err)
	}

	hijackReq, err := reqGenerator.CreateRequest(
		atc.HijackContainer,
		rata.Params{"id": id},
		bytes.NewBuffer(payload),
	)
	if err != nil {
		log.Fatalln("failed to create hijack request:", err)
	}

	if token != nil {
		hijackReq.Header.Add("Authorization", token.Type+" "+token.Value)
	}

	return hijackReq
}

func getContainerIDs(c *HijackCommand) []atc.Container {
	var pipelineName string
	if c.Job.PipelineName != "" {
		pipelineName = c.Job.PipelineName
	} else {
		pipelineName = c.Check.PipelineName
	}

	buildNameOrID := c.Build
	stepName := c.StepName
	jobName := c.Job.JobName
	check := c.Check.ResourceName

	fingerprint := containerFingerprint{
		pipelineName:  pipelineName,
		jobName:       jobName,
		buildNameOrID: buildNameOrID,
		stepName:      stepName,
		checkName:     check,
	}

	connection, err := rc.TargetConnection(Fly.Target)
	if err != nil {
		log.Fatalln("failed to create client:", err)
	}
	client := concourse.NewClient(connection)

	reqValues, err := locateContainer(client, fingerprint)
	if err != nil {
		log.Fatalln(err)
	}

	containers, err := client.ListContainers(reqValues)
	if err != nil {
		log.Fatalln("failed to get containers:", err)
	}
	return containers
}

func (command *HijackCommand) Execute(args []string) error {
	target, err := rc.SelectTarget(Fly.Target)
	if err != nil {
		log.Fatalln(err)
		return nil
	}

	containers := getContainerIDs(command)

	var chosenContainer atc.Container
	if len(containers) == 0 {
		fmt.Fprintln(os.Stderr, "no containers matched your search parameters! they may have expired if your build hasn't recently finished")
		os.Exit(1)
	} else if len(containers) > 1 {
		var choices []interact.Choice
		for _, container := range containers {
			var infos []string
			if container.PipelineName != "" {
				infos = append(infos, fmt.Sprintf("pipeline: %s", container.PipelineName))
			}

			if container.BuildID != 0 {
				infos = append(infos, fmt.Sprintf("build id: %d", container.BuildID))
			}

			infos = append(infos, fmt.Sprintf("type: %s", container.Type))
			infos = append(infos, fmt.Sprintf("name: %s", container.Name))

			choices = append(choices, interact.Choice{
				Display: strings.Join(infos, ", "),
				Value:   container,
			})
		}

		err = interact.NewInteraction("choose a container", choices...).Resolve(&chosenContainer)
		if err == io.EOF {
			os.Exit(0)
		}

		if err != nil {
			return err
		}
	} else {
		chosenContainer = containers[0]
	}

	path, args := remoteCommand(args)
	privileged := true

	reqGenerator := rata.NewRequestGenerator(target.API, atc.Routes)
	tlsConfig := &tls.Config{InsecureSkipVerify: target.Insecure}

	var ttySpec *atc.HijackTTYSpec
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		ttySpec = &atc.HijackTTYSpec{
			WindowSize: atc.HijackWindowSize{
				Columns: cols,
				Rows:    rows,
			},
		}
	}

	envVariables := append(chosenContainer.EnvironmentVariables, "TERM="+os.Getenv("TERM"))

	spec := atc.HijackProcessSpec{
		Path: path,
		Args: args,
		Env:  envVariables,
		User: "root",
		Dir:  chosenContainer.WorkingDirectory,

		Privileged: privileged,
		TTY:        ttySpec,
	}

	hijackReq := constructRequest(reqGenerator, spec, chosenContainer.ID, target.Token)
	hijackResult := performHijack(hijackReq, tlsConfig)
	os.Exit(hijackResult)

	return nil
}

func performHijack(hijackReq *http.Request, tlsConfig *tls.Config) int {
	conn, err := dialEndpoint(hijackReq.URL, tlsConfig)
	if err != nil {
		log.Fatalln("failed to dial hijack endpoint:", err)
	}

	clientConn := httputil.NewClientConn(conn, nil)

	resp, err := clientConn.Do(hijackReq)
	if err != nil {
		log.Fatalln("failed to hijack:", err)
	}

	if resp.StatusCode != http.StatusOK {
		handleBadResponse("hijacking", resp)
	}

	return hijack(clientConn.Hijack())
}

func hijack(conn net.Conn, br *bufio.Reader) int {
	var in io.Reader

	term, err := pty.OpenRawTerm()
	if err == nil {
		defer term.Restore()

		in = term
	} else {
		in = os.Stdin
	}

	encoder := json.NewEncoder(conn)
	decoder := json.NewDecoder(br)

	resized := pty.ResizeNotifier()

	go func() {
		for {
			<-resized
			// TODO json race
			sendSize(encoder)
		}
	}()

	go io.Copy(&stdinWriter{encoder}, in)

	var exitStatus int
	for {
		var output atc.HijackOutput
		err := decoder.Decode(&output)
		if err != nil {
			break
		}

		if output.ExitStatus != nil {
			exitStatus = *output.ExitStatus
		} else if len(output.Error) > 0 {
			fmt.Fprintf(os.Stderr, "%s\n", ansi.Color(output.Error, "red+b"))
			exitStatus = 255
		} else if len(output.Stdout) > 0 {
			os.Stdout.Write(output.Stdout)
		} else if len(output.Stderr) > 0 {
			os.Stderr.Write(output.Stderr)
		}
	}

	return exitStatus
}

func sendSize(enc *json.Encoder) {
	rows, cols, err := pty.Getsize(os.Stdin)
	if err == nil {
		enc.Encode(atc.HijackInput{
			TTYSpec: &atc.HijackTTYSpec{
				WindowSize: atc.HijackWindowSize{
					Columns: cols,
					Rows:    rows,
				},
			},
		})
	}
}

type stdinWriter struct {
	enc *json.Encoder
}

func (w *stdinWriter) Write(d []byte) (int, error) {
	err := w.enc.Encode(atc.HijackInput{
		Stdin: d,
	})
	if err != nil {
		return 0, err
	}

	return len(d), nil
}

var canonicalPortMap = map[string]string{
	"http":  "80",
	"https": "443",
}

func dialEndpoint(url *url.URL, tlsConfig *tls.Config) (net.Conn, error) {
	addr := canonicalAddr(url)

	if url.Scheme == "https" {
		return tls.Dial("tcp", addr, tlsConfig)
	}

	return net.Dial("tcp", addr)
}

func canonicalAddr(url *url.URL) string {
	host, port, err := net.SplitHostPort(url.Host)
	if err != nil {
		if strings.Contains(err.Error(), "missing port in address") {
			host = url.Host
			port = canonicalPortMap[url.Scheme]
		} else {
			log.Fatalln("invalid host:", err)
		}
	}

	return net.JoinHostPort(host, port)
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"testing"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cfpb/rhobot/config"
	"github.com/cfpb/rhobot/database"
	"github.com/cfpb/rhobot/healthcheck"
	"github.com/cfpb/rhobot/report"
)

var conf *config.Config

func init() {
	conf = config.NewConfig()
	conf.SetLogLevel("info")
}

func TestPostgresHealthCheckReporting(t *testing.T) {
	cxn := database.GetPGConnection(conf.DBURI())
	healthChecks, _ := healthcheck.ReadHealthCheckYAMLFromFile("healthcheck/healthchecksAll.yml")
	results, _ := healthChecks.PreformHealthChecks(cxn)
	var elements []report.Element
	for _, val := range results {
		elements = append(elements, val)
	}
	metadata := map[string]interface{}{
		"name":      healthChecks.Name,
		"schema":    "public",
		"table":     "healthchecks",
		"timestamp": time.Now().Format(time.ANSIC),
	}
	rs := report.Set{Elements: elements, Metadata: metadata}

	prr := report.NewPongo2ReportRunnerFromString(healthcheck.TemplateHealthcheckPostgres)
	pgr := report.PGHandler{Cxn: cxn}
	reader, err := prr.ReportReader(rs)
	err = pgr.HandleReport(reader)
	if err != nil {
		t.Fatalf("error writing report to PG database\n%s", err)
	}
}

func TestLogLevelingFiltering(t *testing.T) {
	cxn := database.GetPGConnection(conf.DBURI())
	healthChecks, _ := healthcheck.ReadHealthCheckYAMLFromFile("healthcheck/healthchecksAll.yml")
	results, _ := healthChecks.PreformHealthChecks(cxn)
	var elements []report.Element
	for _, val := range results {
		elements = append(elements, val)
	}
	metadata := map[string]interface{}{
		"name": healthChecks.Name,
	}
	originalReportSet := report.Set{Elements: elements, Metadata: metadata}
	TestLogLevelResults := map[string]int{
		"Debug": 6,
		"Info":  5,
		"Warn":  4,
		"Error": 2,
		"Fatal": 1,
	}

	for _, level := range report.LogLevelArray {
		log.Infof("Report Filter LogLevel %v", level)
		logFilteredSet := report.FilterReportSet(originalReportSet, level)
		prr := report.JSONReportRunner{}
		reader, _ := prr.ReportReader(logFilteredSet)
		phr := report.PrintHandler{}
		_ = phr.HandleReport(reader)

		if len(logFilteredSet.GetElementArray()) != TestLogLevelResults[level] {
			t.Fatalf("wrong number of healthchecks in report")
		}
	}

}

//The following are rhobot CLI test

func TestCLI_Healthchecks(t *testing.T) {
	//TODO create fail verification test

	args := []string{"rhobot", "healthchecks", "healthcheck/healthchecksTest.yml"}
	os.Args = args
	main()
}

func TestCLI_PG_Healthchecks(t *testing.T) {

	//clear hctest
	cxn := database.GetPGConnection(conf.DBURI())
	result, err := cxn.Exec("DROP TABLE IF EXISTS public.hctest")
	if err != nil {
		t.Fatalf("error dropping public.hctest\n%s", err)
	}

	args := []string{"rhobot", "healthchecks", "healthcheck/healthchecksTest.yml",
		"--hcschema", "public", "--hctable", "hctest"}
	os.Args = args
	main()

	//make sure hctest only has 3 rows
	row, err := cxn.Query("SELECT count(*) FROM public.hctest;")
	defer row.Close()
	if err != nil {
		t.Fatalf("error selecting public.hctest\n%s", err)
	} else {

		if row.Next() {
			var count int
			err = row.Scan(&count)

			if count != 3 {
				t.Fatal("public.hctest count should have been 3")
			}

		} else {
			t.Fatalf("no results in selecting public.hctest\n%s", err)
		}

	}

	log.Info(result)
}

func TestCLI(t *testing.T) {
	//All the following should exit with 0

	os.Args = []string{"rhobot"}
	main()

	os.Args = []string{"rhobot", "-V"}
	main()

	os.Args = []string{"healthchecks"}
	main()
}
<commit_msg>change test args to match new names<commit_after>package main

import (
	"os"
	"testing"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cfpb/rhobot/config"
	"github.com/cfpb/rhobot/database"
	"github.com/cfpb/rhobot/healthcheck"
	"github.com/cfpb/rhobot/report"
)

var conf *config.Config

func init() {
	conf = config.NewConfig()
	conf.SetLogLevel("info")
}

func TestPostgresHealthCheckReporting(t *testing.T) {
	cxn := database.GetPGConnection(conf.DBURI())
	healthChecks, _ := healthcheck.ReadHealthCheckYAMLFromFile("healthcheck/healthchecksAll.yml")
	results, _ := healthChecks.PreformHealthChecks(cxn)
	var elements []report.Element
	for _, val := range results {
		elements = append(elements, val)
	}
	metadata := map[string]interface{}{
		"name":      healthChecks.Name,
		"schema":    "public",
		"table":     "healthchecks",
		"timestamp": time.Now().Format(time.ANSIC),
	}
	rs := report.Set{Elements: elements, Metadata: metadata}

	prr := report.NewPongo2ReportRunnerFromString(healthcheck.TemplateHealthcheckPostgres)
	pgr := report.PGHandler{Cxn: cxn}
	reader, err := prr.ReportReader(rs)
	err = pgr.HandleReport(reader)
	if err != nil {
		t.Fatalf("error writing report to PG database\n%s", err)
	}
}

func TestLogLevelingFiltering(t *testing.T) {
	cxn := database.GetPGConnection(conf.DBURI())
	healthChecks, _ := healthcheck.ReadHealthCheckYAMLFromFile("healthcheck/healthchecksAll.yml")
	results, _ := healthChecks.PreformHealthChecks(cxn)
	var elements []report.Element
	for _, val := range results {
		elements = append(elements, val)
	}
	metadata := map[string]interface{}{
		"name": healthChecks.Name,
	}
	originalReportSet := report.Set{Elements: elements, Metadata: metadata}
	TestLogLevelResults := map[string]int{
		"Debug": 6,
		"Info":  5,
		"Warn":  4,
		"Error": 2,
		"Fatal": 1,
	}

	for _, level := range report.LogLevelArray {
		log.Infof("Report Filter LogLevel %v", level)
		logFilteredSet := report.FilterReportSet(originalReportSet, level)
		prr := report.JSONReportRunner{}
		reader, _ := prr.ReportReader(logFilteredSet)
		phr := report.PrintHandler{}
		_ = phr.HandleReport(reader)

		if len(logFilteredSet.GetElementArray()) != TestLogLevelResults[level] {
			t.Fatalf("wrong number of healthchecks in report")
		}
	}

}

//The following are rhobot CLI test

func TestCLI_Healthchecks(t *testing.T) {
	//TODO create fail verification test

	args := []string{"rhobot", "healthchecks", "healthcheck/healthchecksTest.yml"}
	os.Args = args
	main()
}

func TestCLI_PG_Healthchecks(t *testing.T) {

	//clear hctest
	cxn := database.GetPGConnection(conf.DBURI())
	result, err := cxn.Exec("DROP TABLE IF EXISTS public.hctest")
	if err != nil {
		t.Fatalf("error dropping public.hctest\n%s", err)
	}

	args := []string{"rhobot", "healthchecks", "healthcheck/healthchecksTest.yml",
		"--schema", "public", "--table", "hctest"}
	os.Args = args
	main()

	//make sure hctest only has 3 rows
	row, err := cxn.Query("SELECT count(*) FROM public.hctest;")
	defer row.Close()
	if err != nil {
		t.Fatalf("error selecting public.hctest\n%s", err)
	} else {

		if row.Next() {
			var count int
			err = row.Scan(&count)

			if count != 3 {
				t.Fatal("public.hctest count should have been 3")
			}

		} else {
			t.Fatalf("no results in selecting public.hctest\n%s", err)
		}

	}

	log.Info(result)
}

func TestCLI(t *testing.T) {
	//All the following should exit with 0

	os.Args = []string{"rhobot"}
	main()

	os.Args = []string{"rhobot", "-V"}
	main()

	os.Args = []string{"healthchecks"}
	main()
}
<|endoftext|>


<commit_before>package logger

import (
	"context"

	"git.timschuster.info/rls.moe/catgi/snowflakes"
	"github.com/Sirupsen/logrus"
)

// LogFromCtx returns a logger with the given source tag.
// If the context contains a logger it will utilize it as basis.
// If it contains no logger it will create a new logger.
func LogFromCtx(src string, ctx context.Context) logrus.FieldLogger {
	if ctx != nil {
		if val := ctx.Value("logger"); val != nil {
			if log, ok := val.(logrus.FieldLogger); ok {
				if reqId, ok := ctx.Value("logger-req-id").(string); ok {
					return log.WithField("src", src).WithField("req-id", reqId)
				}
				return log.WithField("src", src)
			}
			logrus.WithField("src", "logger").
				Error("Context had an invalid logger")
		}
	}
	return logrus.New().WithField("src", src).WithField("no-ctx", "")
}

// NewLoggingContext returns a background context with a logger
func NewLoggingContext() context.Context {
	return InjectLogToContext(context.Background())
}

// InjectLogToContext injects a logger into the context
func InjectLogToContext(ctx context.Context) context.Context {
	logs := logrus.New()
	return context.WithValue(ctx, "logger", logs)
}

// CreateRequestIDContext creates a snowflake and tags
// all log messages from that context with it.
// cmd/catgi uses this to differentiate HTTP requests
func CreateRequestIDContext(ctx context.Context) context.Context {
	log := LogFromCtx("cr-req-id", ctx)
	sf, err := snowflakes.NewSnowflake()
	if err != nil {
		log.Error("Error on Generating Request ID: ", err)
		return ctx
	}
	return context.WithValue(ctx, "logger-req-id", sf)

}

// SetLoggingLevel sets the logging level of a logger inside
// the context.
// If the logging level is unknown it panics.
func SetLoggingLevel(level string, ctx context.Context) context.Context {
	log, ok := ctx.Value("logger").(*logrus.Logger)
	if !ok {
		return ctx
	}
	lvl, err := logrus.ParseLevel(level)
	if err != nil {
		panic(err)
	}
	log.Level = lvl
	return context.WithValue(ctx, "logger", log)
}
<commit_msg>Added ctx log<commit_after>package logger

import (
	"context"

	"git.timschuster.info/rls.moe/catgi/snowflakes"
	"github.com/Sirupsen/logrus"
)

// Localised Logging interface
type Logger logrus.FieldLogger

// LogFromCtx returns a logger with the given source tag.
// If the context contains a logger it will utilize it as basis.
// If it contains no logger it will create a new logger.
func LogFromCtx(src string, ctx context.Context) Logger {
	if ctx != nil {
		if val := ctx.Value("logger"); val != nil {
			if log, ok := val.(logrus.FieldLogger); ok {
				if reqId, ok := ctx.Value("logger-req-id").(string); ok {
					return log.WithField("src", src).WithField("req-id", reqId)
				}
				return log.WithField("src", src)
			}
			logrus.WithField("src", "logger").
				Error("Context had an invalid logger")
		}
	}
	return logrus.New().WithField("src", src).WithField("no-ctx", "")
}

// NewLoggingContext returns a background context with a logger
func NewLoggingContext() context.Context {
	return InjectLogToContext(context.Background())
}

// InjectLogToContext injects a logger into the context
func InjectLogToContext(ctx context.Context) context.Context {
	logs := logrus.New()
	return context.WithValue(ctx, "logger", logs)
}

// CreateRequestIDContext creates a snowflake and tags
// all log messages from that context with it.
// cmd/catgi uses this to differentiate HTTP requests
func CreateRequestIDContext(ctx context.Context) context.Context {
	log := LogFromCtx("cr-req-id", ctx)
	sf, err := snowflakes.NewSnowflake()
	if err != nil {
		log.Error("Error on Generating Request ID: ", err)
		return ctx
	}
	return context.WithValue(ctx, "logger-req-id", sf)

}

// SetLoggingLevel sets the logging level of a logger inside
// the context.
// If the logging level is unknown it panics.
func SetLoggingLevel(level string, ctx context.Context) context.Context {
	log, ok := ctx.Value("logger").(*logrus.Logger)
	if !ok {
		return ctx
	}
	lvl, err := logrus.ParseLevel(level)
	if err != nil {
		panic(err)
	}
	log.Level = lvl
	return context.WithValue(ctx, "logger", log)
}
<|endoftext|>


<commit_before>package logger

import (
	"bytes"
	"fmt"
	"log/syslog"
	"runtime"
	"strings"

	log "github.com/sirupsen/logrus"
	logrus_syslog "github.com/sirupsen/logrus/hooks/syslog"
	"golang.org/x/net/context"

	"github.com/servicekit/servicekit-go/config"
)

const (
	RequestIDKey = "request-id"
)

func insert(slice []interface{}, insertion interface{}) []interface{} {
	result := make([]interface{}, len(slice)+1)
	result[0] = insertion
	copy(result[1:], slice)
	return result
}

type LogFormatter struct {
}

func (f *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}

	b.WriteString(strings.ToUpper(entry.Level.String()))
	b.WriteString(" ")
	b.WriteString(entry.Message)
	b.WriteString("\n")
	return b.Bytes(), nil
}

type Logger struct {
	logger *log.Logger
}

func NewLogger(serviceName, serviceVersion string, serviceENV config.ServiceENV, network, addr string, priority syslog.Priority) (*Logger, error) {
	logger := &Logger{}
	logger.logger = log.New()

	logger.logger.Formatter = &LogFormatter{}

	hook, err := logrus_syslog.NewSyslogHook(network, addr, priority, fmt.Sprintf("%s_%s_%s", serviceName, serviceVersion, serviceENV))
	if err != nil {
		return nil, err
	} else {
		logger.logger.Hooks.Add(hook)
		return logger, nil
	}
}

func (logger *Logger) Debugf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Debugf(f, args...)
}

func (logger *Logger) Infof(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Infof(f, args...)
}

func (logger *Logger) Printf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Printf(f, args...)
}

func (logger *Logger) Warnf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Warnf(f, args...)
}

func (logger *Logger) Errorf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Errorf(f, args...)
}

func (logger *Logger) Fatalf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Fatalf(f, args...)
}

func (logger *Logger) Panicf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Panicf(f, args...)
}

func (logger *Logger) DebugfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(RequestIDKey))
	}

	logger.logger.Debugf(f, args...)
}

func (logger *Logger) InfofWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(RequestIDKey))
	}

	logger.logger.Infof(f, args...)
}

func (logger *Logger) PrintfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(RequestIDKey))
	}

	logger.logger.Printf(f, args...)
}

func (logger *Logger) WarnfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(RequestIDKey))
	}

	logger.logger.Warnf(f, args...)
}

func (logger *Logger) ErrorfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(RequestIDKey))
	}

	logger.logger.Errorf(f, args...)
}

func (logger *Logger) FatalfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(RequestIDKey))
	}

	logger.logger.Fatalf(f, args...)
}

func (logger *Logger) PanicfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(RequestIDKey))
	}

	logger.logger.Panicf(f, args...)
}

func (logger *Logger) Debug(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Debug(args...)
}

func (logger *Logger) Info(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Info(args...)
}

func (logger *Logger) Print(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Print(args...)
}

func (logger *Logger) Warn(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Warn(args...)
}

func (logger *Logger) Error(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Error(args...)
}

func (logger *Logger) Fatal(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Fatal(args...)
}

func (logger *Logger) Panic(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Panic(args...)
}
<commit_msg>update<commit_after>package logger

import (
	"bytes"
	"fmt"
	"log/syslog"
	"runtime"
	"strings"

	log "github.com/sirupsen/logrus"
	logrus_syslog "github.com/sirupsen/logrus/hooks/syslog"
	"golang.org/x/net/context"

	"github.com/servicekit/servicekit-go/config"
	"github.com/servicekit/servicekit-go/requestid"
)

func insert(slice []interface{}, insertion interface{}) []interface{} {
	result := make([]interface{}, len(slice)+1)
	result[0] = insertion
	copy(result[1:], slice)
	return result
}

type LogFormatter struct {
}

func (f *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
	var b *bytes.Buffer
	if entry.Buffer != nil {
		b = entry.Buffer
	} else {
		b = &bytes.Buffer{}
	}

	b.WriteString(strings.ToUpper(entry.Level.String()))
	b.WriteString(" ")
	b.WriteString(entry.Message)
	b.WriteString("\n")
	return b.Bytes(), nil
}

type Logger struct {
	logger *log.Logger
}

func NewLogger(serviceName, serviceVersion string, serviceENV config.ServiceENV, network, addr string, priority syslog.Priority) (*Logger, error) {
	logger := &Logger{}
	logger.logger = log.New()

	logger.logger.Formatter = &LogFormatter{}

	hook, err := logrus_syslog.NewSyslogHook(network, addr, priority, fmt.Sprintf("%s_%s_%s", serviceName, serviceVersion, serviceENV))
	if err != nil {
		return nil, err
	} else {
		logger.logger.Hooks.Add(hook)
		return logger, nil
	}
}

func (logger *Logger) Debugf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Debugf(f, args...)
}

func (logger *Logger) Infof(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Infof(f, args...)
}

func (logger *Logger) Printf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Printf(f, args...)
}

func (logger *Logger) Warnf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Warnf(f, args...)
}

func (logger *Logger) Errorf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Errorf(f, args...)
}

func (logger *Logger) Fatalf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Fatalf(f, args...)
}

func (logger *Logger) Panicf(format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s", file, line, format)
	}

	logger.logger.Panicf(f, args...)
}

func (logger *Logger) DebugfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(requestid.RequestIDKey))
	}

	logger.logger.Debugf(f, args...)
}

func (logger *Logger) InfofWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(requestid.RequestIDKey))
	}

	logger.logger.Infof(f, args...)
}

func (logger *Logger) PrintfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(requestid.RequestIDKey))
	}

	logger.logger.Printf(f, args...)
}

func (logger *Logger) WarnfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(requestid.RequestIDKey))
	}

	logger.logger.Warnf(f, args...)
}

func (logger *Logger) ErrorfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(requestid.RequestIDKey))
	}

	logger.logger.Errorf(f, args...)
}

func (logger *Logger) FatalfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(requestid.RequestIDKey))
	}

	logger.logger.Fatalf(f, args...)
}

func (logger *Logger) PanicfWithReqID(ctx context.Context, format string, args ...interface{}) {
	f := format

	_, file, line, ok := runtime.Caller(1)
	if ok {
		f = fmt.Sprintf("%s:%d %s reqid: %v", file, line, format, ctx.Value(requestid.RequestIDKey))
	}

	logger.logger.Panicf(f, args...)
}

func (logger *Logger) Debug(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Debug(args...)
}

func (logger *Logger) Info(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Info(args...)
}

func (logger *Logger) Print(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Print(args...)
}

func (logger *Logger) Warn(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Warn(args...)
}

func (logger *Logger) Error(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Error(args...)
}

func (logger *Logger) Fatal(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Fatal(args...)
}

func (logger *Logger) Panic(args ...interface{}) {
	_, file, line, ok := runtime.Caller(1)
	if ok {
		args = insert(args, fmt.Sprintf("%s:%d ", file, line))
	}

	logger.logger.Panic(args...)
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/spf13/cobra"

	"github.com/nanopack/hoarder/config"
)

var updateCmd = &cobra.Command{
	Use:   "update",
	Short: "Update a file in hoarder",
	Long:  ``,

	Run: update,
}

// init
func init() {
	updateCmd.Flags().StringVarP(&key, "key", "k", "", "The key to store the data by")
	updateCmd.Flags().StringVarP(&data, "data", "d", "", "The raw data to be stored")
}

// update
func update(ccmd *cobra.Command, args []string) {

	// handle any missing args
	switch {
	case key == "":
		fmt.Println("Missing key - please provide the key for the record you'd like to update")
	case data == "":
		fmt.Println("Missing data - please provide the data that you would like to update")
	}

	//
	body := bytes.NewBuffer([]byte(args[1]))

	//
	req, err := http.NewRequest("PUT", fmt.Sprintf("%s/blobs/%s", config.URI, args[0]), body)
	if err != nil {
		fmt.Println("ERR!!", err)
	}

	//
	req.Header.Add("X-NANOBOX-TOKEN", config.Token)

	//
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Println("ERR!!", err)
	}
	defer res.Body.Close()

	//
	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		fmt.Println("ERR!!", err)
	}

	fmt.Println("UPDATE??", string(b))
}
<commit_msg>correctly returning from failed flag checks<commit_after>package commands

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/spf13/cobra"

	"github.com/nanopack/hoarder/config"
)

var updateCmd = &cobra.Command{
	Use:   "update",
	Short: "Update a file in hoarder",
	Long:  ``,

	Run: update,
}

// init
func init() {
	updateCmd.Flags().StringVarP(&key, "key", "k", "", "The key to store the data by")
	updateCmd.Flags().StringVarP(&data, "data", "d", "", "The raw data to be stored")
}

// update
func update(ccmd *cobra.Command, args []string) {

	// handle any missing args
	switch {
	case key == "":
		fmt.Println("Missing key - please provide the key for the record you'd like to update")
		return
	case data == "":
		fmt.Println("Missing data - please provide the data that you would like to update")
		return
	}

	//
	body := bytes.NewBuffer([]byte(args[1]))

	//
	req, err := http.NewRequest("PUT", fmt.Sprintf("%s/blobs/%s", config.URI, args[0]), body)
	if err != nil {
		fmt.Println("ERR!!", err)
	}

	//
	req.Header.Add("X-NANOBOX-TOKEN", config.Token)

	//
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Println("ERR!!", err)
	}
	defer res.Body.Close()

	//
	b, err := ioutil.ReadAll(res.Body)
	if err != nil {
		fmt.Println("ERR!!", err)
	}

	fmt.Println("UPDATE??", string(b))
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/slugalisk/overrustlelogs/common"
)

var (
	ErrNotInChannel     = errors.New("not in channel")
	ErrAlreadyInChannel = errors.New("already in channel")
	ErrChannelNotValid  = errors.New("channel not valid")
)

// TwitchLogger ...
type TwitchLogger struct {
	chatLock sync.RWMutex
	chats    map[int]*common.Twitch

	chLock   sync.RWMutex
	channels []string

	logHandler     func(m <-chan *common.Message)
	admins         map[string]struct{}
	commandChannel string
}

// NewTwitchLogger ...
func NewTwitchLogger(f func(m <-chan *common.Message)) *TwitchLogger {
	t := &TwitchLogger{
		chats:      make(map[int]*common.Twitch, 0),
		admins:     make(map[string]struct{}),
		logHandler: f,
	}

	admins := common.GetConfig().Twitch.Admins
	for _, a := range admins {
		t.admins[a] = struct{}{}
	}

	d, err := ioutil.ReadFile(common.GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Fatalf("unable to read channels %s", err)
	}

	if err := json.Unmarshal(d, &t.channels); err != nil {
		log.Fatalf("unable to read channels %s", err)
	}

	t.commandChannel = common.GetConfig().Twitch.CommandChannel

	return t
}

// Start ...
func (t *TwitchLogger) Start() {
	for _, channel := range t.channels {
		err := t.join(channel, false)
		if err != nil {
			log.Printf("failed to join %s err: %s", channel, err.Error())
		}
	}
	log.Println("joined", len(t.channels), "chats, wew lad :^)")
}

// Stop ...
func (t *TwitchLogger) Stop() {
	t.chatLock.Lock()
	defer t.chatLock.Unlock()
	for id, c := range t.chats {
		log.Printf("stopping chat: %d\n", id)
		c.Stop()
	}
}

func (t *TwitchLogger) join(ch string, init bool) error {
	if init {
		if !channelExists(ch) {
			return ErrChannelNotValid
		}
		if inSlice(t.channels, ch) {
			return ErrAlreadyInChannel
		}
		t.addChannel(ch)
		if err := t.saveChannels(); err != nil {
			log.Println(err)
		}
	}
	c := t.getChatToJoin()
	for try := 1; try <= 3; try++ {
		err := c.Join(ch)
		if err == nil {
			log.Printf("joining %s", ch)
			return nil
		}
		if err != nil && try == 3 {
			return fmt.Errorf("failed to join %s :( %v", ch, err)
		}
		log.Println("retrying to join", ch)
	}
	return nil
}

func (t *TwitchLogger) leave(ch string) error {
	t.chLock.Lock()
	if !inSlice(t.channels, ch) {
		t.chLock.Unlock()
		return ErrNotInChannel
	}
	t.chLock.Unlock()

	t.chatLock.Lock()
	defer t.chatLock.Unlock()
	for i, c := range t.chats {
		if inSlice(c.Channels(), ch) {
			c.Leave(ch)
			log.Printf("found channel in chat %d", i)
			break
		}
	}
	t.removeChannel(ch)
	err := t.saveChannels()
	if err != nil {
		log.Println(err)
	}
	log.Println("left", ch)
	return err
}

func (t *TwitchLogger) getChatToJoin() *common.Twitch {
	t.chatLock.Lock()
	for _, c := range t.chats {
		c.ChLock.Lock()
		if len(c.Channels()) < common.MaxChannelsPerChat {
			c.ChLock.Unlock()
			t.chatLock.Unlock()
			return c
		}
		c.ChLock.Unlock()
	}
	c := common.NewTwitch()
	c.Run()
	i := len(t.chats) + 1
	t.chats[i] = c
	go t.msgHandler(i, c.Messages())
	t.chatLock.Unlock()
	return c
}

func (t *TwitchLogger) msgHandler(i int, ch <-chan *common.Message) {
	logCh := make(chan *common.Message, common.MessageBufferSize)
	defer close(logCh)
	go t.logHandler(logCh)
	for m := range ch {
		if t.commandChannel == m.Channel {
			go t.runCommand(i, m)
		}
		select {
		case logCh <- m:
		default:
			log.Println("error buffer is full")
		}
	}
}

func (t *TwitchLogger) runCommand(i int, m *common.Message) {
	c, ok := t.chats[i]
	if !ok {
		return
	}
	if _, ok := t.admins[m.Nick]; !ok && m.Type != "MSG" {
		return
	}
	parts := strings.Split(strings.ToLower(m.Data), " ")
	switch parts[0] {
	case "!join":
		err := t.join(parts[1], true)
		if err != nil {
			c.Message(m.Channel, err.Error())
			return
		}
		c.Message(m.Channel, fmt.Sprintf("Logging %s", parts[1]))
	case "!leave":
		err := t.leave(parts[1])
		if err != nil {
			c.Message(m.Channel, fmt.Sprintf("Not logging %s", parts[1]))
			return
		}
		c.Message(m.Channel, fmt.Sprintf("Leaving %s", parts[1]))
	}
}

func (t *TwitchLogger) addChannel(ch string) {
	t.chLock.Lock()
	t.channels = append(t.channels, ch)
	t.chLock.Unlock()
}

func (t *TwitchLogger) removeChannel(ch string) error {
	t.chLock.Lock()
	defer t.chLock.Unlock()
	sort.Strings(t.channels)
	i := sort.SearchStrings(t.channels, ch)
	if i < len(t.channels) && t.channels[i] == ch {
		t.channels = append(t.channels[:i], t.channels[i+1:]...)
		return nil
	}
	return errors.New("didn't find " + ch + " in the channels list")
}

func (t *TwitchLogger) saveChannels() error {
	f, err := os.Create(common.GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	defer f.Close()
	sort.Strings(t.channels)
	data, err := json.Marshal(t.channels)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	var buf bytes.Buffer
	if err := json.Indent(&buf, data, "", "\t"); err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	f.Write(buf.Bytes())
	return nil
}

// channelExists
func channelExists(ch string) bool {
	req, err := http.NewRequest("GET", "https://api.twitch.tv/kraken/users/"+strings.ToLower(ch), nil)
	if err != nil {
		return false
	}
	req.Header.Add("Client-ID", common.GetConfig().Twitch.ClientID)
	client := http.Client{
		Timeout: 10 * time.Second,
	}
	res, err := client.Do(req)
	if err != nil {
		log.Println(err)
		return false
	}
	res.Body.Close()
	return res.StatusCode == http.StatusOK
}

func inSlice(slice []string, s string) bool {
	for _, v := range slice {
		if strings.EqualFold(s, v) {
			return true
		}
	}
	return false
}
<commit_msg>refactor<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/slugalisk/overrustlelogs/common"
)

var (
	ErrNotInChannel     = errors.New("not in channel")
	ErrAlreadyInChannel = errors.New("already in channel")
	ErrChannelNotValid  = errors.New("channel not valid")
)

// TwitchLogger ...
type TwitchLogger struct {
	chatLock sync.RWMutex
	chats    map[int]*common.Twitch

	chLock   sync.RWMutex
	channels []string

	logHandler     func(m <-chan *common.Message)
	admins         map[string]struct{}
	commandChannel string
}

// NewTwitchLogger ...
func NewTwitchLogger(f func(m <-chan *common.Message)) *TwitchLogger {
	t := &TwitchLogger{
		chats:      make(map[int]*common.Twitch),
		admins:     make(map[string]struct{}),
		logHandler: f,
	}

	admins := common.GetConfig().Twitch.Admins
	for _, a := range admins {
		t.admins[a] = struct{}{}
	}

	d, err := ioutil.ReadFile(common.GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Fatalf("unable to read channels %s", err)
	}

	if err := json.Unmarshal(d, &t.channels); err != nil {
		log.Fatalf("unable to read channels %s", err)
	}

	t.commandChannel = common.GetConfig().Twitch.CommandChannel

	return t
}

// Start ...
func (t *TwitchLogger) Start() {
	t.join(common.GetConfig().Twitch.CommandChannel, false)
	var c int
	for _, channel := range t.channels {
		err := t.join(channel, false)
		if err != nil {
			log.Printf("failed to join %s err: %s", channel, err.Error())
			continue
		}
		c++
	}
	log.Println("joined", c, "chats, wew lad :^)")
}

// Stop ...
func (t *TwitchLogger) Stop() {
	t.saveChannels()
	t.chatLock.Lock()
	for id, c := range t.chats {
		log.Printf("stopping chat: %d\n", id)
		c.Stop()
	}
	t.chatLock.Unlock()
}

func (t *TwitchLogger) join(ch string, init bool) error {
	if init {
		if !channelExists(ch) {
			return ErrChannelNotValid
		}
		if inSlice(t.channels, ch) {
			return ErrAlreadyInChannel
		}
		t.addChannel(ch)
		if err := t.saveChannels(); err != nil {
			log.Println(err)
		}
	}
	c := t.getChatToJoin()
	for try := 1; try <= 3; try++ {
		err := c.Join(ch)
		if err == nil {
			log.Printf("joining %s", ch)
			return nil
		}
		if err != nil && try == 3 {
			return fmt.Errorf("failed to join %s :( %v", ch, err)
		}
		log.Println("retrying to join", ch)
	}
	return nil
}

func (t *TwitchLogger) leave(ch string) error {
	t.chLock.Lock()
	if !inSlice(t.channels, ch) {
		t.chLock.Unlock()
		return ErrNotInChannel
	}
	t.chLock.Unlock()

	t.chatLock.Lock()
	for i, c := range t.chats {
		if inSlice(c.Channels(), ch) {
			c.Leave(ch)
			log.Printf("found channel in chat %d", i)
			break
		}
	}
	t.chatLock.Unlock()
	t.removeChannel(ch)
	err := t.saveChannels()
	if err != nil {
		log.Println(err)
	}
	log.Println("left", ch)
	return err
}

func (t *TwitchLogger) getChatToJoin() *common.Twitch {
	t.chatLock.Lock()
	defer t.chatLock.Unlock()
	for _, c := range t.chats {
		c.ChLock.Lock()
		if len(c.Channels()) < common.MaxChannelsPerChat {
			c.ChLock.Unlock()
			return c
		}
		c.ChLock.Unlock()
	}
	c := common.NewTwitch()
	c.Run()
	i := len(t.chats) + 1
	t.chats[i] = c
	go t.msgHandler(i, c.Messages())
	return c
}

func (t *TwitchLogger) msgHandler(i int, ch <-chan *common.Message) {
	logCh := make(chan *common.Message, common.MessageBufferSize)
	defer close(logCh)
	go t.logHandler(logCh)
	for m := range ch {
		if t.commandChannel == m.Channel {
			go t.runCommand(i, m)
		}
		select {
		case logCh <- m:
		default:
			log.Println("error buffer is full")
		}
	}
}

func (t *TwitchLogger) runCommand(i int, m *common.Message) {
	c, ok := t.chats[i]
	if !ok {
		return
	}
	if _, ok := t.admins[m.Nick]; !ok && m.Type != "MSG" {
		return
	}
	parts := strings.Split(strings.ToLower(m.Data), " ")
	switch parts[0] {
	case "!join":
		err := t.join(parts[1], true)
		if err != nil {
			c.Message(m.Channel, err.Error())
			return
		}
		c.Message(m.Channel, fmt.Sprintf("Logging %s", parts[1]))
	case "!leave":
		err := t.leave(parts[1])
		if err != nil {
			c.Message(m.Channel, fmt.Sprintf("Not logging %s", parts[1]))
			return
		}
		c.Message(m.Channel, fmt.Sprintf("Leaving %s", parts[1]))
	}
}

func (t *TwitchLogger) addChannel(ch string) {
	t.chLock.Lock()
	t.channels = append(t.channels, ch)
	t.chLock.Unlock()
}

func (t *TwitchLogger) removeChannel(ch string) error {
	t.chLock.Lock()
	defer t.chLock.Unlock()
	sort.Strings(t.channels)
	i := sort.SearchStrings(t.channels, ch)
	if i < len(t.channels) && t.channels[i] == ch {
		t.channels = append(t.channels[:i], t.channels[i+1:]...)
		return nil
	}
	return errors.New("didn't find " + ch + " in the channels list")
}

func (t *TwitchLogger) saveChannels() error {
	f, err := os.Create(common.GetConfig().Twitch.ChannelListPath)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	defer f.Close()
	sort.Strings(t.channels)
	data, err := json.Marshal(t.channels)
	if err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	var buf bytes.Buffer
	if err := json.Indent(&buf, data, "", "\t"); err != nil {
		log.Printf("error saving channel list %s", err)
		return err
	}
	f.Write(buf.Bytes())
	return nil
}

// channelExists
func channelExists(ch string) bool {
	req, err := http.NewRequest("GET", "https://api.twitch.tv/kraken/users/"+strings.ToLower(ch), nil)
	if err != nil {
		return false
	}
	req.Header.Add("Client-ID", common.GetConfig().Twitch.ClientID)
	client := http.Client{
		Timeout: 10 * time.Second,
	}
	res, err := client.Do(req)
	if err != nil {
		log.Println(err)
		return false
	}
	res.Body.Close()
	return res.StatusCode == http.StatusOK
}

func inSlice(slice []string, s string) bool {
	for _, v := range slice {
		if strings.EqualFold(s, v) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package transloadit

import (
	"io"
	"os"
	"path/filepath"
	"testing"
)

func TestWatch(t *testing.T) {

	// Clean up testing environment
	remove("./fixtures/input/lol_cat.jpg")
	remove("./fixtures/input/mona_lisa.jpg")
	remove("./fixtures/output/resize_0_lol_cat.jpg")
	remove("./fixtures/output/-original_0_lol_cat.jpg")
	remove("./fixtures/output/resize_0_mona_lisa.jpg")
	remove("./fixtures/output/-original_0_mona_lisa.jpg")
	copyFile("./fixtures/lol_cat.jpg", "./fixtures/input/lol_cat.jpg")

	client := setup(t)

	options := &WatchOptions{
		TemplateId: "68302e702fbd11e4b5fb01f025693443",
		Input:      "./fixtures/input",
		Output:     "./fixtures/output",
		Steps:      make(map[string]map[string]interface{}),
		Watch:      true,
		Preserve:   true,
	}

	watcher := client.Watch(options)

	go func() {
		err, more := <-watcher.Error
		if !more {
			return
		}
		t.Fatal(err)
	}()

	info := <-watcher.Done
	if info.Uploads[0].Name != "lol_cat.jpg" {
		t.Fatal("wrong file uploaded")
	}

	if !exists("./fixtures/output/resize_0_lol_cat.jpg") {
		t.Fatal("output file resize_0_lol_cat.jpg not created")
	}

	if !exists("./fixtures/output/-original_0_lol_cat.jpg") {
		t.Fatal("output file -original_0_lol_cat.jpg not created")
	}

	if exists("./fixtures/input/lol_cat.jpg") {
		t.Fatal("output file lol_cat.jpg not deleted")
	}

	go copyFile("./fixtures/mona_lisa.jpg", "./fixtures/input/mona_lisa.jpg")

	changedFile := <-watcher.Change
	if filepath.ToSlash(changedFile) != "fixtures/input/mona_lisa.jpg" {
		t.Fatal("wrong changed file name")
	}

	info = <-watcher.Done
	if info.Uploads[0].Name != "mona_lisa.jpg" {
		t.Fatal("wrong file uploaded")
	}

	if !exists("./fixtures/output/resize_0_mona_lisa.jpg") {
		t.Fatal("output file resize_0_mona_lisa.jpg not created")
	}

	if !exists("./fixtures/output/-original_0_mona_lisa.jpg") {
		t.Fatal("output file -original_0_mona_lisa.jpg not created")
	}

	if exists("./fixtures/input/mona_lisa.jpg") {
		t.Fatal("output file mona_lisa.jpg not deleted")
	}

	watcher.Stop()

}

func exists(filename string) bool {
	_, err := os.Stat(filename)
	if err == nil {
		return true
	}
	if os.IsNotExist(err) {
		return false
	}
	panic(err)
}

func copyFile(src, dst string) {

	in, err := os.Open(src)
	if err != nil {
		panic(err)
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		panic(err)
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	if err != nil {
		panic(err)
	}
}

func remove(filename string) {

	if err := os.RemoveAll(filename); err != nil {
		panic(err)
	}

}
<commit_msg>Added debugging output (to be removed)<commit_after>package transloadit

import (
	"io"
	"os"
	"path/filepath"
	"testing"
	"fmt"
)

func TestWatch(t *testing.T) {

	// Clean up testing environment
	remove("./fixtures/input/lol_cat.jpg")
	remove("./fixtures/input/mona_lisa.jpg")
	remove("./fixtures/output/resize_0_lol_cat.jpg")
	remove("./fixtures/output/-original_0_lol_cat.jpg")
	remove("./fixtures/output/resize_0_mona_lisa.jpg")
	remove("./fixtures/output/-original_0_mona_lisa.jpg")
	copyFile("./fixtures/lol_cat.jpg", "./fixtures/input/lol_cat.jpg")

	client := setup(t)

	options := &WatchOptions{
		TemplateId: "68302e702fbd11e4b5fb01f025693443",
		Input:      "./fixtures/input",
		Output:     "./fixtures/output",
		Steps:      make(map[string]map[string]interface{}),
		Watch:      true,
		Preserve:   true,
	}

	watcher := client.Watch(options)
	fmt.Println(1)
	go func() {
		err, more := <-watcher.Error
		if !more {
			return
		}
		t.Fatal(err)
	}()
	fmt.Println(2)
	info := <-watcher.Done
	if info.Uploads[0].Name != "lol_cat.jpg" {
		t.Fatal("wrong file uploaded")
	}
	fmt.Println(3)
	if !exists("./fixtures/output/resize_0_lol_cat.jpg") {
		t.Fatal("output file resize_0_lol_cat.jpg not created")
	}
	fmt.Println(4)
	if !exists("./fixtures/output/-original_0_lol_cat.jpg") {
		t.Fatal("output file -original_0_lol_cat.jpg not created")
	}
	fmt.Println(5)
	if exists("./fixtures/input/lol_cat.jpg") {
		t.Fatal("output file lol_cat.jpg not deleted")
	}
	fmt.Println(6)
	go copyFile("./fixtures/mona_lisa.jpg", "./fixtures/input/mona_lisa.jpg")
	fmt.Println(7)
	changedFile := <-watcher.Change
	if filepath.ToSlash(changedFile) != "fixtures/input/mona_lisa.jpg" {
		t.Fatal("wrong changed file name")
	}
	fmt.Println(8)
	info = <-watcher.Done
	if info.Uploads[0].Name != "mona_lisa.jpg" {
		t.Fatal("wrong file uploaded")
	}
	fmt.Println(9)
	if !exists("./fixtures/output/resize_0_mona_lisa.jpg") {
		t.Fatal("output file resize_0_mona_lisa.jpg not created")
	}
	fmt.Println(10)
	if !exists("./fixtures/output/-original_0_mona_lisa.jpg") {
		t.Fatal("output file -original_0_mona_lisa.jpg not created")
	}
	fmt.Println(11)
	if exists("./fixtures/input/mona_lisa.jpg") {
		t.Fatal("output file mona_lisa.jpg not deleted")
	}
	fmt.Println(12)
	watcher.Stop()
	fmt.Println(13)
}

func exists(filename string) bool {
	_, err := os.Stat(filename)
	if err == nil {
		return true
	}
	if os.IsNotExist(err) {
		return false
	}
	panic(err)
}

func copyFile(src, dst string) {

	in, err := os.Open(src)
	if err != nil {
		panic(err)
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		panic(err)
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	if err != nil {
		panic(err)
	}
}

func remove(filename string) {

	if err := os.RemoveAll(filename); err != nil {
		panic(err)
	}

}
<|endoftext|>


<commit_before>package terraform

import (
	"fmt"
	"log"
	"sync"

	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/depgraph"
)

// Terraform is the primary structure that is used to interact with
// Terraform from code, and can perform operations such as returning
// all resources, a resource tree, a specific resource, etc.
type Terraform struct {
	hooks     []Hook
	providers map[string]ResourceProviderFactory
}

// This is a function type used to implement a walker for the resource
// tree internally on the Terraform structure.
type genericWalkFunc func(*Resource) (map[string]string, error)

// Config is the configuration that must be given to instantiate
// a Terraform structure.
type Config struct {
	Hooks     []Hook
	Providers map[string]ResourceProviderFactory
}

// New creates a new Terraform structure, initializes resource providers
// for the given configuration, etc.
//
// Semantic checks of the entire configuration structure are done at this
// time, as well as richer checks such as verifying that the resource providers
// can be properly initialized, can be configured, etc.
func New(c *Config) (*Terraform, error) {
	return &Terraform{
		hooks:     c.Hooks,
		providers: c.Providers,
	}, nil
}

func (t *Terraform) Apply(p *Plan) (*State, error) {
	// Make sure we're working with a plan that doesn't have null pointers
	// everywhere, and is instead just empty otherwise.
	p.init()

	g, err := Graph(&GraphOpts{
		Config:    p.Config,
		Diff:      p.Diff,
		Providers: t.providers,
		State:     p.State,
	})
	if err != nil {
		return nil, err
	}

	return t.apply(g, p)
}

func (t *Terraform) Plan(opts *PlanOpts) (*Plan, error) {
	g, err := Graph(&GraphOpts{
		Config:    opts.Config,
		Providers: t.providers,
		State:     opts.State,
	})
	if err != nil {
		return nil, err
	}

	return t.plan(g, opts)
}

// Refresh goes through all the resources in the state and refreshes them
// to their latest status.
func (t *Terraform) Refresh(c *config.Config, s *State) (*State, error) {
	g, err := Graph(&GraphOpts{
		Config:    c,
		Providers: t.providers,
		State:     s,
	})
	if err != nil {
		return s, err
	}

	return t.refresh(g)
}

func (t *Terraform) apply(
	g *depgraph.Graph,
	p *Plan) (*State, error) {
	s := new(State)
	err := g.Walk(t.applyWalkFn(s, p))
	return s, err
}

func (t *Terraform) plan(g *depgraph.Graph, opts *PlanOpts) (*Plan, error) {
	p := &Plan{
		Config: opts.Config,
		Vars:   opts.Vars,
		State:  opts.State,
	}
	err := g.Walk(t.planWalkFn(p, opts))
	return p, err
}

func (t *Terraform) refresh(g *depgraph.Graph) (*State, error) {
	s := new(State)
	err := g.Walk(t.refreshWalkFn(s))
	return s, err
}

func (t *Terraform) refreshWalkFn(result *State) depgraph.WalkFunc {
	var l sync.Mutex

	// Initialize the result so we don't have to nil check everywhere
	result.init()

	cb := func(r *Resource) (map[string]string, error) {
		for _, h := range t.hooks {
			// TODO: return value
			h.PreRefresh(r.Id, r.State)
		}

		rs, err := r.Provider.Refresh(r.State)
		if err != nil {
			return nil, err
		}
		if rs == nil {
			rs = new(ResourceState)
		}

		// Fix the type to be the type we have
		rs.Type = r.State.Type

		l.Lock()
		result.Resources[r.Id] = rs
		l.Unlock()

		for _, h := range t.hooks {
			// TODO: return value
			h.PostRefresh(r.Id, rs)
		}

		return nil, nil
	}

	return t.genericWalkFn(nil, cb)
}

func (t *Terraform) applyWalkFn(
	result *State,
	p *Plan) depgraph.WalkFunc {
	var l sync.Mutex

	// Initialize the result
	result.init()

	cb := func(r *Resource) (map[string]string, error) {
		diff := r.Diff
		if diff.Empty() {
			return r.Vars(), nil
		}

		if !diff.Destroy {
			var err error
			diff, err = r.Provider.Diff(r.State, r.Config)
			if err != nil {
				return nil, err
			}
		}

		// TODO(mitchellh): we need to verify the diff doesn't change
		// anything and that the diff has no computed values (pre-computed)

		for _, h := range t.hooks {
			// TODO: return value
			h.PreApply(r.Id, r.State, diff)
		}

		// With the completed diff, apply!
		rs, err := r.Provider.Apply(r.State, diff)
		if err != nil {
			return nil, err
		}

		// Make sure the result is instantiated
		if rs == nil {
			rs = new(ResourceState)
		}

		// Force the resource state type to be our type
		rs.Type = r.State.Type

		// If no state was returned, then no variables were updated so
		// just return.
		if rs == nil {
			return nil, nil
		}

		var errs []error
		for ak, av := range rs.Attributes {
			// If the value is the unknown variable value, then it is an error.
			// In this case we record the error and remove it from the state
			if av == config.UnknownVariableValue {
				errs = append(errs, fmt.Errorf(
					"Attribute with unknown value: %s", ak))
				delete(rs.Attributes, ak)
			}
		}

		// Update the resulting diff
		l.Lock()
		result.Resources[r.Id] = rs
		l.Unlock()

		// Update the state for the resource itself
		r.State = rs

		for _, h := range t.hooks {
			// TODO: return value
			h.PostApply(r.Id, r.State)
		}

		// Determine the new state and update variables
		err = nil
		if len(errs) > 0 {
			err = &MultiError{Errors: errs}
		}

		return r.Vars(), err
	}

	return t.genericWalkFn(p.Vars, cb)
}

func (t *Terraform) planWalkFn(result *Plan, opts *PlanOpts) depgraph.WalkFunc {
	var l sync.Mutex

	// Initialize the result
	result.init()

	cb := func(r *Resource) (map[string]string, error) {
		var diff *ResourceDiff

		for _, h := range t.hooks {
			// TODO: return value
			h.PreDiff(r.Id, r.State)
		}

		if opts.Destroy {
			if r.State.ID != "" {
				log.Printf("[DEBUG] %s: Making for destroy", r.Id)
				diff = &ResourceDiff{Destroy: true}
			} else {
				log.Printf("[DEBUG] %s: Not marking for destroy, no ID", r.Id)
			}
		} else if r.Config == nil {
			log.Printf("[DEBUG] %s: Orphan, marking for destroy", r.Id)

			// This is an orphan (no config), so we mark it to be destroyed
			diff = &ResourceDiff{Destroy: true}
		} else {
			log.Printf("[DEBUG] %s: Executing diff", r.Id)

			// Get a diff from the newest state
			var err error
			diff, err = r.Provider.Diff(r.State, r.Config)
			if err != nil {
				return nil, err
			}
		}

		l.Lock()
		if !diff.Empty() {
			result.Diff.Resources[r.Id] = diff
		}
		l.Unlock()

		for _, h := range t.hooks {
			// TODO: return value
			h.PostDiff(r.Id, diff)
		}

		// Determine the new state and update variables
		if !diff.Empty() {
			r.State = r.State.MergeDiff(diff)
		}

		return r.Vars(), nil
	}

	return t.genericWalkFn(opts.Vars, cb)
}

func (t *Terraform) genericWalkFn(
	invars map[string]string,
	cb genericWalkFunc) depgraph.WalkFunc {
	var l sync.Mutex

	// Initialize the variables for application
	vars := make(map[string]string)
	for k, v := range invars {
		vars[fmt.Sprintf("var.%s", k)] = v
	}

	return func(n *depgraph.Noun) error {
		// If it is the root node, ignore
		if n.Name == GraphRootNode {
			return nil
		}

		switch m := n.Meta.(type) {
		case *GraphNodeResource:
		case *GraphNodeResourceProvider:
			var rc *ResourceConfig
			if m.Config != nil {
				if err := m.Config.RawConfig.Interpolate(vars); err != nil {
					panic(err)
				}
				rc = NewResourceConfig(m.Config.RawConfig)
			}

			for k, p := range m.Providers {
				log.Printf("[INFO] Configuring provider: %s", k)
				err := p.Configure(rc)
				if err != nil {
					return err
				}
			}

			return nil
		}

		rn := n.Meta.(*GraphNodeResource)
		if len(vars) > 0 && rn.Config != nil {
			if err := rn.Config.RawConfig.Interpolate(vars); err != nil {
				panic(fmt.Sprintf("Interpolate error: %s", err))
			}

			// Force the config to be set later
			rn.Resource.Config = nil
		}

		// Make sure that at least some resource configuration is set
		if !rn.Orphan {
			if rn.Resource.Config == nil {
				if rn.Config == nil {
					rn.Resource.Config = new(ResourceConfig)
				} else {
					rn.Resource.Config = NewResourceConfig(rn.Config.RawConfig)
				}
			}
		} else {
			rn.Resource.Config = nil
		}

		// Call the callack
		log.Printf("[INFO] Walking: %s", rn.Resource.Id)
		newVars, err := cb(rn.Resource)
		if err != nil {
			return err
		}

		if len(newVars) > 0 {
			// Acquire a lock since this function is called in parallel
			l.Lock()
			defer l.Unlock()

			// Update variables
			for k, v := range newVars {
				vars[k] = v
			}
		}

		return nil
	}
}
<commit_msg>terraform: fix data race<commit_after>package terraform

import (
	"fmt"
	"log"
	"sync"

	"github.com/hashicorp/terraform/config"
	"github.com/hashicorp/terraform/depgraph"
)

// Terraform is the primary structure that is used to interact with
// Terraform from code, and can perform operations such as returning
// all resources, a resource tree, a specific resource, etc.
type Terraform struct {
	hooks     []Hook
	providers map[string]ResourceProviderFactory
}

// This is a function type used to implement a walker for the resource
// tree internally on the Terraform structure.
type genericWalkFunc func(*Resource) (map[string]string, error)

// Config is the configuration that must be given to instantiate
// a Terraform structure.
type Config struct {
	Hooks     []Hook
	Providers map[string]ResourceProviderFactory
}

// New creates a new Terraform structure, initializes resource providers
// for the given configuration, etc.
//
// Semantic checks of the entire configuration structure are done at this
// time, as well as richer checks such as verifying that the resource providers
// can be properly initialized, can be configured, etc.
func New(c *Config) (*Terraform, error) {
	return &Terraform{
		hooks:     c.Hooks,
		providers: c.Providers,
	}, nil
}

func (t *Terraform) Apply(p *Plan) (*State, error) {
	// Make sure we're working with a plan that doesn't have null pointers
	// everywhere, and is instead just empty otherwise.
	p.init()

	g, err := Graph(&GraphOpts{
		Config:    p.Config,
		Diff:      p.Diff,
		Providers: t.providers,
		State:     p.State,
	})
	if err != nil {
		return nil, err
	}

	return t.apply(g, p)
}

func (t *Terraform) Plan(opts *PlanOpts) (*Plan, error) {
	g, err := Graph(&GraphOpts{
		Config:    opts.Config,
		Providers: t.providers,
		State:     opts.State,
	})
	if err != nil {
		return nil, err
	}

	return t.plan(g, opts)
}

// Refresh goes through all the resources in the state and refreshes them
// to their latest status.
func (t *Terraform) Refresh(c *config.Config, s *State) (*State, error) {
	g, err := Graph(&GraphOpts{
		Config:    c,
		Providers: t.providers,
		State:     s,
	})
	if err != nil {
		return s, err
	}

	return t.refresh(g)
}

func (t *Terraform) apply(
	g *depgraph.Graph,
	p *Plan) (*State, error) {
	s := new(State)
	err := g.Walk(t.applyWalkFn(s, p))
	return s, err
}

func (t *Terraform) plan(g *depgraph.Graph, opts *PlanOpts) (*Plan, error) {
	p := &Plan{
		Config: opts.Config,
		Vars:   opts.Vars,
		State:  opts.State,
	}
	err := g.Walk(t.planWalkFn(p, opts))
	return p, err
}

func (t *Terraform) refresh(g *depgraph.Graph) (*State, error) {
	s := new(State)
	err := g.Walk(t.refreshWalkFn(s))
	return s, err
}

func (t *Terraform) refreshWalkFn(result *State) depgraph.WalkFunc {
	var l sync.Mutex

	// Initialize the result so we don't have to nil check everywhere
	result.init()

	cb := func(r *Resource) (map[string]string, error) {
		for _, h := range t.hooks {
			// TODO: return value
			h.PreRefresh(r.Id, r.State)
		}

		rs, err := r.Provider.Refresh(r.State)
		if err != nil {
			return nil, err
		}
		if rs == nil {
			rs = new(ResourceState)
		}

		// Fix the type to be the type we have
		rs.Type = r.State.Type

		l.Lock()
		result.Resources[r.Id] = rs
		l.Unlock()

		for _, h := range t.hooks {
			// TODO: return value
			h.PostRefresh(r.Id, rs)
		}

		return nil, nil
	}

	return t.genericWalkFn(nil, cb)
}

func (t *Terraform) applyWalkFn(
	result *State,
	p *Plan) depgraph.WalkFunc {
	var l sync.Mutex

	// Initialize the result
	result.init()

	cb := func(r *Resource) (map[string]string, error) {
		diff := r.Diff
		if diff.Empty() {
			return r.Vars(), nil
		}

		if !diff.Destroy {
			var err error
			diff, err = r.Provider.Diff(r.State, r.Config)
			if err != nil {
				return nil, err
			}
		}

		// TODO(mitchellh): we need to verify the diff doesn't change
		// anything and that the diff has no computed values (pre-computed)

		for _, h := range t.hooks {
			// TODO: return value
			h.PreApply(r.Id, r.State, diff)
		}

		// With the completed diff, apply!
		rs, err := r.Provider.Apply(r.State, diff)
		if err != nil {
			return nil, err
		}

		// Make sure the result is instantiated
		if rs == nil {
			rs = new(ResourceState)
		}

		// Force the resource state type to be our type
		rs.Type = r.State.Type

		// If no state was returned, then no variables were updated so
		// just return.
		if rs == nil {
			return nil, nil
		}

		var errs []error
		for ak, av := range rs.Attributes {
			// If the value is the unknown variable value, then it is an error.
			// In this case we record the error and remove it from the state
			if av == config.UnknownVariableValue {
				errs = append(errs, fmt.Errorf(
					"Attribute with unknown value: %s", ak))
				delete(rs.Attributes, ak)
			}
		}

		// Update the resulting diff
		l.Lock()
		result.Resources[r.Id] = rs
		l.Unlock()

		// Update the state for the resource itself
		r.State = rs

		for _, h := range t.hooks {
			// TODO: return value
			h.PostApply(r.Id, r.State)
		}

		// Determine the new state and update variables
		err = nil
		if len(errs) > 0 {
			err = &MultiError{Errors: errs}
		}

		return r.Vars(), err
	}

	return t.genericWalkFn(p.Vars, cb)
}

func (t *Terraform) planWalkFn(result *Plan, opts *PlanOpts) depgraph.WalkFunc {
	var l sync.Mutex

	// Initialize the result
	result.init()

	cb := func(r *Resource) (map[string]string, error) {
		var diff *ResourceDiff

		for _, h := range t.hooks {
			// TODO: return value
			h.PreDiff(r.Id, r.State)
		}

		if opts.Destroy {
			if r.State.ID != "" {
				log.Printf("[DEBUG] %s: Making for destroy", r.Id)
				diff = &ResourceDiff{Destroy: true}
			} else {
				log.Printf("[DEBUG] %s: Not marking for destroy, no ID", r.Id)
			}
		} else if r.Config == nil {
			log.Printf("[DEBUG] %s: Orphan, marking for destroy", r.Id)

			// This is an orphan (no config), so we mark it to be destroyed
			diff = &ResourceDiff{Destroy: true}
		} else {
			log.Printf("[DEBUG] %s: Executing diff", r.Id)

			// Get a diff from the newest state
			var err error
			diff, err = r.Provider.Diff(r.State, r.Config)
			if err != nil {
				return nil, err
			}
		}

		l.Lock()
		if !diff.Empty() {
			result.Diff.Resources[r.Id] = diff
		}
		l.Unlock()

		for _, h := range t.hooks {
			// TODO: return value
			h.PostDiff(r.Id, diff)
		}

		// Determine the new state and update variables
		if !diff.Empty() {
			r.State = r.State.MergeDiff(diff)
		}

		return r.Vars(), nil
	}

	return t.genericWalkFn(opts.Vars, cb)
}

func (t *Terraform) genericWalkFn(
	invars map[string]string,
	cb genericWalkFunc) depgraph.WalkFunc {
	var l sync.RWMutex

	// Initialize the variables for application
	vars := make(map[string]string)
	for k, v := range invars {
		vars[fmt.Sprintf("var.%s", k)] = v
	}

	return func(n *depgraph.Noun) error {
		// If it is the root node, ignore
		if n.Name == GraphRootNode {
			return nil
		}

		switch m := n.Meta.(type) {
		case *GraphNodeResource:
		case *GraphNodeResourceProvider:
			var rc *ResourceConfig
			if m.Config != nil {
				if err := m.Config.RawConfig.Interpolate(vars); err != nil {
					panic(err)
				}
				rc = NewResourceConfig(m.Config.RawConfig)
			}

			for k, p := range m.Providers {
				log.Printf("[INFO] Configuring provider: %s", k)
				err := p.Configure(rc)
				if err != nil {
					return err
				}
			}

			return nil
		}

		rn := n.Meta.(*GraphNodeResource)

		l.RLock()
		if len(vars) > 0 && rn.Config != nil {
			if err := rn.Config.RawConfig.Interpolate(vars); err != nil {
				panic(fmt.Sprintf("Interpolate error: %s", err))
			}

			// Force the config to be set later
			rn.Resource.Config = nil
		}
		l.RUnlock()

		// Make sure that at least some resource configuration is set
		if !rn.Orphan {
			if rn.Resource.Config == nil {
				if rn.Config == nil {
					rn.Resource.Config = new(ResourceConfig)
				} else {
					rn.Resource.Config = NewResourceConfig(rn.Config.RawConfig)
				}
			}
		} else {
			rn.Resource.Config = nil
		}

		// Call the callack
		log.Printf("[INFO] Walking: %s", rn.Resource.Id)
		newVars, err := cb(rn.Resource)
		if err != nil {
			return err
		}

		if len(newVars) > 0 {
			// Acquire a lock since this function is called in parallel
			l.Lock()
			defer l.Unlock()

			// Update variables
			for k, v := range newVars {
				vars[k] = v
			}
		}

		return nil
	}
}
<|endoftext|>