<commit_before>package trie

import (
	"bufio"
	"log"
	"os"
	"sort"
	"testing"
)

func addFromFile(t *Trie, path string) {
	file, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}

	reader := bufio.NewScanner(file)

	for reader.Scan() {
		t.Add(reader.Text(), nil)
	}

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

func TestTrieAdd(t *testing.T) {
	trie := New()

	n := trie.Add("foo", 1)

	if n.Meta().(int) != 1 {
		t.Errorf("Expected 1, got: %d", n.Meta().(int))
	}
}

func TestTrieFind(t *testing.T) {
	trie := New()
	trie.Add("foo", 1)

	n, ok := trie.Find("foo")
	if ok != true {
		t.Fatal("Could not find node")
	}

	if n.Meta().(int) != 1 {
		t.Errorf("Expected 1, got: %d", n.Meta().(int))
	}
}

func TestTrieFindMissing(t *testing.T) {
	trie := New()

	n, ok := trie.Find("foo")
	if ok != false {
		t.Errorf("Expected ok to be false")
	}
	if n != nil {
		t.Errorf("Expected nil, got: %v", n)
	}
}

func TestRemove(t *testing.T) {
	trie := New()
	initial := []string{"football", "foostar", "foosball"}

	for _, key := range initial {
		trie.Add(key, nil)
	}

	trie.Remove("foosball")
	keys := trie.Keys()

	if len(keys) != 2 {
		t.Errorf("Expected 2 keys got %d", len(keys))
	}

	for _, k := range keys {
		if k != "football" && k != "foostar" {
			t.Errorf("key was: %s", k)
		}
	}

	keys = trie.FuzzySearch("foo")
	if len(keys) != 2 {
		t.Errorf("Expected 2 keys got %d", len(keys))
	}

	for _, k := range keys {
		if k != "football" && k != "foostar" {
			t.Errorf("Expected football got: %#v", k)
		}
	}
}

func TestTrieKeys(t *testing.T) {
	trie := New()
	expected := []string{"bar", "foo"}

	for _, key := range expected {
		trie.Add(key, nil)
	}

	kl := len(trie.Keys())
	if kl != 2 {
		t.Errorf("Expected 2 keys, got %d, keys were: %v", kl, trie.Keys())
	}

	keys := trie.Keys()

	sort.Strings(keys)
	for i, key := range keys {
		if key != expected[i] {
			t.Errorf("Expected %#v, got %#v", expected[i], key)
		}
	}
}

func TestPrefixSearch(t *testing.T) {
	trie := New()
	expected := []string{
		"foo",
		"foosball",
		"football",
		"foreboding",
		"forementioned",
		"foretold",
		"foreverandeverandeverandever",
		"forbidden",
	}

	defer func() {
		r := recover()
		if r != nil {
			t.Error(r)
		}
	}()

	trie.Add("bar", nil)
	for _, key := range expected {
		trie.Add(key, nil)
	}

	tests := []struct {
		pre      string
		expected []string
		length   int
	}{
		{"fo", expected, len(expected)},
		{"foosbal", []string{"foosball"}, 1},
		{"abc", []string{}, 0},
	}

	for _, test := range tests {
		actual := trie.PrefixSearch(test.pre)
		sort.Strings(actual)
		sort.Strings(test.expected)
		if len(actual) != test.length {
			t.Errorf("Expected len(actual) to == %d for pre %s", test.length, test.pre)
		}

		for i, key := range actual {
			if key != test.expected[i] {
				t.Errorf("Expected %v got: %v", test.expected[i], key)
			}
		}
	}

	trie.PrefixSearch("fsfsdfasdf")
}

func TestFuzzySearch(t *testing.T) {
	trie := New()
	setup := []string{
		"foosball",
		"football",
		"bmerica",
		"ked",
		"kedlock",
		"frosty",
		"bfrza",
		"foo/bart/baz.go",
	}
	tests := []struct {
		partial string
		length  int
	}{
		{"fsb", 1},
		{"footbal", 1},
		{"football", 1},
		{"fs", 2},
		{"oos", 1},
		{"kl", 1},
		{"ft", 3},
		{"fy", 1},
		{"fz", 2},
		{"a", 5},
	}

	for _, key := range setup {
		trie.Add(key, nil)
	}

	for _, test := range tests {
		actual := trie.FuzzySearch(test.partial)
		if len(actual) != test.length {
			t.Errorf("Expected len(actual) to == %d, was %d for %s actual was %#v",
				test.length, len(actual), test.partial, actual)
		}
	}
}

func TestFuzzySearchSorting(t *testing.T) {
	trie := New()
	setup := []string{
		"foosball",
		"football",
		"bmerica",
		"ked",
		"kedlock",
		"frosty",
		"bfrza",
		"foo/bart/baz.go",
	}

	for _, key := range setup {
		trie.Add(key, nil)
	}

	actual := trie.FuzzySearch("fz")
	expected := []string{"bfrza", "foo/bart/baz.go"}

	if len(actual) != len(expected) {
		t.Fatalf("expected len %d got %d", len(expected), len(actual))
	}
	for i, v := range expected {
		if actual[i] != v {
			t.Errorf("Expected %s got %s", v, actual[i])
		}
	}

}

func BenchmarkTieKeys(b *testing.B) {
	trie := New()
	keys := []string{"bar", "foo", "baz", "bur", "zum", "burzum", "bark", "barcelona", "football", "foosball", "footlocker"}

	for _, key := range keys {
		trie.Add(key, nil)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		trie.Keys()
	}
}

func BenchmarkPrefixSearch(b *testing.B) {
	trie := New()
	addFromFile(trie, "/usr/share/dict/words")

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = trie.PrefixSearch("fo")
	}
}

func BenchmarkFuzzySearch(b *testing.B) {
	trie := New()
	addFromFile(trie, "/usr/share/dict/words")

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = trie.FuzzySearch("fs")
	}
}
<commit_msg>add test.<commit_after>package trie

import (
	"bufio"
	"log"
	"os"
	"sort"
	"testing"
)

func addFromFile(t *Trie, path string) {
	file, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}

	reader := bufio.NewScanner(file)

	for reader.Scan() {
		t.Add(reader.Text(), nil)
	}

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

func TestTrieAdd(t *testing.T) {
	trie := New()

	n := trie.Add("foo", 1)

	if n.Meta().(int) != 1 {
		t.Errorf("Expected 1, got: %d", n.Meta().(int))
	}
}

func TestTrieFind(t *testing.T) {
	trie := New()
	trie.Add("foo", 1)

	n, ok := trie.Find("foo")
	if ok != true {
		t.Fatal("Could not find node")
	}

	if n.Meta().(int) != 1 {
		t.Errorf("Expected 1, got: %d", n.Meta().(int))
	}
}

func TestTrieFindMissing(t *testing.T) {
	trie := New()

	n, ok := trie.Find("foo")
	if ok != false {
		t.Errorf("Expected ok to be false")
	}
	if n != nil {
		t.Errorf("Expected nil, got: %v", n)
	}
}

func TestRemove(t *testing.T) {
	trie := New()
	initial := []string{"football", "foostar", "foosball"}

	for _, key := range initial {
		trie.Add(key, nil)
	}

	trie.Remove("foosball")
	keys := trie.Keys()

	if len(keys) != 2 {
		t.Errorf("Expected 2 keys got %d", len(keys))
	}

	for _, k := range keys {
		if k != "football" && k != "foostar" {
			t.Errorf("key was: %s", k)
		}
	}

	keys = trie.FuzzySearch("foo")
	if len(keys) != 2 {
		t.Errorf("Expected 2 keys got %d", len(keys))
	}

	for _, k := range keys {
		if k != "football" && k != "foostar" {
			t.Errorf("Expected football got: %#v", k)
		}
	}
}

func TestTrieKeys(t *testing.T) {
	trie := New()
	expected := []string{"bar", "foo"}

	for _, key := range expected {
		trie.Add(key, nil)
	}

	kl := len(trie.Keys())
	if kl != 2 {
		t.Errorf("Expected 2 keys, got %d, keys were: %v", kl, trie.Keys())
	}

	keys := trie.Keys()

	sort.Strings(keys)
	for i, key := range keys {
		if key != expected[i] {
			t.Errorf("Expected %#v, got %#v", expected[i], key)
		}
	}
}

func TestPrefixSearch(t *testing.T) {
	trie := New()
	expected := []string{
		"foo",
		"foosball",
		"football",
		"foreboding",
		"forementioned",
		"foretold",
		"foreverandeverandeverandever",
		"forbidden",
	}

	defer func() {
		r := recover()
		if r != nil {
			t.Error(r)
		}
	}()

	trie.Add("bar", nil)
	for _, key := range expected {
		trie.Add(key, nil)
	}

	tests := []struct {
		pre      string
		expected []string
		length   int
	}{
		{"fo", expected, len(expected)},
		{"foosbal", []string{"foosball"}, 1},
		{"abc", []string{}, 0},
	}

	for _, test := range tests {
		actual := trie.PrefixSearch(test.pre)
		sort.Strings(actual)
		sort.Strings(test.expected)
		if len(actual) != test.length {
			t.Errorf("Expected len(actual) to == %d for pre %s", test.length, test.pre)
		}

		for i, key := range actual {
			if key != test.expected[i] {
				t.Errorf("Expected %v got: %v", test.expected[i], key)
			}
		}
	}

	trie.PrefixSearch("fsfsdfasdf")
}

func TestFuzzySearch(t *testing.T) {
	trie := New()
	setup := []string{
		"foosball",
		"football",
		"bmerica",
		"ked",
		"kedlock",
		"frosty",
		"bfrza",
		"foo/bart/baz.go",
	}
	tests := []struct {
		partial string
		length  int
	}{
		{"fsb", 1},
		{"footbal", 1},
		{"football", 1},
		{"fs", 2},
		{"oos", 1},
		{"kl", 1},
		{"ft", 3},
		{"fy", 1},
		{"fz", 2},
		{"a", 5},
	}

	for _, key := range setup {
		trie.Add(key, nil)
	}

	for _, test := range tests {
		actual := trie.FuzzySearch(test.partial)
		if len(actual) != test.length {
			t.Errorf("Expected len(actual) to == %d, was %d for %s actual was %#v",
				test.length, len(actual), test.partial, actual)
		}
	}
}

func TestFuzzySearchSorting(t *testing.T) {
	trie := New()
	setup := []string{
		"foosball",
		"football",
		"bmerica",
		"ked",
		"kedlock",
		"frosty",
		"bfrza",
		"foo/bart/baz.go",
	}

	for _, key := range setup {
		trie.Add(key, nil)
	}

	actual := trie.FuzzySearch("fz")
	expected := []string{"bfrza", "foo/bart/baz.go"}

	if len(actual) != len(expected) {
		t.Fatalf("expected len %d got %d", len(expected), len(actual))
	}
	for i, v := range expected {
		if actual[i] != v {
			t.Errorf("Expected %s got %s", v, actual[i])
		}
	}

}

func BenchmarkTieKeys(b *testing.B) {
	trie := New()
	keys := []string{"bar", "foo", "baz", "bur", "zum", "burzum", "bark", "barcelona", "football", "foosball", "footlocker"}

	for _, key := range keys {
		trie.Add(key, nil)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		trie.Keys()
	}
}

func BenchmarkPrefixSearch(b *testing.B) {
	trie := New()
	addFromFile(trie, "/usr/share/dict/words")

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = trie.PrefixSearch("fo")
	}
}

func BenchmarkFuzzySearch(b *testing.B) {
	trie := New()
	addFromFile(trie, "/usr/share/dict/words")

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_ = trie.FuzzySearch("fs")
	}
}

func TestSupportChinese(t *testing.T) {
	trie := New()
	expected := []string{"苹果 沂水县", "苹果", "大蒜", "大豆"}

	for _, key := range expected {
		trie.Add(key, nil)
	}

	tests := []struct {
		pre      string
		expected []string
		length   int
	}{
		{"苹", expected[:2], len(expected[:2])},
		{"大", expected[2:], len(expected[2:])},
		{"大蒜", []string{"大蒜"}, 1},
	}

	for _, test := range tests {
		actual := trie.PrefixSearch(test.pre)
		sort.Strings(actual)
		sort.Strings(test.expected)
		if len(actual) != test.length {
			t.Errorf("Expected len(actual) to == %d for pre %s", test.length, test.pre)
		}

		for i, key := range actual {
			if key != test.expected[i] {
				t.Errorf("Expected %v got: %v", test.expected[i], key)
			}
		}
	}
}
<|endoftext|>


<commit_before>// The design and name of TripleSec is (C) Keybase 2013
// This Go implementation is (C) Filippo Valsorda 2014
// Use of this source code is governed by the MIT License

// Package triplesec implements the TripleSec v3 encryption and authentication scheme.
//
// For details on TripleSec, go to https://keybase.io/triplesec/
package triplesec

import (
	"bytes"
	"code.google.com/p/go.crypto/salsa20"
	"code.google.com/p/go.crypto/scrypt"
	"github.com/keybase/go-triplesec/sha3"
	"code.google.com/p/go.crypto/twofish"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"fmt"
)

// A Cipher is an instance of TripleSec using a particular key.
type Cipher struct {
	passphrase []byte
}

// Overhead is the amount of bytes added to a TripleSec ciphertext.
// 	len(plaintext) + Overhead = len(ciphertext)
// It consists of: magic bytes + version + salt + 2 * MACs + 3 * IVS.
const Overhead = 4 + 4 + 16 + 64 + 64 + 16 + 16 + 24

// The MagicBytes are the four bytes prefixed to every TripleSec
// ciphertext, 1c 94 d7 de.
const MagicBytes = "\x1c\x94\xd7\xde"

var (
	saltSize     = 16
	macKeyLen    = 48
	cipherKeyLen = 32
	dkSize       = 2*macKeyLen + 3*cipherKeyLen
)

// NewCipher creates and returns a Cipher.
// The passphrase can be a human passphrase, and is stretched with scrypt
// and a random salt. However, a long passphrase is strongly recommended.
// There are no limits on passphrase length.
func NewCipher(passphrase []byte) *Cipher {
	c := new(Cipher)
	c.passphrase = append(c.passphrase, passphrase...)

	return c
}

// Encrypt encrypts and signs a plaintext message with TripleSec using a random
// salt and the Cipher passphrase. The dst buffer size must be at least len(src)
// + Overhead. dst and src can not overlap. src is left untouched.
//
// Encrypt returns a error on memory or RNG failures.
func (c *Cipher) Encrypt(dst, src []byte) error {
	if len(src) < 1 {
		return fmt.Errorf("the plaintext cannot be empty")
	}
	if len(dst) < len(src)+Overhead {
		return fmt.Errorf("the destination is shorter than the plaintext plus Overhead")
	}

	buf := bytes.NewBuffer(dst[:0])

	_, err := buf.Write([]byte(MagicBytes))
	if err != nil {
		return err
	}

	// Write version
	err = binary.Write(buf, binary.BigEndian, uint32(3))
	if err != nil {
		return err
	}

	salt := make([]byte, saltSize)
	_, err = rand.Read(salt)
	if err != nil {
		return err
	}
	_, err = buf.Write(salt)
	if err != nil {
		return err
	}

	dk, err := scrypt.Key(c.passphrase, salt, 32768, 8, 1, dkSize)
	if err != nil {
		return err
	}
	macKeys := dk[:macKeyLen*2]
	cipherKeys := dk[macKeyLen*2:]

	// The allocation over here can be made better
	encryptedData, err := encrypt_data(src, cipherKeys)
	if err != nil {
		return err
	}

	authenticatedData := make([]byte, 0, buf.Len()+len(encryptedData))
	authenticatedData = append(authenticatedData, buf.Bytes()...)
	authenticatedData = append(authenticatedData, encryptedData...)
	macsOutput := generate_macs(authenticatedData, macKeys)

	_, err = buf.Write(macsOutput)
	if err != nil {
		return err
	}
	_, err = buf.Write(encryptedData)
	if err != nil {
		return err
	}

	if buf.Len() != len(src)+Overhead {
		panic(fmt.Errorf("something went terribly wrong: output size is not consistent"))
	}

	return nil
}

func encrypt_data(plain, keys []byte) ([]byte, error) {
	var iv, key []byte
	var block cipher.Block
	var stream cipher.Stream

	iv_offset := 16 + 16 + 24
	res := make([]byte, len(plain)+iv_offset)

	iv = res[iv_offset-24 : iv_offset]
	_, err := rand.Read(iv)
	if err != nil {
		return nil, err
	}
	// For some reason salsa20 API is different
	key_array := new([32]byte)
	copy(key_array[:], keys[cipherKeyLen*2:])
	salsa20.XORKeyStream(res[iv_offset:], plain, iv, key_array)
	iv_offset -= 24

	iv = res[iv_offset-16 : iv_offset]
	_, err = rand.Read(iv)
	if err != nil {
		return nil, err
	}
	key = keys[cipherKeyLen : cipherKeyLen*2]
	block, err = twofish.NewCipher(key)
	if err != nil {
		return nil, err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(res[iv_offset:], res[iv_offset:])
	iv_offset -= 16

	iv = res[iv_offset-16 : iv_offset]
	_, err = rand.Read(iv)
	if err != nil {
		return nil, err
	}
	key = keys[:cipherKeyLen]
	block, err = aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(res[iv_offset:], res[iv_offset:])
	iv_offset -= 16

	if iv_offset != 0 {
		panic(fmt.Errorf("something went terribly wrong: iv_offset final value non-zero"))
	}

	return res, nil
}

func generate_macs(data, keys []byte) []byte {
	res := make([]byte, 0, 64*2)

	key := keys[:macKeyLen]
	mac := hmac.New(sha512.New, key)
	mac.Write(data)
	res = mac.Sum(res)

	key = keys[macKeyLen:]
	mac = hmac.New(sha3.NewKeccak512, key)
	mac.Write(data)
	res = mac.Sum(res)

	return res
}

// Decrypt decrypts a TripleSec ciphertext using the Cipher passphrase.
// The dst buffer size must be at least len(src) - Overhead.
// dst and src can not overlap. src is left untouched.
//
// Encrypt returns a error if the ciphertext is not recognized, if
// authentication fails or on memory failures.
func (c *Cipher) Decrypt(dst, src []byte) error {
	if len(src) <= Overhead {
		return fmt.Errorf("the ciphertext is too short to be a TripleSec ciphertext")
	}
	if len(dst) < len(src)-Overhead {
		return fmt.Errorf("the dst buffer is too short to hold the plaintext")
	}

	if !bytes.Equal(src[:4], []byte(MagicBytes)) {
		return fmt.Errorf("the ciphertext does not look like a TripleSec ciphertext")
	}

	v := make([]byte, 4)
	v_b := bytes.NewBuffer(v[:0])
	err := binary.Write(v_b, binary.BigEndian, uint32(3))
	if err != nil {
		return err
	}
	if !bytes.Equal(src[4:8], v) {
		return fmt.Errorf("unknown version")
	}

	salt := src[8:24]
	dk, err := scrypt.Key(c.passphrase, salt, 32768, 8, 1, dkSize)
	if err != nil {
		return err
	}
	macKeys := dk[:macKeyLen*2]
	cipherKeys := dk[macKeyLen*2:]

	macs := src[24 : 24+64*2]
	encryptedData := src[24+64*2:]

	authenticatedData := make([]byte, 0, 24+len(encryptedData))
	authenticatedData = append(authenticatedData, src[:24]...)
	authenticatedData = append(authenticatedData, encryptedData...)

	if !hmac.Equal(macs, generate_macs(authenticatedData, macKeys)) {
		return fmt.Errorf("TripleSec ciphertext authentication FAILED")
	}

	err = decrypt_data(dst, encryptedData, cipherKeys)
	if err != nil {
		return err
	}

	return nil
}

func decrypt_data(dst, data, keys []byte) error {
	var iv, key []byte
	var block cipher.Block
	var stream cipher.Stream
	var err error

	buffer := append([]byte{}, data...)

	iv_offset := 16
	iv = buffer[:iv_offset]
	key = keys[:cipherKeyLen]
	block, err = aes.NewCipher(key)
	if err != nil {
		return err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(buffer[iv_offset:], buffer[iv_offset:])

	iv_offset += 16
	iv = buffer[iv_offset-16 : iv_offset]
	key = keys[cipherKeyLen : cipherKeyLen*2]
	block, err = twofish.NewCipher(key)
	if err != nil {
		return err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(buffer[iv_offset:], buffer[iv_offset:])

	iv_offset += 24
	iv = buffer[iv_offset-24 : iv_offset]
	key_array := new([32]byte)
	copy(key_array[:], keys[cipherKeyLen*2:])
	salsa20.XORKeyStream(dst, buffer[iv_offset:], iv, key_array)

	if len(buffer[iv_offset:]) != len(data)-(16+16+24) {
		panic(fmt.Errorf("something went terribly wrong: buffer size is not consistent"))
	}

	return nil
}
<commit_msg>progress<commit_after>// The design and name of TripleSec is (C) Keybase 2013
// This Go implementation is (C) Filippo Valsorda 2014
// Use of this source code is governed by the MIT License

// Package triplesec implements the TripleSec v3 encryption and authentication scheme.
//
// For details on TripleSec, go to https://keybase.io/triplesec/
package triplesec

import (
	"bytes"
	"code.google.com/p/go.crypto/salsa20"
	"code.google.com/p/go.crypto/scrypt"
	"github.com/keybase/go-triplesec/sha3"
	"code.google.com/p/go.crypto/twofish"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"fmt"
)

const SaltSize = 16

type Cipher struct {
	passphrase []byte
	salt       []byte
	derivedKey []byte
}

// A Cipher is an instance of TripleSec using a particular key and
// a particular salt
func NewCipher(passphrase []byte, salt []byte) (*Cipher, error) {
	if salt != nil && len(salt) != SaltSize {
		return nil, fmt.Errorf("Need a salt of size %d", SaltSize)
	}
	return &Cipher{passphrase, salt, nil}, nil
}

func (c *Cipher) GetSalt() ([]byte, error) {
	if c.salt != nil {
		return c.salt, nil
	}
	c.salt = make([]byte, SaltSize)
	_, err := rand.Read(c.salt)
	if err != nil {
		return nil, err
	}
	return c.salt, nil
}

func (c *Cipher) DeriveKey(dkLen int) ([]byte, error) {
	if c.derivedKey != nil && len(c.derivedKey) >= dkLen {
		return c.derivedKey[0:dkLen], nil
	}
	dk, err := scrypt.Key(c.passphrase, c.salt, 32768, 8, 1, dkLen)
	if err != nil {
		return nil, err
	}
	c.derivedKey = dk
	return dk, err
}

// The MagicBytes are the four bytes prefixed to every TripleSec
// ciphertext, 1c 94 d7 de.
var MagicBytes = [4]byte{ 0x1c, 0x94, 0xd7, 0xde }

var Version uint32 = 3
const MacOutputLen = 64

var (
	macKeyLen    = 48
	cipherKeyLen = 32
	DkSize       = 2*macKeyLen + 3*cipherKeyLen
)

// Overhead is the amount of bytes added to a TripleSec ciphertext.
// 	len(plaintext) + Overhead = len(ciphertext)
// It consists of: magic bytes + version + salt + 2 * MACs + 3 * IVS.
const Overhead = len(MagicBytes) + 4 + SaltSize + 2*MacOutputLen + 16 + 16 + 24

// Encrypt encrypts and signs a plaintext message with TripleSec using a random
// salt and the Cipher passphrase. The dst buffer size must be at least len(src)
// + Overhead. dst and src can not overlap. src is left untouched.
//
// Encrypt returns a error on memory or RNG failures.
func (c *Cipher) Encrypt(src []byte) (dst []byte, err error) {
	if len(src) < 1 {
		return nil, fmt.Errorf("the plaintext cannot be empty")
	}

	dst = make([]byte, len(src) + Overhead)
	buf := bytes.NewBuffer(dst[:0])

	_, err = buf.Write(MagicBytes[0:])
	if err != nil {
		return
	}

	// Write version
	err = binary.Write(buf, binary.BigEndian, Version)
	if err != nil {
		return
	}

	salt, err := c.GetSalt()
	if err != nil {
		return
	}

	_, err = buf.Write(salt)
	if err != nil {
		return
	}

	dk, err := c.DeriveKey(DkSize)
	if err != nil {
		return
	}
	macKeys := dk[:macKeyLen*2]
	cipherKeys := dk[macKeyLen*2:]

	// The allocation over here can be made better
	encryptedData, err := encrypt_data(src, cipherKeys)
	if err != nil {
		return
	}

	authenticatedData := make([]byte, 0, buf.Len()+len(encryptedData))
	authenticatedData = append(authenticatedData, buf.Bytes()...)
	authenticatedData = append(authenticatedData, encryptedData...)
	macsOutput := generate_macs(authenticatedData, macKeys)

	_, err = buf.Write(macsOutput)
	if err != nil {
		return
	}
	_, err = buf.Write(encryptedData)
	if err != nil {
		return
	}

	if buf.Len() != len(src)+Overhead {
		panic(fmt.Errorf("something went terribly wrong: output size wrong"))
	}

	return buf.Bytes(), nil
}

func encrypt_data(plain, keys []byte) ([]byte, error) {
	var iv, key []byte
	var block cipher.Block
	var stream cipher.Stream

	iv_offset := 16 + 16 + 24
	res := make([]byte, len(plain)+iv_offset)

	iv = res[iv_offset-24 : iv_offset]
	_, err := rand.Read(iv)
	if err != nil {
		return nil, err
	}
	// For some reason salsa20 API is different
	key_array := new([32]byte)
	copy(key_array[:], keys[cipherKeyLen*2:])
	salsa20.XORKeyStream(res[iv_offset:], plain, iv, key_array)
	iv_offset -= 24

	iv = res[iv_offset-16 : iv_offset]
	_, err = rand.Read(iv)
	if err != nil {
		return nil, err
	}
	key = keys[cipherKeyLen : cipherKeyLen*2]
	block, err = twofish.NewCipher(key)
	if err != nil {
		return nil, err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(res[iv_offset:], res[iv_offset:])
	iv_offset -= 16

	iv = res[iv_offset-16 : iv_offset]
	_, err = rand.Read(iv)
	if err != nil {
		return nil, err
	}
	key = keys[:cipherKeyLen]
	block, err = aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(res[iv_offset:], res[iv_offset:])
	iv_offset -= 16

	if iv_offset != 0 {
		panic(fmt.Errorf("something went terribly wrong: iv_offset final value non-zero"))
	}

	return res, nil
}

func generate_macs(data, keys []byte) []byte {
	res := make([]byte, 0, 64*2)

	key := keys[:macKeyLen]
	mac := hmac.New(sha512.New, key)
	mac.Write(data)
	res = mac.Sum(res)

	key = keys[macKeyLen:]
	mac = hmac.New(sha3.NewKeccak512, key)
	mac.Write(data)
	res = mac.Sum(res)

	return res
}

// Decrypt decrypts a TripleSec ciphertext using the Cipher passphrase.
// The dst buffer size must be at least len(src) - Overhead.
// dst and src can not overlap. src is left untouched.
//
// Encrypt returns a error if the ciphertext is not recognized, if
// authentication fails or on memory failures.
func (c *Cipher) Decrypt(dst, src []byte) error {
	if len(src) <= Overhead {
		return fmt.Errorf("the ciphertext is too short to be a TripleSec ciphertext")
	}
	if len(dst) < len(src)-Overhead {
		return fmt.Errorf("the dst buffer is too short to hold the plaintext")
	}

	if !bytes.Equal(src[:4], MagicBytes[0:]) {
		return fmt.Errorf("the ciphertext does not look like a TripleSec ciphertext")
	}

	v := make([]byte, 4)
	v_b := bytes.NewBuffer(v[:0])
	err := binary.Write(v_b, binary.BigEndian, uint32(3))
	if err != nil {
		return err
	}
	if !bytes.Equal(src[4:8], v) {
		return fmt.Errorf("unknown version")
	}

	salt := src[8:24]
	dk, err := scrypt.Key(c.passphrase, salt, 32768, 8, 1, DkSize)
	if err != nil {
		return err
	}
	macKeys := dk[:macKeyLen*2]
	cipherKeys := dk[macKeyLen*2:]

	macs := src[24 : 24+64*2]
	encryptedData := src[24+64*2:]

	authenticatedData := make([]byte, 0, 24+len(encryptedData))
	authenticatedData = append(authenticatedData, src[:24]...)
	authenticatedData = append(authenticatedData, encryptedData...)

	if !hmac.Equal(macs, generate_macs(authenticatedData, macKeys)) {
		return fmt.Errorf("TripleSec ciphertext authentication FAILED")
	}

	err = decrypt_data(dst, encryptedData, cipherKeys)
	if err != nil {
		return err
	}

	return nil
}

func decrypt_data(dst, data, keys []byte) error {
	var iv, key []byte
	var block cipher.Block
	var stream cipher.Stream
	var err error

	buffer := append([]byte{}, data...)

	iv_offset := 16
	iv = buffer[:iv_offset]
	key = keys[:cipherKeyLen]
	block, err = aes.NewCipher(key)
	if err != nil {
		return err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(buffer[iv_offset:], buffer[iv_offset:])

	iv_offset += 16
	iv = buffer[iv_offset-16 : iv_offset]
	key = keys[cipherKeyLen : cipherKeyLen*2]
	block, err = twofish.NewCipher(key)
	if err != nil {
		return err
	}
	stream = cipher.NewCTR(block, iv)
	stream.XORKeyStream(buffer[iv_offset:], buffer[iv_offset:])

	iv_offset += 24
	iv = buffer[iv_offset-24 : iv_offset]
	key_array := new([32]byte)
	copy(key_array[:], keys[cipherKeyLen*2:])
	salsa20.XORKeyStream(dst, buffer[iv_offset:], iv, key_array)

	if len(buffer[iv_offset:]) != len(data)-(16+16+24) {
		panic(fmt.Errorf("something went terribly wrong: buffer size is not consistent"))
	}

	return nil
}
<|endoftext|>


<commit_before>package tarGz

/**
 * The code for this package was taken from MadCrazy's question on StackOverflow
 * URL: http://stackoverflow.com/questions/13611100/how-to-write-a-directory-not-just-the-files-in-it-to-a-tar-gz-file-in-golang
 */

import (
	"archive/tar"
	"compress/gzip"
	"io"
	"os"

	"github.com/lfkeitel/verbose"
)

var appLogger *verbose.Logger

func init() {
	appLogger = verbose.New("tarGz-log")

	fileLogger, err := verbose.NewFileHandler("logs/tar.log")
	if err != nil {
		panic("Failed to open logging directory")
	}

	appLogger.AddHandler("file", fileLogger)
}

func handleError(_e error) {
	if _e != nil {
		appLogger.Error(_e.Error())
	}
}

func tarGzWrite(_path string, tw *tar.Writer, fi os.FileInfo) {
	fr, err := os.Open(_path)
	handleError(err)
	defer fr.Close()

	h := new(tar.Header)
	h.Name = _path
	h.Size = fi.Size()
	h.Mode = int64(fi.Mode())
	h.ModTime = fi.ModTime()

	err = tw.WriteHeader(h)
	handleError(err)

	_, err = io.Copy(tw, fr)
	handleError(err)
	return
}

func iterDirectory(dirPath string, tw *tar.Writer) {
	dir, err := os.Open(dirPath)
	handleError(err)
	defer dir.Close()
	fis, err := dir.Readdir(0)
	handleError(err)
	for _, fi := range fis {
		curPath := dirPath + "/" + fi.Name()
		if fi.IsDir() {
			iterDirectory(curPath, tw)
		} else {
			appLogger.Info("adding... %s", curPath)
			tarGzWrite(curPath, tw, fi)
		}
	}
	return
}

func TarGz(outFilePath string, inPath string) {
	// file write
	fw, err := os.Create(outFilePath)
	handleError(err)
	defer fw.Close()

	// gzip write
	gw := gzip.NewWriter(fw)
	defer gw.Close()

	// tar write
	tw := tar.NewWriter(gw)
	defer tw.Close()

	iterDirectory(inPath, tw)
	return
}
<commit_msg>Reduce log noise in Tar package<commit_after>package tarGz

/**
 * The code for this package was taken from MadCrazy's question on StackOverflow
 * URL: http://stackoverflow.com/questions/13611100/how-to-write-a-directory-not-just-the-files-in-it-to-a-tar-gz-file-in-golang
 */

import (
	"archive/tar"
	"compress/gzip"
	"io"
	"os"

	"github.com/lfkeitel/verbose"
)

var appLogger *verbose.Logger

func init() {
	appLogger = verbose.New("tarGz-log")

	fileLogger, err := verbose.NewFileHandler("logs/tar.log")
	if err != nil {
		panic("Failed to open logging directory")
	}

	appLogger.AddHandler("file", fileLogger)
}

func handleError(err error) {
	if err != nil {
		appLogger.Error(err.Error())
	}
}

func tarGzWrite(_path string, tw *tar.Writer, fi os.FileInfo) {
	fr, err := os.Open(_path)
	handleError(err)
	defer fr.Close()

	h := new(tar.Header)
	h.Name = _path
	h.Size = fi.Size()
	h.Mode = int64(fi.Mode())
	h.ModTime = fi.ModTime()

	err = tw.WriteHeader(h)
	handleError(err)

	_, err = io.Copy(tw, fr)
	handleError(err)
	return
}

func iterDirectory(dirPath string, tw *tar.Writer) {
	dir, err := os.Open(dirPath)
	handleError(err)
	defer dir.Close()
	fis, err := dir.Readdir(0)
	handleError(err)
	for _, fi := range fis {
		curPath := dirPath + "/" + fi.Name()
		if fi.IsDir() {
			iterDirectory(curPath, tw)
		} else {
			tarGzWrite(curPath, tw, fi)
		}
	}
	return
}

func TarGz(outFilePath string, inPath string) {
	// file write
	fw, err := os.Create(outFilePath)
	handleError(err)
	defer fw.Close()

	// gzip write
	gw := gzip.NewWriter(fw)
	defer gw.Close()

	// tar write
	tw := tar.NewWriter(gw)
	defer tw.Close()

	iterDirectory(inPath, tw)
	return
}
<|endoftext|>


<commit_before>package task

import "code.google.com/p/go-uuid/uuid"

type MinionTask interface {
	// Gets the UUID of the task
	GetTaskID() uuid.UUID

	// Gets the command to be executed
	GetCommand() (string, error)

	// Gets the command arguments
	GetArgs() ([]string, error)

	// Gets the time the task was sent for processing
	GetTimeReceived() (int64, error)

	// Gets the time when the task has been processed
	GetTimeProcessed() (int64, error)

	// Gets the task result
	GetResult() (string, error)

	// Gets the task error, if any
	GetError() string

	// Whether or not this task can run concurrently with other tasks
	IsConcurrent() bool

	// Sets the flag whether or not this task can run
	// concurrently with other tasks
	SetConcurrent(bool) error

	// Processes the task
	Process() error
}
<commit_msg>Remove "task" package, it has moved back to "minion" package<commit_after><|endoftext|>


<commit_before>package esa

import (
	"encoding/json"
	"fmt"
	"time"
)

type PostMember struct {
	Name       string `json:"name"`
	ScreenName string `json:"screen_name"`
	Icon       string `json:"icon"`
}

type Post struct {
	Number          int         `json:"number"`
	Name            string      `json:"name"`
	FullName        string      `json:"full_name"`
	WIP             bool        `json:"wip"`
	BodyMD          string      `json:"body_md"`
	BodyHTML        string      `json:"body_html"`
	CreatedAt       time.Time   `json:"created_at"`
	Message         string      `json:"message"`
	URL             string      `json:"url"`
	UpdatedAt       time.Time   `json:"updated_at"`
	Tags            []string    `json:"tags"`
	Category        *string     `json:"category"`
	RevisionNumber  int         `json:"revision_number"`
	CreatedBy       PostMember  `json:"created_by"`
	UpdatedBy       PostMember  `json:"updated_by"`
	Kind            string      `json:"kind"`
	CommentsCount   int         `json:"comments_countr"`
	TaskCount       int         `json:"task_count"`
	DoneTasksCount  int         `json:"done_tasks_count"`
	StargazersCount int         `json:"stargazers_count"`
	WatchersCount   int         `json:"watchers_count"`
	Star            bool        `json:"star"`
	Watch           bool        `json:"watch"`
	Comments        []Comment   `json:"comments"`
	Stargazers      []Stargazer `json:"stargazers"`
}

type GetTeamPostsRequest struct {
	Q       *string
	Include *string
	Sort    *string
	Order   *string
	PaginationRequest
}

type GetTeamPostsResponse struct {
	Posts []Post `json:"posts"`
	PaginationResponse
}

func (c *Client) GetTeamPosts(teamName string, req *GetTeamPostsRequest) (*GetTeamPostsResponse, error) {
	buildReq := c.get(fmt.Sprintf("/v1/teams/%s/posts", teamName))

	if req.Q != nil {
		buildReq = buildReq.Param("q", *req.Q)
	}
	if req.Include != nil {
		buildReq = buildReq.Param("include", *req.Include)
	}
	if req.Sort != nil {
		buildReq = buildReq.Param("sort", *req.Sort)
	}
	if req.Order != nil {
		buildReq = buildReq.Param("order", *req.Order)
	}

	resp, body, errs := c.setPaginationParams(buildReq, &req.PaginationRequest).End()
	fmt.Println(body)

	if len(errs) > 0 {
		return nil, errs[0]
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, c.parseError(body)
	}

	var res GetTeamPostsResponse
	if err := json.Unmarshal([]byte(body), &res); err != nil {
		return nil, err
	}

	return &res, nil
}
<commit_msg>delete debug code.<commit_after>package esa

import (
	"encoding/json"
	"fmt"
	"time"
)

type PostMember struct {
	Name       string `json:"name"`
	ScreenName string `json:"screen_name"`
	Icon       string `json:"icon"`
}

type Post struct {
	Number          int         `json:"number"`
	Name            string      `json:"name"`
	FullName        string      `json:"full_name"`
	WIP             bool        `json:"wip"`
	BodyMD          string      `json:"body_md"`
	BodyHTML        string      `json:"body_html"`
	CreatedAt       time.Time   `json:"created_at"`
	Message         string      `json:"message"`
	URL             string      `json:"url"`
	UpdatedAt       time.Time   `json:"updated_at"`
	Tags            []string    `json:"tags"`
	Category        *string     `json:"category"`
	RevisionNumber  int         `json:"revision_number"`
	CreatedBy       PostMember  `json:"created_by"`
	UpdatedBy       PostMember  `json:"updated_by"`
	Kind            string      `json:"kind"`
	CommentsCount   int         `json:"comments_countr"`
	TaskCount       int         `json:"task_count"`
	DoneTasksCount  int         `json:"done_tasks_count"`
	StargazersCount int         `json:"stargazers_count"`
	WatchersCount   int         `json:"watchers_count"`
	Star            bool        `json:"star"`
	Watch           bool        `json:"watch"`
	Comments        []Comment   `json:"comments"`
	Stargazers      []Stargazer `json:"stargazers"`
}

type GetTeamPostsRequest struct {
	Q       *string
	Include *string
	Sort    *string
	Order   *string
	PaginationRequest
}

type GetTeamPostsResponse struct {
	Posts []Post `json:"posts"`
	PaginationResponse
}

func (c *Client) GetTeamPosts(teamName string, req *GetTeamPostsRequest) (*GetTeamPostsResponse, error) {
	buildReq := c.get(fmt.Sprintf("/v1/teams/%s/posts", teamName))

	if req.Q != nil {
		buildReq = buildReq.Param("q", *req.Q)
	}
	if req.Include != nil {
		buildReq = buildReq.Param("include", *req.Include)
	}
	if req.Sort != nil {
		buildReq = buildReq.Param("sort", *req.Sort)
	}
	if req.Order != nil {
		buildReq = buildReq.Param("order", *req.Order)
	}

	resp, body, errs := c.setPaginationParams(buildReq, &req.PaginationRequest).End()

	if len(errs) > 0 {
		return nil, errs[0]
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, c.parseError(body)
	}

	var res GetTeamPostsResponse
	if err := json.Unmarshal([]byte(body), &res); err != nil {
		return nil, err
	}

	return &res, nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package osenv

import (
	"runtime"
	"strings"

	"github.com/juju/utils/featureflag"
)

const (
	JujuEnvEnvKey           = "JUJU_ENV"
	JujuHomeEnvKey          = "JUJU_HOME"
	JujuRepositoryEnvKey    = "JUJU_REPOSITORY"
	JujuLoggingConfigEnvKey = "JUJU_LOGGING_CONFIG"
	JujuFeatureFlagEnvKey   = "JUJU_DEV_FEATURE_FLAGS"

	JujuRegistryKey = `HKLM:\SOFTWARE\Wow6432Node\juju-core`

	// TODO(thumper): 2013-09-02 bug 1219630
	// As much as I'd like to remove JujuContainerType now, it is still
	// needed as MAAS still needs it at this stage, and we can't fix
	// everything at once.
	JujuContainerTypeEnvKey = "JUJU_CONTAINER_TYPE"

	// JujuStatusIsoTimeEnvKey is the env var which if true, will cause status
	// timestamps to be written in RFC3339 format.
	JujuStatusIsoTimeEnvKey = "JUJU_STATUS_ISO_TIME"

	// JujuCLIVersion is a numeric value (1, 2, 3 etc) representing
	// the oldest CLI version which should be adhered to.
	// This includes args and output.
	// Default is 1.
	JujuCLIVersion = "JUJU_CLI_VERSION"
)

// FeatureFlags returns a map that can be merged with os.Environ.
func FeatureFlags() map[string]string {
	result := make(map[string]string)
	if envVar := featureflag.AsEnvironmentValue(); envVar != "" {
		result[JujuFeatureFlagEnvKey] = envVar
	}
	return result
}

// MergeEnvironment will return the current environment updated with
// all the values from newValues.  If current is nil, a new map is
// created.  If current is not nil, it is mutated.
func MergeEnvironment(current, newValues map[string]string) map[string]string {
	if current == nil {
		current = make(map[string]string)
	}
	if runtime.GOOS == "windows" {
		return mergeEnvWin(current, newValues)
	}
	return mergeEnvUnix(current, newValues)
}

// mergeEnvUnix merges the two evironment variable lists in a case sensitive way.
func mergeEnvUnix(current, newValues map[string]string) map[string]string {
	for key, value := range newValues {
		current[key] = value
	}
	return current
}

// mergeEnvWin merges the two environment variable lists in a case insensitive,
// but case preserving way.  Thus, if FOO=bar is set, and newValues has foo=baz,
// then the resultant map will contain FOO=baz.
func mergeEnvWin(current, newValues map[string]string) map[string]string {
	uppers := make(map[string]string, len(current))
	news := map[string]string{}
	for k, v := range current {
		uppers[strings.ToUpper(k)] = v
	}

	for k, v := range newValues {
		up := strings.ToUpper(k)
		if _, ok := uppers[up]; ok {
			uppers[up] = v
		} else {
			news[k] = v
		}
	}

	for k := range current {
		current[k] = uppers[strings.ToUpper(k)]
	}
	for k, v := range news {
		current[k] = v
	}
	return current
}
<commit_msg>Wow6432 might not exist on 32 bit windows<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package osenv

import (
	"runtime"
	"strings"

	"github.com/juju/utils/featureflag"
)

const (
	JujuEnvEnvKey           = "JUJU_ENV"
	JujuHomeEnvKey          = "JUJU_HOME"
	JujuRepositoryEnvKey    = "JUJU_REPOSITORY"
	JujuLoggingConfigEnvKey = "JUJU_LOGGING_CONFIG"
	JujuFeatureFlagEnvKey   = "JUJU_DEV_FEATURE_FLAGS"

	JujuRegistryKey = `HKLM:\SOFTWARE\juju-core`

	// TODO(thumper): 2013-09-02 bug 1219630
	// As much as I'd like to remove JujuContainerType now, it is still
	// needed as MAAS still needs it at this stage, and we can't fix
	// everything at once.
	JujuContainerTypeEnvKey = "JUJU_CONTAINER_TYPE"

	// JujuStatusIsoTimeEnvKey is the env var which if true, will cause status
	// timestamps to be written in RFC3339 format.
	JujuStatusIsoTimeEnvKey = "JUJU_STATUS_ISO_TIME"

	// JujuCLIVersion is a numeric value (1, 2, 3 etc) representing
	// the oldest CLI version which should be adhered to.
	// This includes args and output.
	// Default is 1.
	JujuCLIVersion = "JUJU_CLI_VERSION"
)

// FeatureFlags returns a map that can be merged with os.Environ.
func FeatureFlags() map[string]string {
	result := make(map[string]string)
	if envVar := featureflag.AsEnvironmentValue(); envVar != "" {
		result[JujuFeatureFlagEnvKey] = envVar
	}
	return result
}

// MergeEnvironment will return the current environment updated with
// all the values from newValues.  If current is nil, a new map is
// created.  If current is not nil, it is mutated.
func MergeEnvironment(current, newValues map[string]string) map[string]string {
	if current == nil {
		current = make(map[string]string)
	}
	if runtime.GOOS == "windows" {
		return mergeEnvWin(current, newValues)
	}
	return mergeEnvUnix(current, newValues)
}

// mergeEnvUnix merges the two evironment variable lists in a case sensitive way.
func mergeEnvUnix(current, newValues map[string]string) map[string]string {
	for key, value := range newValues {
		current[key] = value
	}
	return current
}

// mergeEnvWin merges the two environment variable lists in a case insensitive,
// but case preserving way.  Thus, if FOO=bar is set, and newValues has foo=baz,
// then the resultant map will contain FOO=baz.
func mergeEnvWin(current, newValues map[string]string) map[string]string {
	uppers := make(map[string]string, len(current))
	news := map[string]string{}
	for k, v := range current {
		uppers[strings.ToUpper(k)] = v
	}

	for k, v := range newValues {
		up := strings.ToUpper(k)
		if _, ok := uppers[up]; ok {
			uppers[up] = v
		} else {
			news[k] = v
		}
	}

	for k := range current {
		current[k] = uppers[strings.ToUpper(k)]
	}
	for k, v := range news {
		current[k] = v
	}
	return current
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/xml"
	"fmt"
	"io"
	"runtime"
	"strings"

	"github.com/jstemmer/go-junit-report/parser"
)

// JUnitTestSuites is a collection of JUnit test suites.
type JUnitTestSuites struct {
	XMLName xml.Name `xml:"testsuites"`
	Suites  []JUnitTestSuite
}

// JUnitTestSuite is a single JUnit test suite which may contain many
// testcases.
type JUnitTestSuite struct {
	XMLName    xml.Name        `xml:"testsuite"`
	Tests      int             `xml:"tests,attr"`
	Failures   int             `xml:"failures,attr"`
	Time       string          `xml:"time,attr"`
	Name       string          `xml:"name,attr"`
	Properties []JUnitProperty `xml:"properties>property,omitempty"`
	TestCases  []JUnitTestCase
}

// JUnitTestCase is a single test case with its result.
type JUnitTestCase struct {
	XMLName     xml.Name          `xml:"testcase"`
	Classname   string            `xml:"classname,attr"`
	Name        string            `xml:"name,attr"`
	Time        string            `xml:"time,attr"`
	SkipMessage *JUnitSkipMessage `xml:"skipped,omitempty"`
	Failure     *JUnitFailure     `xml:"failure,omitempty"`
}

// JUnitSkipMessage contains the reason why a testcase was skipped.
type JUnitSkipMessage struct {
	Message string `xml:"message,attr"`
}

// JUnitProperty represents a key/value pair used to define properties.
type JUnitProperty struct {
	Name  string `xml:"name,attr"`
	Value string `xml:"value,attr"`
}

// JUnitFailure contains data related to a failed test.
type JUnitFailure struct {
	Message  string `xml:"message,attr"`
	Type     string `xml:"type,attr"`
	Contents string `xml:",chardata"`
}

// JUnitReportXML writes a JUnit xml representation of the given report to w
// in the format described at http://windyroad.org/dl/Open%20Source/JUnit.xsd
func JUnitReportXML(report *parser.Report, noXMLHeader bool, goVersion string, w io.Writer) error {
	suites := JUnitTestSuites{}

	// convert Report to JUnit test suites
	for _, pkg := range report.Packages {
		ts := JUnitTestSuite{
			Tests:      len(pkg.Tests),
			Failures:   0,
			Time:       formatTime(pkg.Time),
			Name:       pkg.Name,
			Properties: []JUnitProperty{},
			TestCases:  []JUnitTestCase{},
		}

		classname := pkg.Name
		if idx := strings.LastIndex(classname, "/"); idx > -1 && idx < len(pkg.Name) {
			classname = pkg.Name[idx+1:]
		}

		// properties
		if goVersion == "" {
			// if goVersion was not specified as a flag, fall back to version reported by runtime
			goVersion = runtime.Version()
		}
		ts.Properties = append(ts.Properties, JUnitProperty{"go.version", goVersion})
		if pkg.CoveragePct != "" {
			ts.Properties = append(ts.Properties, JUnitProperty{"coverage.statements.pct", pkg.CoveragePct})
		}

		// individual test cases
		for _, test := range pkg.Tests {
			testCase := JUnitTestCase{
				Classname: classname,
				Name:      test.Name,
				Time:      formatTime(test.Time),
				Failure:   nil,
			}

			if test.Result == parser.FAIL {
				ts.Failures++
				testCase.Failure = &JUnitFailure{
					Message:  "Failed",
					Type:     "",
					Contents: strings.Join(test.Output, "\n"),
				}
			}

			if test.Result == parser.SKIP {
				testCase.SkipMessage = &JUnitSkipMessage{strings.Join(test.Output, "\n")}
			}

			ts.TestCases = append(ts.TestCases, testCase)
		}

		suites.Suites = append(suites.Suites, ts)
	}

	// to xml
	bytes, err := xml.MarshalIndent(suites, "", "\t")
	if err != nil {
		return err
	}

	writer := bufio.NewWriter(w)

	if !noXMLHeader {
		writer.WriteString(xml.Header)
	}

	writer.Write(bytes)
	writer.WriteByte('\n')
	writer.Flush()

	return nil
}

func countFailures(tests []parser.Test) (result int) {
	for _, test := range tests {
		if test.Result == parser.FAIL {
			result++
		}
	}
	return
}

func formatTime(time int) string {
	return fmt.Sprintf("%.3f", float64(time)/1000.0)
}
<commit_msg>Remove unused function<commit_after>package main

import (
	"bufio"
	"encoding/xml"
	"fmt"
	"io"
	"runtime"
	"strings"

	"github.com/jstemmer/go-junit-report/parser"
)

// JUnitTestSuites is a collection of JUnit test suites.
type JUnitTestSuites struct {
	XMLName xml.Name `xml:"testsuites"`
	Suites  []JUnitTestSuite
}

// JUnitTestSuite is a single JUnit test suite which may contain many
// testcases.
type JUnitTestSuite struct {
	XMLName    xml.Name        `xml:"testsuite"`
	Tests      int             `xml:"tests,attr"`
	Failures   int             `xml:"failures,attr"`
	Time       string          `xml:"time,attr"`
	Name       string          `xml:"name,attr"`
	Properties []JUnitProperty `xml:"properties>property,omitempty"`
	TestCases  []JUnitTestCase
}

// JUnitTestCase is a single test case with its result.
type JUnitTestCase struct {
	XMLName     xml.Name          `xml:"testcase"`
	Classname   string            `xml:"classname,attr"`
	Name        string            `xml:"name,attr"`
	Time        string            `xml:"time,attr"`
	SkipMessage *JUnitSkipMessage `xml:"skipped,omitempty"`
	Failure     *JUnitFailure     `xml:"failure,omitempty"`
}

// JUnitSkipMessage contains the reason why a testcase was skipped.
type JUnitSkipMessage struct {
	Message string `xml:"message,attr"`
}

// JUnitProperty represents a key/value pair used to define properties.
type JUnitProperty struct {
	Name  string `xml:"name,attr"`
	Value string `xml:"value,attr"`
}

// JUnitFailure contains data related to a failed test.
type JUnitFailure struct {
	Message  string `xml:"message,attr"`
	Type     string `xml:"type,attr"`
	Contents string `xml:",chardata"`
}

// JUnitReportXML writes a JUnit xml representation of the given report to w
// in the format described at http://windyroad.org/dl/Open%20Source/JUnit.xsd
func JUnitReportXML(report *parser.Report, noXMLHeader bool, goVersion string, w io.Writer) error {
	suites := JUnitTestSuites{}

	// convert Report to JUnit test suites
	for _, pkg := range report.Packages {
		ts := JUnitTestSuite{
			Tests:      len(pkg.Tests),
			Failures:   0,
			Time:       formatTime(pkg.Time),
			Name:       pkg.Name,
			Properties: []JUnitProperty{},
			TestCases:  []JUnitTestCase{},
		}

		classname := pkg.Name
		if idx := strings.LastIndex(classname, "/"); idx > -1 && idx < len(pkg.Name) {
			classname = pkg.Name[idx+1:]
		}

		// properties
		if goVersion == "" {
			// if goVersion was not specified as a flag, fall back to version reported by runtime
			goVersion = runtime.Version()
		}
		ts.Properties = append(ts.Properties, JUnitProperty{"go.version", goVersion})
		if pkg.CoveragePct != "" {
			ts.Properties = append(ts.Properties, JUnitProperty{"coverage.statements.pct", pkg.CoveragePct})
		}

		// individual test cases
		for _, test := range pkg.Tests {
			testCase := JUnitTestCase{
				Classname: classname,
				Name:      test.Name,
				Time:      formatTime(test.Time),
				Failure:   nil,
			}

			if test.Result == parser.FAIL {
				ts.Failures++
				testCase.Failure = &JUnitFailure{
					Message:  "Failed",
					Type:     "",
					Contents: strings.Join(test.Output, "\n"),
				}
			}

			if test.Result == parser.SKIP {
				testCase.SkipMessage = &JUnitSkipMessage{strings.Join(test.Output, "\n")}
			}

			ts.TestCases = append(ts.TestCases, testCase)
		}

		suites.Suites = append(suites.Suites, ts)
	}

	// to xml
	bytes, err := xml.MarshalIndent(suites, "", "\t")
	if err != nil {
		return err
	}

	writer := bufio.NewWriter(w)

	if !noXMLHeader {
		writer.WriteString(xml.Header)
	}

	writer.Write(bytes)
	writer.WriteByte('\n')
	writer.Flush()

	return nil
}

func formatTime(time int) string {
	return fmt.Sprintf("%.3f", float64(time)/1000.0)
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/cipher"
	"crypto/aes"
	"encoding/base64"
	"fmt"
	"io"
	"crypto/rand"
	"errors"
	"io/ioutil"
	"regexp"
	"bytes"
	"flag"
	"bufio"
	"os"
	"strings"
)

/*
对文本中形如${?:clear text}的文本进行加密，加密后，显示为${AES:密文}，反之解密
*/

func main() {
	mode := flag.String("mode", "encode", "mode(encode/decode)")
	key := flag.String("key", "", "key to encyption or decyption")
	file := flag.String("file", "", "input file")
	flag.Parse()
	if *file == "" {
		fmt.Printf("file argument is required!\n")
		Usage()
		os.Exit(-1)
	}

	keyStr := *key;
	if *key == "" {
		reader := bufio.NewReader(os.Stdin)
		fmt.Print("Please input key: ")
		keyStr, _ = reader.ReadString('\n')
		keyStr = strings.TrimSpace(keyStr)
	}

	keyStr = FixStrLength(keyStr, 16);

	txtBytes, _ := ioutil.ReadFile(*file)
	txt := string(txtBytes);

	if *mode == "encode" {
		regex, _ := regexp.Compile("\\$\\{\\?:(.*?)\\}")
		result := ReplaceAllGroupFunc(regex, txt, func(groups []string) string {
			cipher, _ := CBCEncrypt(keyStr, groups[1])
			return "${AES:" + cipher + "}"
		})

		fmt.Printf("%s\n", result)
	} else if *mode == "decode" {
		regex, _ := regexp.Compile("\\$\\{AES:(.*?)\\}")
		result := ReplaceAllGroupFunc(regex, txt, func(groups []string) string {
			clear, _ := CBCDecrypt(keyStr, groups[1])
			return clear
		})
		fmt.Printf("%s\n", result)
	} else {
		fmt.Printf("mode argument should be ecode or decode!\n")
		Usage();
		os.Exit(-1)
	}
}
func FixStrLength(s string, fixLen int) string {
	slen := len(s)
	if slen < fixLen {
		return s + strings.Repeat("0", fixLen - slen)
	}

	if (slen > fixLen) {
		return s[:fixLen]
	}

	return s
}

var Usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
}

func ReplaceAllGroupFunc(re *regexp.Regexp, str string, repl func([]string) string) string {
	result := ""
	lastIndex := 0

	for _, v := range re.FindAllSubmatchIndex([]byte(str), -1) {
		groups := []string{}
		for i := 0; i < len(v); i += 2 {
			groups = append(groups, str[v[i]:v[i+1]])
		}

		result += str[lastIndex:v[0]] + repl(groups)
		lastIndex = v[1]
	}

	return result + str[lastIndex:]
}

func CBCEncrypt(strKey, strPlaintext string) (string, error) {
	key := []byte(strKey)
	plaintext := []byte(strPlaintext)

	// CBC mode works on blocks so plaintexts may need to be padded to the
	// next whole block. For an example of such padding, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. Here we'll
	// assume that the plaintext is already of the correct length.
	//if len(plaintext) % aes.BlockSize != 0 {
	//	return "", errors.New("plaintext is not a multiple of the block size")
	//}
	plaintext = PKCS5Padding(plaintext, aes.BlockSize)

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize + len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	base64Text := base64.StdEncoding.EncodeToString(ciphertext)

	return base64Text, nil
}

func CBCDecrypt(strKey, strCiphertext string) (string, error) {
	key := []byte(strKey)
	ciphertext, _ := base64.StdEncoding.DecodeString(strCiphertext)

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	if len(ciphertext) < aes.BlockSize {
		return "", errors.New("ciphertext too short")
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	// CBC mode always works in whole blocks.
	if len(ciphertext) % aes.BlockSize != 0 {
		return "", errors.New("ciphertext is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)

	// CryptBlocks can work in-place if the two arguments are the same.
	mode.CryptBlocks(ciphertext, ciphertext)

	ciphertext = PKCS5UnPadding(ciphertext)

	// If the original plaintext lengths are not a multiple of the block
	// size, padding would have to be added when encrypting, which would be
	// removed at this point. For an example, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. However, it's
	// critical to note that ciphertexts must be authenticated (i.e. by
	// using crypto/hmac) before being decrypted in order to avoid creating
	// a padding oracle.
	return string(ciphertext), nil
}

func PKCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func PKCS5UnPadding(origData []byte) []byte {
	length := len(origData)
	// 去掉最后一个字节 unpadding 次
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

func ZeroPadding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{0}, padding)
	return append(ciphertext, padtext...)
}<commit_msg>use GetPass to get pass from terminal<commit_after>package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"flag"
	"fmt"
	"github.com/dgiagio/getpass"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
)

/*
对文本中形如${?:clear text}的文本进行加密，加密后，显示为${AES:密文}，反之解密
*/

func main() {
	mode := flag.String("mode", "encode", "mode(encode/decode/src)")
	key := flag.String("key", "", "key to encyption or decyption")
	infile := flag.String("infile", "", "input file")
	outfile := flag.String("outfile", "", "output file")
	flag.Parse()
	if *infile == "" {
		msg := "file argument is required!\n"
		printErrorAndExit(msg)
	}

	keyStr := *key
	if *key == "" {
		keyStr, _ = getpass.GetPassword("Please input the key: ")
	}
	keyStr = FixStrLength(keyStr, 16)

	txtBytes, err := ioutil.ReadFile(*infile)
	checkError(err)

	txt := string(txtBytes)

	var regex *regexp.Regexp
	var replaceFunc func(groups []string) string

	switch *mode {
	default:
		printErrorAndExit("mode argument should be ecode/decode/src!\n")

	case "encode":
		regex, _ = regexp.Compile("\\$\\{\\?:(.*?)\\}")
		replaceFunc = func(groups []string) string {
			cipher, _ := CBCEncrypt(keyStr, groups[1])
			return "${AES:" + cipher + "}"
		}
	case "decode":
		regex, _ = regexp.Compile("\\$\\{AES:(.*?)\\}")
		replaceFunc = func(groups []string) string {
			clear, _ := CBCDecrypt(keyStr, groups[1])
			return clear
		}
	case "src":
		regex, _ = regexp.Compile("\\$\\{AES:(.*?)\\}")
		replaceFunc = func(groups []string) string {
			clear, _ := CBCDecrypt(keyStr, groups[1])
			return "${?:" + clear + "}"
		}
	}

	result := ReplaceAllGroupFunc(regex, txt, replaceFunc)
	WriteOutput(*outfile, result)
}

func printErrorAndExit(msg string) {
	fmt.Printf(msg)
	Usage()
	os.Exit(-1)
}

func WriteOutput(outfile, result string) {
	if outfile == "" {
		fmt.Printf("%s\n", result)
	} else {
		err := ioutil.WriteFile(outfile, []byte(result), 0644)
		checkError(err)
	}
}
func checkError(err error) {
	if err != nil {
		panic(err)
	}
}

func FixStrLength(s string, fixLen int) string {
	slen := len(s)
	if slen < fixLen {
		return s + strings.Repeat("0", fixLen-slen)
	}

	if slen > fixLen {
		return s[:fixLen]
	}

	return s
}

var Usage = func() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	flag.PrintDefaults()
}

func ReplaceAllGroupFunc(re *regexp.Regexp, str string, repl func([]string) string) string {
	result := ""
	lastIndex := 0

	for _, v := range re.FindAllSubmatchIndex([]byte(str), -1) {
		groups := []string{}
		for i := 0; i < len(v); i += 2 {
			groups = append(groups, str[v[i]:v[i+1]])
		}

		result += str[lastIndex:v[0]] + repl(groups)
		lastIndex = v[1]
	}

	return result + str[lastIndex:]
}

func CBCEncrypt(strKey, strPlaintext string) (string, error) {
	key := []byte(strKey)
	plaintext := []byte(strPlaintext)

	// CBC mode works on blocks so plaintexts may need to be padded to the
	// next whole block. For an example of such padding, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. Here we'll
	// assume that the plaintext is already of the correct length.
	//if len(plaintext) % aes.BlockSize != 0 {
	//	return "", errors.New("plaintext is not a multiple of the block size")
	//}
	plaintext = PKCS5Padding(plaintext, aes.BlockSize)

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	ciphertext := make([]byte, aes.BlockSize+len(plaintext))
	iv := ciphertext[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		return "", err
	}

	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(ciphertext[aes.BlockSize:], plaintext)

	// It's important to remember that ciphertexts must be authenticated
	// (i.e. by using crypto/hmac) as well as being encrypted in order to
	// be secure.

	base64Text := base64.StdEncoding.EncodeToString(ciphertext)

	return base64Text, nil
}

func CBCDecrypt(strKey, strCiphertext string) (string, error) {
	key := []byte(strKey)
	ciphertext, _ := base64.StdEncoding.DecodeString(strCiphertext)

	block, err := aes.NewCipher(key)
	if err != nil {
		panic(err)
	}

	// The IV needs to be unique, but not secure. Therefore it's common to
	// include it at the beginning of the ciphertext.
	if len(ciphertext) < aes.BlockSize {
		return "", errors.New("ciphertext too short")
	}
	iv := ciphertext[:aes.BlockSize]
	ciphertext = ciphertext[aes.BlockSize:]

	// CBC mode always works in whole blocks.
	if len(ciphertext)%aes.BlockSize != 0 {
		return "", errors.New("ciphertext is not a multiple of the block size")
	}

	mode := cipher.NewCBCDecrypter(block, iv)

	// CryptBlocks can work in-place if the two arguments are the same.
	mode.CryptBlocks(ciphertext, ciphertext)

	ciphertext = PKCS5UnPadding(ciphertext)

	// If the original plaintext lengths are not a multiple of the block
	// size, padding would have to be added when encrypting, which would be
	// removed at this point. For an example, see
	// https://tools.ietf.org/html/rfc5246#section-6.2.3.2. However, it's
	// critical to note that ciphertexts must be authenticated (i.e. by
	// using crypto/hmac) before being decrypted in order to avoid creating
	// a padding oracle.
	return string(ciphertext), nil
}

func PKCS5Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padtext...)
}

func PKCS5UnPadding(origData []byte) []byte {
	length := len(origData)
	// 去掉最后一个字节 unpadding 次
	unpadding := int(origData[length-1])
	return origData[:(length - unpadding)]
}

func ZeroPadding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padtext := bytes.Repeat([]byte{0}, padding)
	return append(ciphertext, padtext...)
}
<|endoftext|>


<commit_before>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package check

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/circonus-agent/internal/config/defaults"
	"github.com/circonus-labs/circonus-agent/internal/release"
	"github.com/circonus-labs/go-apiclient"
	apiconf "github.com/circonus-labs/go-apiclient/config"
	"github.com/pkg/errors"
	"github.com/spf13/viper"
)

// initCheck initializes the check for the agent.
// 1. fetch a check explicitly provided via CID
// 2. search for a check matching the current system
// 3. create a check for the system if --check-create specified
// if fetched, found, or created - set Check.bundle
// otherwise, return an error
func (c *Check) initCheck(cid string, create bool) error {
	var bundle *apiclient.CheckBundle

	// if explicit cid configured, attempt to fetch check bundle using cid
	if cid != "" {
		b, err := c.fetchCheck(cid)
		if err != nil {
			return errors.Wrapf(err, "fetching check for cid %s", cid)
		}
		bundle = b
	} else {
		// if no cid configured, attempt to find check bundle matching this system
		b, found, err := c.findCheck()
		if err != nil {
			if !create || found != 0 {
				return errors.Wrap(err, "unable to find a check for this system")
			}
			c.logger.Info().Msg("no existing check found, creating")
			// attempt to create if not found and create flag ON
			b, err = c.createCheck()
			if err != nil {
				return errors.Wrap(err, "creating new check for this system")
			}
		}
		bundle = b
	}

	if bundle == nil {
		return errors.New("invalid Check object state, bundle is nil")
	}

	c.bundle = bundle

	return nil
}

// func (c *Check) setCheck() error {
// 	// retrieve the check via the Circonus API or create a new check (if configured to do so)
// 	isCreate := viper.GetBool(config.KeyCheckCreate)
// 	isManaged := viper.GetBool(config.KeyCheckEnableNewMetrics)
// 	isReverse := viper.GetBool(config.KeyReverse)
// 	cid := viper.GetString(config.KeyCheckBundleID)
//
// 	var bundle *apiclient.CheckBundle
//
// 	// if explicit cid configured, attempt to fetch check bundle using cid
// 	if cid != "" {
// 		b, err := c.fetchCheck(cid)
// 		if err != nil {
// 			return errors.Wrapf(err, "fetching check for cid %s", cid)
// 		}
// 		bundle = b
// 	} else {
// 		// if no cid configured, attempt to find check bundle matching this system
// 		b, found, err := c.findCheck()
// 		if err != nil {
// 			if !isCreate || found != 0 {
// 				return errors.Wrap(err, "unable to find a check for this system")
// 			}
// 			c.logger.Info().Msg("no existing check found, creating")
// 			// attempt to create if not found and create flag ON
// 			b, err = c.createCheck()
// 			if err != nil {
// 				return errors.Wrap(err, "creating new check for this system")
// 			}
// 		}
// 		bundle = b
// 	}
//
// 	if bundle == nil {
// 		return errors.New("invalid Check object state, bundle is nil")
// 	}
//
// 	c.bundle = bundle
//
// 	if isManaged {
// 		c.logger.Debug().Msg("setting metric states")
// 		err := c.setMetricStates(&bundle.Metrics)
// 		if err != nil {
// 			return errors.Wrap(err, "setting metric states")
// 		}
// 	}
//
// 	// the metrics from the reference bundle are not needed in memory
// 	// as they will never be used again.
// 	c.bundle.Metrics = []apiclient.CheckBundleMetric{}
//
// 	if isReverse {
// 		// populate reverse configuration
// 		c.logger.Debug().Msg("setting reverse config")
// 		err := c.setReverseConfig()
// 		if err != nil {
// 			return errors.Wrap(err, "setting up reverse configuration")
// 		}
// 	}
// 	c.logger.Debug().Msg("done updating check")
//
// 	return nil
// }

func (c *Check) fetchCheck(cid string) (*apiclient.CheckBundle, error) {
	if cid == "" {
		return nil, errors.New("invalid cid (empty)")
	}

	if ok, _ := regexp.MatchString(`^[0-9]+$`, cid); ok {
		cid = "/check_bundle/" + cid
	}

	if ok, _ := regexp.MatchString(`^/check_bundle/[0-9]+$`, cid); !ok {
		return nil, errors.Errorf("invalid cid (%s)", cid)
	}

	bundle, err := c.client.FetchCheckBundle(apiclient.CIDType(&cid))
	if err != nil {
		return nil, errors.Wrapf(err, "unable to retrieve check bundle (%s)", cid)
	}

	return bundle, nil
}

func (c *Check) findCheck() (*apiclient.CheckBundle, int, error) {
	target := viper.GetString(config.KeyCheckTarget)
	if target == "" {
		return nil, -1, errors.New("invalid check target (empty)")
	}

	criteria := apiclient.SearchQueryType(fmt.Sprintf(`(active:1)(type:"json:nad")(target:"%s")`, target))
	bundles, err := c.client.SearchCheckBundles(&criteria, nil)
	if err != nil {
		return nil, -1, errors.Wrap(err, "searching for check bundle")
	}

	found := len(*bundles)

	if found == 0 {
		return nil, found, errors.Errorf("no check bundles matched criteria (%s)", string(criteria))
	}

	if found > 1 {
		return nil, found, errors.Errorf("more than one (%d) check bundle matched criteria (%s)", len(*bundles), string(criteria))
	}

	return &(*bundles)[0], found, nil
}

func (c *Check) createCheck() (*apiclient.CheckBundle, error) {

	// parse the first listen address to use as the required
	// URL in the check config
	var targetAddr string
	{
		serverList := viper.GetStringSlice(config.KeyListen)
		if len(serverList) == 0 {
			serverList = []string{defaults.Listen}
		}
		if serverList[0][0:1] == ":" {
			serverList[0] = "localhost" + serverList[0]
		}
		ta, err := config.ParseListen(serverList[0])
		if err != nil {
			c.logger.Error().Err(err).Str("addr", serverList[0]).Msg("resolving address")
			return nil, errors.Wrap(err, "parsing listen address")
		}
		targetAddr = ta.String()
	}

	target := viper.GetString(config.KeyCheckTarget)
	if target == "" {
		return nil, errors.New("invalid check target (empty)")
	}

	cfg := apiclient.NewCheckBundle()
	cfg.Target = target
	cfg.DisplayName = viper.GetString(config.KeyCheckTitle)
	if cfg.DisplayName == "" {
		cfg.DisplayName = cfg.Target + " /agent"
	}
	note := fmt.Sprintf("created by %s %s", release.NAME, release.VERSION)
	cfg.Notes = &note
	cfg.Type = "json:nad"
	cfg.Config = apiclient.CheckBundleConfig{apiconf.URL: "http://" + targetAddr + "/"}
	cfg.Metrics = []apiclient.CheckBundleMetric{
		{Name: "placeholder", Type: "text", Status: c.statusActiveMetric}, // one metric is required again
	}

	tags := viper.GetString(config.KeyCheckTags)
	if tags != "" {
		cfg.Tags = strings.Split(tags, ",")
	}

	brokerCID := viper.GetString(config.KeyCheckBroker)
	if brokerCID == "" || strings.ToLower(brokerCID) == "select" {
		broker, err := c.selectBroker("json:nad")
		if err != nil {
			return nil, errors.Wrap(err, "selecting broker to create check")
		}

		brokerCID = broker.CID
	}

	if ok, _ := regexp.MatchString(`^[0-9]+$`, brokerCID); ok {
		brokerCID = "/broker/" + brokerCID
	}

	cfg.Brokers = []string{brokerCID}

	if viper.GetBool(config.KeyCheckEnableNewMetrics) {
		cfg.MetricFilters = defaults.CheckMetricFilters
		if viper.GetString(config.KeyCheckMetricFilters) != "" {
			var filters [][]string
			if err := json.Unmarshal([]byte(viper.GetString(config.KeyCheckMetricFilters)), &filters); err != nil {
				return nil, errors.Wrap(err, "parsing check metric filters")
			}
			cfg.MetricFilters = filters
		}
	}

	bundle, err := c.client.CreateCheckBundle(cfg)
	if err != nil {
		return nil, errors.Wrap(err, "creating check bundle")
	}

	return bundle, nil
}
<commit_msg>fix: remove placeholder metrics when creating a check with metric_filters<commit_after>// Copyright © 2017 Circonus, Inc. <support@circonus.com>
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//

package check

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strings"

	"github.com/circonus-labs/circonus-agent/internal/config"
	"github.com/circonus-labs/circonus-agent/internal/config/defaults"
	"github.com/circonus-labs/circonus-agent/internal/release"
	"github.com/circonus-labs/go-apiclient"
	apiconf "github.com/circonus-labs/go-apiclient/config"
	"github.com/pkg/errors"
	"github.com/spf13/viper"
)

// initCheck initializes the check for the agent.
// 1. fetch a check explicitly provided via CID
// 2. search for a check matching the current system
// 3. create a check for the system if --check-create specified
// if fetched, found, or created - set Check.bundle
// otherwise, return an error
func (c *Check) initCheck(cid string, create bool) error {
	var bundle *apiclient.CheckBundle

	// if explicit cid configured, attempt to fetch check bundle using cid
	if cid != "" {
		b, err := c.fetchCheck(cid)
		if err != nil {
			return errors.Wrapf(err, "fetching check for cid %s", cid)
		}
		bundle = b
	} else {
		// if no cid configured, attempt to find check bundle matching this system
		b, found, err := c.findCheck()
		if err != nil {
			if !create || found != 0 {
				return errors.Wrap(err, "unable to find a check for this system")
			}
			c.logger.Info().Msg("no existing check found, creating")
			// attempt to create if not found and create flag ON
			b, err = c.createCheck()
			if err != nil {
				return errors.Wrap(err, "creating new check for this system")
			}
		}
		bundle = b
	}

	if bundle == nil {
		return errors.New("invalid Check object state, bundle is nil")
	}

	c.bundle = bundle

	return nil
}

// func (c *Check) setCheck() error {
// 	// retrieve the check via the Circonus API or create a new check (if configured to do so)
// 	isCreate := viper.GetBool(config.KeyCheckCreate)
// 	isManaged := viper.GetBool(config.KeyCheckEnableNewMetrics)
// 	isReverse := viper.GetBool(config.KeyReverse)
// 	cid := viper.GetString(config.KeyCheckBundleID)
//
// 	var bundle *apiclient.CheckBundle
//
// 	// if explicit cid configured, attempt to fetch check bundle using cid
// 	if cid != "" {
// 		b, err := c.fetchCheck(cid)
// 		if err != nil {
// 			return errors.Wrapf(err, "fetching check for cid %s", cid)
// 		}
// 		bundle = b
// 	} else {
// 		// if no cid configured, attempt to find check bundle matching this system
// 		b, found, err := c.findCheck()
// 		if err != nil {
// 			if !isCreate || found != 0 {
// 				return errors.Wrap(err, "unable to find a check for this system")
// 			}
// 			c.logger.Info().Msg("no existing check found, creating")
// 			// attempt to create if not found and create flag ON
// 			b, err = c.createCheck()
// 			if err != nil {
// 				return errors.Wrap(err, "creating new check for this system")
// 			}
// 		}
// 		bundle = b
// 	}
//
// 	if bundle == nil {
// 		return errors.New("invalid Check object state, bundle is nil")
// 	}
//
// 	c.bundle = bundle
//
// 	if isManaged {
// 		c.logger.Debug().Msg("setting metric states")
// 		err := c.setMetricStates(&bundle.Metrics)
// 		if err != nil {
// 			return errors.Wrap(err, "setting metric states")
// 		}
// 	}
//
// 	// the metrics from the reference bundle are not needed in memory
// 	// as they will never be used again.
// 	c.bundle.Metrics = []apiclient.CheckBundleMetric{}
//
// 	if isReverse {
// 		// populate reverse configuration
// 		c.logger.Debug().Msg("setting reverse config")
// 		err := c.setReverseConfig()
// 		if err != nil {
// 			return errors.Wrap(err, "setting up reverse configuration")
// 		}
// 	}
// 	c.logger.Debug().Msg("done updating check")
//
// 	return nil
// }

func (c *Check) fetchCheck(cid string) (*apiclient.CheckBundle, error) {
	if cid == "" {
		return nil, errors.New("invalid cid (empty)")
	}

	if ok, _ := regexp.MatchString(`^[0-9]+$`, cid); ok {
		cid = "/check_bundle/" + cid
	}

	if ok, _ := regexp.MatchString(`^/check_bundle/[0-9]+$`, cid); !ok {
		return nil, errors.Errorf("invalid cid (%s)", cid)
	}

	bundle, err := c.client.FetchCheckBundle(apiclient.CIDType(&cid))
	if err != nil {
		return nil, errors.Wrapf(err, "unable to retrieve check bundle (%s)", cid)
	}

	return bundle, nil
}

func (c *Check) findCheck() (*apiclient.CheckBundle, int, error) {
	target := viper.GetString(config.KeyCheckTarget)
	if target == "" {
		return nil, -1, errors.New("invalid check target (empty)")
	}

	criteria := apiclient.SearchQueryType(fmt.Sprintf(`(active:1)(type:"json:nad")(target:"%s")`, target))
	bundles, err := c.client.SearchCheckBundles(&criteria, nil)
	if err != nil {
		return nil, -1, errors.Wrap(err, "searching for check bundle")
	}

	found := len(*bundles)

	if found == 0 {
		return nil, found, errors.Errorf("no check bundles matched criteria (%s)", string(criteria))
	}

	if found > 1 {
		return nil, found, errors.Errorf("more than one (%d) check bundle matched criteria (%s)", len(*bundles), string(criteria))
	}

	return &(*bundles)[0], found, nil
}

func (c *Check) createCheck() (*apiclient.CheckBundle, error) {

	// parse the first listen address to use as the required
	// URL in the check config
	var targetAddr string
	{
		serverList := viper.GetStringSlice(config.KeyListen)
		if len(serverList) == 0 {
			serverList = []string{defaults.Listen}
		}
		if serverList[0][0:1] == ":" {
			serverList[0] = "localhost" + serverList[0]
		}
		ta, err := config.ParseListen(serverList[0])
		if err != nil {
			c.logger.Error().Err(err).Str("addr", serverList[0]).Msg("resolving address")
			return nil, errors.Wrap(err, "parsing listen address")
		}
		targetAddr = ta.String()
	}

	target := viper.GetString(config.KeyCheckTarget)
	if target == "" {
		return nil, errors.New("invalid check target (empty)")
	}

	cfg := apiclient.NewCheckBundle()
	cfg.Target = target
	cfg.DisplayName = viper.GetString(config.KeyCheckTitle)
	if cfg.DisplayName == "" {
		cfg.DisplayName = cfg.Target + " /agent"
	}
	note := fmt.Sprintf("created by %s %s", release.NAME, release.VERSION)
	cfg.Notes = &note
	cfg.Type = "json:nad"
	cfg.Config = apiclient.CheckBundleConfig{apiconf.URL: "http://" + targetAddr + "/"}
	cfg.Metrics = []apiclient.CheckBundleMetric{
		{Name: "placeholder", Type: "text", Status: c.statusActiveMetric}, // one metric is required again
	}

	tags := viper.GetString(config.KeyCheckTags)
	if tags != "" {
		cfg.Tags = strings.Split(tags, ",")
	}

	brokerCID := viper.GetString(config.KeyCheckBroker)
	if brokerCID == "" || strings.ToLower(brokerCID) == "select" {
		broker, err := c.selectBroker("json:nad")
		if err != nil {
			return nil, errors.Wrap(err, "selecting broker to create check")
		}

		brokerCID = broker.CID
	}

	if ok, _ := regexp.MatchString(`^[0-9]+$`, brokerCID); ok {
		brokerCID = "/broker/" + brokerCID
	}

	cfg.Brokers = []string{brokerCID}

	if viper.GetBool(config.KeyCheckEnableNewMetrics) {
		cfg.Metrics = []apiclient.CheckBundleMetric{}
		cfg.MetricFilters = defaults.CheckMetricFilters
		if viper.GetString(config.KeyCheckMetricFilters) != "" {
			var filters [][]string
			if err := json.Unmarshal([]byte(viper.GetString(config.KeyCheckMetricFilters)), &filters); err != nil {
				return nil, errors.Wrap(err, "parsing check metric filters")
			}
			cfg.MetricFilters = filters
		}
	}

	bundle, err := c.client.CreateCheckBundle(cfg)
	if err != nil {
		return nil, errors.Wrap(err, "creating check bundle")
	}

	return bundle, nil
}
<|endoftext|>


<commit_before>package flags

import (
	"os"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// DockerAPIMinVersion is the minimum version of the docker api required to
// use watchtower
const DockerAPIMinVersion string = "1.24"

// RegisterDockerFlags that are used directly by the docker api client
func RegisterDockerFlags(rootCmd *cobra.Command) {
	flags := rootCmd.PersistentFlags()
	flags.StringP("host", "H", viper.GetString("DOCKER_HOST"), "daemon socket to connect to")
	flags.BoolP("tlsverify", "v", viper.GetBool("DOCKER_TLS_VERIFY"), "use TLS and verify the remote")
	flags.StringP("api-version", "a", viper.GetString("DOCKER_API_VERSION"), "api version to use by docker client")
}

// RegisterSystemFlags that are used by watchtower to modify the program flow
func RegisterSystemFlags(rootCmd *cobra.Command) {
	flags := rootCmd.PersistentFlags()
	flags.IntP(
		"interval",
		"i",
		viper.GetInt("WATCHTOWER_POLL_INTERVAL"),
		"poll interval (in seconds)")

	flags.StringP("schedule",
		"s",
		viper.GetString("WATCHTOWER_SCHEDULE"),
		"the cron expression which defines when to update")

	flags.DurationP("stop-timeout",
		"t",
		viper.GetDuration("WATCHTOWER_TIMEOUT"),
		"timeout before a container is forcefully stopped")

	flags.BoolP(
		"no-pull",
		"",
		viper.GetBool("WATCHTOWER_NO_PULL"),
		"do not pull any new images")

	flags.BoolP(
		"no-restart",
		"",
		viper.GetBool("WATCHTOWER_NO_RESTART"),
		"do not restart any containers")

	flags.BoolP(
		"cleanup",
		"c",
		viper.GetBool("WATCHTOWER_CLEANUP"),
		"remove previously used images after updating")

	flags.BoolP(
		"remove-volumes",
		"",
		viper.GetBool("WATCHTOWER_REMOVE_VOLUMES"),
		"remove attached volumes before updating")

	flags.BoolP(
		"label-enable",
		"e",
		viper.GetBool("WATCHTOWER_LABEL_ENABLE"),
		"watch containers where the com.centurylinklabs.watchtower.enable label is true")

	flags.BoolP(
		"debug",
		"d",
		viper.GetBool("WATCHTOWER_DEBUG"),
		"enable debug mode with verbose logging")

	flags.BoolP(
		"monitor-only",
		"m",
		viper.GetBool("WATCHTOWER_MONITOR_ONLY"),
		"Will only monitor for new images, not update the containers")

	flags.BoolP(
		"run-once",
		"R",
		viper.GetBool("WATCHTOWER_RUN_ONCE"),
		"Run once now and exit")

	flags.BoolP(
		"include-stopped",
		"S",
		viper.GetBool("WATCHTOWER_INCLUDE_STOPPED"),
		"Will also include created and exited containers")

	flags.BoolP(
		"enable-lifecycle-hooks",
		"",
		viper.GetBool("WATCHTOWER_LIFECYCLE_HOOKS"),
		"Enable the execution of commands triggered by pre- and post-update lifecycle hooks")
}

// RegisterNotificationFlags that are used by watchtower to send notifications
func RegisterNotificationFlags(rootCmd *cobra.Command) {
	flags := rootCmd.PersistentFlags()

	flags.StringSliceP(
		"notifications",
		"n",
		viper.GetStringSlice("WATCHTOWER_NOTIFICATIONS"),
		" notification types to send (valid: email, slack, msteams, gotify)")

	flags.StringP(
		"notifications-level",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATIONS_LEVEL"),
		"The log level used for sending notifications. Possible values: panic, fatal, error, warn, info or debug")

	flags.StringP(
		"notification-email-from",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_FROM"),
		"Address to send notification emails from")

	flags.StringP(
		"notification-email-to",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_TO"),
		"Address to send notification emails to")

	flags.IntP(
		"notification-email-delay",
		"",
		viper.GetInt("WATCHTOWER_NOTIFICATION_EMAIL_DELAY"),
		"Delay before sending notifications, expressed in seconds")

	flags.StringP(
		"notification-email-server",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SERVER"),
		"SMTP server to send notification emails through")

	flags.IntP(
		"notification-email-server-port",
		"",
		viper.GetInt("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_PORT"),
		"SMTP server port to send notification emails through")

	flags.BoolP(
		"notification-email-server-tls-skip-verify",
		"",
		viper.GetBool("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_TLS_SKIP_VERIFY"),
		`
Controls whether watchtower verifies the SMTP server's certificate chain and host name.
Should only be used for testing.
`)

	flags.StringP(
		"notification-email-server-user",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_USER"),
		"SMTP server user for sending notifications")

	flags.StringP(
		"notification-email-server-password",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_PASSWORD"),
		"SMTP server password for sending notifications")

	flags.StringP(
		"notification-email-subjecttag",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SUBJECTTAG"),
		"Subject prefix tag for notifications via mail")

	flags.StringP(
		"notification-slack-hook-url",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_HOOK_URL"),
		"The Slack Hook URL to send notifications to")

	flags.StringP(
		"notification-slack-identifier",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_IDENTIFIER"),
		"A string which will be used to identify the messages coming from this watchtower instance")

	flags.StringP(
		"notification-slack-channel",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_CHANNEL"),
		"A string which overrides the webhook's default channel. Example: #my-custom-channel")

	flags.StringP(
		"notification-slack-icon-emoji",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_ICON_EMOJI"),
		"An emoji code string to use in place of the default icon")

	flags.StringP(
		"notification-slack-icon-url",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_ICON_URL"),
		"An icon image URL string to use in place of the default icon")

	flags.StringP(
		"notification-msteams-hook",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_MSTEAMS_HOOK_URL"),
		"The MSTeams WebHook URL to send notifications to")

	flags.BoolP(
		"notification-msteams-data",
		"",
		viper.GetBool("WATCHTOWER_NOTIFICATION_MSTEAMS_USE_LOG_DATA"),
		"The MSTeams notifier will try to extract log entry fields as MSTeams message facts")

	flags.StringP(
		"notification-gotify-url",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_GOTIFY_URL"),
		"The Gotify URL to send notifications to")
	flags.StringP(
		"notification-gotify-token",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_GOTIFY_TOKEN"),
		"The Gotify Application required to query the Gotify API")
}

// SetDefaults provides default values for environment variables
func SetDefaults() {
	viper.AutomaticEnv()
	viper.SetDefault("DOCKER_HOST", "unix:///var/run/docker.sock")
	viper.SetDefault("DOCKER_API_VERSION", DockerAPIMinVersion)
	viper.SetDefault("WATCHTOWER_POLL_INTERVAL", 300)
	viper.SetDefault("WATCHTOWER_TIMEOUT", time.Second*10)
	viper.SetDefault("WATCHTOWER_NOTIFICATIONS", []string{})
	viper.SetDefault("WATCHTOWER_NOTIFICATIONS_LEVEL", "info")
	viper.SetDefault("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_PORT", 25)
	viper.SetDefault("WATCHTOWER_NOTIFICATION_EMAIL_SUBJECTTAG", "")
	viper.SetDefault("WATCHTOWER_NOTIFICATION_SLACK_IDENTIFIER", "watchtower")
}

// EnvConfig translates the command-line options into environment variables
// that will initialize the api client
func EnvConfig(cmd *cobra.Command) error {
	var err error
	var host string
	var tls bool
	var version string

	flags := cmd.PersistentFlags()

	if host, err = flags.GetString("host"); err != nil {
		return err
	}
	if tls, err = flags.GetBool("tlsverify"); err != nil {
		return err
	}
	if version, err = flags.GetString("api-version"); err != nil {
		return err
	}
	if err = setEnvOptStr("DOCKER_HOST", host); err != nil {
		return err
	}
	if err = setEnvOptBool("DOCKER_TLS_VERIFY", tls); err != nil {
		return err
	}
	if err = setEnvOptStr("DOCKER_API_VERSION", version); err != nil {
		return err
	}
	return nil
}

// ReadFlags reads common flags used in the main program flow of watchtower
func ReadFlags(cmd *cobra.Command) (bool, bool, bool, time.Duration) {
	flags := cmd.PersistentFlags()

	var err error
	var cleanup bool
	var noRestart bool
	var monitorOnly bool
	var timeout time.Duration

	if cleanup, err = flags.GetBool("cleanup"); err != nil {
		log.Fatal(err)
	}
	if noRestart, err = flags.GetBool("no-restart"); err != nil {
		log.Fatal(err)
	}
	if monitorOnly, err = flags.GetBool("monitor-only"); err != nil {
		log.Fatal(err)
	}
	if timeout, err = flags.GetDuration("stop-timeout"); err != nil {
		log.Fatal(err)
	}

	return cleanup, noRestart, monitorOnly, timeout
}

func setEnvOptStr(env string, opt string) error {
	if opt == "" || opt == os.Getenv(env) {
		return nil
	}
	err := os.Setenv(env, opt)
	if err != nil {
		return err
	}
	return nil
}

func setEnvOptBool(env string, opt bool) error {
	if opt {
		return setEnvOptStr(env, "1")
	}
	return nil
}
<commit_msg>Update flags.go<commit_after>package flags

import (
	"os"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// DockerAPIMinVersion is the minimum version of the docker api required to
// use watchtower
const DockerAPIMinVersion string = "1.24"

// RegisterDockerFlags that are used directly by the docker api client
func RegisterDockerFlags(rootCmd *cobra.Command) {
	flags := rootCmd.PersistentFlags()
	flags.StringP("host", "H", viper.GetString("DOCKER_HOST"), "daemon socket to connect to")
	flags.BoolP("tlsverify", "v", viper.GetBool("DOCKER_TLS_VERIFY"), "use TLS and verify the remote")
	flags.StringP("api-version", "a", viper.GetString("DOCKER_API_VERSION"), "api version to use by docker client")
}

// RegisterSystemFlags that are used by watchtower to modify the program flow
func RegisterSystemFlags(rootCmd *cobra.Command) {
	flags := rootCmd.PersistentFlags()
	flags.IntP(
		"interval",
		"i",
		viper.GetInt("WATCHTOWER_POLL_INTERVAL"),
		"poll interval (in seconds)")

	flags.StringP("schedule",
		"s",
		viper.GetString("WATCHTOWER_SCHEDULE"),
		"the cron expression which defines when to update")

	flags.DurationP("stop-timeout",
		"t",
		viper.GetDuration("WATCHTOWER_TIMEOUT"),
		"timeout before a container is forcefully stopped")

	flags.BoolP(
		"no-pull",
		"",
		viper.GetBool("WATCHTOWER_NO_PULL"),
		"do not pull any new images")

	flags.BoolP(
		"no-restart",
		"",
		viper.GetBool("WATCHTOWER_NO_RESTART"),
		"do not restart any containers")

	flags.BoolP(
		"cleanup",
		"c",
		viper.GetBool("WATCHTOWER_CLEANUP"),
		"remove previously used images after updating")

	flags.BoolP(
		"remove-volumes",
		"",
		viper.GetBool("WATCHTOWER_REMOVE_VOLUMES"),
		"remove attached volumes before updating")

	flags.BoolP(
		"label-enable",
		"e",
		viper.GetBool("WATCHTOWER_LABEL_ENABLE"),
		"watch containers where the com.centurylinklabs.watchtower.enable label is true")

	flags.BoolP(
		"debug",
		"d",
		viper.GetBool("WATCHTOWER_DEBUG"),
		"enable debug mode with verbose logging")

	flags.BoolP(
		"monitor-only",
		"m",
		viper.GetBool("WATCHTOWER_MONITOR_ONLY"),
		"Will only monitor for new images, not update the containers")

	flags.BoolP(
		"run-once",
		"R",
		viper.GetBool("WATCHTOWER_RUN_ONCE"),
		"Run once now and exit")

	flags.BoolP(
		"include-stopped",
		"S",
		viper.GetBool("WATCHTOWER_INCLUDE_STOPPED"),
		"Will also include created and exited containers")

	flags.BoolP(
		"enable-lifecycle-hooks",
		"",
		viper.GetBool("WATCHTOWER_LIFECYCLE_HOOKS"),
		"Enable the execution of commands triggered by pre- and post-update lifecycle hooks")
}

// RegisterNotificationFlags that are used by watchtower to send notifications
func RegisterNotificationFlags(rootCmd *cobra.Command) {
	flags := rootCmd.PersistentFlags()

	flags.StringSliceP(
		"notifications",
		"n",
		viper.GetStringSlice("WATCHTOWER_NOTIFICATIONS"),
		" notification types to send (valid: email, slack, msteams, gotify)")

	flags.StringP(
		"notifications-level",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATIONS_LEVEL"),
		"The log level used for sending notifications. Possible values: panic, fatal, error, warn, info or debug")

	flags.StringP(
		"notification-email-from",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_FROM"),
		"Address to send notification emails from")

	flags.StringP(
		"notification-email-to",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_TO"),
		"Address to send notification emails to")

	flags.IntP(
		"notification-email-delay",
		"",
		viper.GetInt("WATCHTOWER_NOTIFICATION_EMAIL_DELAY"),
		"Delay before sending notifications, expressed in seconds")

	flags.StringP(
		"notification-email-server",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SERVER"),
		"SMTP server to send notification emails through")

	flags.IntP(
		"notification-email-server-port",
		"",
		viper.GetInt("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_PORT"),
		"SMTP server port to send notification emails through")

	flags.BoolP(
		"notification-email-server-tls-skip-verify",
		"",
		viper.GetBool("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_TLS_SKIP_VERIFY"),
		`
Controls whether watchtower verifies the SMTP server's certificate chain and host name.
Should only be used for testing.
`)

	flags.StringP(
		"notification-email-server-user",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_USER"),
		"SMTP server user for sending notifications")

	flags.StringP(
		"notification-email-server-password",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_PASSWORD"),
		"SMTP server password for sending notifications")

	flags.StringP(
		"notification-email-subjecttag",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_EMAIL_SUBJECTTAG"),
		"Subject prefix tag for notifications via mail")

	flags.StringP(
		"notification-slack-hook-url",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_HOOK_URL"),
		"The Slack Hook URL to send notifications to")

	flags.StringP(
		"notification-slack-identifier",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_IDENTIFIER"),
		"A string which will be used to identify the messages coming from this watchtower instance")

	flags.StringP(
		"notification-slack-channel",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_CHANNEL"),
		"A string which overrides the webhook's default channel. Example: #my-custom-channel")

	flags.StringP(
		"notification-slack-icon-emoji",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_ICON_EMOJI"),
		"An emoji code string to use in place of the default icon")

	flags.StringP(
		"notification-slack-icon-url",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_SLACK_ICON_URL"),
		"An icon image URL string to use in place of the default icon")

	flags.StringP(
		"notification-msteams-hook",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_MSTEAMS_HOOK_URL"),
		"The MSTeams WebHook URL to send notifications to")

	flags.BoolP(
		"notification-msteams-data",
		"",
		viper.GetBool("WATCHTOWER_NOTIFICATION_MSTEAMS_USE_LOG_DATA"),
		"The MSTeams notifier will try to extract log entry fields as MSTeams message facts")

	flags.StringP(
		"notification-gotify-url",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_GOTIFY_URL"),
		"The Gotify URL to send notifications to")
	flags.StringP(
		"notification-gotify-token",
		"",
		viper.GetString("WATCHTOWER_NOTIFICATION_GOTIFY_TOKEN"),
		"The Gotify Application required to query the Gotify API")
}

// SetDefaults provides default values for environment variables
func SetDefaults() {
	viper.AutomaticEnv()
	viper.SetDefault("DOCKER_HOST", "unix:///var/run/docker.sock")
	viper.SetDefault("DOCKER_API_VERSION", DockerAPIMinVersion)
	viper.SetDefault("WATCHTOWER_POLL_INTERVAL", 300)
	viper.SetDefault("WATCHTOWER_TIMEOUT", time.Second*10)
	viper.SetDefault("WATCHTOWER_NOTIFICATIONS", []string{})
	viper.SetDefault("WATCHTOWER_NOTIFICATIONS_LEVEL", "info")
	viper.SetDefault("WATCHTOWER_NOTIFICATION_EMAIL_SERVER_PORT", 25)
	viper.SetDefault("WATCHTOWER_NOTIFICATION_EMAIL_SUBJECTTAG", "")
	viper.SetDefault("WATCHTOWER_NOTIFICATION_SLACK_IDENTIFIER", "watchtower")
}

// EnvConfig translates the command-line options into environment variables
// that will initialize the api client
func EnvConfig(cmd *cobra.Command) error {
	var err error
	var host string
	var tls bool
	var version string

	flags := cmd.PersistentFlags()

	if host, err = flags.GetString("host"); err != nil {
		return err
	}
	if tls, err = flags.GetBool("tlsverify"); err != nil {
		return err
	}
	if version, err = flags.GetString("api-version"); err != nil {
		return err
	}
	if err = setEnvOptStr("DOCKER_HOST", host); err != nil {
		return err
	}
	if err = setEnvOptBool("DOCKER_TLS_VERIFY", tls); err != nil {
		return err
	}
	if err = setEnvOptStr("DOCKER_API_VERSION", version); err != nil {
		return err
	}
	return nil
}

// ReadFlags reads common flags used in the main program flow of watchtower
func ReadFlags(cmd *cobra.Command) (bool, bool, bool, time.Duration) {
	flags := cmd.PersistentFlags()

	var err error
	var cleanup bool
	var noRestart bool
	var monitorOnly bool
	var timeout time.Duration

	if cleanup, err = flags.GetBool("cleanup"); err != nil {
		log.Fatal(err)
	}
	if noRestart, err = flags.GetBool("no-restart"); err != nil {
		log.Fatal(err)
	}
	if monitorOnly, err = flags.GetBool("monitor-only"); err != nil {
		log.Fatal(err)
	}
	if timeout, err = flags.GetDuration("stop-timeout"); err != nil {
		log.Fatal(err)
	}

	return cleanup, noRestart, monitorOnly, timeout
}

func setEnvOptStr(env string, opt string) error {
	if opt == "" || opt == os.Getenv(env) {
		return nil
	}
	err := os.Setenv(env, opt)
	if err != nil {
		return err
	}
	return nil
}

func setEnvOptBool(env string, opt bool) error {
	if opt {
		return setEnvOptStr(env, "1")
	}
	return nil
}
<|endoftext|>