<commit_before>// Copyright 2016 Richard Lehane. All rights reserved.
//
// 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 loc

import (
	"fmt"
	"sort"
	"strings"

	"github.com/richardlehane/siegfried/internal/identifier"
	"github.com/richardlehane/siegfried/internal/persist"
	"github.com/richardlehane/siegfried/pkg/config"
	"github.com/richardlehane/siegfried/pkg/core"
)

func init() {
	core.RegisterIdentifier(core.LOC, Load)
}

type Identifier struct {
	infos map[string]formatInfo
	*identifier.Base
}

func (i *Identifier) Save(ls *persist.LoadSaver) {
	ls.SaveByte(core.LOC)
	ls.SaveSmallInt(len(i.infos))
	for k, v := range i.infos {
		ls.SaveString(k)
		ls.SaveString(v.name)
		ls.SaveString(v.longName)
		ls.SaveString(v.mimeType)
	}
	i.Base.Save(ls)
}

func Load(ls *persist.LoadSaver) core.Identifier {
	i := &Identifier{}
	i.infos = make(map[string]formatInfo)
	le := ls.LoadSmallInt()
	for j := 0; j < le; j++ {
		i.infos[ls.LoadString()] = formatInfo{
			ls.LoadString(),
			ls.LoadString(),
			ls.LoadString(),
		}
	}
	i.Base = identifier.Load(ls)
	return i
}

func New(opts ...config.Option) (core.Identifier, error) {
	for _, v := range opts {
		v()
	}
	loc, err := newLOC(config.LOC())
	if err != nil {
		return nil, err
	}
	// set updated
	updated := loc.(fdds).Updated().Format(dateFmt)
	// add extensions
	for _, v := range config.Extend() {
		e, err := newLOC(v)
		if err != nil {
			return nil, fmt.Errorf("LOC: error loading extension file %s; got %s", v, err)
		}
		loc = identifier.Join(loc, e)
	}
	// apply config
	loc = identifier.ApplyConfig(loc)
	// return identifier
	return &Identifier{
		infos: infos(loc.Infos()),
		Base:  identifier.New(loc, config.ZipLOC(), updated),
	}, nil
}

func (i *Identifier) Fields() []string {
	return []string{"namespace", "id", "format", "full", "mime", "basis", "warning"}
}

func (i *Identifier) Recorder() core.Recorder {
	return &Recorder{
		Identifier: i,
		ids:        make(pids, 0, 1),
	}
}

type Recorder struct {
	*Identifier
	ids        pids
	cscore     int
	satisfied  bool
	extActive  bool
	mimeActive bool
	textActive bool
}

const (
	extScore = 1 << iota
	mimeScore
	textScore
	incScore
)

func (r *Recorder) Active(m core.MatcherType) {
	if r.Identifier.Active(m) {
		switch m {
		case core.NameMatcher:
			r.extActive = true
		case core.MIMEMatcher:
			r.mimeActive = true
		case core.TextMatcher:
			r.textActive = true
		}
	}
}

func (r *Recorder) Record(m core.MatcherType, res core.Result) bool {
	switch m {
	default:
		return false
	case core.NameMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			r.ids = add(r.ids, r.Name(), id, r.infos[id], res.Basis(), extScore)
			return true
		} else {
			return false
		}
	case core.MIMEMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			r.ids = add(r.ids, r.Name(), id, r.infos[id], res.Basis(), mimeScore)
			return true
		} else {
			return false
		}
	case core.ContainerMatcher:
		// add zip default
		if res.Index() < 0 {
			if r.ZipDefault() {
				r.cscore += incScore
				r.ids = add(r.ids, r.Name(), config.ZipLOC(), r.infos[config.ZipLOC()], res.Basis(), r.cscore)
			}
			return false
		}
		if hit, id := r.Hit(m, res.Index()); hit {
			r.cscore += incScore
			basis := res.Basis()
			p, t := r.Place(core.ContainerMatcher, res.Index())
			if t > 1 {
				basis = basis + fmt.Sprintf(" (signature %d/%d)", p, t)
			}
			r.ids = add(r.ids, r.Name(), id, r.infos[id], basis, r.cscore)
			return true
		} else {
			return false
		}
	case core.RIFFMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			if r.satisfied {
				return true
			}
			r.cscore += incScore
			r.ids = add(r.ids, r.Name(), id, r.infos[id], res.Basis(), r.cscore)
			return true
		} else {
			return false
		}
	case core.ByteMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			if r.satisfied {
				return true
			}
			r.cscore += incScore
			basis := res.Basis()
			p, t := r.Place(core.ByteMatcher, res.Index())
			if t > 1 {
				basis = basis + fmt.Sprintf(" (signature %d/%d)", p, t)
			}
			r.ids = add(r.ids, r.Name(), id, r.infos[id], basis, r.cscore)
			return true
		} else {
			return false
		}
	}
}

func (r *Recorder) Satisfied(mt core.MatcherType) (bool, core.Hint) {
	if r.NoPriority() {
		return false, core.Hint{}
	}
	if r.cscore < incScore {
		if mt == core.ByteMatcher || mt == core.XMLMatcher || mt == core.RIFFMatcher {
			return false, core.Hint{}
		}
		if len(r.ids) == 0 {
			return false, core.Hint{}
		}
	}
	r.satisfied = true
	if mt == core.ByteMatcher {
		return true, core.Hint{r.Start(mt), nil}
	}
	return true, core.Hint{}
}

func lowConfidence(conf int) string {
	var ls = make([]string, 0, 1)
	if conf&extScore == extScore {
		ls = append(ls, "extension")
	}
	if conf&mimeScore == mimeScore {
		ls = append(ls, "MIME")
	}
	if conf&textScore == textScore {
		ls = append(ls, "text")
	}
	switch len(ls) {
	case 0:
		return ""
	case 1:
		return ls[0]
	case 2:
		return ls[0] + " and " + ls[1]
	default:
		return strings.Join(ls[:len(ls)-1], ", ") + " and " + ls[len(ls)-1]
	}
}

func (r *Recorder) Report() []core.Identification {
	// no results
	if len(r.ids) == 0 {
		return []core.Identification{Identification{
			Namespace: r.Name(),
			ID:        "UNKNOWN",
			Warning:   "no match",
		}}
	}
	sort.Sort(r.ids)
	// exhaustive
	if r.Multi() == config.Exhaustive {
		ret := make([]core.Identification, len(r.ids))
		for i, v := range r.ids {
			ret[i] = r.updateWarning(v)
		}
		return ret
	}
	conf := r.ids[0].confidence
	// if we've only got extension / mime matches, check if those matches are ruled out by lack of byte match
	// only permit a single extension or mime only match
	// add warnings too
	if conf <= textScore {
		nids := make([]Identification, 0, 1)
		for _, v := range r.ids {
			// if overall confidence is greater than mime or ext only, then rule out any lesser confident matches
			if conf > mimeScore && v.confidence != conf {
				break
			}
			// if the match has no corresponding byte or RIFF signature...
			if ok := r.HasSig(v.ID, core.RIFFMatcher, core.ByteMatcher); !ok {
				// break immediately if more than one match
				if len(nids) > 0 {
					nids = nids[:0]
					break
				}
				nids = append(nids, v)
			}
		}
		if len(nids) != 1 {
			poss := make([]string, len(r.ids))
			for i, v := range r.ids {
				poss[i] = v.ID
				conf = conf | v.confidence
			}
			return []core.Identification{Identification{
				Namespace: r.Name(),
				ID:        "UNKNOWN",
				Warning:   fmt.Sprintf("no match; possibilities based on %v are %v", lowConfidence(conf), strings.Join(poss, ", ")),
			}}
		}
		r.ids = nids
	}
	// handle single result only
	if r.Multi() == config.Single && len(r.ids) > 1 && r.ids[0].confidence == r.ids[1].confidence {
		poss := make([]string, 0, len(r.ids))
		for _, v := range r.ids {
			if v.confidence < conf {
				break
			}
			poss = append(poss, v.ID)
		}
		return []core.Identification{Identification{
			Namespace: r.Name(),
			ID:        "UNKNOWN",
			Warning:   fmt.Sprintf("multiple matches %v", strings.Join(poss, ", ")),
		}}
	}
	ret := make([]core.Identification, len(r.ids))
	for i, v := range r.ids {
		if i > 0 {
			switch r.Multi() {
			case config.Single:
				return ret[:i]
			case config.Conclusive:
				if v.confidence < conf {
					return ret[:i]
				}
			default:
				if v.confidence < incScore {
					return ret[:i]
				}
			}
		}
		ret[i] = r.updateWarning(v)
	}
	return ret
}

func (r *Recorder) updateWarning(i Identification) Identification {
	// apply low confidence
	if i.confidence <= textScore {
		if len(i.Warning) > 0 {
			i.Warning += "; " + "match on " + lowConfidence(i.confidence) + " only"
		} else {
			i.Warning = "match on " + lowConfidence(i.confidence) + " only"
		}
	}
	// apply mismatches
	if r.extActive && (i.confidence&extScore != extScore) {
		for _, v := range r.IDs(core.NameMatcher) {
			if i.ID == v {
				if len(i.Warning) > 0 {
					i.Warning += "; extension mismatch"
				} else {
					i.Warning = "extension mismatch"
				}
				break
			}
		}
	}
	if r.mimeActive && (i.confidence&mimeScore != mimeScore) {
		for _, v := range r.IDs(core.MIMEMatcher) {
			if i.ID == v {
				if len(i.Warning) > 0 {
					i.Warning += "; MIME mismatch"
				} else {
					i.Warning = "MIME mismatch"
				}
				break
			}
		}
	}
	return i
}

type Identification struct {
	Namespace  string
	ID         string
	Name       string
	LongName   string
	MIME       string
	Basis      []string
	Warning    string
	archive    config.Archive
	confidence int
}

func (id Identification) String() string {
	return id.ID
}

func (id Identification) Known() bool {
	return id.ID != "UNKNOWN"
}

func (id Identification) Warn() string {
	return id.Warning
}

func (id Identification) Values() []string {
	var basis string
	if len(id.Basis) > 0 {
		basis = strings.Join(id.Basis, "; ")
	}
	return []string{
		id.Namespace,
		id.ID,
		id.Name,
		id.LongName,
		id.MIME,
		basis,
		id.Warning,
	}
}

func (id Identification) Archive() config.Archive {
	return id.archive
}

type pids []Identification

func (p pids) Len() int { return len(p) }

func (p pids) Less(i, j int) bool { return p[j].confidence < p[i].confidence }

func (p pids) Swap(i, j int) { p[i], p[j] = p[j], p[i] }

func add(p pids, id string, f string, info formatInfo, basis string, c int) pids {
	for i, v := range p {
		if v.ID == f {
			p[i].confidence += c
			p[i].Basis = append(p[i].Basis, basis)
			return p
		}
	}
	return append(p, Identification{id, f, info.name, info.longName, info.mimeType, []string{basis}, "", config.IsArchive(f), c})
}
<commit_msg>fix LOC identifier<commit_after>// Copyright 2016 Richard Lehane. All rights reserved.
//
// 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 loc

import (
	"fmt"
	"sort"
	"strings"

	"github.com/richardlehane/siegfried/internal/identifier"
	"github.com/richardlehane/siegfried/internal/persist"
	"github.com/richardlehane/siegfried/pkg/config"
	"github.com/richardlehane/siegfried/pkg/core"
)

func init() {
	core.RegisterIdentifier(core.LOC, Load)
}

type Identifier struct {
	infos map[string]formatInfo
	*identifier.Base
}

func (i *Identifier) Save(ls *persist.LoadSaver) {
	ls.SaveByte(core.LOC)
	ls.SaveSmallInt(len(i.infos))
	for k, v := range i.infos {
		ls.SaveString(k)
		ls.SaveString(v.name)
		ls.SaveString(v.longName)
		ls.SaveString(v.mimeType)
	}
	i.Base.Save(ls)
}

func Load(ls *persist.LoadSaver) core.Identifier {
	i := &Identifier{}
	i.infos = make(map[string]formatInfo)
	le := ls.LoadSmallInt()
	for j := 0; j < le; j++ {
		i.infos[ls.LoadString()] = formatInfo{
			ls.LoadString(),
			ls.LoadString(),
			ls.LoadString(),
		}
	}
	i.Base = identifier.Load(ls)
	return i
}

func New(opts ...config.Option) (core.Identifier, error) {
	for _, v := range opts {
		v()
	}
	loc, err := newLOC(config.LOC())
	if err != nil {
		return nil, err
	}
	// set updated
	updated := loc.(fdds).Updated().Format(dateFmt)
	// add extensions
	for _, v := range config.Extend() {
		e, err := newLOC(v)
		if err != nil {
			return nil, fmt.Errorf("LOC: error loading extension file %s; got %s", v, err)
		}
		loc = identifier.Join(loc, e)
	}
	// apply config
	loc = identifier.ApplyConfig(loc)
	// return identifier
	return &Identifier{
		infos: infos(loc.Infos()),
		Base:  identifier.New(loc, config.ZipLOC(), updated),
	}, nil
}

func (i *Identifier) Fields() []string {
	return []string{"namespace", "id", "format", "full", "mime", "basis", "warning"}
}

func (i *Identifier) Recorder() core.Recorder {
	return &Recorder{
		Identifier: i,
		ids:        make(pids, 0, 1),
	}
}

type Recorder struct {
	*Identifier
	ids        pids
	cscore     int
	satisfied  bool
	extActive  bool
	mimeActive bool
	textActive bool
}

const (
	extScore = 1 << iota
	mimeScore
	textScore
	incScore
)

func (r *Recorder) Active(m core.MatcherType) {
	if r.Identifier.Active(m) {
		switch m {
		case core.NameMatcher:
			r.extActive = true
		case core.MIMEMatcher:
			r.mimeActive = true
		case core.TextMatcher:
			r.textActive = true
		}
	}
}

func (r *Recorder) Record(m core.MatcherType, res core.Result) bool {
	switch m {
	default:
		return false
	case core.NameMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			r.ids = add(r.ids, r.Name(), id, r.infos[id], res.Basis(), extScore)
			return true
		} else {
			return false
		}
	case core.MIMEMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			r.ids = add(r.ids, r.Name(), id, r.infos[id], res.Basis(), mimeScore)
			return true
		} else {
			return false
		}
	case core.ContainerMatcher:
		// add zip default
		if res.Index() < 0 {
			if r.ZipDefault() {
				r.cscore += incScore
				r.ids = add(r.ids, r.Name(), config.ZipLOC(), r.infos[config.ZipLOC()], res.Basis(), r.cscore)
			}
			return false
		}
		if hit, id := r.Hit(m, res.Index()); hit {
			r.cscore += incScore
			basis := res.Basis()
			p, t := r.Place(core.ContainerMatcher, res.Index())
			if t > 1 {
				basis = basis + fmt.Sprintf(" (signature %d/%d)", p, t)
			}
			r.ids = add(r.ids, r.Name(), id, r.infos[id], basis, r.cscore)
			return true
		} else {
			return false
		}
	case core.RIFFMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			if r.satisfied {
				return true
			}
			r.cscore += incScore
			r.ids = add(r.ids, r.Name(), id, r.infos[id], res.Basis(), r.cscore)
			return true
		} else {
			return false
		}
	case core.ByteMatcher:
		if hit, id := r.Hit(m, res.Index()); hit {
			if r.satisfied {
				return true
			}
			r.cscore += incScore
			basis := res.Basis()
			p, t := r.Place(core.ByteMatcher, res.Index())
			if t > 1 {
				basis = basis + fmt.Sprintf(" (signature %d/%d)", p, t)
			}
			r.ids = add(r.ids, r.Name(), id, r.infos[id], basis, r.cscore)
			return true
		} else {
			return false
		}
	}
}

func (r *Recorder) Satisfied(mt core.MatcherType) (bool, core.Hint) {
	if r.NoPriority() {
		return false, core.Hint{}
	}
	if r.cscore < incScore {
		if mt == core.ContainerMatcher || mt == core.ByteMatcher || mt == core.XMLMatcher || mt == core.RIFFMatcher {
			return false, core.Hint{}
		}
		if len(r.ids) == 0 {
			return false, core.Hint{}
		}
	}
	r.satisfied = true
	if mt == core.ByteMatcher {
		return true, core.Hint{r.Start(mt), nil}
	}
	return true, core.Hint{}
}

func lowConfidence(conf int) string {
	var ls = make([]string, 0, 1)
	if conf&extScore == extScore {
		ls = append(ls, "extension")
	}
	if conf&mimeScore == mimeScore {
		ls = append(ls, "MIME")
	}
	if conf&textScore == textScore {
		ls = append(ls, "text")
	}
	switch len(ls) {
	case 0:
		return ""
	case 1:
		return ls[0]
	case 2:
		return ls[0] + " and " + ls[1]
	default:
		return strings.Join(ls[:len(ls)-1], ", ") + " and " + ls[len(ls)-1]
	}
}

func (r *Recorder) Report() []core.Identification {
	// no results
	if len(r.ids) == 0 {
		return []core.Identification{Identification{
			Namespace: r.Name(),
			ID:        "UNKNOWN",
			Warning:   "no match",
		}}
	}
	sort.Sort(r.ids)
	// exhaustive
	if r.Multi() == config.Exhaustive {
		ret := make([]core.Identification, len(r.ids))
		for i, v := range r.ids {
			ret[i] = r.updateWarning(v)
		}
		return ret
	}
	conf := r.ids[0].confidence
	// if we've only got extension / mime matches, check if those matches are ruled out by lack of byte match
	// only permit a single extension or mime only match
	// add warnings too
	if conf <= textScore {
		nids := make([]Identification, 0, 1)
		for _, v := range r.ids {
			// if overall confidence is greater than mime or ext only, then rule out any lesser confident matches
			if conf > mimeScore && v.confidence != conf {
				break
			}
			// if the match has no corresponding byte or RIFF signature...
			if ok := r.HasSig(v.ID, core.RIFFMatcher, core.ByteMatcher); !ok {
				// break immediately if more than one match
				if len(nids) > 0 {
					nids = nids[:0]
					break
				}
				nids = append(nids, v)
			}
		}
		if len(nids) != 1 {
			poss := make([]string, len(r.ids))
			for i, v := range r.ids {
				poss[i] = v.ID
				conf = conf | v.confidence
			}
			return []core.Identification{Identification{
				Namespace: r.Name(),
				ID:        "UNKNOWN",
				Warning:   fmt.Sprintf("no match; possibilities based on %v are %v", lowConfidence(conf), strings.Join(poss, ", ")),
			}}
		}
		r.ids = nids
	}
	// handle single result only
	if r.Multi() == config.Single && len(r.ids) > 1 && r.ids[0].confidence == r.ids[1].confidence {
		poss := make([]string, 0, len(r.ids))
		for _, v := range r.ids {
			if v.confidence < conf {
				break
			}
			poss = append(poss, v.ID)
		}
		return []core.Identification{Identification{
			Namespace: r.Name(),
			ID:        "UNKNOWN",
			Warning:   fmt.Sprintf("multiple matches %v", strings.Join(poss, ", ")),
		}}
	}
	ret := make([]core.Identification, len(r.ids))
	for i, v := range r.ids {
		if i > 0 {
			switch r.Multi() {
			case config.Single:
				return ret[:i]
			case config.Conclusive:
				if v.confidence < conf {
					return ret[:i]
				}
			default:
				if v.confidence < incScore {
					return ret[:i]
				}
			}
		}
		ret[i] = r.updateWarning(v)
	}
	return ret
}

func (r *Recorder) updateWarning(i Identification) Identification {
	// apply low confidence
	if i.confidence <= textScore {
		if len(i.Warning) > 0 {
			i.Warning += "; " + "match on " + lowConfidence(i.confidence) + " only"
		} else {
			i.Warning = "match on " + lowConfidence(i.confidence) + " only"
		}
	}
	// apply mismatches
	if r.extActive && (i.confidence&extScore != extScore) {
		for _, v := range r.IDs(core.NameMatcher) {
			if i.ID == v {
				if len(i.Warning) > 0 {
					i.Warning += "; extension mismatch"
				} else {
					i.Warning = "extension mismatch"
				}
				break
			}
		}
	}
	if r.mimeActive && (i.confidence&mimeScore != mimeScore) {
		for _, v := range r.IDs(core.MIMEMatcher) {
			if i.ID == v {
				if len(i.Warning) > 0 {
					i.Warning += "; MIME mismatch"
				} else {
					i.Warning = "MIME mismatch"
				}
				break
			}
		}
	}
	return i
}

type Identification struct {
	Namespace  string
	ID         string
	Name       string
	LongName   string
	MIME       string
	Basis      []string
	Warning    string
	archive    config.Archive
	confidence int
}

func (id Identification) String() string {
	return id.ID
}

func (id Identification) Known() bool {
	return id.ID != "UNKNOWN"
}

func (id Identification) Warn() string {
	return id.Warning
}

func (id Identification) Values() []string {
	var basis string
	if len(id.Basis) > 0 {
		basis = strings.Join(id.Basis, "; ")
	}
	return []string{
		id.Namespace,
		id.ID,
		id.Name,
		id.LongName,
		id.MIME,
		basis,
		id.Warning,
	}
}

func (id Identification) Archive() config.Archive {
	return id.archive
}

type pids []Identification

func (p pids) Len() int { return len(p) }

func (p pids) Less(i, j int) bool { return p[j].confidence < p[i].confidence }

func (p pids) Swap(i, j int) { p[i], p[j] = p[j], p[i] }

func add(p pids, id string, f string, info formatInfo, basis string, c int) pids {
	for i, v := range p {
		if v.ID == f {
			p[i].confidence += c
			p[i].Basis = append(p[i].Basis, basis)
			return p
		}
	}
	return append(p, Identification{id, f, info.name, info.longName, info.mimeType, []string{basis}, "", config.IsArchive(f), c})
}
<|endoftext|>


<commit_before>// Copyright © 2016 Benjamin Martensson <benjamin.martensson@nrk.no>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var cfgFile string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "tpm",
	Short: "TPM Command-line interface",
	Long:  "Team Password Manager CLI tool to easily access your passwords.",
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	// Here you will define your flags and configuration settings.
	// Cobra supports Persistent Flags, which, if defined here,
	// will be global for your application.

	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.tpm.yaml)")
	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	// RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" { // enable ability to specify config file via flag
		viper.SetConfigFile(cfgFile)
	}

	viper.SetConfigName(".tpm")  // name of config file (without extension)
	viper.AddConfigPath("$HOME") // adding home directory as first search path
	viper.AutomaticEnv()         // read in environment variables that match

	// If a config file is found, read it in.
	viper.ReadInConfig()
}
<commit_msg>correct should be manage<commit_after>// Copyright © 2016 Benjamin Martensson <benjamin.martensson@nrk.no>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var cfgFile string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "tpm",
	Short: "TPM Command-line interface",
	Long:  "Team Password Manager CLI tool to easily manage your passwords.",
}

// Execute adds all child commands to the root command sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(-1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	// Here you will define your flags and configuration settings.
	// Cobra supports Persistent Flags, which, if defined here,
	// will be global for your application.

	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.tpm.yaml)")
	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	// RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" { // enable ability to specify config file via flag
		viper.SetConfigFile(cfgFile)
	}

	viper.SetConfigName(".tpm")  // name of config file (without extension)
	viper.AddConfigPath("$HOME") // adding home directory as first search path
	viper.AutomaticEnv()         // read in environment variables that match

	// If a config file is found, read it in.
	viper.ReadInConfig()
}
<|endoftext|>


<commit_before>// Copyright 2016 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"time"

	"github.com/la5nta/wl2k-go/fbb"
)

type ex struct {
	conn   net.Conn
	target string
	master bool
	errors chan error
}

func exchangeLoop() (ce chan ex) {
	ce = make(chan ex)
	go func() {
		for ex := range ce {
			ex.errors <- sessionExchange(ex.conn, ex.target, ex.master)
			close(ex.errors)
		}
	}()
	return ce
}

func exchange(conn net.Conn, targetCall string, master bool) error {
	e := ex{
		conn:   conn,
		target: targetCall,
		master: master,
		errors: make(chan error),
	}
	exchangeChan <- e
	return <-e.errors
}

type NotifyMBox struct{ fbb.MBoxHandler }

func (m NotifyMBox) ProcessInbound(msgs ...*fbb.Message) error {
	if err := m.MBoxHandler.ProcessInbound(msgs...); err != nil {
		return err
	}
	for _, msg := range msgs {
		websocketHub.WriteJSON(struct{ Notification Notification }{
			Notification{
				Title: fmt.Sprintf("New message from %s", msg.From().Addr),
				Body:  msg.Subject(),
			},
		})
	}
	return nil
}

func sessionExchange(conn net.Conn, targetCall string, master bool) error {
	exchangeConn = conn
	websocketHub.UpdateStatus()
	defer func() { exchangeConn = nil; websocketHub.UpdateStatus() }()

	// New wl2k Session
	targetCall = strings.Split(targetCall, ` `)[0]
	session := fbb.NewSession(
		fOptions.MyCall,
		targetCall,
		config.Locator,
		NotifyMBox{mbox},
	)

	session.SetUserAgent(fbb.UserAgent{
		Name:    AppName,
		Version: Version,
	})

	if len(config.MOTD) > 0 {
		session.SetMOTD(config.MOTD...)
	}

	// Handle secure login
	session.SetSecureLoginHandleFunc(func() (string, error) {
		if config.SecureLoginPassword != "" {
			return config.SecureLoginPassword, nil
		}
		resp := <-promptHub.Prompt("password", "Enter secure login password")
		return resp.Value, resp.Err
	})

	for _, addr := range config.AuxAddrs {
		session.AddAuxiliaryAddress(fbb.AddressFromString(addr))
	}

	session.IsMaster(master)
	session.SetLogger(log.New(logWriter, "", 0))

	session.SetStatusUpdater(new(StatusUpdate))

	if fOptions.Robust {
		session.SetRobustMode(fbb.RobustForced)
	}

	log.Printf("Connected to %s (%s)", conn.RemoteAddr(), conn.RemoteAddr().Network())

	// Close connection on os.Interrupt
	stop := handleInterrupt()
	defer close(stop)

	startTs := time.Now()

	stats, err := session.Exchange(conn)
	if fbb.IsLoginFailure(err) {
		fmt.Println("NOTE: A new password scheme for Winlink is being implemented as of 2018-01-31.")
		fmt.Println("      Users with passwords created/changed prior to January 31, 2018 should be")
		fmt.Println("      aware that their password MUST be entered in ALL-UPPERCASE letters. Only")
		fmt.Println("      passwords created/changed/issued after January 31, 2018 should/may contain")
		fmt.Println("      lowercase letters. - https://github.com/la5nta/pat/issues/113")
	}

	event := map[string]interface{}{
		"mycall":              session.Mycall(),
		"targetcall":          session.Targetcall(),
		"remote_fw":           session.RemoteForwarders(),
		"remote_sid":          session.RemoteSID(),
		"master":              master,
		"local_locator":       config.Locator,
		"auxiliary_addresses": config.AuxAddrs,
		"network":             conn.RemoteAddr().Network(),
		"remote_addr":         conn.RemoteAddr().String(),
		"local_addr":          conn.LocalAddr().String(),
		"sent":                stats.Sent,
		"received":            stats.Received,
		"start":               startTs.Unix(),
		"end":                 time.Now().Unix(),
		"success":             err == nil,
	}
	if err != nil {
		event["error"] = err.Error()
	}

	eventLog.Log("exchange", event)

	return err
}

func handleInterrupt() (stop chan struct{}) {
	stop = make(chan struct{})

	go func() {
		sig := make(chan os.Signal)
		signal.Notify(sig, os.Interrupt)
		defer func() { signal.Stop(sig); close(sig) }()

		wmDisc := false // So we can DirtyDisconnect on second interrupt
		adDisc := false // So we can Abort on second interrupt
		for {
			select {
			case <-stop:
				return
			case s := <-sig:
				if exchangeConn != nil {
					log.Printf("Got %s, disconnecting...", s)
					exchangeConn.Close()
					break
				}
				if pModem != nil {
					log.Println("Disconnecting pactor...")
					if err := pModem.Close(); err != nil {
						log.Println(err)
					}
					break
				}
				if wmTNC != nil && !wmTNC.Idle() {
					if wmDisc {
						log.Println("Dirty disconnecting winmor...")
						wmTNC.DirtyDisconnect()
						wmDisc = false
					} else {
						log.Println("Disconnecting winmor...")
						wmDisc = true
						go func() {
							if err := wmTNC.Disconnect(); err != nil {
								log.Println(err)
							} else {
								wmDisc = false
							}
						}()
					}
				}
				if adTNC != nil && !adTNC.Idle() {
					if adDisc {
						log.Println("Dirty disconnecting ardop...")
						adTNC.Abort()
						adDisc = false
					} else {
						log.Println("Disconnecting ardop...")
						adDisc = true
						go func() {
							if err := adTNC.Disconnect(); err != nil {
								log.Println(err)
							} else {
								adDisc = false
							}
						}()
					}
				}
			}
		}
	}()

	return stop
}

type StatusUpdate int

func (s *StatusUpdate) UpdateStatus(stat fbb.Status) {
	var prop fbb.Proposal
	switch {
	case stat.Receiving != nil:
		prop = *stat.Receiving
	case stat.Sending != nil:
		prop = *stat.Sending
	}

	websocketHub.WriteProgress(Progress{
		MID:              prop.MID(),
		BytesTotal:       stat.BytesTotal,
		BytesTransferred: stat.BytesTransferred,
		Subject:          prop.Title(),
		Receiving:        stat.Receiving != nil,
		Sending:          stat.Sending != nil,
		Done:             stat.Done,
	})

	percent := float64(stat.BytesTransferred) / float64(stat.BytesTotal) * 100
	fmt.Printf("\r%s: %3.0f%%", prop.Title(), percent)

	if stat.Done {
		fmt.Println("")
	}
	os.Stdout.Sync()
}
<commit_msg>Refactor handleInterrupt to be re-used it for web<commit_after>// Copyright 2016 Martin Hebnes Pedersen (LA5NTA). All rights reserved.
// Use of this source code is governed by the MIT-license that can be
// found in the LICENSE file.

package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"time"

	"github.com/la5nta/wl2k-go/fbb"
)

type ex struct {
	conn   net.Conn
	target string
	master bool
	errors chan error
}

func exchangeLoop() (ce chan ex) {
	ce = make(chan ex)
	go func() {
		for ex := range ce {
			ex.errors <- sessionExchange(ex.conn, ex.target, ex.master)
			close(ex.errors)
		}
	}()
	return ce
}

func exchange(conn net.Conn, targetCall string, master bool) error {
	e := ex{
		conn:   conn,
		target: targetCall,
		master: master,
		errors: make(chan error),
	}
	exchangeChan <- e
	return <-e.errors
}

type NotifyMBox struct{ fbb.MBoxHandler }

func (m NotifyMBox) ProcessInbound(msgs ...*fbb.Message) error {
	if err := m.MBoxHandler.ProcessInbound(msgs...); err != nil {
		return err
	}
	for _, msg := range msgs {
		websocketHub.WriteJSON(struct{ Notification Notification }{
			Notification{
				Title: fmt.Sprintf("New message from %s", msg.From().Addr),
				Body:  msg.Subject(),
			},
		})
	}
	return nil
}

func sessionExchange(conn net.Conn, targetCall string, master bool) error {
	exchangeConn = conn
	websocketHub.UpdateStatus()
	defer func() { exchangeConn = nil; websocketHub.UpdateStatus() }()

	// New wl2k Session
	targetCall = strings.Split(targetCall, ` `)[0]
	session := fbb.NewSession(
		fOptions.MyCall,
		targetCall,
		config.Locator,
		NotifyMBox{mbox},
	)

	session.SetUserAgent(fbb.UserAgent{
		Name:    AppName,
		Version: Version,
	})

	if len(config.MOTD) > 0 {
		session.SetMOTD(config.MOTD...)
	}

	// Handle secure login
	session.SetSecureLoginHandleFunc(func() (string, error) {
		if config.SecureLoginPassword != "" {
			return config.SecureLoginPassword, nil
		}
		resp := <-promptHub.Prompt("password", "Enter secure login password")
		return resp.Value, resp.Err
	})

	for _, addr := range config.AuxAddrs {
		session.AddAuxiliaryAddress(fbb.AddressFromString(addr))
	}

	session.IsMaster(master)
	session.SetLogger(log.New(logWriter, "", 0))

	session.SetStatusUpdater(new(StatusUpdate))

	if fOptions.Robust {
		session.SetRobustMode(fbb.RobustForced)
	}

	log.Printf("Connected to %s (%s)", conn.RemoteAddr(), conn.RemoteAddr().Network())

	// Close connection on os.Interrupt
	stop := handleInterrupt()
	defer close(stop)

	startTs := time.Now()

	stats, err := session.Exchange(conn)
	if fbb.IsLoginFailure(err) {
		fmt.Println("NOTE: A new password scheme for Winlink is being implemented as of 2018-01-31.")
		fmt.Println("      Users with passwords created/changed prior to January 31, 2018 should be")
		fmt.Println("      aware that their password MUST be entered in ALL-UPPERCASE letters. Only")
		fmt.Println("      passwords created/changed/issued after January 31, 2018 should/may contain")
		fmt.Println("      lowercase letters. - https://github.com/la5nta/pat/issues/113")
	}

	event := map[string]interface{}{
		"mycall":              session.Mycall(),
		"targetcall":          session.Targetcall(),
		"remote_fw":           session.RemoteForwarders(),
		"remote_sid":          session.RemoteSID(),
		"master":              master,
		"local_locator":       config.Locator,
		"auxiliary_addresses": config.AuxAddrs,
		"network":             conn.RemoteAddr().Network(),
		"remote_addr":         conn.RemoteAddr().String(),
		"local_addr":          conn.LocalAddr().String(),
		"sent":                stats.Sent,
		"received":            stats.Received,
		"start":               startTs.Unix(),
		"end":                 time.Now().Unix(),
		"success":             err == nil,
	}
	if err != nil {
		event["error"] = err.Error()
	}

	eventLog.Log("exchange", event)

	return err
}

func handleInterrupt() (stop chan struct{}) {
	stop = make(chan struct{})

	go func() {
		sig := make(chan os.Signal)
		signal.Notify(sig, os.Interrupt)
		defer func() { signal.Stop(sig); close(sig) }()

		dirtyDisconnectNext := false // So we can do a dirty disconnect on the second interrupt
		for {
			select {
			case <-stop:
				return
			case <-sig:
				abortActiveConnection(dirtyDisconnectNext)
				dirtyDisconnectNext = !dirtyDisconnectNext
			}
		}
	}()

	return stop
}

func abortActiveConnection(dirty bool) (ok bool) {
	switch {
	case exchangeConn != nil:
		log.Println("Got abort signal, disconnecting...")
		exchangeConn.Close()
		return true
	case pModem != nil:
		log.Println("Disconnecting pactor...")
		err := pModem.Close()
		if err != nil {
			log.Println(err)
		}
		return err == nil
	case wmTNC != nil && !wmTNC.Idle():
		if dirty {
			log.Println("Dirty disconnecting winmor...")
			wmTNC.DirtyDisconnect()
			return true
		}
		log.Println("Disconnecting winmor...")
		go func() {
			if err := wmTNC.Disconnect(); err != nil {
				log.Println(err)
			}
		}()
		return true
	case adTNC != nil && !adTNC.Idle():
		if dirty {
			log.Println("Dirty disconnecting ardop...")
			adTNC.Abort()
			return true
		}
		log.Println("Disconnecting ardop...")
		go func() {
			if err := adTNC.Disconnect(); err != nil {
				log.Println(err)
			}
		}()
		return true
	default:
		return false
	}
}

type StatusUpdate int

func (s *StatusUpdate) UpdateStatus(stat fbb.Status) {
	var prop fbb.Proposal
	switch {
	case stat.Receiving != nil:
		prop = *stat.Receiving
	case stat.Sending != nil:
		prop = *stat.Sending
	}

	websocketHub.WriteProgress(Progress{
		MID:              prop.MID(),
		BytesTotal:       stat.BytesTotal,
		BytesTransferred: stat.BytesTransferred,
		Subject:          prop.Title(),
		Receiving:        stat.Receiving != nil,
		Sending:          stat.Sending != nil,
		Done:             stat.Done,
	})

	percent := float64(stat.BytesTransferred) / float64(stat.BytesTotal) * 100
	fmt.Printf("\r%s: %3.0f%%", prop.Title(), percent)

	if stat.Done {
		fmt.Println("")
	}
	os.Stdout.Sync()
}
<|endoftext|>


<commit_before>package cmd

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/PuerkitoBio/goquery"
	"github.com/fatih/color"
	"github.com/spf13/cobra"
	"golang.org/x/text/encoding/japanese"
	"golang.org/x/text/transform"
)

var (
	errColor   func(string, ...interface{}) string = color.HiYellowString
	countColor func(string, ...interface{}) string = color.HiYellowString
)

type exitCode int

const (
	normal exitCode = iota
	abnormal
)

func (c exitCode) Exit() {
	os.Exit(int(c))
}

func newRootCmd(newOut, newErr io.Writer, args []string) *cobra.Command {

	cmd := &cobra.Command{
		Use:           "kuroneko [flags] 伝票番号",
		Short:         "ヤマト運輸のステータス取得",
		SilenceErrors: true,
		SilenceUsage:  true,

		Args: func(cmd *cobra.Command, args []string) error {
			if len(args) > 1 {
				count := strconv.Itoa(len(args))
				return fmt.Errorf("accepts at most 1 arg(s), received %s", errColor(count))
			}

			if len(args) == 0 {
				return errors.New(errColor("伝票番号を入力してください"))
			}

			return nil
		},

		PreRunE: func(cmd *cobra.Command, args []string) error {
			flagCount := cmd.Flags().NFlag()
			if flagCount > 1 {
				count := strconv.Itoa(flagCount)
				return fmt.Errorf("accepte at most 1 flag(s), received %s", errColor(count))
			}

			serial, err := cmd.Flags().GetInt("serial")
			if err != nil {
				return err
			}

			if serial < 1 || serial > 10 {
				return errors.New(errColor("連番で取得できるのは 1~10件 までです"))
			}

			return nil
		},

		RunE: func(cmd *cobra.Command, args []string) error {
			trackingNumber := args[0]
			tracker := newTracker(cmd)
			return tracker.track(trackingNumber)
		},
	}

	cmd.Flags().IntP("serial", "s", 1, "連番取得(10件まで)")
	cmd.SetArgs(args)
	cmd.SetOut(newOut)
	cmd.SetErr(newErr)

	return cmd
}

func Execute(newOut, newErr io.Writer, args []string) exitCode {
	cmd := newRootCmd(newOut, newErr, args)
	if err := cmd.Execute(); err != nil {
		cmd.PrintErrf("Error: %+v\n", err)
		return abnormal
	}
	return normal
}

func init() {}

func makeSpace(count int) string {
	// 注:全角スペース
	s := "　"
	return strings.Repeat(s, count)
}

type tracker interface {
	track(s string) error
}

func newTracker(cmd *cobra.Command) tracker {
	flagCount := cmd.Flags().NFlag()
	switch flagCount {
	case 0:
		return &trackShipmentsOne{
			cmd: cmd,
		}
	default:
		// PreRunEでエラーチェック済み
		serial, _ := cmd.Flags().GetInt("serial")
		return &trackShipmentsMultiple{
			cmd:    cmd,
			serial: serial,
		}
	}
}

type trackShipmentsOne struct {
	cmd *cobra.Command
}

func (t *trackShipmentsOne) track(s string) error {
	values := url.Values{}
	values.Add("number00", "1")
	values.Add("number01", s)

	contactUrl := "http://toi.kuronekoyamato.co.jp/cgi-bin/tneko"
	resp, err := http.PostForm(contactUrl, values)
	if err != nil {
		return err
	}

	defer resp.Body.Close()

	utfBody := transform.NewReader(resp.Body, japanese.ShiftJIS.NewDecoder())

	doc, err := goquery.NewDocumentFromReader(utfBody)
	if err != nil {
		return err
	}

	w := t.cmd.OutOrStdout()
	doc.Find(".saisin td").Each(func(_ int, args *goquery.Selection) {
		if args.HasClass("bold") || args.HasClass("font14") {
			text := args.Text()
			fmt.Fprintf(w, " %s\n", text)
		}
	})

	fmt.Fprintf(w, "\n")

	doc.Find(".meisai tr").Each(func(i int, args *goquery.Selection) {
		if i != 0 {
			information := args.Find("td").Map(func(_ int, s *goquery.Selection) string {
				text := s.Text()
				return text
			})
			detailInfo := information[1:6]
			statusLength := utf8.RuneCountInString(detailInfo[0])
			whitespace := 15 - statusLength
			space := makeSpace(whitespace)
			status := detailInfo[0] + space
			branchLength := utf8.RuneCountInString(detailInfo[3])
			whitespace = 20 - branchLength
			space = makeSpace(whitespace)
			branch := detailInfo[3] + space
			date, times, code := detailInfo[1], detailInfo[2], detailInfo[4]
			if date == "" {
				date = "     "
			}
			if times == "" {
				times = "     "
			}
			fmt.Fprintf(w, " %s| %s | %s | %s| %s |\n", status, date, times, branch, code)
		}
	})

	underLine := strings.Repeat("-", 99)
	fmt.Fprintf(w, "%s\n", underLine)

	return nil

}

type trackShipmentsMultiple struct {
	cmd    *cobra.Command
	serial int
}

func (t *trackShipmentsMultiple) track(s string) error {
	trackingNumber := removeHyphen(s)
	if !isInt(trackingNumber) {
		return fmt.Errorf("%s", errColor("不正な数値です"))
	}

	if !is12or11Digits(trackingNumber) {
		return fmt.Errorf("%s", errColor("12 or 11桁の伝票番号を入力してください"))
	}

	if !isCorrectNumber(trackingNumber) {
		return fmt.Errorf("%s", errColor("伝票番号に誤りがあります"))
	}

	ctx, cancel := context.WithCancel(context.Background())
	ch := sevenCheckCalculate(ctx, trackingNumber[:len(trackingNumber)-1])
	values := url.Values{}
	values.Add("number00", "1")

	var i int
	for i = 0; i < t.serial; i++ {
		querykey := fmt.Sprintf("number%02d", i+1)
		values.Add(querykey, <-ch)
	}
	cancel()

	contactUrl := "http://toi.kuronekoyamato.co.jp/cgi-bin/tneko"
	resp, err := http.PostForm(contactUrl, values)
	if err != nil {
		return err
	}

	defer resp.Body.Close()

	utfBody := transform.NewReader(resp.Body, japanese.ShiftJIS.NewDecoder())

	doc, err := goquery.NewDocumentFromReader(utfBody)
	if err != nil {
		return err
	}

	w := t.cmd.OutOrStdout()
	doc.Find("center").Each(func(_ int, s *goquery.Selection) {
		hasDetail := false
		s.Find(".saisin td").Each(func(_ int, args *goquery.Selection) {
			if args.HasClass("number") {
				hasDetail = true
				subject := args.Text()
				fmt.Fprintf(w, " %s\n", countColor(subject))
			}

			if args.HasClass("bold") || args.HasClass("font14") {
				text := args.Text()
				fmt.Fprintf(w, " %s\n", text)
			}
		})

		if hasDetail {
			fmt.Fprintf(w, "\n")
		}

		s.Find(".meisai tr").Each(func(i int, args *goquery.Selection) {
			if i != 0 {
				information := args.Find("td").Map(func(_ int, s *goquery.Selection) string {
					text := s.Text()
					return text
				})
				detailInfo := information[1:6]
				statusLength := utf8.RuneCountInString(detailInfo[0])
				whitespace := 15 - statusLength
				space := makeSpace(whitespace)
				status := detailInfo[0] + space
				branchLength := utf8.RuneCountInString(detailInfo[3])
				whitespace = 20 - branchLength
				space = makeSpace(whitespace)
				branch := detailInfo[3] + space
				date, times, code := detailInfo[1], detailInfo[2], detailInfo[4]
				if date == "" {
					date = "     "
				}
				if times == "" {
					times = "     "
				}
				fmt.Fprintf(w, " %s| %s | %s | %s| %s |\n", status, date, times, branch, code)
			}
		})

		if hasDetail {
			underLine := strings.Repeat("-", 99)
			fmt.Fprintf(w, "%s\n", underLine)
		}
	})

	return nil
}

func removeHyphen(s string) string {
	if strings.Contains(s, "-") {
		removed := strings.Replace(s, "-", "", -1)
		return removed
	}
	return s
}

func sevenCheckCalculate(ctx context.Context, n string) <-chan string {
	ch := make(chan string)
	const coef = 7
	var format = "%012s"
	if len(n) == 10 {
		format = "%011s"
	}
	go func() {
		sign, _ := strconv.ParseInt(n, 10, 64)
	LOOP:
		for {
			select {
			case <-ctx.Done():
				break LOOP
			default:
				digit := sign % coef
				digitStr := strconv.FormatInt(digit, 10)
				trackingNumber := strconv.FormatInt(sign, 10) + digitStr
				zeroPaddingNumber := fmt.Sprintf(format, trackingNumber)
				ch <- zeroPaddingNumber
				sign++
			}
		}
		close(ch)
	}()
	return ch
}

func isCorrectNumber(s string) bool {
	const coef = 7
	lastDigits := s[len(s)-1:]
	otherDigits := s[:len(s)-1]
	sign, _ := strconv.ParseInt(otherDigits, 10, 64)
	digit := sign % coef
	return lastDigits == fmt.Sprint(digit)
}

func isInt(s string) bool {
	for _, c := range s {
		if !unicode.IsDigit(c) {
			return false
		}
	}
	return true
}

func is12or11Digits(s string) bool {
	if len(s) == 12 || len(s) == 11 {
		return true
	}
	return false
}
<commit_msg>refactor<commit_after>package cmd

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"

	"github.com/PuerkitoBio/goquery"
	"github.com/fatih/color"
	"github.com/spf13/cobra"
	"golang.org/x/text/encoding/japanese"
	"golang.org/x/text/transform"
)

var (
	errColor   func(string, ...interface{}) string = color.HiYellowString
	countColor func(string, ...interface{}) string = color.HiYellowString
)

type exitCode int

const (
	normal exitCode = iota
	abnormal
)

func (c exitCode) Exit() {
	os.Exit(int(c))
}

func newRootCmd(newOut, newErr io.Writer, args []string) *cobra.Command {

	cmd := &cobra.Command{
		Use:           "kuroneko [flags] 伝票番号",
		Short:         "ヤマト運輸のステータス取得",
		SilenceErrors: true,
		SilenceUsage:  true,

		Args: func(cmd *cobra.Command, args []string) error {
			if len(args) > 1 {
				count := strconv.Itoa(len(args))
				return fmt.Errorf("accepts at most 1 arg(s), received %s", errColor(count))
			}

			if len(args) == 0 {
				return errors.New(errColor("伝票番号を入力してください"))
			}

			return nil
		},

		PreRunE: func(cmd *cobra.Command, args []string) error {
			flagCount := cmd.Flags().NFlag()
			if flagCount > 1 {
				count := strconv.Itoa(flagCount)
				return fmt.Errorf("accepte at most 1 flag(s), received %s", errColor(count))
			}

			serial, err := cmd.Flags().GetInt("serial")
			if err != nil {
				return err
			}

			if serial < 1 || serial > 10 {
				return errors.New(errColor("連番で取得できるのは 1~10件 までです"))
			}

			return nil
		},

		RunE: func(cmd *cobra.Command, args []string) error {
			trackingNumber := args[0]
			tracker := newTracker(cmd)
			return tracker.track(trackingNumber)
		},
	}

	cmd.Flags().IntP("serial", "s", 1, "連番取得(10件まで)")
	cmd.SetArgs(args)
	cmd.SetOut(newOut)
	cmd.SetErr(newErr)

	return cmd
}

func Execute(newOut, newErr io.Writer, args []string) exitCode {
	cmd := newRootCmd(newOut, newErr, args)
	if err := cmd.Execute(); err != nil {
		cmd.PrintErrf("Error: %+v\n", err)
		return abnormal
	}
	return normal
}

func init() {}

func makeSpace(count int) string {
	// 注:全角スペース
	s := "　"
	return strings.Repeat(s, count)
}

type tracker interface {
	track(s string) error
}

func newTracker(cmd *cobra.Command) tracker {
	flagCount := cmd.Flags().NFlag()
	switch flagCount {
	case 0:
		return &trackShipmentsOne{
			cmd: cmd,
		}
	default:
		// PreRunEでエラーチェック済み
		serial, _ := cmd.Flags().GetInt("serial")
		return &trackShipmentsMultiple{
			cmd:    cmd,
			serial: serial,
		}
	}
}

type trackShipmentsOne struct {
	cmd *cobra.Command
}

func (t *trackShipmentsOne) track(s string) error {
	values := url.Values{}
	values.Add("number00", "1")
	values.Add("number01", s)

	contactUrl := "http://toi.kuronekoyamato.co.jp/cgi-bin/tneko"
	resp, err := http.PostForm(contactUrl, values)
	if err != nil {
		return err
	}

	defer resp.Body.Close()

	utfBody := transform.NewReader(resp.Body, japanese.ShiftJIS.NewDecoder())

	doc, err := goquery.NewDocumentFromReader(utfBody)
	if err != nil {
		return err
	}

	w := t.cmd.OutOrStdout()
	doc.Find(".saisin td").Each(func(_ int, args *goquery.Selection) {
		if args.HasClass("bold") || args.HasClass("font14") {
			text := args.Text()
			fmt.Fprintf(w, " %s\n", text)
		}
	})

	fmt.Fprintf(w, "\n")

	doc.Find(".meisai tr").Each(func(i int, args *goquery.Selection) {
		if i != 0 {
			information := args.Find("td").Map(func(_ int, s *goquery.Selection) string {
				text := s.Text()
				return text
			})
			detailInfo := information[1:6]
			statusLength := utf8.RuneCountInString(detailInfo[0])
			whitespace := 15 - statusLength
			space := makeSpace(whitespace)
			status := detailInfo[0] + space
			branchLength := utf8.RuneCountInString(detailInfo[3])
			whitespace = 20 - branchLength
			space = makeSpace(whitespace)
			branch := detailInfo[3] + space
			date, times, code := detailInfo[1], detailInfo[2], detailInfo[4]
			if date == "" {
				date = "     "
			}
			if times == "" {
				times = "     "
			}
			fmt.Fprintf(w, " %s| %s | %s | %s| %s |\n", status, date, times, branch, code)
		}
	})

	underLine := strings.Repeat("-", 99)
	fmt.Fprintf(w, "%s\n", underLine)

	return nil

}

type trackShipmentsMultiple struct {
	cmd    *cobra.Command
	serial int
}

func (t *trackShipmentsMultiple) track(s string) error {
	trackingNumber := removeHyphen(s)
	if !isInt(trackingNumber) {
		return errors.New(errColor("不正な数値です"))
	}

	if !is12or11Digits(trackingNumber) {
		return errors.New(errColor("12 or 11桁の伝票番号を入力してください"))
	}

	if !isCorrectNumber(trackingNumber) {
		return errors.New(errColor("伝票番号に誤りがあります"))
	}

	ctx, cancel := context.WithCancel(context.Background())
	ch := sevenCheckCalculate(ctx, trackingNumber[:len(trackingNumber)-1])
	values := url.Values{}
	values.Add("number00", "1")

	var i int
	for i = 0; i < t.serial; i++ {
		querykey := fmt.Sprintf("number%02d", i+1)
		values.Add(querykey, <-ch)
	}
	cancel()

	contactUrl := "http://toi.kuronekoyamato.co.jp/cgi-bin/tneko"
	resp, err := http.PostForm(contactUrl, values)
	if err != nil {
		return err
	}

	defer resp.Body.Close()

	utfBody := transform.NewReader(resp.Body, japanese.ShiftJIS.NewDecoder())

	doc, err := goquery.NewDocumentFromReader(utfBody)
	if err != nil {
		return err
	}

	w := t.cmd.OutOrStdout()
	doc.Find("center").Each(func(_ int, s *goquery.Selection) {
		hasDetail := false
		s.Find(".saisin td").Each(func(_ int, args *goquery.Selection) {
			if args.HasClass("number") {
				hasDetail = true
				subject := args.Text()
				fmt.Fprintf(w, " %s\n", countColor(subject))
			}

			if args.HasClass("bold") || args.HasClass("font14") {
				text := args.Text()
				fmt.Fprintf(w, " %s\n", text)
			}
		})

		if hasDetail {
			fmt.Fprintf(w, "\n")
		}

		s.Find(".meisai tr").Each(func(i int, args *goquery.Selection) {
			if i != 0 {
				information := args.Find("td").Map(func(_ int, s *goquery.Selection) string {
					text := s.Text()
					return text
				})
				detailInfo := information[1:6]
				statusLength := utf8.RuneCountInString(detailInfo[0])
				whitespace := 15 - statusLength
				space := makeSpace(whitespace)
				status := detailInfo[0] + space
				branchLength := utf8.RuneCountInString(detailInfo[3])
				whitespace = 20 - branchLength
				space = makeSpace(whitespace)
				branch := detailInfo[3] + space
				date, times, code := detailInfo[1], detailInfo[2], detailInfo[4]
				if date == "" {
					date = "     "
				}
				if times == "" {
					times = "     "
				}
				fmt.Fprintf(w, " %s| %s | %s | %s| %s |\n", status, date, times, branch, code)
			}
		})

		if hasDetail {
			underLine := strings.Repeat("-", 99)
			fmt.Fprintf(w, "%s\n", underLine)
		}
	})

	return nil
}

func removeHyphen(s string) string {
	if strings.Contains(s, "-") {
		removed := strings.Replace(s, "-", "", -1)
		return removed
	}
	return s
}

func sevenCheckCalculate(ctx context.Context, n string) <-chan string {
	ch := make(chan string)
	const coef = 7
	var format = "%012s"
	if len(n) == 10 {
		format = "%011s"
	}
	go func() {
		sign, _ := strconv.ParseInt(n, 10, 64)
	LOOP:
		for {
			select {
			case <-ctx.Done():
				break LOOP
			default:
				digit := sign % coef
				digitStr := strconv.FormatInt(digit, 10)
				trackingNumber := strconv.FormatInt(sign, 10) + digitStr
				zeroPaddingNumber := fmt.Sprintf(format, trackingNumber)
				ch <- zeroPaddingNumber
				sign++
			}
		}
		close(ch)
	}()
	return ch
}

func isCorrectNumber(s string) bool {
	const coef = 7
	lastDigits := s[len(s)-1:]
	otherDigits := s[:len(s)-1]
	sign, _ := strconv.ParseInt(otherDigits, 10, 64)
	digit := sign % coef
	return lastDigits == fmt.Sprint(digit)
}

func isInt(s string) bool {
	for _, c := range s {
		if !unicode.IsDigit(c) {
			return false
		}
	}
	return true
}

func is12or11Digits(s string) bool {
	if len(s) == 12 || len(s) == 11 {
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package model

import (
	"net/http"
	"testing"
	"time"

	"github.com/labstack/echo"
	sd "github.com/labstack/echo/engine/standard"
)

const (
	etcdAddr   = "http://192.168.70.13:2379"
	etcdPrefix = "/gateway2"
)

var (
	serverAddr    = "127.0.0.1:12345"
	apiURL        = "/api/test"
	apiMethod     = "GET"
	checkDuration = 3
	checkTimeout  = 2
	clusterName   = "app"
	lbName        = "ROUNDROBIN"
	sleep         = false
)

var rt *RouteTable

func createRouteTable(t *testing.T) {
	store, err := NewEtcdStore([]string{etcdAddr}, etcdPrefix)

	if nil != err {
		t.Fatalf("create etcd store err.addr:<%s>", err)
	}

	store.Clean()

	rt = NewRouteTable(store)
	time.Sleep(time.Second * 1)
}

func createLocalServer() {
	e := echo.New()

	e.Get("/check", func() echo.HandlerFunc {
		return func(c echo.Context) error {
			if sleep {
				time.Sleep(time.Second * time.Duration(checkTimeout+1))
			}

			return c.String(http.StatusOK, "OK")
		}
	}())

	e.Run(sd.New(serverAddr))
}

func waitNotify() {
	time.Sleep(time.Second * 1)
}

func TestCreateRouteTable(t *testing.T) {
	createRouteTable(t)
}

func TestEtcdWatchNewServer(t *testing.T) {
	go createLocalServer()

	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout,
		MaxQPS:          1500,
		HalfToOpen:      10,
		HalfTrafficRate: 10,
		CloseCount:      100,
	}

	err := rt.store.SaveServer(server)

	if nil != err {
		t.Error("add server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 1 {
		t.Errorf("expect:<1>, acture:<%d>", len(rt.svrs))
		return
	}

	if rt.svrs[serverAddr].lock == nil {
		t.Error("server init error.")
		return
	}
}

func TestServerCheckOk(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration))

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check ok err.expect:<UP>, acture:<%v>", Down)
	}
}

func TestServerCheckTimeout(t *testing.T) {
	defer func() {
		sleep = false
	}()

	sleep = true
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Up {
		t.Errorf("status check timeout err.expect:<DOWN>, acture:<%v>", Up)
		return
	}
}

func TestServerCheckTimeoutRecovery(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check timeout recovery err.expect:<UP>, acture:<%v>", Up)
		return
	}
}

func TestEtcdWatchNewCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.SaveCluster(cluster)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.clusters))
}

func TestEtcdWatchNewBind(t *testing.T) {
	bind := &Bind{
		ClusterName: clusterName,
		ServerAddr:  serverAddr,
	}

	err := rt.store.SaveBind(bind)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.mapping) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>. %+v", len(rt.mapping), rt.mapping)
}

func TestEtcdWatchNewAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	err := rt.store.SaveAPI(&API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n},
	})

	if nil != err {
		t.Error("add api err.")
		return
	}

	waitNotify()

	if len(rt.apis) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.apis))
}

func TestEtcdWatchUpdateServer(t *testing.T) {
	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout * 2,
		MaxQPS:          3000,
		HalfToOpen:      100,
		HalfTrafficRate: 30,
		CloseCount:      200,
	}

	err := rt.store.UpdateServer(server)

	if nil != err {
		t.Error("update server err.")
		return
	}

	waitNotify()

	svr := rt.svrs[serverAddr]

	if svr.MaxQPS != server.MaxQPS {
		t.Errorf("MaxQPS expect:<%d>, acture:<%d>. ", server.MaxQPS, svr.MaxQPS)
		return
	}

	if svr.HalfToOpen != server.HalfToOpen {
		t.Errorf("HalfToOpen expect:<%d>, acture:<%d>. ", server.HalfToOpen, svr.HalfToOpen)
		return
	}

	if svr.HalfTrafficRate != server.HalfTrafficRate {
		t.Errorf("HalfTrafficRate expect:<%d>, acture:<%d>. ", server.HalfTrafficRate, svr.HalfTrafficRate)
		return
	}

	if svr.CloseCount != server.CloseCount {
		t.Errorf("CloseCount expect:<%d>, acture:<%d>. ", server.CloseCount, svr.CloseCount)
		return
	}

	if svr.CheckTimeout == server.CheckTimeout {
		t.Errorf("CheckTimeout expect:<%d>, acture:<%d>. ", svr.CheckTimeout, server.CheckTimeout)
		return
	}
}

func TestEtcdWatchUpdateCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.UpdateCluster(cluster)

	if nil != err {
		t.Error("update cluster err.")
		return
	}

	waitNotify()

	existCluster := rt.clusters[clusterName]

	if existCluster.LbName != cluster.LbName {
		t.Errorf("LbName expect:<%s>, acture:<%s>. ", cluster.LbName, existCluster.LbName)
		return
	}
}

func TestEtcdWatchUpdateAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	n2 := &Node{
		AttrName:    "tes2t",
		ClusterName: clusterName,
	}

	api := &API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n, n2},
	}

	err := rt.store.UpdateAPI(api)

	if nil != err {
		t.Error("update api err.")
		return
	}

	waitNotify()

	existAPI, _ := rt.apis[getAPIKey(api.URL, api.Method)]

	if len(existAPI.Nodes) != len(api.Nodes) {
		t.Errorf("Nodes expect:<%s>, acture:<%d>. ", len(existAPI.Nodes), len(api.Nodes))
		return
	}
}

func TestEtcdWatchDeleteCluster(t *testing.T) {
	err := rt.store.DeleteCluster(clusterName)

	if nil != err {
		t.Error("delete cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) != 0 {
		t.Errorf("clusters expect:<0>, acture:<%d>", len(rt.clusters))
		return
	}

	banded, _ := rt.mapping[serverAddr]

	if len(banded) != 0 {
		t.Errorf("banded expect:<0>, acture:<%d>", len(banded))
		return
	}
}

func TestEtcdWatchDeleteServer(t *testing.T) {
	err := rt.store.DeleteServer(serverAddr)

	if nil != err {
		t.Error("delete server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 0 {
		t.Errorf("svrs expect:<0>, acture:<%d>", len(rt.svrs))
		return
	}

	if len(rt.mapping) != 0 {
		t.Errorf("mapping expect:<0>, acture:<%d>", len(rt.mapping))
		return
	}
}

func TestEtcdWatchDeleteAPI(t *testing.T) {
	err := rt.store.DeleteAPI(apiURL, apiMethod)

	if nil != err {
		t.Error("delete api err.")
		return
	}

	waitNotify()

	if len(rt.apis) != 0 {
		t.Errorf("apis expect:<0>, acture:<%d>", len(rt.apis))
		return
	}
}

func TestEtcdWatchNewRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 100; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	waitNotify()

	if len(rt.routings) == 1 {
		delete(rt.routings, r.ID)
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.routings))
}

func TestEtcdWatchDeleteRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 3; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	time.Sleep(time.Second * 30)

	if len(rt.routings) == 0 {
		return
	}

	t.Errorf("expect:<0>, acture:<%d>", len(rt.routings))
}
<commit_msg>fix warn about go_vet<commit_after>package model

import (
	"net/http"
	"testing"
	"time"

	"github.com/labstack/echo"
	sd "github.com/labstack/echo/engine/standard"
)

const (
	etcdAddr   = "http://192.168.70.13:2379"
	etcdPrefix = "/gateway2"
)

var (
	serverAddr    = "127.0.0.1:12345"
	apiURL        = "/api/test"
	apiMethod     = "GET"
	checkDuration = 3
	checkTimeout  = 2
	clusterName   = "app"
	lbName        = "ROUNDROBIN"
	sleep         = false
)

var rt *RouteTable

func createRouteTable(t *testing.T) {
	store, err := NewEtcdStore([]string{etcdAddr}, etcdPrefix)

	if nil != err {
		t.Fatalf("create etcd store err.addr:<%s>", err)
	}

	store.Clean()

	rt = NewRouteTable(store)
	time.Sleep(time.Second * 1)
}

func createLocalServer() {
	e := echo.New()

	e.Get("/check", func() echo.HandlerFunc {
		return func(c echo.Context) error {
			if sleep {
				time.Sleep(time.Second * time.Duration(checkTimeout+1))
			}

			return c.String(http.StatusOK, "OK")
		}
	}())

	e.Run(sd.New(serverAddr))
}

func waitNotify() {
	time.Sleep(time.Second * 1)
}

func TestCreateRouteTable(t *testing.T) {
	createRouteTable(t)
}

func TestEtcdWatchNewServer(t *testing.T) {
	go createLocalServer()

	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout,
		MaxQPS:          1500,
		HalfToOpen:      10,
		HalfTrafficRate: 10,
		CloseCount:      100,
	}

	err := rt.store.SaveServer(server)

	if nil != err {
		t.Error("add server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 1 {
		t.Errorf("expect:<1>, acture:<%d>", len(rt.svrs))
		return
	}

	if rt.svrs[serverAddr].lock == nil {
		t.Error("server init error.")
		return
	}
}

func TestServerCheckOk(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration))

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check ok err.expect:<UP>, acture:<%v>", Down)
	}
}

func TestServerCheckTimeout(t *testing.T) {
	defer func() {
		sleep = false
	}()

	sleep = true
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Up {
		t.Errorf("status check timeout err.expect:<DOWN>, acture:<%v>", Up)
		return
	}
}

func TestServerCheckTimeoutRecovery(t *testing.T) {
	time.Sleep(time.Second * time.Duration(checkDuration*2+1)) // 等待两个周期

	if rt.svrs[serverAddr].Status == Down {
		t.Errorf("status check timeout recovery err.expect:<UP>, acture:<%v>", Up)
		return
	}
}

func TestEtcdWatchNewCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.SaveCluster(cluster)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.clusters))
}

func TestEtcdWatchNewBind(t *testing.T) {
	bind := &Bind{
		ClusterName: clusterName,
		ServerAddr:  serverAddr,
	}

	err := rt.store.SaveBind(bind)

	if nil != err {
		t.Error("add cluster err.")
		return
	}

	waitNotify()

	if len(rt.mapping) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>. %+v", len(rt.mapping), rt.mapping)
}

func TestEtcdWatchNewAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	err := rt.store.SaveAPI(&API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n},
	})

	if nil != err {
		t.Error("add api err.")
		return
	}

	waitNotify()

	if len(rt.apis) == 1 {
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.apis))
}

func TestEtcdWatchUpdateServer(t *testing.T) {
	server := &Server{
		Schema:          "http",
		Addr:            serverAddr,
		CheckPath:       "/check",
		CheckDuration:   checkDuration,
		CheckTimeout:    checkTimeout * 2,
		MaxQPS:          3000,
		HalfToOpen:      100,
		HalfTrafficRate: 30,
		CloseCount:      200,
	}

	err := rt.store.UpdateServer(server)

	if nil != err {
		t.Error("update server err.")
		return
	}

	waitNotify()

	svr := rt.svrs[serverAddr]

	if svr.MaxQPS != server.MaxQPS {
		t.Errorf("MaxQPS expect:<%d>, acture:<%d>. ", server.MaxQPS, svr.MaxQPS)
		return
	}

	if svr.HalfToOpen != server.HalfToOpen {
		t.Errorf("HalfToOpen expect:<%d>, acture:<%d>. ", server.HalfToOpen, svr.HalfToOpen)
		return
	}

	if svr.HalfTrafficRate != server.HalfTrafficRate {
		t.Errorf("HalfTrafficRate expect:<%d>, acture:<%d>. ", server.HalfTrafficRate, svr.HalfTrafficRate)
		return
	}

	if svr.CloseCount != server.CloseCount {
		t.Errorf("CloseCount expect:<%d>, acture:<%d>. ", server.CloseCount, svr.CloseCount)
		return
	}

	if svr.CheckTimeout == server.CheckTimeout {
		t.Errorf("CheckTimeout expect:<%d>, acture:<%d>. ", svr.CheckTimeout, server.CheckTimeout)
		return
	}
}

func TestEtcdWatchUpdateCluster(t *testing.T) {
	cluster := &Cluster{
		Name:   clusterName,
		LbName: lbName,
	}

	err := rt.store.UpdateCluster(cluster)

	if nil != err {
		t.Error("update cluster err.")
		return
	}

	waitNotify()

	existCluster := rt.clusters[clusterName]

	if existCluster.LbName != cluster.LbName {
		t.Errorf("LbName expect:<%s>, acture:<%s>. ", cluster.LbName, existCluster.LbName)
		return
	}
}

func TestEtcdWatchUpdateAPI(t *testing.T) {
	n := &Node{
		AttrName:    "test",
		ClusterName: clusterName,
	}

	n2 := &Node{
		AttrName:    "tes2t",
		ClusterName: clusterName,
	}

	api := &API{
		URL:    apiURL,
		Method: apiMethod,
		Nodes:  []*Node{n, n2},
	}

	err := rt.store.UpdateAPI(api)

	if nil != err {
		t.Error("update api err.")
		return
	}

	waitNotify()

	existAPI, _ := rt.apis[getAPIKey(api.URL, api.Method)]

	if len(existAPI.Nodes) != len(api.Nodes) {
		t.Errorf("Nodes expect:<%d>, acture:<%d>. ", len(existAPI.Nodes), len(api.Nodes))
		return
	}
}

func TestEtcdWatchDeleteCluster(t *testing.T) {
	err := rt.store.DeleteCluster(clusterName)

	if nil != err {
		t.Error("delete cluster err.")
		return
	}

	waitNotify()

	if len(rt.clusters) != 0 {
		t.Errorf("clusters expect:<0>, acture:<%d>", len(rt.clusters))
		return
	}

	banded, _ := rt.mapping[serverAddr]

	if len(banded) != 0 {
		t.Errorf("banded expect:<0>, acture:<%d>", len(banded))
		return
	}
}

func TestEtcdWatchDeleteServer(t *testing.T) {
	err := rt.store.DeleteServer(serverAddr)

	if nil != err {
		t.Error("delete server err.")
		return
	}

	waitNotify()

	if len(rt.svrs) != 0 {
		t.Errorf("svrs expect:<0>, acture:<%d>", len(rt.svrs))
		return
	}

	if len(rt.mapping) != 0 {
		t.Errorf("mapping expect:<0>, acture:<%d>", len(rt.mapping))
		return
	}
}

func TestEtcdWatchDeleteAPI(t *testing.T) {
	err := rt.store.DeleteAPI(apiURL, apiMethod)

	if nil != err {
		t.Error("delete api err.")
		return
	}

	waitNotify()

	if len(rt.apis) != 0 {
		t.Errorf("apis expect:<0>, acture:<%d>", len(rt.apis))
		return
	}
}

func TestEtcdWatchNewRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 100; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	waitNotify()

	if len(rt.routings) == 1 {
		delete(rt.routings, r.ID)
		return
	}

	t.Errorf("expect:<1>, acture:<%d>", len(rt.routings))
}

func TestEtcdWatchDeleteRouting(t *testing.T) {
	r, err := NewRouting(`desc = "test"; deadline = 3; rule = ["$query_abc == 10", "$query_123 == 20"];`, clusterName, "")

	if nil != err {
		t.Error("add routing err.")
		return
	}

	err = rt.store.SaveRouting(r)

	if nil != err {
		t.Error("add routing err.")
		return
	}

	time.Sleep(time.Second * 30)

	if len(rt.routings) == 0 {
		return
	}

	t.Errorf("expect:<0>, acture:<%d>", len(rt.routings))
}
<|endoftext|>


<commit_before>// This file contains the MinQuery interface and its implementation.

package minquery

import (
	"errors"

	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"
)

// DefaultCursorCodec is the default CursorCodec value that is used if none
// is specified. The default implementation produces web-safe cursor strings.
var DefaultCursorCodec cursorCodec

// MinQuery is an mgo-like Query that supports cursors to continue listing documents
// where we left off. If a cursor is set, it specifies the last index entry
// that was already returned, and result documents will be listed after this.
type MinQuery interface {
	// Sort asks the database to order returned documents according to
	// the provided field names.
	Sort(fields ...string) MinQuery

	// Select enables selecting which fields should be retrieved for
	// the results found.
	Select(selector interface{}) MinQuery

	// Limit restricts the maximum number of documents retrieved to n,
	// and also changes the batch size to the same value.
	Limit(n int) MinQuery

	// Cursor sets the cursor, which specifies the last index entry
	// that was already returned, and result documents will be listed after this.
	// Parsing a cursor may fail which is not returned. If an invalid cursor
	// is specified, All() will fail and return the error.
	Cursor(c string) MinQuery

	// CursorCoded sets the CursorCodec to be used to parse and to create cursors.
	// This gives you the possibility to implement your own logic to create cursors,
	// including encryption should you need it.
	CursorCodec(cc CursorCodec) MinQuery

	// All retrieves all documents from the result set into the provided slice.
	// cursorFields lists the fields (in order) to be used to generate
	// the returned cursor.
	All(result interface{}, cursorFields ...string) (cursor string, err error)
}

// errTestValue is the error value returned for testing purposes.
var errTestValue = errors.New("Intentional testing error")

// minQuery is the MinQuery implementation.
type minQuery struct {
	// db is the mgo Database to use
	db *mgo.Database

	// Name of the collection
	coll string

	// filter document (query)
	filter interface{}

	// sort document
	sort bson.D

	// projection document (to retrieve only selected fields)
	projection interface{}

	// limit is the max number of results
	limit int

	// Cursor, need to store and supply it if query returns no results
	cursor string

	// cursorCodec to be used to parse and to create cursors
	cursorCodec CursorCodec

	// cursorErr contains an error if an invalid cursor is supplied
	cursorErr error

	// min specifies the last index entry
	min bson.D

	// testError is a helper field to aid testing errors to reach 100% coverage.
	// May only be changed from tests! Zero value means normal operation.
	testError bool
}

// New returns a new MinQuery.
func New(db *mgo.Database, coll string, query interface{}) MinQuery {
	return &minQuery{
		db:          db,
		coll:        coll,
		filter:      query,
		cursorCodec: DefaultCursorCodec,
	}
}

// Sort implements MinQuery.Sort().
func (mq *minQuery) Sort(fields ...string) MinQuery {
	mq.sort = make(bson.D, 0, len(fields))
	for _, field := range fields {
		if field == "" {
			continue
		}
		n := 1
		if field[0] == '+' {
			field = field[1:]
		} else if field[0] == '-' {
			n, field = -1, field[1:]
		}
		mq.sort = append(mq.sort, bson.DocElem{Name: field, Value: n})
	}
	return mq
}

// Select implements MinQuery.Select().
func (mq *minQuery) Select(selector interface{}) MinQuery {
	mq.projection = selector
	return mq
}

// Limit implements MinQuery.Limit().
func (mq *minQuery) Limit(n int) MinQuery {
	mq.limit = n
	return mq
}

// Cursor implements MinQuery.Cursor().
func (mq *minQuery) Cursor(c string) MinQuery {
	mq.cursor = c
	if c != "" {
		mq.min, mq.cursorErr = mq.cursorCodec.ParseCursor(c)
	} else {
		mq.min, mq.cursorErr = nil, nil
	}
	return mq
}

// CursorCodec implements MinQuery.CursorCodec().
func (mq *minQuery) CursorCodec(cc CursorCodec) MinQuery {
	mq.cursorCodec = cc
	return mq
}

// All implements MinQuery.All().
func (mq *minQuery) All(result interface{}, cursorFields ...string) (cursor string, err error) {
	if mq.cursorErr != nil {
		return "", mq.cursorErr
	}

	// Mongodb "find" reference:
	// https://docs.mongodb.com/manual/reference/command/find/

	cmd := bson.D{
		{Name: "find", Value: mq.coll},
		{Name: "limit", Value: mq.limit},
		{Name: "batchSize", Value: mq.limit},
		{Name: "singleBatch", Value: true},
	}
	if mq.filter != nil {
		cmd = append(cmd, bson.DocElem{Name: "filter", Value: mq.filter})
	}
	if mq.sort != nil {
		cmd = append(cmd, bson.DocElem{Name: "sort", Value: mq.sort})
	}
	if mq.projection != nil {
		cmd = append(cmd, bson.DocElem{Name: "projection", Value: mq.projection})
	}
	if mq.min != nil {
		// min is inclusive, skip the first (which is the previous last)
		cmd = append(cmd,
			bson.DocElem{Name: "skip", Value: 1},
			bson.DocElem{Name: "min", Value: mq.min},
		)
	}

	var res struct {
		OK       int `bson:"ok"`
		WaitedMS int `bson:"waitedMS"`
		Cursor   struct {
			ID         interface{} `bson:"id"`
			NS         string      `bson:"ns"`
			FirstBatch []bson.Raw  `bson:"firstBatch"`
		} `bson:"cursor"`
	}

	if err = mq.db.Run(cmd, &res); err != nil {
		return
	}

	firstBatch := res.Cursor.FirstBatch
	if len(firstBatch) > 0 {
		if len(cursorFields) > 0 {
			// create cursor from the last document
			var doc bson.M
			err = firstBatch[len(firstBatch)-1].Unmarshal(&doc)
			if mq.testError {
				err = errTestValue
			}
			if err != nil {
				return
			}
			cursorData := make(bson.D, len(cursorFields))
			for i, cf := range cursorFields {
				cursorData[i] = bson.DocElem{Name: cf, Value: doc[cf]}
			}
			cursor, err = mq.cursorCodec.CreateCursor(cursorData)
			if err != nil {
				return
			}
		}
	} else {
		// No more results. Use the same cursor that was used for the query.
		// It's possible that the last doc was returned previously, and there
		// are no more.
		cursor = mq.cursor
	}

	// Unmarshal results (FirstBatch) into the user-provided value:
	err = mq.db.C(mq.coll).NewIter(nil, firstBatch, 0, nil).All(result)
	return
}
<commit_msg>added hint<commit_after>// This file contains the MinQuery interface and its implementation.

package minquery

import (
	"errors"
	"fmt"

	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"
)

// DefaultCursorCodec is the default CursorCodec value that is used if none
// is specified. The default implementation produces web-safe cursor strings.
var DefaultCursorCodec cursorCodec

// MinQuery is an mgo-like Query that supports cursors to continue listing documents
// where we left off. If a cursor is set, it specifies the last index entry
// that was already returned, and result documents will be listed after this.
type MinQuery interface {
	// Sort asks the database to order returned documents according to
	// the provided field names.
	Sort(fields ...string) MinQuery

	// Select enables selecting which fields should be retrieved for
	// the results found.
	Select(selector interface{}) MinQuery

	// Limit restricts the maximum number of documents retrieved to n,
	// and also changes the batch size to the same value.
	Limit(n int) MinQuery

	// Cursor sets the cursor, which specifies the last index entry
	// that was already returned, and result documents will be listed after this.
	// Parsing a cursor may fail which is not returned. If an invalid cursor
	// is specified, All() will fail and return the error.
	Cursor(c string) MinQuery

	// CursorCoded sets the CursorCodec to be used to parse and to create cursors.
	// This gives you the possibility to implement your own logic to create cursors,
	// including encryption should you need it.
	CursorCodec(cc CursorCodec) MinQuery

	// All retrieves all documents from the result set into the provided slice.
	// cursorFields lists the fields (in order) to be used to generate
	// the returned cursor.
	All(result interface{}, cursorFields ...string) (cursor string, err error)
}

// errTestValue is the error value returned for testing purposes.
var errTestValue = errors.New("Intentional testing error")

// minQuery is the MinQuery implementation.
type minQuery struct {
	// db is the mgo Database to use
	db   *mgo.Database
	hint map[string]int
	// Name of the collection
	coll string

	// filter document (query)
	filter interface{}

	// sort document
	sort bson.D

	// projection document (to retrieve only selected fields)
	projection interface{}

	// limit is the max number of results
	limit int

	// Cursor, need to store and supply it if query returns no results
	cursor string

	// cursorCodec to be used to parse and to create cursors
	cursorCodec CursorCodec

	// cursorErr contains an error if an invalid cursor is supplied
	cursorErr error

	// min specifies the last index entry
	min bson.D

	// testError is a helper field to aid testing errors to reach 100% coverage.
	// May only be changed from tests! Zero value means normal operation.
	testError bool
}

// New returns a new MinQuery.
func New(db *mgo.Database, coll string, query interface{}, hint map[string]int) MinQuery {
	return &minQuery{
		db:          db,
		coll:        coll,
		filter:      query,
		hint:        hint,
		cursorCodec: DefaultCursorCodec,
	}
}

// Sort implements MinQuery.Sort().
func (mq *minQuery) Sort(fields ...string) MinQuery {
	mq.sort = make(bson.D, 0, len(fields))
	for _, field := range fields {
		if field == "" {
			continue
		}
		n := 1
		if field[0] == '+' {
			field = field[1:]
		} else if field[0] == '-' {
			n, field = -1, field[1:]
		}
		mq.sort = append(mq.sort, bson.DocElem{Name: field, Value: n})
	}
	return mq
}

// Select implements MinQuery.Select().
func (mq *minQuery) Select(selector interface{}) MinQuery {
	mq.projection = selector
	return mq
}

// Limit implements MinQuery.Limit().
func (mq *minQuery) Limit(n int) MinQuery {
	mq.limit = n
	return mq
}

// Cursor implements MinQuery.Cursor().
func (mq *minQuery) Cursor(c string) MinQuery {
	mq.cursor = c
	if c != "" {
		mq.min, mq.cursorErr = mq.cursorCodec.ParseCursor(c)
	} else {
		mq.min, mq.cursorErr = nil, nil
	}
	return mq
}

// CursorCodec implements MinQuery.CursorCodec().
func (mq *minQuery) CursorCodec(cc CursorCodec) MinQuery {
	mq.cursorCodec = cc
	return mq
}

// All implements MinQuery.All().
func (mq *minQuery) All(result interface{}, cursorFields ...string) (cursor string, err error) {
	if mq.cursorErr != nil {
		fmt.Println("Error in cusror ", err.Error())
		return "", mq.cursorErr
	}

	// Mongodb "find" reference:
	// https://docs.mongodb.com/manual/reference/command/find/

	cmd := bson.D{
		{Name: "find", Value: mq.coll},
		{Name: "limit", Value: mq.limit},
		{Name: "batchSize", Value: mq.limit},
		{Name: "singleBatch", Value: true},
	}
	if mq.filter != nil {
		cmd = append(cmd, bson.DocElem{Name: "filter", Value: mq.filter})
	}
	if mq.sort != nil {
		cmd = append(cmd, bson.DocElem{Name: "sort", Value: mq.sort})
	}
	if mq.projection != nil {
		cmd = append(cmd, bson.DocElem{Name: "projection", Value: mq.projection})
	}
	if mq.min != nil {
		// min is inclusive, skip the first (which is the previous last)
		cmd = append(cmd,
			bson.DocElem{Name: "skip", Value: 1},
			bson.DocElem{Name: "min", Value: mq.min},
			bson.DocElem{Name: "hint", Value: mq.hint},
		)
	}

	var res struct {
		OK       int `bson:"ok"`
		WaitedMS int `bson:"waitedMS"`
		Cursor   struct {
			ID         interface{} `bson:"id"`
			NS         string      `bson:"ns"`
			FirstBatch []bson.Raw  `bson:"firstBatch"`
		} `bson:"cursor"`
	}

	if err = mq.db.Run(cmd, &res); err != nil {
		return
	}

	firstBatch := res.Cursor.FirstBatch
	if len(firstBatch) > 0 {
		if len(cursorFields) > 0 {
			// create cursor from the last document
			var doc bson.M
			err = firstBatch[len(firstBatch)-1].Unmarshal(&doc)
			if mq.testError {
				err = errTestValue
			}
			if err != nil {
				return
			}
			cursorData := make(bson.D, len(cursorFields))
			for i, cf := range cursorFields {
				cursorData[i] = bson.DocElem{Name: cf, Value: doc[cf]}
			}
			cursor, err = mq.cursorCodec.CreateCursor(cursorData)
			if err != nil {
				return
			}
		}
	} else {
		// No more results. Use the same cursor that was used for the query.
		// It's possible that the last doc was returned previously, and there
		// are no more.
		cursor = mq.cursor
	}

	// Unmarshal results (FirstBatch) into the user-provided value:
	err = mq.db.C(mq.coll).NewIter(nil, firstBatch, 0, nil).All(result)

	return
}

<|endoftext|>