<commit_before>package client

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/distribution/digest"
)

type httpLayerUpload struct {
	repo   distribution.Repository
	client *http.Client

	uuid      string
	startedAt time.Time

	location string // always the last value of the location header.
	offset   int64
	closed   bool
}

func (hlu *httpLayerUpload) handleErrorResponse(resp *http.Response) error {
	if resp.StatusCode == http.StatusNotFound {
		return &BlobUploadNotFoundError{Location: hlu.location}
	}
	return handleErrorResponse(resp)
}

func (hlu *httpLayerUpload) ReadFrom(r io.Reader) (n int64, err error) {
	req, err := http.NewRequest("PATCH", hlu.location, r)
	if err != nil {
		return 0, err
	}
	defer req.Body.Close()

	resp, err := hlu.client.Do(req)
	if err != nil {
		return 0, err
	}

	if resp.StatusCode != http.StatusAccepted {
		return 0, hlu.handleErrorResponse(resp)
	}

	// TODO(dmcgowan): Validate headers
	hlu.uuid = resp.Header.Get("Docker-Upload-UUID")
	hlu.location, err = sanitizeLocation(resp.Header.Get("Location"), hlu.location)
	if err != nil {
		return 0, err
	}
	rng := resp.Header.Get("Range")
	var start, end int64
	if n, err := fmt.Sscanf(rng, "%d-%d", &start, &end); err != nil {
		return 0, err
	} else if n != 2 || end < start {
		return 0, fmt.Errorf("bad range format: %s", rng)
	}

	return (end - start + 1), nil

}

func (hlu *httpLayerUpload) Write(p []byte) (n int, err error) {
	req, err := http.NewRequest("PATCH", hlu.location, bytes.NewReader(p))
	if err != nil {
		return 0, err
	}
	req.Header.Set("Content-Range", fmt.Sprintf("%d-%d", hlu.offset, hlu.offset+int64(len(p)-1)))
	req.Header.Set("Content-Length", fmt.Sprintf("%d", len(p)))
	req.Header.Set("Content-Type", "application/octet-stream")

	resp, err := hlu.client.Do(req)
	if err != nil {
		return 0, err
	}

	if resp.StatusCode != http.StatusAccepted {
		return 0, hlu.handleErrorResponse(resp)
	}

	// TODO(dmcgowan): Validate headers
	hlu.uuid = resp.Header.Get("Docker-Upload-UUID")
	hlu.location, err = sanitizeLocation(resp.Header.Get("Location"), hlu.location)
	if err != nil {
		return 0, err
	}
	rng := resp.Header.Get("Range")
	var start, end int
	if n, err := fmt.Sscanf(rng, "%d-%d", &start, &end); err != nil {
		return 0, err
	} else if n != 2 || end < start {
		return 0, fmt.Errorf("bad range format: %s", rng)
	}

	return (end - start + 1), nil

}

func (hlu *httpLayerUpload) Seek(offset int64, whence int) (int64, error) {
	newOffset := hlu.offset

	switch whence {
	case os.SEEK_CUR:
		newOffset += int64(offset)
	case os.SEEK_END:
		return newOffset, errors.New("Cannot seek from end on incomplete upload")
	case os.SEEK_SET:
		newOffset = int64(offset)
	}

	hlu.offset = newOffset

	return hlu.offset, nil
}

func (hlu *httpLayerUpload) UUID() string {
	return hlu.uuid
}

func (hlu *httpLayerUpload) StartedAt() time.Time {
	return hlu.startedAt
}

func (hlu *httpLayerUpload) Finish(digest digest.Digest) (distribution.Layer, error) {
	// TODO(dmcgowan): Check if already finished, if so just fetch
	req, err := http.NewRequest("PUT", hlu.location, nil)
	if err != nil {
		return nil, err
	}

	values := req.URL.Query()
	values.Set("digest", digest.String())
	req.URL.RawQuery = values.Encode()

	resp, err := hlu.client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusCreated {
		return nil, hlu.handleErrorResponse(resp)
	}

	return hlu.repo.Layers().Fetch(digest)
}

func (hlu *httpLayerUpload) Cancel() error {
	panic("not implemented")
}

func (hlu *httpLayerUpload) Close() error {
	hlu.closed = true
	return nil
}
<commit_msg>Update ReadFrom to wrap reader in NopCloser<commit_after>package client

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"time"

	"github.com/docker/distribution"
	"github.com/docker/distribution/digest"
)

type httpLayerUpload struct {
	repo   distribution.Repository
	client *http.Client

	uuid      string
	startedAt time.Time

	location string // always the last value of the location header.
	offset   int64
	closed   bool
}

func (hlu *httpLayerUpload) handleErrorResponse(resp *http.Response) error {
	if resp.StatusCode == http.StatusNotFound {
		return &BlobUploadNotFoundError{Location: hlu.location}
	}
	return handleErrorResponse(resp)
}

func (hlu *httpLayerUpload) ReadFrom(r io.Reader) (n int64, err error) {
	req, err := http.NewRequest("PATCH", hlu.location, ioutil.NopCloser(r))
	if err != nil {
		return 0, err
	}
	defer req.Body.Close()

	resp, err := hlu.client.Do(req)
	if err != nil {
		return 0, err
	}

	if resp.StatusCode != http.StatusAccepted {
		return 0, hlu.handleErrorResponse(resp)
	}

	// TODO(dmcgowan): Validate headers
	hlu.uuid = resp.Header.Get("Docker-Upload-UUID")
	hlu.location, err = sanitizeLocation(resp.Header.Get("Location"), hlu.location)
	if err != nil {
		return 0, err
	}
	rng := resp.Header.Get("Range")
	var start, end int64
	if n, err := fmt.Sscanf(rng, "%d-%d", &start, &end); err != nil {
		return 0, err
	} else if n != 2 || end < start {
		return 0, fmt.Errorf("bad range format: %s", rng)
	}

	return (end - start + 1), nil

}

func (hlu *httpLayerUpload) Write(p []byte) (n int, err error) {
	req, err := http.NewRequest("PATCH", hlu.location, bytes.NewReader(p))
	if err != nil {
		return 0, err
	}
	req.Header.Set("Content-Range", fmt.Sprintf("%d-%d", hlu.offset, hlu.offset+int64(len(p)-1)))
	req.Header.Set("Content-Length", fmt.Sprintf("%d", len(p)))
	req.Header.Set("Content-Type", "application/octet-stream")

	resp, err := hlu.client.Do(req)
	if err != nil {
		return 0, err
	}

	if resp.StatusCode != http.StatusAccepted {
		return 0, hlu.handleErrorResponse(resp)
	}

	// TODO(dmcgowan): Validate headers
	hlu.uuid = resp.Header.Get("Docker-Upload-UUID")
	hlu.location, err = sanitizeLocation(resp.Header.Get("Location"), hlu.location)
	if err != nil {
		return 0, err
	}
	rng := resp.Header.Get("Range")
	var start, end int
	if n, err := fmt.Sscanf(rng, "%d-%d", &start, &end); err != nil {
		return 0, err
	} else if n != 2 || end < start {
		return 0, fmt.Errorf("bad range format: %s", rng)
	}

	return (end - start + 1), nil

}

func (hlu *httpLayerUpload) Seek(offset int64, whence int) (int64, error) {
	newOffset := hlu.offset

	switch whence {
	case os.SEEK_CUR:
		newOffset += int64(offset)
	case os.SEEK_END:
		return newOffset, errors.New("Cannot seek from end on incomplete upload")
	case os.SEEK_SET:
		newOffset = int64(offset)
	}

	hlu.offset = newOffset

	return hlu.offset, nil
}

func (hlu *httpLayerUpload) UUID() string {
	return hlu.uuid
}

func (hlu *httpLayerUpload) StartedAt() time.Time {
	return hlu.startedAt
}

func (hlu *httpLayerUpload) Finish(digest digest.Digest) (distribution.Layer, error) {
	// TODO(dmcgowan): Check if already finished, if so just fetch
	req, err := http.NewRequest("PUT", hlu.location, nil)
	if err != nil {
		return nil, err
	}

	values := req.URL.Query()
	values.Set("digest", digest.String())
	req.URL.RawQuery = values.Encode()

	resp, err := hlu.client.Do(req)
	if err != nil {
		return nil, err
	}

	if resp.StatusCode != http.StatusCreated {
		return nil, hlu.handleErrorResponse(resp)
	}

	return hlu.repo.Layers().Fetch(digest)
}

func (hlu *httpLayerUpload) Cancel() error {
	panic("not implemented")
}

func (hlu *httpLayerUpload) Close() error {
	hlu.closed = true
	return nil
}
<|endoftext|>


<commit_before>// Package bitutils provides a collection of utilities to deal with bits.
package bitutils

import (
	"fmt"
	"strconv"
)

// W is the length of a machine word.
const W = 64

// Magic constants.
const (
	Lsh2 = 0x5555555555555555
	Lsh4 = 0x3333333333333333
	Lsh8 = 0x0f0f0f0f0f0f0f0f

	Msh2 = 0xaaaaaaaaaaaaaaaa
	Msh4 = 0xcccccccccccccccc
	Msh8 = 0xf0f0f0f0f0f0f0f0

	Lsb2 = 0x5555555555555555
	Lsb4 = 0x1111111111111111
	Lsb8 = 0x0101010101010101

	Msb2 = 0xaaaaaaaaaaaaaaaa
	Msb4 = 0x8888888888888888
	Msb8 = 0x8080808080808080
)

// Word represents a 64-bit binary string.
type Word uint64

var (
	Pos  [W + 1]Word // Pos[i] has a 1 only at i.
	PosC [W + 1]Word // PosC[i] has a 0 only at i.
	LsbN [W + 1]Word // LsbN[i] has 1s in its i LSBs.
	MsbN [W + 1]Word // MsbN[i] has 1s in its i MSBs.
)

func init() {
	for i := 0; i < len(Pos); i++ {
		Pos[i] = Word(1) << uint(i)
		PosC[i] = ^Pos[i]
		LsbN[i] = Pos[i] - 1
		MsbN[i] = LsbN[i] << uint(W-i)
	}
}

// ParseWord returns a Word from a string.
func ParseWord(s string) (Word, error) {
	w, err := strconv.ParseUint(s, 2, 64)
	return Word(w), err
}

// String returns binary string w[0]w[1]...w[63].
func (w Word) String() string {
	return fmt.Sprintf("%064b", w)
}

// Count1 returns the number of ones contained in w.
func (w Word) Count1() int {
	w -= (w >> 1) & Lsh2
	w = (w & Lsh4) + ((w >> 2) & Lsh4)
	w = (w + (w >> 4)) & Lsh8
	return int((w * Lsb8) >> 56)
}

// Count0 returns the number of zeros contained in w.
func (w Word) Count0() int {
	return (^w).Count1()
}

// Count returns the number of b[0]'s contained in w.
func (w Word) Count(b int) int {
	w = w ^ (^Word(0) + Word(b))
	return w.Count1()
}

// Get returns w[i].
func (w Word) Get(i int) Word {
	w = w >> uint(i)
	return w & Pos[0]
}

// Set1 sets w[i] to 1.
func (w Word) Set1(i int) Word {
	return w | Pos[i]
}

// Set0 sets w[i] to 0.
func (w Word) Set0(i int) Word {
	return w & PosC[i]
}

// Flip flips w[i].
func (w Word) Flip(i int) Word {
	return w ^ Pos[i]
}

// Least1 returns a word that indicates the least 1 in w.
func (w Word) Least1() Word {
	if w == 0 {
		return 0
	}
	w = ((w - 1) ^ w) & w
	return w
}

// LeastIndex1 returns the index of the least 1 in w if exists and -1
// otherwise.
func (w Word) LeastIndex1() int {
	if w == 0 {
		return -1
	}
	w = (w - 1) ^ w
	return w.Count1() - 1
}

// Rank1 returns the number of ones in w[0]...w[i].
func (w Word) Rank1(i int) int {
	return (LsbN[i+1] & w).Count1()
}

// Rank0 returns the number of zeros in w[0]...w[i].
func (w Word) Rank0(i int) int {
	return (MsbN[W-i-1] | w).Count0()
}

func (w Word) zcmp8() Word {
	w = w | ((w | Msb8) - Lsb8)
	return (w & Msb8) >> 7
}

func (w Word) leq8(v Word) Word {
	w = (((v | Msb8) - (w & ^Word(Lsb8))) ^ w) ^ v
	return (w & Msb8) >> 7
}

// Select1 returns the ith 1 in w.
func (w Word) Select1(i int) int {
	s := w - ((w & Msb2) >> 1)
	s = (s & Lsh4) + ((s >> 2) & Lsh4)
	s = ((s + (s >> 4)) & Lsh8) * Lsb8
	b := ((s.leq8(Word(i)*Lsb8) * Lsb8) >> 53) & ^Word(0x0111)
	l := Word(i) - (((s << 8) >> b) & 0xff)
	s = ((((w >> b) & 0xff) * Lsb8) & 0x8040201008040201).zcmp8() * Lsb8
	if w = b + ((s.leq8(l*Lsb8) * Lsb8) >> 56); w != 0x48 {
		return int(w)
	} else {
		return -1
	}
}

// Select0 returns the ith 0 in w.
func (w Word) Select0(i int) int {
	w = ^w
	return w.Select1(i)
}
<commit_msg>Refactor Count<commit_after>// Package bitutils provides a collection of utilities to deal with bits.
package bitutils

import (
	"fmt"
	"strconv"
)

// W is the length of a machine word.
const W = 64

// Magic constants.
const (
	Lsh2 = 0x5555555555555555
	Lsh4 = 0x3333333333333333
	Lsh8 = 0x0f0f0f0f0f0f0f0f

	Msh2 = 0xaaaaaaaaaaaaaaaa
	Msh4 = 0xcccccccccccccccc
	Msh8 = 0xf0f0f0f0f0f0f0f0

	Lsb2 = 0x5555555555555555
	Lsb4 = 0x1111111111111111
	Lsb8 = 0x0101010101010101

	Msb2 = 0xaaaaaaaaaaaaaaaa
	Msb4 = 0x8888888888888888
	Msb8 = 0x8080808080808080
)

// Word represents a 64-bit binary string.
type Word uint64

var (
	Pos  [W + 1]Word // Pos[i] has a 1 only at i.
	PosC [W + 1]Word // PosC[i] has a 0 only at i.
	LsbN [W + 1]Word // LsbN[i] has 1s in its i LSBs.
	MsbN [W + 1]Word // MsbN[i] has 1s in its i MSBs.
)

func init() {
	for i := 0; i < len(Pos); i++ {
		Pos[i] = Word(1) << uint(i)
		PosC[i] = ^Pos[i]
		LsbN[i] = Pos[i] - 1
		MsbN[i] = LsbN[i] << uint(W-i)
	}
}

// ParseWord returns a Word from a string.
func ParseWord(s string) (Word, error) {
	w, err := strconv.ParseUint(s, 2, 64)
	return Word(w), err
}

// String returns binary string w[0]w[1]...w[63].
func (w Word) String() string {
	return fmt.Sprintf("%064b", w)
}

// Count1 returns the number of ones contained in w.
func (w Word) Count1() int {
	w -= (w >> 1) & Lsh2
	w = (w & Lsh4) + ((w >> 2) & Lsh4)
	w = (w + (w >> 4)) & Lsh8
	return int((w * Lsb8) >> 56)
}

// Count0 returns the number of zeros contained in w.
func (w Word) Count0() int {
	return (^w).Count1()
}

// Count returns the number of b[0]'s contained in w.
func (w Word) Count(b int) int {
	return ((LsbN[W] + Word(b)) ^ w).Count1()
}

// Get returns w[i].
func (w Word) Get(i int) Word {
	w = w >> uint(i)
	return w & Pos[0]
}

// Set1 sets w[i] to 1.
func (w Word) Set1(i int) Word {
	return w | Pos[i]
}

// Set0 sets w[i] to 0.
func (w Word) Set0(i int) Word {
	return w & PosC[i]
}

// Flip flips w[i].
func (w Word) Flip(i int) Word {
	return w ^ Pos[i]
}

// Least1 returns a word that indicates the least 1 in w.
func (w Word) Least1() Word {
	if w == 0 {
		return 0
	}
	w = ((w - 1) ^ w) & w
	return w
}

// LeastIndex1 returns the index of the least 1 in w if exists and -1
// otherwise.
func (w Word) LeastIndex1() int {
	if w == 0 {
		return -1
	}
	w = (w - 1) ^ w
	return w.Count1() - 1
}

// Rank1 returns the number of ones in w[0]...w[i].
func (w Word) Rank1(i int) int {
	return (LsbN[i+1] & w).Count1()
}

// Rank0 returns the number of zeros in w[0]...w[i].
func (w Word) Rank0(i int) int {
	return (MsbN[W-i-1] | w).Count0()
}

func (w Word) zcmp8() Word {
	w = w | ((w | Msb8) - Lsb8)
	return (w & Msb8) >> 7
}

func (w Word) leq8(v Word) Word {
	w = (((v | Msb8) - (w & ^Word(Lsb8))) ^ w) ^ v
	return (w & Msb8) >> 7
}

// Select1 returns the ith 1 in w.
func (w Word) Select1(i int) int {
	s := w - ((w & Msb2) >> 1)
	s = (s & Lsh4) + ((s >> 2) & Lsh4)
	s = ((s + (s >> 4)) & Lsh8) * Lsb8
	b := ((s.leq8(Word(i)*Lsb8) * Lsb8) >> 53) & ^Word(0x0111)
	l := Word(i) - (((s << 8) >> b) & 0xff)
	s = ((((w >> b) & 0xff) * Lsb8) & 0x8040201008040201).zcmp8() * Lsb8
	if w = b + ((s.leq8(l*Lsb8) * Lsb8) >> 56); w != 0x48 {
		return int(w)
	} else {
		return -1
	}
}

// Select0 returns the ith 0 in w.
func (w Word) Select0(i int) int {
	w = ^w
	return w.Select1(i)
}
<|endoftext|>


<commit_before>package infrastructure

import (
	"log"
	"os"

	"github.com/marcusolsson/goddd/domain/location"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type locationRepositoryMongoDB struct {
}

func (r *locationRepositoryMongoDB) Store(l location.Location) {
	session, err := mgo.Dial(os.Getenv("MONGOHQ_URL"))

	if err != nil {
		panic(err)
	}
	defer session.Close()

	c := session.DB("app30695645").C("locations")
	err = c.Insert(&l)

	if err != nil {
		log.Fatal(err)
	}
}

func (r *locationRepositoryMongoDB) Find(locode location.UNLocode) (location.Location, error) {
	return location.Location{}, nil
}

func (r *locationRepositoryMongoDB) FindAll() []location.Location {
	session, err := mgo.Dial(os.Getenv("MONGOHQ_URL"))

	if err != nil {
		panic(err)
	}
	defer session.Close()

	c := session.DB("app30695645").C("locations")

	var result []location.Location
	err = c.Find(bson.M{}).All(&result)

	if err != nil {
		log.Fatal(err)
	}

	return result
}

func NewLocationRepositoryMongoDB() location.LocationRepository {
	session, err := mgo.Dial(os.Getenv("MONGOHQ_URL"))

	if err != nil {
		panic(err)
	}
	defer session.Close()

	c := session.DB("app30695645").C("locations")
	err = c.EnsureIndexKey("unlocode")

	if err != nil {
		panic(err)
	}

	r := &locationRepositoryMongoDB{}

	r.Store(location.Stockholm)
	r.Store(location.Hamburg)
	r.Store(location.Chicago)

	return r
}
<commit_msg>Ensure index is unique.<commit_after>package infrastructure

import (
	"log"
	"os"

	"github.com/marcusolsson/goddd/domain/location"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type locationRepositoryMongoDB struct {
}

func (r *locationRepositoryMongoDB) Store(l location.Location) {
	session, err := mgo.Dial(os.Getenv("MONGOHQ_URL"))

	if err != nil {
		panic(err)
	}
	defer session.Close()

	c := session.DB("app30695645").C("locations")
	err = c.Insert(&l)

	if err != nil {
		log.Fatal(err)
	}
}

func (r *locationRepositoryMongoDB) Find(locode location.UNLocode) (location.Location, error) {
	return location.Location{}, nil
}

func (r *locationRepositoryMongoDB) FindAll() []location.Location {
	session, err := mgo.Dial(os.Getenv("MONGOHQ_URL"))

	if err != nil {
		panic(err)
	}
	defer session.Close()

	c := session.DB("app30695645").C("locations")

	var result []location.Location
	err = c.Find(bson.M{}).All(&result)

	if err != nil {
		log.Fatal(err)
	}

	return result
}

func NewLocationRepositoryMongoDB() location.LocationRepository {
	session, err := mgo.Dial(os.Getenv("MONGOHQ_URL"))

	if err != nil {
		panic(err)
	}
	defer session.Close()

	c := session.DB("app30695645").C("locations")

	index := mgo.Index{
		Key:    []string{"unlocode"},
		Unique: true,
	}
	err = c.EnsureIndex(index)

	if err != nil {
		panic(err)
	}

	r := &locationRepositoryMongoDB{}

	r.Store(location.Stockholm)
	r.Store(location.Hamburg)
	r.Store(location.Chicago)

	return r
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"

	app "github.com/luistm/banksaurus/cmd/bscli/application"
)

func version(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, app.Version+"\n")
}

func main() {
	http.HandleFunc("/version", version)
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<commit_msg>Errors explicitly handled<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"

	app "github.com/luistm/banksaurus/cmd/bscli/application"
)

func version(w http.ResponseWriter, _ *http.Request) {
	_, _ = fmt.Fprintf(w, app.Version+"\n")
}

func main() {
	http.HandleFunc("/version", version)
	log.Fatal(http.ListenAndServe(":8080", nil))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"net"

	"github.com/packet-guardian/pg-dhcp/rpcclient"
)

var (
	serverAddress string
)

func init() {
	flag.StringVar(&serverAddress, "h", "localhost:8677", "DHCP managment host address")
}

func main() {
	flag.Parse()

	client, err := rpcclient.Connect("tcp", serverAddress)
	if err != nil {
		log.Fatal(err)
	}
	defer client.Close()

	command := flag.Arg(0)
	var args []string
	if len(flag.Args()) > 0 {
		args = flag.Args()[1:]
	}

	if command == "device-test" {
		mac, _ := net.ParseMAC(args[0])
		if err := client.Device().Register(mac); err != nil {
			log.Fatal(err)
		}

		d, err := client.Device().Get(mac)
		if err != nil {
			log.Fatal(err)
		}

		log.Printf("%#v\n", d)
	}
}
<commit_msg>Expanded RPC test script with all available calls<commit_after>package main

import (
	"flag"
	"log"
	"net"
	"time"

	"github.com/packet-guardian/pg-dhcp/rpcclient"
)

var (
	serverAddress string
)

func init() {
	flag.StringVar(&serverAddress, "h", "localhost:8677", "DHCP managment host address")
}

func main() {
	flag.Parse()

	client, err := rpcclient.Connect("tcp", serverAddress)
	if err != nil {
		log.Fatal(err)
	}
	defer client.Close()

	command := flag.Arg(0)
	var args []string
	if len(flag.Args()) > 0 {
		args = flag.Args()[1:]
	}

	t := &timer{}

	if command == "device-test" {
		mac, _ := net.ParseMAC(args[0])

		t.start()
		if err := client.Device().Register(mac); err != nil {
			log.Fatal(err)
		}
		t.print()

		t.start()
		if err := getAndPrintDevice(client, mac); err != nil {
			log.Fatal(err)
		}
		t.print()

		t.start()
		if err := client.Device().Blacklist(mac); err != nil {
			log.Fatal(err)
		}
		t.print()

		t.start()
		if err := getAndPrintDevice(client, mac); err != nil {
			log.Fatal(err)
		}
		t.print()

		t.start()
		if err := client.Device().Unregister(mac); err != nil {
			log.Fatal(err)
		}
		t.print()

		t.start()
		if err := getAndPrintDevice(client, mac); err != nil {
			log.Fatal(err)
		}
		t.print()

		t.start()
		if err := client.Device().RemoveBlacklist(mac); err != nil {
			log.Fatal(err)
		}
		t.print()

		t.start()
		if err := getAndPrintDevice(client, mac); err != nil {
			log.Fatal(err)
		}
		t.print()
	} else if command == "network-test" {
		list, err := client.Network().GetNameList()
		if err != nil {
			log.Fatal(err)
		}
		log.Printf("%#v\n", list)
	} else if command == "lease-test" {
		leases, err := client.Lease().GetAllFromNetwork("network1")
		if err != nil {
			log.Fatal(err)
		}
		for _, lease := range leases {
			log.Printf("%#v\n", lease)
		}

		lease, err := client.Lease().Get(net.ParseIP("10.0.2.12"))
		if err != nil {
			log.Fatal(err)
		}
		log.Printf("\n%#v\n", lease)

		lease, err = client.Lease().Get(net.ParseIP("10.0.2.3"))
		if err != nil {
			log.Fatal(err)
		}
		log.Printf("\n%#v\n", lease)
	} else if command == "stat-test" {
		stats, err := client.Server().GetPoolStats()
		if err != nil {
			log.Fatal(err)
		}
		for _, stat := range stats {
			log.Printf("%#v\n", stat)
		}
	}
}

func getAndPrintDevice(c *rpcclient.Client, mac net.HardwareAddr) error {
	d, err := c.Device().Get(mac)
	if err != nil {
		return err
	}

	log.Printf("%#v\n", d)
	return nil
}

type timer struct {
	s time.Time
}

func (t *timer) start() {
	t.s = time.Now()
}

func (t *timer) print() {
	log.Printf("%s\n", time.Now().Sub(t.s).String())
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"os/exec"
	"strings"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju"
)

// SCPCommand is responsible for launching a scp command to copy files to/from remote machine(s)
type SCPCommand struct {
	SSHCommon
}

func (c *SCPCommand) Info() *cmd.Info {
	return &cmd.Info{"scp", "", "launch a scp command to copy files to/from remote machine(s)", ""}
}

func (c *SCPCommand) Init(f *gnuflag.FlagSet, args []string) error {
	addEnvironFlags(&c.EnvName, f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	switch len(f.Args()) {
	case 0, 1:
		return errors.New("at least two arguments required")
	default:
		c.Args = f.Args()
	}
	return nil
}

// Run resolves c.Target to a machine, or host of a unit and
// forks ssh with c.Args, if provided.
func (c *SCPCommand) Run(ctx *cmd.Context) error {
	var err error
	c.Conn, err = juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer c.Close()

	// translate arguments in the form 0:/somepath or service/0:/somepath into
	// ubuntu@machine:/somepath so they can be presented to scp.
	for i := range c.Args {
		// BUG(dfc) This will not work for IPv6 addresses if the user passes
		// 2001:db8::1:2:/somepath.
		// BUG(dfc) This will also not work if the local path has a : in it.
		if v := strings.SplitN(c.Args[i], ":", 2); len(v) > 1 {
			host, err := c.hostFromTarget(v[0])
			if err != nil {
				return err
			}
			c.Args[i] = "ubuntu@" + host + ":" + v[1]
		}
	}

	args := []string{"-o", "StrictHostKeyChecking no", "-o", "PasswordAuthentication no"}
	args = append(args, c.Args...)
	cmd := exec.Command("scp", args...)
	cmd.Stdin = ctx.Stdin
	cmd.Stdout = ctx.Stdout
	cmd.Stderr = ctx.Stderr
	return cmd.Run()
}
<commit_msg>final tweaks<commit_after>package main

import (
	"errors"
	"os/exec"
	"strings"

	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/juju"
)

// SCPCommand is responsible for launching a scp command to copy files to/from remote machine(s)
type SCPCommand struct {
	SSHCommon
}

func (c *SCPCommand) Info() *cmd.Info {
	return &cmd.Info{"scp", "", "launch a scp command to copy files to/from remote machine(s)", ""}
}

func (c *SCPCommand) Init(f *gnuflag.FlagSet, args []string) error {
	addEnvironFlags(&c.EnvName, f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	switch len(f.Args()) {
	case 0, 1:
		return errors.New("at least two arguments required")
	default:
		c.Args = f.Args()
	}
	return nil
}

// Run resolves c.Target to a machine, or host of a unit and
// forks ssh with c.Args, if provided.
func (c *SCPCommand) Run(ctx *cmd.Context) error {
	var err error
	c.Conn, err = juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer c.Close()

	// translate arguments in the form 0:/somepath or service/0:/somepath into
	// ubuntu@machine:/somepath so they can be presented to scp.
	for i := range c.Args {
		// BUG(dfc) This will not work for IPv6 addresses like 2001:db8::1:2:/somepath.
		if v := strings.SplitN(c.Args[i], ":", 2); len(v) > 1 {
			host, err := c.hostFromTarget(v[0])
			if err != nil {
				return err
			}
			c.Args[i] = "ubuntu@" + host + ":" + v[1]
		}
	}

	args := []string{"-o", "StrictHostKeyChecking no", "-o", "PasswordAuthentication no"}
	args = append(args, c.Args...)
	cmd := exec.Command("scp", args...)
	cmd.Stdin = ctx.Stdin
	cmd.Stdout = ctx.Stdout
	cmd.Stderr = ctx.Stderr
	return cmd.Run()
}
<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	"os"

	"github.com/coreos/etcd-ca/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/coreos/etcd-ca/depot"
	"github.com/coreos/etcd-ca/pkix"
)

func NewNewCertCommand() cli.Command {
	return cli.Command{
		Name:        "new-cert",
		Usage:       "Create certificate request for host",
		Description: "Create certificate for host, including certificate signing request and key. Certificate could be generated by signing the request.",
		Flags: []cli.Flag{
			cli.StringFlag{"passphrase", "", "Passphrase to encrypt private-key PEM block", ""},
			cli.StringFlag{"ip", "127.0.0.1", "IP address of the host", ""},
			cli.IntFlag{"key-bits", 4096, "Bit size of RSA keypair to generate", ""},
			cli.StringFlag{"domain", "", "Use domain instead of IP address for SAN", ""},
			cli.StringFlag{"organization", "etcd-ca", "Certificate organization", ""},
			cli.StringFlag{"country", "USA", "Certificate country", ""},
		},
		Action: newCertAction,
	}
}

func newCertAction(c *cli.Context) {
	if len(c.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "One host name must be provided.")
		os.Exit(1)
	}
	name := c.Args()[0]

	if depot.CheckCertificateSigningRequest(d, name) || depot.CheckPrivateKeyHost(d, name) {
		fmt.Fprintln(os.Stderr, "Certificate request has existed!")
		os.Exit(1)
	}

	var passphrase []byte
	var err error
	if c.IsSet("passphrase") {
		passphrase = []byte(c.String("passphrase"))
	} else {
		passphrase, err = createPassPhrase()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}

	key, err := pkix.CreateRSAKey(c.Int("key-bits"))
	if err != nil {
		fmt.Fprintln(os.Stderr, "Create RSA Key error:", err)
		os.Exit(1)
	} else {
		fmt.Printf("Created %s/key\n", name)
	}

	csr, err := pkix.CreateCertificateSigningRequest(key, name, c.String("ip"), c.String("domain"), c.String("organization"), c.String("country"))
	if err != nil {
		fmt.Fprintln(os.Stderr, "Create certificate request error:", err)
		os.Exit(1)
	} else {
		fmt.Printf("Created %s/crt\n", name)
	}

	if err = depot.PutCertificateSigningRequest(d, name, csr); err != nil {
		fmt.Fprintln(os.Stderr, "Save certificate request error:", err)
	}
	if err = depot.PutEncryptedPrivateKeyHost(d, name, key, passphrase); err != nil {
		fmt.Fprintln(os.Stderr, "Save key error:", err)
	}
}
<commit_msg>Fixed csr typo<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	"os"

	"github.com/coreos/etcd-ca/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/coreos/etcd-ca/depot"
	"github.com/coreos/etcd-ca/pkix"
)

func NewNewCertCommand() cli.Command {
	return cli.Command{
		Name:        "new-cert",
		Usage:       "Create certificate request for host",
		Description: "Create certificate for host, including certificate signing request and key. Certificate could be generated by signing the request.",
		Flags: []cli.Flag{
			cli.StringFlag{"passphrase", "", "Passphrase to encrypt private-key PEM block", ""},
			cli.StringFlag{"ip", "127.0.0.1", "IP address of the host", ""},
			cli.IntFlag{"key-bits", 4096, "Bit size of RSA keypair to generate", ""},
			cli.StringFlag{"domain", "", "Use domain instead of IP address for SAN", ""},
			cli.StringFlag{"organization", "etcd-ca", "Certificate organization", ""},
			cli.StringFlag{"country", "USA", "Certificate country", ""},
		},
		Action: newCertAction,
	}
}

func newCertAction(c *cli.Context) {
	if len(c.Args()) != 1 {
		fmt.Fprintln(os.Stderr, "One host name must be provided.")
		os.Exit(1)
	}
	name := c.Args()[0]

	if depot.CheckCertificateSigningRequest(d, name) || depot.CheckPrivateKeyHost(d, name) {
		fmt.Fprintln(os.Stderr, "Certificate request has existed!")
		os.Exit(1)
	}

	var passphrase []byte
	var err error
	if c.IsSet("passphrase") {
		passphrase = []byte(c.String("passphrase"))
	} else {
		passphrase, err = createPassPhrase()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}

	key, err := pkix.CreateRSAKey(c.Int("key-bits"))
	if err != nil {
		fmt.Fprintln(os.Stderr, "Create RSA Key error:", err)
		os.Exit(1)
	} else {
		fmt.Printf("Created %s/key\n", name)
	}

	csr, err := pkix.CreateCertificateSigningRequest(key, name, c.String("ip"), c.String("domain"), c.String("organization"), c.String("country"))
	if err != nil {
		fmt.Fprintln(os.Stderr, "Create certificate request error:", err)
		os.Exit(1)
	} else {
		fmt.Printf("Created %s/csr\n", name)
	}

	if err = depot.PutCertificateSigningRequest(d, name, csr); err != nil {
		fmt.Fprintln(os.Stderr, "Save certificate request error:", err)
	}
	if err = depot.PutEncryptedPrivateKeyHost(d, name, key, passphrase); err != nil {
		fmt.Fprintln(os.Stderr, "Save key error:", err)
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright 2016 Frank Wessels <fwessels@xs4all.nl>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cmd

import (
	"fmt"
	"github.com/s3git/s3git-go"
	"github.com/spf13/cobra"
)

// snapshotCmd represents the snapshot command
var snapshotCmd = &cobra.Command{
	Use:   "snapshot",
	Short: "Manage snapshots",
	Long: "Create, checkout and list snapshots",
}

var snapshotCreateCmd = &cobra.Command{
	Use:   "create [directory]",
	Short: "Create a snapshot",
	Long: "Create a snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Directory for snapshot must be specified")
		} else if message == "" {
			er("Commit message for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		key, nothing, err := repo.SnapshotCreate(args[0], message)
		if err != nil {
			er(err)
		}
		if nothing {
			fmt.Println("No changes to snapshot")
		} else {
			fmt.Printf("[commit %s]\n", key)
		}
	},
}

var snapshotCheckoutCmd = &cobra.Command{
	Use:   "checkout [directory] ([commit])",
	Short: "Checkout a snapshot",
	Long: "Checkout a snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		// TODO: Partial checkout would be nice (eg specify path as filter)

		if len(args) == 0 {
			er("Directory for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		var commit string
		if len(args) == 2 {
			commit = args[1]
		}

		err = repo.SnapshotCheckout(args[0], commit, hydrate)
		if err != nil {
			er(err)
		}
	},
}

var presignedUrls bool
var jsonOutput bool

var snapshotListCmd = &cobra.Command{
	Use:   "ls ([commit])",
	Short: "List a snapshot",
	Long: "List a snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Commit for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		var commit string
		if len(args) == 1 {
			commit = args[0]
		}

		// TODO: Dump result in JSON format
		err = repo.SnapshotList(commit, presignedUrls)
		if err != nil {
			er(err)
		}

	},
}

var snapshotLogCmd = &cobra.Command{
	Use:   "log",
	Short: "Show commit log for snapshots",
	Long: "Show commit log for snapshots",
	Run: func(cmd *cobra.Command, args []string) {

		_/*repo*/, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		// TODO: Implement log
	},
}

var snapshotStatusCmd = &cobra.Command{
	Use:   "status [directory] ([commit])",
	Short: "Show changes for snapshot",
	Long: "Show changes for snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Directory for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		var commit string
		if len(args) == 2 {
			commit = args[1]
		}

		err = repo.SnapshotStatus(args[0], commit)
		if err != nil {
			er(err)
		}
	},
}

func init() {
	RootCmd.AddCommand(snapshotCmd)
	snapshotCmd.AddCommand(snapshotCreateCmd)
	snapshotCmd.AddCommand(snapshotCheckoutCmd)
	snapshotCmd.AddCommand(snapshotListCmd)
	snapshotCmd.AddCommand(snapshotLogCmd)
	snapshotCmd.AddCommand(snapshotStatusCmd)

	// Local flags for create
	snapshotCreateCmd.Flags().StringVarP(&message, "message", "m", "", "Message for the commit of create snapshot")

	// Local flags for checkout
	snapshotCheckoutCmd.Flags().BoolVar(&hydrate, "hydrate", false, "Checkout in hydrated (original) format")

	// Local flags for list
	snapshotListCmd.Flags().BoolVar(&presignedUrls, "presigned", false, "Generate presigned urls for direct access from S3")
	snapshotListCmd.Flags().BoolVar(&presignedUrls, "json", false, "Output result in JSON")
}
<commit_msg>Change hydrate flag into dedupe flag for snapshot checkout (so default is hydrated storage)<commit_after>/*
 * Copyright 2016 Frank Wessels <fwessels@xs4all.nl>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package cmd

import (
	"fmt"
	"github.com/s3git/s3git-go"
	"github.com/spf13/cobra"
)

var dedupe bool
// snapshotCmd represents the snapshot command
var snapshotCmd = &cobra.Command{
	Use:   "snapshot",
	Short: "Manage snapshots",
	Long: "Create, checkout and list snapshots",
}

var snapshotCreateCmd = &cobra.Command{
	Use:   "create [directory]",
	Short: "Create a snapshot",
	Long: "Create a snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Directory for snapshot must be specified")
		} else if message == "" {
			er("Commit message for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		key, nothing, err := repo.SnapshotCreate(args[0], message)
		if err != nil {
			er(err)
		}
		if nothing {
			fmt.Println("No changes to snapshot")
		} else {
			fmt.Printf("[commit %s]\n", key)
		}
	},
}

var snapshotCheckoutCmd = &cobra.Command{
	Use:   "checkout [directory] ([commit])",
	Short: "Checkout a snapshot",
	Long: "Checkout a snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		// TODO: Partial checkout would be nice (eg specify path as filter)

		if len(args) == 0 {
			er("Directory for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		var commit string
		if len(args) == 2 {
			commit = args[1]
		}

		err = repo.SnapshotCheckout(args[0], commit, dedupe)
		if err != nil {
			er(err)
		}
	},
}

var presignedUrls bool
var jsonOutput bool

var snapshotListCmd = &cobra.Command{
	Use:   "ls ([commit])",
	Short: "List a snapshot",
	Long: "List a snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Commit for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		var commit string
		if len(args) == 1 {
			commit = args[0]
		}

		// TODO: Dump result in JSON format
		err = repo.SnapshotList(commit, presignedUrls)
		if err != nil {
			er(err)
		}

	},
}

var snapshotLogCmd = &cobra.Command{
	Use:   "log",
	Short: "Show commit log for snapshots",
	Long: "Show commit log for snapshots",
	Run: func(cmd *cobra.Command, args []string) {

		_/*repo*/, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		// TODO: Implement log
	},
}

var snapshotStatusCmd = &cobra.Command{
	Use:   "status [directory] ([commit])",
	Short: "Show changes for snapshot",
	Long: "Show changes for snapshot",
	Run: func(cmd *cobra.Command, args []string) {

		if len(args) == 0 {
			er("Directory for snapshot must be specified")
		}

		repo, err := s3git.OpenRepository(".")
		if err != nil {
			er(err)
		}

		var commit string
		if len(args) == 2 {
			commit = args[1]
		}

		err = repo.SnapshotStatus(args[0], commit)
		if err != nil {
			er(err)
		}
	},
}

func init() {
	RootCmd.AddCommand(snapshotCmd)
	snapshotCmd.AddCommand(snapshotCreateCmd)
	snapshotCmd.AddCommand(snapshotCheckoutCmd)
	snapshotCmd.AddCommand(snapshotListCmd)
	snapshotCmd.AddCommand(snapshotLogCmd)
	snapshotCmd.AddCommand(snapshotStatusCmd)

	// Local flags for create
	snapshotCreateCmd.Flags().StringVarP(&message, "message", "m", "", "Message for the commit of create snapshot")

	// Local flags for checkout
	snapshotCheckoutCmd.Flags().BoolVar(&dedupe, "dedupe", false, "Checkout in deduped (pointers) format")

	// Local flags for list
	snapshotListCmd.Flags().BoolVar(&presignedUrls, "presigned", false, "Generate presigned urls for direct access from S3")
	snapshotListCmd.Flags().BoolVar(&presignedUrls, "json", false, "Output result in JSON")
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"regexp"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/pressly/sup"
)

var (
	supfile     string
	envVars     flagStringSlice
	onlyHosts   string
	exceptHosts string

	showVersion bool
	showHelp    bool

	ErrUsage            = errors.New("Usage: sup [OPTIONS] NETWORK COMMAND [...]\n       sup [ --help | -v | --version ]")
	ErrUnknownNetwork   = errors.New("Unknown network")
	ErrNetworkNoHosts   = errors.New("No hosts defined for a given network")
	ErrCmd              = errors.New("Unknown command/target")
	ErrTargetNoCommands = errors.New("No commands defined for a given target")
)

type flagStringSlice []string

func (f *flagStringSlice) String() string {
	return fmt.Sprintf("%v", *f)
}

func (f *flagStringSlice) Set(value string) error {
	*f = append(*f, value)
	return nil
}

func init() {
	flag.StringVar(&supfile, "f", "./Supfile", "Custom path to Supfile")
	flag.Var(&envVars, "e", "Set environment variables")
	flag.Var(&envVars, "env", "Set environment variables")
	flag.StringVar(&onlyHosts, "only", "", "Filter hosts using regexp")
	flag.StringVar(&exceptHosts, "except", "", "Filter out hosts using regexp")

	flag.BoolVar(&showVersion, "v", false, "Print version")
	flag.BoolVar(&showVersion, "version", false, "Print version")
	flag.BoolVar(&showHelp, "h", false, "Show help")
	flag.BoolVar(&showHelp, "help", false, "Show help")
}

func networkUsage(conf *sup.Supfile) {
	w := &tabwriter.Writer{}
	w.Init(os.Stderr, 4, 4, 2, ' ', 0)
	defer w.Flush()

	// Print available networks/hosts.
	fmt.Fprintln(w, "Networks:\t")
	for name, network := range conf.Networks {
		fmt.Fprintf(w, "- %v\n", name)
		for _, host := range network.Hosts {
			fmt.Fprintf(w, "\t- %v\n", host)
		}
	}
	fmt.Fprintln(w)
}

func cmdUsage(conf *sup.Supfile) {
	w := &tabwriter.Writer{}
	w.Init(os.Stderr, 4, 4, 2, ' ', 0)
	defer w.Flush()

	// Print available targets/commands.
	fmt.Fprintln(w, "Targets:\t")
	for name, commands := range conf.Targets {
		fmt.Fprintf(w, "- %v\t%v\n", name, strings.Join(commands, " "))
	}
	fmt.Fprintln(w, "\t")
	fmt.Fprintln(w, "Commands:\t")
	for name, cmd := range conf.Commands {
		fmt.Fprintf(w, "- %v\t%v\n", name, cmd.Desc)
	}
	fmt.Fprintln(w)
}

// parseArgs parses args and returns network and commands to be run.
// On error, it prints usage and exits.
func parseArgs(conf *sup.Supfile) (*sup.Network, []*sup.Command, error) {
	var commands []*sup.Command

	args := flag.Args()

	if len(args) < 1 {
		networkUsage(conf)
		return nil, nil, ErrUsage
	}

	// Does the <network> exist?
	network, ok := conf.Networks[args[0]]
	if !ok {
		networkUsage(conf)
		return nil, nil, ErrUnknownNetwork
	}

	// Does the <network> have at least one host?
	if len(network.Hosts) == 0 {
		networkUsage(conf)
		return nil, nil, ErrNetworkNoHosts
	}

	// Check for the second argument
	if len(args) < 2 {
		cmdUsage(conf)
		return nil, nil, ErrUsage
	}

	// In case of the network.Env needs an initialization
	if network.Env == nil {
		network.Env = make(sup.EnvList, 0)
	}

	// Add default env variable with current network
	network.Env.Set("SUP_NETWORK", args[0])

	// Add default nonce
	network.Env.Set("SUP_TIME", time.Now().UTC().Format(time.RFC3339))
	if os.Getenv("SUP_TIME") != "" {
		network.Env.Set("SUP_TIME", os.Getenv("SUP_TIME"))
	}

	// Add user
	if os.Getenv("SUP_USER") != "" {
		network.Env.Set("SUP_USER", os.Getenv("SUP_USER"))
	} else {
		network.Env.Set("SUP_USER", os.Getenv("USER"))
	}

	for _, cmd := range args[1:] {
		// Target?
		target, isTarget := conf.Targets[cmd]
		if isTarget {
			// Loop over target's commands.
			for _, cmd := range target {
				command, isCommand := conf.Commands[cmd]
				if !isCommand {
					cmdUsage(conf)
					return nil, nil, fmt.Errorf("%v: %v", ErrCmd, cmd)
				}
				command.Name = cmd
				commands = append(commands, &command)
			}
		}

		// Command?
		command, isCommand := conf.Commands[cmd]
		if isCommand {
			command.Name = cmd
			commands = append(commands, &command)
		}

		if !isTarget && !isCommand {
			cmdUsage(conf)
			return nil, nil, fmt.Errorf("%v: %v", ErrCmd, cmd)
		}
	}

	return &network, commands, nil
}

func main() {
	flag.Parse()

	if showHelp {
		fmt.Fprintln(os.Stderr, ErrUsage, "\n\nOptions:")
		flag.PrintDefaults()
		return
	}

	if showVersion {
		fmt.Fprintln(os.Stderr, sup.VERSION)
		return
	}

	conf, err := sup.NewSupfile(supfile)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Parse network and commands to be run from args.
	network, commands, err := parseArgs(conf)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// --only flag filters hosts
	if onlyHosts != "" {
		expr, err := regexp.CompilePOSIX(onlyHosts)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		var hosts []string
		for _, host := range network.Hosts {
			if expr.MatchString(host) {
				hosts = append(hosts, host)
			}
		}
		if len(hosts) == 0 {
			fmt.Fprintln(os.Stderr, fmt.Errorf("no hosts match --only '%v' regexp", onlyHosts))
			os.Exit(1)
		}
		network.Hosts = hosts
	}

	// --except flag filters out hosts
	if exceptHosts != "" {
		expr, err := regexp.CompilePOSIX(exceptHosts)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		var hosts []string
		for _, host := range network.Hosts {
			if !expr.MatchString(host) {
				hosts = append(hosts, host)
			}
		}
		if len(hosts) == 0 {
			fmt.Fprintln(os.Stderr, fmt.Errorf("no hosts left after --except '%v' regexp", onlyHosts))
			os.Exit(1)
		}
		network.Hosts = hosts
	}

	// --env flag sets environment variables and overrides values defined in Supfile
	var vars sup.EnvList
	for _, env := range envVars {
		if len(env) == 0 {
			continue
		}
		i := strings.Index(env, "=")
		if i < 0 {
			if len(env) > 0 {
				vars.Set(env, "")
			}
			continue
		}
		vars.Set(env[:i], env[i+1:])
	}
	network.Env = append(network.Env, vars...)

	// Create new Stackup app.
	app, err := sup.New(conf)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Run all the commands in the given network.
	err = app.Run(network, commands...)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}
<commit_msg>Implement $SUP_ENV env var<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"regexp"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/pressly/sup"
)

var (
	supfile     string
	envVars     flagStringSlice
	onlyHosts   string
	exceptHosts string

	showVersion bool
	showHelp    bool

	ErrUsage            = errors.New("Usage: sup [OPTIONS] NETWORK COMMAND [...]\n       sup [ --help | -v | --version ]")
	ErrUnknownNetwork   = errors.New("Unknown network")
	ErrNetworkNoHosts   = errors.New("No hosts defined for a given network")
	ErrCmd              = errors.New("Unknown command/target")
	ErrTargetNoCommands = errors.New("No commands defined for a given target")
)

type flagStringSlice []string

func (f *flagStringSlice) String() string {
	return fmt.Sprintf("%v", *f)
}

func (f *flagStringSlice) Set(value string) error {
	*f = append(*f, value)
	return nil
}

func init() {
	flag.StringVar(&supfile, "f", "./Supfile", "Custom path to Supfile")
	flag.Var(&envVars, "e", "Set environment variables")
	flag.Var(&envVars, "env", "Set environment variables")
	flag.StringVar(&onlyHosts, "only", "", "Filter hosts using regexp")
	flag.StringVar(&exceptHosts, "except", "", "Filter out hosts using regexp")

	flag.BoolVar(&showVersion, "v", false, "Print version")
	flag.BoolVar(&showVersion, "version", false, "Print version")
	flag.BoolVar(&showHelp, "h", false, "Show help")
	flag.BoolVar(&showHelp, "help", false, "Show help")
}

func networkUsage(conf *sup.Supfile) {
	w := &tabwriter.Writer{}
	w.Init(os.Stderr, 4, 4, 2, ' ', 0)
	defer w.Flush()

	// Print available networks/hosts.
	fmt.Fprintln(w, "Networks:\t")
	for name, network := range conf.Networks {
		fmt.Fprintf(w, "- %v\n", name)
		for _, host := range network.Hosts {
			fmt.Fprintf(w, "\t- %v\n", host)
		}
	}
	fmt.Fprintln(w)
}

func cmdUsage(conf *sup.Supfile) {
	w := &tabwriter.Writer{}
	w.Init(os.Stderr, 4, 4, 2, ' ', 0)
	defer w.Flush()

	// Print available targets/commands.
	fmt.Fprintln(w, "Targets:\t")
	for name, commands := range conf.Targets {
		fmt.Fprintf(w, "- %v\t%v\n", name, strings.Join(commands, " "))
	}
	fmt.Fprintln(w, "\t")
	fmt.Fprintln(w, "Commands:\t")
	for name, cmd := range conf.Commands {
		fmt.Fprintf(w, "- %v\t%v\n", name, cmd.Desc)
	}
	fmt.Fprintln(w)
}

// parseArgs parses args and returns network and commands to be run.
// On error, it prints usage and exits.
func parseArgs(conf *sup.Supfile) (*sup.Network, []*sup.Command, error) {
	var commands []*sup.Command

	args := flag.Args()

	if len(args) < 1 {
		networkUsage(conf)
		return nil, nil, ErrUsage
	}

	// Does the <network> exist?
	network, ok := conf.Networks[args[0]]
	if !ok {
		networkUsage(conf)
		return nil, nil, ErrUnknownNetwork
	}

	// Does the <network> have at least one host?
	if len(network.Hosts) == 0 {
		networkUsage(conf)
		return nil, nil, ErrNetworkNoHosts
	}

	// Check for the second argument
	if len(args) < 2 {
		cmdUsage(conf)
		return nil, nil, ErrUsage
	}

	// In case of the network.Env needs an initialization
	if network.Env == nil {
		network.Env = make(sup.EnvList, 0)
	}

	// Add default env variable with current network
	network.Env.Set("SUP_NETWORK", args[0])

	// Add default nonce
	network.Env.Set("SUP_TIME", time.Now().UTC().Format(time.RFC3339))
	if os.Getenv("SUP_TIME") != "" {
		network.Env.Set("SUP_TIME", os.Getenv("SUP_TIME"))
	}

	// Add user
	if os.Getenv("SUP_USER") != "" {
		network.Env.Set("SUP_USER", os.Getenv("SUP_USER"))
	} else {
		network.Env.Set("SUP_USER", os.Getenv("USER"))
	}

	for _, cmd := range args[1:] {
		// Target?
		target, isTarget := conf.Targets[cmd]
		if isTarget {
			// Loop over target's commands.
			for _, cmd := range target {
				command, isCommand := conf.Commands[cmd]
				if !isCommand {
					cmdUsage(conf)
					return nil, nil, fmt.Errorf("%v: %v", ErrCmd, cmd)
				}
				command.Name = cmd
				commands = append(commands, &command)
			}
		}

		// Command?
		command, isCommand := conf.Commands[cmd]
		if isCommand {
			command.Name = cmd
			commands = append(commands, &command)
		}

		if !isTarget && !isCommand {
			cmdUsage(conf)
			return nil, nil, fmt.Errorf("%v: %v", ErrCmd, cmd)
		}
	}

	return &network, commands, nil
}

func main() {
	flag.Parse()

	if showHelp {
		fmt.Fprintln(os.Stderr, ErrUsage, "\n\nOptions:")
		flag.PrintDefaults()
		return
	}

	if showVersion {
		fmt.Fprintln(os.Stderr, sup.VERSION)
		return
	}

	conf, err := sup.NewSupfile(supfile)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Parse network and commands to be run from args.
	network, commands, err := parseArgs(conf)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// --only flag filters hosts
	if onlyHosts != "" {
		expr, err := regexp.CompilePOSIX(onlyHosts)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		var hosts []string
		for _, host := range network.Hosts {
			if expr.MatchString(host) {
				hosts = append(hosts, host)
			}
		}
		if len(hosts) == 0 {
			fmt.Fprintln(os.Stderr, fmt.Errorf("no hosts match --only '%v' regexp", onlyHosts))
			os.Exit(1)
		}
		network.Hosts = hosts
	}

	// --except flag filters out hosts
	if exceptHosts != "" {
		expr, err := regexp.CompilePOSIX(exceptHosts)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}

		var hosts []string
		for _, host := range network.Hosts {
			if !expr.MatchString(host) {
				hosts = append(hosts, host)
			}
		}
		if len(hosts) == 0 {
			fmt.Fprintln(os.Stderr, fmt.Errorf("no hosts left after --except '%v' regexp", onlyHosts))
			os.Exit(1)
		}
		network.Hosts = hosts
	}

	// --env flag sets environment variables and overrides values defined in Supfile
	var vars sup.EnvList
	var supEnv string
	for _, env := range envVars {
		if len(env) == 0 {
			continue
		}
		i := strings.Index(env, "=")
		if i < 0 {
			if len(env) > 0 {
				vars.Set(env, "")
			}
			continue
		}
		vars.Set(env[:i], env[i+1:])
	}
	if len(vars) > 0 {
		network.Env = append(network.Env, vars...)

		// Separate loop to omit duplicated keys.
		for _, v := range vars {
			supEnv += fmt.Sprintf(" -e %v=%q", v.Key, v.Value)
		}
		network.Env.Set("SUP_ENV", strings.TrimSpace(supEnv))
	}

	// Create new Stackup app.
	app, err := sup.New(conf)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	// Run all the commands in the given network.
	err = app.Run(network, commands...)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package kutil

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
)

type UtilError string

func (err UtilError) Error() string { return "Utils Error: " + string(err) }

func ProcessName() string {
	file := os.Args[0]
	i := strings.LastIndex(file, "\\")
	j := strings.LastIndex(file, "/")
	if j < i {
		file = file[i+1:]
	} else if j > i {
		file = file[j+1:]
	}
	i = strings.LastIndex(file, ".")
	if i > 0 {
		file = file[0:i]
	}
	return file
}

func DefaultPidFileName() string {
	file := ProcessName()
	pidFile := "/var/run/"
	if runtime.GOOS == "windows" {
		pidFile = "c:\\" + file + ".pid"
	} else {
		pidFile += file + ".pid"
	}
	return pidFile
}

func WritePidFile(myFile string, pid int) (err error) {
	return ioutil.WriteFile(myFile, []byte(fmt.Sprintf("%d", pid)), 0644)
}

func CheckWritePidPermission(pidFile string) error {
	if len(pidFile) <= 0 {
		pidFile = DefaultPidFileName()
	}
	if err := ioutil.WriteFile(pidFile, []byte(fmt.Sprintf("%d", 0)), 0644); err != nil {
		fmt.Sprintln("had no permission to write pid file: ", pidFile, err)
		return UtilError("sdfdsf")
	}
	return nil
}

func CreateProcess(background bool, file string, args []string) error {
	filePath, _ := filepath.Abs(file)
	cmd := exec.Command(filePath, args...)
	if background {
		cmd.Stdin = nil //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = nil
		cmd.Stderr = nil
	} else {
		cmd.Stdin = os.Stdin //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	return cmd.Start()
}

func CreateProcess2(background bool, file string, args []string) (*os.Process, error) {
	filePath, _ := filepath.Abs(file)
	if background {
		return os.StartProcess(filePath, args, &os.ProcAttr{Files: []*os.File{nil, nil, nil}})
	}
	return os.StartProcess(filePath, args, &os.ProcAttr{Files: []*os.File{os.Stdin, os.Stdout, os.Stderr}})
}

func WaitProcess(background bool, file string, args []string) ([]byte, error) {
	filePath, _ := filepath.Abs(file)
	cmd := exec.Command(filePath, args...)
	if background {
		cmd.Stdin = nil //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = nil
		cmd.Stderr = nil
	} else {
		cmd.Stdin = os.Stdin //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	//cmd.Run()不能返回内容，而Output可以返回内容。
	return cmd.Output()
}
<commit_msg>更新的Fork<commit_after>package kutil

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
)

type UtilError string

func (err UtilError) Error() string { return "Utils Error: " + string(err) }

func ProcessName() string {
	file := os.Args[0]
	i := strings.LastIndex(file, "\\")
	j := strings.LastIndex(file, "/")
	if j < i {
		file = file[i+1:]
	} else if j > i {
		file = file[j+1:]
	}
	i = strings.LastIndex(file, ".")
	if i > 0 {
		file = file[0:i]
	}
	return file
}

func DefaultPidFileName() string {
	file := ProcessName()
	pidFile := "/var/run/"
	if runtime.GOOS == "windows" {
		pidFile = "c:\\" + file + ".pid"
	} else {
		pidFile += file + ".pid"
	}
	return pidFile
}

func ProcessFile() string {
	file, _ := exec.LookPath(os.Args[0])
	path, _ := filepath.Abs(file)
	return path
}

func ProcessPath() string {
	path := filepath.Dir(ProcessFile())
	return path
}

func WritePidFile(myFile string, pid int) (err error) {
	return ioutil.WriteFile(myFile, []byte(fmt.Sprintf("%d", pid)), 0644)
}

func CheckWritePidPermission(pidFile string) error {
	if len(pidFile) <= 0 {
		pidFile = DefaultPidFileName()
	}
	if err := ioutil.WriteFile(pidFile, []byte(fmt.Sprintf("%d", 0)), 0644); err != nil {
		fmt.Sprintln("had no permission to write pid file: ", pidFile, err)
		return UtilError("sdfdsf")
	}
	return nil
}

func ExecProcess(background bool, file string, args ...string) (int, error) {
	fmt.Println("args:", args, "-Len:", len(args))
	filePath, _ := filepath.Abs(file)
	cmd := exec.Command(filePath, args...)
	if background {
		cmd.Stdin = nil //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = nil
		cmd.Stderr = nil
	} else {
		cmd.Stdin = os.Stdin //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	err := cmd.Start()
	if err == nil {
		return cmd.Process.Pid, nil
	}
	return -1, err
}

func StartProcess(background bool, file string, args []string) (*os.Process, error) {
	filePath, _ := filepath.Abs(file)
	if background {
		return os.StartProcess(filePath, args, &os.ProcAttr{Files: []*os.File{nil, nil, nil}})
	}
	return os.StartProcess(filePath, args, &os.ProcAttr{Files: []*os.File{os.Stdin, os.Stdout, os.Stderr}})
}

func WaitProcess(background bool, file string, args []string) ([]byte, error) {
	filePath, _ := filepath.Abs(file)
	cmd := exec.Command(filePath, args...)
	if background {
		cmd.Stdin = nil //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = nil
		cmd.Stderr = nil
	} else {
		cmd.Stdin = os.Stdin //给新进程设置文件描述符，可以重定向到文件中
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
	}
	//cmd.Run()不能返回内容，而Output可以返回内容。
	return cmd.Output()
}
<|endoftext|>


<commit_before>package krpc

import (
	"fmt"
)

// Msg represents messages that nodes in the network send to each other as specified by the protocol.
// They are also refered to as the KRPC messages.
// There are three types of messages: QUERY, RESPONSE, ERROR
// The message is a dictonary that is then
// "bencoded" (serialization & compression format adopted by the BitTorrent)
// and sent via the UDP connection to peers.
//
// A KRPC message is a single dictionary with two keys common to every message and additional keys depending on the type of message.
// Every message has a key "t" with a string value representing a transaction ID.
// This transaction ID is generated by the querying node and is echoed in the response, so responses
// may be correlated with multiple queries to the same node. The transaction ID should be encoded as a short string of binary numbers, typically 2 characters are enough as they cover 2^16 outstanding queries. The other key contained in every KRPC message is "y" with a single character value describing the type of message. The value of the "y" key is one of "q" for query, "r" for response, or "e" for error.
// 3 message types:  QUERY, RESPONSE, ERROR
type Msg struct {
	Q        string   `bencode:"q,omitempty"` // Query method (one of 4: "ping", "find_node", "get_peers", "announce_peer")
	A        *MsgArgs `bencode:"a,omitempty"` // named arguments sent with a query
	T        string   `bencode:"t"`           // required: transaction ID
	Y        string   `bencode:"y"`           // required: type of the message: q for QUERY, r for RESPONSE, e for ERROR
	R        *Return  `bencode:"r,omitempty"` // RESPONSE type only
	E        *Error   `bencode:"e,omitempty"` // ERROR type only
	IP       NodeAddr `bencode:"ip,omitempty"`
	ReadOnly bool     `bencode:"ro,omitempty"`
}

type MsgArgs struct {
	ID          ID     `bencode:"id"`                     // ID of the querying Node
	InfoHash    ID     `bencode:"info_hash,omitempty"`    // InfoHash of the torrent
	Target      ID     `bencode:"target,omitempty"`       // ID of the node sought
	Token       string `bencode:"token,omitempty"`        // Token received from an earlier get_peers query
	Port        int    `bencode:"port,omitempty"`         // Senders torrent port
	ImpliedPort bool   `bencode:"implied_port,omitempty"` // Use senders apparent DHT port
	Want        []Want `bencode:"want,omitempty"`         // Contains strings like "n4" and "n6" from BEP 32.
	NoSeed      int    `bencode:"noseed,omitempty"`       // BEP 33
	Scrape      int    `bencode:"scrape,omitempty"`       // BEP 33
}

type Want string

const (
	WantNodes  Want = "n4"
	WantNodes6 Want = "n6"
)

type Return struct {
	ID     ID                  `bencode:"id"`               // ID of the querying node
	Nodes  CompactIPv4NodeInfo `bencode:"nodes,omitempty"`  // K closest nodes to the requested target
	Nodes6 CompactIPv6NodeInfo `bencode:"nodes6,omitempty"` // K closest nodes to the requested target
	Token  *string             `bencode:"token,omitempty"`  // Token for future announce_peer
	Values []NodeAddr          `bencode:"values,omitempty"` // Torrent peers
}

func (r Return) ForAllNodes(f func(NodeInfo)) {
	for _, n := range r.Nodes {
		f(n)
	}
	for _, n := range r.Nodes6 {
		f(n)
	}
}

var _ fmt.Stringer = Msg{}

func (m Msg) String() string {
	return fmt.Sprintf("%#v", m)
}

// The node ID of the source of this Msg. Returns nil if it isn't present.
// TODO: Can we verify Msgs more aggressively so this is guaranteed to return
// a valid ID for a checked Msg?
func (m Msg) SenderID() *ID {
	switch m.Y {
	case "q":
		if m.A == nil {
			return nil
		}
		return &m.A.ID
	case "r":
		if m.R == nil {
			return nil
		}
		return &m.R.ID
	}
	return nil
}

func (m Msg) Error() *Error {
	if m.Y != "e" {
		return nil
	}
	return m.E
}
<commit_msg>Fix typo (#29)<commit_after>package krpc

import (
	"fmt"
)

// Msg represents messages that nodes in the network send to each other as specified by the protocol.
// They are also referred to as the KRPC messages.
// There are three types of messages: QUERY, RESPONSE, ERROR
// The message is a dictonary that is then
// "bencoded" (serialization & compression format adopted by the BitTorrent)
// and sent via the UDP connection to peers.
//
// A KRPC message is a single dictionary with two keys common to every message and additional keys depending on the type of message.
// Every message has a key "t" with a string value representing a transaction ID.
// This transaction ID is generated by the querying node and is echoed in the response, so responses
// may be correlated with multiple queries to the same node. The transaction ID should be encoded as a short string of binary numbers, typically 2 characters are enough as they cover 2^16 outstanding queries. The other key contained in every KRPC message is "y" with a single character value describing the type of message. The value of the "y" key is one of "q" for query, "r" for response, or "e" for error.
// 3 message types:  QUERY, RESPONSE, ERROR
type Msg struct {
	Q        string   `bencode:"q,omitempty"` // Query method (one of 4: "ping", "find_node", "get_peers", "announce_peer")
	A        *MsgArgs `bencode:"a,omitempty"` // named arguments sent with a query
	T        string   `bencode:"t"`           // required: transaction ID
	Y        string   `bencode:"y"`           // required: type of the message: q for QUERY, r for RESPONSE, e for ERROR
	R        *Return  `bencode:"r,omitempty"` // RESPONSE type only
	E        *Error   `bencode:"e,omitempty"` // ERROR type only
	IP       NodeAddr `bencode:"ip,omitempty"`
	ReadOnly bool     `bencode:"ro,omitempty"`
}

type MsgArgs struct {
	ID          ID     `bencode:"id"`                     // ID of the querying Node
	InfoHash    ID     `bencode:"info_hash,omitempty"`    // InfoHash of the torrent
	Target      ID     `bencode:"target,omitempty"`       // ID of the node sought
	Token       string `bencode:"token,omitempty"`        // Token received from an earlier get_peers query
	Port        int    `bencode:"port,omitempty"`         // Senders torrent port
	ImpliedPort bool   `bencode:"implied_port,omitempty"` // Use senders apparent DHT port
	Want        []Want `bencode:"want,omitempty"`         // Contains strings like "n4" and "n6" from BEP 32.
	NoSeed      int    `bencode:"noseed,omitempty"`       // BEP 33
	Scrape      int    `bencode:"scrape,omitempty"`       // BEP 33
}

type Want string

const (
	WantNodes  Want = "n4"
	WantNodes6 Want = "n6"
)

type Return struct {
	ID     ID                  `bencode:"id"`               // ID of the querying node
	Nodes  CompactIPv4NodeInfo `bencode:"nodes,omitempty"`  // K closest nodes to the requested target
	Nodes6 CompactIPv6NodeInfo `bencode:"nodes6,omitempty"` // K closest nodes to the requested target
	Token  *string             `bencode:"token,omitempty"`  // Token for future announce_peer
	Values []NodeAddr          `bencode:"values,omitempty"` // Torrent peers
}

func (r Return) ForAllNodes(f func(NodeInfo)) {
	for _, n := range r.Nodes {
		f(n)
	}
	for _, n := range r.Nodes6 {
		f(n)
	}
}

var _ fmt.Stringer = Msg{}

func (m Msg) String() string {
	return fmt.Sprintf("%#v", m)
}

// The node ID of the source of this Msg. Returns nil if it isn't present.
// TODO: Can we verify Msgs more aggressively so this is guaranteed to return
// a valid ID for a checked Msg?
func (m Msg) SenderID() *ID {
	switch m.Y {
	case "q":
		if m.A == nil {
			return nil
		}
		return &m.A.ID
	case "r":
		if m.R == nil {
			return nil
		}
		return &m.R.ID
	}
	return nil
}

func (m Msg) Error() *Error {
	if m.Y != "e" {
		return nil
	}
	return m.E
}
<|endoftext|>


<commit_before>package datadog

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/zorkian/go-datadog-api"

	"github.com/hashicorp/terraform/helper/schema"
	//"github.com/mitchellh/reflectwalk"
)

// resourceDatadogMonitor is a Datadog monitor resource
func resourceDatadogMonitor() *schema.Resource {
	return &schema.Resource{
		Create: resourceDatadogMonitorCreate,
		Read:   resourceDatadogMonitorRead,
		Update: resourceDatadogMonitorUpdate,
		Delete: resourceDatadogMonitorDelete,
		Exists: resourceDatadogMonitorExists,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"type": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"check": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"check_count": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			// Metric and Monitor settings
			"metric": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"tags": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "*",
			},
			"time_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"time_window": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"space_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"operator": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"message": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// Alert Settings
			"warning": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
			},
			"critical": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
			},

			// Additional Settings
			"notify_no_data": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"no_data_timeframe": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
			"renotify_interval": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  0,
			},
		},
	}
}

// buildMonitorStruct returns a monitor struct
func buildMonitorStruct(d *schema.ResourceData, typeStr string) *datadog.Monitor {
	// TODO: This is WIP, it might make more sense to have 2 separate resources.
	// A little duplication VS being able to handle if configs are optional/mandatory, etc

	log.Print("[DEBUG] building monitor struct")
	name := d.Get("name").(string)
	message := d.Get("message").(string)
	tags := d.Get("tags").(string)
	monitorType := d.Get("type").(string)
	var monitorName string
	var query string

	if d.Get("type").(string) == "service check" {
		check := d.Get("check").(string)
		checkCount := d.Get("check_count").(string)
		query = fmt.Sprintf("\"%s\".over(\"%s\").last(%s).count_by_status()", check, tags, checkCount)
		monitorName = name
	} else {
		operator := d.Get("operator").(string)
		timeAggr := d.Get("time_aggr").(string)
		timeWindow := d.Get("time_window").(string)
		spaceAggr := d.Get("space_aggr").(string)
		metric := d.Get("metric").(string)
		query = fmt.Sprintf("%s(%s):%s:%s{%s} %s %s", timeAggr,
			timeWindow,
			spaceAggr,
			metric,
			tags,
			operator,
			d.Get(fmt.Sprintf("%s.threshold", typeStr)))
		monitorName = fmt.Sprintf("[%s] %s", typeStr, name)
	}

	o := datadog.Options{
		NotifyNoData:    d.Get("notify_no_data").(bool),
		NoDataTimeframe: d.Get("no_data_timeframe").(int),
		RenotifyInterval: d.Get("renotify_interval").(int),
	}

	// TODO: handle notifications for both 'metric alerts' and 'service checks'.

	m := datadog.Monitor{
		Type:    monitorType,
		Query:   query,
		Name:    monitorName,
		Message: fmt.Sprintf("%s", message),
		Options: o,
	}

	return &m
}

// resourceDatadogMonitorCreate creates a monitor.
func resourceDatadogMonitorCreate(d *schema.ResourceData, meta interface{}) error {
	log.Print("[DEBUG] creating monitor")
	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		log.Print("[DEBUG] Creating service check")
		m, err := client.CreateMonitor(buildMonitorStruct(d, ""))

		if err != nil {
			return fmt.Errorf("error creating service check: %s", err)
		}

		d.SetId(strconv.Itoa(m.Id))
		return nil
	}

	log.Print("[DEBUG] Creating metrics check")

	w, err := client.CreateMonitor(buildMonitorStruct(d, "warning"))

	if err != nil {
		return fmt.Errorf("error creating warning: %s", err)
	}

	c, cErr := client.CreateMonitor(buildMonitorStruct(d, "critical"))

	if cErr != nil {
		return fmt.Errorf("error creating warning: %s", cErr)
	}

	log.Printf("[DEBUG] Saving IDs: %s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id))

	d.SetId(fmt.Sprintf("%s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id)))

	return nil
}

// resourceDatadogMonitorDelete deletes a monitor.
func resourceDatadogMonitorDelete(d *schema.ResourceData, meta interface{}) error {
	log.Print("[DEBUG] deleting monitor")
	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		log.Print("[DEBUG] Deleting service check")
		ID, err := strconv.Atoi(d.Id())
		if err != nil {
			return err
		}

		err = client.DeleteMonitor(ID)
		if err != nil {
			return err
		}
		return nil
	}

	log.Print("[DEBUG] Deleting metrics check")
	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			return fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			return iErr
		}

		err := client.DeleteMonitor(ID)
		if err != nil {
			return err
		}
	}
	return nil
}

// resourceDatadogMonitorExists verifies a monitor exists.
func resourceDatadogMonitorExists(d *schema.ResourceData, meta interface{}) (b bool, e error) {
	// Exists - This is called to verify a resource still exists. It is called prior to Read,
	// and lowers the burden of Read to be able to assume the resource exists.

	log.Print("[DEBUG] verifying monitor exists")
	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		log.Print("[DEBUG] verifying service check exists")
		ID, err := strconv.Atoi(d.Id())
		if err != nil {
			return false, err
		}
		_, err = client.GetMonitor(ID)

		if err != nil {
			if strings.EqualFold(err.Error(), "API error: 404 Not Found") {
				log.Printf("[DEBUG] monitor does not exist: %s", err)
				return false, err
			}
			return false, err
		}

		return true, nil
	}

	log.Print("[DEBUG] verifying metrics check exists")

	exists := false
	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			log.Printf("[DEBUG] Could not parse IDs: %s", v)
			return false, fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			log.Printf("[DEBUG] Received error converting string: %s", iErr)
			return false, iErr
		}
		_, err := client.GetMonitor(ID)
		if err != nil {
			if strings.EqualFold(err.Error(), "API error: 404 Not Found") {
				log.Printf("[DEBUG] monitor does not exist: %s", err)
				exists = false
				continue
			} else {
				e = err
				continue
			}
		}
		exists = true
	}

	if exists == false {
		return false, nil
	}

	return true, nil
}

// resourceDatadogMonitorRead synchronises Datadog and local state .
func resourceDatadogMonitorRead(d *schema.ResourceData, meta interface{}) error {
	// TODO: add support for this a read function.
	/* Read - This is called to resync the local state with the remote state.
	Terraform guarantees that an existing ID will be set. This ID should be
	used to look up the resource. Any remote data should be updated into the
	local data. No changes to the remote resource are to be made.
	*/

	return nil
}

// resourceDatadogMonitorUpdate updates a monitor.
func resourceDatadogMonitorUpdate(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] running update.")

	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		body := buildMonitorStruct(d, "")

		ID, err := strconv.Atoi(d.Id())
		if err != nil {
			return err
		}

		body.Id = ID
		err = client.UpdateMonitor(body)

		if err != nil {
			return fmt.Errorf("error updating warning: %s", err.Error())
		}
		return nil
	}

	split := strings.Split(d.Id(), "__")

	wID, cID := split[0], split[1]

	if wID == "" {
		return fmt.Errorf("Id not set.")
	}

	if cID == "" {
		return fmt.Errorf("Id not set.")
	}

	warningID, iErr := strconv.Atoi(wID)

	if iErr != nil {
		return iErr
	}

	criticalID, iErr := strconv.Atoi(cID)

	if iErr != nil {
		return iErr
	}

	warningBody := buildMonitorStruct(d, "warning")
	criticalBody := buildMonitorStruct(d, "critical")

	warningBody.Id = warningID
	criticalBody.Id = criticalID

	wErr := client.UpdateMonitor(warningBody)

	if wErr != nil {
		return fmt.Errorf("error updating warning: %s", wErr.Error())
	}

	cErr := client.UpdateMonitor(criticalBody)

	if cErr != nil {
		return fmt.Errorf("error updating critical: %s", cErr.Error())
	}

	return nil
}
<commit_msg>Gofmt.<commit_after>package datadog

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/zorkian/go-datadog-api"

	"github.com/hashicorp/terraform/helper/schema"
	//"github.com/mitchellh/reflectwalk"
)

// resourceDatadogMonitor is a Datadog monitor resource
func resourceDatadogMonitor() *schema.Resource {
	return &schema.Resource{
		Create: resourceDatadogMonitorCreate,
		Read:   resourceDatadogMonitorRead,
		Update: resourceDatadogMonitorUpdate,
		Delete: resourceDatadogMonitorDelete,
		Exists: resourceDatadogMonitorExists,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"type": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"check": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"check_count": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			// Metric and Monitor settings
			"metric": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"tags": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "*",
			},
			"time_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"time_window": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"space_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"operator": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},
			"message": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// Alert Settings
			"warning": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
			},
			"critical": &schema.Schema{
				Type:     schema.TypeMap,
				Optional: true,
			},

			// Additional Settings
			"notify_no_data": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"no_data_timeframe": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
			"renotify_interval": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
				Default:  0,
			},
		},
	}
}

// buildMonitorStruct returns a monitor struct
func buildMonitorStruct(d *schema.ResourceData, typeStr string) *datadog.Monitor {
	// TODO: This is WIP, it might make more sense to have 2 separate resources.
	// A little duplication VS being able to handle if configs are optional/mandatory, etc

	log.Print("[DEBUG] building monitor struct")
	name := d.Get("name").(string)
	message := d.Get("message").(string)
	tags := d.Get("tags").(string)
	monitorType := d.Get("type").(string)
	var monitorName string
	var query string

	if d.Get("type").(string) == "service check" {
		check := d.Get("check").(string)
		checkCount := d.Get("check_count").(string)
		query = fmt.Sprintf("\"%s\".over(\"%s\").last(%s).count_by_status()", check, tags, checkCount)
		monitorName = name
	} else {
		operator := d.Get("operator").(string)
		timeAggr := d.Get("time_aggr").(string)
		timeWindow := d.Get("time_window").(string)
		spaceAggr := d.Get("space_aggr").(string)
		metric := d.Get("metric").(string)
		query = fmt.Sprintf("%s(%s):%s:%s{%s} %s %s", timeAggr,
			timeWindow,
			spaceAggr,
			metric,
			tags,
			operator,
			d.Get(fmt.Sprintf("%s.threshold", typeStr)))
		monitorName = fmt.Sprintf("[%s] %s", typeStr, name)
	}

	o := datadog.Options{
		NotifyNoData:     d.Get("notify_no_data").(bool),
		NoDataTimeframe:  d.Get("no_data_timeframe").(int),
		RenotifyInterval: d.Get("renotify_interval").(int),
	}

	// TODO: handle notifications for both 'metric alerts' and 'service checks'.

	m := datadog.Monitor{
		Type:    monitorType,
		Query:   query,
		Name:    monitorName,
		Message: fmt.Sprintf("%s", message),
		Options: o,
	}

	return &m
}

// resourceDatadogMonitorCreate creates a monitor.
func resourceDatadogMonitorCreate(d *schema.ResourceData, meta interface{}) error {
	log.Print("[DEBUG] creating monitor")
	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		log.Print("[DEBUG] Creating service check")
		m, err := client.CreateMonitor(buildMonitorStruct(d, ""))

		if err != nil {
			return fmt.Errorf("error creating service check: %s", err)
		}

		d.SetId(strconv.Itoa(m.Id))
		return nil
	}

	log.Print("[DEBUG] Creating metrics check")

	w, err := client.CreateMonitor(buildMonitorStruct(d, "warning"))

	if err != nil {
		return fmt.Errorf("error creating warning: %s", err)
	}

	c, cErr := client.CreateMonitor(buildMonitorStruct(d, "critical"))

	if cErr != nil {
		return fmt.Errorf("error creating warning: %s", cErr)
	}

	log.Printf("[DEBUG] Saving IDs: %s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id))

	d.SetId(fmt.Sprintf("%s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id)))

	return nil
}

// resourceDatadogMonitorDelete deletes a monitor.
func resourceDatadogMonitorDelete(d *schema.ResourceData, meta interface{}) error {
	log.Print("[DEBUG] deleting monitor")
	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		log.Print("[DEBUG] Deleting service check")
		ID, err := strconv.Atoi(d.Id())
		if err != nil {
			return err
		}

		err = client.DeleteMonitor(ID)
		if err != nil {
			return err
		}
		return nil
	}

	log.Print("[DEBUG] Deleting metrics check")
	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			return fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			return iErr
		}

		err := client.DeleteMonitor(ID)
		if err != nil {
			return err
		}
	}
	return nil
}

// resourceDatadogMonitorExists verifies a monitor exists.
func resourceDatadogMonitorExists(d *schema.ResourceData, meta interface{}) (b bool, e error) {
	// Exists - This is called to verify a resource still exists. It is called prior to Read,
	// and lowers the burden of Read to be able to assume the resource exists.

	log.Print("[DEBUG] verifying monitor exists")
	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		log.Print("[DEBUG] verifying service check exists")
		ID, err := strconv.Atoi(d.Id())
		if err != nil {
			return false, err
		}
		_, err = client.GetMonitor(ID)

		if err != nil {
			if strings.EqualFold(err.Error(), "API error: 404 Not Found") {
				log.Printf("[DEBUG] monitor does not exist: %s", err)
				return false, err
			}
			return false, err
		}

		return true, nil
	}

	log.Print("[DEBUG] verifying metrics check exists")

	exists := false
	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			log.Printf("[DEBUG] Could not parse IDs: %s", v)
			return false, fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			log.Printf("[DEBUG] Received error converting string: %s", iErr)
			return false, iErr
		}
		_, err := client.GetMonitor(ID)
		if err != nil {
			if strings.EqualFold(err.Error(), "API error: 404 Not Found") {
				log.Printf("[DEBUG] monitor does not exist: %s", err)
				exists = false
				continue
			} else {
				e = err
				continue
			}
		}
		exists = true
	}

	if exists == false {
		return false, nil
	}

	return true, nil
}

// resourceDatadogMonitorRead synchronises Datadog and local state .
func resourceDatadogMonitorRead(d *schema.ResourceData, meta interface{}) error {
	// TODO: add support for this a read function.
	/* Read - This is called to resync the local state with the remote state.
	Terraform guarantees that an existing ID will be set. This ID should be
	used to look up the resource. Any remote data should be updated into the
	local data. No changes to the remote resource are to be made.
	*/

	return nil
}

// resourceDatadogMonitorUpdate updates a monitor.
func resourceDatadogMonitorUpdate(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] running update.")

	client := meta.(*datadog.Client)

	if d.Get("type").(string) == "service check" {
		body := buildMonitorStruct(d, "")

		ID, err := strconv.Atoi(d.Id())
		if err != nil {
			return err
		}

		body.Id = ID
		err = client.UpdateMonitor(body)

		if err != nil {
			return fmt.Errorf("error updating warning: %s", err.Error())
		}
		return nil
	}

	split := strings.Split(d.Id(), "__")

	wID, cID := split[0], split[1]

	if wID == "" {
		return fmt.Errorf("Id not set.")
	}

	if cID == "" {
		return fmt.Errorf("Id not set.")
	}

	warningID, iErr := strconv.Atoi(wID)

	if iErr != nil {
		return iErr
	}

	criticalID, iErr := strconv.Atoi(cID)

	if iErr != nil {
		return iErr
	}

	warningBody := buildMonitorStruct(d, "warning")
	criticalBody := buildMonitorStruct(d, "critical")

	warningBody.Id = warningID
	criticalBody.Id = criticalID

	wErr := client.UpdateMonitor(warningBody)

	if wErr != nil {
		return fmt.Errorf("error updating warning: %s", wErr.Error())
	}

	cErr := client.UpdateMonitor(criticalBody)

	if cErr != nil {
		return fmt.Errorf("error updating critical: %s", cErr.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package datadog

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/zorkian/go-datadog-api"

	"github.com/hashicorp/terraform/helper/schema"
)

// resourceDatadogMonitor is a Datadog monitor resource
func resourceDatadogMonitor() *schema.Resource {
	return &schema.Resource{
		Create: resourceDatadogMonitorCreate,
		Read:   resourceDatadogMonitorRead,
		Update: resourceDatadogMonitorUpdate,
		Delete: resourceDatadogMonitorDelete,
		Exists: resourceDatadogMonitorExists,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			// Metric and Monitor settings
			"metric": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"metric_tags": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "*",
			},
			"time_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"time_window": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"space_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"operator": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"message": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// Alert Settings
			"warning": &schema.Schema{
				Type:     schema.TypeMap,
				Required: true,
			},
			"critical": &schema.Schema{
				Type:     schema.TypeMap,
				Required: true,
			},

			// Additional Settings
			"notify_no_data": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"no_data_timeframe": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
		},
	}
}

// buildMonitorStruct returns a monitor struct
func buildMonitorStruct(d *schema.ResourceData, typeStr string) *datadog.Monitor {
	name := d.Get("name").(string)
	message := d.Get("message").(string)
	timeAggr := d.Get("time_aggr").(string)
	timeWindow := d.Get("time_window").(string)
	spaceAggr := d.Get("space_aggr").(string)
	metric := d.Get("metric").(string)
	tags := d.Get("metric_tags").(string)
	operator := d.Get("operator").(string)
	query := fmt.Sprintf("%s(%s):%s:%s{%s} %s %s", timeAggr,
		timeWindow,
		spaceAggr,
		metric,
		tags,
		operator,
		d.Get(fmt.Sprintf("%s.threshold", typeStr)))

	o := datadog.Options{
		NotifyNoData:    d.Get("notify_no_data").(bool),
		NoDataTimeframe: d.Get("no_data_timeframe").(int),
	}

	m := datadog.Monitor{
		Type:    "metric alert",
		Query:   query,
		Name:    fmt.Sprintf("[%s] %s", typeStr, name),
		Message: fmt.Sprintf("%s %s", message, d.Get(fmt.Sprintf("%s.notify", typeStr))),
		Options: o,
	}

	return &m
}

// resourceDatadogMonitorCreate creates a monitor.
func resourceDatadogMonitorCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*datadog.Client)

	w, err := client.CreateMonitor(buildMonitorStruct(d, "warning"))

	if err != nil {
		return fmt.Errorf("error creating warning: %s", err)
	}

	c, cErr := client.CreateMonitor(buildMonitorStruct(d, "critical"))

	if cErr != nil {
		return fmt.Errorf("error creating warning: %s", cErr)
	}

	log.Printf("[DEBUG] Saving IDs: %s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id))

	d.SetId(fmt.Sprintf("%s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id)))

	return nil
}

// resourceDatadogMonitorDelete deletes a monitor.
func resourceDatadogMonitorDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*datadog.Client)

	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			return fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			return iErr
		}

		err := client.DeleteMonitor(ID)
		if err != nil {
			return err
		}
	}
	return nil
}

// resourceDatadogMonitorExists verifies a monitor exists.
func resourceDatadogMonitorExists(d *schema.ResourceData, meta interface{}) (b bool, e error) {
	// Exists - This is called to verify a resource still exists. It is called prior to Read,
	// and lowers the burden of Read to be able to assume the resource exists.

	client := meta.(*datadog.Client)

	exists := false
	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			log.Printf("[DEBUG] Could not parse IDs: %s", v)
			return false, fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			log.Printf("[DEBUG] Received error converting string: %s", iErr)
			return false, iErr
		}
		_, err := client.GetMonitor(ID)
		if err != nil {
			if strings.EqualFold(err.Error(), "API error: 404 Not Found") {
				log.Printf("[DEBUG] monitor does not exist: %s", err)
				exists = false
				continue
			} else {
				e = err
				continue
			}
		}
		exists = true
	}

	if exists == false {
		return false, nil
	}

	return true, nil
}

// resourceDatadogMonitorRead synchronises Datadog and local state .
func resourceDatadogMonitorRead(d *schema.ResourceData, meta interface{}) error {
	// TODO: add support for this a read function.
	/* Read - This is called to resync the local state with the remote state.
	Terraform guarantees that an existing ID will be set. This ID should be
	used to look up the resource. Any remote data should be updated into the
	local data. No changes to the remote resource are to be made.
	*/

	return nil
}

// resourceDatadogMonitorUpdate updates a monitor.
func resourceDatadogMonitorUpdate(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] running update.")

	split := strings.Split(d.Id(), "__")

	wID, cID := split[0], split[1]

	if wID == "" {
		return fmt.Errorf("Id not set.")
	}

	if cID == "" {
		return fmt.Errorf("Id not set.")
	}

	warningID, iErr := strconv.Atoi(wID)

	if iErr != nil {
		return iErr
	}

	criticalID, iErr := strconv.Atoi(cID)

	if iErr != nil {
		return iErr
	}

	client := meta.(*datadog.Client)

	warningBody := buildMonitorStruct(d, "warning")
	criticalBody := buildMonitorStruct(d, "critical")

	warningBody.Id = warningID
	criticalBody.Id = criticalID

	wErr := client.UpdateMonitor(warningBody)

	if wErr != nil {
		return fmt.Errorf("error updating warning: %s", wErr.Error())
	}

	cErr := client.UpdateMonitor(criticalBody)

	if cErr != nil {
		return fmt.Errorf("error updating critical: %s", cErr.Error())
	}

	return nil
}
<commit_msg>Add notes to prepare support.<commit_after>package datadog

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/zorkian/go-datadog-api"

	"github.com/hashicorp/terraform/helper/schema"
)

// resourceDatadogMonitor is a Datadog monitor resource
func resourceDatadogMonitor() *schema.Resource {
	return &schema.Resource{
		Create: resourceDatadogMonitorCreate,
		Read:   resourceDatadogMonitorRead,
		Update: resourceDatadogMonitorUpdate,
		Delete: resourceDatadogMonitorDelete,
		Exists: resourceDatadogMonitorExists,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			// Metric and Monitor settings
			"metric": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"metric_tags": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Default:  "*",
			},
			"time_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"time_window": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"space_aggr": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"operator": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},
			"message": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			// Alert Settings
			"warning": &schema.Schema{
				Type:     schema.TypeMap,
				Required: true,
			},
			"critical": &schema.Schema{
				Type:     schema.TypeMap,
				Required: true,
			},

			// Additional Settings
			"notify_no_data": &schema.Schema{
				Type:     schema.TypeBool,
				Optional: true,
				Default:  true,
			},

			"no_data_timeframe": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},
		},
	}
}

// buildMonitorStruct returns a monitor struct
func buildMonitorStruct(d *schema.ResourceData, typeStr string) *datadog.Monitor {
	// TODO: Add support for service checks
	/*
	For service checks, query should be:

	"check".over(tags).last(count).count_by_status()
		* "check" name of the check, e.g. datadog.agent.up
	    * "tags" one or more quoted tags (comma-separated), or "*". e.g.: .over("env:prod", "role:db")
		* "count" must be at >= your max threshold (defined in the options). e.g. if you want to notify on 1 critical, 3 ok and 2 warn statuses count should be 3.

    For metric checks:

	* time_aggr(time_window):space_aggr:metric{tags} [by {key}] operator #
	* time_aggr avg, sum, max, min, change, or pct_change
	* time_window last_#m (5, 10, 15, or 30), last_#h (1, 2, or 4), or last_1d
	* space_aggr avg, sum, min, or max
	* tags one or more tags (comma-separated), or *
	* key a 'key' in key:value tag syntax; defines a separate alert for each tag in the group (multi-alert)
	* operator <, <=, >, >=, ==, or !=
	* # an integer or decimal number used to set the threshold
	If you are using the change or pct_change time aggregator, you can instead use change_aggr(time_aggr(time_window), timeshift):space_aggr:metric{tags} [by {key}] operator # with:
	* change_aggr change, pct_change
	* time_aggr avg, sum, max, min
	* time_window last_#m (1, 5, 10, 15, or 30), last_#h (1, 2, or 4), or last_#d (1 or 2)
	* timeshift #m_ago (5, 10, 15, or 30), #h_ago (1, 2, or 4), or 1d_ago
	 */
	name := d.Get("name").(string)
	message := d.Get("message").(string)
	timeAggr := d.Get("time_aggr").(string)
	timeWindow := d.Get("time_window").(string)
	spaceAggr := d.Get("space_aggr").(string)
	metric := d.Get("metric").(string)
	tags := d.Get("metric_tags").(string)
	operator := d.Get("operator").(string)
	query := fmt.Sprintf("%s(%s):%s:%s{%s} %s %s", timeAggr,
		timeWindow,
		spaceAggr,
		metric,
		tags,
		operator,
		d.Get(fmt.Sprintf("%s.threshold", typeStr)))

	o := datadog.Options{
		NotifyNoData:    d.Get("notify_no_data").(bool),
		NoDataTimeframe: d.Get("no_data_timeframe").(int),
	}

	m := datadog.Monitor{
		Type:    "metric alert",
		Query:   query,
		Name:    fmt.Sprintf("[%s] %s", typeStr, name),
		Message: fmt.Sprintf("%s %s", message, d.Get(fmt.Sprintf("%s.notify", typeStr))),
		Options: o,
	}

	return &m
}

// resourceDatadogMonitorCreate creates a monitor.
func resourceDatadogMonitorCreate(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*datadog.Client)

	w, err := client.CreateMonitor(buildMonitorStruct(d, "warning"))

	if err != nil {
		return fmt.Errorf("error creating warning: %s", err)
	}

	c, cErr := client.CreateMonitor(buildMonitorStruct(d, "critical"))

	if cErr != nil {
		return fmt.Errorf("error creating warning: %s", cErr)
	}

	log.Printf("[DEBUG] Saving IDs: %s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id))

	d.SetId(fmt.Sprintf("%s__%s", strconv.Itoa(w.Id), strconv.Itoa(c.Id)))

	return nil
}

// resourceDatadogMonitorDelete deletes a monitor.
func resourceDatadogMonitorDelete(d *schema.ResourceData, meta interface{}) error {
	client := meta.(*datadog.Client)

	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			return fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			return iErr
		}

		err := client.DeleteMonitor(ID)
		if err != nil {
			return err
		}
	}
	return nil
}

// resourceDatadogMonitorExists verifies a monitor exists.
func resourceDatadogMonitorExists(d *schema.ResourceData, meta interface{}) (b bool, e error) {
	// Exists - This is called to verify a resource still exists. It is called prior to Read,
	// and lowers the burden of Read to be able to assume the resource exists.

	client := meta.(*datadog.Client)

	exists := false
	for _, v := range strings.Split(d.Id(), "__") {
		if v == "" {
			log.Printf("[DEBUG] Could not parse IDs: %s", v)
			return false, fmt.Errorf("Id not set.")
		}
		ID, iErr := strconv.Atoi(v)

		if iErr != nil {
			log.Printf("[DEBUG] Received error converting string: %s", iErr)
			return false, iErr
		}
		_, err := client.GetMonitor(ID)
		if err != nil {
			if strings.EqualFold(err.Error(), "API error: 404 Not Found") {
				log.Printf("[DEBUG] monitor does not exist: %s", err)
				exists = false
				continue
			} else {
				e = err
				continue
			}
		}
		exists = true
	}

	if exists == false {
		return false, nil
	}

	return true, nil
}

// resourceDatadogMonitorRead synchronises Datadog and local state .
func resourceDatadogMonitorRead(d *schema.ResourceData, meta interface{}) error {
	// TODO: add support for this a read function.
	/* Read - This is called to resync the local state with the remote state.
	Terraform guarantees that an existing ID will be set. This ID should be
	used to look up the resource. Any remote data should be updated into the
	local data. No changes to the remote resource are to be made.
	*/

	return nil
}

// resourceDatadogMonitorUpdate updates a monitor.
func resourceDatadogMonitorUpdate(d *schema.ResourceData, meta interface{}) error {
	log.Printf("[DEBUG] running update.")

	split := strings.Split(d.Id(), "__")

	wID, cID := split[0], split[1]

	if wID == "" {
		return fmt.Errorf("Id not set.")
	}

	if cID == "" {
		return fmt.Errorf("Id not set.")
	}

	warningID, iErr := strconv.Atoi(wID)

	if iErr != nil {
		return iErr
	}

	criticalID, iErr := strconv.Atoi(cID)

	if iErr != nil {
		return iErr
	}

	client := meta.(*datadog.Client)

	warningBody := buildMonitorStruct(d, "warning")
	criticalBody := buildMonitorStruct(d, "critical")

	warningBody.Id = warningID
	criticalBody.Id = criticalID

	wErr := client.UpdateMonitor(warningBody)

	if wErr != nil {
		return fmt.Errorf("error updating warning: %s", wErr.Error())
	}

	cErr := client.UpdateMonitor(criticalBody)

	if cErr != nil {
		return fmt.Errorf("error updating critical: %s", cErr.Error())
	}

	return nil
}
<|endoftext|>