<commit_before>package generator

import (
	"strings"
	"text/template"

	_gen "github.com/golang/protobuf/protoc-gen-go/generator"
)

type printerProxy struct {
	printer *Printer
}

type handlerParams struct {
	Service        string
	Method         string
	Query          string
	Request        string
	Response       string
	RespMethodCall string
	ZeroResponse   bool
	Before         bool
	After          bool
}

func OneOrZero(hp handlerParams) string {
	if hp.ZeroResponse {
		return strings.Join([]string{`
err := result.Zero()
res := &`, hp.Response, `{}
        `}, "")
	}
	return "res, err := result.One()." + hp.RespMethodCall + "()"
}

func (h *printerProxy) Write(data []byte) (int, error) {
	h.printer.Q(string(data))
	return len(data), nil
}

func NewPrinterProxy(printer *Printer) *printerProxy {
	return &printerProxy{
		printer: printer,
	}
}

func WritePersistServerStruct(printer *Printer, service, db string) error {
	printerProxy := NewPrinterProxy(printer)
	structFormat := `
type Opts_{{.Service}} struct {
    MAPPINGS TypeMappings_{{.Service}}
    HOOKS    Hooks_{{.Service}}
}

func Opts{{.Service}}(hooks Hooks_{{.Service}}, mappings TypeMappings_{{.Service}}) Opts_{{.Service}} {
    opts := Opts_{{.Service}}{
        HOOKS: &DefaultHooks_{{.Service}}{},
        MAPPINGS: &DefaultTypeMappings_{{.Service}}{},
    }
    if hooks != nil {
        opts.HOOKS = hooks
    }
    if mappings != nil {
        opts.MAPPINGS = mappings
    }
    return opts
}

type Impl_{{.Service}} struct {
    opts    *Opts_{{.Service}}
    QUERIES *Queries_{{.Service}}
    HANDLERS RestOfHandlers_{{.Service}}
    DB      *{{.DB}}
}

func Impl{{.Service}}(db *{{.DB}}, handlers RestOfHandlers_{{.Service}}, opts ...Opts_{{.Service}}) *Impl_{{.Service}} {
    var myOpts Opts_{{.Service}}
    if len(opts) > 0 {
        myOpts = opts[0]
    } else {
        myOpts = Opts{{.Service}}(&DefaultHooks_{{.Service}}{}, &DefaultTypeMappings_{{.Service}}{})
    }
    return &Impl_{{.Service}}{
        opts:    &myOpts,
        QUERIES: Queries{{.Service}}(myOpts),
        DB:      db,
        HANDLERS: handlers,
    }
}
    `
	t := template.Must(template.New("PersistServerStruct").Parse(structFormat))
	return t.Execute(printerProxy, map[string]string{
		"Service": service,
		"DB":      db,
	})
}

func WriteClientStreaming(printer *Printer, params *handlerParams, isSql bool) error {
	printerProxy := NewPrinterProxy(printer)
	sqlClientStreamingFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(stream {{.Service}}_{{.Method}}Server) error {
    tx, err := DefaultClientStreamingPersistTx(stream.Context(), this.DB)
    if err != nil {
        return gstatus.Errorf(codes.Unknown, "error creating persist tx: %v", err)
    }
    if err := this.{{.Method}}Tx(stream, tx); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(stream {{.Service}}_{{.Method}}Server, tx PersistTx) error {
    query := this.QUERIES.{{camelCase .Query}}(stream.Context(), tx)
    var first *{{.Request}}
    for {
        req, err := stream.Recv()
        if err == io.EOF {
            break
        } else if err != nil {
            return gstatus.Errorf(codes.Unknown, "error receiving request: %v", err)
        }
        if first == nil {
            first = req
        }
        {{if .Before}}
        beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(stream.Context(), req)
        if err != nil {
            return gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
        } else if beforeRes != nil {
            continue
        }
        {{end}}
        result := query.Execute(req)
        if err := result.Zero(); err != nil {
            return err
        }
    }
    if err := tx.Commit(); err != nil {
        if rollbackErr := tx.Rollback(); rollbackErr != nil {
            return fmt.Errorf("error executing '{{.Query}}' query :::AND COULD NOT ROLLBACK::: rollback err: %v, query err: %v", rollbackErr, err)
        }
    }
    res := &{{.Response}}{}

    {{if .After}}
    if err := this.opts.HOOKS.{{.Method}}AfterHook(stream.Context(), first, res); err != nil {
        return gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
    }
    {{end}}
    if err := stream.SendAndClose(res); err != nil {
        return gstatus.Errorf(codes.Unknown, "error sending back response: %v", err)
    }

    return nil
}
        `

	spannerClientStreamingFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(stream {{.Service}}_{{.Method}}Server) error {
    if err := this.{{.Method}}Tx(stream); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(stream {{.Service}}_{{.Method}}Server) error {
    items := make([]*{{.Request}}, 0)
    var first *{{.Request}}
    for {
        req, err := stream.Recv()
        if err == io.EOF {
            break
        } else if err != nil {
            return gstatus.Errorf(codes.Unknown, "error receiving request: %v", err)
        }
        if first == nil {
            first = req
        }
        {{if .Before}}
        beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(stream.Context(), req)
        if err != nil {
            return gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
        } else if beforeRes != nil {
            continue
        }
        {{end}}

        items = append(items, req)
    }

    _, err := this.DB.ReadWriteTransaction(stream.Context(), func(ctx context.Context, tx *spanner.ReadWriteTransaction) error {
        for _, item := range items {
            query := this.QUERIES.{{camelCase .Query}}(ctx, tx)
            result := query.Execute(item)
            if err := result.Zero(); err != nil {
                return err
            }
        }
        return nil
    })
	if err != nil {
		return gstatus.Errorf(codes.Unknown, "error in read write transaction: %v", err)
	}

    res := &{{.Response}}{}

    {{if .After}}
    if err := this.opts.HOOKS.{{.Method}}AfterHook(stream.Context(), first, res); err != nil {
        return gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
    }
    {{end}}
    if err := stream.SendAndClose(res); err != nil {
        return gstatus.Errorf(codes.Unknown, "error sending back response: %v", err)
    }

    return nil
}
        `

	var clientStreamingFormat string
	if isSql {
		clientStreamingFormat = sqlClientStreamingFormat
	} else {
		clientStreamingFormat = spannerClientStreamingFormat
	}

	funcMap := template.FuncMap{
		"camelCase": _gen.CamelCase,
	}
	t := template.Must(template.New("ClientStreaming").Funcs(funcMap).Parse(clientStreamingFormat))
	return t.Execute(printerProxy, params)
}

func WriteUnary(printer *Printer, params *handlerParams, isSql bool) error {
	printerProxy := NewPrinterProxy(printer)
	sqlUnaryFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(ctx context.Context, req *{{.Request}}) (*{{.Response}}, error) {
    query := this.QUERIES.{{camelCase .Query}}(ctx, this.DB)
    {{if .Before}}
    beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(ctx, req)
    if err != nil {
        return nil, gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
    } else if beforeRes != nil {
        return beforeRes, nil
    }
    {{end}}

    result := query.Execute(req)
    {{oneOrZero .}}
    if err != nil {
        return nil, err
    }

    {{if .After}}
    if err := this.opts.HOOKS.{{.Method}}AfterHook(ctx, req, res); err != nil {
        return nil, gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
    }
    {{end}}

    return res, nil
}
  `
	spannerUnaryFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(ctx context.Context, req *{{.Request}}) (*{{.Response}}, error) {
    query := this.QUERIES.{{camelCase .Query}}(ctx, this.DB.Single())
    {{if .Before}}
    beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(ctx, req)
    if err != nil {
        return nil, gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
    } else if beforeRes != nil {
        return beforeRes, nil
    }
    {{end}}

    result := query.Execute(req)
    {{oneOrZero .}}
    if err != nil {
        return nil, err
    }

    {{if .After}}
    if err := this.opts.HOOKS.{{.Method}}AfterHook(ctx, req, res); err != nil {
        return nil, gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
    }
    {{end}}

    return res, nil
}
    `
	funcMap := template.FuncMap{
		"camelCase": _gen.CamelCase,
		"oneOrZero": OneOrZero,
	}

	var unaryFormat string
	if isSql {
		unaryFormat = sqlUnaryFormat
	} else {
		unaryFormat = spannerUnaryFormat
	}

	t := template.Must(template.New("UnaryRequest").Funcs(funcMap).Parse(unaryFormat))
	return t.Execute(printerProxy, params)
}

func WriteServerStream(printer *Printer, params *handlerParams, isSql bool) error {
	printerProxy := NewPrinterProxy(printer)
	sqlServerFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server) error {
    tx, err := DefaultServerStreamingPersistTx(stream.Context(), this.DB)
    if err != nil {
        return gstatus.Errorf(codes.Unknown, "error creating persist tx: %v", err)
    }
    if err := this.{{.Method}}Tx(req, stream, tx); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server, tx PersistTx) error {
    ctx := stream.Context()
    query := this.QUERIES.{{camelCase .Query}}(ctx, tx)

    iter := query.Execute(req)
    return iter.Each(func(row *Row_{{.Service}}_{{camelCase .Query}}) error {
        res, err := row.{{.RespMethodCall}}()
        if err != nil {
            return err
        }
        return stream.Send(res)
    })
}
  `

	spannerServerFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server) error {
    if err := this.{{.Method}}Tx(req, stream, this.DB.Single()); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server, tx PersistTx) error {
    ctx := stream.Context()
    query := this.QUERIES.{{camelCase .Query}}(ctx, tx)

    iter := query.Execute(req)
    return iter.Each(func(row *Row_{{.Service}}_{{camelCase .Query}}) error {
        res, err := row.{{.RespMethodCall}}()
        if err != nil {
            return err
        }
        return stream.Send(res)
    })
}
    `
	funcMap := template.FuncMap{
		"camelCase": _gen.CamelCase,
	}

	var serverFormat string
	if isSql {
		serverFormat = sqlServerFormat
	} else {
		serverFormat = spannerServerFormat
	}

	t := template.Must(template.New("ServerStream").Funcs(funcMap).Parse(serverFormat))
	return t.Execute(printerProxy, params)
}
<commit_msg>fixed bug generating wrong var declaration under certain conditions<commit_after>package generator

import (
	"strings"
	"text/template"

	_gen "github.com/golang/protobuf/protoc-gen-go/generator"
)

type printerProxy struct {
	printer *Printer
}

type handlerParams struct {
	Service        string
	Method         string
	Query          string
	Request        string
	Response       string
	RespMethodCall string
	ZeroResponse   bool
	Before         bool
	After          bool
}

func OneOrZero(hp handlerParams) string {
	if hp.ZeroResponse {
		return strings.Join([]string{`
err := result.Zero()
res := &`, hp.Response, `{}
        `}, "")
	}
	return "res, err := result.One()." + hp.RespMethodCall + "()"
}

func (h *printerProxy) Write(data []byte) (int, error) {
	h.printer.Q(string(data))
	return len(data), nil
}

func NewPrinterProxy(printer *Printer) *printerProxy {
	return &printerProxy{
		printer: printer,
	}
}

func WritePersistServerStruct(printer *Printer, service, db string) error {
	printerProxy := NewPrinterProxy(printer)
	structFormat := `
type Opts_{{.Service}} struct {
    MAPPINGS TypeMappings_{{.Service}}
    HOOKS    Hooks_{{.Service}}
}

func Opts{{.Service}}(hooks Hooks_{{.Service}}, mappings TypeMappings_{{.Service}}) Opts_{{.Service}} {
    opts := Opts_{{.Service}}{
        HOOKS: &DefaultHooks_{{.Service}}{},
        MAPPINGS: &DefaultTypeMappings_{{.Service}}{},
    }
    if hooks != nil {
        opts.HOOKS = hooks
    }
    if mappings != nil {
        opts.MAPPINGS = mappings
    }
    return opts
}

type Impl_{{.Service}} struct {
    opts    *Opts_{{.Service}}
    QUERIES *Queries_{{.Service}}
    HANDLERS RestOfHandlers_{{.Service}}
    DB      *{{.DB}}
}

func Impl{{.Service}}(db *{{.DB}}, handlers RestOfHandlers_{{.Service}}, opts ...Opts_{{.Service}}) *Impl_{{.Service}} {
    var myOpts Opts_{{.Service}}
    if len(opts) > 0 {
        myOpts = opts[0]
    } else {
        myOpts = Opts{{.Service}}(&DefaultHooks_{{.Service}}{}, &DefaultTypeMappings_{{.Service}}{})
    }
    return &Impl_{{.Service}}{
        opts:    &myOpts,
        QUERIES: Queries{{.Service}}(myOpts),
        DB:      db,
        HANDLERS: handlers,
    }
}
    `
	t := template.Must(template.New("PersistServerStruct").Parse(structFormat))
	return t.Execute(printerProxy, map[string]string{
		"Service": service,
		"DB":      db,
	})
}

func WriteClientStreaming(printer *Printer, params *handlerParams, isSql bool) error {
	printerProxy := NewPrinterProxy(printer)
	sqlClientStreamingFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(stream {{.Service}}_{{.Method}}Server) error {
    tx, err := DefaultClientStreamingPersistTx(stream.Context(), this.DB)
    if err != nil {
        return gstatus.Errorf(codes.Unknown, "error creating persist tx: %v", err)
    }
    if err := this.{{.Method}}Tx(stream, tx); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(stream {{.Service}}_{{.Method}}Server, tx PersistTx) error {
    query := this.QUERIES.{{camelCase .Query}}(stream.Context(), tx)
    var first *{{.Request}}
    for {
        req, err := stream.Recv()
        if err == io.EOF {
            break
        } else if err != nil {
            return gstatus.Errorf(codes.Unknown, "error receiving request: %v", err)
        }
        if first == nil {
            first = req
        }
        {{if .Before}}
        {
            beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(stream.Context(), req)
            if err != nil {
                return gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
            } else if beforeRes != nil {
                continue
            }
        }
        {{end}}
        result := query.Execute(req)
        if err := result.Zero(); err != nil {
            return err
        }
    }
    if err := tx.Commit(); err != nil {
        if rollbackErr := tx.Rollback(); rollbackErr != nil {
            return fmt.Errorf("error executing '{{.Query}}' query :::AND COULD NOT ROLLBACK::: rollback err: %v, query err: %v", rollbackErr, err)
        }
    }
    res := &{{.Response}}{}

    {{if .After}}
    {
        if err := this.opts.HOOKS.{{.Method}}AfterHook(stream.Context(), first, res); err != nil {
            return gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
        }
    }
    {{end}}
    if err := stream.SendAndClose(res); err != nil {
        return gstatus.Errorf(codes.Unknown, "error sending back response: %v", err)
    }

    return nil
}
        `

	spannerClientStreamingFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(stream {{.Service}}_{{.Method}}Server) error {
    if err := this.{{.Method}}Tx(stream); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(stream {{.Service}}_{{.Method}}Server) error {
    items := make([]*{{.Request}}, 0)
    var first *{{.Request}}
    for {
        req, err := stream.Recv()
        if err == io.EOF {
            break
        } else if err != nil {
            return gstatus.Errorf(codes.Unknown, "error receiving request: %v", err)
        }
        if first == nil {
            first = req
        }
        {{if .Before}}
        {
            beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(stream.Context(), req)
            if err != nil {
                return gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
            } else if beforeRes != nil {
                continue
            }
        }
        {{end}}

        items = append(items, req)
    }

    _, err := this.DB.ReadWriteTransaction(stream.Context(), func(ctx context.Context, tx *spanner.ReadWriteTransaction) error {
        for _, item := range items {
            query := this.QUERIES.{{camelCase .Query}}(ctx, tx)
            result := query.Execute(item)
            if err := result.Zero(); err != nil {
                return err
            }
        }
        return nil
    })
	if err != nil {
		return gstatus.Errorf(codes.Unknown, "error in read write transaction: %v", err)
	}

    res := &{{.Response}}{}

    {{if .After}}
    {
        if err := this.opts.HOOKS.{{.Method}}AfterHook(stream.Context(), first, res); err != nil {
            return gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
        }
    }
    {{end}}
    if err := stream.SendAndClose(res); err != nil {
        return gstatus.Errorf(codes.Unknown, "error sending back response: %v", err)
    }

    return nil
}
        `

	var clientStreamingFormat string
	if isSql {
		clientStreamingFormat = sqlClientStreamingFormat
	} else {
		clientStreamingFormat = spannerClientStreamingFormat
	}

	funcMap := template.FuncMap{
		"camelCase": _gen.CamelCase,
	}
	t := template.Must(template.New("ClientStreaming").Funcs(funcMap).Parse(clientStreamingFormat))
	return t.Execute(printerProxy, params)
}

func WriteUnary(printer *Printer, params *handlerParams, isSql bool) error {
	printerProxy := NewPrinterProxy(printer)
	sqlUnaryFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(ctx context.Context, req *{{.Request}}) (*{{.Response}}, error) {
    query := this.QUERIES.{{camelCase .Query}}(ctx, this.DB)
    {{if .Before}}
    {
        beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(ctx, req)
        if err != nil {
            return nil, gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
        } else if beforeRes != nil {
            return beforeRes, nil
        }
    }
    {{end}}

    result := query.Execute(req)
    {{oneOrZero .}}
    if err != nil {
        return nil, err
    }

    {{if .After}}
    {
        if err := this.opts.HOOKS.{{.Method}}AfterHook(ctx, req, res); err != nil {
            return nil, gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
        }
    }
    {{end}}

    return res, nil
}
  `
	spannerUnaryFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(ctx context.Context, req *{{.Request}}) (*{{.Response}}, error) {
    query := this.QUERIES.{{camelCase .Query}}(ctx, this.DB.Single())
    {{if .Before}}
    beforeRes, err := this.opts.HOOKS.{{.Method}}BeforeHook(ctx, req)
    if err != nil {
        return nil, gstatus.Errorf(codes.Unknown, "error in before hook: %v", err)
    } else if beforeRes != nil {
        return beforeRes, nil
    }
    {{end}}

    result := query.Execute(req)
    {{oneOrZero .}}
    if err != nil {
        return nil, err
    }

    {{if .After}}
    {
        if err := this.opts.HOOKS.{{.Method}}AfterHook(ctx, req, res); err != nil {
            return nil, gstatus.Errorf(codes.Unknown, "error in after hook: %v", err)
        }
    }
    {{end}}

    return res, nil
}
    `
	funcMap := template.FuncMap{
		"camelCase": _gen.CamelCase,
		"oneOrZero": OneOrZero,
	}

	var unaryFormat string
	if isSql {
		unaryFormat = sqlUnaryFormat
	} else {
		unaryFormat = spannerUnaryFormat
	}

	t := template.Must(template.New("UnaryRequest").Funcs(funcMap).Parse(unaryFormat))
	return t.Execute(printerProxy, params)
}

func WriteServerStream(printer *Printer, params *handlerParams, isSql bool) error {
	printerProxy := NewPrinterProxy(printer)
	sqlServerFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server) error {
    tx, err := DefaultServerStreamingPersistTx(stream.Context(), this.DB)
    if err != nil {
        return gstatus.Errorf(codes.Unknown, "error creating persist tx: %v", err)
    }
    if err := this.{{.Method}}Tx(req, stream, tx); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server, tx PersistTx) error {
    ctx := stream.Context()
    query := this.QUERIES.{{camelCase .Query}}(ctx, tx)

    iter := query.Execute(req)
    return iter.Each(func(row *Row_{{.Service}}_{{camelCase .Query}}) error {
        res, err := row.{{.RespMethodCall}}()
        if err != nil {
            return err
        }
        return stream.Send(res)
    })
}
  `

	spannerServerFormat := `
func (this *Impl_{{.Service}}) {{.Method}}(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server) error {
    if err := this.{{.Method}}Tx(req, stream, this.DB.Single()); err != nil {
        return gstatus.Errorf(codes.Unknown, "error executing '{{.Query}}' query: %v", err)
    }
    return nil
}

func (this *Impl_{{.Service}}) {{.Method}}Tx(req *{{.Request}}, stream {{.Service}}_{{.Method}}Server, tx PersistTx) error {
    ctx := stream.Context()
    query := this.QUERIES.{{camelCase .Query}}(ctx, tx)

    iter := query.Execute(req)
    return iter.Each(func(row *Row_{{.Service}}_{{camelCase .Query}}) error {
        res, err := row.{{.RespMethodCall}}()
        if err != nil {
            return err
        }
        return stream.Send(res)
    })
}
    `
	funcMap := template.FuncMap{
		"camelCase": _gen.CamelCase,
	}

	var serverFormat string
	if isSql {
		serverFormat = sqlServerFormat
	} else {
		serverFormat = spannerServerFormat
	}

	t := template.Must(template.New("ServerStream").Funcs(funcMap).Parse(serverFormat))
	return t.Execute(printerProxy, params)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/codegangsta/cli"
)

var infoCommand = cli.Command{
	Name:   "info",
	Usage:  "show cluster info",
	Action: infoAction,
}

func infoAction(c *cli.Context) {
	cfg, err := loadConfig()
	if err != nil {
		logger.Fatal(err)
	}
	m := NewManager(cfg)
	info, err := m.Info()
	if err != nil {
		logger.Fatalf("error getting cluster info: %s", err)
	}
	w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintf(w, "Cpus: %.2f\n", info.Cpus)
	fmt.Fprintf(w, "Memory: %.2f MB\n", info.Memory)
	fmt.Fprintf(w, "Containers: %d\n", info.ContainerCount)
	fmt.Fprintf(w, "Images: %d\n", info.ImageCount)
	fmt.Fprintf(w, "Engines: %d\n", info.EngineCount)
	fmt.Fprintf(w, "Reserved Cpus: %.2f\n", info.ReservedCpus)
	fmt.Fprintf(w, "Reserved Memory: %.2f\n", info.ReservedMemory)
	w.Flush()
}
<commit_msg>show cluster usage as percentage as well as actual<commit_after>package main

import (
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/codegangsta/cli"
)

var infoCommand = cli.Command{
	Name:   "info",
	Usage:  "show cluster info",
	Action: infoAction,
}

func infoAction(c *cli.Context) {
	cfg, err := loadConfig()
	if err != nil {
		logger.Fatal(err)
	}
	m := NewManager(cfg)
	info, err := m.Info()
	if err != nil {
		logger.Fatalf("error getting cluster info: %s", err)
	}
	cpuPercentage := (info.ReservedCpus / info.Cpus) * 100
	memPercentage := (info.ReservedMemory / info.Memory) * 100

	w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintf(w, "Cpus: %.2f\n", info.Cpus)
	fmt.Fprintf(w, "Memory: %.2f MB\n", info.Memory)
	fmt.Fprintf(w, "Containers: %d\n", info.ContainerCount)
	fmt.Fprintf(w, "Images: %d\n", info.ImageCount)
	fmt.Fprintf(w, "Engines: %d\n", info.EngineCount)
	fmt.Fprintf(w, "Reserved Cpus: %.2f%% (%.2f)\n", cpuPercentage, info.ReservedCpus)
	fmt.Fprintf(w, "Reserved Memory: %.2f%% (%.2f) MB\n", memPercentage, info.ReservedMemory)
	w.Flush()
}
<|endoftext|>


<commit_before>package docker

import (
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/iron-io/titan/runner/agent"
)

// wrap docker client calls so we can retry 500s, kind of sucks but fsouza doesn't
// bake in retries we can use internally, could contribute it at some point, would
// be much more convenient if we didn't have to do this, but it's better than ad hoc retries.
// also adds timeouts to many operations, varying by operation
// TODO could generate this, maybe not worth it, may not change often
type dockerClient interface {
	// Each of these are github.com/fsouza/go-dockerclient methods

	AttachToContainer(opts docker.AttachToContainerOptions) error
	StartContainer(id string, hostConfig *docker.HostConfig) error
	CreateContainer(opts docker.CreateContainerOptions) (*docker.Container, error)
	RemoveContainer(opts docker.RemoveContainerOptions) error
	PullImage(opts docker.PullImageOptions, auth docker.AuthConfiguration) error
	InspectImage(name string) (*docker.Image, error)
	InspectContainer(id string) (*docker.Container, error)
	StopContainer(id string, timeout uint) error
	Stats(opts docker.StatsOptions) error
}

// TODO: switch to github.com/docker/engine-api
func newClient() dockerClient {
	// docker, err := docker.NewClient(conf.Docker)
	client, err := docker.NewClientFromEnv()
	if err != nil {
		logrus.WithError(err).Fatal("couldn't create docker client")
	}

	// NOTE add granularity to things like pull, should not effect
	// hijacked / streaming endpoints
	client.SetTimeout(120 * time.Second)

	return &dockerWrap{client}
}

type dockerWrap struct {
	docker *docker.Client
}

func retry(f func() error) {
	var b agent.Backoff
	then := time.Now()
	for time.Now().Sub(then) < 10*time.Minute { // retry for 10 minutes
		err := f()
		if isTemporary(err) || isDocker500(err) {
			logrus.WithError(err).Warn("docker temporary error, retrying")
			b.Sleep()
			continue
		}
		return
	}
	logrus.Warn("retrying on docker errors exceeded 2 minutes, restart docker or rotate this instance?")
}

func isTemporary(err error) bool {
	terr, ok := err.(interface {
		Temporary() bool
	})
	return ok && terr.Temporary()
}

func isDocker500(err error) bool {
	derr, ok := err.(*docker.Error)
	return ok && derr.Status >= 500
}

func (d *dockerWrap) AttachToContainer(opts docker.AttachToContainerOptions) (err error) {
	retry(func() error {
		err = d.docker.AttachToContainer(opts)
		return err
	})
	return err
}

func (d *dockerWrap) StartContainer(id string, hostConfig *docker.HostConfig) (err error) {
	retry(func() error {
		err = d.docker.StartContainer(id, hostConfig)
		return err
	})
	return err
}

func (d *dockerWrap) CreateContainer(opts docker.CreateContainerOptions) (c *docker.Container, err error) {
	retry(func() error {
		c, err = d.docker.CreateContainer(opts)
		return err
	})
	return c, err
}

func (d *dockerWrap) RemoveContainer(opts docker.RemoveContainerOptions) (err error) {
	retry(func() error {
		err = d.docker.RemoveContainer(opts)
		return err
	})
	return err
}

func (d *dockerWrap) PullImage(opts docker.PullImageOptions, auth docker.AuthConfiguration) (err error) {
	retry(func() error {
		err = d.docker.PullImage(opts, auth)
		return err
	})
	return err
}

func (d *dockerWrap) InspectImage(name string) (i *docker.Image, err error) {
	retry(func() error {
		i, err = d.docker.InspectImage(name)
		return err
	})
	return i, err
}

func (d *dockerWrap) InspectContainer(id string) (c *docker.Container, err error) {
	retry(func() error {
		c, err = d.docker.InspectContainer(id)
		return err
	})
	return c, err
}

func (d *dockerWrap) StopContainer(id string, timeout uint) (err error) {
	retry(func() error {
		err = d.docker.StopContainer(id, timeout)
		return err
	})
	return err
}

func (d *dockerWrap) Stats(opts docker.StatsOptions) (err error) {
	retry(func() error {
		err = d.docker.Stats(opts)
		return err
	})
	return err
}
<commit_msg>appease the vet overlords<commit_after>package docker

import (
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/fsouza/go-dockerclient"
	"github.com/iron-io/titan/runner/agent"
)

// wrap docker client calls so we can retry 500s, kind of sucks but fsouza doesn't
// bake in retries we can use internally, could contribute it at some point, would
// be much more convenient if we didn't have to do this, but it's better than ad hoc retries.
// also adds timeouts to many operations, varying by operation
// TODO could generate this, maybe not worth it, may not change often
type dockerClient interface {
	// Each of these are github.com/fsouza/go-dockerclient methods

	AttachToContainer(opts docker.AttachToContainerOptions) error
	StartContainer(id string, hostConfig *docker.HostConfig) error
	CreateContainer(opts docker.CreateContainerOptions) (*docker.Container, error)
	RemoveContainer(opts docker.RemoveContainerOptions) error
	PullImage(opts docker.PullImageOptions, auth docker.AuthConfiguration) error
	InspectImage(name string) (*docker.Image, error)
	InspectContainer(id string) (*docker.Container, error)
	StopContainer(id string, timeout uint) error
	Stats(opts docker.StatsOptions) error
}

// TODO: switch to github.com/docker/engine-api
func newClient() dockerClient {
	// docker, err := docker.NewClient(conf.Docker)
	client, err := docker.NewClientFromEnv()
	if err != nil {
		logrus.WithError(err).Fatal("couldn't create docker client")
	}

	// NOTE add granularity to things like pull, should not effect
	// hijacked / streaming endpoints
	client.SetTimeout(120 * time.Second)

	return &dockerWrap{client}
}

type dockerWrap struct {
	docker *docker.Client
}

func retry(f func() error) {
	var b agent.Backoff
	then := time.Now()
	for time.Since(then) < 10*time.Minute { // retry for 10 minutes
		err := f()
		if isTemporary(err) || isDocker500(err) {
			logrus.WithError(err).Warn("docker temporary error, retrying")
			b.Sleep()
			continue
		}
		return
	}
	logrus.Warn("retrying on docker errors exceeded 2 minutes, restart docker or rotate this instance?")
}

func isTemporary(err error) bool {
	terr, ok := err.(interface {
		Temporary() bool
	})
	return ok && terr.Temporary()
}

func isDocker500(err error) bool {
	derr, ok := err.(*docker.Error)
	return ok && derr.Status >= 500
}

func (d *dockerWrap) AttachToContainer(opts docker.AttachToContainerOptions) (err error) {
	retry(func() error {
		err = d.docker.AttachToContainer(opts)
		return err
	})
	return err
}

func (d *dockerWrap) StartContainer(id string, hostConfig *docker.HostConfig) (err error) {
	retry(func() error {
		err = d.docker.StartContainer(id, hostConfig)
		return err
	})
	return err
}

func (d *dockerWrap) CreateContainer(opts docker.CreateContainerOptions) (c *docker.Container, err error) {
	retry(func() error {
		c, err = d.docker.CreateContainer(opts)
		return err
	})
	return c, err
}

func (d *dockerWrap) RemoveContainer(opts docker.RemoveContainerOptions) (err error) {
	retry(func() error {
		err = d.docker.RemoveContainer(opts)
		return err
	})
	return err
}

func (d *dockerWrap) PullImage(opts docker.PullImageOptions, auth docker.AuthConfiguration) (err error) {
	retry(func() error {
		err = d.docker.PullImage(opts, auth)
		return err
	})
	return err
}

func (d *dockerWrap) InspectImage(name string) (i *docker.Image, err error) {
	retry(func() error {
		i, err = d.docker.InspectImage(name)
		return err
	})
	return i, err
}

func (d *dockerWrap) InspectContainer(id string) (c *docker.Container, err error) {
	retry(func() error {
		c, err = d.docker.InspectContainer(id)
		return err
	})
	return c, err
}

func (d *dockerWrap) StopContainer(id string, timeout uint) (err error) {
	retry(func() error {
		err = d.docker.StopContainer(id, timeout)
		return err
	})
	return err
}

func (d *dockerWrap) Stats(opts docker.StatsOptions) (err error) {
	retry(func() error {
		err = d.docker.Stats(opts)
		return err
	})
	return err
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/dnephin/dobi/config"
	"github.com/dnephin/dobi/logging"
	"github.com/dnephin/dobi/tasks"
	"github.com/dnephin/dobi/tasks/client"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/spf13/cobra"
)

const (
	// DefaultDockerAPIVersion is the default version of the docker API to use
	DefaultDockerAPIVersion = "1.23"
)

var (
	version   = "0.8"
	gitsha    = "unknown"
	buildDate = ""
)

type dobiOptions struct {
	filename string
	verbose  bool
	quiet    bool
	tasks    []string
	version  bool
}

// NewRootCommand returns a new root command
func NewRootCommand() *cobra.Command {
	var opts dobiOptions

	cmd := &cobra.Command{
		Use:              "dobi [flags] RESOURCE[:ACTION] [RESOURCE[:ACTION]...]",
		Short:            "A build automation tool for Docker applications",
		SilenceUsage:     true,
		SilenceErrors:    true,
		TraverseChildren: true,
		Args:             cobra.ArbitraryArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.tasks = args
			return runDobi(opts)
		},
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			initLogging(opts.verbose, opts.quiet)
			return nil
		},
	}

	flags := cmd.Flags()
	flags.StringVarP(&opts.filename, "filename", "f", "dobi.yaml", "Path to config file")
	flags.BoolVarP(&opts.verbose, "verbose", "v", false, "Verbose")
	flags.BoolVarP(&opts.quiet, "quiet", "q", false, "Quiet")
	flags.BoolVar(&opts.version, "version", false, "Print version and exit")

	flags.SetInterspersed(false)
	cmd.AddCommand(
		newListCommand(&opts),
		newCleanCommand(&opts),
	)
	return cmd
}

func runDobi(opts dobiOptions) error {
	if opts.version {
		printVersion()
		return nil
	}

	conf, err := config.Load(opts.filename)
	if err != nil {
		return err
	}

	client, err := buildClient()
	if err != nil {
		return fmt.Errorf("failed to create client: %s", err)
	}

	return tasks.Run(tasks.RunOptions{
		Client: client,
		Config: conf,
		Tasks:  opts.tasks,
		Quiet:  opts.quiet,
	})
}

func initLogging(verbose, quiet bool) {
	logger := logging.Log
	if verbose {
		logger.Level = log.DebugLevel
	}
	if quiet {
		logger.Level = log.WarnLevel
	}
	logger.Out = os.Stderr

	formatter := &logging.Formatter{}
	log.SetFormatter(formatter)
	logger.Formatter = formatter
}

func buildClient() (client.DockerClient, error) {
	apiVersion := os.Getenv("DOCKER_API_VERSION")
	if apiVersion == "" {
		apiVersion = DefaultDockerAPIVersion
	}
	// TODO: args for client
	client, err := docker.NewVersionedClientFromEnv(apiVersion)
	if err != nil {
		return nil, err
	}
	log.Debug("Docker client created")
	return client, nil
}

func printVersion() {
	fmt.Printf("dobi version %v (build: %v, date: %s)\n", version, gitsha, buildDate)
}
<commit_msg>Bump version to 0.9<commit_after>package cmd

import (
	"fmt"
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/dnephin/dobi/config"
	"github.com/dnephin/dobi/logging"
	"github.com/dnephin/dobi/tasks"
	"github.com/dnephin/dobi/tasks/client"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/spf13/cobra"
)

const (
	// DefaultDockerAPIVersion is the default version of the docker API to use
	DefaultDockerAPIVersion = "1.23"
)

var (
	version   = "0.9"
	gitsha    = "unknown"
	buildDate = ""
)

type dobiOptions struct {
	filename string
	verbose  bool
	quiet    bool
	tasks    []string
	version  bool
}

// NewRootCommand returns a new root command
func NewRootCommand() *cobra.Command {
	var opts dobiOptions

	cmd := &cobra.Command{
		Use:              "dobi [flags] RESOURCE[:ACTION] [RESOURCE[:ACTION]...]",
		Short:            "A build automation tool for Docker applications",
		SilenceUsage:     true,
		SilenceErrors:    true,
		TraverseChildren: true,
		Args:             cobra.ArbitraryArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			opts.tasks = args
			return runDobi(opts)
		},
		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
			initLogging(opts.verbose, opts.quiet)
			return nil
		},
	}

	flags := cmd.Flags()
	flags.StringVarP(&opts.filename, "filename", "f", "dobi.yaml", "Path to config file")
	flags.BoolVarP(&opts.verbose, "verbose", "v", false, "Verbose")
	flags.BoolVarP(&opts.quiet, "quiet", "q", false, "Quiet")
	flags.BoolVar(&opts.version, "version", false, "Print version and exit")

	flags.SetInterspersed(false)
	cmd.AddCommand(
		newListCommand(&opts),
		newCleanCommand(&opts),
	)
	return cmd
}

func runDobi(opts dobiOptions) error {
	if opts.version {
		printVersion()
		return nil
	}

	conf, err := config.Load(opts.filename)
	if err != nil {
		return err
	}

	client, err := buildClient()
	if err != nil {
		return fmt.Errorf("failed to create client: %s", err)
	}

	return tasks.Run(tasks.RunOptions{
		Client: client,
		Config: conf,
		Tasks:  opts.tasks,
		Quiet:  opts.quiet,
	})
}

func initLogging(verbose, quiet bool) {
	logger := logging.Log
	if verbose {
		logger.Level = log.DebugLevel
	}
	if quiet {
		logger.Level = log.WarnLevel
	}
	logger.Out = os.Stderr

	formatter := &logging.Formatter{}
	log.SetFormatter(formatter)
	logger.Formatter = formatter
}

func buildClient() (client.DockerClient, error) {
	apiVersion := os.Getenv("DOCKER_API_VERSION")
	if apiVersion == "" {
		apiVersion = DefaultDockerAPIVersion
	}
	// TODO: args for client
	client, err := docker.NewVersionedClientFromEnv(apiVersion)
	if err != nil {
		return nil, err
	}
	log.Debug("Docker client created")
	return client, nil
}

func printVersion() {
	fmt.Printf("dobi version %v (build: %v, date: %s)\n", version, gitsha, buildDate)
}
<|endoftext|>


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

package cmd

import (
	"fmt"

	"github.com/jconard3/docore/client"
	"github.com/spf13/cobra"
)

// listCmd represents the list command
var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List all dropets",
	Long:  `List all droplets.`,
	Run: func(cmd *cobra.Command, args []string) {
		client, _ := CreateClient()
		names := ListDroplets(client)
		fmt.Println(names)
	},
}

func init() {
	dropletCmd.AddCommand(listCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// listCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// listCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

}

func ListDroplets(client godo.Client) []string {
	opt := &godo.ListOptions{
		Page:    1,
		PerPage: 25,
	}
	var dropletNames []string
	droplets, _, _ := client.Droplets.List(opt)
	for _, element := range droplets {
		dropletNames = append(dropletNames, element.Name)
	}
	return dropletNames
}
<commit_msg>Fixed up new client package importing.<commit_after>// Copyright © 2017 NAME HERE <EMAIL ADDRESS>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"

	"github.com/digitalocean/godo"
	"github.com/jconard3/docore/client"
	"github.com/spf13/cobra"
)

// listCmd represents the list command
var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List all dropets",
	Long:  `List all droplets.`,
	Run: func(cmd *cobra.Command, args []string) {
		my_client, _ := client.CreateClient()
		names := ListDroplets(my_client)
		fmt.Println(names)
	},
}

func init() {
	dropletCmd.AddCommand(listCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// listCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// listCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

}

func ListDroplets(client godo.Client) []string {
	opt := &godo.ListOptions{
		Page:    1,
		PerPage: 25,
	}
	var dropletNames []string
	droplets, _, _ := client.Droplets.List(opt)
	for _, element := range droplets {
		dropletNames = append(dropletNames, element.Name)
	}
	return dropletNames
}
<|endoftext|>


<commit_before>/*
Copyright 2016 Skippbox, Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"encoding/json"
	"fmt"

	"github.com/gosuri/uitable"
	"github.com/skippbox/kubeless/pkg/controller"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
	"github.com/skippbox/kubeless/pkg/spec"
)

var listCmd = &cobra.Command{
	Use:   "ls FLAG",
	Short: "list all functions deployed to Kubeless",
	Long:  `list all functions deployed to Kubeless`,
	Run: func(cmd *cobra.Command, args []string) {
		master, err := cmd.Flags().GetString("master")
		if master == "" {
			master = "localhost"
		}

		output, err := cmd.Flags().GetString("out")
		cfg := newControllerConfig(master)
		c := controller.New(cfg)
		_, err = c.FindResourceVersion()
		if err != nil {
			fmt.Errorf("Can not list functions: %v", err)
		}

		if len(args) == 0 {
			for k, _ := range c.Functions {
				args = append(args, k)
			}
		}

		printFunctions(args, c.Functions, output)
	},
}

func init() {
	listCmd.Flags().StringP("out", "o", "", "Output format. One of: json|yaml")
}

func printFunctions(args []string, functions map[string]*spec.Function, output string) {
	if output == "" {
		table := uitable.New()
		table.MaxColWidth = 30
		table.AddRow("NAME", "HANDLER", "RUNTIME")
		for _, f := range args {
			n := fmt.Sprintf(f)
			h := fmt.Sprintf(functions[f].Spec.Handler)
			r := fmt.Sprintf(functions[f].Spec.Runtime)
			table.AddRow(n, h, r)
		}
		fmt.Println(table.String())
	} else {
		for _, f := range args {
			switch output {
			case "json":
				b, _ := json.MarshalIndent(functions[f].Spec, "", "  ")
				fmt.Println(string(b))
			case "yaml":
				b, _ := yaml.Marshal(functions[f].Spec)
				fmt.Println(string(b))
			default:
				fmt.Errorf("Wrong output format. Please use only json|yaml.")
			}
		}
	}
}<commit_msg>display function type: http / pubsub<commit_after>/*
Copyright 2016 Skippbox, Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"encoding/json"
	"fmt"

	"github.com/gosuri/uitable"
	"github.com/skippbox/kubeless/pkg/controller"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
	"github.com/skippbox/kubeless/pkg/spec"
)

var listCmd = &cobra.Command{
	Use:   "ls FLAG",
	Short: "list all functions deployed to Kubeless",
	Long:  `list all functions deployed to Kubeless`,
	Run: func(cmd *cobra.Command, args []string) {
		master, err := cmd.Flags().GetString("master")
		if master == "" {
			master = "localhost"
		}

		output, err := cmd.Flags().GetString("out")
		cfg := newControllerConfig(master)
		c := controller.New(cfg)
		_, err = c.FindResourceVersion()
		if err != nil {
			fmt.Errorf("Can not list functions: %v", err)
		}

		if len(args) == 0 {
			for k, _ := range c.Functions {
				args = append(args, k)
			}
		}

		printFunctions(args, c.Functions, output)
	},
}

func init() {
	listCmd.Flags().StringP("out", "o", "", "Output format. One of: json|yaml")
}

func printFunctions(args []string, functions map[string]*spec.Function, output string) {
	if output == "" {
		table := uitable.New()
		table.MaxColWidth = 30
		table.AddRow("NAME", "HANDLER", "RUNTIME", "TYPE")
		for _, f := range args {
			n := fmt.Sprintf(f)
			h := fmt.Sprintf(functions[f].Spec.Handler)
			r := fmt.Sprintf(functions[f].Spec.Runtime)
			t := fmt.Sprintf(functions[f].Spec.Type)
			table.AddRow(n, h, r, t)
		}
		fmt.Println(table.String())
	} else {
		for _, f := range args {
			switch output {
			case "json":
				b, _ := json.MarshalIndent(functions[f].Spec, "", "  ")
				fmt.Println(string(b))
			case "yaml":
				b, _ := yaml.Marshal(functions[f].Spec)
				fmt.Println(string(b))
			default:
				fmt.Errorf("Wrong output format. Please use only json|yaml.")
			}
		}
	}
}<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	rtc "github.com/shiguredo/sora-webrtc-build"
	y "github.com/shiguredo/yspata"
	"os"
	"path/filepath"
	"runtime"
)

func printHelp() {
	y.PrintLines(
		"Usage: build [options] <command>",
		"",
		"Commands:",
		"  fetch",
		"        Get depot_tools and source files",
		"  build",
		"        Build libraries",
		"  archive",
		"        Archive libraries",
		"  build-clean",
		"        Remove/restore files generated build command",
		"  clean",
		"        Remove all built files and discard all changes",
		"  help",
		"        Print this message",
		"  version",
		"        Print version",
		"",
		"Options:")
	flag.PrintDefaults()
}

var confFile = "config.json"

var confOpt = flag.String("config", confFile, "configuration file")

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		printHelp()
		y.Fail()
	}

	conf, err := rtc.LoadConfig(*confOpt)
	if err != nil {
		y.Eprintf("cannot load config: %s", err.Error())
		y.Fail()
	}

	path := os.Getenv("PATH")
	os.Setenv("PATH", conf.DepotToolsDir+":"+path)

	var native rtc.Native
	if y.IsMac {
		native = rtc.NewIOS(conf)
	} else if y.IsLinux {
		native = rtc.NewAndroid(conf)
	} else {
		y.Eprintf("%s OS is not supported\n", runtime.GOOS)
		y.Fail()
	}

	bld := rtc.NewBuilder(conf, native)

	subcmd := flag.Arg(0)
	switch subcmd {
	case "fetch":
		bld.GetDepotTools()
		bld.Fetch()

	case "build":
		if !y.Exists(conf.GclientConf) || !y.Exists(conf.GclientEntries) {
			y.Eprintf("%s or %s are not found. Do './webrtc-build fetch'.",
				conf.GclientConf, conf.GclientEntries)
			y.Fail()
		}
		bld.Build()

	case "archive":
		bld.Archive()

	case "clean":
		bld.Clean()
		bld.Reset()

	case "build-clean":
		bld.BuildClean()

	case "help":
		printHelp()

	case "version":
		y.Printf("webrtc-build %s, library %s",
			rtc.FullVersion, conf.WebRTCVersion())

	case "selfdist":
		dist := fmt.Sprintf("sora-webrtc-build-%s", rtc.FullVersion)

		wd, _ := os.Getwd()
		curPatchDir, _ := filepath.Rel(wd, conf.PatchDir)
		distPatchDir := y.Join(dist, curPatchDir)
		y.Execf("rm -rf %s %s.zip", dist, dist)
		y.Execf("mkdir %s", dist)
		y.Execf("go build -o webrtc-build cmd/main.go")
		y.Execf("cp webrtc-build %s", dist)
		y.Execf("cp config.json %s", dist)
		y.Execf("cp config-ios-dev.json %s", dist)
		os.MkdirAll(distPatchDir, 0755)
		for _, p := range conf.Patches {
			path := y.Join(curPatchDir, p.Patch)
			y.Execf("cp %s %s", path, distPatchDir)
		}
		y.Execf("tar czf %s.tar.gz %s", dist, dist)

	default:
		y.Eprintf("Unknown command: %s", subcmd)
		y.Fail()
	}
}
<commit_msg>build-clean は消す その2<commit_after>package main

import (
	"flag"
	"fmt"
	rtc "github.com/shiguredo/sora-webrtc-build"
	y "github.com/shiguredo/yspata"
	"os"
	"path/filepath"
	"runtime"
)

func printHelp() {
	y.PrintLines(
		"Usage: build [options] <command>",
		"",
		"Commands:",
		"  fetch",
		"        Get depot_tools and source files",
		"  build",
		"        Build libraries",
		"  archive",
		"        Archive libraries",
		"  build-clean",
		"        Remove/restore files generated build command",
		"  clean",
		"        Remove all built files and discard all changes",
		"  help",
		"        Print this message",
		"  version",
		"        Print version",
		"",
		"Options:")
	flag.PrintDefaults()
}

var confFile = "config.json"

var confOpt = flag.String("config", confFile, "configuration file")

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		printHelp()
		y.Fail()
	}

	conf, err := rtc.LoadConfig(*confOpt)
	if err != nil {
		y.Eprintf("cannot load config: %s", err.Error())
		y.Fail()
	}

	path := os.Getenv("PATH")
	os.Setenv("PATH", conf.DepotToolsDir+":"+path)

	var native rtc.Native
	if y.IsMac {
		native = rtc.NewIOS(conf)
	} else if y.IsLinux {
		native = rtc.NewAndroid(conf)
	} else {
		y.Eprintf("%s OS is not supported\n", runtime.GOOS)
		y.Fail()
	}

	bld := rtc.NewBuilder(conf, native)

	subcmd := flag.Arg(0)
	switch subcmd {
	case "fetch":
		bld.GetDepotTools()
		bld.Fetch()

	case "build":
		if !y.Exists(conf.GclientConf) || !y.Exists(conf.GclientEntries) {
			y.Eprintf("%s or %s are not found. Do './webrtc-build fetch'.",
				conf.GclientConf, conf.GclientEntries)
			y.Fail()
		}
		bld.Build()

	case "archive":
		bld.Archive()

	case "clean":
		bld.Clean()
		bld.Reset()

	case "help":
		printHelp()

	case "version":
		y.Printf("webrtc-build %s, library %s",
			rtc.FullVersion, conf.WebRTCVersion())

	case "selfdist":
		dist := fmt.Sprintf("sora-webrtc-build-%s", rtc.FullVersion)

		wd, _ := os.Getwd()
		curPatchDir, _ := filepath.Rel(wd, conf.PatchDir)
		distPatchDir := y.Join(dist, curPatchDir)
		y.Execf("rm -rf %s %s.zip", dist, dist)
		y.Execf("mkdir %s", dist)
		y.Execf("go build -o webrtc-build cmd/main.go")
		y.Execf("cp webrtc-build %s", dist)
		y.Execf("cp config.json %s", dist)
		y.Execf("cp config-ios-dev.json %s", dist)
		os.MkdirAll(distPatchDir, 0755)
		for _, p := range conf.Patches {
			path := y.Join(curPatchDir, p.Patch)
			y.Execf("cp %s %s", path, distPatchDir)
		}
		y.Execf("tar czf %s.tar.gz %s", dist, dist)

	default:
		y.Eprintf("Unknown command: %s", subcmd)
		y.Fail()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"

	"github.com/megapk/fake_http_server/database"
	"encoding/json"
	"io/ioutil"
	"os"
	"github.com/gorilla/handlers"
)

var db *database.DatabaseConnection

func handler(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()
	link := r.URL.String()
	fmt.Fprintf(w, "My url: %s", link)

	header, _ := json.Marshal(r.Header)
	body, _ := ioutil.ReadAll(r.Body)
	form, _ := json.Marshal(r.Form)
	postForm, _ := json.Marshal(r.PostForm)
	db.AddUrl(&database.Bot{
		Link: link,
		Header: string(header),
		Body: string(body),
		Form: string(form),
		PostForm: string(postForm),
	})
}

func GetUrls(w http.ResponseWriter, r *http.Request) {
	js, err := json.Marshal(db.GetUrls())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func main() {
	db = database.New("db.sqlite3")
	http.HandleFunc("/get_urls", GetUrls)
	http.Handle("/", handlers.LoggingHandler(os.Stdout, http.HandlerFunc(handler)))
	http.ListenAndServe(":8080", nil)
}<commit_msg>Add pretty json<commit_after>package main

import (
	"fmt"
	"net/http"

	"github.com/megapk/fake_http_server/database"
	"encoding/json"
	"io/ioutil"
	"os"
	"github.com/gorilla/handlers"
)

var db *database.DatabaseConnection

func handler(w http.ResponseWriter, r *http.Request) {
	defer r.Body.Close()
	link := r.URL.String()
	fmt.Fprintf(w, "My url: %s", link)

	header, _ := json.Marshal(r.Header)
	body, _ := ioutil.ReadAll(r.Body)
	form, _ := json.Marshal(r.Form)
	postForm, _ := json.Marshal(r.PostForm)
	db.AddUrl(&database.Bot{
		Link: link,
		Header: string(header),
		Body: string(body),
		Form: string(form),
		PostForm: string(postForm),
	})
}

func GetUrls(w http.ResponseWriter, r *http.Request) {
	js, err := json.MarshalIndent(db.GetUrls(), "", "\t")
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write(js)
}

func main() {
	db = database.New("db.sqlite3")
	r := http.NewServeMux()
	r.HandleFunc("/get_urls", GetUrls)
	r.Handle("/", handlers.LoggingHandler(os.Stdout, http.HandlerFunc(handler)))
	http.ListenAndServe(":8080", handlers.CompressHandler(r))
}<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"sync"

	"github.com/skratchdot/open-golang/open"
	"github.com/spf13/cobra"

	"github.com/Shopify/themekit/kit"
)

var openCmd = &cobra.Command{
	Use:   "open",
	Short: "Open the preview for your store.",
	Long: `Open will open the preview page in your browser as well as print out
url for your reference`,
	RunE: forEachClient(preview),
}

func preview(client kit.ThemeClient, filenames []string, wg *sync.WaitGroup) {
	defer wg.Done()
	previewURL := fmt.Sprintf("https://%s?preview_theme_id=%s",
		client.Config.Domain,
		client.Config.ThemeID)

	kit.Printf("[%s] opening %s",
		kit.GreenText(environment),
		kit.GreenText(previewURL))

	err := open.Run(previewURL)
	if err != nil {
		kit.LogErrorf("[%s] %s",
			kit.GreenText(environment),
			kit.RedText(err))
	}
}
<commit_msg>Logging the environment that the open command is opening<commit_after>package cmd

import (
	"fmt"
	"sync"

	"github.com/skratchdot/open-golang/open"
	"github.com/spf13/cobra"

	"github.com/Shopify/themekit/kit"
)

var openCmd = &cobra.Command{
	Use:   "open",
	Short: "Open the preview for your store.",
	Long: `Open will open the preview page in your browser as well as print out
url for your reference`,
	RunE: forEachClient(preview),
}

func preview(client kit.ThemeClient, filenames []string, wg *sync.WaitGroup) {
	defer wg.Done()
	previewURL := fmt.Sprintf("https://%s?preview_theme_id=%s",
		client.Config.Domain,
		client.Config.ThemeID)

	kit.Printf("[%s] opening %s",
		kit.GreenText(client.Config.Environment),
		kit.GreenText(previewURL))

	err := open.Run(previewURL)
	if err != nil {
		kit.LogErrorf("[%s] %s",
			kit.GreenText(client.Config.Environment),
			kit.RedText(err))
	}
}
<|endoftext|>


<commit_before>package brokerintegration_test

import (
	"encoding/json"
	"net/http"

	"github.com/pivotal-cf/brokerapi"
	"github.com/pivotal-cf/cf-redis-broker/integration"
	"github.com/pivotal-cf/cf-redis-broker/integration/helpers"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Catalog", func() {

	It("returns HTTP 200", func() {
		code, _ := brokerClient.MakeCatalogRequest()
		Ω(code).To(Equal(http.StatusOK))
	})

	var plans []brokerapi.ServicePlan
	var service brokerapi.Service

	Describe("Service", func() {

		BeforeEach(func() {
			_, body := brokerClient.MakeCatalogRequest()

			catalog := struct {
				Services []brokerapi.Service `json:"services"`
			}{}

			json.Unmarshal(body, &catalog)
			Ω(len(catalog.Services)).Should(Equal(1))

			service = catalog.Services[0]
			plans = service.Plans
		})

		It("displays the correct service name and id", func() {
			Ω(service.Name).Should(Equal("my-redis"))
			Ω(service.ID).Should(Equal("123456abcdef"))
		})

		It("displays the correct documentation URL", func() {
			Ω(service.Metadata.DocumentationUrl).Should(Equal("http://docs.pivotal.io/p1-services/Redis.html"))
		})

		It("displays the correct support URL", func() {
			Ω(service.Metadata.SupportUrl).Should(Equal("http://support.pivotal.io"))
		})

		It("displays the description", func() {
			Ω(service.Description).Should(Equal("Redis service to provide a key-value store"))
		})

		Describe("Shared-vm plan", func() {
			var plan brokerapi.ServicePlan

			BeforeEach(func() {
				for _, p := range plans {
					if p.Name == "shared-vm" {
						plan = p
					}
				}
			})

			It("has the correct id from the config file", func() {
				Ω(plan.ID).Should(Equal("C210CA06-E7E5-4F5D-A5AA-7A2C51CC290E"))
			})

			It("displays the correct description", func() {
				Ω(plan.Description).Should(Equal("This plan provides a single Redis process on a shared VM, which is suitable for development and testing workloads"))
			})

			It("displays the correct metadata bullet points", func() {
				Ω(plan.Metadata.Bullets).Should(Equal([]string{
					"Each instance shares the same VM",
					"Single dedicated Redis process",
					"Suitable for development & testing workloads",
				}))
			})
		})

		Describe("Dedicated-vm plan", func() {
			var plan brokerapi.ServicePlan

			BeforeEach(func() {
				for _, p := range plans {
					if p.Name == "dedicated-vm" {
						plan = p
					}
				}
			})

			It("has the correct id from the config file", func() {
				Ω(plan.ID).Should(Equal("74E8984C-5F8C-11E4-86BE-07807B3B2589"))
			})

			It("displays the correct description", func() {
				Ω(plan.Description).Should(Equal("This plan provides a single Redis process on a dedicated VM, which is suitable for production workloads"))
			})

			It("displays the correct metadata bullet points", func() {
				Ω(plan.Metadata.Bullets).Should(Equal([]string{
					"Dedicated VM per instance",
					"Single dedicated Redis process",
					"Suitable for production workloads",
				}))
			})
		})
	})

	Context("When there are no dedicated nodes", func() {
		BeforeEach(func() {
			switchBroker("broker.yml-no-dedicated")

			_, body := brokerClient.MakeCatalogRequest()

			catalog := struct {
				Services []brokerapi.Service `json:"services"`
			}{}

			json.Unmarshal(body, &catalog)
			Ω(len(catalog.Services)).Should(Equal(1))

			plans = catalog.Services[0].Plans
		})

		AfterEach(func() {
			switchBroker("broker.yml")
		})

		It("only shows the shared plan", func() {
			Ω(len(plans)).Should(Equal(1))

			sharedPlan := plans[0]
			Ω(sharedPlan.Name).Should(Equal("shared-vm"))
		})
	})

	Context("When there are dedicated nodes", func() {

		BeforeEach(func() {
			_, body := brokerClient.MakeCatalogRequest()

			catalog := struct {
				Services []brokerapi.Service `json:"services"`
			}{}

			json.Unmarshal(body, &catalog)
			Ω(len(catalog.Services)).Should(Equal(1))

			plans = catalog.Services[0].Plans
		})

		It("shows both plans", func() {
			Ω(len(plans)).Should(Equal(2))

			planNames := []string{}
			for _, plan := range plans {
				planNames = append(planNames, plan.Name)
			}

			Ω(planNames).Should(ContainElement("shared-vm"))
			Ω(planNames).Should(ContainElement("dedicated-vm"))
		})

		Context("When the service instance limit is set to zero", func() {
			BeforeEach(func() {
				switchBroker("broker.yml-no-shared")

				_, body := brokerClient.MakeCatalogRequest()

				catalog := struct {
					Services []brokerapi.Service `json:"services"`
				}{}

				json.Unmarshal(body, &catalog)
				Ω(len(catalog.Services)).Should(Equal(1))

				plans = catalog.Services[0].Plans
			})

			AfterEach(func() {
				switchBroker("broker.yml")
			})

			It("Only shows the dedicated plan", func() {
				Ω(len(plans)).Should(Equal(1))

				dedicatedPlan := plans[0]
				Ω(dedicatedPlan.Name).Should(Equal("dedicated-vm"))
			})
		})
	})
})

func switchBroker(config string) {
	helpers.KillProcess(brokerSession)
	helpers.ResetTestDirs()
	brokerSession = integration.LaunchProcessWithBrokerConfig(brokerExecutablePath, config)
	Ω(helpers.ServiceAvailable(brokerPort)).Should(BeTrue())
}
<commit_msg>Update documentation url in catalog test<commit_after>package brokerintegration_test

import (
	"encoding/json"
	"net/http"

	"github.com/pivotal-cf/brokerapi"
	"github.com/pivotal-cf/cf-redis-broker/integration"
	"github.com/pivotal-cf/cf-redis-broker/integration/helpers"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Catalog", func() {

	It("returns HTTP 200", func() {
		code, _ := brokerClient.MakeCatalogRequest()
		Ω(code).To(Equal(http.StatusOK))
	})

	var plans []brokerapi.ServicePlan
	var service brokerapi.Service

	Describe("Service", func() {

		BeforeEach(func() {
			_, body := brokerClient.MakeCatalogRequest()

			catalog := struct {
				Services []brokerapi.Service `json:"services"`
			}{}

			json.Unmarshal(body, &catalog)
			Ω(len(catalog.Services)).Should(Equal(1))

			service = catalog.Services[0]
			plans = service.Plans
		})

		It("displays the correct service name and id", func() {
			Ω(service.Name).Should(Equal("my-redis"))
			Ω(service.ID).Should(Equal("123456abcdef"))
		})

		It("displays the correct documentation URL", func() {
			Ω(service.Metadata.DocumentationUrl).Should(Equal("http://docs.pivotal.io/redis/index.html"))
		})

		It("displays the correct support URL", func() {
			Ω(service.Metadata.SupportUrl).Should(Equal("http://support.pivotal.io"))
		})

		It("displays the description", func() {
			Ω(service.Description).Should(Equal("Redis service to provide a key-value store"))
		})

		Describe("Shared-vm plan", func() {
			var plan brokerapi.ServicePlan

			BeforeEach(func() {
				for _, p := range plans {
					if p.Name == "shared-vm" {
						plan = p
					}
				}
			})

			It("has the correct id from the config file", func() {
				Ω(plan.ID).Should(Equal("C210CA06-E7E5-4F5D-A5AA-7A2C51CC290E"))
			})

			It("displays the correct description", func() {
				Ω(plan.Description).Should(Equal("This plan provides a single Redis process on a shared VM, which is suitable for development and testing workloads"))
			})

			It("displays the correct metadata bullet points", func() {
				Ω(plan.Metadata.Bullets).Should(Equal([]string{
					"Each instance shares the same VM",
					"Single dedicated Redis process",
					"Suitable for development & testing workloads",
				}))
			})
		})

		Describe("Dedicated-vm plan", func() {
			var plan brokerapi.ServicePlan

			BeforeEach(func() {
				for _, p := range plans {
					if p.Name == "dedicated-vm" {
						plan = p
					}
				}
			})

			It("has the correct id from the config file", func() {
				Ω(plan.ID).Should(Equal("74E8984C-5F8C-11E4-86BE-07807B3B2589"))
			})

			It("displays the correct description", func() {
				Ω(plan.Description).Should(Equal("This plan provides a single Redis process on a dedicated VM, which is suitable for production workloads"))
			})

			It("displays the correct metadata bullet points", func() {
				Ω(plan.Metadata.Bullets).Should(Equal([]string{
					"Dedicated VM per instance",
					"Single dedicated Redis process",
					"Suitable for production workloads",
				}))
			})
		})
	})

	Context("When there are no dedicated nodes", func() {
		BeforeEach(func() {
			switchBroker("broker.yml-no-dedicated")

			_, body := brokerClient.MakeCatalogRequest()

			catalog := struct {
				Services []brokerapi.Service `json:"services"`
			}{}

			json.Unmarshal(body, &catalog)
			Ω(len(catalog.Services)).Should(Equal(1))

			plans = catalog.Services[0].Plans
		})

		AfterEach(func() {
			switchBroker("broker.yml")
		})

		It("only shows the shared plan", func() {
			Ω(len(plans)).Should(Equal(1))

			sharedPlan := plans[0]
			Ω(sharedPlan.Name).Should(Equal("shared-vm"))
		})
	})

	Context("When there are dedicated nodes", func() {

		BeforeEach(func() {
			_, body := brokerClient.MakeCatalogRequest()

			catalog := struct {
				Services []brokerapi.Service `json:"services"`
			}{}

			json.Unmarshal(body, &catalog)
			Ω(len(catalog.Services)).Should(Equal(1))

			plans = catalog.Services[0].Plans
		})

		It("shows both plans", func() {
			Ω(len(plans)).Should(Equal(2))

			planNames := []string{}
			for _, plan := range plans {
				planNames = append(planNames, plan.Name)
			}

			Ω(planNames).Should(ContainElement("shared-vm"))
			Ω(planNames).Should(ContainElement("dedicated-vm"))
		})

		Context("When the service instance limit is set to zero", func() {
			BeforeEach(func() {
				switchBroker("broker.yml-no-shared")

				_, body := brokerClient.MakeCatalogRequest()

				catalog := struct {
					Services []brokerapi.Service `json:"services"`
				}{}

				json.Unmarshal(body, &catalog)
				Ω(len(catalog.Services)).Should(Equal(1))

				plans = catalog.Services[0].Plans
			})

			AfterEach(func() {
				switchBroker("broker.yml")
			})

			It("Only shows the dedicated plan", func() {
				Ω(len(plans)).Should(Equal(1))

				dedicatedPlan := plans[0]
				Ω(dedicatedPlan.Name).Should(Equal("dedicated-vm"))
			})
		})
	})
})

func switchBroker(config string) {
	helpers.KillProcess(brokerSession)
	helpers.ResetTestDirs()
	brokerSession = integration.LaunchProcessWithBrokerConfig(brokerExecutablePath, config)
	Ω(helpers.ServiceAvailable(brokerPort)).Should(BeTrue())
}
<|endoftext|>


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

package cmd

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"regexp"

	"github.com/mpopadic/go_n_find/colors"
	"github.com/spf13/cobra"
)

var (
	pathFlag              string
	nameFlag              string
	replaceFlag           string
	ignoreCaseFlag        bool
	showAbsolutePathsFlag bool
	forceReplaceFlag      bool
	contentFlag           string
)

var (
	_numberOfResults int
	_renameMap       map[string]string
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "go_n_find",
	Short: "CLI for finding files and folders",
	Long:  `CLI tool for finding files and folders by name or content`,
	PreRunE: func(cmd *cobra.Command, args []string) error {
		if pathFlag == "" {
			return fmt.Errorf("path flag is required")
		}
		if nameFlag == "" && contentFlag == "" {
			return fmt.Errorf("name flag or content flag are required")
		}
		return nil
	},
	// Uncomment the following line if your bare application
	// has an action associated with it:
	RunE: func(cmd *cobra.Command, args []string) error {

		// Set findOptions
		options := &findOptions{
			Path:              pathFlag,
			Name:              nameFlag,
			Content:           contentFlag,
			ReplaceWith:       replaceFlag,
			IgnoreCase:        ignoreCaseFlag,
			ShowAbsolutePaths: showAbsolutePathsFlag,
			ForceReplace:      forceReplaceFlag,
		}

		_numberOfResults = 0

		if options.ReplaceWith != "" && !options.ForceReplace {
			_renameMap = make(map[string]string)
		}

		if err := findInTree(options); err != nil {
			return err
		}

		colors.CYAN.Printf("Number of results: %d\n", _numberOfResults)
		if options.ReplaceWith != "" && !options.ForceReplace {
			response := waitResponse("Are you sure? [Yes/No] ", map[string][]string{
				"Yes": []string{"Yes", "Y", "y"},
				"No":  []string{"No", "N", "n"},
			})
			switch response {
			case "Yes":
				renamePaths(_renameMap)
			case "No":
				colors.RED.Print(response)
			}
		}

		return nil
	},
}

// Execute adds all child commands to the root command and 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() {
	colors.InitColors()

	RootCmd.Flags().StringVarP(&pathFlag, "path", "p", "", "path to directory")
	RootCmd.Flags().StringVarP(&nameFlag, "name", "n", "", "regular expression for matching file or directory name")
	RootCmd.Flags().StringVarP(&replaceFlag, "replace", "r", "", "replaces mached regular expression parts with given value")
	RootCmd.Flags().BoolVarP(&ignoreCaseFlag, "ignore-case", "i", false, "ignore case")
	RootCmd.Flags().BoolVarP(&showAbsolutePathsFlag, "absolute-paths", "a", false, "print absolute paths in result")
	RootCmd.Flags().BoolVarP(&forceReplaceFlag, "force-replace", "f", false, "Force replace without responding")

	RootCmd.Flags().StringVarP(&contentFlag, "content", "c", "", "regular expression for matching file content")

}

func findInTree(options *findOptions) error {
	fileInfo, err := os.Stat(options.Path)
	if err != nil {
		return fmt.Errorf("could not get fileInfo for %s: %v", options.Path, err)
	}

	if fileInfo.IsDir() {
		files, err := ioutil.ReadDir(options.Path)
		if err != nil {
			return fmt.Errorf("could not read directory %s: %v", options.Path, err)
		}
		for _, file := range files {
			childOptions := options.CreateCopy()
			childOptions.Path = path.Join(options.Path, file.Name())
			findInTree(childOptions)
		}
	}

	doAction(options, fileInfo.Name())
	return nil
}

func doAction(options *findOptions, fileName string) {
	if options.Name != "" {
		var finalPathPrint = ""
		absolutePath, err := filepath.Abs(options.Path)
		if err != nil {
			log.Fatalf("could not get absolute path: %v", err)
		}
		if options.ShowAbsolutePaths {
			finalPathPrint = absolutePath
		} else {
			finalPathPrint = options.Path
		}
		finalPathPrint = filepath.Clean(finalPathPrint)

		// re := regexp.MustCompile(options.Name)
		re, err := regexp.Compile(options.Name)
		if err != nil {
			colors.RED.Printf("regular expresion for name flag is not valid\n")
			os.Exit(1)
		}
		if options.IgnoreCase {
			re, err = regexp.Compile("(?i)" + options.Name)
			if err != nil {
				colors.RED.Printf("regular expresion for name flag is not valid\n")
				os.Exit(1)
			}
		}

		if re.MatchString(fileName) {
			_numberOfResults++
			if options.ReplaceWith != "" {
				pathDir := filepath.Dir(absolutePath)
				newFileName := re.ReplaceAllString(fileName, options.ReplaceWith)

				if options.ForceReplace {
					err := os.Rename(absolutePath, filepath.FromSlash(path.Join(pathDir, newFileName)))
					if err != nil {
						fmt.Printf("could not rename file: %v", err)
					}
					colors.RED.Print(absolutePath)
					colors.CYAN.Print(" => ")
					colors.GREEN.Println(filepath.FromSlash(path.Join(pathDir, newFileName)))
				} else {
					_renameMap[absolutePath] = filepath.FromSlash(path.Join(pathDir, newFileName))

					fmt.Print(absolutePath)
					colors.CYAN.Print(" => ")
					fmt.Println(filepath.FromSlash(path.Join(pathDir, newFileName)))
				}
			} else {
				fmt.Println(filepath.FromSlash(finalPathPrint))
			}
		}
	}
	if options.Content != "" {

	}
}

type findOptions struct {
	Path              string
	Name              string
	Content           string
	ReplaceWith       string
	IgnoreCase        bool
	ShowAbsolutePaths bool
	ForceReplace      bool
}

func (o *findOptions) CreateCopy() *findOptions {
	newFindOptions := &findOptions{
		Path:              o.Path,
		Name:              o.Name,
		Content:           o.Content,
		ReplaceWith:       o.ReplaceWith,
		IgnoreCase:        o.IgnoreCase,
		ShowAbsolutePaths: o.ShowAbsolutePaths,
		ForceReplace:      o.ForceReplace,
	}
	return newFindOptions
}

func waitResponse(question string, responseAliases map[string][]string) string {
	colors.YELLOW.Printf("%s ", question)
	var respond string

	for {
		fmt.Scanf("%s\n", &respond)

		for response, aliases := range responseAliases {
			for _, alias := range aliases {
				if respond == alias {
					return response
				}
			}
		}

		colors.YELLOW.Printf("%s ", question)
	}
}

func renamePaths(paths map[string]string) error {
	for oldPath, newPath := range paths {
		err := os.Rename(oldPath, newPath)
		if err != nil {
			return fmt.Errorf("could not rename file: %v", err)
		}
		colors.RED.Print(oldPath)
		colors.CYAN.Print(" => ")
		colors.GREEN.Println(newPath)
	}
	return nil
}
<commit_msg>added pritty print for find content<commit_after>// Copyright © 2017 NAME HERE <EMAIL ADDRESS>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/mpopadic/go_n_find/colors"
	"github.com/spf13/cobra"
)

var (
	pathFlag              string
	nameFlag              string
	replaceFlag           string
	ignoreCaseFlag        bool
	showAbsolutePathsFlag bool
	forceReplaceFlag      bool
	contentFlag           string
)

var (
	_numberOfResults int
	_renameMap       map[string]string
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "go_n_find",
	Short: "CLI for finding files and folders",
	Long:  `CLI tool for finding files and folders by name or content`,
	PreRunE: func(cmd *cobra.Command, args []string) error {
		if pathFlag == "" {
			return fmt.Errorf("path flag is required")
		}
		if nameFlag == "" && contentFlag == "" {
			return fmt.Errorf("name flag or content flag are required")
		}
		return nil
	},
	// Uncomment the following line if your bare application
	// has an action associated with it:
	RunE: func(cmd *cobra.Command, args []string) error {

		// Set findOptions
		options := &findOptions{
			Path:              pathFlag,
			Name:              nameFlag,
			Content:           contentFlag,
			ReplaceWith:       replaceFlag,
			IgnoreCase:        ignoreCaseFlag,
			ShowAbsolutePaths: showAbsolutePathsFlag,
			ForceReplace:      forceReplaceFlag,
		}

		_numberOfResults = 0

		if options.ReplaceWith != "" && !options.ForceReplace {
			_renameMap = make(map[string]string)
		}

		if err := findInTree(options); err != nil {
			return err
		}

		colors.CYAN.Printf("Number of results: %d\n", _numberOfResults)
		if options.ReplaceWith != "" && !options.ForceReplace {
			response := waitResponse("Are you sure? [Yes/No] ", map[string][]string{
				"Yes": []string{"Yes", "Y", "y"},
				"No":  []string{"No", "N", "n"},
			})
			switch response {
			case "Yes":
				renamePaths(_renameMap)
			case "No":
				colors.RED.Print(response)
			}
		}

		return nil
	},
}

// Execute adds all child commands to the root command and 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() {
	colors.InitColors()

	RootCmd.Flags().StringVarP(&pathFlag, "path", "p", "", "path to directory")
	RootCmd.Flags().StringVarP(&nameFlag, "name", "n", "", "regular expression for matching file or directory name")
	RootCmd.Flags().StringVarP(&replaceFlag, "replace", "r", "", "replaces mached regular expression parts with given value")
	RootCmd.Flags().BoolVarP(&ignoreCaseFlag, "ignore-case", "i", false, "ignore case")
	RootCmd.Flags().BoolVarP(&showAbsolutePathsFlag, "absolute-paths", "a", false, "print absolute paths in result")
	RootCmd.Flags().BoolVarP(&forceReplaceFlag, "force-replace", "f", false, "Force replace without responding")

	RootCmd.Flags().StringVarP(&contentFlag, "content", "c", "", "regular expression for matching file content")

}

func findInTree(options *findOptions) error {
	fileInfo, err := os.Stat(options.Path)
	if err != nil {
		return fmt.Errorf("could not get fileInfo for %s: %v", options.Path, err)
	}

	if fileInfo.IsDir() {
		files, err := ioutil.ReadDir(options.Path)
		if err != nil {
			return fmt.Errorf("could not read directory %s: %v", options.Path, err)
		}
		for _, file := range files {
			childOptions := options.CreateCopy()
			childOptions.Path = path.Join(options.Path, file.Name())
			findInTree(childOptions)
		}
	}

	doAction(options, fileInfo)
	return nil
}

func doAction(options *findOptions, fileInfo os.FileInfo) {
	absolutePath, err := filepath.Abs(options.Path)
	if err != nil {
		log.Fatalf("could not get absolute path: %v", err)
	}
	finalPathPrint := getPathPrintFormat(options.Path, absolutePath, options.ShowAbsolutePaths)

	if options.Name != "" {
		re := createRegex(options.Name, options.IgnoreCase)

		if re.MatchString(fileInfo.Name()) {
			_numberOfResults++
			if options.ReplaceWith != "" {
				pathDir := filepath.Dir(absolutePath)
				newFileName := re.ReplaceAllString(fileInfo.Name(), options.ReplaceWith)

				if options.ForceReplace {
					err := os.Rename(absolutePath, filepath.FromSlash(path.Join(pathDir, newFileName)))
					if err != nil {
						fmt.Printf("could not rename file: %v", err)
					}
					colors.RED.Print(absolutePath)
					colors.CYAN.Print(" => ")
					colors.GREEN.Println(filepath.FromSlash(path.Join(pathDir, newFileName)))
				} else {
					_renameMap[absolutePath] = filepath.FromSlash(path.Join(pathDir, newFileName))

					fmt.Print(absolutePath)
					colors.CYAN.Print(" => ")
					fmt.Println(filepath.FromSlash(path.Join(pathDir, newFileName)))
				}
			} else {
				fmt.Println(filepath.FromSlash(finalPathPrint))
			}
		}
	}
	if options.Content != "" {
		if !fileInfo.IsDir() {
			re := createRegex(options.Content, options.IgnoreCase)

			fileBytes, err := ioutil.ReadFile(absolutePath)
			if err != nil {
				log.Fatalf("could not read file content: %v", err)
			}
			fileString := string(fileBytes)

			fileLines := strings.Split(fileString, "\n")

			printedFileName := false
			for lineNumber, line := range fileLines {
				if re.MatchString(line) {
					_numberOfResults++
					if !printedFileName {
						colors.CYAN.Printf("%s:\n", finalPathPrint)
						printedFileName = !printedFileName
					}
					allIndexes := re.FindAllStringIndex(line, -1)

					colors.YELLOW.Printf("%v:", lineNumber+1)
					location := 0
					for _, match := range allIndexes {
						fmt.Printf("%s", line[location:match[0]])
						colors.GREEN.Printf("%s", line[match[0]:match[1]])
						location = match[1]
					}
					fmt.Println()
				}
			}
		}
	}
}

type findOptions struct {
	Path              string
	Name              string
	Content           string
	ReplaceWith       string
	IgnoreCase        bool
	ShowAbsolutePaths bool
	ForceReplace      bool
}

func (o *findOptions) CreateCopy() *findOptions {
	newFindOptions := &findOptions{
		Path:              o.Path,
		Name:              o.Name,
		Content:           o.Content,
		ReplaceWith:       o.ReplaceWith,
		IgnoreCase:        o.IgnoreCase,
		ShowAbsolutePaths: o.ShowAbsolutePaths,
		ForceReplace:      o.ForceReplace,
	}
	return newFindOptions
}

func waitResponse(question string, responseAliases map[string][]string) string {
	colors.YELLOW.Printf("%s ", question)
	var respond string

	for {
		fmt.Scanf("%s\n", &respond)

		for response, aliases := range responseAliases {
			for _, alias := range aliases {
				if respond == alias {
					return response
				}
			}
		}

		colors.YELLOW.Printf("%s ", question)
	}
}

func renamePaths(paths map[string]string) error {
	for oldPath, newPath := range paths {
		err := os.Rename(oldPath, newPath)
		if err != nil {
			return fmt.Errorf("could not rename file: %v", err)
		}
		colors.RED.Print(oldPath)
		colors.CYAN.Print(" => ")
		colors.GREEN.Println(newPath)
	}
	return nil
}

func getPathPrintFormat(filePath, absolutePath string, showAbsolute bool) string {
	var result = ""
	if showAbsolute {
		result = absolutePath
	} else {
		result = filePath
	}
	return filepath.Clean(result)
}

func createRegex(text string, ignoreCase bool) *regexp.Regexp {
	re, err := regexp.Compile(text)
	if err != nil {
		colors.RED.Printf("regular expresion for name flag is not valid\n")
		os.Exit(1)
	}
	if ignoreCase {
		re, err = regexp.Compile("(?i)" + text)
		if err != nil {
			colors.RED.Printf("regular expresion for name flag is not valid\n")
			os.Exit(1)
		}
	}
	return re
}
<|endoftext|>


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

import (
	"fmt"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/zmap/zdns/internal/util"
	"github.com/zmap/zdns/pkg/zdns"
)

var cfgFile string
var GC zdns.GlobalConf

//TODO: these options may need to be set as flags or in GC, to standardize.
var (
	Servers_string   string
	Localaddr_string string
	Localif_string   string
	Config_file      string
	Timeout          int
	IterationTimeout int
	Class_string     string
	NanoSeconds      bool
)

// rootCmd represents the base command when called without any subcommands
var rootCmd = &cobra.Command{
	Use:   "zdns",
	Short: "High-speed, low-drag DNS lookups",
	Long: `ZDNS is a library and CLI tool for making very fast DNS requests. It's built upon
https://github.com/zmap/dns (and in turn https://github.com/miekg/dns) for constructing
and parsing raw DNS packets.

ZDNS also includes its own recursive resolution and a cache to further optimize performance.`,
	ValidArgs: zdns.Validlookups(),
	Args:      cobra.ExactValidArgs(1),
	Run: func(cmd *cobra.Command, args []string) {
		GC.Module = strings.ToUpper(args[0])
		zdns.Run(GC, cmd.Flags(),
			&Timeout, &IterationTimeout,
			&Class_string, &Servers_string,
			&Config_file, &Localaddr_string,
			&Localif_string, &NanoSeconds)
	},
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	err := rootCmd.Execute()
	if err != nil {
		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/.zdns.yaml)")

	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	rootCmd.PersistentFlags().IntVar(&GC.Threads, "threads", 1000, "number of lightweight go threads")
	rootCmd.PersistentFlags().IntVar(&GC.GoMaxProcs, "go-processes", 0, "number of OS processes (GOMAXPROCS)")
	rootCmd.PersistentFlags().StringVar(&GC.NamePrefix, "prefix", "", "name to be prepended to what's passed in (e.g., www.)")
	rootCmd.PersistentFlags().StringVar(&GC.NameOverride, "override-name", "", "name overrides all passed in names")
	rootCmd.PersistentFlags().BoolVar(&GC.AlexaFormat, "alexa", false, "is input file from Alexa Top Million download")
	rootCmd.PersistentFlags().BoolVar(&GC.MetadataFormat, "metadata-passthrough", false, "if input records have the form 'name,METADATA', METADATA will be propagated to the output")
	rootCmd.PersistentFlags().BoolVar(&GC.IterativeResolution, "iterative", false, "Perform own iteration instead of relying on recursive resolver")
	rootCmd.PersistentFlags().StringVar(&GC.InputFilePath, "input-file", "-", "names to read")
	rootCmd.PersistentFlags().StringVar(&GC.OutputFilePath, "output-file", "-", "where should JSON output be saved")
	rootCmd.PersistentFlags().StringVar(&GC.MetadataFilePath, "metadata-file", "", "where should JSON metadata be saved")
	rootCmd.PersistentFlags().StringVar(&GC.LogFilePath, "log-file", "", "where should JSON logs be saved")

	rootCmd.PersistentFlags().StringVar(&GC.ResultVerbosity, "result-verbosity", "normal", "Sets verbosity of each output record. Options: short, normal, long, trace")
	rootCmd.PersistentFlags().StringVar(&GC.IncludeInOutput, "include-fields", "", "Comma separated list of fields to additionally output beyond result verbosity. Options: class, protocol, ttl, resolver, flags")

	rootCmd.PersistentFlags().IntVar(&GC.Verbosity, "verbosity", 3, "log verbosity: 1 (lowest)--5 (highest)")
	rootCmd.PersistentFlags().IntVar(&GC.Retries, "retries", 1, "how many times should zdns retry query if timeout or temporary failure")
	rootCmd.PersistentFlags().IntVar(&GC.MaxDepth, "max-depth", 10, "how deep should we recurse when performing iterative lookups")
	rootCmd.PersistentFlags().IntVar(&GC.CacheSize, "cache-size", 10000, "how many items can be stored in internal recursive cache")
	rootCmd.PersistentFlags().BoolVar(&GC.TCPOnly, "tcp-only", false, "Only perform lookups over TCP")
	rootCmd.PersistentFlags().BoolVar(&GC.UDPOnly, "udp-only", false, "Only perform lookups over UDP")
	rootCmd.PersistentFlags().BoolVar(&GC.RecycleSockets, "recycle-sockets", true, "Create long-lived UDP for each thread at launch and reuse for all queries")
	rootCmd.PersistentFlags().BoolVar(&GC.NameServerMode, "name-server-mode", false, "Treats input as nameservers to query with a static query rather than queries to send to a static name server")

	rootCmd.PersistentFlags().StringVar(&Servers_string, "name-servers", "", "List of DNS servers to use. Can be passed as comma-delimited string or via @/path/to/file. If no port is specified, defaults to 53.")
	rootCmd.PersistentFlags().StringVar(&Localaddr_string, "local-addr", "", "comma-delimited list of local addresses to use")
	rootCmd.PersistentFlags().StringVar(&Localif_string, "local-interface", "", "local interface to use")
	rootCmd.PersistentFlags().StringVar(&Config_file, "conf-file", "/etc/resolv.conf", "config file for DNS servers")
	rootCmd.PersistentFlags().IntVar(&Timeout, "timeout", 15, "timeout for resolving an individual name")
	rootCmd.PersistentFlags().IntVar(&IterationTimeout, "iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
	rootCmd.PersistentFlags().StringVar(&Class_string, "class", "INET", "DNS class to query. Options: INET, CSNET, CHAOS, HESIOD, NONE, ANY. Default: INET.")
	rootCmd.PersistentFlags().BoolVar(&NanoSeconds, "nanoseconds", false, "Use nanosecond resolution timestamps")

	rootCmd.PersistentFlags().Bool("ipv4-lookup", false, "Perform an IPv4 Lookup in modules")
	rootCmd.PersistentFlags().Bool("ipv6-lookup", false, "Perform an IPv6 Lookup in modules")
	rootCmd.PersistentFlags().String("blacklist-file", "", "blacklist file for servers to exclude from lookups")
	rootCmd.PersistentFlags().Int("mx-cache-size", 1000, "number of records to store in MX -> A/AAAA cache")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" {
		// Use config file from the flag.
		viper.SetConfigFile(cfgFile)
	} else {
		// Find home directory.
		home, err := os.UserHomeDir()
		cobra.CheckErr(err)

		// Search config in home directory with name ".zdns" (without extension).
		viper.AddConfigPath(home)
		viper.SetConfigType("yaml")
		viper.SetConfigName(".zdns")
	}

	viper.SetEnvPrefix(util.EnvPrefix)
	viper.AutomaticEnv()

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed())
	}
	// Bind the current command's flags to viper
	util.BindFlags(rootCmd, viper.GetViper(), util.EnvPrefix)
}
<commit_msg>clarification<commit_after>/*
 * ZDNS Copyright 2016 Regents of the University of Michigan
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */
package cmd

import (
	"fmt"
	"os"
	"strings"

	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"github.com/zmap/zdns/internal/util"
	"github.com/zmap/zdns/pkg/zdns"
)

var cfgFile string
var GC zdns.GlobalConf

//TODO: these options may need to be set as flags or in GC, to standardize.
var (
	Servers_string   string
	Localaddr_string string
	Localif_string   string
	Config_file      string
	Timeout          int
	IterationTimeout int
	Class_string     string
	NanoSeconds      bool
)

// rootCmd represents the base command when called without any subcommands
var rootCmd = &cobra.Command{
	Use:   "zdns",
	Short: "High-speed, low-drag DNS lookups",
	Long: `ZDNS is a library and CLI tool for making very fast DNS requests. It's built upon
https://github.com/zmap/dns (and in turn https://github.com/miekg/dns) for constructing
and parsing raw DNS packets.

ZDNS also includes its own recursive resolution and a cache to further optimize performance.`,
	ValidArgs: zdns.Validlookups(),
	Args:      cobra.ExactValidArgs(1),
	Run: func(cmd *cobra.Command, args []string) {
		GC.Module = strings.ToUpper(args[0])
		zdns.Run(GC, cmd.Flags(),
			&Timeout, &IterationTimeout,
			&Class_string, &Servers_string,
			&Config_file, &Localaddr_string,
			&Localif_string, &NanoSeconds)
	},
}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	err := rootCmd.Execute()
	if err != nil {
		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/.zdns.yaml)")

	// Cobra also supports local flags, which will only run
	// when this action is called directly.
	rootCmd.PersistentFlags().IntVar(&GC.Threads, "threads", 1000, "number of lightweight go threads")
	rootCmd.PersistentFlags().IntVar(&GC.GoMaxProcs, "go-processes", 0, "number of OS processes (GOMAXPROCS)")
	rootCmd.PersistentFlags().StringVar(&GC.NamePrefix, "prefix", "", "name to be prepended to what's passed in (e.g., www.)")
	rootCmd.PersistentFlags().StringVar(&GC.NameOverride, "override-name", "", "name overrides all passed in names")
	rootCmd.PersistentFlags().BoolVar(&GC.AlexaFormat, "alexa", false, "is input file from Alexa Top Million download")
	rootCmd.PersistentFlags().BoolVar(&GC.MetadataFormat, "metadata-passthrough", false, "if input records have the form 'name,METADATA', METADATA will be propagated to the output")
	rootCmd.PersistentFlags().BoolVar(&GC.IterativeResolution, "iterative", false, "Perform own iteration instead of relying on recursive resolver")
	rootCmd.PersistentFlags().StringVar(&GC.InputFilePath, "input-file", "-", "names to read")
	rootCmd.PersistentFlags().StringVar(&GC.OutputFilePath, "output-file", "-", "where should JSON output be saved")
	rootCmd.PersistentFlags().StringVar(&GC.MetadataFilePath, "metadata-file", "", "where should JSON metadata be saved")
	rootCmd.PersistentFlags().StringVar(&GC.LogFilePath, "log-file", "", "where should JSON logs be saved")

	rootCmd.PersistentFlags().StringVar(&GC.ResultVerbosity, "result-verbosity", "normal", "Sets verbosity of each output record. Options: short, normal, long, trace")
	rootCmd.PersistentFlags().StringVar(&GC.IncludeInOutput, "include-fields", "", "Comma separated list of fields to additionally output beyond result verbosity. Options: class, protocol, ttl, resolver, flags")

	rootCmd.PersistentFlags().IntVar(&GC.Verbosity, "verbosity", 3, "log verbosity: 1 (lowest)--5 (highest)")
	rootCmd.PersistentFlags().IntVar(&GC.Retries, "retries", 1, "how many times should zdns retry query if timeout or temporary failure")
	rootCmd.PersistentFlags().IntVar(&GC.MaxDepth, "max-depth", 10, "how deep should we recurse when performing iterative lookups")
	rootCmd.PersistentFlags().IntVar(&GC.CacheSize, "cache-size", 10000, "how many items can be stored in internal recursive cache")
	rootCmd.PersistentFlags().BoolVar(&GC.TCPOnly, "tcp-only", false, "Only perform lookups over TCP")
	rootCmd.PersistentFlags().BoolVar(&GC.UDPOnly, "udp-only", false, "Only perform lookups over UDP")
	rootCmd.PersistentFlags().BoolVar(&GC.RecycleSockets, "recycle-sockets", true, "Create long-lived unbound UDP socket for each thread at launch and reuse for all (UDP) queries")
	rootCmd.PersistentFlags().BoolVar(&GC.NameServerMode, "name-server-mode", false, "Treats input as nameservers to query with a static query rather than queries to send to a static name server")

	rootCmd.PersistentFlags().StringVar(&Servers_string, "name-servers", "", "List of DNS servers to use. Can be passed as comma-delimited string or via @/path/to/file. If no port is specified, defaults to 53.")
	rootCmd.PersistentFlags().StringVar(&Localaddr_string, "local-addr", "", "comma-delimited list of local addresses to use")
	rootCmd.PersistentFlags().StringVar(&Localif_string, "local-interface", "", "local interface to use")
	rootCmd.PersistentFlags().StringVar(&Config_file, "conf-file", "/etc/resolv.conf", "config file for DNS servers")
	rootCmd.PersistentFlags().IntVar(&Timeout, "timeout", 15, "timeout for resolving an individual name")
	rootCmd.PersistentFlags().IntVar(&IterationTimeout, "iteration-timeout", 4, "timeout for resolving a single iteration in an iterative query")
	rootCmd.PersistentFlags().StringVar(&Class_string, "class", "INET", "DNS class to query. Options: INET, CSNET, CHAOS, HESIOD, NONE, ANY. Default: INET.")
	rootCmd.PersistentFlags().BoolVar(&NanoSeconds, "nanoseconds", false, "Use nanosecond resolution timestamps")

	rootCmd.PersistentFlags().Bool("ipv4-lookup", false, "Perform an IPv4 Lookup in modules")
	rootCmd.PersistentFlags().Bool("ipv6-lookup", false, "Perform an IPv6 Lookup in modules")
	rootCmd.PersistentFlags().String("blacklist-file", "", "blacklist file for servers to exclude from lookups")
	rootCmd.PersistentFlags().Int("mx-cache-size", 1000, "number of records to store in MX -> A/AAAA cache")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" {
		// Use config file from the flag.
		viper.SetConfigFile(cfgFile)
	} else {
		// Find home directory.
		home, err := os.UserHomeDir()
		cobra.CheckErr(err)

		// Search config in home directory with name ".zdns" (without extension).
		viper.AddConfigPath(home)
		viper.SetConfigType("yaml")
		viper.SetConfigName(".zdns")
	}

	viper.SetEnvPrefix(util.EnvPrefix)
	viper.AutomaticEnv()

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Fprintln(os.Stderr, "Using config file:", viper.ConfigFileUsed())
	}
	// Bind the current command's flags to viper
	util.BindFlags(rootCmd, viper.GetViper(), util.EnvPrefix)
}
<|endoftext|>