<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/aws/signer/v4"
)

type proxy struct {
	Scheme      string
	Host        string
	Region      string
	Service     string
	Verbose     bool
	Prettify    bool
	Credentials *credentials.Credentials
}

func getSigner(p *proxy) *v4.Signer {
	if p.Credentials == nil {
		p.Credentials = getCredentials()
	}
	return v4.NewSigner(p.Credentials)
}

func getCredentials() *credentials.Credentials {
	sess := session.Must(session.NewSession())
	credentials := sess.Config.Credentials
	log.Print("Generated fresh AWS Credentials object")
	return credentials
}

func copyHeaders(dst, src http.Header) {
	for k, vals := range src {
		for _, v := range vals {
			dst.Add(k, v)
		}
	}
}

func replaceBody(req *http.Request) []byte {
	if req.Body == nil {
		return []byte{}
	}
	payload, _ := ioutil.ReadAll(req.Body)
	req.Body = ioutil.NopCloser(bytes.NewReader(payload))
	return payload
}

func parseEndpoint(endpoint string, p *proxy) {
	link, err := url.Parse(endpoint)
	if err != nil {
		log.Fatalf("ERROR: Failed parsing endpoint: %s\n", endpoint)
	}

	// Only http/https are supported schemes
	scheme := func(x string) string {
		switch x {
		case "http", "https":
			return x
		}
		return "https"
	}
	link.Scheme = scheme(link.Scheme)

	// Unkown schemes sometimes result in empty host value
	if link.Host == "" {
		log.Fatalf("ERROR: Empty host information in submitted endpoint (%s)\n", endpoint)
	}

	// Extract region and service from link
	parts := strings.Split(link.Host, ".")
	var region, service string

	if len(parts) == 5 {
		region, service = parts[1], parts[2]
	} else {
		log.Fatalln("ERROR: Submitted endpoint is not a valid Amazon ElasticSearch Endpoint")
	}

	// Build proxy struct
	p.Scheme = link.Scheme
	p.Host = link.Host
	p.Region = region
	p.Service = service
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	requestStarted := time.Now()
	dump, err := httputil.DumpRequest(r, true)
	defer r.Body.Close()

	respondError := func(err error) {
		w.WriteHeader(http.StatusBadRequest)
		w.Write([]byte(err.Error()))
	}

	endpoint := *r.URL
	endpoint.Host = p.Host
	endpoint.Scheme = p.Scheme

	req, err := http.NewRequest(r.Method, endpoint.String(), r.Body)
	if err != nil {
		respondError(err)
		return
	}

	// Workaround for ES 5.1 and Kibana 5.1.1
	if val, ok := r.Header["Kbn-Version"]; ok {
		req.Header.Set("Kbn-Version", val[0])
	}

	// Workaround for ES 6.0
	if val, ok := r.Header["Content-Type"]; ok {
		if val[0] == "application/json" {
			req.Header.Set("Content-Type", "application/json")
		}
	}

	// Start AWS session from ENV, Shared Creds or EC2Role
	signer := getSigner(p)

	// Sign the request with AWSv4
	payload := bytes.NewReader(replaceBody(req))
	signer.Sign(req, payload, p.Service, p.Region, time.Now())

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		log.Println(err)
		respondError(err)
		return
	}

	if resp.StatusCode == 403 {
		p.Credentials = nil
		return
	}

	defer resp.Body.Close()

	// Write back received headers
	copyHeaders(w.Header(), resp.Header)

	buf := bytes.Buffer{}
	if _, err := io.Copy(&buf, resp.Body); err != nil {
		log.Fatal(err)
	}

	// Send response back
	w.WriteHeader(resp.StatusCode)
	w.Write(buf.Bytes())

	// Log everything
	remoteAddr := r.RemoteAddr
	rawQuery := string(dump)
	rawQuery = strings.Replace(rawQuery, "\n", " ", -1)
	regex, _ := regexp.Compile("{.*}")
	regEx, _ := regexp.Compile("_msearch|_bulk")
	queryEx := regEx.FindString(rawQuery)

	var query string

	if len(queryEx) == 0 {
		query = regex.FindString(rawQuery)
	} else {
		query = ""
	}

	if p.Verbose {
		requestEnded := time.Since(requestStarted)

		if p.Prettify {
			var prettyBody bytes.Buffer
			json.Indent(&prettyBody, []byte(query), "", "  ")
			t := time.Now()

			fmt.Println()
			fmt.Println("========================")
			fmt.Println(t.Format("2006/01/02 15:04:05"))
			fmt.Println("Remote Address: ", remoteAddr)
			fmt.Println("Request URI: ", endpoint.RequestURI())
			fmt.Println("Method: ", r.Method)
			fmt.Println("Status: ", resp.StatusCode)
			fmt.Printf("Took: %.3fs\n", requestEnded.Seconds())
			fmt.Println("Body: ")
			fmt.Println(string(prettyBody.Bytes()))
			fmt.Println("========================")

		} else {
			log.Printf(" -> %s; %s; %s; %s; %d; %.3fs\n",
				r.Method, remoteAddr, endpoint.RequestURI(), query, resp.StatusCode, requestEnded.Seconds())
		}
	}
}

func main() {
	var endpoint, listenAddress string
	var verbose bool
	var prettify bool

	// TODO: Use a more sophisticated args parser that can enforce arguments
	flag.StringVar(&endpoint, "endpoint", "", "Amazon ElasticSearch Endpoint (e.g: https://dummy-host.eu-west-1.es.amazonaws.com)")
	flag.StringVar(&listenAddress, "listen", "127.0.0.1:9200", "Local TCP port to listen on")
	flag.BoolVar(&verbose, "verbose", false, "Print user requests")
	flag.BoolVar(&prettify, "pretty", false, "Prettify verbose output")

	flag.Parse()

	if len(os.Args) < 3 {
		fmt.Println("You need to specify Amazon ElasticSearch endpoint.")
		fmt.Println("Please run with '-h' for a list of available arguments.")
		os.Exit(1)
	}

	mux := &proxy{Verbose: verbose, Prettify: prettify}
	parseEndpoint(endpoint, mux)

	fmt.Printf("Listening on %s\n", listenAddress)
	log.Fatal(http.ListenAndServe(listenAddress, mux))
}
<commit_msg>- Refactor most of the code. - Added ability to stop signing requests. This feature makes aws-es-proxy work with any ElasticSearch installation. - Added ability to log requests and responses to log files.<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	// log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/aws/signer/v4"
	uuid "github.com/satori/go.uuid"
)

type requestStruct struct {
	Requestid  string
	Datetime   string
	Remoteaddr string
	Requesturi string
	Method     string
	Statuscode int
	Elapsed    float64
	Body       string
}

type responseStruct struct {
	Requestid string
	Body      string
}

type proxy struct {
	scheme       string
	host         string
	region       string
	service      string
	endpoint     string
	verbose      bool
	prettify     bool
	logtofile    bool
	nosignreq    bool
	fileRequest  *os.File
	fileResponse *os.File
	credentials  *credentials.Credentials
}

func newProxy(args ...interface{}) *proxy {
	return &proxy{
		endpoint:  args[0].(string),
		verbose:   args[1].(bool),
		prettify:  args[2].(bool),
		logtofile: args[3].(bool),
		nosignreq: args[4].(bool),
	}
}

func (p *proxy) parseEndpoint() error {
	var link *url.URL
	var err error

	if link, err = url.Parse(p.endpoint); err != nil {
		return fmt.Errorf("error: failure while parsing endpoint: %s. Error: %s",
			p.endpoint, err.Error())
	}

	// Only http/https are supported schemes
	switch link.Scheme {
	case "http", "https":
	default:
		link.Scheme = "https"
	}

	// Unknown schemes sometimes result in empty host value
	if link.Host == "" {
		return fmt.Errorf("error: empty host or protocol information in submitted endpoint (%s)",
			p.endpoint)
	}

	// AWS SignV4 enabled, extract required parts for signing process
	if !p.nosignreq {
		// Extract region and service from link
		parts := strings.Split(link.Host, ".")

		if len(parts) == 5 {
			p.region, p.service = parts[1], parts[2]
		} else {
			return fmt.Errorf("error: submitted endpoint is not a valid Amazon ElasticSearch Endpoint")
		}
	}

	// Update proxy struct
	p.scheme = link.Scheme
	p.host = link.Host

	return nil
}

func (p *proxy) getSigner() *v4.Signer {
	// Refresh credentials after expiration. Required for STS
	if p.credentials == nil {
		sess := session.Must(session.NewSession())
		credentials := sess.Config.Credentials
		p.credentials = credentials
		log.Println("Generated fresh AWS Credentials object")
	}
	return v4.NewSigner(p.credentials)
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	requestStarted := time.Now()
	dump, err := httputil.DumpRequest(r, true)
	if err != nil {
		log.Fatalln("error while dumping request. Error: ", err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer r.Body.Close()

	ep := *r.URL
	ep.Host = p.host
	ep.Scheme = p.scheme

	req, err := http.NewRequest(r.Method, ep.String(), r.Body)
	if err != nil {
		log.Fatalln("error creating new request. ", err.Error())
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	addHeaders(r.Header, req.Header)

	// Make signV4 optional
	if !p.nosignreq {
		// Start AWS session from ENV, Shared Creds or EC2Role
		signer := p.getSigner()

		// Sign the request with AWSv4
		payload := bytes.NewReader(replaceBody(req))
		signer.Sign(req, payload, p.service, p.region, time.Now())
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		log.Fatalln(err.Error())
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	if !p.nosignreq {
		// AWS credentials expired, need to generate fresh ones
		if resp.StatusCode == 403 {
			p.credentials = nil
			return
		}
	}

	defer resp.Body.Close()

	// Write back headers to requesting client
	copyHeaders(w.Header(), resp.Header)

	// Send response back to requesting client
	body := bytes.Buffer{}
	if _, err := io.Copy(&body, resp.Body); err != nil {
		log.Fatalln(err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	w.WriteHeader(resp.StatusCode)
	w.Write(body.Bytes())

	requestEnded := time.Since(requestStarted)

	/*############################
	## Logging
	############################*/

	rawQuery := string(dump)
	rawQuery = strings.Replace(rawQuery, "\n", " ", -1)
	regex, _ := regexp.Compile("{.*}")
	regEx, _ := regexp.Compile("_msearch|_bulk")
	queryEx := regEx.FindString(rawQuery)

	var query string

	if len(queryEx) == 0 {
		query = regex.FindString(rawQuery)
	} else {
		query = ""
	}

	if p.verbose {
		if p.prettify {
			var prettyBody bytes.Buffer
			json.Indent(&prettyBody, []byte(query), "", "  ")
			t := time.Now()

			fmt.Println()
			fmt.Println("========================")
			fmt.Println(t.Format("2006/01/02 15:04:05"))
			fmt.Println("Remote Address: ", r.RemoteAddr)
			fmt.Println("Request URI: ", ep.RequestURI())
			fmt.Println("Method: ", r.Method)
			fmt.Println("Status: ", resp.StatusCode)
			fmt.Printf("Took: %.3fs\n", requestEnded.Seconds())
			fmt.Println("Body: ")
			fmt.Println(string(prettyBody.Bytes()))
		} else {
			log.Printf(" -> %s; %s; %s; %s; %d; %.3fs\n",
				r.Method, r.RemoteAddr,
				ep.RequestURI(), query,
				resp.StatusCode, requestEnded.Seconds())
		}
	}

	if p.logtofile {
		requestID := fmt.Sprintf("%s", uuid.NewV4().String())

		reqStruct := &requestStruct{
			Requestid:  requestID,
			Datetime:   time.Now().Format("2006/01/02 15:04:05"),
			Remoteaddr: r.RemoteAddr,
			Requesturi: ep.RequestURI(),
			Method:     r.Method,
			Statuscode: resp.StatusCode,
			Elapsed:    requestEnded.Seconds(),
			Body:       query,
		}

		respStruct := &responseStruct{
			Requestid: requestID,
			Body:      string(body.Bytes()),
		}

		y, _ := json.Marshal(reqStruct)
		z, _ := json.Marshal(respStruct)
		p.fileRequest.Write(y)
		p.fileRequest.WriteString("\n")
		p.fileResponse.Write(z)
		p.fileResponse.WriteString("\n")

	}

}

// Recent versions of ES/Kibana require
// "kbn-version" and "content-type: application/json"
// headers to exist in the request.
// If missing requests fails.
func addHeaders(src, dest http.Header) {
	if val, ok := src["Kbn-Version"]; ok {
		dest.Add("Kbn-Version", val[0])
	}

	if val, ok := src["Content-Type"]; ok {
		if val[0] == "application/json" {
			dest.Add("Content-Type", "application/json")
		}
	}
}

// Signer.Sign requires a "seekable" body to sum body's sha256
func replaceBody(req *http.Request) []byte {
	if req.Body == nil {
		return []byte{}
	}
	payload, _ := ioutil.ReadAll(req.Body)
	req.Body = ioutil.NopCloser(bytes.NewReader(payload))
	return payload
}

func copyHeaders(dst, src http.Header) {
	for k, vals := range src {
		for _, v := range vals {
			dst.Add(k, v)
		}
	}
}

func main() {

	var (
		verbose       bool
		prettify      bool
		logtofile     bool
		nosignreq     bool
		endpoint      string
		listenAddress string
		fileRequest   *os.File
		fileResponse  *os.File
		err           error
	)

	flag.StringVar(&endpoint, "endpoint", "", "Amazon ElasticSearch Endpoint (e.g: https://dummy-host.eu-west-1.es.amazonaws.com)")
	flag.StringVar(&listenAddress, "listen", "127.0.0.1:9200", "Local TCP port to listen on")
	flag.BoolVar(&verbose, "verbose", false, "Print user requests")
	flag.BoolVar(&logtofile, "log-to-file", false, "Log user requests and ElasticSearch responses to files")
	flag.BoolVar(&prettify, "pretty", false, "Prettify verbose and file output")
	flag.BoolVar(&nosignreq, "no-sign-reqs", false, "Disable AWS Signature v4")
	flag.Parse()

	if len(os.Args) < 3 {
		fmt.Println("You need to specify Amazon ElasticSearch endpoint.")
		fmt.Println("Please run with '-h' for a list of available arguments.")
		os.Exit(1)
	}

	p := newProxy(
		endpoint,
		verbose,
		prettify,
		logtofile,
		nosignreq,
	)

	if err = p.parseEndpoint(); err != nil {
		log.Fatalln(err)
		os.Exit(1)
	}

	if p.logtofile {
		requestFname := fmt.Sprintf("request-%s.log", uuid.NewV4().String())
		responseFname := fmt.Sprintf("response-%s.log", uuid.NewV4().String())

		if fileRequest, err = os.Create(requestFname); err != nil {
			log.Fatalln(err.Error())
		}
		if fileResponse, err = os.Create(responseFname); err != nil {
			log.Fatalln(err.Error())
		}

		defer fileRequest.Close()
		defer fileResponse.Close()

		p.fileRequest = fileRequest
		p.fileResponse = fileResponse

	}

	log.Printf("Listening on %s...\n", listenAddress)
	log.Fatal(http.ListenAndServe(listenAddress, p))
}
<|endoftext|>


<commit_before>package tui

import (
	"image"
	"math"
)

var _ Widget = &Box{}

// Alignment is used to set the direction in which widgets are laid out.
type Alignment int

// Available alignment options.
const (
	Horizontal Alignment = iota
	Vertical
)

// Box is a layout for placing widgets.
type Box struct {
	WidgetBase

	children []Widget

	border bool
	title  string

	alignment Alignment
}

// NewVBox returns a new vertical Box.
func NewVBox(c ...Widget) *Box {
	return &Box{
		children:  c,
		alignment: Vertical,
	}
}

// NewHBox returns a new horizontal Box.
func NewHBox(c ...Widget) *Box {
	return &Box{
		children:  c,
		alignment: Horizontal,
	}
}

// Append adds a new widget to the layout.
func (b *Box) Append(w Widget) {
	b.children = append(b.children, w)
}

// SetBorder sets whether the border is visible or not.
func (b *Box) SetBorder(enabled bool) {
	b.border = enabled
}

// SetTitle sets the title of the box.
func (b *Box) SetTitle(title string) {
	b.title = title
}

// Alignment returns the currently set alignment or the Box.
func (b *Box) Alignment() Alignment {
	return b.alignment
}

// IsFocused return true if one of the children is focused.
func (b *Box) IsFocused() bool {
	for _, w := range b.children {
		if w.IsFocused() {
			return true
		}
	}
	return false
}

// Draw recursively draws the children it contains.
func (b *Box) Draw(p *Painter) {
	style := "box"
	if b.IsFocused() {
		style += ".focused"
	}

	sz := b.Size()

	p.WithStyle(style, func(p *Painter) {
		if b.border {
			p.DrawRect(0, 0, sz.X, sz.Y)
			p.WithMask(image.Rect(2, 0, sz.X-3, 0)).DrawText(2, 0, b.title)

			p.Translate(1, 1)
			defer p.Restore()
		}

		var off image.Point
		for _, child := range b.children {
			switch b.Alignment() {
			case Horizontal:
				p.Translate(off.X, 0)
			case Vertical:
				p.Translate(0, off.Y)
			}

			child.Draw(p.WithMask(image.Rectangle{
				Min: image.ZP,
				Max: child.Size().Sub(image.Point{1, 1}),
			}))

			p.Restore()

			off = off.Add(child.Size())
		}
	})
}

// MinSizeHint returns the minimum size for the layout.
func (b *Box) MinSizeHint() image.Point {
	var minSize image.Point

	for _, child := range b.children {
		size := child.MinSizeHint()
		if b.Alignment() == Horizontal {
			minSize.X += size.X
			if size.Y > minSize.Y {
				minSize.Y = size.Y
			}
		} else {
			minSize.Y += size.Y
			if size.X > minSize.X {
				minSize.X = size.X
			}
		}
	}

	if b.border {
		minSize = minSize.Add(image.Point{2, 2})
	}

	return minSize
}

// SizeHint returns the recommended size for the layout.
func (b *Box) SizeHint() image.Point {
	var sizeHint image.Point

	for _, child := range b.children {
		size := child.SizeHint()
		if b.Alignment() == Horizontal {
			sizeHint.X += size.X
			if size.Y > sizeHint.Y {
				sizeHint.Y = size.Y
			}
		} else {
			sizeHint.Y += size.Y
			if size.X > sizeHint.X {
				sizeHint.X = size.X
			}
		}
	}

	if b.border {
		sizeHint = sizeHint.Add(image.Point{2, 2})
	}

	return sizeHint
}

// OnEvent handles an event and propagates it to all children.
func (b *Box) OnEvent(ev Event) {
	for _, child := range b.children {
		child.OnEvent(ev)
	}
}

// Resize updates the size of the layout.
func (b *Box) Resize(size image.Point) {
	b.size = size
	inner := b.size
	if b.border {
		inner = b.size.Sub(image.Point{2, 2})
	}
	b.layoutChildren(inner)
}

func (b *Box) layoutChildren(size image.Point) {
	space := doLayout(b.children, dim(b.Alignment(), size), b.Alignment())

	for i, s := range space {
		switch b.Alignment() {
		case Horizontal:
			b.children[i].Resize(image.Point{s, size.Y})
		case Vertical:
			b.children[i].Resize(image.Point{size.X, s})
		}
	}
}

func doLayout(ws []Widget, space int, a Alignment) []int {
	sizes := make([]int, len(ws))

	if len(sizes) == 0 {
		return sizes
	}

	remaining := space

	// Distribute MinSizeHint
	for {
		var changed bool
		for i, sz := range sizes {
			if sz < dim(a, ws[i].MinSizeHint()) {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute Minimum
	for {
		var changed bool
		for i, sz := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if p == Minimum && sz < dim(a, ws[i].SizeHint()) {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute Preferred
	for {
		var changed bool
		for i, sz := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if (p == Preferred || p == Maximum) && sz < dim(a, ws[i].SizeHint()) {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute Expanding
	for {
		var changed bool
		for i, sz := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if p == Expanding {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute remaining space
	for {
		min := math.MaxInt8
		for i, s := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if (p == Preferred || p == Minimum) && s <= min {
				min = s
			}
		}
		var changed bool
		for i, sz := range sizes {
			if sz != min {
				continue
			}
			p := alignedSizePolicy(a, ws[i])
			if p == Preferred || p == Minimum {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

Resize:

	return sizes
}

func dim(a Alignment, pt image.Point) int {
	if a == Horizontal {
		return pt.X
	}
	return pt.Y
}

func alignedSizePolicy(a Alignment, w Widget) SizePolicy {
	hpol, vpol := w.SizePolicy()
	if a == Horizontal {
		return hpol
	}
	return vpol
}
<commit_msg>Do not apply the box.focused to the children widget<commit_after>package tui

import (
	"image"
	"math"
)

var _ Widget = &Box{}

// Alignment is used to set the direction in which widgets are laid out.
type Alignment int

// Available alignment options.
const (
	Horizontal Alignment = iota
	Vertical
)

// Box is a layout for placing widgets.
type Box struct {
	WidgetBase

	children []Widget

	border bool
	title  string

	alignment Alignment
}

// NewVBox returns a new vertical Box.
func NewVBox(c ...Widget) *Box {
	return &Box{
		children:  c,
		alignment: Vertical,
	}
}

// NewHBox returns a new horizontal Box.
func NewHBox(c ...Widget) *Box {
	return &Box{
		children:  c,
		alignment: Horizontal,
	}
}

// Append adds a new widget to the layout.
func (b *Box) Append(w Widget) {
	b.children = append(b.children, w)
}

// SetBorder sets whether the border is visible or not.
func (b *Box) SetBorder(enabled bool) {
	b.border = enabled
}

// SetTitle sets the title of the box.
func (b *Box) SetTitle(title string) {
	b.title = title
}

// Alignment returns the currently set alignment or the Box.
func (b *Box) Alignment() Alignment {
	return b.alignment
}

// IsFocused return true if one of the children is focused.
func (b *Box) IsFocused() bool {
	for _, w := range b.children {
		if w.IsFocused() {
			return true
		}
	}
	return false
}

// Draw recursively draws the children it contains.
func (b *Box) Draw(p *Painter) {
	style := "box"
	if b.IsFocused() {
		style += ".focused"
	}

	sz := b.Size()

	if b.border {
		p.WithStyle(style, func(p *Painter) {
			p.DrawRect(0, 0, sz.X, sz.Y)
			p.WithMask(image.Rect(2, 0, sz.X-3, 0)).DrawText(2, 0, b.title)
		})

		p.Translate(1, 1)
		defer p.Restore()
	}

	var off image.Point
	for _, child := range b.children {
		switch b.Alignment() {
		case Horizontal:
			p.Translate(off.X, 0)
		case Vertical:
			p.Translate(0, off.Y)
		}

		child.Draw(p.WithMask(image.Rectangle{
			Min: image.ZP,
			Max: child.Size().Sub(image.Point{1, 1}),
		}))

		p.Restore()

		off = off.Add(child.Size())
	}
}

// MinSizeHint returns the minimum size for the layout.
func (b *Box) MinSizeHint() image.Point {
	var minSize image.Point

	for _, child := range b.children {
		size := child.MinSizeHint()
		if b.Alignment() == Horizontal {
			minSize.X += size.X
			if size.Y > minSize.Y {
				minSize.Y = size.Y
			}
		} else {
			minSize.Y += size.Y
			if size.X > minSize.X {
				minSize.X = size.X
			}
		}
	}

	if b.border {
		minSize = minSize.Add(image.Point{2, 2})
	}

	return minSize
}

// SizeHint returns the recommended size for the layout.
func (b *Box) SizeHint() image.Point {
	var sizeHint image.Point

	for _, child := range b.children {
		size := child.SizeHint()
		if b.Alignment() == Horizontal {
			sizeHint.X += size.X
			if size.Y > sizeHint.Y {
				sizeHint.Y = size.Y
			}
		} else {
			sizeHint.Y += size.Y
			if size.X > sizeHint.X {
				sizeHint.X = size.X
			}
		}
	}

	if b.border {
		sizeHint = sizeHint.Add(image.Point{2, 2})
	}

	return sizeHint
}

// OnEvent handles an event and propagates it to all children.
func (b *Box) OnEvent(ev Event) {
	for _, child := range b.children {
		child.OnEvent(ev)
	}
}

// Resize updates the size of the layout.
func (b *Box) Resize(size image.Point) {
	b.size = size
	inner := b.size
	if b.border {
		inner = b.size.Sub(image.Point{2, 2})
	}
	b.layoutChildren(inner)
}

func (b *Box) layoutChildren(size image.Point) {
	space := doLayout(b.children, dim(b.Alignment(), size), b.Alignment())

	for i, s := range space {
		switch b.Alignment() {
		case Horizontal:
			b.children[i].Resize(image.Point{s, size.Y})
		case Vertical:
			b.children[i].Resize(image.Point{size.X, s})
		}
	}
}

func doLayout(ws []Widget, space int, a Alignment) []int {
	sizes := make([]int, len(ws))

	if len(sizes) == 0 {
		return sizes
	}

	remaining := space

	// Distribute MinSizeHint
	for {
		var changed bool
		for i, sz := range sizes {
			if sz < dim(a, ws[i].MinSizeHint()) {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute Minimum
	for {
		var changed bool
		for i, sz := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if p == Minimum && sz < dim(a, ws[i].SizeHint()) {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute Preferred
	for {
		var changed bool
		for i, sz := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if (p == Preferred || p == Maximum) && sz < dim(a, ws[i].SizeHint()) {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute Expanding
	for {
		var changed bool
		for i, sz := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if p == Expanding {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

	// Distribute remaining space
	for {
		min := math.MaxInt8
		for i, s := range sizes {
			p := alignedSizePolicy(a, ws[i])
			if (p == Preferred || p == Minimum) && s <= min {
				min = s
			}
		}
		var changed bool
		for i, sz := range sizes {
			if sz != min {
				continue
			}
			p := alignedSizePolicy(a, ws[i])
			if p == Preferred || p == Minimum {
				sizes[i] = sz + 1
				remaining--
				if remaining <= 0 {
					goto Resize
				}
				changed = true
			}
		}
		if !changed {
			break
		}
	}

Resize:

	return sizes
}

func dim(a Alignment, pt image.Point) int {
	if a == Horizontal {
		return pt.X
	}
	return pt.Y
}

func alignedSizePolicy(a Alignment, w Widget) SizePolicy {
	hpol, vpol := w.SizePolicy()
	if a == Horizontal {
		return hpol
	}
	return vpol
}
<|endoftext|>


<commit_before>package sigupdate

import (
	"errors"
	"fmt"
	"github.com/dekobon/clamav-mirror/utils"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"
)

import (
	"github.com/hashicorp/errwrap"
)

// Function that downloads a file from the mirror URL and moves it into the
// data directory if it was successfully downloaded.
func downloadFile(filename string, localFilePath string,
	downloadMirrorURL string) (int, error) {

	unknownStatus := -1
	downloadURL := downloadMirrorURL + "/" + filename

	output, err := ioutil.TempFile(os.TempDir(), filename+"-")

	/* For all non-cvd files, skip downloading the file if our local copy is
	 * newer than the remote copy. For .cvd files, the only authoritative way
	 * to know what is newer is to use sigtool. */
	if utils.Exists(localFilePath) {
		newer, err := checkIfRemoteIsNewer(localFilePath, downloadURL)

		if err != nil {
			return unknownStatus, err
		}

		if !newer {
			logger.Printf("Not downloading [%v] because remote copy is "+
				"older than local copy", filename)
			return unknownStatus, nil
		}
	}

	if verboseMode {
		logger.Printf("Downloading to temporary file: [%v]", output.Name())
	}

	if err != nil {
		msg := fmt.Sprintf("Unable to create file: [%v]. {{err}}", output.Name())
		return unknownStatus, errwrap.Wrapf(msg, err)
	}

	defer output.Close()

	response, err := http.Get(downloadURL)

	if err != nil {
		msg := fmt.Sprintf("Unable to retrieve file from: [%v]. {{err}}", downloadURL)
		return unknownStatus, errwrap.Wrapf(msg, err)
	}

	if response.StatusCode != http.StatusOK {
		msg := fmt.Sprintf("Unable to download file: [%v]", response.Status)
		return response.StatusCode, errors.New(msg)
	}

	lastModified, err := http.ParseTime(response.Header.Get("Last-Modified"))

	if err != nil {
		logger.Printf("Error parsing last-modified header [%v] for file: %v",
			response.Header.Get("Last-Modified"), downloadURL)
		lastModified = time.Now()
	}

	defer response.Body.Close()

	n, err := io.Copy(output, response.Body)

	if err != nil {
		msg := fmt.Sprintf("Error copying data from URL [%v] to local file [%v]. {{err}}",
			downloadURL, localFilePath)
		return response.StatusCode, errwrap.Wrapf(msg, err)
	}

	if isItOkToOverwrite(filename, localFilePath, output.Name()) {
		/* Change the last modified time so that we have a record that corresponds to the
		 * server's timestamps. */
		os.Chtimes(output.Name(), lastModified, lastModified)
		os.Rename(output.Name(), localFilePath)

		logger.Printf("Download complete: %v --> %v [%v bytes]", downloadURL, localFilePath, n)
	} else {
		logger.Println("Downloaded file an older signature version than the current file")

		err := os.Remove(output.Name())

		if err != nil {
			logger.Printf("Unable to delete temporary file: %v", output.Name())
		}
	}

	return response.StatusCode, nil
}

// Function that checks to see if the remote file is newer than the locally stored
// file.
func checkIfRemoteIsNewer(localFilePath string, downloadURL string) (bool, error) {
	localFileStat, err := os.Stat(localFilePath)

	if err != nil {
		return true, errwrap.Wrapf("Unable to stat file. {{err}}", err)
	}

	localModTime := localFileStat.ModTime()
	response, err := http.Head(downloadURL)

	if err != nil {
		msg := fmt.Sprintf("Unable to complete HEAD request: [%v]. {{err}}", downloadURL)
		return true, errwrap.Wrapf(msg, err)
	}

	remoteModTime, err := http.ParseTime(response.Header.Get("Last-Modified"))

	if verboseMode {
		logger.Printf("Local file [%v] last-modified: %v", downloadURL, localModTime)
		logger.Printf("Remote file [%v] last-modified: %v", downloadURL, remoteModTime)
	}

	if err != nil {
		msg := fmt.Sprintf("Error parsing last-modified header [%v] for file [%v]. {{err}}",
			response.Header.Get("Last-Modified"), downloadURL)
		return true, errwrap.Wrapf(msg, err)
	}

	if localModTime.After(remoteModTime) {
		logger.Printf("Skipping download of [%v] because local copy is newer", downloadURL)
		return false, nil
	}

	return true, nil
}

// Function that checks to see if we can overwrite a file with a newly downloaded file
func isItOkToOverwrite(filename string, originalFilePath string, newFileTempPath string) bool {
	if !strings.HasSuffix(filename, ".cvd") {
		return true
	}

	oldVersion, err := findLocalVersion(originalFilePath)

	// If there is a problem with the original file, we just overwrite it
	if err != nil {
		return true
	}

	newVersion, err := findLocalVersion(newFileTempPath)

	// If there is a problem with the new file, we don't overwrite the original
	if err != nil {
		return false
	}

	isNewer := newVersion > oldVersion

	if verboseMode {
		logger.Printf("Current file [%v] version [%v]. New file version [%v]. "+
			"Will overwrite: %v",
			filename, oldVersion, newVersion, isNewer)
	}

	return isNewer
}
<commit_msg>Added user-agent and if-modified-since to client headers<commit_after>package sigupdate

import (
	"errors"
	"fmt"
	"github.com/dekobon/clamav-mirror/utils"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"
)

import (
	"github.com/hashicorp/errwrap"
)

// Function that downloads a file from the mirror URL and moves it into the
// data directory if it was successfully downloaded.
func downloadFile(filename string, localFilePath string,
	downloadMirrorURL string) (int, error) {

	unknownStatus := -1
	downloadURL := downloadMirrorURL + "/" + filename

	output, err := ioutil.TempFile(os.TempDir(), filename+"-")

	if verboseMode {
		logger.Printf("Downloading to temporary file: [%v]", output.Name())
	}

	if err != nil {
		msg := fmt.Sprintf("Unable to create file: [%v]. {{err}}", output.Name())
		return unknownStatus, errwrap.Wrapf(msg, err)
	}

	defer output.Close()

	request, err := http.NewRequest("GET", downloadURL, nil)

	if err != nil {
		msg := fmt.Sprintf("Unable to create request for: [GEt %v]. {{err}}", downloadURL)
		return unknownStatus, errwrap.Wrapf(msg, err)
	}

	request.Header.Add("User-Agent", "github.com/dekobon/clamav-mirror")

	/* For all non-cvd files, skip downloading the file if our local copy is
	 * newer than the remote copy. For .cvd files, the only authoritative way
	 * to know what is newer is to use sigtool. */
	if utils.Exists(localFilePath) {
		stat, err := os.Stat(localFilePath)

		if err == nil {
			localModTime := stat.ModTime().UTC().Format(http.TimeFormat)
			request.Header.Add("If-Modified-Since", localModTime)
		} else {
			logger.Printf("Unable to stat local file [%v]. %v", localFilePath, err)
		}
	}

	response, err := http.DefaultClient.Do(request)

	if err != nil {
		msg := fmt.Sprintf("Unable to retrieve file from: [%v]. {{err}}", downloadURL)
		return unknownStatus, errwrap.Wrapf(msg, err)
	}

	if response.StatusCode == http.StatusNotModified {
		logger.Printf("Not downloading [%v] because local copy is newer or the same as remote",
			filename)
		return response.StatusCode, nil
	}

	if response.StatusCode != http.StatusOK {
		msg := fmt.Sprintf("Unable to download file: [%v]", response.Status)
		return response.StatusCode, errors.New(msg)
	}

	lastModified, err := http.ParseTime(response.Header.Get("Last-Modified"))

	if err != nil {
		logger.Printf("Error parsing last-modified header [%v] for file: %v",
			response.Header.Get("Last-Modified"), downloadURL)
		lastModified = time.Now()
	}

	defer response.Body.Close()

	n, err := io.Copy(output, response.Body)

	if err != nil {
		msg := fmt.Sprintf("Error copying data from URL [%v] to local file [%v]. {{err}}",
			downloadURL, localFilePath)
		return response.StatusCode, errwrap.Wrapf(msg, err)
	}

	if isItOkToOverwrite(filename, localFilePath, output.Name()) {
		/* Change the last modified time so that we have a record that corresponds to the
		 * server's timestamps. */
		os.Chtimes(output.Name(), lastModified, lastModified)
		os.Rename(output.Name(), localFilePath)

		logger.Printf("Download complete: %v --> %v [%v bytes]", downloadURL, localFilePath, n)
	} else {
		logger.Println("Downloaded file an older signature version than the current file")

		err := os.Remove(output.Name())

		if err != nil {
			logger.Printf("Unable to delete temporary file: %v", output.Name())
		}
	}

	return response.StatusCode, nil
}

// Function that checks to see if we can overwrite a file with a newly downloaded file
func isItOkToOverwrite(filename string, originalFilePath string, newFileTempPath string) bool {
	if !strings.HasSuffix(filename, ".cvd") {
		return true
	}

	oldVersion, err := findLocalVersion(originalFilePath)

	// If there is a problem with the original file, we just overwrite it
	if err != nil {
		return true
	}

	newVersion, err := findLocalVersion(newFileTempPath)

	// If there is a problem with the new file, we don't overwrite the original
	if err != nil {
		return false
	}

	isNewer := newVersion > oldVersion

	if verboseMode {
		logger.Printf("Current file [%v] version [%v]. New file version [%v]. "+
			"Will overwrite: %v",
			filename, oldVersion, newVersion, isNewer)
	}

	return isNewer
}
<|endoftext|>


<commit_before><commit_msg>base types: remove unused method<commit_after><|endoftext|>


<commit_before>package ssh

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/term"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnutils"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

type Client interface {
	Output(command string) (string, error)
	Shell(args ...string) error

	// Start starts the specified command without waiting for it to finish. You
	// have to call the Wait function for that.
	//
	// The first two io.ReadCloser are the standard output and the standard
	// error of the executing command respectively. The returned error follows
	// the same logic as in the exec.Cmd.Start function.
	Start(command string) (io.ReadCloser, io.ReadCloser, error)

	// Wait waits for the command started by the Start function to exit. The
	// returned error follows the same logic as in the exec.Cmd.Wait function.
	Wait() error
}

type ExternalClient struct {
	BaseArgs   []string
	BinaryPath string
	cmd        *exec.Cmd
}

type NativeClient struct {
	Config      ssh.ClientConfig
	Hostname    string
	Port        int
	openSession *ssh.Session
	openClient  *ssh.Client
}

type Auth struct {
	Passwords []string
	Keys      []string
}

type ClientType string

const (
	maxDialAttempts = 10
)

const (
	External ClientType = "external"
	Native   ClientType = "native"
)

var (
	baseSSHArgs = []string{
		"-F", "/dev/null",
		"-o", "PasswordAuthentication=no",
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
		"-o", "LogLevel=quiet", // suppress "Warning: Permanently added '[localhost]:2022' (ECDSA) to the list of known hosts."
		"-o", "ConnectionAttempts=3", // retry 3 times if SSH connection fails
		"-o", "ConnectTimeout=10", // timeout after 10 seconds
		"-o", "ControlMaster=no", // disable ssh multiplexing
		"-o", "ControlPath=none",
	}
	defaultClientType = External
)

func SetDefaultClient(clientType ClientType) {
	// Allow over-riding of default client type, so that even if ssh binary
	// is found in PATH we can still use the Go native implementation if
	// desired.
	switch clientType {
	case External:
		defaultClientType = External
	case Native:
		defaultClientType = Native
	}
}

func NewClient(user string, host string, port int, auth *Auth) (Client, error) {
	sshBinaryPath, err := exec.LookPath("ssh")
	if err != nil {
		log.Debug("SSH binary not found, using native Go implementation")
		client, err := NewNativeClient(user, host, port, auth)
		log.Debug(client)
		return client, err
	}

	if defaultClientType == Native {
		log.Debug("Using SSH client type: native")
		client, err := NewNativeClient(user, host, port, auth)
		log.Debug(client)
		return client, err
	}

	log.Debug("Using SSH client type: external")
	client, err := NewExternalClient(sshBinaryPath, user, host, port, auth)
	log.Debug(client)
	return client, err
}

func NewNativeClient(user, host string, port int, auth *Auth) (Client, error) {
	config, err := NewNativeConfig(user, auth)
	if err != nil {
		return nil, fmt.Errorf("Error getting config for native Go SSH: %s", err)
	}

	return &NativeClient{
		Config:   config,
		Hostname: host,
		Port:     port,
	}, nil
}

func NewNativeConfig(user string, auth *Auth) (ssh.ClientConfig, error) {
	var (
		authMethods []ssh.AuthMethod
	)

	for _, k := range auth.Keys {
		key, err := ioutil.ReadFile(k)
		if err != nil {
			return ssh.ClientConfig{}, err
		}

		privateKey, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return ssh.ClientConfig{}, err
		}

		authMethods = append(authMethods, ssh.PublicKeys(privateKey))
	}

	for _, p := range auth.Passwords {
		authMethods = append(authMethods, ssh.Password(p))
	}

	return ssh.ClientConfig{
		User:            user,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}, nil
}

func (client *NativeClient) dialSuccess() bool {
	conn, err := ssh.Dial("tcp", net.JoinHostPort(client.Hostname, strconv.Itoa(client.Port)), &client.Config)
	if err != nil {
		log.Debugf("Error dialing TCP: %s", err)
		return false
	}
	closeConn(conn)
	return true
}

func (client *NativeClient) session(command string) (*ssh.Client, *ssh.Session, error) {
	if err := mcnutils.WaitFor(client.dialSuccess); err != nil {
		return nil, nil, fmt.Errorf("Error attempting SSH client dial: %s", err)
	}

	conn, err := ssh.Dial("tcp", net.JoinHostPort(client.Hostname, strconv.Itoa(client.Port)), &client.Config)
	if err != nil {
		return nil, nil, fmt.Errorf("Mysterious error dialing TCP for SSH (we already succeeded at least once) : %s", err)
	}
	session, err := conn.NewSession()

	return conn, session, err
}

func (client *NativeClient) Output(command string) (string, error) {
	conn, session, err := client.session(command)
	if err != nil {
		return "", nil
	}
	defer closeConn(conn)
	defer session.Close()

	output, err := session.CombinedOutput(command)

	return string(output), err
}

func (client *NativeClient) OutputWithPty(command string) (string, error) {
	conn, session, err := client.session(command)
	if err != nil {
		return "", nil
	}
	defer closeConn(conn)
	defer session.Close()

	fd := int(os.Stdout.Fd())

	termWidth, termHeight, err := terminal.GetSize(fd)
	if err != nil {
		return "", err
	}

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	// request tty -- fixes error with hosts that use
	// "Defaults requiretty" in /etc/sudoers - I'm looking at you RedHat
	if err := session.RequestPty("xterm", termHeight, termWidth, modes); err != nil {
		return "", err
	}

	output, err := session.CombinedOutput(command)

	return string(output), err
}

func (client *NativeClient) Start(command string) (io.ReadCloser, io.ReadCloser, error) {
	conn, session, err := client.session(command)
	if err != nil {
		return nil, nil, err
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}
	stderr, err := session.StderrPipe()
	if err != nil {
		return nil, nil, err
	}
	if err := session.Start(command); err != nil {
		return nil, nil, err
	}

	client.openClient = conn
	client.openSession = session
	return ioutil.NopCloser(stdout), ioutil.NopCloser(stderr), nil
}

func (client *NativeClient) Wait() error {
	err := client.openSession.Wait()
	if err != nil {
		return err
	}

	_ = client.openSession.Close()

	err = client.openClient.Close()
	if err != nil {
		return err
	}

	client.openSession = nil
	client.openClient = nil
	return nil
}

func (client *NativeClient) Shell(args ...string) error {
	var (
		termWidth, termHeight int
	)
	conn, err := ssh.Dial("tcp", net.JoinHostPort(client.Hostname, strconv.Itoa(client.Port)), &client.Config)
	if err != nil {
		return err
	}
	defer closeConn(conn)

	session, err := conn.NewSession()
	if err != nil {
		return err
	}

	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO: 1,
	}

	fd := os.Stdin.Fd()

	if term.IsTerminal(fd) {
		oldState, err := term.MakeRaw(fd)
		if err != nil {
			return err
		}

		defer term.RestoreTerminal(fd, oldState)

		winsize, err := term.GetWinsize(fd)
		if err != nil {
			termWidth = 80
			termHeight = 24
		} else {
			termWidth = int(winsize.Width)
			termHeight = int(winsize.Height)
		}
	}

	if err := session.RequestPty("xterm", termHeight, termWidth, modes); err != nil {
		return err
	}

	if len(args) == 0 {
		if err := session.Shell(); err != nil {
			return err
		}
		session.Wait()
	} else {
		session.Run(strings.Join(args, " "))
	}

	return nil
}

func NewExternalClient(sshBinaryPath, user, host string, port int, auth *Auth) (*ExternalClient, error) {
	client := &ExternalClient{
		BinaryPath: sshBinaryPath,
	}

	args := append(baseSSHArgs, fmt.Sprintf("%s@%s", user, host))

	// If no identities are explicitly provided, also look at the identities
	// offered by ssh-agent
	if len(auth.Keys) > 0 {
		args = append(args, "-o", "IdentitiesOnly=yes")
	}

	// Specify which private keys to use to authorize the SSH request.
	for _, privateKeyPath := range auth.Keys {
		if privateKeyPath != "" {
			// Check each private key before use it
			fi, err := os.Stat(privateKeyPath)
			if err != nil {
				// Abort if key not accessible
				return nil, err
			}
			if runtime.GOOS != "windows" {
				mode := fi.Mode()
				log.Debugf("Using SSH private key: %s (%s)", privateKeyPath, mode)
				// Private key file should have strict permissions
				perm := mode.Perm()
				if perm&0400 == 0 {
					return nil, fmt.Errorf("'%s' is not readable", privateKeyPath)
				}
				if perm&0077 != 0 {
					return nil, fmt.Errorf("permissions %#o for '%s' are too open", perm, privateKeyPath)
				}
			}
			args = append(args, "-i", privateKeyPath)
		}
	}

	// Set which port to use for SSH.
	args = append(args, "-p", fmt.Sprintf("%d", port))

	client.BaseArgs = args

	return client, nil
}

func getSSHCmd(binaryPath string, args ...string) *exec.Cmd {
	return exec.Command(binaryPath, args...)
}

func (client *ExternalClient) Output(command string) (string, error) {
	args := append(client.BaseArgs, command)
	cmd := getSSHCmd(client.BinaryPath, args...)
	output, err := cmd.CombinedOutput()
	return string(output), err
}

func (client *ExternalClient) Shell(args ...string) error {
	args = append(client.BaseArgs, args...)
	cmd := getSSHCmd(client.BinaryPath, args...)

	log.Debug(cmd)

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

	return cmd.Run()
}

func (client *ExternalClient) Start(command string) (io.ReadCloser, io.ReadCloser, error) {
	args := append(client.BaseArgs, command)
	cmd := getSSHCmd(client.BinaryPath, args...)

	log.Debug(cmd)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		if closeErr := stdout.Close(); closeErr != nil {
			return nil, nil, fmt.Errorf("%s, %s", err, closeErr)
		}
		return nil, nil, err
	}
	if err := cmd.Start(); err != nil {
		stdOutCloseErr := stdout.Close()
		stdErrCloseErr := stderr.Close()
		if stdOutCloseErr != nil || stdErrCloseErr != nil {
			return nil, nil, fmt.Errorf("%s, %s, %s",
				err, stdOutCloseErr, stdErrCloseErr)
		}
		return nil, nil, err
	}

	client.cmd = cmd
	return stdout, stderr, nil
}

func (client *ExternalClient) Wait() error {
	err := client.cmd.Wait()
	client.cmd = nil
	return err
}

func closeConn(c io.Closer) {
	err := c.Close()
	if err != nil {
		log.Debugf("Error closing SSH Client: %s", err)
	}
}
<commit_msg>ISSUE-4450 Added ServerAliveInterval among the default ssh options and sorted them alphabetically to improve readability<commit_after>package ssh

import (
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"

	"github.com/docker/docker/pkg/term"
	"github.com/docker/machine/libmachine/log"
	"github.com/docker/machine/libmachine/mcnutils"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/terminal"
)

type Client interface {
	Output(command string) (string, error)
	Shell(args ...string) error

	// Start starts the specified command without waiting for it to finish. You
	// have to call the Wait function for that.
	//
	// The first two io.ReadCloser are the standard output and the standard
	// error of the executing command respectively. The returned error follows
	// the same logic as in the exec.Cmd.Start function.
	Start(command string) (io.ReadCloser, io.ReadCloser, error)

	// Wait waits for the command started by the Start function to exit. The
	// returned error follows the same logic as in the exec.Cmd.Wait function.
	Wait() error
}

type ExternalClient struct {
	BaseArgs   []string
	BinaryPath string
	cmd        *exec.Cmd
}

type NativeClient struct {
	Config      ssh.ClientConfig
	Hostname    string
	Port        int
	openSession *ssh.Session
	openClient  *ssh.Client
}

type Auth struct {
	Passwords []string
	Keys      []string
}

type ClientType string

const (
	maxDialAttempts = 10
)

const (
	External ClientType = "external"
	Native   ClientType = "native"
)

var (
	baseSSHArgs = []string{
		"-F", "/dev/null",
		"-o", "ConnectionAttempts=3", // retry 3 times if SSH connection fails
		"-o", "ConnectTimeout=10", // timeout after 10 seconds
		"-o", "ControlMaster=no", // disable ssh multiplexing
		"-o", "ControlPath=none",
		"-o", "LogLevel=quiet", // suppress "Warning: Permanently added '[localhost]:2022' (ECDSA) to the list of known hosts."
		"-o", "PasswordAuthentication=no",
		"-o", "ServerAliveInterval=60", // prevents connection to be dropped if command takes too long
		"-o", "StrictHostKeyChecking=no",
		"-o", "UserKnownHostsFile=/dev/null",
	}
	defaultClientType = External
)

func SetDefaultClient(clientType ClientType) {
	// Allow over-riding of default client type, so that even if ssh binary
	// is found in PATH we can still use the Go native implementation if
	// desired.
	switch clientType {
	case External:
		defaultClientType = External
	case Native:
		defaultClientType = Native
	}
}

func NewClient(user string, host string, port int, auth *Auth) (Client, error) {
	sshBinaryPath, err := exec.LookPath("ssh")
	if err != nil {
		log.Debug("SSH binary not found, using native Go implementation")
		client, err := NewNativeClient(user, host, port, auth)
		log.Debug(client)
		return client, err
	}

	if defaultClientType == Native {
		log.Debug("Using SSH client type: native")
		client, err := NewNativeClient(user, host, port, auth)
		log.Debug(client)
		return client, err
	}

	log.Debug("Using SSH client type: external")
	client, err := NewExternalClient(sshBinaryPath, user, host, port, auth)
	log.Debug(client)
	return client, err
}

func NewNativeClient(user, host string, port int, auth *Auth) (Client, error) {
	config, err := NewNativeConfig(user, auth)
	if err != nil {
		return nil, fmt.Errorf("Error getting config for native Go SSH: %s", err)
	}

	return &NativeClient{
		Config:   config,
		Hostname: host,
		Port:     port,
	}, nil
}

func NewNativeConfig(user string, auth *Auth) (ssh.ClientConfig, error) {
	var (
		authMethods []ssh.AuthMethod
	)

	for _, k := range auth.Keys {
		key, err := ioutil.ReadFile(k)
		if err != nil {
			return ssh.ClientConfig{}, err
		}

		privateKey, err := ssh.ParsePrivateKey(key)
		if err != nil {
			return ssh.ClientConfig{}, err
		}

		authMethods = append(authMethods, ssh.PublicKeys(privateKey))
	}

	for _, p := range auth.Passwords {
		authMethods = append(authMethods, ssh.Password(p))
	}

	return ssh.ClientConfig{
		User:            user,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}, nil
}

func (client *NativeClient) dialSuccess() bool {
	conn, err := ssh.Dial("tcp", net.JoinHostPort(client.Hostname, strconv.Itoa(client.Port)), &client.Config)
	if err != nil {
		log.Debugf("Error dialing TCP: %s", err)
		return false
	}
	closeConn(conn)
	return true
}

func (client *NativeClient) session(command string) (*ssh.Client, *ssh.Session, error) {
	if err := mcnutils.WaitFor(client.dialSuccess); err != nil {
		return nil, nil, fmt.Errorf("Error attempting SSH client dial: %s", err)
	}

	conn, err := ssh.Dial("tcp", net.JoinHostPort(client.Hostname, strconv.Itoa(client.Port)), &client.Config)
	if err != nil {
		return nil, nil, fmt.Errorf("Mysterious error dialing TCP for SSH (we already succeeded at least once) : %s", err)
	}
	session, err := conn.NewSession()

	return conn, session, err
}

func (client *NativeClient) Output(command string) (string, error) {
	conn, session, err := client.session(command)
	if err != nil {
		return "", nil
	}
	defer closeConn(conn)
	defer session.Close()

	output, err := session.CombinedOutput(command)

	return string(output), err
}

func (client *NativeClient) OutputWithPty(command string) (string, error) {
	conn, session, err := client.session(command)
	if err != nil {
		return "", nil
	}
	defer closeConn(conn)
	defer session.Close()

	fd := int(os.Stdout.Fd())

	termWidth, termHeight, err := terminal.GetSize(fd)
	if err != nil {
		return "", err
	}

	modes := ssh.TerminalModes{
		ssh.ECHO:          0,
		ssh.TTY_OP_ISPEED: 14400,
		ssh.TTY_OP_OSPEED: 14400,
	}

	// request tty -- fixes error with hosts that use
	// "Defaults requiretty" in /etc/sudoers - I'm looking at you RedHat
	if err := session.RequestPty("xterm", termHeight, termWidth, modes); err != nil {
		return "", err
	}

	output, err := session.CombinedOutput(command)

	return string(output), err
}

func (client *NativeClient) Start(command string) (io.ReadCloser, io.ReadCloser, error) {
	conn, session, err := client.session(command)
	if err != nil {
		return nil, nil, err
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}
	stderr, err := session.StderrPipe()
	if err != nil {
		return nil, nil, err
	}
	if err := session.Start(command); err != nil {
		return nil, nil, err
	}

	client.openClient = conn
	client.openSession = session
	return ioutil.NopCloser(stdout), ioutil.NopCloser(stderr), nil
}

func (client *NativeClient) Wait() error {
	err := client.openSession.Wait()
	if err != nil {
		return err
	}

	_ = client.openSession.Close()

	err = client.openClient.Close()
	if err != nil {
		return err
	}

	client.openSession = nil
	client.openClient = nil
	return nil
}

func (client *NativeClient) Shell(args ...string) error {
	var (
		termWidth, termHeight int
	)
	conn, err := ssh.Dial("tcp", net.JoinHostPort(client.Hostname, strconv.Itoa(client.Port)), &client.Config)
	if err != nil {
		return err
	}
	defer closeConn(conn)

	session, err := conn.NewSession()
	if err != nil {
		return err
	}

	defer session.Close()

	session.Stdout = os.Stdout
	session.Stderr = os.Stderr
	session.Stdin = os.Stdin

	modes := ssh.TerminalModes{
		ssh.ECHO: 1,
	}

	fd := os.Stdin.Fd()

	if term.IsTerminal(fd) {
		oldState, err := term.MakeRaw(fd)
		if err != nil {
			return err
		}

		defer term.RestoreTerminal(fd, oldState)

		winsize, err := term.GetWinsize(fd)
		if err != nil {
			termWidth = 80
			termHeight = 24
		} else {
			termWidth = int(winsize.Width)
			termHeight = int(winsize.Height)
		}
	}

	if err := session.RequestPty("xterm", termHeight, termWidth, modes); err != nil {
		return err
	}

	if len(args) == 0 {
		if err := session.Shell(); err != nil {
			return err
		}
		session.Wait()
	} else {
		session.Run(strings.Join(args, " "))
	}

	return nil
}

func NewExternalClient(sshBinaryPath, user, host string, port int, auth *Auth) (*ExternalClient, error) {
	client := &ExternalClient{
		BinaryPath: sshBinaryPath,
	}

	args := append(baseSSHArgs, fmt.Sprintf("%s@%s", user, host))

	// If no identities are explicitly provided, also look at the identities
	// offered by ssh-agent
	if len(auth.Keys) > 0 {
		args = append(args, "-o", "IdentitiesOnly=yes")
	}

	// Specify which private keys to use to authorize the SSH request.
	for _, privateKeyPath := range auth.Keys {
		if privateKeyPath != "" {
			// Check each private key before use it
			fi, err := os.Stat(privateKeyPath)
			if err != nil {
				// Abort if key not accessible
				return nil, err
			}
			if runtime.GOOS != "windows" {
				mode := fi.Mode()
				log.Debugf("Using SSH private key: %s (%s)", privateKeyPath, mode)
				// Private key file should have strict permissions
				perm := mode.Perm()
				if perm&0400 == 0 {
					return nil, fmt.Errorf("'%s' is not readable", privateKeyPath)
				}
				if perm&0077 != 0 {
					return nil, fmt.Errorf("permissions %#o for '%s' are too open", perm, privateKeyPath)
				}
			}
			args = append(args, "-i", privateKeyPath)
		}
	}

	// Set which port to use for SSH.
	args = append(args, "-p", fmt.Sprintf("%d", port))

	client.BaseArgs = args

	return client, nil
}

func getSSHCmd(binaryPath string, args ...string) *exec.Cmd {
	return exec.Command(binaryPath, args...)
}

func (client *ExternalClient) Output(command string) (string, error) {
	args := append(client.BaseArgs, command)
	cmd := getSSHCmd(client.BinaryPath, args...)
	output, err := cmd.CombinedOutput()
	return string(output), err
}

func (client *ExternalClient) Shell(args ...string) error {
	args = append(client.BaseArgs, args...)
	cmd := getSSHCmd(client.BinaryPath, args...)

	log.Debug(cmd)

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

	return cmd.Run()
}

func (client *ExternalClient) Start(command string) (io.ReadCloser, io.ReadCloser, error) {
	args := append(client.BaseArgs, command)
	cmd := getSSHCmd(client.BinaryPath, args...)

	log.Debug(cmd)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		if closeErr := stdout.Close(); closeErr != nil {
			return nil, nil, fmt.Errorf("%s, %s", err, closeErr)
		}
		return nil, nil, err
	}
	if err := cmd.Start(); err != nil {
		stdOutCloseErr := stdout.Close()
		stdErrCloseErr := stderr.Close()
		if stdOutCloseErr != nil || stdErrCloseErr != nil {
			return nil, nil, fmt.Errorf("%s, %s, %s",
				err, stdOutCloseErr, stdErrCloseErr)
		}
		return nil, nil, err
	}

	client.cmd = cmd
	return stdout, stderr, nil
}

func (client *ExternalClient) Wait() error {
	err := client.cmd.Wait()
	client.cmd = nil
	return err
}

func closeConn(c io.Closer) {
	err := c.Close()
	if err != nil {
		log.Debugf("Error closing SSH Client: %s", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/go-martini/martini"
	"github.com/johnernaut/webhog/webhog"
	"github.com/martini-contrib/binding"
	_ "github.com/martini-contrib/gzip"
	"github.com/martini-contrib/render"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"runtime"
)

type Url struct {
	Url  string `json:"url"`
	UUID string `json:"uuid"`
}

func main() {
	// All the parallelism are belong to us!
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Load configuration file.
	webhog.LoadConfig()

	// Load DB instance.
	webhog.LoadDB()

	// Start the server.
	m := martini.Classic()

	m.Use(render.Renderer())
	// m.Use(gzip.All())
	m.Use(martini.Recovery())
	m.Use(martini.Static("public"))

	m.Group("/api", func(r martini.Router) {
		r.Post("/scrape", binding.Json(Url{}), func(url Url, r render.Render) {
			log.Println("SUP")
			entity, err := webhog.NewScraper(url.Url)
			log.Println(entity)
			if err != nil {
				r.JSON(400, map[string]interface{}{"errors": err.Error()})
			} else {
				r.JSON(200, entity)
			}
		})

		r.Get("/entity/:uuid", func(params martini.Params, r render.Render) {
			entity := new(webhog.Entity)
			err := entity.Find(bson.M{"uuid": params["uuid"]})

			if err != nil {
				r.JSON(200, map[string]interface{}{"errors": "Entity not found."})
			} else {
				r.JSON(200, entity)
			}
		})
	})

	m.Run()
}

func KeyRequired() martini.Handler {
	return func(res http.ResponseWriter, req *http.Request, r render.Render) {
		if req.Header.Get("X-API-KEY") != webhog.Config.ApiKey {
			r.JSON(401, map[string]interface{}{"error": "Invalid API key."})
		}
	}
}
<commit_msg>rm unique<commit_after>package main

import (
	"github.com/go-martini/martini"
	"github.com/johnernaut/webhog/webhog"
	"github.com/martini-contrib/binding"
	_ "github.com/martini-contrib/gzip"
	"github.com/martini-contrib/render"
	"labix.org/v2/mgo/bson"
	"log"
	"net/http"
	"runtime"
)

type Url struct {
	Url  string `json:"url"`
	UUID string `json:"uuid"`
}

func main() {
	// All the parallelism are belong to us!
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Load configuration file.
	webhog.LoadConfig()

	// Load DB instance.
	webhog.LoadDB()

	// Start the server.
	m := martini.Classic()

	m.Use(render.Renderer())
	// m.Use(gzip.All())
	m.Use(martini.Recovery())
	m.Use(martini.Static("public"))

	m.Group("/api", func(r martini.Router) {
		r.Post("/scrape", binding.Json(Url{}), func(url Url, r render.Render) {
			entity, err := webhog.NewScraper(url.Url)
			if err != nil {
				r.JSON(400, map[string]interface{}{"errors": err.Error()})
			} else {
				r.JSON(200, entity)
			}
		})

		r.Get("/entity/:uuid", func(params martini.Params, r render.Render) {
			entity := new(webhog.Entity)
			err := entity.Find(bson.M{"uuid": params["uuid"]})

			if err != nil {
				r.JSON(200, map[string]interface{}{"errors": "Entity not found."})
			} else {
				r.JSON(200, entity)
			}
		})
	})

	m.Run()
}

func KeyRequired() martini.Handler {
	return func(res http.ResponseWriter, req *http.Request, r render.Render) {
		if req.Header.Get("X-API-KEY") != webhog.Config.ApiKey {
			r.JSON(401, map[string]interface{}{"error": "Invalid API key."})
		}
	}
}
<|endoftext|>


<commit_before>package cdl

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"unicode"
)

// type Template is a user-provided uncompiled template.
//
// See the overview for how these work.
type Template map[string]interface{}

// type CompiledTemplate is a compiled template.
//
// It is opaque to the user in operations.
type CompiledTemplate struct {
	s map[string]interface{}
}

type options map[string]interface{}

type optrange struct {
	Min int
	Max int
}

type array struct {
	name string
	r    optrange
}

type requirement struct {
	mandatory bool
	array     bool
	r         optrange
}

type CdlError struct {
	Type          int
	Supplementary string
	Context       []string
}

const (
	_                              = iota
	ErrInternal                    = iota
	ErrMissingRoot                 = iota
	ErrBadOptionValue              = iota
	ErrBadRangeOptionModifier      = iota
	ErrBadRangeOptionModifierValue = iota
	ErrBadOptionModifier           = iota
	ErrBadKey                      = iota
	ErrBadValue                    = iota
	ErrUnknownKey                  = iota
	ErrExpectedMap                 = iota
	ErrExpectedArray               = iota
	ErrOutOfRange                  = iota
	ErrBadType                     = iota
	ErrMissingMandatory            = iota
)

var ErrorMap map[int]string = map[int]string{
	ErrInternal:                    "Internal error",
	ErrMissingRoot:                 "No root key in template",
	ErrBadOptionValue:              "Bad option value",
	ErrBadRangeOptionModifier:      "Bad range option modifer",
	ErrBadRangeOptionModifierValue: "Bad range option modifier value",
	ErrBadOptionModifier:           "Bad option modifier",
	ErrBadKey:                      "Bad key",
	ErrBadValue:                    "Bad value",
	ErrUnknownKey:                  "Unknown key",
	ErrExpectedMap:                 "Expected map",
	ErrExpectedArray:               "Expected array",
	ErrOutOfRange:                  "Number of array items outside permissible range",
	ErrBadType:                     "Bad type",
	ErrMissingMandatory:            "Missing mandatory key",
}

// type ValidatorFunc allows user specified validation functions to be passed to cdl.
type ValidatorFunc func(obj interface{}) (err *CdlError)

// func Error implements the Error() function of the error interface.
//
// An error string is returned in context.
func (e CdlError) Error() string {
	main := ErrorMap[e.Type]
	if e.Supplementary != "" {
		main = fmt.Sprintf("%s; %s", main, e.Supplementary)
	}
	if len(e.Context) == 0 {
		return fmt.Sprintf("%s (code %d)", main, e.Type)
	} else {
		return fmt.Sprintf("%s (code %d) near %s", main, e.Type, strings.Join(e.Context, " at "))
	}
}

// func NewError returns a new CdlError of a given type.
//
// The type should be a type starting with `Err` in the constants section.
func NewError(t int) *CdlError {
	return &CdlError{Type: t}
}

// func NewErrorContext creates a new CdlError with the specified context string.
//
// The type should be a type starting with `Err` in the constants section.
func NewErrorContext(t int, c string) *CdlError {
	return (&CdlError{Type: t}).AddContext(c)
}

// func NewErrorContext creates a new CdlError with the specified context string.
//
// The type should be a type starting with `Err` in the constants section.
// The context string will be quoted.
func NewErrorContextQuoted(t int, c string) *CdlError {
	return (&CdlError{Type: t}).AddContextQuoted(c)
}

// func AddContext adds the specified context to an existing cdl error.
func (e *CdlError) AddContext(c string) *CdlError {
	e.Context = append(e.Context, c)
	return e
}

// func AddContextQuoted adds the specified context to an existing cdl error.
//
// The context will be quoted.
func (e *CdlError) AddContextQuoted(c string) *CdlError {
	return e.AddContext(fmt.Sprintf("'%s'", c))
}
// func SetSupplementary adds the specified supplementary data to an existing cdl error.
func (e *CdlError) SetSupplementary(s string) *CdlError {
	e.Supplementary = s
	return e
}

func (r *optrange) describeError(value int) string {
	min := r.Min
	if min < 0 {
		min = 0
	}
	if r.Max < 0 {
		return fmt.Sprintf("got %d, expecting at least %d", value, min)
	} else {
		return fmt.Sprintf("got %d, expecting between %d and %d", value, min, r.Max)
	}
}

func (r *optrange) contains(value int) bool {
	return (value >= r.Min || r.Min == -1) && (value <= r.Max || r.Max == -1)
}

func makeoptions(optString string) (*options, *CdlError) {
	opts := make(options)
	spaceOrBar := func(r rune) bool {
		return unicode.IsSpace(r) || (r == '|')
	}
	for _, o := range strings.FieldsFunc(optString, spaceOrBar) {
		s := regexp.MustCompile("^(\\w+)(.*)$").FindStringSubmatch(o)
		if len(s) < 3 || s[1] == "" {
			return nil, NewErrorContextQuoted(ErrBadOptionValue, o)
		}
		req := requirement{mandatory: true, array: false, r: optrange{-1, -1}}
		if s[2] != "" {
			optslice := regexp.MustCompile("[*+!?]|\\{\\d+,\\d*\\}").FindAllStringSubmatch(s[2], -1)
			if len(optslice) == 0 {
				return nil, NewErrorContextQuoted(ErrBadOptionModifier, o)
			}
			for _, c := range optslice {
				if len(c) != 1 {
					return nil, NewErrorContextQuoted(ErrBadOptionModifier, o)
				}
				switch {
				case c[0] == "?":
					req.mandatory = false
				case c[0] == "!":
					req.mandatory = true
				case c[0] == "+":
					req.r = optrange{1, -1}
					req.array = true
				case c[0] == "*":
					req.array = true
					req.r = optrange{0, -1}
				case strings.HasPrefix(c[0], "{"):
					minMax := regexp.MustCompile("^\\{(\\d+),(\\d*)\\}$").FindStringSubmatch(c[0])
					if len(minMax) != 3 {
						return nil, NewErrorContextQuoted(ErrBadRangeOptionModifier, o)
					}
					min, err1 := strconv.Atoi(minMax[1])
					if err1 != nil {
						return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, o)
					}
					max := -1
					if minMax[2] != "" {
						max, err2 := strconv.Atoi(minMax[2])
						if (err2 != nil) || (min > max) {
							return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, o)
						}
					}
					req.array = true
					req.r = optrange{min, max}
				default:
					return nil, NewErrorContextQuoted(ErrBadOptionModifier, o)
				}
			}
		}
		opts[s[1]] = req
	}

	return &opts, nil
}

func newCompiledTemplate() *CompiledTemplate {
	return &CompiledTemplate{s: make(map[string]interface{})}
}

// func Compile compiles a specified cdl template.
func Compile(t Template) (*CompiledTemplate, error) {
	ct := newCompiledTemplate()
	for k, v := range t {
		if match, err := regexp.MatchString("^(/|(\\w+))?$", k); !match || err != nil {
			return nil, NewErrorContextQuoted(ErrBadKey, k)
		}
		switch t := v.(type) {
		case string:
			if t == "" {
				t = "/"
			}
			switch {
			case strings.HasPrefix(t, "{}"):
				if o, err := makeoptions(strings.TrimPrefix(t, "{}")); err != nil {
					return nil, err.AddContextQuoted(k)
				} else {
					ct.s[k] = o
				}
			case strings.HasPrefix(t, "[]"):
				arr := strings.TrimPrefix(t, "[]")
				rng := optrange{-1, -1}
				minMax := regexp.MustCompile("^(\\w+)(\\{(\\d+),(\\d*)\\})?$").FindStringSubmatch(arr)
				if len(minMax) != 5 {
					return nil, NewErrorContextQuoted(ErrBadRangeOptionModifier, arr)
				}
				if minMax[3] != "" {
					min, err1 := strconv.Atoi(minMax[3])
					if err1 != nil {
						return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, arr)
					}
					max := -1
					if minMax[4] != "" {
						var err2 error
						max, err2 = strconv.Atoi(minMax[4])
						if (err2 != nil) || (min > max) {
							return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, arr)
						}
					}
					rng = optrange{min, max}
				}
				ct.s[k] = &array{name: minMax[1], r: rng}
			default:
				ct.s[k] = t
			}
		case ValidatorFunc:
			ct.s[k] = t
		case func(interface{}) *CdlError: // in case they didn't cast it
			ct.s[k] = ValidatorFunc(t)
		default:
			return nil, NewErrorContextQuoted(ErrBadValue, fmt.Sprintf("%T", t)).AddContextQuoted(k)
		}
	}
	for _, v := range ct.s {
		switch t := v.(type) {
		case *options:
			for optk, _ := range *t {
				if _, ok := ct.s[optk]; !ok {
					ct.s[optk] = 0 // autodiscovered
				}
			}
		}
	}
	if _, ok := ct.s["/"]; !ok {
		return nil, NewError(ErrMissingRoot)
	}
	return ct, nil
}

func (ct *CompiledTemplate) validateRange(o interface{}, pos string, r optrange) *CdlError {
	slice, ok := o.([]interface{})
	if !ok {
		return NewError(ErrExpectedArray)
	}
	if !r.contains(len(slice)) {
		return NewError(ErrOutOfRange).SetSupplementary(r.describeError(len(slice)))
	}
	for i, v := range slice {
		if err := ct.validateItem(v, pos); err != nil {
			return err.AddContext(fmt.Sprintf("index %d", i))
		}
	}
	return nil
}

func (ct *CompiledTemplate) validateMap(o interface{}, pos string, opts *options) *CdlError {
	m, ok := o.(map[string]interface{})
	if !ok {
		return NewError(ErrExpectedMap)
	}
	mand := make(map[string]bool)
	for k, v := range *opts {
		switch t := v.(type) {
		case requirement:
			if t.mandatory {
				mand[k] = true
			}
		}
	}
	for k, v := range m {
		if o, ok := (*opts)[k]; !ok {
			return NewErrorContextQuoted(ErrBadKey, k)
		} else {
			switch t := o.(type) {
			case requirement:
				if t.array {
					if err := ct.validateRange(v, k, t.r); err != nil {
						return err.AddContextQuoted(k)
					}
				} else {
					if err := ct.validateItem(v, k); err != nil {
						return err.AddContextQuoted(k)
					}
				}
				if t.mandatory {
					delete(mand, k)
				}
			}
		}
	}
	if len(mand) != 0 {
		missing := make([]string, len(mand))
		i := 0
		for k := range mand {
			missing[i] = fmt.Sprintf("'%s'", k)
			i++
		}
		return NewError(ErrMissingMandatory).SetSupplementary(fmt.Sprintf("missing %s", strings.Join(missing, ", ")))
	}
	return nil
}

func (ct *CompiledTemplate) validateItem(o interface{}, pos string) *CdlError {
	if val, ok := ct.s[pos]; !ok {
		return NewError(ErrUnknownKey)
	} else {
		switch t := val.(type) {
		case ValidatorFunc:
			return t(o)
		case *options:
			return ct.validateMap(o, pos, t)
		case *array:
			return ct.validateRange(o, t.name, t.r)
		case string:
			ok := false
			switch t {
			case "number":
				switch o.(type) {
				case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64:
					ok = true
				}
			case "integer":
				switch o.(type) {
				case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
					ok = true
				}
			default:
				otype := fmt.Sprintf("%T", o)
				if otype == t {
					ok = true
				}
			}
			if !ok {
				return NewError(ErrBadType).SetSupplementary(fmt.Sprintf("got %T expected %s", o, t))
			}
		case int:
			// autodiscovered
		default:
			return NewError(ErrInternal).SetSupplementary(fmt.Sprintf("type is neither validator func nor options: %T", val))
		}
	}
	return nil
}

// func Validate validates an object against a cdl template.
func (ct *CompiledTemplate) Validate(o interface{}) error {
	if err := ct.validateItem(o, "/"); err != nil {
		return err
	}
	return nil
}
<commit_msg>Add MustCompile<commit_after>package cdl

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"unicode"
)

// type Template is a user-provided uncompiled template.
//
// See the overview for how these work.
type Template map[string]interface{}

// type CompiledTemplate is a compiled template.
//
// It is opaque to the user in operations.
type CompiledTemplate struct {
	s map[string]interface{}
}

type options map[string]interface{}

type optrange struct {
	Min int
	Max int
}

type array struct {
	name string
	r    optrange
}

type requirement struct {
	mandatory bool
	array     bool
	r         optrange
}

type CdlError struct {
	Type          int
	Supplementary string
	Context       []string
}

const (
	_                              = iota
	ErrInternal                    = iota
	ErrMissingRoot                 = iota
	ErrBadOptionValue              = iota
	ErrBadRangeOptionModifier      = iota
	ErrBadRangeOptionModifierValue = iota
	ErrBadOptionModifier           = iota
	ErrBadKey                      = iota
	ErrBadValue                    = iota
	ErrUnknownKey                  = iota
	ErrExpectedMap                 = iota
	ErrExpectedArray               = iota
	ErrOutOfRange                  = iota
	ErrBadType                     = iota
	ErrMissingMandatory            = iota
)

var ErrorMap map[int]string = map[int]string{
	ErrInternal:                    "Internal error",
	ErrMissingRoot:                 "No root key in template",
	ErrBadOptionValue:              "Bad option value",
	ErrBadRangeOptionModifier:      "Bad range option modifer",
	ErrBadRangeOptionModifierValue: "Bad range option modifier value",
	ErrBadOptionModifier:           "Bad option modifier",
	ErrBadKey:                      "Bad key",
	ErrBadValue:                    "Bad value",
	ErrUnknownKey:                  "Unknown key",
	ErrExpectedMap:                 "Expected map",
	ErrExpectedArray:               "Expected array",
	ErrOutOfRange:                  "Number of array items outside permissible range",
	ErrBadType:                     "Bad type",
	ErrMissingMandatory:            "Missing mandatory key",
}

// type ValidatorFunc allows user specified validation functions to be passed to cdl.
type ValidatorFunc func(obj interface{}) (err *CdlError)

// func Error implements the Error() function of the error interface.
//
// An error string is returned in context.
func (e CdlError) Error() string {
	main := ErrorMap[e.Type]
	if e.Supplementary != "" {
		main = fmt.Sprintf("%s; %s", main, e.Supplementary)
	}
	if len(e.Context) == 0 {
		return fmt.Sprintf("%s (code %d)", main, e.Type)
	} else {
		return fmt.Sprintf("%s (code %d) near %s", main, e.Type, strings.Join(e.Context, " at "))
	}
}

// func NewError returns a new CdlError of a given type.
//
// The type should be a type starting with `Err` in the constants section.
func NewError(t int) *CdlError {
	return &CdlError{Type: t}
}

// func NewErrorContext creates a new CdlError with the specified context string.
//
// The type should be a type starting with `Err` in the constants section.
func NewErrorContext(t int, c string) *CdlError {
	return (&CdlError{Type: t}).AddContext(c)
}

// func NewErrorContext creates a new CdlError with the specified context string.
//
// The type should be a type starting with `Err` in the constants section.
// The context string will be quoted.
func NewErrorContextQuoted(t int, c string) *CdlError {
	return (&CdlError{Type: t}).AddContextQuoted(c)
}

// func AddContext adds the specified context to an existing cdl error.
func (e *CdlError) AddContext(c string) *CdlError {
	e.Context = append(e.Context, c)
	return e
}

// func AddContextQuoted adds the specified context to an existing cdl error.
//
// The context will be quoted.
func (e *CdlError) AddContextQuoted(c string) *CdlError {
	return e.AddContext(fmt.Sprintf("'%s'", c))
}

// func SetSupplementary adds the specified supplementary data to an existing cdl error.
func (e *CdlError) SetSupplementary(s string) *CdlError {
	e.Supplementary = s
	return e
}

func (r *optrange) describeError(value int) string {
	min := r.Min
	if min < 0 {
		min = 0
	}
	if r.Max < 0 {
		return fmt.Sprintf("got %d, expecting at least %d", value, min)
	} else {
		return fmt.Sprintf("got %d, expecting between %d and %d", value, min, r.Max)
	}
}

func (r *optrange) contains(value int) bool {
	return (value >= r.Min || r.Min == -1) && (value <= r.Max || r.Max == -1)
}

func makeoptions(optString string) (*options, *CdlError) {
	opts := make(options)
	spaceOrBar := func(r rune) bool {
		return unicode.IsSpace(r) || (r == '|')
	}
	for _, o := range strings.FieldsFunc(optString, spaceOrBar) {
		s := regexp.MustCompile("^(\\w+)(.*)$").FindStringSubmatch(o)
		if len(s) < 3 || s[1] == "" {
			return nil, NewErrorContextQuoted(ErrBadOptionValue, o)
		}
		req := requirement{mandatory: true, array: false, r: optrange{-1, -1}}
		if s[2] != "" {
			optslice := regexp.MustCompile("[*+!?]|\\{\\d+,\\d*\\}").FindAllStringSubmatch(s[2], -1)
			if len(optslice) == 0 {
				return nil, NewErrorContextQuoted(ErrBadOptionModifier, o)
			}
			for _, c := range optslice {
				if len(c) != 1 {
					return nil, NewErrorContextQuoted(ErrBadOptionModifier, o)
				}
				switch {
				case c[0] == "?":
					req.mandatory = false
				case c[0] == "!":
					req.mandatory = true
				case c[0] == "+":
					req.r = optrange{1, -1}
					req.array = true
				case c[0] == "*":
					req.array = true
					req.r = optrange{0, -1}
				case strings.HasPrefix(c[0], "{"):
					minMax := regexp.MustCompile("^\\{(\\d+),(\\d*)\\}$").FindStringSubmatch(c[0])
					if len(minMax) != 3 {
						return nil, NewErrorContextQuoted(ErrBadRangeOptionModifier, o)
					}
					min, err1 := strconv.Atoi(minMax[1])
					if err1 != nil {
						return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, o)
					}
					max := -1
					if minMax[2] != "" {
						max, err2 := strconv.Atoi(minMax[2])
						if (err2 != nil) || (min > max) {
							return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, o)
						}
					}
					req.array = true
					req.r = optrange{min, max}
				default:
					return nil, NewErrorContextQuoted(ErrBadOptionModifier, o)
				}
			}
		}
		opts[s[1]] = req
	}

	return &opts, nil
}

func newCompiledTemplate() *CompiledTemplate {
	return &CompiledTemplate{s: make(map[string]interface{})}
}

// func Compile compiles a specified cdl template.
func Compile(t Template) (*CompiledTemplate, error) {
	ct := newCompiledTemplate()
	for k, v := range t {
		if match, err := regexp.MatchString("^(/|(\\w+))?$", k); !match || err != nil {
			return nil, NewErrorContextQuoted(ErrBadKey, k)
		}
		switch t := v.(type) {
		case string:
			if t == "" {
				t = "/"
			}
			switch {
			case strings.HasPrefix(t, "{}"):
				if o, err := makeoptions(strings.TrimPrefix(t, "{}")); err != nil {
					return nil, err.AddContextQuoted(k)
				} else {
					ct.s[k] = o
				}
			case strings.HasPrefix(t, "[]"):
				arr := strings.TrimPrefix(t, "[]")
				rng := optrange{-1, -1}
				minMax := regexp.MustCompile("^(\\w+)(\\{(\\d+),(\\d*)\\})?$").FindStringSubmatch(arr)
				if len(minMax) != 5 {
					return nil, NewErrorContextQuoted(ErrBadRangeOptionModifier, arr)
				}
				if minMax[3] != "" {
					min, err1 := strconv.Atoi(minMax[3])
					if err1 != nil {
						return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, arr)
					}
					max := -1
					if minMax[4] != "" {
						var err2 error
						max, err2 = strconv.Atoi(minMax[4])
						if (err2 != nil) || (min > max) {
							return nil, NewErrorContextQuoted(ErrBadRangeOptionModifierValue, arr)
						}
					}
					rng = optrange{min, max}
				}
				ct.s[k] = &array{name: minMax[1], r: rng}
			default:
				ct.s[k] = t
			}
		case ValidatorFunc:
			ct.s[k] = t
		case func(interface{}) *CdlError: // in case they didn't cast it
			ct.s[k] = ValidatorFunc(t)
		default:
			return nil, NewErrorContextQuoted(ErrBadValue, fmt.Sprintf("%T", t)).AddContextQuoted(k)
		}
	}
	for _, v := range ct.s {
		switch t := v.(type) {
		case *options:
			for optk, _ := range *t {
				if _, ok := ct.s[optk]; !ok {
					ct.s[optk] = 0 // autodiscovered
				}
			}
		}
	}
	if _, ok := ct.s["/"]; !ok {
		return nil, NewError(ErrMissingRoot)
	}
	return ct, nil
}

// MustCompile is like Compile but panics if the expression cannot be parsed.
// It simplifies safe initialization of global variables holding compiled templates
func MustCompile(t Template) *CompiledTemplate {
	ct, error := Compile(t)
	if error != nil {
		panic(`cdl: Compile failed: ` + error.Error())
	}
	return ct
}

func (ct *CompiledTemplate) validateRange(o interface{}, pos string, r optrange) *CdlError {
	slice, ok := o.([]interface{})
	if !ok {
		return NewError(ErrExpectedArray)
	}
	if !r.contains(len(slice)) {
		return NewError(ErrOutOfRange).SetSupplementary(r.describeError(len(slice)))
	}
	for i, v := range slice {
		if err := ct.validateItem(v, pos); err != nil {
			return err.AddContext(fmt.Sprintf("index %d", i))
		}
	}
	return nil
}

func (ct *CompiledTemplate) validateMap(o interface{}, pos string, opts *options) *CdlError {
	m, ok := o.(map[string]interface{})
	if !ok {
		return NewError(ErrExpectedMap)
	}
	mand := make(map[string]bool)
	for k, v := range *opts {
		switch t := v.(type) {
		case requirement:
			if t.mandatory {
				mand[k] = true
			}
		}
	}
	for k, v := range m {
		if o, ok := (*opts)[k]; !ok {
			return NewErrorContextQuoted(ErrBadKey, k)
		} else {
			switch t := o.(type) {
			case requirement:
				if t.array {
					if err := ct.validateRange(v, k, t.r); err != nil {
						return err.AddContextQuoted(k)
					}
				} else {
					if err := ct.validateItem(v, k); err != nil {
						return err.AddContextQuoted(k)
					}
				}
				if t.mandatory {
					delete(mand, k)
				}
			}
		}
	}
	if len(mand) != 0 {
		missing := make([]string, len(mand))
		i := 0
		for k := range mand {
			missing[i] = fmt.Sprintf("'%s'", k)
			i++
		}
		return NewError(ErrMissingMandatory).SetSupplementary(fmt.Sprintf("missing %s", strings.Join(missing, ", ")))
	}
	return nil
}

func (ct *CompiledTemplate) validateItem(o interface{}, pos string) *CdlError {
	if val, ok := ct.s[pos]; !ok {
		return NewError(ErrUnknownKey)
	} else {
		switch t := val.(type) {
		case ValidatorFunc:
			return t(o)
		case *options:
			return ct.validateMap(o, pos, t)
		case *array:
			return ct.validateRange(o, t.name, t.r)
		case string:
			ok := false
			switch t {
			case "number":
				switch o.(type) {
				case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64:
					ok = true
				}
			case "integer":
				switch o.(type) {
				case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
					ok = true
				}
			default:
				otype := fmt.Sprintf("%T", o)
				if otype == t {
					ok = true
				}
			}
			if !ok {
				return NewError(ErrBadType).SetSupplementary(fmt.Sprintf("got %T expected %s", o, t))
			}
		case int:
			// autodiscovered
		default:
			return NewError(ErrInternal).SetSupplementary(fmt.Sprintf("type is neither validator func nor options: %T", val))
		}
	}
	return nil
}

// func Validate validates an object against a cdl template.
func (ct *CompiledTemplate) Validate(o interface{}) error {
	if err := ct.validateItem(o, "/"); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package cli

import (
	"bytes"
	"fmt"
	"os"
	"reflect"
	"strings"

	"github.com/labstack/gommon/color"
)

// Run runs a single command app
func Run(argv interface{}, fn CommandFunc) {
	err := (&Command{
		Name: os.Args[0],
		Argv: func() interface{} { return argv },
		Fn:   fn,
	}).Run(os.Args[1:])
	if err != nil {
		fmt.Println(err)
	}
}

// Root registers forest for root and return root
func Root(root *Command, forest ...*CommandTree) *Command {
	root.RegisterTree(forest...)
	return root
}

// Tree creates a CommandTree
func Tree(cmd *Command, forest ...*CommandTree) *CommandTree {
	return &CommandTree{
		command: cmd,
		forest:  forest,
	}
}

//-----------------------------
// Implements parse and others
//-----------------------------

func parseArgv(args []string, argv interface{}, clr color.Color) *flagSet {
	var (
		typ     = reflect.TypeOf(argv)
		val     = reflect.ValueOf(argv)
		flagSet = newFlagSet()
	)
	switch typ.Kind() {
	case reflect.Ptr:
		if reflect.Indirect(val).Type().Kind() != reflect.Struct {
			flagSet.err = errNotPointToStruct
			return flagSet
		}
		parse(args, typ, val, flagSet, clr)
		return flagSet
	default:
		flagSet.err = errNotAPointer
		return flagSet
	}
}

func usage(v interface{}, clr color.Color) string {
	var (
		typ     = reflect.TypeOf(v)
		val     = reflect.ValueOf(v)
		flagSet = newFlagSet()
	)
	if typ.Kind() == reflect.Ptr {
		if reflect.Indirect(val).Type().Kind() == reflect.Struct {
			initFlagSet(typ, val, flagSet, clr)
			if flagSet.err != nil {
				return ""
			}
			return flagSlice(flagSet.flags).String(clr)
		}
	}
	return ""
}

func initFlagSet(typ reflect.Type, val reflect.Value, flagSet *flagSet, clr color.Color) {
	var (
		tm       = typ.Elem()
		vm       = val.Elem()
		fieldNum = vm.NumField()
	)
	for i := 0; i < fieldNum; i++ {
		tfield := tm.Field(i)
		vfield := vm.Field(i)
		tag, isEmpty := parseTag(tfield.Name, tfield.Tag)
		if tag == nil {
			continue
		}
		// if `cli` tag is empty and the field is a struct
		if isEmpty && vfield.Kind() == reflect.Struct {
			subObj := vfield.Addr().Interface()
			initFlagSet(reflect.TypeOf(subObj), reflect.ValueOf(subObj), flagSet, clr)
			if flagSet.err != nil {
				return
			}
			continue
		}
		fl, err := newFlag(tfield, vfield, tag, clr)
		if flagSet.err = err; err != nil {
			return
		}
		// Ignored flag
		if fl == nil {
			continue
		}
		flagSet.flags = append(flagSet.flags, fl)
		value := ""
		if fl.assigned {
			value = fmt.Sprintf("%v", vfield.Interface())
		}

		names := append(fl.tag.shortNames, fl.tag.longNames...)
		for i, name := range names {
			if _, ok := flagSet.flagMap[name]; ok {
				flagSet.err = fmt.Errorf("flag %s repeat", clr.Bold(name))
				return
			}
			flagSet.flagMap[name] = fl
			if fl.assigned && i == 0 {
				flagSet.values[name] = []string{value}
			}
		}
	}
}

func parse(args []string, typ reflect.Type, val reflect.Value, flagSet *flagSet, clr color.Color) {
	initFlagSet(typ, val, flagSet, clr)
	if flagSet.err != nil {
		return
	}

	size := len(args)
	for i := 0; i < size; i++ {
		arg := args[i]
		if !strings.HasPrefix(arg, dashOne) {
			continue
		}
		values := []string{}
		for j := i + 1; j < size; j++ {
			if strings.HasPrefix(args[j], dashOne) {
				break
			}
			values = append(values, args[j])
		}
		i += len(values)

		strs := strings.Split(arg, "=")
		if strs == nil || len(strs) == 0 {
			continue
		}

		arg = strs[0]
		fl, ok := flagSet.flagMap[arg]
		if !ok {
			// If has prefix `--`
			if strings.HasPrefix(arg, dashTwo) {
				flagSet.err = fmt.Errorf("undefined flag %s", clr.Bold(arg))
				return
			}
			// Else find arg char by char
			chars := []byte(strings.TrimPrefix(arg, dashOne))
			for _, c := range chars {
				tmp := dashOne + string([]byte{c})
				fl, ok := flagSet.flagMap[tmp]
				if !ok {
					flagSet.err = fmt.Errorf("undefined flag %s", clr.Bold(tmp))
					return
				}

				if flagSet.err = fl.set(tmp, "", clr); flagSet.err != nil {
					return
				}
				if fl.err == nil {
					flagSet.values[tmp] = []string{fmt.Sprintf("%v", fl.v.Interface())}
				}

			}
			continue
		}

		values = append(strs[1:], values...)
		if len(values) == 0 {
			flagSet.err = fl.set(arg, "", clr)
		} else if len(values) == 1 {
			flagSet.err = fl.set(arg, values[0], clr)
		} else {
			flagSet.err = fmt.Errorf("too many(%d) value for flag %s", len(values), clr.Bold(arg))
		}
		if flagSet.err != nil {
			return
		}
		if fl.err == nil {
			flagSet.values[arg] = []string{fmt.Sprintf("%v", fl.v.Interface())}
		} else if fl.assigned {
			flagSet.err = fmt.Errorf("assigned argument %s invalid: %v", clr.Bold(fl.name()), fl.err)
			return
		}
	}

	buff := bytes.NewBufferString("")
	for _, fl := range flagSet.flags {
		if !fl.assigned && fl.tag.required {
			if buff.Len() > 0 {
				buff.WriteByte('\n')
			}
			fmt.Fprintf(buff, "required argument %s missing", clr.Bold(fl.name()))
		}
	}
	if buff.Len() > 0 {
		for _, fl := range flagSet.flags {
			if fl.tag.isHelp && fl.getBool() {
				flagSet.dontValidate = true
				break
			}
		}
		if !flagSet.dontValidate {
			flagSet.err = fmt.Errorf(buff.String())
		}
	}
}
<commit_msg>fix cli flag error prevent<commit_after>package cli

import (
	"bytes"
	"fmt"
	"os"
	"reflect"
	"strings"

	"github.com/labstack/gommon/color"
)

// Run runs a single command app
func Run(argv interface{}, fn CommandFunc) {
	err := (&Command{
		Name: os.Args[0],
		Argv: func() interface{} { return argv },
		Fn:   fn,
	}).Run(os.Args[1:])
	if err != nil {
		fmt.Println(err)
	}
}

// Root registers forest for root and return root
func Root(root *Command, forest ...*CommandTree) *Command {
	root.RegisterTree(forest...)
	return root
}

// Tree creates a CommandTree
func Tree(cmd *Command, forest ...*CommandTree) *CommandTree {
	return &CommandTree{
		command: cmd,
		forest:  forest,
	}
}

//-----------------------------
// Implements parse and others
//-----------------------------

func parseArgv(args []string, argv interface{}, clr color.Color) *flagSet {
	var (
		typ     = reflect.TypeOf(argv)
		val     = reflect.ValueOf(argv)
		flagSet = newFlagSet()
	)
	switch typ.Kind() {
	case reflect.Ptr:
		if reflect.Indirect(val).Type().Kind() != reflect.Struct {
			flagSet.err = errNotPointToStruct
			return flagSet
		}
		parse(args, typ, val, flagSet, clr)
		return flagSet
	default:
		flagSet.err = errNotAPointer
		return flagSet
	}
}

func usage(v interface{}, clr color.Color) string {
	var (
		typ     = reflect.TypeOf(v)
		val     = reflect.ValueOf(v)
		flagSet = newFlagSet()
	)
	if typ.Kind() == reflect.Ptr {
		if reflect.Indirect(val).Type().Kind() == reflect.Struct {
			initFlagSet(typ, val, flagSet, clr)
			if flagSet.err != nil {
				return ""
			}
			return flagSlice(flagSet.flags).String(clr)
		}
	}
	return ""
}

func initFlagSet(typ reflect.Type, val reflect.Value, flagSet *flagSet, clr color.Color) {
	var (
		tm       = typ.Elem()
		vm       = val.Elem()
		fieldNum = vm.NumField()
	)
	for i := 0; i < fieldNum; i++ {
		tfield := tm.Field(i)
		vfield := vm.Field(i)
		tag, isEmpty := parseTag(tfield.Name, tfield.Tag)
		if tag == nil {
			continue
		}
		// if `cli` tag is empty and the field is a struct
		if isEmpty && vfield.Kind() == reflect.Struct {
			subObj := vfield.Addr().Interface()
			initFlagSet(reflect.TypeOf(subObj), reflect.ValueOf(subObj), flagSet, clr)
			if flagSet.err != nil {
				return
			}
			continue
		}
		fl, err := newFlag(tfield, vfield, tag, clr)
		if flagSet.err = err; err != nil {
			return
		}
		// Ignored flag
		if fl == nil {
			continue
		}
		flagSet.flags = append(flagSet.flags, fl)
		value := ""
		if fl.assigned {
			value = fmt.Sprintf("%v", vfield.Interface())
		}

		names := append(fl.tag.shortNames, fl.tag.longNames...)
		for i, name := range names {
			if _, ok := flagSet.flagMap[name]; ok {
				flagSet.err = fmt.Errorf("flag %s repeat", clr.Bold(name))
				return
			}
			flagSet.flagMap[name] = fl
			if fl.assigned && i == 0 {
				flagSet.values[name] = []string{value}
			}
		}
	}
}

func parse(args []string, typ reflect.Type, val reflect.Value, flagSet *flagSet, clr color.Color) {
	initFlagSet(typ, val, flagSet, clr)
	if flagSet.err != nil {
		return
	}

	size := len(args)
	for i := 0; i < size; i++ {
		arg := args[i]
		if !strings.HasPrefix(arg, dashOne) {
			continue
		}
		values := []string{}
		for j := i + 1; j < size; j++ {
			if strings.HasPrefix(args[j], dashOne) {
				break
			}
			values = append(values, args[j])
		}
		i += len(values)

		strs := strings.Split(arg, "=")
		if strs == nil || len(strs) == 0 {
			continue
		}

		arg = strs[0]
		fl, ok := flagSet.flagMap[arg]
		if !ok {
			// If has prefix `--`
			if strings.HasPrefix(arg, dashTwo) {
				flagSet.err = fmt.Errorf("undefined flag %s", clr.Bold(arg))
				return
			}
			// Else find arg char by char
			chars := []byte(strings.TrimPrefix(arg, dashOne))
			for _, c := range chars {
				tmp := dashOne + string([]byte{c})
				fl, ok := flagSet.flagMap[tmp]
				if !ok {
					flagSet.err = fmt.Errorf("undefined flag %s", clr.Bold(tmp))
					return
				}

				if flagSet.err = fl.set(tmp, "", clr); flagSet.err != nil {
					return
				}
				if fl.err == nil {
					flagSet.values[tmp] = []string{fmt.Sprintf("%v", fl.v.Interface())}
				}

			}
			continue
		}

		values = append(strs[1:], values...)
		if len(values) == 0 {
			flagSet.err = fl.set(arg, "", clr)
		} else if len(values) == 1 {
			flagSet.err = fl.set(arg, values[0], clr)
		} else {
			flagSet.err = fmt.Errorf("too many(%d) value for flag %s", len(values), clr.Bold(arg))
		}
		if flagSet.err != nil {
			return
		}
		if fl.err == nil {
			flagSet.values[arg] = []string{fmt.Sprintf("%v", fl.v.Interface())}
		} else if fl.assigned {
			flagSet.err = fmt.Errorf("assigned argument %s invalid: %v", clr.Bold(fl.name()), fl.err)
			return
		}
	}

	buff := bytes.NewBufferString("")
	for _, fl := range flagSet.flags {
		if !fl.assigned && fl.tag.required {
			if buff.Len() > 0 {
				buff.WriteByte('\n')
			}
			fmt.Fprintf(buff, "required argument %s missing", clr.Bold(fl.name()))
		}
	}
	for _, fl := range flagSet.flags {
		if fl.tag.isHelp && fl.getBool() {
			flagSet.dontValidate = true
			break
		}
	}
	if buff.Len() > 0 && !flagSet.dontValidate {
		flagSet.err = fmt.Errorf(buff.String())
	}
}
<|endoftext|>