<commit_before>package amazonproduct

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"
)

type AmazonProductAPI struct {
	AccessKey    string
	SecretKey    string
	AssociateTag string
	Host         string
}

func (api AmazonProductAPI) ItemLookup(ItemId string) (string, error) {
	params := map[string]string{
		"ItemId":        ItemId,
		"ResponseGroup": "Images,ItemAttributes,Small,EditorialReview",
	}

	return api.genSignAndFetch("ItemLookup", params)
}

func (api AmazonProductAPI) ItemSearchByKeyword(Keywords string) (string, error) {
	params := map[string]string{
		"Keywords":      Keywords,
		"ResponseGroup": "Images,ItemAttributes,Small,EditorialReview",
	}
	return api.ItemSearch("All", params)
}

func (api AmazonProductAPI) ItemSearchByKeywordWithResponseGroup(Keywords string, ResponseGroup string) (string, error) {
	params := map[string]string{
		"Keywords":      Keywords,
		"ResponseGroup": ResponseGroup,
	}
	return api.ItemSearch("All", params)
}

func (api AmazonProductAPI) ItemSearch(SearchIndex string, Parameters map[string]string) (string, error) {
	Parameters["SearchIndex"] = SearchIndex
	return api.genSignAndFetch("ItemSearch", Parameters)
}

func (api AmazonProductAPI) genSignAndFetch(Operation string, Parameters map[string]string) (string, error) {
	genUrl, err := GenerateAmazonUrl(api, Operation, Parameters)
	if err != nil {
		return "", err
	}

	SetTimestamp(genUrl)

	signedurl, err := SignAmazonUrl(genUrl, api)
	if err != nil {
		return "", err
	}

	resp, err := http.Get(signedurl)
	if err != nil {
		return "", err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return string(body), nil
}

func GenerateAmazonUrl(api AmazonProductAPI, Operation string, Parameters map[string]string) (finalUrl *url.URL, err error) {

	result, err := url.Parse(api.Host)
	if err != nil {
		return nil, err
	}

	result.Host = api.Host
	result.Scheme = "http"
	result.Path = "/onca/xml"

	values := url.Values{}
	values.Add("Operation", Operation)
	values.Add("Service", "AWSECommerceService")
	values.Add("AWSAccessKeyId", api.AccessKey)
	values.Add("Version", "2009-01-01")
	values.Add("AssociateTag", api.AssociateTag)

	for k, v := range Parameters {
		values.Set(k, v)
	}

	params := values.Encode()
	result.RawQuery = params

	return result, nil
}

func SetTimestamp(origUrl *url.URL) (err error) {
	values, err := url.ParseQuery(origUrl.RawQuery)
	if err != nil {
		return err
	}
	values.Set("Timestamp", time.Now().UTC().Format(time.RFC3339))
	origUrl.RawQuery = values.Encode()

	return nil
}

func SignAmazonUrl(origUrl *url.URL, api AmazonProductAPI) (signedUrl string, err error) {

	escapeUrl := strings.Replace(origUrl.RawQuery, ",", "%2C", -1)
	escapeUrl = strings.Replace(escapeUrl, ":", "%3A", -1)

	params := strings.Split(escapeUrl, "&")
	sort.Strings(params)
	sortedParams := strings.Join(params, "&")

	toSign := fmt.Sprintf("GET\n%s\n%s\n%s", origUrl.Host, origUrl.Path, sortedParams)

	hasher := hmac.New(sha256.New, []byte(api.SecretKey))
	_, err = hasher.Write([]byte(toSign))
	if err != nil {
		return "", err
	}

	hash := base64.StdEncoding.EncodeToString(hasher.Sum(nil))

	hash = url.QueryEscape(hash)

	newParams := fmt.Sprintf("%s&Signature=%s", sortedParams, hash)

	origUrl.RawQuery = newParams

	return origUrl.String(), nil
}
<commit_msg>Add package description<commit_after>//Package amazonproduct provides methods for interacting with the Amazon Product Advertising API
package amazonproduct

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"
)

type AmazonProductAPI struct {
	AccessKey    string
	SecretKey    string
	AssociateTag string
	Host         string
}

func (api AmazonProductAPI) ItemLookup(ItemId string) (string, error) {
	params := map[string]string{
		"ItemId":        ItemId,
		"ResponseGroup": "Images,ItemAttributes,Small,EditorialReview",
	}

	return api.genSignAndFetch("ItemLookup", params)
}

func (api AmazonProductAPI) ItemSearchByKeyword(Keywords string) (string, error) {
	params := map[string]string{
		"Keywords":      Keywords,
		"ResponseGroup": "Images,ItemAttributes,Small,EditorialReview",
	}
	return api.ItemSearch("All", params)
}

func (api AmazonProductAPI) ItemSearchByKeywordWithResponseGroup(Keywords string, ResponseGroup string) (string, error) {
	params := map[string]string{
		"Keywords":      Keywords,
		"ResponseGroup": ResponseGroup,
	}
	return api.ItemSearch("All", params)
}

func (api AmazonProductAPI) ItemSearch(SearchIndex string, Parameters map[string]string) (string, error) {
	Parameters["SearchIndex"] = SearchIndex
	return api.genSignAndFetch("ItemSearch", Parameters)
}

func (api AmazonProductAPI) genSignAndFetch(Operation string, Parameters map[string]string) (string, error) {
	genUrl, err := GenerateAmazonUrl(api, Operation, Parameters)
	if err != nil {
		return "", err
	}

	SetTimestamp(genUrl)

	signedurl, err := SignAmazonUrl(genUrl, api)
	if err != nil {
		return "", err
	}

	resp, err := http.Get(signedurl)
	if err != nil {
		return "", err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return string(body), nil
}

func GenerateAmazonUrl(api AmazonProductAPI, Operation string, Parameters map[string]string) (finalUrl *url.URL, err error) {

	result, err := url.Parse(api.Host)
	if err != nil {
		return nil, err
	}

	result.Host = api.Host
	result.Scheme = "http"
	result.Path = "/onca/xml"

	values := url.Values{}
	values.Add("Operation", Operation)
	values.Add("Service", "AWSECommerceService")
	values.Add("AWSAccessKeyId", api.AccessKey)
	values.Add("Version", "2009-01-01")
	values.Add("AssociateTag", api.AssociateTag)

	for k, v := range Parameters {
		values.Set(k, v)
	}

	params := values.Encode()
	result.RawQuery = params

	return result, nil
}

func SetTimestamp(origUrl *url.URL) (err error) {
	values, err := url.ParseQuery(origUrl.RawQuery)
	if err != nil {
		return err
	}
	values.Set("Timestamp", time.Now().UTC().Format(time.RFC3339))
	origUrl.RawQuery = values.Encode()

	return nil
}

func SignAmazonUrl(origUrl *url.URL, api AmazonProductAPI) (signedUrl string, err error) {

	escapeUrl := strings.Replace(origUrl.RawQuery, ",", "%2C", -1)
	escapeUrl = strings.Replace(escapeUrl, ":", "%3A", -1)

	params := strings.Split(escapeUrl, "&")
	sort.Strings(params)
	sortedParams := strings.Join(params, "&")

	toSign := fmt.Sprintf("GET\n%s\n%s\n%s", origUrl.Host, origUrl.Path, sortedParams)

	hasher := hmac.New(sha256.New, []byte(api.SecretKey))
	_, err = hasher.Write([]byte(toSign))
	if err != nil {
		return "", err
	}

	hash := base64.StdEncoding.EncodeToString(hasher.Sum(nil))

	hash = url.QueryEscape(hash)

	newParams := fmt.Sprintf("%s&Signature=%s", sortedParams, hash)

	origUrl.RawQuery = newParams

	return origUrl.String(), nil
}
<|endoftext|>


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

package zdns

import (
	"bufio"
	"encoding/json"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
)

type routineMetadata struct {
	Names  int
	Status map[Status]int
}

func GetDNSServers(path string) ([]string, error) {
	c, err := dns.ClientConfigFromFile(path)
	if err != nil {
		return []string{}, err
	}
	var servers []string
	for _, s := range c.Servers {
		if s[0:1] != "[" && strings.Contains(s, ":") {
			s = "[" + s + "]"
		}
		full := strings.Join([]string{s, c.Port}, ":")
		servers = append(servers, full)
	}
	return servers, nil
}

func parseAlexa(line string) (string, int) {
	s := strings.SplitN(line, ",", 2)
	rank, err := strconv.Atoi(s[0])
	if err != nil {
		log.Fatal("Malformed Alexa Top Million file")
	}
	return s[1], rank
}

func makeName(name string, prefix string) (string, bool) {
	if prefix == "" {
		return name, false
	} else {
		return strings.Join([]string{prefix, name}, ""), true
	}
}

func doLookup(g *GlobalLookupFactory, gc *GlobalConf, input <-chan interface{}, output chan<- string, metaChan chan<- routineMetadata, wg *sync.WaitGroup, threadID int) error {
	f, err := (*g).MakeRoutineFactory(threadID)
	if err != nil {
		log.Fatal("Unable to create new routine factory", err.Error())
	}
	var metadata routineMetadata
	metadata.Status = make(map[Status]int)
	for genericInput := range input {
		var res Result
		var innerRes interface{}
		var trace []interface{}
		var status Status
		var err error
		l, err := f.MakeLookup()
		if err != nil {
			log.Fatal("Unable to build lookup instance", err)
		}
		if (*g).ZonefileInput() {
			length := len(genericInput.(*dns.Token).RR.Header().Name)
			if length == 0 {
				continue
			}
			res.Name = genericInput.(*dns.Token).RR.Header().Name[0 : length-1]
			res.Class = dns.Class(gc.Class).String()
			switch typ := genericInput.(*dns.Token).RR.(type) {
			case *dns.NS:
				ns := strings.ToLower(typ.Ns)
				res.Nameserver = ns[:len(ns)-1]
			}
			innerRes, status, err = l.DoZonefileLookup(genericInput.(*dns.Token))
		} else {
			line := genericInput.(string)
			var changed bool
			var rawName string
			var rank int
			if gc.AlexaFormat == true {
				rawName, rank = parseAlexa(line)
				res.AlexaRank = rank
			} else {
				rawName = line
			}
			lookupName, changed := makeName(rawName, gc.NamePrefix)
			if changed {
				res.AlteredName = lookupName
			}
			res.Name = rawName
			res.Class = dns.Class(gc.Class).String()
			innerRes, trace, status, err = l.DoLookup(lookupName)
		}
		res.Timestamp = time.Now().Format(gc.TimeFormat)
		if status != STATUS_NO_OUTPUT {
			res.Status = string(status)
			res.Data = innerRes
			res.Trace = trace
			if err != nil {
				res.Error = err.Error()
			}
			jsonRes, err := json.Marshal(res)
			if err != nil {
				log.Fatal("Unable to marshal JSON result", err)
			}
			output <- string(jsonRes)
		}
		metadata.Names++
		metadata.Status[status]++
	}
	metaChan <- metadata
	(*wg).Done()
	return nil
}

// write results from lookup to output file
func doOutput(out <-chan string, path string, wg *sync.WaitGroup) error {
	var f *os.File
	if path == "" || path == "-" {
		f = os.Stdout
	} else {
		var err error
		f, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			log.Fatal("unable to open output file:", err.Error())
		}
		defer f.Close()
	}
	for n := range out {
		f.WriteString(n + "\n")
	}
	(*wg).Done()
	return nil
}

// read input file and put results into channel
func doInput(in chan<- interface{}, path string, wg *sync.WaitGroup, zonefileInput bool, passedName string) error {

	defer func() {
		close(in)
		(*wg).Done()
	}()

	if len(passedName) > 0 {
		in <- passedName
		return nil
	}

	var f *os.File
	if path == "" || path == "-" {
		f = os.Stdin
	} else {
		var err error
		f, err = os.Open(path)
		if err != nil {
			log.Fatal("unable to open input file:", err.Error())
		}
	}
	if zonefileInput {
		tokens := dns.ParseZone(f, ".", path)
		for t := range tokens {
			in <- t
		}
	} else {
		s := bufio.NewScanner(f)
		for s.Scan() {
			in <- s.Text()
		}
		if err := s.Err(); err != nil {
			log.Fatal("input unable to read file", err)
		}
	}

	return nil
}

func aggregateMetadata(c <-chan routineMetadata) Metadata {
	var meta Metadata
	meta.Status = make(map[string]int)
	for m := range c {
		meta.Names += m.Names
		for k, v := range m.Status {
			meta.Status[string(k)] += v
		}
	}
	return meta
}

func DoLookups(g *GlobalLookupFactory, c *GlobalConf) error {
	// doInput: take lines from input -> inChan
	//	- closes channel when done processing
	// doOutput: take serialized JSON from outChan and write
	// DoLookup:
	//	- n threads that do processing from in and place results in out
	//	- process until inChan closes, then wg.done()
	// Once we processing threads have all finished, wait until the
	// output and metadata threads have completed
	inChan := make(chan interface{})
	outChan := make(chan string)
	metaChan := make(chan routineMetadata, c.Threads)
	var routineWG sync.WaitGroup
	go doOutput(outChan, c.OutputFilePath, &routineWG)
	go doInput(inChan, c.InputFilePath, &routineWG, (*g).ZonefileInput(), c.PassedName)
	routineWG.Add(2)
	// create pool of worker goroutines
	var lookupWG sync.WaitGroup
	lookupWG.Add(c.Threads)
	startTime := time.Now().Format(c.TimeFormat)
	for i := 0; i < c.Threads; i++ {
		go doLookup(g, c, inChan, outChan, metaChan, &lookupWG, i)
	}
	lookupWG.Wait()
	close(outChan)
	close(metaChan)
	routineWG.Wait()
	if c.MetadataFilePath != "" {
		// we're done processing data. aggregate all the data from individual routines
		metaData := aggregateMetadata(metaChan)
		metaData.StartTime = startTime
		metaData.EndTime = time.Now().Format(c.TimeFormat)
		metaData.NameServers = c.NameServers
		metaData.Retries = c.Retries
		// Seconds() returns a float. However, timeout is passed in as an integer
		// command line argument, so there should be no loss of data when casting
		// back to an integer here.
		metaData.Timeout = int(c.Timeout.Seconds())
		metaData.Conf = c
		// add global lookup-related metadata
		// write out metadata
		var f *os.File
		if c.MetadataFilePath == "-" {
			f = os.Stderr
		} else {
			var err error
			f, err = os.OpenFile(c.MetadataFilePath, os.O_WRONLY|os.O_CREATE, 0666)
			if err != nil {
				log.Fatal("unable to open metadata file:", err.Error())
			}
			defer f.Close()
		}
		j, err := json.Marshal(metaData)
		if err != nil {
			log.Fatal("unable to JSON encode metadata:", err.Error())
		}
		f.WriteString(string(j))
	}
	return nil
}
<commit_msg>Fix metadata file creation and permissions (#140)<commit_after>/*
 * ZDNS Copyright 2016 Regents of the University of Michigan
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package zdns

import (
	"bufio"
	"encoding/json"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/miekg/dns"
	log "github.com/sirupsen/logrus"
)

type routineMetadata struct {
	Names  int
	Status map[Status]int
}

func GetDNSServers(path string) ([]string, error) {
	c, err := dns.ClientConfigFromFile(path)
	if err != nil {
		return []string{}, err
	}
	var servers []string
	for _, s := range c.Servers {
		if s[0:1] != "[" && strings.Contains(s, ":") {
			s = "[" + s + "]"
		}
		full := strings.Join([]string{s, c.Port}, ":")
		servers = append(servers, full)
	}
	return servers, nil
}

func parseAlexa(line string) (string, int) {
	s := strings.SplitN(line, ",", 2)
	rank, err := strconv.Atoi(s[0])
	if err != nil {
		log.Fatal("Malformed Alexa Top Million file")
	}
	return s[1], rank
}

func makeName(name string, prefix string) (string, bool) {
	if prefix == "" {
		return name, false
	} else {
		return strings.Join([]string{prefix, name}, ""), true
	}
}

func doLookup(g *GlobalLookupFactory, gc *GlobalConf, input <-chan interface{}, output chan<- string, metaChan chan<- routineMetadata, wg *sync.WaitGroup, threadID int) error {
	f, err := (*g).MakeRoutineFactory(threadID)
	if err != nil {
		log.Fatal("Unable to create new routine factory", err.Error())
	}
	var metadata routineMetadata
	metadata.Status = make(map[Status]int)
	for genericInput := range input {
		var res Result
		var innerRes interface{}
		var trace []interface{}
		var status Status
		var err error
		l, err := f.MakeLookup()
		if err != nil {
			log.Fatal("Unable to build lookup instance", err)
		}
		if (*g).ZonefileInput() {
			length := len(genericInput.(*dns.Token).RR.Header().Name)
			if length == 0 {
				continue
			}
			res.Name = genericInput.(*dns.Token).RR.Header().Name[0 : length-1]
			res.Class = dns.Class(gc.Class).String()
			switch typ := genericInput.(*dns.Token).RR.(type) {
			case *dns.NS:
				ns := strings.ToLower(typ.Ns)
				res.Nameserver = ns[:len(ns)-1]
			}
			innerRes, status, err = l.DoZonefileLookup(genericInput.(*dns.Token))
		} else {
			line := genericInput.(string)
			var changed bool
			var rawName string
			var rank int
			if gc.AlexaFormat == true {
				rawName, rank = parseAlexa(line)
				res.AlexaRank = rank
			} else {
				rawName = line
			}
			lookupName, changed := makeName(rawName, gc.NamePrefix)
			if changed {
				res.AlteredName = lookupName
			}
			res.Name = rawName
			res.Class = dns.Class(gc.Class).String()
			innerRes, trace, status, err = l.DoLookup(lookupName)
		}
		res.Timestamp = time.Now().Format(gc.TimeFormat)
		if status != STATUS_NO_OUTPUT {
			res.Status = string(status)
			res.Data = innerRes
			res.Trace = trace
			if err != nil {
				res.Error = err.Error()
			}
			jsonRes, err := json.Marshal(res)
			if err != nil {
				log.Fatal("Unable to marshal JSON result", err)
			}
			output <- string(jsonRes)
		}
		metadata.Names++
		metadata.Status[status]++
	}
	metaChan <- metadata
	(*wg).Done()
	return nil
}

// write results from lookup to output file
func doOutput(out <-chan string, path string, wg *sync.WaitGroup) error {
	var f *os.File
	if path == "" || path == "-" {
		f = os.Stdout
	} else {
		var err error
		f, err = os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
		if err != nil {
			log.Fatal("unable to open output file:", err.Error())
		}
		defer f.Close()
	}
	for n := range out {
		f.WriteString(n + "\n")
	}
	(*wg).Done()
	return nil
}

// read input file and put results into channel
func doInput(in chan<- interface{}, path string, wg *sync.WaitGroup, zonefileInput bool, passedName string) error {

	defer func() {
		close(in)
		(*wg).Done()
	}()

	if len(passedName) > 0 {
		in <- passedName
		return nil
	}

	var f *os.File
	if path == "" || path == "-" {
		f = os.Stdin
	} else {
		var err error
		f, err = os.Open(path)
		if err != nil {
			log.Fatal("unable to open input file:", err.Error())
		}
	}
	if zonefileInput {
		tokens := dns.ParseZone(f, ".", path)
		for t := range tokens {
			in <- t
		}
	} else {
		s := bufio.NewScanner(f)
		for s.Scan() {
			in <- s.Text()
		}
		if err := s.Err(); err != nil {
			log.Fatal("input unable to read file", err)
		}
	}

	return nil
}

func aggregateMetadata(c <-chan routineMetadata) Metadata {
	var meta Metadata
	meta.Status = make(map[string]int)
	for m := range c {
		meta.Names += m.Names
		for k, v := range m.Status {
			meta.Status[string(k)] += v
		}
	}
	return meta
}

func DoLookups(g *GlobalLookupFactory, c *GlobalConf) error {
	// doInput: take lines from input -> inChan
	//	- closes channel when done processing
	// doOutput: take serialized JSON from outChan and write
	// DoLookup:
	//	- n threads that do processing from in and place results in out
	//	- process until inChan closes, then wg.done()
	// Once we processing threads have all finished, wait until the
	// output and metadata threads have completed
	inChan := make(chan interface{})
	outChan := make(chan string)
	metaChan := make(chan routineMetadata, c.Threads)
	var routineWG sync.WaitGroup
	go doOutput(outChan, c.OutputFilePath, &routineWG)
	go doInput(inChan, c.InputFilePath, &routineWG, (*g).ZonefileInput(), c.PassedName)
	routineWG.Add(2)
	// create pool of worker goroutines
	var lookupWG sync.WaitGroup
	lookupWG.Add(c.Threads)
	startTime := time.Now().Format(c.TimeFormat)
	for i := 0; i < c.Threads; i++ {
		go doLookup(g, c, inChan, outChan, metaChan, &lookupWG, i)
	}
	lookupWG.Wait()
	close(outChan)
	close(metaChan)
	routineWG.Wait()
	if c.MetadataFilePath != "" {
		// we're done processing data. aggregate all the data from individual routines
		metaData := aggregateMetadata(metaChan)
		metaData.StartTime = startTime
		metaData.EndTime = time.Now().Format(c.TimeFormat)
		metaData.NameServers = c.NameServers
		metaData.Retries = c.Retries
		// Seconds() returns a float. However, timeout is passed in as an integer
		// command line argument, so there should be no loss of data when casting
		// back to an integer here.
		metaData.Timeout = int(c.Timeout.Seconds())
		metaData.Conf = c
		// add global lookup-related metadata
		// write out metadata
		var f *os.File
		if c.MetadataFilePath == "-" {
			f = os.Stderr
		} else {
			var err error
			f, err = os.OpenFile(c.MetadataFilePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
			if err != nil {
				log.Fatal("unable to open metadata file:", err.Error())
			}
			defer f.Close()
		}
		j, err := json.Marshal(metaData)
		if err != nil {
			log.Fatal("unable to JSON encode metadata:", err.Error())
		}
		f.WriteString(string(j))
	}
	return nil
}
<|endoftext|>


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

package cloudimagemetadata

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	jujutxn "github.com/juju/txn"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/mongo"
)

var logger = loggo.GetLogger("juju.state.cloudimagemetadata")

type storage struct {
	envuuid        string
	collection     string
	runTransaction func(jujutxn.TransactionSource) error
	getCollection  func(string) (_ mongo.Collection, closer func())
}

var _ Storage = (*storage)(nil)

// NewStorage constructs a new Storage that stores image metadata
// in the provided collection using the provided transaction runner.
func NewStorage(
	envuuid string,
	collectionName string,
	runTransaction func(jujutxn.TransactionSource) error,
	getCollection func(string) (_ mongo.Collection, closer func()),
) Storage {
	return &storage{envuuid, collectionName, runTransaction, getCollection}
}

var emptyMetadata = Metadata{}

// SaveMetadata implements Storage.SaveMetadata and behaves as save-or-update:
// if desired metadata does not exist - we insert it; if it exists - we update.
// It throws an error if metadata did not change.
func (s *storage) SaveMetadata(metadata Metadata) error {
	newDoc := s.mongoDoc(metadata)
	all, closeAll := s.getCollection(s.collection)
	defer closeAll()

	buildTxn := func(attempt int) ([]txn.Op, error) {
		existing, err := s.getMetadata(newDoc.Id)
		if err != nil {
			return nil, errors.Trace(err)
		}

		op := txn.Op{
			C:  all.Name(),
			Id: newDoc.Id,
		}
		if existing == emptyMetadata {
			op.Assert = txn.DocMissing
			op.Insert = &newDoc
			logger.Debugf("inserting cloud image metadata for %v", newDoc.Id)
		} else {
			same := isSameMetadata(existing, metadata)
			if same {
				logger.Debugf("cloud image metadata for %v has not changed", newDoc.Id)
				return []txn.Op{}, errors.Annotate(jujutxn.ErrNoOperations, "no changes were made")
			}
			op.Assert = txn.DocExists
			op.Update = bson.D{{"$set", newDoc.updates()}}
			logger.Debugf("updating cloud image metadata for %v", newDoc.Id)
		}

		return []txn.Op{op}, nil
	}

	err := s.runTransaction(buildTxn)
	if err != nil {
		return errors.Annotatef(err, "cannot add cloud image metadata for %v", newDoc.Id)
	}
	return nil
}

func (s *storage) getMetadata(id string) (Metadata, error) {
	coll, closer := s.getCollection(s.collection)
	defer closer()

	var old imagesMetadataDoc
	err := coll.Find(bson.D{{"_id", id}}).One(&old)
	if err != nil {
		if err == mgo.ErrNotFound {
			return emptyMetadata, nil
		}
		return emptyMetadata, errors.Trace(err)
	}
	return old.metadata(), nil
}

func isSameMetadata(old, new Metadata) bool {
	return old.ImageId == new.ImageId &&
		areSameAttributes(old.MetadataAttributes, new.MetadataAttributes)
}

func areSameAttributes(old, new MetadataAttributes) bool {
	return old.Arch == new.Arch &&
		old.Region == new.Region &&
		old.RootStorageType == new.RootStorageType &&
		old.RootStorageSize == new.RootStorageSize &&
		old.Series == new.Series &&
		old.Stream == new.Stream &&
		old.VirtualType == new.VirtualType
}

// AllMetadata implements Storage.AllMetadata.
func (s *storage) AllMetadata() ([]Metadata, error) {
	return s.FindMetadata(MetadataAttributes{})
}

// FindMetadata implements Storage.FindMetadata.
func (s *storage) FindMetadata(criteria MetadataAttributes) ([]Metadata, error) {
	coll, closer := s.getCollection(s.collection)
	defer closer()

	searchCriteria := buildSearchClauses(criteria)
	var docs []imagesMetadataDoc
	if err := coll.Find(searchCriteria).All(&docs); err != nil {
		return nil, errors.Trace(err)
	}
	if searchCriteria != nil && len(docs) == 0 {
		// If criteria had values and no metadata was found, err
		return nil, errors.NotFoundf("matching cloud image metadata")
	}
	metadata := make([]Metadata, len(docs))
	for i, doc := range docs {
		metadata[i] = doc.metadata()
	}
	return metadata, nil
}

func buildSearchClauses(criteria MetadataAttributes) bson.D {
	all := bson.D{}

	if criteria.Stream != "" {
		all = append(all, bson.DocElem{"stream", criteria.Stream})
	}

	if criteria.Region != "" {
		all = append(all, bson.DocElem{"region", criteria.Region})
	}

	if criteria.Series != "" {
		all = append(all, bson.DocElem{"series", criteria.Series})
	}

	if criteria.Arch != "" {
		all = append(all, bson.DocElem{"arch", criteria.Arch})
	}

	if criteria.VirtualType != "" {
		all = append(all, bson.DocElem{"virtual_type", criteria.VirtualType})
	}

	if criteria.RootStorageType != "" {
		all = append(all, bson.DocElem{"root_storage_type", criteria.RootStorageType})
	}

	if criteria.RootStorageSize != "" {
		all = append(all, bson.DocElem{"root_storage_size", criteria.RootStorageSize})
	}

	if len(all.Map()) == 0 {
		return nil
	}
	return all
}

type imagesMetadataDoc struct {
	// EnvUUID is the environment identifier.
	EnvUUID string `bson:"env-uuid"`

	// Id contains unique natural key for cloud image metadata
	Id string `bson:"_id"`

	// Stream contains reference to a particular stream,
	// for e.g. "daily" or "released"
	Stream string `bson:"stream"`

	// Region is the name of cloud region associated with the image.
	Region string `bson:"region"`

	// Series is Os version, for e.g. "quantal".
	Series string `bson:"series"`

	// Arch is the architecture for this cloud image, for e.g. "amd64"
	Arch string `bson:"arch"`

	// ImageId is an image identifier.
	ImageId string `bson:"image_id"`

	// VirtualType contains the type of the cloud image, for e.g. "pv", "hvm". "kvm".
	VirtualType string `bson:"virtual_type,omitempty"`

	// RootStorageType contains type of root storage, for e.g. "ebs", "instance".
	RootStorageType string `bson:"root_storage_type,omitempty"`

	// RootStorageSize contains size of root storage, for e.g. "30GB", "8GB".
	RootStorageSize string `bson:"root_storage_size,omitempty"`
}

func (m imagesMetadataDoc) metadata() Metadata {
	return Metadata{
		MetadataAttributes{
			Stream:          m.Stream,
			Region:          m.Region,
			Series:          m.Series,
			Arch:            m.Arch,
			RootStorageType: m.RootStorageType,
			RootStorageSize: m.RootStorageSize,
			VirtualType:     m.VirtualType},
		m.ImageId,
	}
}

func (m imagesMetadataDoc) updates() bson.D {
	return bson.D{
		{"stream", m.Stream},
		{"region", m.Region},
		{"series", m.Series},
		{"arch", m.Arch},
		{"virtual_type", m.VirtualType},
		{"root_storage_type", m.RootStorageType},
		{"root_storage_size", m.RootStorageSize},
		{"image_id", m.ImageId},
	}
}

func (s *storage) mongoDoc(m Metadata) imagesMetadataDoc {
	return imagesMetadataDoc{
		EnvUUID:         s.envuuid,
		Id:              buildKey(m),
		Stream:          m.Stream,
		Region:          m.Region,
		Series:          m.Series,
		Arch:            m.Arch,
		VirtualType:     m.VirtualType,
		RootStorageType: m.RootStorageType,
		RootStorageSize: m.RootStorageSize,
		ImageId:         m.ImageId,
	}
}

func buildKey(m Metadata) string {
	return fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s",
		m.Stream,
		m.Region,
		m.Series,
		m.Arch,
		m.VirtualType,
		m.RootStorageType,
		m.RootStorageSize)
}
<commit_msg>Behave nicer when metadata cannot be found.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cloudimagemetadata

import (
	"fmt"

	"github.com/juju/errors"
	"github.com/juju/loggo"
	jujutxn "github.com/juju/txn"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/mongo"
)

var logger = loggo.GetLogger("juju.state.cloudimagemetadata")

type storage struct {
	envuuid        string
	collection     string
	runTransaction func(jujutxn.TransactionSource) error
	getCollection  func(string) (_ mongo.Collection, closer func())
}

var _ Storage = (*storage)(nil)

// NewStorage constructs a new Storage that stores image metadata
// in the provided collection using the provided transaction runner.
func NewStorage(
	envuuid string,
	collectionName string,
	runTransaction func(jujutxn.TransactionSource) error,
	getCollection func(string) (_ mongo.Collection, closer func()),
) Storage {
	return &storage{envuuid, collectionName, runTransaction, getCollection}
}

var emptyMetadata = Metadata{}

// SaveMetadata implements Storage.SaveMetadata and behaves as save-or-update:
// if desired metadata does not exist - we insert it; if it exists - we update.
// It throws an error if metadata did not change.
func (s *storage) SaveMetadata(metadata Metadata) error {
	newDoc := s.mongoDoc(metadata)
	all, closeAll := s.getCollection(s.collection)
	defer closeAll()

	buildTxn := func(attempt int) ([]txn.Op, error) {
		existing, err := s.getMetadata(newDoc.Id)
		if err != nil {
			if err != mgo.ErrNotFound {
				return nil, errors.Trace(err)
			}
		}

		op := txn.Op{
			C:  all.Name(),
			Id: newDoc.Id,
		}
		if existing == emptyMetadata {
			op.Assert = txn.DocMissing
			op.Insert = &newDoc
			logger.Debugf("inserting cloud image metadata for %v", newDoc.Id)
		} else {
			same := isSameMetadata(existing, metadata)
			if same {
				logger.Debugf("cloud image metadata for %v has not changed", newDoc.Id)
				return []txn.Op{}, errors.Annotate(jujutxn.ErrNoOperations, "no changes were made")
			}
			op.Assert = txn.DocExists
			op.Update = bson.D{{"$set", newDoc.updates()}}
			logger.Debugf("updating cloud image metadata for %v", newDoc.Id)
		}

		return []txn.Op{op}, nil
	}

	err := s.runTransaction(buildTxn)
	if err != nil {
		return errors.Annotatef(err, "cannot add cloud image metadata for %v", newDoc.Id)
	}
	return nil
}

func (s *storage) getMetadata(id string) (Metadata, error) {
	coll, closer := s.getCollection(s.collection)
	defer closer()

	var old imagesMetadataDoc
	err := coll.Find(bson.D{{"_id", id}}).One(&old)
	if err != nil {
		return emptyMetadata, err
	}
	return old.metadata(), nil
}

func isSameMetadata(old, new Metadata) bool {
	return old.ImageId == new.ImageId &&
		areSameAttributes(old.MetadataAttributes, new.MetadataAttributes)
}

func areSameAttributes(old, new MetadataAttributes) bool {
	return old.Arch == new.Arch &&
		old.Region == new.Region &&
		old.RootStorageType == new.RootStorageType &&
		old.RootStorageSize == new.RootStorageSize &&
		old.Series == new.Series &&
		old.Stream == new.Stream &&
		old.VirtualType == new.VirtualType
}

// AllMetadata implements Storage.AllMetadata.
func (s *storage) AllMetadata() ([]Metadata, error) {
	return s.FindMetadata(MetadataAttributes{})
}

// FindMetadata implements Storage.FindMetadata.
func (s *storage) FindMetadata(criteria MetadataAttributes) ([]Metadata, error) {
	coll, closer := s.getCollection(s.collection)
	defer closer()

	searchCriteria := buildSearchClauses(criteria)
	var docs []imagesMetadataDoc
	if err := coll.Find(searchCriteria).All(&docs); err != nil {
		return nil, errors.Trace(err)
	}
	if searchCriteria != nil && len(docs) == 0 {
		// If criteria had values and no metadata was found, err
		return nil, errors.NotFoundf("matching cloud image metadata")
	}
	metadata := make([]Metadata, len(docs))
	for i, doc := range docs {
		metadata[i] = doc.metadata()
	}
	return metadata, nil
}

func buildSearchClauses(criteria MetadataAttributes) bson.D {
	all := bson.D{}

	if criteria.Stream != "" {
		all = append(all, bson.DocElem{"stream", criteria.Stream})
	}

	if criteria.Region != "" {
		all = append(all, bson.DocElem{"region", criteria.Region})
	}

	if criteria.Series != "" {
		all = append(all, bson.DocElem{"series", criteria.Series})
	}

	if criteria.Arch != "" {
		all = append(all, bson.DocElem{"arch", criteria.Arch})
	}

	if criteria.VirtualType != "" {
		all = append(all, bson.DocElem{"virtual_type", criteria.VirtualType})
	}

	if criteria.RootStorageType != "" {
		all = append(all, bson.DocElem{"root_storage_type", criteria.RootStorageType})
	}

	if criteria.RootStorageSize != "" {
		all = append(all, bson.DocElem{"root_storage_size", criteria.RootStorageSize})
	}

	if len(all.Map()) == 0 {
		return nil
	}
	return all
}

type imagesMetadataDoc struct {
	// EnvUUID is the environment identifier.
	EnvUUID string `bson:"env-uuid"`

	// Id contains unique natural key for cloud image metadata
	Id string `bson:"_id"`

	// Stream contains reference to a particular stream,
	// for e.g. "daily" or "released"
	Stream string `bson:"stream"`

	// Region is the name of cloud region associated with the image.
	Region string `bson:"region"`

	// Series is Os version, for e.g. "quantal".
	Series string `bson:"series"`

	// Arch is the architecture for this cloud image, for e.g. "amd64"
	Arch string `bson:"arch"`

	// ImageId is an image identifier.
	ImageId string `bson:"image_id"`

	// VirtualType contains the type of the cloud image, for e.g. "pv", "hvm". "kvm".
	VirtualType string `bson:"virtual_type,omitempty"`

	// RootStorageType contains type of root storage, for e.g. "ebs", "instance".
	RootStorageType string `bson:"root_storage_type,omitempty"`

	// RootStorageSize contains size of root storage, for e.g. "30GB", "8GB".
	RootStorageSize string `bson:"root_storage_size,omitempty"`
}

func (m imagesMetadataDoc) metadata() Metadata {
	return Metadata{
		MetadataAttributes{
			Stream:          m.Stream,
			Region:          m.Region,
			Series:          m.Series,
			Arch:            m.Arch,
			RootStorageType: m.RootStorageType,
			RootStorageSize: m.RootStorageSize,
			VirtualType:     m.VirtualType},
		m.ImageId,
	}
}

func (m imagesMetadataDoc) updates() bson.D {
	return bson.D{
		{"stream", m.Stream},
		{"region", m.Region},
		{"series", m.Series},
		{"arch", m.Arch},
		{"virtual_type", m.VirtualType},
		{"root_storage_type", m.RootStorageType},
		{"root_storage_size", m.RootStorageSize},
		{"image_id", m.ImageId},
	}
}

func (s *storage) mongoDoc(m Metadata) imagesMetadataDoc {
	return imagesMetadataDoc{
		EnvUUID:         s.envuuid,
		Id:              buildKey(m),
		Stream:          m.Stream,
		Region:          m.Region,
		Series:          m.Series,
		Arch:            m.Arch,
		VirtualType:     m.VirtualType,
		RootStorageType: m.RootStorageType,
		RootStorageSize: m.RootStorageSize,
		ImageId:         m.ImageId,
	}
}

func buildKey(m Metadata) string {
	return fmt.Sprintf("%s:%s:%s:%s:%s:%s:%s",
		m.Stream,
		m.Region,
		m.Series,
		m.Arch,
		m.VirtualType,
		m.RootStorageType,
		m.RootStorageSize)
}
<|endoftext|>


<commit_before>package main

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

const (
	LUNCHY_VERSION = "0.1.4"
)

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func fileCopy(src string, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}

	defer s.Close()

	d, err := os.Create(dst)
	if err != nil {
		return err
	}

	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}

	return d.Close()
}

func findPlists(path string) []string {
	result := []string{}
	files, err := ioutil.ReadDir(path)

	if err != nil {
		return result
	}

	for _, file := range files {
		if !file.IsDir() {
			if (filepath.Ext(file.Name())) == ".plist" {
				name := strings.Replace(file.Name(), ".plist", "", -1)
				result = append(result, name)
			}
		}
	}

	return result
}

func getPlists() []string {
	path := fmt.Sprintf("%s/Library/LaunchAgents", os.Getenv("HOME"))
	files := findPlists(path)

	return files
}

func getPlist(name string) string {
	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			return plist
		}
	}

	return ""
}

func sliceIncludes(slice []string, match string) bool {
	for _, val := range slice {
		if val == match {
			return true
		}
	}

	return false
}

func printUsage() {
	fmt.Printf("Lunchy %s, the friendly launchctl wrapper\n", LUNCHY_VERSION)
	fmt.Println("Usage: lunchy [start|stop|restart|list|status|install|show|edit|remove|scan] [options]")
}

func printList() {
	for _, file := range getPlists() {
		fmt.Println(file)
	}
}

func printStatus(args []string) {
	out, err := exec.Command("launchctl", "list").Output()

	if err != nil {
		fatal("failed to get process list")
	}

	pattern := ""

	if len(args) == 3 {
		pattern = args[2]
	}

	installed := getPlists()
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")

	for _, line := range lines {
		chunks := strings.Split(line, "\t")
		clean_line := strings.Replace(line, "\t", " ", -1)

		if len(pattern) > 0 {
			if strings.Index(chunks[2], pattern) != -1 {
				if sliceIncludes(installed, chunks[2]) {
					fmt.Println(clean_line)
				}
			}
		} else {
			if sliceIncludes(installed, chunks[2]) {
				fmt.Println(clean_line)
			}
		}
	}
}

func exitWithInvalidArgs(args []string, msg string) {
	if len(args) < 3 {
		fmt.Println(msg)
		os.Exit(1)
	}
}

func startDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			startDaemon(plist)
		}
	}
}

func startDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "load", path).Output()

	if err != nil {
		fmt.Println("failed to start", name)
		return
	}

	fmt.Println("started", name)
}

func stopDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
		}
	}
}

func stopDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "unload", path).Output()

	if err != nil {
		fmt.Println("failed to stop", name)
		return
	}

	fmt.Println("stopped", name)
}

func restartDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
			startDaemon(plist)
		}
	}
}

func showPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			printPlistContent(plist)
			return
		}
	}
}

func printPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	contents, err := ioutil.ReadFile(path)

	if err != nil {
		fatal("unable to read plist")
	}

	fmt.Printf(string(contents))
}

func editPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			editPlistContent(plist)
			return
		}
	}
}

func editPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	editor := os.Getenv("EDITOR")

	if len(editor) == 0 {
		fatal("EDITOR environment variable is not set")
	}

	cmd := exec.Command(editor, path)

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout

	cmd.Start()
	cmd.Wait()
}

func installPlist(args []string) {
	exitWithInvalidArgs(args, "path required")

	path := args[2]

	if !fileExists(path) {
		fatal("source file does not exist")
	}

	info, _ := os.Stat(path)
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")
	new_path := fmt.Sprintf("%s/%s", base_path, info.Name())

	if fileExists(new_path) && os.Remove(new_path) != nil {
		fatal("unable to delete existing plist")
	}

	if fileCopy(path, new_path) != nil {
		fatal("failed to copy file")
	}

	fmt.Println(path, "installed to", base_path)
}

func removePlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			path := fmt.Sprintf("%s/%s.plist", base_path, plist)

			if os.Remove(path) == nil {
				fmt.Println("removed", path)
			} else {
				fmt.Println("failed to remove", path)
			}
		}
	}
}

func scanPath(args []string) {
	path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")

	if len(args) >= 3 {
		path = args[2]
	}

	out, err := exec.Command("find", path, "-name", "*.plist").Output()
	if err != nil {
		fatal("failed to scan files")
	}

	fmt.Println(strings.TrimSpace(string(out)))
}

func fatal(message string) {
	fmt.Println(message)
	os.Exit(1)
}

func main() {
	args := os.Args

	if len(args) == 1 {
		printUsage()
		os.Exit(1)
	}

	switch args[1] {
	default:
		printUsage()
		os.Exit(1)
	case "help":
		printUsage()
		return
	case "list", "ls":
		printList()
		return
	case "status", "ps":
		printStatus(args)
		return
	case "start":
		startDaemons(args)
		return
	case "stop":
		stopDaemons(args)
		return
	case "restart":
		restartDaemons(args)
		return
	case "show":
		showPlist(args)
		return
	case "edit":
		editPlist(args)
		return
	case "install", "add":
		installPlist(args)
		return
	case "remove", "rm":
		removePlist(args)
		return
	case "scan":
		scanPath(args)
		return
	}
}
<commit_msg>Version bump: 0.1.5<commit_after>package main

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

const (
	LUNCHY_VERSION = "0.1.5"
)

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

func fileCopy(src string, dst string) error {
	s, err := os.Open(src)
	if err != nil {
		return err
	}

	defer s.Close()

	d, err := os.Create(dst)
	if err != nil {
		return err
	}

	if _, err := io.Copy(d, s); err != nil {
		d.Close()
		return err
	}

	return d.Close()
}

func findPlists(path string) []string {
	result := []string{}
	files, err := ioutil.ReadDir(path)

	if err != nil {
		return result
	}

	for _, file := range files {
		if !file.IsDir() {
			if (filepath.Ext(file.Name())) == ".plist" {
				name := strings.Replace(file.Name(), ".plist", "", -1)
				result = append(result, name)
			}
		}
	}

	return result
}

func getPlists() []string {
	path := fmt.Sprintf("%s/Library/LaunchAgents", os.Getenv("HOME"))
	files := findPlists(path)

	return files
}

func getPlist(name string) string {
	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			return plist
		}
	}

	return ""
}

func sliceIncludes(slice []string, match string) bool {
	for _, val := range slice {
		if val == match {
			return true
		}
	}

	return false
}

func printUsage() {
	fmt.Printf("Lunchy %s, the friendly launchctl wrapper\n", LUNCHY_VERSION)
	fmt.Println("Usage: lunchy [start|stop|restart|list|status|install|show|edit|remove|scan] [options]")
}

func printList() {
	for _, file := range getPlists() {
		fmt.Println(file)
	}
}

func printStatus(args []string) {
	out, err := exec.Command("launchctl", "list").Output()

	if err != nil {
		fatal("failed to get process list")
	}

	pattern := ""

	if len(args) == 3 {
		pattern = args[2]
	}

	installed := getPlists()
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")

	for _, line := range lines {
		chunks := strings.Split(line, "\t")
		clean_line := strings.Replace(line, "\t", " ", -1)

		if len(pattern) > 0 {
			if strings.Index(chunks[2], pattern) != -1 {
				if sliceIncludes(installed, chunks[2]) {
					fmt.Println(clean_line)
				}
			}
		} else {
			if sliceIncludes(installed, chunks[2]) {
				fmt.Println(clean_line)
			}
		}
	}
}

func exitWithInvalidArgs(args []string, msg string) {
	if len(args) < 3 {
		fmt.Println(msg)
		os.Exit(1)
	}
}

func startDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			startDaemon(plist)
		}
	}
}

func startDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "load", path).Output()

	if err != nil {
		fmt.Println("failed to start", name)
		return
	}

	fmt.Println("started", name)
}

func stopDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
		}
	}
}

func stopDaemon(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	_, err := exec.Command("launchctl", "unload", path).Output()

	if err != nil {
		fmt.Println("failed to stop", name)
		return
	}

	fmt.Println("stopped", name)
}

func restartDaemons(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			stopDaemon(plist)
			startDaemon(plist)
		}
	}
}

func showPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			printPlistContent(plist)
			return
		}
	}
}

func printPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	contents, err := ioutil.ReadFile(path)

	if err != nil {
		fatal("unable to read plist")
	}

	fmt.Printf(string(contents))
}

func editPlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			editPlistContent(plist)
			return
		}
	}
}

func editPlistContent(name string) {
	path := fmt.Sprintf("%s/Library/LaunchAgents/%s.plist", os.Getenv("HOME"), name)
	editor := os.Getenv("EDITOR")

	if len(editor) == 0 {
		fatal("EDITOR environment variable is not set")
	}

	cmd := exec.Command(editor, path)

	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout

	cmd.Start()
	cmd.Wait()
}

func installPlist(args []string) {
	exitWithInvalidArgs(args, "path required")

	path := args[2]

	if !fileExists(path) {
		fatal("source file does not exist")
	}

	info, _ := os.Stat(path)
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")
	new_path := fmt.Sprintf("%s/%s", base_path, info.Name())

	if fileExists(new_path) && os.Remove(new_path) != nil {
		fatal("unable to delete existing plist")
	}

	if fileCopy(path, new_path) != nil {
		fatal("failed to copy file")
	}

	fmt.Println(path, "installed to", base_path)
}

func removePlist(args []string) {
	exitWithInvalidArgs(args, "name required")

	name := args[2]
	base_path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")

	for _, plist := range getPlists() {
		if strings.Index(plist, name) != -1 {
			path := fmt.Sprintf("%s/%s.plist", base_path, plist)

			if os.Remove(path) == nil {
				fmt.Println("removed", path)
			} else {
				fmt.Println("failed to remove", path)
			}
		}
	}
}

func scanPath(args []string) {
	path := fmt.Sprintf("%s/%s", os.Getenv("HOME"), "Library/LaunchAgents")

	if len(args) >= 3 {
		path = args[2]
	}

	out, err := exec.Command("find", path, "-name", "*.plist").Output()
	if err != nil {
		fatal("failed to scan files")
	}

	fmt.Println(strings.TrimSpace(string(out)))
}

func fatal(message string) {
	fmt.Println(message)
	os.Exit(1)
}

func main() {
	args := os.Args

	if len(args) == 1 {
		printUsage()
		os.Exit(1)
	}

	switch args[1] {
	default:
		printUsage()
		os.Exit(1)
	case "help":
		printUsage()
		return
	case "list", "ls":
		printList()
		return
	case "status", "ps":
		printStatus(args)
		return
	case "start":
		startDaemons(args)
		return
	case "stop":
		stopDaemons(args)
		return
	case "restart":
		restartDaemons(args)
		return
	case "show":
		showPlist(args)
		return
	case "edit":
		editPlist(args)
		return
	case "install", "add":
		installPlist(args)
		return
	case "remove", "rm":
		removePlist(args)
		return
	case "scan":
		scanPath(args)
		return
	}
}
<|endoftext|>


<commit_before>package keychain

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
	"time"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/wallet"
	"github.com/btcsuite/btcwallet/walletdb"

	_ "github.com/btcsuite/btcwallet/walletdb/bdb" // Required in order to create the default database.
)

// versionZeroKeyFamilies is a slice of all the known key families for first
// version of the key derivation schema defined in this package.
var versionZeroKeyFamilies = []KeyFamily{
	KeyFamilyMultiSig,
	KeyFamilyRevocationBase,
	KeyFamilyHtlcBase,
	KeyFamilyPaymentBase,
	KeyFamilyDelayBase,
	KeyFamilyRevocationRoot,
	KeyFamilyNodeKey,
}

var (
	testHDSeed = chainhash.Hash{
		0xb7, 0x94, 0x38, 0x5f, 0x2d, 0x1e, 0xf7, 0xab,
		0x4d, 0x92, 0x73, 0xd1, 0x90, 0x63, 0x81, 0xb4,
		0x4f, 0x2f, 0x6f, 0x25, 0x98, 0xa3, 0xef, 0xb9,
		0x69, 0x49, 0x18, 0x83, 0x31, 0x98, 0x47, 0x53,
	}
)

func createTestBtcWallet(coinType uint32) (func(), *wallet.Wallet, error) {
	tempDir, err := ioutil.TempDir("", "keyring-lnwallet")
	if err != nil {
		return nil, nil, err
	}
	loader := wallet.NewLoader(&chaincfg.SimNetParams, tempDir, 0)

	pass := []byte("test")

	baseWallet, err := loader.CreateNewWallet(
		pass, pass, testHDSeed[:], time.Time{},
	)
	if err != nil {
		return nil, nil, err
	}

	if err := baseWallet.Unlock(pass, nil); err != nil {
		return nil, nil, err
	}

	// Construct the key scope required to derive keys for the chose
	// coinType.
	chainKeyScope := waddrmgr.KeyScope{
		Purpose: BIP0043Purpose,
		Coin:    coinType,
	}

	// We'll now ensure that the KeyScope: (1017, coinType) exists within
	// the internal waddrmgr. We'll need this in order to properly generate
	// the keys required for signing various contracts.
	_, err = baseWallet.Manager.FetchScopedKeyManager(chainKeyScope)
	if err != nil {
		err := walletdb.Update(baseWallet.Database(), func(tx walletdb.ReadWriteTx) error {
			addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)

			_, err := baseWallet.Manager.NewScopedKeyManager(
				addrmgrNs, chainKeyScope, lightningAddrSchema,
			)
			return err
		})
		if err != nil {
			return nil, nil, err
		}
	}

	cleanUp := func() {
		baseWallet.Lock()
		os.RemoveAll(tempDir)
	}

	return cleanUp, baseWallet, nil
}

// secretKeyRingConstructor is a function signature that's used as a generic
// constructor for various implementations of the KeyRing interface. A string
// naming the returned interface, a function closure that cleans up any
// resources, and the clean up interface itself are to be returned.
type keyRingConstructor func() (string, func(), KeyRing, error)

// TestKeyRingDerivation tests that each known KeyRing implementation properly
// adheres to the expected behavior of the set of interfaces.
func TestKeyRingDerivation(t *testing.T) {
	t.Parallel()

	keyRingImplementations := []keyRingConstructor{
		func() (string, func(), KeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeBitcoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeBitcoin)

			return "btcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), KeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeLitecoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeLitecoin)

			return "ltcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), KeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeTestnet,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeTestnet)

			return "testwallet", cleanUp, keyRing, nil
		},
	}

	// For each implementation constructor registered above, we'll execute
	// an identical set of tests in order to ensure that the interface
	// adheres to our nominal specification.
	for _, keyRingConstructor := range keyRingImplementations {
		keyRingName, cleanUp, keyRing, err := keyRingConstructor()
		if err != nil {
			t.Fatalf("unable to create key ring %v: %v", keyRingName,
				err)
		}
		defer cleanUp()

		success := t.Run(fmt.Sprintf("%v", keyRingName), func(t *testing.T) {
			// First, we'll ensure that we're able to derive keys
			// from each of the known key families.
			for _, keyFam := range versionZeroKeyFamilies {
				// First, we'll ensure that we can derive the
				// *next* key in the keychain.
				keyDesc, err := keyRing.DeriveNextKey(keyFam)
				if err != nil {
					t.Fatalf("unable to derive next for "+
						"keyFam=%v: %v", keyFam, err)
				}

				// If we now try to manually derive the *first*
				// key, then we should get an identical public
				// key back.
				keyLoc := KeyLocator{
					Family: keyFam,
					Index:  0,
				}
				firstKeyDesc, err := keyRing.DeriveKey(keyLoc)
				if err != nil {
					t.Fatalf("unable to derive first key for "+
						"keyFam=%v: %v", keyFam, err)
				}

				if !keyDesc.PubKey.IsEqual(firstKeyDesc.PubKey) {
					t.Fatalf("mismatched keys: expected %v, "+
						"got %x",
						keyDesc.PubKey.SerializeCompressed(),
						firstKeyDesc.PubKey.SerializeCompressed())
				}

				// If this succeeds, then we'll also try to
				// derive a random index within the range.
				randKeyIndex := uint32(rand.Int31())
				keyLoc = KeyLocator{
					Family: keyFam,
					Index:  randKeyIndex,
				}
				_, err = keyRing.DeriveKey(keyLoc)
				if err != nil {
					t.Fatalf("unable to derive key_index=%v "+
						"for keyFam=%v: %v",
						randKeyIndex, keyFam, err)
				}
			}
		})
		if !success {
			break
		}
	}
}

// secretKeyRingConstructor is a function signature that's used as a generic
// constructor for various implementations of the SecretKeyRing interface. A
// string naming the returned interface, a function closure that cleans up any
// resources, and the clean up interface itself are to be returned.
type secretKeyRingConstructor func() (string, func(), SecretKeyRing, error)

// TestSecretKeyRingDerivation tests that each known SecretKeyRing
// implementation properly adheres to the expected behavior of the set of
// interface.
func TestSecretKeyRingDerivation(t *testing.T) {
	t.Parallel()

	secretKeyRingImplementations := []secretKeyRingConstructor{
		func() (string, func(), SecretKeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeBitcoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeBitcoin)

			return "btcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), SecretKeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeLitecoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeLitecoin)

			return "ltcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), SecretKeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeTestnet,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeTestnet)

			return "testwallet", cleanUp, keyRing, nil
		},
	}

	// For each implementation constructor registered above, we'll execute
	// an identical set of tests in order to ensure that the interface
	// adheres to our nominal specification.
	for _, secretKeyRingConstructor := range secretKeyRingImplementations {
		keyRingName, cleanUp, secretKeyRing, err := secretKeyRingConstructor()
		if err != nil {
			t.Fatalf("unable to create secret key ring %v: %v",
				keyRingName, err)
		}
		defer cleanUp()

		success := t.Run(fmt.Sprintf("%v", keyRingName), func(t *testing.T) {
			// First, each key family, we'll ensure that we're able
			// to obtain the private key of a randomly select child
			// index within the key family.
			for _, keyFam := range versionZeroKeyFamilies {
				randKeyIndex := uint32(rand.Int31())
				keyLoc := KeyLocator{
					Family: keyFam,
					Index:  randKeyIndex,
				}

				// First, we'll query for the public key for
				// this target key locator.
				pubKeyDesc, err := secretKeyRing.DeriveKey(keyLoc)
				if err != nil {
					t.Fatalf("unable to derive pubkey "+
						"(fam=%v, index=%v): %v",
						keyLoc.Family,
						keyLoc.Index, err)
				}

				// With the public key derive, ensure that
				// we're able to obtain the corresponding
				// private key correctly.
				privKey, err := secretKeyRing.DerivePrivKey(KeyDescriptor{
					KeyLocator: keyLoc,
				})
				if err != nil {
					t.Fatalf("unable to derive priv "+
						"(fam=%v, index=%v): %v", keyLoc.Family,
						keyLoc.Index, err)
				}

				// Finally, ensure that the keys match up
				// properly.
				if !pubKeyDesc.PubKey.IsEqual(privKey.PubKey()) {
					t.Fatalf("pubkeys mismatched: expected %x, got %x",
						pubKeyDesc.PubKey.SerializeCompressed(),
						privKey.PubKey().SerializeCompressed())
				}

				// TODO(roasbeef): scalar mult once integrated
			}
		})
		if !success {
			break
		}
	}
}
<commit_msg>keychain: extend TestKeyRingDerivation to check KeyLocators of derived keys<commit_after>package keychain

import (
	"fmt"
	"io/ioutil"
	"math/rand"
	"os"
	"runtime"
	"testing"
	"time"

	"github.com/btcsuite/btcd/chaincfg"
	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcwallet/waddrmgr"
	"github.com/btcsuite/btcwallet/wallet"
	"github.com/btcsuite/btcwallet/walletdb"
	"github.com/davecgh/go-spew/spew"

	_ "github.com/btcsuite/btcwallet/walletdb/bdb" // Required in order to create the default database.
)

// versionZeroKeyFamilies is a slice of all the known key families for first
// version of the key derivation schema defined in this package.
var versionZeroKeyFamilies = []KeyFamily{
	KeyFamilyMultiSig,
	KeyFamilyRevocationBase,
	KeyFamilyHtlcBase,
	KeyFamilyPaymentBase,
	KeyFamilyDelayBase,
	KeyFamilyRevocationRoot,
	KeyFamilyNodeKey,
}

var (
	testHDSeed = chainhash.Hash{
		0xb7, 0x94, 0x38, 0x5f, 0x2d, 0x1e, 0xf7, 0xab,
		0x4d, 0x92, 0x73, 0xd1, 0x90, 0x63, 0x81, 0xb4,
		0x4f, 0x2f, 0x6f, 0x25, 0x98, 0xa3, 0xef, 0xb9,
		0x69, 0x49, 0x18, 0x83, 0x31, 0x98, 0x47, 0x53,
	}
)

func createTestBtcWallet(coinType uint32) (func(), *wallet.Wallet, error) {
	tempDir, err := ioutil.TempDir("", "keyring-lnwallet")
	if err != nil {
		return nil, nil, err
	}
	loader := wallet.NewLoader(&chaincfg.SimNetParams, tempDir, 0)

	pass := []byte("test")

	baseWallet, err := loader.CreateNewWallet(
		pass, pass, testHDSeed[:], time.Time{},
	)
	if err != nil {
		return nil, nil, err
	}

	if err := baseWallet.Unlock(pass, nil); err != nil {
		return nil, nil, err
	}

	// Construct the key scope required to derive keys for the chose
	// coinType.
	chainKeyScope := waddrmgr.KeyScope{
		Purpose: BIP0043Purpose,
		Coin:    coinType,
	}

	// We'll now ensure that the KeyScope: (1017, coinType) exists within
	// the internal waddrmgr. We'll need this in order to properly generate
	// the keys required for signing various contracts.
	_, err = baseWallet.Manager.FetchScopedKeyManager(chainKeyScope)
	if err != nil {
		err := walletdb.Update(baseWallet.Database(), func(tx walletdb.ReadWriteTx) error {
			addrmgrNs := tx.ReadWriteBucket(waddrmgrNamespaceKey)

			_, err := baseWallet.Manager.NewScopedKeyManager(
				addrmgrNs, chainKeyScope, lightningAddrSchema,
			)
			return err
		})
		if err != nil {
			return nil, nil, err
		}
	}

	cleanUp := func() {
		baseWallet.Lock()
		os.RemoveAll(tempDir)
	}

	return cleanUp, baseWallet, nil
}

func assertEqualKeyLocator(t *testing.T, a, b KeyLocator) {
	_, _, line, _ := runtime.Caller(1)
	if a != b {
		t.Fatalf("line #%v: mismatched key locators: expected %v, "+
			"got %v", line, spew.Sdump(a), spew.Sdump(b))
	}
}

// secretKeyRingConstructor is a function signature that's used as a generic
// constructor for various implementations of the KeyRing interface. A string
// naming the returned interface, a function closure that cleans up any
// resources, and the clean up interface itself are to be returned.
type keyRingConstructor func() (string, func(), KeyRing, error)

// TestKeyRingDerivation tests that each known KeyRing implementation properly
// adheres to the expected behavior of the set of interfaces.
func TestKeyRingDerivation(t *testing.T) {
	t.Parallel()

	keyRingImplementations := []keyRingConstructor{
		func() (string, func(), KeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeBitcoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeBitcoin)

			return "btcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), KeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeLitecoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeLitecoin)

			return "ltcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), KeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeTestnet,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeTestnet)

			return "testwallet", cleanUp, keyRing, nil
		},
	}

	const numKeysToDerive = 10

	// For each implementation constructor registered above, we'll execute
	// an identical set of tests in order to ensure that the interface
	// adheres to our nominal specification.
	for _, keyRingConstructor := range keyRingImplementations {
		keyRingName, cleanUp, keyRing, err := keyRingConstructor()
		if err != nil {
			t.Fatalf("unable to create key ring %v: %v", keyRingName,
				err)
		}
		defer cleanUp()

		success := t.Run(fmt.Sprintf("%v", keyRingName), func(t *testing.T) {
			// First, we'll ensure that we're able to derive keys
			// from each of the known key families.
			for _, keyFam := range versionZeroKeyFamilies {
				// First, we'll ensure that we can derive the
				// *next* key in the keychain.
				keyDesc, err := keyRing.DeriveNextKey(keyFam)
				if err != nil {
					t.Fatalf("unable to derive next for "+
						"keyFam=%v: %v", keyFam, err)
				}
				assertEqualKeyLocator(t,
					KeyLocator{
						Family: keyFam,
						Index:  0,
					}, keyDesc.KeyLocator,
				)

				// We'll now re-derive that key to ensure that
				// we're able to properly access the key via
				// the random access derivation methods.
				keyLoc := KeyLocator{
					Family: keyFam,
					Index:  0,
				}
				firstKeyDesc, err := keyRing.DeriveKey(keyLoc)
				if err != nil {
					t.Fatalf("unable to derive first key for "+
						"keyFam=%v: %v", keyFam, err)
				}
				if !keyDesc.PubKey.IsEqual(firstKeyDesc.PubKey) {
					t.Fatalf("mismatched keys: expected %x, "+
						"got %x",
						keyDesc.PubKey.SerializeCompressed(),
						firstKeyDesc.PubKey.SerializeCompressed())
				}
				assertEqualKeyLocator(t,
					KeyLocator{
						Family: keyFam,
						Index:  0,
					}, firstKeyDesc.KeyLocator,
				)

				// If we now try to manually derive the next 10
				// keys (including the original key), then we
				// should get an identical public key back and
				// their KeyLocator information
				// should be set properly.
				for i := 0; i < numKeysToDerive+1; i++ {
					keyLoc := KeyLocator{
						Family: keyFam,
						Index:  uint32(i),
					}
					keyDesc, err := keyRing.DeriveKey(keyLoc)
					if err != nil {
						t.Fatalf("unable to derive first key for "+
							"keyFam=%v: %v", keyFam, err)
					}

					// Ensure that the key locator matches
					// up as well.
					assertEqualKeyLocator(
						t, keyLoc, keyDesc.KeyLocator,
					)
				}

				// If this succeeds, then we'll also try to
				// derive a random index within the range.
				randKeyIndex := uint32(rand.Int31())
				keyLoc = KeyLocator{
					Family: keyFam,
					Index:  randKeyIndex,
				}
				keyDesc, err = keyRing.DeriveKey(keyLoc)
				if err != nil {
					t.Fatalf("unable to derive key_index=%v "+
						"for keyFam=%v: %v",
						randKeyIndex, keyFam, err)
				}
				assertEqualKeyLocator(
					t, keyLoc, keyDesc.KeyLocator,
				)
			}
		})
		if !success {
			break
		}
	}
}

// secretKeyRingConstructor is a function signature that's used as a generic
// constructor for various implementations of the SecretKeyRing interface. A
// string naming the returned interface, a function closure that cleans up any
// resources, and the clean up interface itself are to be returned.
type secretKeyRingConstructor func() (string, func(), SecretKeyRing, error)

// TestSecretKeyRingDerivation tests that each known SecretKeyRing
// implementation properly adheres to the expected behavior of the set of
// interface.
func TestSecretKeyRingDerivation(t *testing.T) {
	t.Parallel()

	secretKeyRingImplementations := []secretKeyRingConstructor{
		func() (string, func(), SecretKeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeBitcoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeBitcoin)

			return "btcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), SecretKeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeLitecoin,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeLitecoin)

			return "ltcwallet", cleanUp, keyRing, nil
		},
		func() (string, func(), SecretKeyRing, error) {
			cleanUp, wallet, err := createTestBtcWallet(
				CoinTypeTestnet,
			)
			if err != nil {
				t.Fatalf("unable to create wallet: %v", err)
			}

			keyRing := NewBtcWalletKeyRing(wallet, CoinTypeTestnet)

			return "testwallet", cleanUp, keyRing, nil
		},
	}

	// For each implementation constructor registered above, we'll execute
	// an identical set of tests in order to ensure that the interface
	// adheres to our nominal specification.
	for _, secretKeyRingConstructor := range secretKeyRingImplementations {
		keyRingName, cleanUp, secretKeyRing, err := secretKeyRingConstructor()
		if err != nil {
			t.Fatalf("unable to create secret key ring %v: %v",
				keyRingName, err)
		}
		defer cleanUp()

		success := t.Run(fmt.Sprintf("%v", keyRingName), func(t *testing.T) {
			// First, each key family, we'll ensure that we're able
			// to obtain the private key of a randomly select child
			// index within the key family.
			for _, keyFam := range versionZeroKeyFamilies {
				randKeyIndex := uint32(rand.Int31())
				keyLoc := KeyLocator{
					Family: keyFam,
					Index:  randKeyIndex,
				}

				// First, we'll query for the public key for
				// this target key locator.
				pubKeyDesc, err := secretKeyRing.DeriveKey(keyLoc)
				if err != nil {
					t.Fatalf("unable to derive pubkey "+
						"(fam=%v, index=%v): %v",
						keyLoc.Family,
						keyLoc.Index, err)
				}

				// With the public key derive, ensure that
				// we're able to obtain the corresponding
				// private key correctly.
				privKey, err := secretKeyRing.DerivePrivKey(KeyDescriptor{
					KeyLocator: keyLoc,
				})
				if err != nil {
					t.Fatalf("unable to derive priv "+
						"(fam=%v, index=%v): %v", keyLoc.Family,
						keyLoc.Index, err)
				}

				// Finally, ensure that the keys match up
				// properly.
				if !pubKeyDesc.PubKey.IsEqual(privKey.PubKey()) {
					t.Fatalf("pubkeys mismatched: expected %x, got %x",
						pubKeyDesc.PubKey.SerializeCompressed(),
						privKey.PubKey().SerializeCompressed())
				}

				// TODO(roasbeef): scalar mult once integrated
			}
		})
		if !success {
			break
		}
	}
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

// CmdDatacenter subcommand
import (
	"fmt"
	"os"

	"github.com/fatih/color"
	"github.com/urfave/cli"
)

// ListServices ...
var ListServices = cli.Command{
	Name:      "list",
	Usage:     "List available services.",
	ArgsUsage: " ",
	Description: `List available services and shows its most relevant information.

   Example:
    $ ernest service list
	`,
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}
		services, err := m.ListServices(cfg.Token)
		if err != nil {
			color.Red(err.Error())
			return err
		}

		printServiceList(services)
		return nil
	},
}

// ApplyService command
// Applies changes described on a YAML file to a service
var ApplyService = cli.Command{
	Name:      "apply",
	Aliases:   []string{"a"},
	ArgsUsage: "<file.yml>",
	Usage:     "Builds or changes infrastructure.",
	Description: `Sends a service YAML description file to Ernest to be executed.
   You must be logged in to execute this command.

   If the file is not provided, ernest.yml will be used by default.

   Example:
    $ ernest apply myservice.yml
	`,
	Action: func(c *cli.Context) error {
		file := "ernest.yml"
		if len(c.Args()) == 1 {
			file = c.Args()[0]
		}
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}
		_, err := m.Apply(cfg.Token, file, true)
		if err != nil {
			color.Red(err.Error())
		}
		return nil
	},
}

// DestroyService command
var DestroyService = cli.Command{
	Name:      "destroy",
	Aliases:   []string{"d"},
	ArgsUsage: "<service_name>",
	Usage:     "Destroy a service.",
	Description: `Destroys a service by its name.

   Example:
    $ ernest destroy myservice
  `,
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name:  "force,f",
			Usage: "Force destroy command without asking for permission.",
		},
	},
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}

		if len(c.Args()) < 1 {
			color.Red("You should specify an existing service name")
			return nil
		}
		name := c.Args()[0]

		if c.Bool("force") {
			err := m.Destroy(cfg.Token, name, true)
			if err != nil {
				color.Red(err.Error())
			}
		} else {
			fmt.Print("Are you sure? Please type yes or no and then press enter: ")
			if askForConfirmation() {
				err := m.Destroy(cfg.Token, name, true)
				if err != nil {
					color.Red(err.Error())
					return err
				}
			}
		}
		return nil
	},
}

// HistoryService command
// Shows the history of a service, a list of builds
var HistoryService = cli.Command{
	Name:      "history",
	Usage:     "Shows the history of a service, a list of builds",
	ArgsUsage: "<service_name>",
	Description: `Shows the history of a service, a list of builds and its status and basic information.

   Example:
    $ ernest history myservice
	`,
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}

		if len(c.Args()) < 1 {
			color.Red("You should specify an existing service name")
			return nil
		}

		serviceName := c.Args()[0]

		services, _ := m.ListBuilds(serviceName, cfg.Token)
		printServiceHistory(services)
		return nil
	},
}

// ResetService command
var ResetService = cli.Command{
	Name:      "reset",
	ArgsUsage: "<service_name>",
	Usage:     "Reset an in progress service.",
	Description: `Reseting a service creation may cause problems, please make sure you know what are you doing.

   Example:
    $ ernest reset myservice
  `,
	Action: func(c *cli.Context) error {
		if len(c.Args()) < 1 {
			color.Red("You should specify the service name")
		} else {
			serviceName := c.Args()[0]
			m, cfg := setup(c)
			err := m.ResetService(serviceName, cfg.Token)
			if err != nil {
				fmt.Println(err)
			}
		}
		return nil
	},
}

// DefinitionService command
// Shows the current definition of a service by its name
var DefinitionService = cli.Command{
	Name:      "definition",
	Aliases:   []string{"s"},
	ArgsUsage: "<service_name>",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "build",
			Value: "",
			Usage: "Build ID",
		},
	},
	Usage: "Show the current definition of a service by its name",
	Description: `Show the current definition of a service by its name getting the definition about the build.

   Example:
    $ ernest service definition myservice
	`,
	Action: func(c *cli.Context) error {
		if len(c.Args()) < 1 {
			color.Red("You should specify the service name")
		} else {
			m, cfg := setup(c)
			serviceName := c.Args()[0]
			if c.String("build") != "" {
				service, err := m.ServiceBuildStatus(cfg.Token, serviceName, c.String("build"))
				if err != nil {
					color.Red(err.Error())
					os.Exit(1)
				}
				fmt.Println(service.Definition)
			} else {
				service, err := m.ServiceStatus(cfg.Token, serviceName)
				if err != nil {
					color.Red(err.Error())
					os.Exit(1)
				}

				fmt.Println(service.Definition)
			}
		}
		return nil
	},
}

// InfoService : Shows detailed information of a service by its name
var InfoService = cli.Command{
	Name:      "info",
	Aliases:   []string{"i"},
	ArgsUsage: "<service_name>",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "build",
			Value: "",
			Usage: "Build ID",
		},
	},
	Usage: "$ ernest service info <my_service> --build <specific build>",
	Description: `Will show detailed information of the last build of a specified service.
	In case you specify --build option you will be able to output the detailed information of specific build of a service.

   Examples:
    $ ernest service definition myservice
    $ ernest service definition myservice --build build1
	`,
	Action: func(c *cli.Context) error {
		var err error
		var service Service
		if len(c.Args()) < 1 {
			color.Red("You should specify the service name")
		} else {
			m, cfg := setup(c)
			name := c.Args()[0]
			if c.String("build") != "" {
				build := c.String("build")
				service, err = m.ServiceBuildStatus(cfg.Token, name, build)
			} else {
				service, err = m.ServiceStatus(cfg.Token, name)
			}
			if err != nil {
				color.Red(err.Error())
				os.Exit(1)
			}
			printServiceInfo(&service)
		}
		return nil
	},
}

// CmdService ...
var CmdService = cli.Command{
	Name:  "service",
	Usage: "Service related subcommands",
	Subcommands: []cli.Command{
		ListServices,
		ApplyService,
		DestroyService,
		HistoryService,
		ResetService,
		DefinitionService,
		InfoService,
		MonitorService,
	},
}
<commit_msg>Scenario : service info<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

// CmdDatacenter subcommand
import (
	"fmt"
	"os"

	"github.com/fatih/color"
	"github.com/urfave/cli"
)

// ListServices ...
var ListServices = cli.Command{
	Name:      "list",
	Usage:     "List available services.",
	ArgsUsage: " ",
	Description: `List available services and shows its most relevant information.

   Example:
    $ ernest service list
	`,
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}
		services, err := m.ListServices(cfg.Token)
		if err != nil {
			color.Red(err.Error())
			return err
		}

		printServiceList(services)
		return nil
	},
}

// ApplyService command
// Applies changes described on a YAML file to a service
var ApplyService = cli.Command{
	Name:      "apply",
	Aliases:   []string{"a"},
	ArgsUsage: "<file.yml>",
	Usage:     "Builds or changes infrastructure.",
	Description: `Sends a service YAML description file to Ernest to be executed.
   You must be logged in to execute this command.

   If the file is not provided, ernest.yml will be used by default.

   Example:
    $ ernest apply myservice.yml
	`,
	Action: func(c *cli.Context) error {
		file := "ernest.yml"
		if len(c.Args()) == 1 {
			file = c.Args()[0]
		}
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}
		_, err := m.Apply(cfg.Token, file, true)
		if err != nil {
			color.Red(err.Error())
		}
		return nil
	},
}

// DestroyService command
var DestroyService = cli.Command{
	Name:      "destroy",
	Aliases:   []string{"d"},
	ArgsUsage: "<service_name>",
	Usage:     "Destroy a service.",
	Description: `Destroys a service by its name.

   Example:
    $ ernest destroy myservice
  `,
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name:  "force,f",
			Usage: "Force destroy command without asking for permission.",
		},
	},
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}

		if len(c.Args()) < 1 {
			color.Red("You should specify an existing service name")
			return nil
		}
		name := c.Args()[0]

		if c.Bool("force") {
			err := m.Destroy(cfg.Token, name, true)
			if err != nil {
				color.Red(err.Error())
			}
		} else {
			fmt.Print("Are you sure? Please type yes or no and then press enter: ")
			if askForConfirmation() {
				err := m.Destroy(cfg.Token, name, true)
				if err != nil {
					color.Red(err.Error())
					return err
				}
			}
		}
		return nil
	},
}

// HistoryService command
// Shows the history of a service, a list of builds
var HistoryService = cli.Command{
	Name:      "history",
	Usage:     "Shows the history of a service, a list of builds",
	ArgsUsage: "<service_name>",
	Description: `Shows the history of a service, a list of builds and its status and basic information.

   Example:
    $ ernest history myservice
	`,
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}

		if len(c.Args()) < 1 {
			color.Red("You should specify an existing service name")
			return nil
		}

		serviceName := c.Args()[0]

		services, _ := m.ListBuilds(serviceName, cfg.Token)
		printServiceHistory(services)
		return nil
	},
}

// ResetService command
var ResetService = cli.Command{
	Name:      "reset",
	ArgsUsage: "<service_name>",
	Usage:     "Reset an in progress service.",
	Description: `Reseting a service creation may cause problems, please make sure you know what are you doing.

   Example:
    $ ernest reset myservice
  `,
	Action: func(c *cli.Context) error {
		if len(c.Args()) < 1 {
			color.Red("You should specify the service name")
		} else {
			serviceName := c.Args()[0]
			m, cfg := setup(c)
			err := m.ResetService(serviceName, cfg.Token)
			if err != nil {
				fmt.Println(err)
			}
		}
		return nil
	},
}

// DefinitionService command
// Shows the current definition of a service by its name
var DefinitionService = cli.Command{
	Name:      "definition",
	Aliases:   []string{"s"},
	ArgsUsage: "<service_name>",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "build",
			Value: "",
			Usage: "Build ID",
		},
	},
	Usage: "Show the current definition of a service by its name",
	Description: `Show the current definition of a service by its name getting the definition about the build.

   Example:
    $ ernest service definition myservice
	`,
	Action: func(c *cli.Context) error {
		if len(c.Args()) < 1 {
			color.Red("You should specify the service name")
		} else {
			m, cfg := setup(c)
			serviceName := c.Args()[0]
			if c.String("build") != "" {
				service, err := m.ServiceBuildStatus(cfg.Token, serviceName, c.String("build"))
				if err != nil {
					color.Red(err.Error())
					os.Exit(1)
				}
				fmt.Println(service.Definition)
			} else {
				service, err := m.ServiceStatus(cfg.Token, serviceName)
				if err != nil {
					color.Red(err.Error())
					os.Exit(1)
				}

				fmt.Println(service.Definition)
			}
		}
		return nil
	},
}

// InfoService : Shows detailed information of a service by its name
var InfoService = cli.Command{
	Name:      "info",
	Aliases:   []string{"i"},
	ArgsUsage: "<service_name>",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "build",
			Value: "",
			Usage: "Build ID",
		},
	},
	Usage: "$ ernest service info <my_service> --build <specific build>",
	Description: `Will show detailed information of the last build of a specified service.
	In case you specify --build option you will be able to output the detailed information of specific build of a service.

   Examples:
    $ ernest service definition myservice
    $ ernest service definition myservice --build build1
	`,
	Action: func(c *cli.Context) error {
		var err error
		var service Service

		m, cfg := setup(c)
		if cfg.Token == "" {
			color.Red("You're not allowed to perform this action, please log in")
			return nil
		}

		if len(c.Args()) == 0 {
			color.Red("You should specify an existing service name")
			return nil
		}

		name := c.Args()[0]
		if c.String("build") != "" {
			build := c.String("build")
			service, err = m.ServiceBuildStatus(cfg.Token, name, build)
		} else {
			service, err = m.ServiceStatus(cfg.Token, name)
		}

		if err != nil {
			color.Red(err.Error())
			os.Exit(1)
		}
		printServiceInfo(&service)
		return nil
	},
}

// CmdService ...
var CmdService = cli.Command{
	Name:  "service",
	Usage: "Service related subcommands",
	Subcommands: []cli.Command{
		ListServices,
		ApplyService,
		DestroyService,
		HistoryService,
		ResetService,
		DefinitionService,
		InfoService,
		MonitorService,
	},
}
<|endoftext|>


<commit_before>package sechat

import (
	"encoding/json"
	"time"
)

// wsRoom represents data from a specific chat room.
type wsRoom struct {
	Events []*Event `json:"e"`
}

// run is the main loop. It continually connects to the chat server, sleeping,
// and reconnecting upon error. It runs continually until stopped.
func (c *Conn) run(ch chan<- *Event) {
	defer close(c.closedCh)
	defer close(c.connectedCh)
	defer close(ch)
	defer c.log.Info("closing event channel")
	for {
		// Use the stored credentials to authenticate
		if err := c.auth(); err != nil {
			c.log.Error(err)
			goto retry
		}
		// Connect to to the websocket server
		if err := c.connectWebSocket(); err != nil {
			c.log.Error(err)
			goto retry
		}
		c.log.Info("connected to WebSocket")
		select {
		case c.connectedCh <- true:
		default:
		}
		// Event receiving loop
	loop:
		for {
			_, r, err := c.conn.NextReader()
			if err != nil {
				// Check to see if the error was caused by a shutdown or if it
				// is an actual error (in which case, leave the loop)
				select {
				case <-c.closeCh:
					return
				default:
					c.log.Error(err)
					break loop
				}
			}
			// Partially decode the message
			msg := map[string]json.RawMessage{}
			if err := json.NewDecoder(r).Decode(&msg); err != nil {
				continue
			}
			// Use a "set" to prevent duplicate events from being sent
			msgIDs := map[int]struct{}{}
			for _, v := range msg {
				room := &wsRoom{}
				if err := json.Unmarshal(v, &room); err != nil {
					continue
				}
				for _, e := range room.Events {
					if _, exists := msgIDs[e.ID]; !exists {
						e.precompute()
						// Use non-blocking send
						select {
						case ch <- e:
						default:
						}
						msgIDs[e.ID] = struct{}{}
					}
				}
			}
		}
	retry:
		select {
		case c.connectedCh <- false:
		default:
		}
		select {
		case <-c.closeCh:
			return
		default:
		}
		c.log.Info("reconnecting in 30 seconds")
		select {
		case <-time.After(30 * time.Second):
		case <-c.closeCh:
			return
		}
	}
}
<commit_msg>Add fields to connection messages.<commit_after>package sechat

import (
	"encoding/json"
	"time"

	"github.com/sirupsen/logrus"
)

// wsRoom represents data from a specific chat room.
type wsRoom struct {
	Events []*Event `json:"e"`
}

// run is the main loop. It continually connects to the chat server, sleeping,
// and reconnecting upon error. It runs continually until stopped.
func (c *Conn) run(ch chan<- *Event) {
	defer close(c.closedCh)
	defer close(c.connectedCh)
	defer close(ch)
	defer c.log.Info("closing event channel")
	for {
		// Use the stored credentials to authenticate
		if err := c.auth(); err != nil {
			c.log.Error(err)
			goto retry
		}
		// Connect to to the websocket server
		if err := c.connectWebSocket(); err != nil {
			c.log.Error(err)
			goto retry
		}
		c.log.WithFields(logrus.Fields{
			"connected": true,
		}).Info("connected to WebSocket")
		select {
		case c.connectedCh <- true:
		default:
		}
		// Event receiving loop
	loop:
		for {
			_, r, err := c.conn.NextReader()
			if err != nil {
				// Check to see if the error was caused by a shutdown or if it
				// is an actual error (in which case, leave the loop)
				select {
				case <-c.closeCh:
					return
				default:
					c.log.Error(err)
					break loop
				}
			}
			// Partially decode the message
			msg := map[string]json.RawMessage{}
			if err := json.NewDecoder(r).Decode(&msg); err != nil {
				continue
			}
			// Use a "set" to prevent duplicate events from being sent
			msgIDs := map[int]struct{}{}
			for _, v := range msg {
				room := &wsRoom{}
				if err := json.Unmarshal(v, &room); err != nil {
					continue
				}
				for _, e := range room.Events {
					if _, exists := msgIDs[e.ID]; !exists {
						e.precompute()
						// Use non-blocking send
						select {
						case ch <- e:
						default:
						}
						msgIDs[e.ID] = struct{}{}
					}
				}
			}
		}
	retry:
		c.log.WithFields(logrus.Fields{
			"connected": false,
		}).Info("disconnected from WebSocket")
		select {
		case c.connectedCh <- false:
		default:
		}
		select {
		case <-c.closeCh:
			return
		default:
		}
		c.log.Info("reconnecting in 30 seconds")
		select {
		case <-time.After(30 * time.Second):
		case <-c.closeCh:
			return
		}
	}
}
<|endoftext|>


<commit_before>package dna

import "testing"

func (h Histogram) Equal(o Histogram) bool {
	return h.sameLength(o) && h.sameMappings(o)
}

func (h Histogram) sameLength(o Histogram) bool {
	return len(h) == len(o)
}

func (h Histogram) sameMappings(o Histogram) (res bool) {
	res = true
	for k := range h {
		if h[k] != o[k] {
			res = false
		}
	}
	return
}

var tallyTests = []struct {
	strand     string
	nucleotide string
	expected   int
}{
	{"", "A", 0},
	{"ACT", "G", 0},
	{"CCCCC", "C", 5},
	{"GGGGGTAACCCGG", "T", 1},
}

func TestNucleotideCounts(t *testing.T) {
	for _, tt := range tallyTests {
		dna := DNA{tt.strand}
		count, _ := dna.Count(tt.nucleotide)
		if count != tt.expected {
			t.Fatalf("Got \"%v\", expected \"%v\"", count, tt.expected)
		}
	}
}

func TestHasErrorForInvalidNucleotides(t *testing.T) {
	dna := DNA{"GATTACA"}
	count, err := dna.Count("X")
	if count != 0 {
		t.Fatalf("Got \"%v\", expected \"%v\"", count, 0)
	}
	if err == nil {
		t.Fatalf("X is an invalid nucleotide, but no error was raised")
	}
}

// In most cases, this test is pointless.
// Very occasionally it matters.
// Just roll with it.
func TestCountingDoesntChangeCount(t *testing.T) {
	dna := DNA{"CGATTGGG"}
	dna.Count("T")
	count, _ := dna.Count("T")
	if count != 2 {
		t.Fatalf("Got \"%v\", expected \"%v\"", count, 2)
	}
}

type histogramTest struct {
	strand   string
	expected Histogram
}

var histogramTests = []histogramTest{
	{
		"",
		Histogram{"A": 0, "C": 0, "T": 0, "G": 0},
	},
	{
		"GGGGGGGG",
		Histogram{"A": 0, "C": 0, "T": 0, "G": 8},
	},
	{
		"AGCTTTTCATTCTGACTGCAACGGGCAATATGTCTCTGTGTGGATTAAAAAAAGAGTGTCTGATAGCAGC",
		Histogram{"A": 20, "C": 12, "T": 21, "G": 17},
	},
}

func TestSequenceHistograms(t *testing.T) {
	for _, tt := range histogramTests {
		dna := DNA{tt.strand}
		if !dna.Counts().Equal(tt.expected) {
			t.Fatalf("DNA{ \"%v\" }: Got \"%v\", expected \"%v\"", tt.strand, dna.Counts(), tt.expected)
		}
	}
}
<commit_msg>Added benchmark in go/nucleotide-count<commit_after>package dna

import "testing"

func (h Histogram) Equal(o Histogram) bool {
	return h.sameLength(o) && h.sameMappings(o)
}

func (h Histogram) sameLength(o Histogram) bool {
	return len(h) == len(o)
}

func (h Histogram) sameMappings(o Histogram) (res bool) {
	res = true
	for k := range h {
		if h[k] != o[k] {
			res = false
		}
	}
	return
}

var tallyTests = []struct {
	strand     string
	nucleotide string
	expected   int
}{
	{"", "A", 0},
	{"ACT", "G", 0},
	{"CCCCC", "C", 5},
	{"GGGGGTAACCCGG", "T", 1},
}

func TestNucleotideCounts(t *testing.T) {
	for _, tt := range tallyTests {
		dna := DNA{tt.strand}
		count, _ := dna.Count(tt.nucleotide)
		if count != tt.expected {
			t.Fatalf("Got \"%v\", expected \"%v\"", count, tt.expected)
		}
	}
}

func TestHasErrorForInvalidNucleotides(t *testing.T) {
	dna := DNA{"GATTACA"}
	count, err := dna.Count("X")
	if count != 0 {
		t.Fatalf("Got \"%v\", expected \"%v\"", count, 0)
	}
	if err == nil {
		t.Fatalf("X is an invalid nucleotide, but no error was raised")
	}
}

// In most cases, this test is pointless.
// Very occasionally it matters.
// Just roll with it.
func TestCountingDoesntChangeCount(t *testing.T) {
	dna := DNA{"CGATTGGG"}
	dna.Count("T")
	count, _ := dna.Count("T")
	if count != 2 {
		t.Fatalf("Got \"%v\", expected \"%v\"", count, 2)
	}
}

type histogramTest struct {
	strand   string
	expected Histogram
}

var histogramTests = []histogramTest{
	{
		"",
		Histogram{"A": 0, "C": 0, "T": 0, "G": 0},
	},
	{
		"GGGGGGGG",
		Histogram{"A": 0, "C": 0, "T": 0, "G": 8},
	},
	{
		"AGCTTTTCATTCTGACTGCAACGGGCAATATGTCTCTGTGTGGATTAAAAAAAGAGTGTCTGATAGCAGC",
		Histogram{"A": 20, "C": 12, "T": 21, "G": 17},
	},
}

func TestSequenceHistograms(t *testing.T) {
	for _, tt := range histogramTests {
		dna := DNA{tt.strand}
		if !dna.Counts().Equal(tt.expected) {
			t.Fatalf("DNA{ \"%v\" }: Got \"%v\", expected \"%v\"", tt.strand, dna.Counts(), tt.expected)
		}
	}
}

func BenchmarkSequenceHistograms(b *testing.B) {
	b.StopTimer()
	for _, tt := range histogramTests {
		for i := 0; i < b.N; i++ {
			dna := DNA{tt.strand}
			b.StartTimer()

			dna.Counts()

			b.StopTimer()
		}
	}
}
<|endoftext|>