<commit_before>// Copyright 2016 CoreOS, Inc.
//
// 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 tcpproxy

import (
	"io"
	"net"
	"sync"
	"time"
)

type remote struct {
	mu       sync.Mutex
	addr     string
	inactive bool
}

func (r *remote) inactivate() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.inactive = true
}

func (r *remote) tryReactivate() {
	conn, err := net.Dial("tcp", r.addr)
	if err != nil {
		return
	}
	conn.Close()
	r.mu.Lock()
	defer r.mu.Unlock()
	r.inactive = false
	return
}

func (r *remote) isActive() bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return !r.inactive
}

type TCPProxy struct {
	Listener        net.Listener
	Endpoints       []string
	MonitorInterval time.Duration

	donec chan struct{}

	mu         sync.Mutex // guards the following fields
	remotes    []*remote
	nextRemote int
}

func (tp *TCPProxy) Run() error {
	tp.donec = make(chan struct{})
	if tp.MonitorInterval == 0 {
		tp.MonitorInterval = 5 * time.Minute
	}
	for _, ep := range tp.Endpoints {
		tp.remotes = append(tp.remotes, &remote{addr: ep})
	}

	go tp.runMonitor()
	for {
		in, err := tp.Listener.Accept()
		if err != nil {
			return err
		}

		go tp.serve(in)
	}
}

func (tp *TCPProxy) numRemotes() int {
	tp.mu.Lock()
	defer tp.mu.Unlock()
	return len(tp.remotes)
}

func (tp *TCPProxy) serve(in net.Conn) {
	var (
		err error
		out net.Conn
	)

	for i := 0; i < tp.numRemotes(); i++ {
		remote := tp.pick()
		if !remote.isActive() {
			continue
		}
		// TODO: add timeout
		out, err = net.Dial("tcp", remote.addr)
		if err == nil {
			break
		}
		remote.inactivate()
	}

	if out == nil {
		in.Close()
		return
	}

	go func() {
		io.Copy(in, out)
		in.Close()
		out.Close()
	}()

	io.Copy(out, in)
	out.Close()
	in.Close()
}

// pick picks a remote in round-robin fashion
func (tp *TCPProxy) pick() *remote {
	tp.mu.Lock()
	defer tp.mu.Unlock()

	picked := tp.remotes[tp.nextRemote]
	tp.nextRemote = (tp.nextRemote + 1) % len(tp.remotes)
	return picked
}

func (tp *TCPProxy) runMonitor() {
	for {
		select {
		case <-time.After(tp.MonitorInterval):
			tp.mu.Lock()
			for _, r := range tp.remotes {
				if !r.isActive() {
					go r.tryReactivate()
				}
			}
			tp.mu.Unlock()
		case <-tp.donec:
			return
		}
	}
}

func (tp *TCPProxy) Stop() {
	// graceful shutdown?
	// shutdown current connections?
	tp.Listener.Close()
	close(tp.donec)
}
<commit_msg>proxy/tcpproxy: add more logs<commit_after>// Copyright 2016 CoreOS, Inc.
//
// 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 tcpproxy

import (
	"io"
	"net"
	"sync"
	"time"

	"github.com/coreos/pkg/capnslog"
)

var (
	plog = capnslog.NewPackageLogger("github.com/coreos/etcd/proxy", "tcpproxy")
)

type remote struct {
	mu       sync.Mutex
	addr     string
	inactive bool
}

func (r *remote) inactivate() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.inactive = true
}

func (r *remote) tryReactivate() error {
	conn, err := net.Dial("tcp", r.addr)
	if err != nil {
		return err
	}
	conn.Close()
	r.mu.Lock()
	defer r.mu.Unlock()
	r.inactive = false
	return nil
}

func (r *remote) isActive() bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return !r.inactive
}

type TCPProxy struct {
	Listener        net.Listener
	Endpoints       []string
	MonitorInterval time.Duration

	donec chan struct{}

	mu         sync.Mutex // guards the following fields
	remotes    []*remote
	nextRemote int
}

func (tp *TCPProxy) Run() error {
	tp.donec = make(chan struct{})
	if tp.MonitorInterval == 0 {
		tp.MonitorInterval = 5 * time.Minute
	}
	for _, ep := range tp.Endpoints {
		tp.remotes = append(tp.remotes, &remote{addr: ep})
	}

	go tp.runMonitor()
	for {
		in, err := tp.Listener.Accept()
		if err != nil {
			return err
		}

		go tp.serve(in)
	}
}

func (tp *TCPProxy) numRemotes() int {
	tp.mu.Lock()
	defer tp.mu.Unlock()
	return len(tp.remotes)
}

func (tp *TCPProxy) serve(in net.Conn) {
	var (
		err error
		out net.Conn
	)

	for i := 0; i < tp.numRemotes(); i++ {
		remote := tp.pick()
		if !remote.isActive() {
			continue
		}
		// TODO: add timeout
		out, err = net.Dial("tcp", remote.addr)
		if err == nil {
			break
		}
		remote.inactivate()
		plog.Warningf("deactivated endpoint [%s] due to %v for %v", remote.addr, err, tp.MonitorInterval)
	}

	if out == nil {
		in.Close()
		return
	}

	go func() {
		io.Copy(in, out)
		in.Close()
		out.Close()
	}()

	io.Copy(out, in)
	out.Close()
	in.Close()
}

// pick picks a remote in round-robin fashion
func (tp *TCPProxy) pick() *remote {
	tp.mu.Lock()
	defer tp.mu.Unlock()

	picked := tp.remotes[tp.nextRemote]
	tp.nextRemote = (tp.nextRemote + 1) % len(tp.remotes)
	return picked
}

func (tp *TCPProxy) runMonitor() {
	for {
		select {
		case <-time.After(tp.MonitorInterval):
			tp.mu.Lock()
			for _, r := range tp.remotes {
				if !r.isActive() {
					go func() {
						if err := r.tryReactivate(); err != nil {
							plog.Warningf("failed to activate endpoint [%s] due to %v (stay inactive for another %v)", r.addr, err, tp.MonitorInterval)
						} else {
							plog.Printf("activated %s", r.addr)
						}
					}()
				}
			}
			tp.mu.Unlock()
		case <-tp.donec:
			return
		}
	}
}

func (tp *TCPProxy) Stop() {
	// graceful shutdown?
	// shutdown current connections?
	tp.Listener.Close()
	close(tp.donec)
}
<|endoftext|>


<commit_before>package translator_test

// Note: Test tests depend on a functioning parser, test it first.

import (
	"reflect"
	"strings"
	"testing"

	"github.com/bkidney/EQ2Dot/parser"
	"github.com/bkidney/EQ2Dot/translator"
)

func TestTranlator_ProduceDot(t *testing.T) {
	var tests = []struct {
		s   string
		out string
	}{
		// Simple pattern match
		{
			s:   `?node:ip:send(?srcIP,?destIP) and destIP in [203.0.113.12,192.168.1.100]`,
			out: "digraph query {\n\tnode [shape = none];\n\t0 [label = \"start\"];\n\tnode [shape = circle];\n\t0 -> 1;\n\t1 -> 2 [label = \"?node:ip:send(?srcIP,?destIP)\"];\n\tnode [shape = doublecircle];\n\t2 -> 3 [label = \"destIP in [203.0.113.12,192.168.1.100]\"];\n}",
		},
		// A more complete example
		{
			s: `
		 {DBServerNode:myDB:openSession(_):?sessionID Within
		  DBServerNode:myDb:userAuthenticate:(?user)} Precedes
			DBServerNode:myDB:sqlQuery(sessionID):?resultData Precedes
			  ?egressNode:ip::send(?outData,203.0.113.12)
		 and resultData FlowsTo* outData
		 `,
			out: "digraph query {\n\tnode [shape = none];\n\t0 [label = \"start\"];\n\tnode [shape = circle];\n\t0 -> 1;\n\t1 -> 2 [label = \"Call - DBServerNode:myDB:openSession(_):?sessionID\"];\n\tsubgraph cluster_0 {\n\trank = same;\n\tstyle=\"dashed\";\n\t2 -> 3 [label = \"DBServerNode:myDb:userAuthenticate:(?user)\"];\n\t}\t3 -> 4 [label = \"Ret - DBServerNode:myDB:openSession(_):?sessionID\"];\n\t4 -> 5 [label = \"DBServerNode:myDB:sqlQuery(sessionID):?resultData\"];\n\t5 -> 6 [label = \"?egressNode:ip::send(?outData,203.0.113.12)\"];\n\tnode [shape = doublecircle];\n\t6 -> 7 [label = \"resultData FlowsTo* outData\"];\n}",
		},
	}

	for i, tt := range tests {
		ast, _ := parser.NewParser(strings.NewReader(tt.s)).Parse()
		out := translator.NewTranslator(ast).Dot()
		if !reflect.DeepEqual(tt.out, out) {
			t.Errorf("%d. %q\noutput mismatch:\n\nexp=%#v\n\ngot=%#v\n\n", i, tt.s, tt.out, out)
		}
	}
}

// errstring returns the string representation of an error
func errstring(err error) string {
	if err != nil {
		return err.Error()
	}
	return ""
}
<commit_msg>Adds simple test case using or.<commit_after>package translator_test

// Note: Test tests depend on a functioning parser, test it first.

import (
	"reflect"
	"strings"
	"testing"

	"github.com/bkidney/EQ2Dot/parser"
	"github.com/bkidney/EQ2Dot/translator"
)

func TestTranlator_ProduceDot(t *testing.T) {
	var tests = []struct {
		s   string
		out string
	}{
		// Simple pattern match
		{
			s:   `?node:ip:send(?srcIP,?destIP) and destIP in [203.0.113.12,192.168.1.100]`,
			out: "digraph query {\n\tnode [shape = none];\n\t0 [label = \"start\"];\n\tnode [shape = circle];\n\t0 -> 1;\n\t1 -> 2 [label = \"?node:ip:send(?srcIP,?destIP)\"];\n\t2 -> 3 [label = \"destIP in [203.0.113.12,192.168.1.100]\"];\n\t3 [shape = doublecircle];\n}",
		},
		// A more complete example
		{
			s: `
		 {DBServerNode:myDB:openSession(_):?sessionID Within
		  DBServerNode:myDb:userAuthenticate:(?user)} Precedes
			DBServerNode:myDB:sqlQuery(sessionID):?resultData Precedes
			  ?egressNode:ip::send(?outData,203.0.113.12)
		 and resultData FlowsTo* outData
		 `,
			out: "digraph query {\n\tnode [shape = none];\n\t0 [label = \"start\"];\n\tnode [shape = circle];\n\t0 -> 1;\n\t1 -> 2 [label = \"Call - DBServerNode:myDB:openSession(_):?sessionID\"];\n\tsubgraph cluster_0 {\n\trank = same;\n\tstyle=\"dashed\";\n\t2 -> 3 [label = \"DBServerNode:myDb:userAuthenticate:(?user)\"];\n\t}\t3 -> 4 [label = \"Ret - DBServerNode:myDB:openSession(_):?sessionID\"];\n\t4 -> 5 [label = \"DBServerNode:myDB:sqlQuery(sessionID):?resultData\"];\n\t5 -> 6 [label = \"?egressNode:ip::send(?outData,203.0.113.12)\"];\n\t6 -> 7 [label = \"resultData FlowsTo* outData\"];\n\t7 [shape = doublecircle];\n}",
		},
		// A simple example using 'or' with simplified identifiers.
		{
			s:   `a or {b and c}`,
			out: "digraph query {\n\tnode [shape = none];\n\t0 [label = \"start\"];\n\tnode [shape = circle];\n\t0 -> 1;\n\t1 -> 3 [label = \"a\"];\n\t1 -> 2 [label = \"b\"];\n\t2 -> 3 [label = \"c\"];\n\t3 [shape = doublecircle];\n}",
		},
	}

	for i, tt := range tests {
		ast, _ := parser.NewParser(strings.NewReader(tt.s)).Parse()
		out := translator.NewTranslator(ast).Dot()
		if !reflect.DeepEqual(tt.out, out) {
			t.Errorf("%d. %q\noutput mismatch:\n\nexp=%#v\n\ngot=%#v\n\n", i, tt.s, tt.out, out)
		}
	}
}

// errstring returns the string representation of an error
func errstring(err error) string {
	if err != nil {
		return err.Error()
	}
	return ""
}
<|endoftext|>


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

package tchannel

import (
	"net"
	"time"

	"github.com/yarpc/yarpc-go/transport"

	"github.com/uber/tchannel-go"
	"golang.org/x/net/context"
)

// Inbound represents a TChannel Inbound. It is the same as a transport
// Inbound except it exposes the address on which the system is listening for
// connections.
type Inbound interface {
	transport.Inbound

	// Address on which the server is listening. Returns nil if Start has not
	// been called yet.
	Addr() net.Addr
}

// InboundOption configures Inbound.
type InboundOption func(*inbound)

// ListenAddr changes the address on which the TChannel server will listen for
// connections. By default, the server listens on an OS-assigned port.
//
// This option has no effect if the Chanel provided to NewInbound is already
// listening for connections when Start() is called.
func ListenAddr(addr string) InboundOption {
	return func(i *inbound) { i.addr = addr }
}

// NewInbound builds a new TChannel inbound from the given Channel.
func NewInbound(ch *tchannel.Channel, opts ...InboundOption) Inbound {
	i := &inbound{ch: ch}
	for _, opt := range opts {
		opt(i)
	}
	return i
}

type inbound struct {
	ch       *tchannel.Channel
	addr     string
	listener net.Listener
}

func (i *inbound) Start(h transport.Handler) error {
	i.ch.GetSubChannel(i.ch.ServiceName()).SetHandler(handler{h})

	if i.ch.State() == tchannel.ChannelListening {
		// Channel.Start() was called before RPC.Start(). We still want to
		// update the Handler and what i.addr means, but nothing else.
		i.addr = i.listener.Addr().String()
		return nil
	}

	// Default to ListenIP if addr wasn't given.
	addr := i.addr
	if addr == "" {
		listenIP, err := tchannel.ListenIP()
		if err != nil {
			return err
		}

		addr = listenIP.String() + ":0"
		// TODO(abg): Find a way to export this to users
	}

	// TODO(abg): If addr was just the port (":4040"), we want to use
	// ListenIP() + ":4040" rather than just ":4040".

	var err error
	i.listener, err = net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	i.addr = i.listener.Addr().String() // in case it changed

	if err := i.ch.Serve(i.listener); err != nil {
		return err
	}

	return nil
}

func (i *inbound) Stop() error {
	i.ch.Close()
	return nil
}

func (i *inbound) Addr() net.Addr {
	if i.listener == nil {
		return nil
	}
	return i.listener.Addr()
}

type handler struct{ Handler transport.Handler }

func (h handler) Handle(ctx context.Context, call *tchannel.InboundCall) {
	deadline, ok := ctx.Deadline()
	if !ok {
		call.Response().SendSystemError(tchannel.ErrTimeoutRequired)
		return
	}

	headers, err := readHeaders(call.Format(), call.Arg2Reader)
	if err != nil {
		call.Response().SendSystemError(tchannel.NewSystemError(
			tchannel.ErrCodeUnexpected, "failed to read headers: %v", err))
		return
	}

	body, err := call.Arg3Reader()
	if err != nil {
		call.Response().SendSystemError(tchannel.NewSystemError(
			tchannel.ErrCodeUnexpected, "failed to read body: %v", err))
		return
	}
	defer body.Close()

	rw := newResponseWriter(call)
	defer rw.Close() // TODO(abg): log if this errors

	treq := &transport.Request{
		Caller:    call.CallerName(),
		Service:   call.ServiceName(),
		Encoding:  transport.Encoding(call.Format()),
		Procedure: call.MethodString(),
		Headers:   headers,
		Body:      body,
		TTL:       deadline.Sub(time.Now()),
	}

	if err := h.Handler.Handle(ctx, treq, rw); err != nil {
		call.Response().SendSystemError(tchannel.NewSystemError(
			tchannel.ErrCodeUnexpected, "internal error: %v", err))
		return
	}
}

type responseWriter struct {
	failedWith   error
	bodyWriter   tchannel.ArgWriter
	format       tchannel.Format
	headers      transport.Headers
	response     *tchannel.InboundCallResponse
	wroteHeaders bool
}

func newResponseWriter(call *tchannel.InboundCall) *responseWriter {
	return &responseWriter{
		response: call.Response(),
		headers:  make(transport.Headers),
		format:   call.Format(),
	}
}

func (rw *responseWriter) AddHeaders(h transport.Headers) {
	if rw.wroteHeaders {
		panic("AddHeaders() cannot be called after calling Write().")
	}
	for k, v := range h {
		rw.headers.Set(k, v)
	}
}

func (rw *responseWriter) Write(s []byte) (int, error) {
	if rw.failedWith != nil {
		return 0, rw.failedWith
	}

	if !rw.wroteHeaders {
		rw.wroteHeaders = true
		if err := writeHeaders(rw.format, rw.headers, rw.response.Arg2Writer); err != nil {
			rw.failedWith = err
			return 0, err
		}
	}

	if rw.bodyWriter == nil {
		var err error
		rw.bodyWriter, err = rw.response.Arg3Writer()
		if err != nil {
			rw.failedWith = err
			return 0, err
		}
	}

	return rw.bodyWriter.Write(s)
}

func (rw *responseWriter) Close() error {
	if rw.bodyWriter != nil {
		return rw.bodyWriter.Close()
	}
	if rw.failedWith != nil {
		return rw.failedWith
	}
	return nil
}
<commit_msg>tchannel/inbound: fix nil pointer deref when channel is already listening<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"net"
	"time"

	"github.com/yarpc/yarpc-go/transport"

	"github.com/uber/tchannel-go"
	"golang.org/x/net/context"
)

// Inbound represents a TChannel Inbound. It is the same as a transport
// Inbound except it exposes the address on which the system is listening for
// connections.
type Inbound interface {
	transport.Inbound

	// Address on which the server is listening. Returns nil if Start has not
	// been called yet.
	Addr() net.Addr
}

// InboundOption configures Inbound.
type InboundOption func(*inbound)

// ListenAddr changes the address on which the TChannel server will listen for
// connections. By default, the server listens on an OS-assigned port.
//
// This option has no effect if the Chanel provided to NewInbound is already
// listening for connections when Start() is called.
func ListenAddr(addr string) InboundOption {
	return func(i *inbound) { i.addr = addr }
}

// NewInbound builds a new TChannel inbound from the given Channel.
func NewInbound(ch *tchannel.Channel, opts ...InboundOption) Inbound {
	i := &inbound{ch: ch}
	for _, opt := range opts {
		opt(i)
	}
	return i
}

type inbound struct {
	ch       *tchannel.Channel
	addr     string
	listener net.Listener
}

func (i *inbound) Start(h transport.Handler) error {
	i.ch.GetSubChannel(i.ch.ServiceName()).SetHandler(handler{h})

	if i.ch.State() == tchannel.ChannelListening {
		// Channel.Start() was called before RPC.Start(). We still want to
		// update the Handler and what i.addr means, but nothing else.
		i.addr = i.ch.PeerInfo().HostPort
		return nil
	}

	// Default to ListenIP if addr wasn't given.
	addr := i.addr
	if addr == "" {
		listenIP, err := tchannel.ListenIP()
		if err != nil {
			return err
		}

		addr = listenIP.String() + ":0"
		// TODO(abg): Find a way to export this to users
	}

	// TODO(abg): If addr was just the port (":4040"), we want to use
	// ListenIP() + ":4040" rather than just ":4040".

	var err error
	i.listener, err = net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	i.addr = i.listener.Addr().String() // in case it changed

	if err := i.ch.Serve(i.listener); err != nil {
		return err
	}

	return nil
}

func (i *inbound) Stop() error {
	i.ch.Close()
	return nil
}

func (i *inbound) Addr() net.Addr {
	if i.listener == nil {
		return nil
	}
	return i.listener.Addr()
}

type handler struct{ Handler transport.Handler }

func (h handler) Handle(ctx context.Context, call *tchannel.InboundCall) {
	deadline, ok := ctx.Deadline()
	if !ok {
		call.Response().SendSystemError(tchannel.ErrTimeoutRequired)
		return
	}

	headers, err := readHeaders(call.Format(), call.Arg2Reader)
	if err != nil {
		call.Response().SendSystemError(tchannel.NewSystemError(
			tchannel.ErrCodeUnexpected, "failed to read headers: %v", err))
		return
	}

	body, err := call.Arg3Reader()
	if err != nil {
		call.Response().SendSystemError(tchannel.NewSystemError(
			tchannel.ErrCodeUnexpected, "failed to read body: %v", err))
		return
	}
	defer body.Close()

	rw := newResponseWriter(call)
	defer rw.Close() // TODO(abg): log if this errors

	treq := &transport.Request{
		Caller:    call.CallerName(),
		Service:   call.ServiceName(),
		Encoding:  transport.Encoding(call.Format()),
		Procedure: call.MethodString(),
		Headers:   headers,
		Body:      body,
		TTL:       deadline.Sub(time.Now()),
	}

	if err := h.Handler.Handle(ctx, treq, rw); err != nil {
		call.Response().SendSystemError(tchannel.NewSystemError(
			tchannel.ErrCodeUnexpected, "internal error: %v", err))
		return
	}
}

type responseWriter struct {
	failedWith   error
	bodyWriter   tchannel.ArgWriter
	format       tchannel.Format
	headers      transport.Headers
	response     *tchannel.InboundCallResponse
	wroteHeaders bool
}

func newResponseWriter(call *tchannel.InboundCall) *responseWriter {
	return &responseWriter{
		response: call.Response(),
		headers:  make(transport.Headers),
		format:   call.Format(),
	}
}

func (rw *responseWriter) AddHeaders(h transport.Headers) {
	if rw.wroteHeaders {
		panic("AddHeaders() cannot be called after calling Write().")
	}
	for k, v := range h {
		rw.headers.Set(k, v)
	}
}

func (rw *responseWriter) Write(s []byte) (int, error) {
	if rw.failedWith != nil {
		return 0, rw.failedWith
	}

	if !rw.wroteHeaders {
		rw.wroteHeaders = true
		if err := writeHeaders(rw.format, rw.headers, rw.response.Arg2Writer); err != nil {
			rw.failedWith = err
			return 0, err
		}
	}

	if rw.bodyWriter == nil {
		var err error
		rw.bodyWriter, err = rw.response.Arg3Writer()
		if err != nil {
			rw.failedWith = err
			return 0, err
		}
	}

	return rw.bodyWriter.Write(s)
}

func (rw *responseWriter) Close() error {
	if rw.bodyWriter != nil {
		return rw.bodyWriter.Close()
	}
	if rw.failedWith != nil {
		return rw.failedWith
	}
	return nil
}
<|endoftext|>


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

package refresh

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/pkg/errors"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/allegro/bigcache"

	"github.com/fellou89/caddy-reauth/backend"

	. "github.com/startsmartlabs/caddy-secrets"
)

// Backend name
const Backend = "refresh"

// DefaultTimeout for sub requests
const DefaultTimeout = time.Minute

// Refresh backend provides authentication against a refresh token endpoint.
// If the refresh request returns a http 200 status code then the user
// is considered logged in.
type Refresh struct {
	refreshUrl         string
	refreshCache       *bigcache.BigCache
	timeout            time.Duration
	insecureSkipVerify bool
	followRedirects    bool
	passCookies        bool
}

func init() {
	err := backend.Register(Backend, constructor)
	if err != nil {
		panic(err)
	}
}

func noRedirectsPolicy(req *http.Request, via []*http.Request) error {
	return errors.New("follow redirects disabled")
}

func constructor(config string) (backend.Backend, error) {
	options, err := backend.ParseOptions(config)
	if err != nil {
		return nil, err
	}

	s, found := options["url"]
	if !found {
		return nil, errors.New("url is a required parameter")
	}

	u, err := url.Parse(s)
	if err != nil {
		return nil, errors.Wrap(err, "unable to parse url "+s)
	}

	life, err := parseDurationOption(options, "lifewindow")
	if err != nil {
		return nil, err
	}

	clean, err := parseDurationOption(options, "cleanwindow")
	if err != nil {
		return nil, err
	}

	cacheConfig := bigcache.DefaultConfig(life)
	cacheConfig.CleanWindow = clean
	cache, err := bigcache.NewBigCache(cacheConfig)
	if err != nil {
		return nil, err
	}

	rf := &Refresh{
		refreshUrl:   u.String(),
		refreshCache: cache,
	}

	val, err := parseDurationOption(options, "timeout")
	if err != nil {
		return nil, err
	}
	rf.timeout = val

	bval, err := parseBoolOption(options, "skipverify")
	if err != nil {
		return nil, err
	}
	rf.insecureSkipVerify = bval

	bval, err = parseBoolOption(options, "follow")
	if err != nil {
		return nil, err
	}
	rf.followRedirects = bval

	bval, err = parseBoolOption(options, "cookies")
	if err != nil {
		return nil, err
	}
	rf.passCookies = bval

	return rf, nil
}

func parseBoolOption(options map[string]string, key string) (bool, error) {
	if s, found := options[key]; found {
		return strconv.ParseBool(s)
	}
	return false, nil
}

func parseDurationOption(options map[string]string, key string) (time.Duration, error) {
	if s, found := options[key]; found {
		return time.ParseDuration(s)
	}
	return DefaultTimeout, nil
}

func (h Refresh) refreshRequestObject(c *http.Client, requestToAuth *http.Request, endpoint Endpoint, inputMap map[string]string) ([]byte, error) {
	data := url.Values{}
	for _, d := range endpoint.Data {
		if len(d.Input) > 0 {
			data.Set(d.Key, inputMap[d.Input])
		} else {
			data.Set(d.Key, d.Value)
		}
	}

	// In case endpoints at different urls need to be used,
	// otherwise the url set in the refresh Caddyfile entry is used
	var url string
	if len(endpoint.Url) == 0 {
		url = h.refreshUrl
	} else {
		url = endpoint.Url
	}

	var refreshTokenReq *http.Request
	var err error
	if endpoint.Method == "POST" {
		refreshTokenReq, err = http.NewRequest(endpoint.Method, url+endpoint.Path, strings.NewReader(data.Encode()))

	} else if endpoint.Method == "GET" {
		refreshTokenReq, err = http.NewRequest(endpoint.Method, url+endpoint.Path+"?"+data.Encode(), nil)
	}
	if err != nil {
		return nil, err
	}

	if endpoint.Skipverify {
		if refreshTokenReq.URL.Scheme == "https" && h.insecureSkipVerify {
			c.Transport = &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			}
		}
	}

	if endpoint.Cookies {
		if h.passCookies {
			for _, c := range requestToAuth.Cookies() {
				refreshTokenReq.AddCookie(c)
			}
		}
	}

	for _, h := range endpoint.Headers {
		if len(h.Value) > 0 {
			refreshTokenReq.Header.Add(h.Key, h.Value)

		} else {
			keyCheck := regexp.MustCompile(`(#[[:alnum:]+\-*_*]+#)`)
			keyMatch := keyCheck.FindStringSubmatch(h.Input)
			if len(keyMatch) > 0 {
				for _, m := range keyMatch[1:] {
					replace := m[1 : len(m)-1]
					replaced := strings.Replace(h.Input, m, inputMap[replace], -1)
					refreshTokenReq.Header.Add(h.Key, replaced)
				}
			}
		}
	}

	if refreshResp, err := c.Do(refreshTokenReq); err != nil {
		return nil, errors.Wrap(err, "Error requesting access token")

	} else {
		if refreshBody, err := ioutil.ReadAll(refreshResp.Body); err != nil {
			return nil, errors.Wrap(err, "Error reading response body from access token refresh")

		} else {
			var body map[string]interface{}
			json.Unmarshal(refreshBody, &body)

			for _, f := range endpoint.Failures {
				if f.Validation == "equals" {
					if body[f.Key] == f.Value {
						fmt.Println(url + endpoint.Path + ": " + f.Message)
						return nil, nil
					}
				} else if f.Validation == "presence" {
					if body[f.Key] != nil {
						fmt.Println(url + endpoint.Path + ": " + f.Message + body[f.Key].(string))
						return nil, nil
					}
				}
			}

			if len(endpoint.Responsekey) > 0 {
				if body[endpoint.Responsekey] != nil {
					return []byte(body[endpoint.Responsekey].(string)), nil
				}
			}
			return refreshBody, nil
		}
	}
}

func getObject(mapslice yaml.MapSlice, key string) yaml.MapSlice {
	for _, s := range mapslice {
		if s.Key == key {
			return s.Value.(yaml.MapSlice)
		}
	}
	return nil
}

func getArray(mapslice yaml.MapSlice, key string) []interface{} {
	for _, s := range mapslice {
		if s.Key == key {
			return s.Value.([]interface{})
		}
	}
	return nil
}

func getValue(mapslice yaml.MapSlice, key string) interface{} {
	for _, s := range mapslice {
		if s.Key == key {
			return s.Value
		}
	}
	return nil
}

// Authenticate fulfils the backend interface
func (h Refresh) Authenticate(requestToAuth *http.Request) (bool, error) {
	if requestToAuth.Header.Get("Authorization") == "" {
		// No Token, Unauthorized response
		return failAuth(false, nil)
	}
	authHeader := strings.Split(requestToAuth.Header.Get("Authorization"), " ")
	if len(authHeader) != 2 || authHeader[0] != "Bearer" {
		return failAuth(false, errors.New("Authorization token not properly formatted"))
	}
	resultsMap := map[string]string{}
	resultsMap["client_token"] = authHeader[1]

	c := &http.Client{Timeout: h.timeout}
	if !h.followRedirects {
		c.CheckRedirect = noRedirectsPolicy
	}

	reauth := getObject(SecretsMap, "reauth")
	reauthEndpoints := getArray(reauth, "endpoints")
	endpointData, err := yaml.Marshal(reauthEndpoints)
	if err != nil {
		return failAuth(false, errors.New("Endpoints yaml not setup properly in secrets file"))
	}
	var endpoints []Endpoint
	yaml.Unmarshal(endpointData, &endpoints)

	// this specific structure is needed in the secrets file to have a refresh token available
	for _, e := range endpoints {
		if e.Name == "refresh" {
			for _, d := range e.Data {
				if d.Key == "refresh_token" {
					resultsMap["refresh_token"] = d.Value
				}
			}
		}
	}

	for _, endpoint := range endpoints {
		// check cache for saved response
		entry, err := h.refreshCache.Get(string(resultsMap[endpoint.Cachekey]))
		if err != nil {
			if strings.Contains(err.Error(), "not found") {
				// request data to put in cache when entry is not found
				if responseData, err := h.refreshRequestObject(c, requestToAuth, endpoint, resultsMap); err != nil {
					return failAuth(false, err)

				} else {
					resultsMap[endpoint.Name] = string(responseData)
					h.refreshCache.Set(resultsMap[endpoint.Cachekey], responseData)
				}
			} else {
				return failAuth(false, err)
			}
		} else {
			resultsMap[endpoint.Name] = string(entry)
		}
	}

	requestToAuth.ParseForm()
	requestToAuth.Form[getValue(reauth, "resultkey").(string)] = []string{resultsMap[endpoints[len(endpoints)-1].Name]}

	return true, nil
}

type Endpoint struct {
	Name        string
	Url         string
	Path        string
	Method      string
	Data        []DataObject
	Headers     []DataObject
	Skipverify  bool
	Cookies     bool
	Cachekey    string
	Responsekey string
	Failures    []Failure
}

type Failure struct {
	Validation string
	Key        string
	Value      string
	Message    string
}

type DataObject struct {
	Key   string
	Value string
	Input string
}

func failAuth(result bool, err error) (bool, error) {
	if err != nil {
		log.Println(err.Error())
	}
	return result, err
}
<commit_msg>fixed logic to get 401s back<commit_after>/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2017 Shannon Wynter
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package refresh

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"github.com/pkg/errors"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/allegro/bigcache"

	"github.com/fellou89/caddy-reauth/backend"

	. "github.com/startsmartlabs/caddy-secrets"
)

// Backend name
const Backend = "refresh"

// DefaultTimeout for sub requests
const DefaultTimeout = time.Minute

// Refresh backend provides authentication against a refresh token endpoint.
// If the refresh request returns a http 200 status code then the user
// is considered logged in.
type Refresh struct {
	refreshUrl         string
	refreshCache       *bigcache.BigCache
	timeout            time.Duration
	insecureSkipVerify bool
	followRedirects    bool
	passCookies        bool
}

func init() {
	err := backend.Register(Backend, constructor)
	if err != nil {
		panic(err)
	}
}

func noRedirectsPolicy(req *http.Request, via []*http.Request) error {
	return errors.New("follow redirects disabled")
}

func constructor(config string) (backend.Backend, error) {
	options, err := backend.ParseOptions(config)
	if err != nil {
		return nil, err
	}

	s, found := options["url"]
	if !found {
		return nil, errors.New("url is a required parameter")
	}

	u, err := url.Parse(s)
	if err != nil {
		return nil, errors.Wrap(err, "unable to parse url "+s)
	}

	life, err := parseDurationOption(options, "lifewindow")
	if err != nil {
		return nil, err
	}

	clean, err := parseDurationOption(options, "cleanwindow")
	if err != nil {
		return nil, err
	}

	cacheConfig := bigcache.DefaultConfig(life)
	cacheConfig.CleanWindow = clean
	cache, err := bigcache.NewBigCache(cacheConfig)
	if err != nil {
		return nil, err
	}

	rf := &Refresh{
		refreshUrl:   u.String(),
		refreshCache: cache,
	}

	val, err := parseDurationOption(options, "timeout")
	if err != nil {
		return nil, err
	}
	rf.timeout = val

	bval, err := parseBoolOption(options, "skipverify")
	if err != nil {
		return nil, err
	}
	rf.insecureSkipVerify = bval

	bval, err = parseBoolOption(options, "follow")
	if err != nil {
		return nil, err
	}
	rf.followRedirects = bval

	bval, err = parseBoolOption(options, "cookies")
	if err != nil {
		return nil, err
	}
	rf.passCookies = bval

	return rf, nil
}

func parseBoolOption(options map[string]string, key string) (bool, error) {
	if s, found := options[key]; found {
		return strconv.ParseBool(s)
	}
	return false, nil
}

func parseDurationOption(options map[string]string, key string) (time.Duration, error) {
	if s, found := options[key]; found {
		return time.ParseDuration(s)
	}
	return DefaultTimeout, nil
}

func (h Refresh) refreshRequestObject(c *http.Client, requestToAuth *http.Request, endpoint Endpoint, inputMap map[string]string) ([]byte, error) {
	data := url.Values{}
	for _, d := range endpoint.Data {
		if len(d.Input) > 0 {
			data.Set(d.Key, inputMap[d.Input])
		} else {
			data.Set(d.Key, d.Value)
		}
	}

	// In case endpoints at different urls need to be used,
	// otherwise the url set in the refresh Caddyfile entry is used
	var url string
	if len(endpoint.Url) == 0 {
		url = h.refreshUrl
	} else {
		url = endpoint.Url
	}

	var refreshTokenReq *http.Request
	var err error
	if endpoint.Method == "POST" {
		refreshTokenReq, err = http.NewRequest(endpoint.Method, url+endpoint.Path, strings.NewReader(data.Encode()))

	} else if endpoint.Method == "GET" {
		refreshTokenReq, err = http.NewRequest(endpoint.Method, url+endpoint.Path+"?"+data.Encode(), nil)
	}
	if err != nil {
		return nil, err
	}

	if endpoint.Skipverify {
		if refreshTokenReq.URL.Scheme == "https" && h.insecureSkipVerify {
			c.Transport = &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
			}
		}
	}

	if endpoint.Cookies {
		if h.passCookies {
			for _, c := range requestToAuth.Cookies() {
				refreshTokenReq.AddCookie(c)
			}
		}
	}

	for _, h := range endpoint.Headers {
		if len(h.Value) > 0 {
			refreshTokenReq.Header.Add(h.Key, h.Value)

		} else {
			keyCheck := regexp.MustCompile(`(#[[:alnum:]+\-*_*]+#)`)
			keyMatch := keyCheck.FindStringSubmatch(h.Input)
			if len(keyMatch) > 0 {
				for _, m := range keyMatch[1:] {
					replace := m[1 : len(m)-1]
					replaced := strings.Replace(h.Input, m, inputMap[replace], -1)
					refreshTokenReq.Header.Add(h.Key, replaced)
				}
			}
		}
	}

	if refreshResp, err := c.Do(refreshTokenReq); err != nil {
		return nil, errors.Wrap(err, "Error requesting access token")

	} else {
		if refreshBody, err := ioutil.ReadAll(refreshResp.Body); err != nil {
			return nil, errors.Wrap(err, "Error reading response body from access token refresh")

		} else {
			var body map[string]interface{}
			json.Unmarshal(refreshBody, &body)

			for _, f := range endpoint.Failures {
				if f.Validation == "equals" {
					if body[f.Key] == f.Value {
						fmt.Println(url + endpoint.Path + ": " + f.Message)
						return nil, nil
					}
				} else if f.Validation == "presence" {
					if body[f.Key] != nil {
						fmt.Println(url + endpoint.Path + ": " + f.Message + body[f.Key].(string))
						return nil, nil
					}
				}
			}

			if len(endpoint.Responsekey) > 0 {
				if body[endpoint.Responsekey] != nil {
					return []byte(body[endpoint.Responsekey].(string)), nil
				}
			}
			return refreshBody, nil
		}
	}
}

func getObject(mapslice yaml.MapSlice, key string) yaml.MapSlice {
	for _, s := range mapslice {
		if s.Key == key {
			return s.Value.(yaml.MapSlice)
		}
	}
	return nil
}

func getArray(mapslice yaml.MapSlice, key string) []interface{} {
	for _, s := range mapslice {
		if s.Key == key {
			return s.Value.([]interface{})
		}
	}
	return nil
}

func getValue(mapslice yaml.MapSlice, key string) interface{} {
	for _, s := range mapslice {
		if s.Key == key {
			return s.Value
		}
	}
	return nil
}

// Authenticate fulfils the backend interface
func (h Refresh) Authenticate(requestToAuth *http.Request) (bool, error) {
	if requestToAuth.Header.Get("Authorization") == "" {
		// No Token, Unauthorized response
		return failAuth(false, nil)
	}
	authHeader := strings.Split(requestToAuth.Header.Get("Authorization"), " ")
	if len(authHeader) != 2 || authHeader[0] != "Bearer" {
		return failAuth(false, errors.New("Authorization token not properly formatted"))
	}
	resultsMap := map[string]string{}
	resultsMap["client_token"] = authHeader[1]

	c := &http.Client{Timeout: h.timeout}
	if !h.followRedirects {
		c.CheckRedirect = noRedirectsPolicy
	}

	reauth := getObject(SecretsMap, "reauth")
	reauthEndpoints := getArray(reauth, "endpoints")
	endpointData, err := yaml.Marshal(reauthEndpoints)
	if err != nil {
		return failAuth(false, errors.New("Endpoints yaml not setup properly in secrets file"))
	}
	var endpoints []Endpoint
	yaml.Unmarshal(endpointData, &endpoints)

	// this specific structure is needed in the secrets file to have a refresh token available
	for _, e := range endpoints {
		if e.Name == "refresh" {
			for _, d := range e.Data {
				if d.Key == "refresh_token" {
					resultsMap["refresh_token"] = d.Value
				}
			}
		}
	}

	for _, endpoint := range endpoints {
		// check cache for saved response
		entry, err := h.refreshCache.Get(string(resultsMap[endpoint.Cachekey]))
		if err != nil {
			if strings.Contains(err.Error(), "not found") {
				// request data to put in cache when entry is not found
				if responseData, err := h.refreshRequestObject(c, requestToAuth, endpoint, resultsMap); err != nil {
					return failAuth(false, err)

				} else {
					if responseData == nil {
						return false, nil

					} else {
						resultsMap[endpoint.Name] = string(responseData)
						h.refreshCache.Set(resultsMap[endpoint.Cachekey], responseData)
					}
				}
			} else {
				return failAuth(false, err)
			}
		} else {
			resultsMap[endpoint.Name] = string(entry)
		}
	}

	resultkey := getValue(reauth, "resultkey").(string)
	if len(resultkey) > 0 {
		requestToAuth.ParseForm()
		requestToAuth.Form[resultkey] = []string{resultsMap[endpoints[len(endpoints)-1].Name]}
	}

	return true, nil
}

type Endpoint struct {
	Name        string
	Url         string
	Path        string
	Method      string
	Data        []DataObject
	Headers     []DataObject
	Skipverify  bool
	Cookies     bool
	Cachekey    string
	Responsekey string
	Failures    []Failure
}

type Failure struct {
	Validation string
	Key        string
	Value      string
	Message    string
}

type DataObject struct {
	Key   string
	Value string
	Input string
}

func failAuth(result bool, err error) (bool, error) {
	if err != nil {
		log.Println(err.Error())
	}
	return result, err
}
<|endoftext|>


<commit_before>package imap

import (
	"fmt"
	"regexp"
)

type command struct {
	match   *regexp.Regexp
	handler func(commandArgs, *Conn)
}

type commandArgs []string

func (a commandArgs) FullCommand() string {
	return a[0]
}

func (a commandArgs) ID() string {
	return a[1]
}

func (a commandArgs) Arg(i int) string {
	return a[i+2]
}

var commands []command

// Register all supported client command handlers
// with the server. This function is run on server startup and
// panics if a command regex is invalid.
func init() {
	commands = make([]command, 0)

	registerCommand("(?i:CAPABILITY)", cmdCapability)
	registerCommand("(?i:LOGIN) \"([A-z0-9]+)\" \"([A-z0-9]+)\"", cmdLogin)
	registerCommand("(?i:AUTHENTICATE PLAIN)", cmdAuthPlain)
	registerCommand("(?i:LIST) \"?([A-z0-9]+)?\"? \"?([A-z0-9*]+)?\"?", cmdList)
	registerCommand("(?i:LSUB)", cmdLSub)
	registerCommand("(?i:LOGOUT)", cmdLogout)
	registerCommand("(?i:NOOP)", cmdNoop)
	registerCommand("(?i:CLOSE)", cmdClose)
	registerCommand("(?i:SELECT) \"?([A-z0-9]+)?\"?", cmdSelect)
	registerCommand("(?i:EXAMINE) \"?([A-z0-9]+)\"?", cmdExamine)
	registerCommand("(?i:STATUS) \"?([A-z0-9/]+)\"? \\(([A-z\\s]+)\\)", cmdStatus)
	registerCommand("((?i)UID )?(?i:FETCH) (?:(\\d+)(?:\\:([\\*\\d]+))?) \\(([A-z0-9\\s\\(\\)\\[\\]\\.-]+)\\)", cmdFetch)
	// STORE 2:4 +FLAGS (\Deleted)       Mark messages as deleted
	// STORE 2:4 -FLAGS (\Seen)          Mark messages as unseen
	// STORE 2:4 FLAGS (\Seen \Deleted)  Replace flags
	registerCommand("((?i)UID )?(?i:STORE) (?:(\\d+)(?:\\:([\\*\\d]+))?) ([\\+\\-])?(?i:FLAGS(\\.SILENT)?) \\(?([\\\\A-z0-9]+)\\)?", cmdStoreFlags)
}

func registerCommand(matchExpr string, handleFunc func(commandArgs, *Conn)) error {
	// Add command identifier to beginning of command
	matchExpr = "([A-z0-9\\.]+) " + matchExpr

	newRE := regexp.MustCompile(matchExpr)
	c := command{match: newRE, handler: handleFunc}
	commands = append(commands, c)
	return nil
}

// Write out the info for a mailbox (used in both SELECT and EXAMINE)
func writeMailboxInfo(c *Conn, m Mailbox) {
	fmt.Fprintf(c, "* %d EXISTS\r\n", m.Messages())
	fmt.Fprintf(c, "* %d RECENT\r\n", m.Recent())
	fmt.Fprintf(c, "* OK [UNSEEN %d]\r\n", m.Unseen())
	fmt.Fprintf(c, "* OK [UIDNEXT %d]\r\n", m.NextUid())
	fmt.Fprintf(c, "* OK [UIDVALIDITY %d]\r\n", 250)
	fmt.Fprintf(c, "* FLAGS (\\Answered \\Flagged \\Deleted \\Seen \\Draft)\r\n")
}

func messageFlags(msg Message) []string {
	flags := make([]string, 0)
	if msg.IsAnswered() {
		flags = append(flags, "\\Answered")
	}
	if msg.IsSeen() {
		flags = append(flags, "\\Seen")
	}
	if msg.IsRecent() {
		flags = append(flags, "\\Recent")
	}
	if msg.IsDeleted() {
		flags = append(flags, "\\Deleted")
	}
	if msg.IsDraft() {
		flags = append(flags, "\\Draft")
	}
	if msg.IsFlagged() {
		flags = append(flags, "\\Flagged")
	}
	return flags
}

func cmdNA(args commandArgs, c *Conn) {
	c.writeResponse(args.ID(), "BAD Not implemented")
}
<commit_msg>Preallocate flags array #8<commit_after>package imap

import (
	"fmt"
	"regexp"
)

type command struct {
	match   *regexp.Regexp
	handler func(commandArgs, *Conn)
}

type commandArgs []string

func (a commandArgs) FullCommand() string {
	return a[0]
}

func (a commandArgs) ID() string {
	return a[1]
}

func (a commandArgs) Arg(i int) string {
	return a[i+2]
}

var commands []command

// Register all supported client command handlers
// with the server. This function is run on server startup and
// panics if a command regex is invalid.
func init() {
	commands = make([]command, 0)

	registerCommand("(?i:CAPABILITY)", cmdCapability)
	registerCommand("(?i:LOGIN) \"([A-z0-9]+)\" \"([A-z0-9]+)\"", cmdLogin)
	registerCommand("(?i:AUTHENTICATE PLAIN)", cmdAuthPlain)
	registerCommand("(?i:LIST) \"?([A-z0-9]+)?\"? \"?([A-z0-9*]+)?\"?", cmdList)
	registerCommand("(?i:LSUB)", cmdLSub)
	registerCommand("(?i:LOGOUT)", cmdLogout)
	registerCommand("(?i:NOOP)", cmdNoop)
	registerCommand("(?i:CLOSE)", cmdClose)
	registerCommand("(?i:SELECT) \"?([A-z0-9]+)?\"?", cmdSelect)
	registerCommand("(?i:EXAMINE) \"?([A-z0-9]+)\"?", cmdExamine)
	registerCommand("(?i:STATUS) \"?([A-z0-9/]+)\"? \\(([A-z\\s]+)\\)", cmdStatus)
	registerCommand("((?i)UID )?(?i:FETCH) (?:(\\d+)(?:\\:([\\*\\d]+))?) \\(([A-z0-9\\s\\(\\)\\[\\]\\.-]+)\\)", cmdFetch)
	// STORE 2:4 +FLAGS (\Deleted)       Mark messages as deleted
	// STORE 2:4 -FLAGS (\Seen)          Mark messages as unseen
	// STORE 2:4 FLAGS (\Seen \Deleted)  Replace flags
	registerCommand("((?i)UID )?(?i:STORE) (?:(\\d+)(?:\\:([\\*\\d]+))?) ([\\+\\-])?(?i:FLAGS(\\.SILENT)?) \\(?([\\\\A-z0-9]+)\\)?", cmdStoreFlags)
}

func registerCommand(matchExpr string, handleFunc func(commandArgs, *Conn)) error {
	// Add command identifier to beginning of command
	matchExpr = "([A-z0-9\\.]+) " + matchExpr

	newRE := regexp.MustCompile(matchExpr)
	c := command{match: newRE, handler: handleFunc}
	commands = append(commands, c)
	return nil
}

// Write out the info for a mailbox (used in both SELECT and EXAMINE)
func writeMailboxInfo(c *Conn, m Mailbox) {
	fmt.Fprintf(c, "* %d EXISTS\r\n", m.Messages())
	fmt.Fprintf(c, "* %d RECENT\r\n", m.Recent())
	fmt.Fprintf(c, "* OK [UNSEEN %d]\r\n", m.Unseen())
	fmt.Fprintf(c, "* OK [UIDNEXT %d]\r\n", m.NextUid())
	fmt.Fprintf(c, "* OK [UIDVALIDITY %d]\r\n", 250)
	fmt.Fprintf(c, "* FLAGS (\\Answered \\Flagged \\Deleted \\Seen \\Draft)\r\n")
}

func messageFlags(msg Message) []string {
	flags := make([]string, 0, 6) // Up to 6 flags
	if msg.IsAnswered() {
		flags = append(flags, "\\Answered")
	}
	if msg.IsSeen() {
		flags = append(flags, "\\Seen")
	}
	if msg.IsRecent() {
		flags = append(flags, "\\Recent")
	}
	if msg.IsDeleted() {
		flags = append(flags, "\\Deleted")
	}
	if msg.IsDraft() {
		flags = append(flags, "\\Draft")
	}
	if msg.IsFlagged() {
		flags = append(flags, "\\Flagged")
	}
	return flags
}

func cmdNA(args commandArgs, c *Conn) {
	c.writeResponse(args.ID(), "BAD Not implemented")
}
<|endoftext|>


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

package cmd

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"os/user"
	"strings"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var accountId  string
var userName   string
var tokenCode  string
var profile    string
var noMfa      bool

// authCmd represents the auth command
var authCmd = &cobra.Command{
	Use:   "auth",
	Short: "establishes an MFA session via STS",
	Long:  `The auth command helps you authenticate via MFA.`,
	Run: func(cmd *cobra.Command, args []string) {
		err := viper.ReadInConfig() // Find and read the config file
		if accountId == "" && err != nil {
			fmt.Println("No config file available and no account details supplied")
			os.Exit(1)
		}
		if accountId != "" {
			viper.Set("AccountId", accountId)
		}
		if userName != "" {
			viper.Set("UserName", userName)
		}
		usr, err := user.Current()
		checkError(err)
		fileName := usr.HomeDir + "/.aws/portray-session-" + profile + ".json"
		awsCreds := getCredsFromFile(fileName)
		if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
			if tokenCode != "" {
				awsCreds = getNewSession(accountId, userName, tokenCode)
			} else if viper.GetString("AccountId") != "" {
				reader := bufio.NewReader(os.Stdin)
				fmt.Print("Enter token: ")
				token, _ := reader.ReadString('\n')
				token = strings.TrimSpace(token)
				awsCreds = getNewSession(
					viper.GetString("AccountId"),
					viper.GetString("UserName"),
					token)
			} else {
				fmt.Println("You need a valid session!")
				os.Exit(1)
			}
			//writeSessionFile(awsCreds, fileName)
		}
		if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
			fmt.Println("You need a valid session!")
			os.Exit(1)
		}

		accountId := awsCreds.AccountId

		sessionToEnvVars(awsCreds, accountId, "", profile)
		startShell(accountId)
	},
}

func init() {
	RootCmd.AddCommand(authCmd)

	authCmd.Flags().StringVarP(&accountId, "account", "a", "", "the AWS account number")
	authCmd.Flags().StringVarP(&userName, "username", "u", "", "the AWS user name")
	authCmd.Flags().StringVarP(&tokenCode, "token", "t", "", "an MFA token")
	authCmd.Flags().StringVarP(&profile, "profile", "p", "default", "a name for your profile")
	authCmd.Flags().BoolP("no-mfa", "n", false, "disable MFA")
    viper.BindPFlag("noMfa", configCmd.Flags().Lookup("no-mfa"))
}

func checkError(err error) {
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
}

func getCredsFromFile(fileName string) (awsCreds AwsCreds) {
	file, err := ioutil.ReadFile(fileName)
	if err != nil {
		return
	}

	json.Unmarshal(file, &awsCreds)
	return
}

func validateSession(awsCreds AwsCreds) (valid bool) {
	valid = false
	timestamp := int64(time.Now().Unix())
	if timestamp < awsCreds.Expiration {
		valid = true
	}
	return
}

// AwsCreds represents a set of AWS credentials
type AwsCreds struct {
	AccessKeyID     string
	SecretAccessKey string
	SessionToken    string
	Expiration      int64
	AccountId       string
}

func getNewSession(accountId string, userName string, tokenCode string) (awsCreds AwsCreds) {
	sess, err := session.NewSession(&aws.Config{Region: aws.String("us-east-1")})
	checkError(err)
	svc := sts.New(sess)

	params := &sts.GetSessionTokenInput{
		DurationSeconds: aws.Int64(43200),
		SerialNumber:    aws.String("arn:aws:iam::" + accountId + ":mfa/" + userName),
		TokenCode:       aws.String(tokenCode),
	}

	resp, err := svc.GetSessionToken(params)

	checkError(err)

	awsCreds = AwsCreds{
        *resp.Credentials.AccessKeyId,
		*resp.Credentials.SecretAccessKey,
		*resp.Credentials.SessionToken,
		resp.Credentials.Expiration.Unix(),
		accountId,
	}

	return
}

func sessionToEnvVars(awsCreds AwsCreds, account string, role string, profile string) {
	prompt := account
	if role != "" {
		prompt = prompt + ":" + role
	}
	if profile != "" {
		prompt = prompt + ":" + profile
	}

	fmt.Println("Setting ENV VARS")
	os.Setenv("AWS_ACCESS_KEY_ID", awsCreds.AccessKeyID)
	os.Setenv("AWS_SECRET_ACCESS_KEY", awsCreds.SecretAccessKey)
	os.Setenv("AWS_SESSION_TOKEN", awsCreds.SessionToken)
	os.Setenv("PORTRAY_PROMPT", prompt)

}

func startShell(account string) {
	fmt.Println("Starting shell with Session in: " + account)
	syscall.Exec(os.Getenv("SHELL"), []string{os.Getenv("SHELL")}, syscall.Environ())
}
<commit_msg>Update auth command to support new config format, add support for --no-mfa<commit_after>// Copyright © 2017 Jason Myers <jason@mailthemyers.com>
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"syscall"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sts"
	homedir "github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var accountId string
var userName string
var tokenCode string
var profile string
var noMfa bool

// authCmd represents the auth command
var authCmd = &cobra.Command{
	Use:   "auth",
	Short: "establishes an MFA session via STS",
	Long:  `The auth command helps you authenticate via MFA.`,
	Run: func(cmd *cobra.Command, args []string) {
		err := viper.ReadInConfig() // Find and read the config file
		checkError(err)

		// If the user doesn't pass in a specific account, try to find it from
		// the default auth profile.
		if accountId == "" {
			defaultAccountId := viper.GetString("AuthProfiles.default.AccountId")
			if defaultAccountId != "" {
				viper.Set("AccountId", defaultAccountId)
			} else {
				fmt.Println("Couldn't find default profile and no account details specified!")
				os.Exit(1)
			}
		} else {
			viper.Set("UserName", userName)
		}

		// If the user doesn't pass in a specific username, try to find it from
		// the default auth profile.
		if userName == "" {
			defaultUserName := viper.GetString("AuthProfiles.default.UserName")
			if defaultUserName != "" {
				viper.Set("UserName", defaultUserName)
			} else {
				fmt.Println("Couldn't find default username!")
				os.Exit(1)
			}
		} else {
			viper.Set("UserName", userName)
		}

		// If the user doesn't pass in a specific profile, try to find it from
		// the default auth profile.
		if profile == "" {
			defaultProfile := viper.GetString("AuthProfiles.default.Name")
			if defaultProfile != "" {
				viper.Set("Profile", defaultProfile)
			} else {
				// Default to zee default
				fmt.Println("Couldn't find default auth profile via config and none specified. Trying \"default\"")
				viper.Set("Profile", "default")
			}
		} else {
			viper.Set("Profile", profile)
		}

		home, err := homedir.Dir()
		checkError(err)
		fileName := home + "/.aws/portray-session-" + viper.GetString("Profile") + ".json"
		awsCreds := getCredsFromFile(fileName)

		// If there's no valid session cache, generate a new session. Prompt
		// for MFA token if it's not passed, unless the --no-mfa flag is set.
		if awsCreds.SessionToken == "" || !validateSession(awsCreds) {
			if tokenCode == "" {
				if viper.GetBool("noMfa") {
					fmt.Println("Skipping MFA token prompting")
				} else {
					// Prompt for MFA token
					reader := bufio.NewReader(os.Stdin)
					fmt.Print("Enter token: ")
					token, _ := reader.ReadString('\n')
					tokenCode = strings.TrimSpace(token)
				}
			}

			awsCreds = getNewSession(
				viper.GetString("AccountId"),
				viper.GetString("UserName"),
				tokenCode)

			//writeSessionFile(awsCreds, fileName)
		}

		sessionToEnvVars(
			awsCreds,
			viper.GetString("AccountId"),
			"",
			viper.GetString("Profile"))

		startShell(viper.GetString("AccountId"))
	},
}

func init() {
	RootCmd.AddCommand(authCmd)

	authCmd.Flags().StringVarP(&accountId, "account", "a", "", "the AWS account number")
	authCmd.Flags().StringVarP(&userName, "username", "u", "", "the AWS user name")
	authCmd.Flags().StringVarP(&tokenCode, "token", "t", "", "an MFA token")
	authCmd.Flags().StringVarP(&profile, "profile", "p", "", "a name for your profile")
	authCmd.Flags().BoolP("no-mfa", "n", false, "disable MFA")
	viper.BindPFlag("noMfa", authCmd.Flags().Lookup("no-mfa"))
}

func checkError(err error) {
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
}

func getCredsFromFile(fileName string) (awsCreds AwsCreds) {
	file, err := ioutil.ReadFile(fileName)
	if err != nil {
		return
	}

	json.Unmarshal(file, &awsCreds)
	return
}

func validateSession(awsCreds AwsCreds) (valid bool) {
	valid = false
	timestamp := int64(time.Now().Unix())
	if timestamp < awsCreds.Expiration {
		valid = true
	}
	return
}

// AwsCreds represents a set of AWS credentials
type AwsCreds struct {
	AccessKeyID     string
	SecretAccessKey string
	SessionToken    string
	Expiration      int64
	AccountId       string
}

func getNewSession(accountId string, userName string, tokenCode string) (awsCreds AwsCreds) {
	sess, err := session.NewSession(&aws.Config{Region: aws.String("us-east-1")})
	checkError(err)
	svc := sts.New(sess)

	// If no tokenCode is passed, assume MFA has been disabled by a flag
	var params *sts.GetSessionTokenInput
	if tokenCode == "" {
		params = &sts.GetSessionTokenInput{
			DurationSeconds: aws.Int64(43200),
		}
	} else {
		params = &sts.GetSessionTokenInput{
			DurationSeconds: aws.Int64(43200),
			SerialNumber:    aws.String("arn:aws:iam::" + accountId + ":mfa/" + userName),
			TokenCode:       aws.String(tokenCode),
		}
	}

	resp, err := svc.GetSessionToken(params)
	checkError(err)

	awsCreds = AwsCreds{
		*resp.Credentials.AccessKeyId,
		*resp.Credentials.SecretAccessKey,
		*resp.Credentials.SessionToken,
		resp.Credentials.Expiration.Unix(),
		accountId,
	}

	return
}

func sessionToEnvVars(awsCreds AwsCreds, account string, role string, profile string) {
	prompt := account
	if role != "" {
		prompt = prompt + ":" + role
	}
	if profile != "" {
		prompt = prompt + ":" + profile
	}

	fmt.Println("Setting ENV VARS")
	os.Setenv("AWS_ACCESS_KEY_ID", awsCreds.AccessKeyID)
	os.Setenv("AWS_SECRET_ACCESS_KEY", awsCreds.SecretAccessKey)
	os.Setenv("AWS_SESSION_TOKEN", awsCreds.SessionToken)
	os.Setenv("PORTRAY_PROMPT", prompt)

}

func startShell(account string) {
	fmt.Println("Starting shell with Session in: " + account)
	syscall.Exec(os.Getenv("SHELL"), []string{os.Getenv("SHELL")}, syscall.Environ())
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/evandroflores/claimr/model"
	"github.com/shomali11/slacker"
)

func init() {
	Register("free <container-name>", "Makes a container available for use.", free)
}

func free(request *slacker.Request, response slacker.ResponseWriter) {
	response.Typing()

	event := getEvent(request)
	if direct, err := isDirect(event.Channel); direct {
		response.Reply(err.Error())
		return
	}

	containerName := request.Param("container-name")

	container, err := model.GetContainer(event.Team, event.Channel, containerName)

	if err != nil {
		response.Reply(err.Error())
		return
	}

	if container == (model.Container{}) {
		response.Reply(fmt.Sprintf(Messages["container-not-found-on-channel"], containerName, event.Channel))
		return
	}

	if container.InUseBy != event.User {
		response.Reply(fmt.Sprintf(Messages["container-in-use-by-other"], containerName))
		return
	}

	container.InUseBy = ""
	container.InUseForReason = ""

	err = container.Update()
	if err != nil {
		response.Reply(Messages["fail-to-update"])
		return
	}

	response.Reply(fmt.Sprintf(Messages["container-free"], containerName))
}
<commit_msg>Grouping validations on cmd/free<commit_after>package cmd

import (
	"fmt"

	"github.com/evandroflores/claimr/model"
	"github.com/shomali11/slacker"
)

func init() {
	Register("free <container-name>", "Makes a container available for use.", free)
}

func free(request *slacker.Request, response slacker.ResponseWriter) {
	response.Typing()

	event := getEvent(request)
	if direct, err := isDirect(event.Channel); direct {
		response.Reply(err.Error())
		return
	}

	containerName := request.Param("container-name")

	container, err := model.GetContainer(event.Team, event.Channel, containerName)

	if err != nil {
		response.Reply(err.Error())
		return
	}

	checks := []Check{
		{container == (model.Container{}), fmt.Sprintf(Messages["container-not-found-on-channel"], containerName, event.Channel)},
		{container.InUseBy != event.User, fmt.Sprintf(Messages["container-in-use-by-other"], containerName)},
	}

	err = RunChecks(checks)
	if err != nil {
		response.Reply(err.Error())
		return
	}

	container.InUseBy = ""
	container.InUseForReason = ""

	err = container.Update()
	if err != nil {
		response.Reply(Messages["fail-to-update"])
		return
	}

	response.Reply(fmt.Sprintf(Messages["container-free"], containerName))
}
<|endoftext|>


<commit_before>package cmd

import (
	"os"
	"path"

	"github.com/blang/semver"

	"github.com/containerum/chkit/pkg/model"
	"github.com/sirupsen/logrus"
	cli "gopkg.in/urfave/cli.v2"
)

const (
	Version        = "3.0.0-alpha"
	FlagAPIaddr    = "apiaddr"
	FlagConfigFile = "config"
)

func Run(args []string) error {
	Configuration := &model.Config{}
	log := &logrus.Logger{
		Formatter: &logrus.TextFormatter{},
	}
	err := initConfig(Configuration)
	if err != nil {
		log.WithError(err).
			Errorf("error while getting homedir path")
		return err
	}
	var App = &cli.App{
		Name:    "chkit",
		Version: semver.MustParse(Version).String(),
		Action: func(ctx *cli.Context) error {
			err := loadConfig(&Configuration.Client, ctx.String("config"))
			if err != nil && !os.IsNotExist(err) {
				log.WithError(err).
					Errorf("error while loading config file")
				return err
			} else if os.IsNotExist(err) {
				log.Info("You are not logged in!")
				err = ctx.App.Command("login").Run(ctx)
				if err != nil {
					return err
				}
			}
			return nil
		},
		Commands: []*cli.Command{
			commandLogin(log, Configuration),
		},
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    "config",
				Usage:   "config file",
				Aliases: []string{"c"},
				Value:   path.Join(Configuration.ConfigPath, "config.toml"),
			},
		},
	}
	return App.Run(args)
}
<commit_msg>add logger level and app info<commit_after>package cmd

import (
	"os"
	"path"

	"github.com/blang/semver"

	"github.com/containerum/chkit/pkg/model"
	"github.com/sirupsen/logrus"
	cli "gopkg.in/urfave/cli.v2"
)

const (
	Version        = "3.0.0-alpha"
	FlagAPIaddr    = "apiaddr"
	FlagConfigFile = "config"
)

func Run(args []string) error {
	Configuration := &model.Config{}
	log := &logrus.Logger{
		Formatter: &logrus.TextFormatter{},
		Level:     logrus.DebugLevel,
	}
	err := initConfig(Configuration)
	if err != nil {
		log.WithError(err).
			Errorf("error while getting homedir path")
		return err
	}
	var App = &cli.App{
		Name:    "chkit",
		Usage:   "containerum cli",
		Version: semver.MustParse(Version).String(),
		Action: func(ctx *cli.Context) error {
			err := loadConfig(&Configuration.Client, ctx.String("config"))
			if err != nil && !os.IsNotExist(err) {
				log.WithError(err).
					Errorf("error while loading config file")
				return err
			} else if os.IsNotExist(err) {
				log.Info("You are not logged in!")
				err = ctx.App.Command("login").Run(ctx)
				if err != nil {
					return err
				}
			}
			return nil
		},
		Commands: []*cli.Command{
			commandLogin(log, Configuration),
		},
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:    "config",
				Usage:   "config file",
				Aliases: []string{"c"},
				Value:   path.Join(Configuration.ConfigPath, "config.toml"),
			},
		},
	}
	return App.Run(args)
}
<|endoftext|>


<commit_before>// Copyright © 2016 Nick Klauer <klauer@gmail.com>
//
// 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"

	"log"
	"os"
	"time"

	"github.com/VojtechVitek/go-trello"
	"github.com/klauern/trackello/rest"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// listCmd represents the list command
var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List activities on a board",
	Long: `List will pull all the activities for a particular
Trello board and list them in descending order.  This is useful
if you find yourself having to see what you've been working on`,
	Run: func(cmd *cobra.Command, args []string) {
		Track()
	},
}

func init() {
	RootCmd.AddCommand(listCmd)
}

// trelloConnection repesents the connection to Trello and your preferred Board.
type trelloConnection struct {
	token string
	appKey string
	board trello.Board
}

func newTrelloConnection() (*trelloConnection, error) {
	token := viper.GetString("token")
	appKey := viper.GetString("appkey")
	// New Trello Client
	tr, err := trello.NewAuthClient(appKey, &token)
	if err != nil {
		log.Fatal(err)
		return nil, err
	}
	board, err := tr.Board(viper.GetString("board"))
	if err != nil {
		log.Fatal(err)
		return nil, err
	}

	return &trelloConnection{
		token:token,
		appKey:appKey,
		board:*board,
	}, nil
}

// Track pulls all the latest activity from your Trello board given you've set the token, appkey, and preferred board
// ID to use.
// TODO: cmd\list.go:51::warning: cyclomatic complexity 13 of function Track() is high (> 10) (gocyclo)
func Track() {

	conn, err := newTrelloConnection()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	args := rest.CreateArgsForBoardActions()
	actions, err := conn.board.Actions(args...)
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	cardsWorkedOn := make(map[string]time.Time)
	oldestDate := time.Now()
	boardActions := make(map[string][]trello.Action)

	for _, action := range actions {
		switch boardActions[action.Data.Card.Name] {
		case nil:
			boardActions[action.Data.Card.Name] = []trello.Action{action}
		default:
			boardActions[action.Data.Card.Name] = append(boardActions[action.Data.Card.Name], action)
		}
		actionDate, err := time.Parse(rest.DateLayout, action.Date)
		if err != nil {
			continue // skip this one
		}
		if actionDate.Before(oldestDate) {
			oldestDate = actionDate
		}
		cardDate := cardsWorkedOn[action.Data.Card.Name]
		if cardDate.IsZero() || cardDate.After(actionDate) {
			cardsWorkedOn[action.Data.Card.Name] = actionDate
		}
	}

	fmt.Printf("Cards Worked from %s to now:\n", oldestDate.Format(time.ANSIC))
	for k, v := range boardActions {
		fmt.Printf("* %s\n", k)
		for _, vv := range v {
			fmt.Printf("  - %-24s %ss\n", vv.Date, vv.Type)
		}
	}
}
<commit_msg>slight improvement for gocyclo<commit_after>// Copyright © 2016 Nick Klauer <klauer@gmail.com>
//
// 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"

	"log"
	"os"
	"time"

	"github.com/VojtechVitek/go-trello"
	"github.com/klauern/trackello/rest"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

// listCmd represents the list command
var listCmd = &cobra.Command{
	Use:   "list",
	Short: "List activities on a board",
	Long: `List will pull all the activities for a particular
Trello board and list them in descending order.  This is useful
if you find yourself having to see what you've been working on`,
	Run: func(cmd *cobra.Command, args []string) {
		Track()
	},
}

func init() {
	RootCmd.AddCommand(listCmd)
}

// trelloConnection repesents the connection to Trello and your preferred Board.
type trelloConnection struct {
	token string
	appKey string
	board trello.Board
}

func newTrelloConnection() (*trelloConnection, error) {
	token := viper.GetString("token")
	appKey := viper.GetString("appkey")
	// New Trello Client
	tr, err := trello.NewAuthClient(appKey, &token)
	if err != nil {
		log.Fatal(err)
		return nil, err
	}
	board, err := tr.Board(viper.GetString("board"))
	if err != nil {
		log.Fatal(err)
		return nil, err
	}

	return &trelloConnection{
		token:token,
		appKey:appKey,
		board:*board,
	}, nil
}

// Track pulls all the latest activity from your Trello board given you've set the token, appkey, and preferred board
// ID to use.
// TODO: cmd\list.go:78::warning: cyclomatic complexity 12 of function Track() is high (> 10) (gocyclo)
func Track() {

	conn, err := newTrelloConnection()
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	args := rest.CreateArgsForBoardActions()
	actions, err := conn.board.Actions(args...)
	if err != nil {
		log.Fatal(err)
		os.Exit(1)
	}

	cardsWorkedOn := make(map[string]time.Time)
	oldestDate := time.Now()
	boardActions := make(map[string][]trello.Action)

	for _, action := range actions {
		switch boardActions[action.Data.Card.Name] {
		case nil:
			boardActions[action.Data.Card.Name] = []trello.Action{action}
		default:
			boardActions[action.Data.Card.Name] = append(boardActions[action.Data.Card.Name], action)
		}
		actionDate, err := time.Parse(rest.DateLayout, action.Date)
		if err != nil {
			continue // skip this one
		}
		if actionDate.Before(oldestDate) {
			oldestDate = actionDate
		}
		cardDate := cardsWorkedOn[action.Data.Card.Name]
		if cardDate.IsZero() || cardDate.After(actionDate) {
			cardsWorkedOn[action.Data.Card.Name] = actionDate
		}
	}

	fmt.Printf("Cards Worked from %s to now:\n", oldestDate.Format(time.ANSIC))
	for k, v := range boardActions {
		fmt.Printf("* %s\n", k)
		for _, vv := range v {
			fmt.Printf("  - %-24s %ss\n", vv.Date, vv.Type)
		}
	}
}
<|endoftext|>


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

package cmd

import (
	"fmt"
	"os"

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

var cfgFile string
var logLevel string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "k8s-updater",
	Short: "A brief description of your application",
	Long: `A longer description that spans multiple lines and likely contains
examples and usage of using your application. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.`,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	Run: func(cmd *cobra.Command, args []string) {
	},
}

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

func init() {
	cobra.OnInitialize(initConfig, initLogging)

	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.k8s-updater.yaml)")
	RootCmd.PersistentFlags().StringVarP(&logLevel, "loglevel", "l", log.DebugLevel.String(), "Log level")
}

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

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

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Println("Using config file:", viper.ConfigFileUsed())
	}
}

func initLogging() {
	level, err := log.ParseLevel(logLevel)
	if err != nil {
		log.Fatalln(err)
	}
	log.SetLevel(level)
	log.Debugln("Loglevel set to", log.GetLevel().String())
}
<commit_msg>init k8s api client<commit_after>// Copyright © 2016 Sabaka OU <hello@sabaka.io>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/sabakaio/k8s-updater/pkg/util"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
)

var cfgFile string
var logLevel string

// Kubernetes API host
var kHost string

// Kubernetes API client
var k *client.Client

// Default namespace for Kubernetes API resources
var namespace string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "k8s-updater",
	Short: "A brief description of your application",
	Long: `A longer description that spans multiple lines and likely contains
examples and usage of using your application. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.`,
	// Uncomment the following line if your bare application
	// has an action associated with it:
	Run: func(cmd *cobra.Command, args []string) {
	},
}

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

func init() {
	cobra.OnInitialize(initConfig, initLogging, initClient)

	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.k8s-updater.yaml)")
	RootCmd.PersistentFlags().StringVarP(&logLevel, "loglevel", "l", log.DebugLevel.String(), "Log level")
	RootCmd.PersistentFlags().StringVarP(&kHost, "host", "H", "", "Kubernetes host to connect to")
	RootCmd.PersistentFlags().StringVarP(&namespace, "namespace", "n", api.NamespaceDefault, "Kubernetes namespace")
}

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

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

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err == nil {
		fmt.Println("Using config file:", viper.ConfigFileUsed())
	}
}

func initLogging() {
	level, err := log.ParseLevel(logLevel)
	if err != nil {
		log.Fatalln(err)
	}
	log.SetLevel(level)
	log.Debugln("Loglevel set to", log.GetLevel().String())
}

// Create Kubernetes API client for given host. Shared for all subcommands
func initClient() {
	var err error
	k, err = util.CreateClient(kHost)
	if err != nil {
		log.Fatalln("Can't connect to Kubernetes API:", err)
	}
}
<|endoftext|>


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

package cmd

import (
	"fmt"
	"os"
	"sync"

	g "github.com/jaxxstorm/unseal/gpg"
	v "github.com/jaxxstorm/unseal/vault"

	log "github.com/Sirupsen/logrus"
	"github.com/bgentry/speakeasy"
	"github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var cfgFile string
var unsealKey string
var vaultHost string
var vaultPort int

var caPath string
var gpgPub string
var gpgSecret string
var gpgPass string

type Host struct {
	Name string
	Port int
	Key  string
}

var hosts []Host

var Version string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "unseal",
	Short: "Unseal a set of vault servers",
	Long:  `Unseal allows you to unseal a large set of vault servers using the HTTP API.`,

	// Uncomment the following line if your bare application
	// has an action associated with it:
	Run: func(cmd *cobra.Command, args []string) {

		// unmarshal config file
		err := viper.UnmarshalKey("hosts", &hosts)

		// check for valid config file
		if err != nil {
			log.Fatal("Unable to read hosts key in config file: %s", err)
		}

		caPath = viper.GetString("capath")

		gpg := viper.GetBool("gpg")

		if gpg == true {
			log.Info("Using GPG")
			gpgSecret = viper.GetString("gpgsecretkeyring")
			gpgPub = viper.GetString("gpgpublickeyring")
			gpgPass, err = speakeasy.Ask("Please enter your password: ")
			if err != nil {
				log.Fatal("Password error")
			}
		}

		var wg sync.WaitGroup

		for _, h := range hosts {

			hostName := h.Name
			hostPort := h.Port
			key := h.Key

			var vaultKey string

			if gpg == true {
				vaultKey, err = g.Decrypt(gpgPub, gpgSecret, key, gpgPass)
				if err != nil {
					log.Fatal("GPG Decrypt Error: ", err)
				}
			} else {
				vaultKey = key
			}

			wg.Add(1)

			go func(hostName string, hostPort int) {
				defer wg.Done()

				// create a vault client
				client, err := v.VaultClient(hostName, hostPort, caPath)
				if err != nil {
					log.WithFields(log.Fields{"host": hostName, "port": hostPort}).Error(err)
				}
				// get the current status
				init := v.InitStatus(client)
				if init.Ready == true {
					result, err := client.Sys().Unseal(vaultKey)
					// error while unsealing
					if err != nil {
						log.WithFields(log.Fields{"host": hostName}).Error("Error running unseal operation")
					}

					// if it's still sealed, print the progress
					if result.Sealed == true {
						log.WithFields(log.Fields{"host": hostName, "progress": result.Progress, "threshold": result.T}).Info("Unseal operation performed")
						// otherwise, tell us it's unsealed!
					} else {
						log.WithFields(log.Fields{"host": hostName, "progress": result.Progress, "threshold": result.T}).Info("Vault is unsealed!")
					}
					// zero out the key
					// FIXME: is this the best way to do this?
					// Is it safe?
					key = ""
				} else {
					// sad times, not ready to be unsealed
					log.WithFields(log.Fields{"host": hostName}).Error("Vault is not ready to be unsealed")
				}

			}(hostName, hostPort)

		}
		wg.Wait()
	},
}

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

func init() {
	cobra.OnInitialize(initConfig)

	// define flags
	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.unseal/config.yaml)")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" { // enable ability to specify config file via flag
		viper.SetConfigFile(cfgFile)
	} else {
		viper.SetConfigName("config") // name of config file (without extension)
		viper.AddConfigPath("/etc/unseal")
		viper.AddConfigPath("$HOME/.unseal") // adding home directory as first search path
		viper.AddConfigPath(".")
		viper.AutomaticEnv() // read in environment variables that match
	}

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err != nil {
		log.Fatal("Error reading config file: ", err)
	}

	home, err := homedir.Dir()
	if err != nil {
		log.Error("Error getting home directory: ", err)
	}
	viper.SetDefault("gpgsecretkeyring", home+"/.gnupg/secring.gpg")
	viper.SetDefault("gpgpublickeyring", home+"/.gnupg/pubring.gpg")
}
<commit_msg>Validate key exists before sending unseal op<commit_after>// Copyright © 2017 Lee Briggs <lee@leebriggs.co.uk>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package cmd

import (
	"fmt"
	"os"
	"sync"

	g "github.com/jaxxstorm/unseal/gpg"
	v "github.com/jaxxstorm/unseal/vault"

	log "github.com/Sirupsen/logrus"
	"github.com/bgentry/speakeasy"
	"github.com/mitchellh/go-homedir"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var cfgFile string
var unsealKey string
var vaultHost string
var vaultPort int

var caPath string
var gpgPub string
var gpgSecret string
var gpgPass string

type Host struct {
	Name string
	Port int
	Key  string
}

var hosts []Host

var Version string

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "unseal",
	Short: "Unseal a set of vault servers",
	Long:  `Unseal allows you to unseal a large set of vault servers using the HTTP API.`,

	// Uncomment the following line if your bare application
	// has an action associated with it:
	Run: func(cmd *cobra.Command, args []string) {

		// unmarshal config file
		err := viper.UnmarshalKey("hosts", &hosts)

		// check for valid config file
		if err != nil {
			log.Fatal("Unable to read hosts key in config file: %s", err)
		}

		caPath = viper.GetString("capath")

		gpg := viper.GetBool("gpg")

		if gpg == true {
			log.Info("Using GPG")
			gpgSecret = viper.GetString("gpgsecretkeyring")
			gpgPub = viper.GetString("gpgpublickeyring")
			gpgPass, err = speakeasy.Ask("Please enter your password: ")
			if err != nil {
				log.Fatal("Password error")
			}
		}

		var wg sync.WaitGroup

		for _, h := range hosts {

			hostName := h.Name
			hostPort := h.Port
			key := h.Key

			var vaultKey string

			if gpg == true {
				vaultKey, err = g.Decrypt(gpgPub, gpgSecret, key, gpgPass)
				if err != nil {
					log.Fatal("GPG Decrypt Error: ", err)
				}
			} else {
				vaultKey = key
			}

			wg.Add(1)

			go func(hostName string, hostPort int) {
				defer wg.Done()

				// create a vault client
				client, err := v.VaultClient(hostName, hostPort, caPath)
				if err != nil {
					log.WithFields(log.Fields{"host": hostName, "port": hostPort}).Error(err)
				}
				// get the current status
				init := v.InitStatus(client)
				if init.Ready == true {
					if vaultKey != "" {
						result, err := client.Sys().Unseal(vaultKey)
						// error while unsealing
						if err != nil {
							log.WithFields(log.Fields{"host": hostName}).Error("Error running unseal operation")
						}

						// if it's still sealed, print the progress
						if result.Sealed == true {
							log.WithFields(log.Fields{"host": hostName, "progress": result.Progress, "threshold": result.T}).Info("Unseal operation performed")
							// otherwise, tell us it's unsealed!
						} else {
							log.WithFields(log.Fields{"host": hostName, "progress": result.Progress, "threshold": result.T}).Info("Vault is unsealed!")
						}
						// zero out the key
						// FIXME: is this the best way to do this?
						// Is it safe?
						key = ""
					} else {
						log.WithFields(log.Fields{"host": hostName}).Error("No Key Provided")
					}
				} else {
					// sad times, not ready to be unsealed
					log.WithFields(log.Fields{"host": hostName}).Error("Vault is not ready to be unsealed")
				}

			}(hostName, hostPort)

		}
		wg.Wait()
	},
}

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

func init() {
	cobra.OnInitialize(initConfig)

	// define flags
	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.unseal/config.yaml)")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if cfgFile != "" { // enable ability to specify config file via flag
		viper.SetConfigFile(cfgFile)
	} else {
		viper.SetConfigName("config") // name of config file (without extension)
		viper.AddConfigPath("/etc/unseal")
		viper.AddConfigPath("$HOME/.unseal") // adding home directory as first search path
		viper.AddConfigPath(".")
		viper.AutomaticEnv() // read in environment variables that match
	}

	// If a config file is found, read it in.
	if err := viper.ReadInConfig(); err != nil {
		log.Fatal("Error reading config file: ", err)
	}

	home, err := homedir.Dir()
	if err != nil {
		log.Error("Error getting home directory: ", err)
	}
	viper.SetDefault("gpgsecretkeyring", home+"/.gnupg/secring.gpg")
	viper.SetDefault("gpgpublickeyring", home+"/.gnupg/pubring.gpg")
}
<|endoftext|>


<commit_before>package cacher

import (
	"sync"
	"time"
)

type (
	CacheFiller func() (data interface{}, err error)

	Cacher struct {
		data map[string]*cacheItem
		mux  sync.Mutex
	}
)

// New returns a new cacher
// purgeEvery sets the timer for purging old cache, set to 0 to never purge.
func New(purgeEvery time.Duration) *Cacher {
	c := &Cacher{
		data: map[string]*cacheItem{},
	}
	if purgeEvery > 0 {
		go c.purge(purgeEvery)
	}
	return c
}

func (c *Cacher) Get(key string, fn CacheFiller, ttl time.Duration) (data interface{}, err error) {
	var ci *cacheItem
	c.mux.Lock()
	if ci = c.data[key]; ci == nil {
		ci = &cacheItem{
			fn:  fn,
			ttl: int64(ttl / time.Second),
		}
		c.data[key] = ci
	}
	c.mux.Unlock()

	return ci.call()
}

func (c *Cacher) Delete(key string) {
	c.mux.Lock()
	delete(c.data, key)
	c.mux.Unlock()
}

func (c *Cacher) purge(ttl time.Duration) {
	for {
		ts := time.Now().Add(ttl).Unix()
		c.mux.Lock()
		for key, ci := range c.data {
			if ci.expiresAt > ts {
				delete(c.data, key)
			}
		}
		c.mux.Unlock()
		time.Sleep(ttl)
	}
}

type cacheItem struct {
	fn        CacheFiller
	expiresAt int64
	ttl       int64

	data interface{}
	err  error

	sync.Mutex
}

func (ci *cacheItem) call() (data interface{}, err error) {
	ts := time.Now().Unix()
	ci.Lock()
	if ci.expiresAt == 0 || ci.expiresAt < ts {
		ci.data, ci.err = ci.fn()
		ci.expiresAt = ts + ci.ttl
	}
	data, err = ci.data, ci.err
	ci.Unlock()
	return
}
<commit_msg>sleep first before trying to purge right away<commit_after>package cacher

import (
	"sync"
	"time"
)

type (
	CacheFiller func() (data interface{}, err error)

	Cacher struct {
		data map[string]*cacheItem
		mux  sync.Mutex
	}
)

// New returns a new cacher
// purgeEvery sets the timer for purging old cache, set to 0 to never purge.
func New(purgeEvery time.Duration) *Cacher {
	c := &Cacher{
		data: map[string]*cacheItem{},
	}
	if purgeEvery > 0 {
		go c.purge(purgeEvery)
	}
	return c
}

func (c *Cacher) Get(key string, fn CacheFiller, ttl time.Duration) (data interface{}, err error) {
	var ci *cacheItem
	c.mux.Lock()
	if ci = c.data[key]; ci == nil {
		ci = &cacheItem{
			fn:  fn,
			ttl: int64(ttl / time.Second),
		}
		c.data[key] = ci
	}
	c.mux.Unlock()

	return ci.call()
}

func (c *Cacher) Delete(key string) {
	c.mux.Lock()
	delete(c.data, key)
	c.mux.Unlock()
}

func (c *Cacher) purge(ttl time.Duration) {
	for {
		time.Sleep(ttl)
		ts := time.Now().Add(ttl).Unix()
		c.mux.Lock()
		for key, ci := range c.data {
			if ci.expiresAt > ts {
				delete(c.data, key)
			}
		}
		c.mux.Unlock()
	}
}

type cacheItem struct {
	fn        CacheFiller
	expiresAt int64
	ttl       int64

	data interface{}
	err  error

	sync.Mutex
}

func (ci *cacheItem) call() (data interface{}, err error) {
	ts := time.Now().Unix()
	ci.Lock()
	if ci.expiresAt == 0 || ci.expiresAt < ts {
		ci.data, ci.err = ci.fn()
		ci.expiresAt = ts + ci.ttl
	}
	data, err = ci.data, ci.err
	ci.Unlock()
	return
}
<|endoftext|>