<commit_before>package sdp

import (
	"io"
	"strconv"
	"time"
)

// An Encoder writes a session description to a buffer.
type Encoder struct {
	w io.Writer
	b writer
}

// NewEncoder returns a new encoder that writes to w.
func NewEncoder(w io.Writer) *Encoder {
	return NewEncoderSize(w, 1024)
}

// NewEncoderSize returns a new encoder that writes to w.
func NewEncoderSize(w io.Writer, size int) *Encoder {
	return &Encoder{w: w, b: make([]byte, 0, size)}
}

// Encode encodes the session description.
func (e *Encoder) Encode(s *Session) error {
	e.Reset()
	e.b = e.b.session(s)
	if e.w != nil {
		return e.Flush()
	}
	return nil
}

// Flush writes encoded bytes to w.
func (e *Encoder) Flush() error {
	if b := e.b; len(b) > 0 {
		n, err := e.w.Write(b)
		e.b = b[:copy(b, b[n:])]
		return err
	}
	return nil
}

// Reset resets encoder state to be empty.
func (e *Encoder) Reset() {
	e.b = e.b[:0]
	//	e.pos, e.newline = 0, false
}

// Bytes returns encoded bytes of the last session description.
// The bytes stop being valid at the next encoder call.
func (e *Encoder) Bytes() []byte {
	return e.b
}

// Bytes returns the encoded session description as str.
func (e *Encoder) String() string {
	return string(e.b)
}

type writer []byte

func (w writer) char(v byte) writer {
	return append(w, v)
}

func (w writer) add(v byte) writer {
	return append(w, '\r', '\n', v, '=')
}

func (w writer) str(v string) writer {
	if v == "" {
		return append(w, '-')
	}
	return append(w, v...)
}

func (w writer) int(v int64) writer {
	return strconv.AppendInt(w, v, 10)
}

func (w writer) sp() writer {
	return append(w, ' ')
}

func (w writer) crlf() writer {
	return append(w, '\r', '\n')
}

func (w writer) session(s *Session) writer {
	w = w.str("v=").int(int64(s.Version))
	if s.Origin != nil {
		w = w.add('o').origin(s.Origin)
	}
	w = w.add('s').str(s.Name)
	if s.Information != "" {
		w = w.add('i').str(s.Information)
	}
	if s.URI != "" {
		w = w.add('u').str(s.URI)
	}
	for _, it := range s.Email {
		w = w.add('e').str(it)
	}
	for _, it := range s.Phone {
		w = w.add('p').str(it)
	}
	if s.Connection != nil {
		w = w.add('c').connection(s.Connection)
	}
	for _, b := range s.Bandwidth {
		w = w.add('b').bandwidth(b)
	}
	if len(s.TimeZone) > 0 {
		w = w.add('z').timezone(s.TimeZone)
	}
	for _, it := range s.Key {
		w = w.add('k').key(it)
	}
	if s.Mode != "" {
		w = w.add('a').str(s.Mode)
	}
	for _, it := range s.Attributes {
		w = w.add('a').attr(it)
	}
	w = w.add('t').timing(s.Timing)
	for _, it := range s.Repeat {
		w = w.add('r').repeat(it)
	}
	for _, it := range s.Media {
		w = w.media(it)
	}
	return w.crlf()
}

func (w writer) origin(o *Origin) writer {
	return w.str(strdef(o.Username, "-")).sp().int(o.SessionID).sp().int(o.SessionVersion).sp().transport(o.Network, o.Type, o.Address)
}

func (w writer) media(m *Media) writer {
	w = w.add('m').str(m.Type).sp().int(int64(m.Port))
	if m.PortNum > 0 {
		w = w.char('/').int(int64(m.PortNum))
	}
	w = w.sp().str(m.Proto)
	if f := m.FormatDescr; f != "" {
		w = w.sp().str(f)
	} else {
		for _, it := range m.Format {
			w = w.sp().int(int64(it.Payload))
		}
	}
	if m.Information != "" {
		w = w.add('i').str(m.Information)
	}
	for _, it := range m.Connection {
		w = w.add('c').connection(it)
	}
	for _, b := range m.Bandwidth {
		w = w.add('b').bandwidth(b)
	}
	for _, it := range m.Key {
		w = w.add('k').key(it)
	}
	for _, it := range m.Format {
		w = w.format(it)
	}
	if m.Mode != "" {
		w = w.add('a').str(m.Mode)
	}
	for _, it := range m.Attributes {
		w = w.add('a').attr(it)
	}
	return w
}

func (w writer) format(f *Format) writer {
	p := int64(f.Payload)
	if f.Name != "" {
		w = w.add('a').str("rtpmap:").int(p).sp().str(f.Name).char('/').int(int64(f.ClockRate))
		if f.Channels > 0 {
			w = w.char('/').int(int64(f.Channels))
		}
	}
	for _, it := range f.Feedback {
		w = w.add('a').str("rtcp-fb:").int(p).sp().str(it)
	}
	for _, it := range f.Params {
		w = w.add('a').str("fmtp:").int(p).sp().str(it)
	}
	return w
}

func (w writer) attr(a *Attr) writer {
	if a.Value == "" {
		return w.str(a.Name)
	}
	return w.str(a.Name).char(':').str(a.Value)
}

func (w writer) timezone(z []*TimeZone) writer {
	for i, it := range z {
		if i > 0 {
			w = w.char(' ')
		}
		w = w.time(it.Time).sp().duration(it.Offset)
	}
	return w
}

func (w writer) timing(t *Timing) writer {
	if t == nil {
		return w.str("0 0")
	}
	return w.time(t.Start).sp().time(t.Stop)
}

func (w writer) repeat(r *Repeat) writer {
	w = w.duration(r.Interval).sp().duration(r.Duration)
	for _, it := range r.Offsets {
		w = w.sp().duration(it)
	}
	return w
}

func (w writer) time(t time.Time) writer {
	if t.IsZero() {
		return w.char('0')
	}
	return w.int(int64(t.Sub(epoch).Seconds()))
}

func (w writer) duration(d time.Duration) writer {
	v := int64(d.Seconds())
	switch {
	case v == 0:
		return w.char('0')
	case v%86400 == 0:
		return w.int(v / 86400).char('d')
	case v%3600 == 0:
		return w.int(v / 3600).char('h')
	case v%60 == 0:
		return w.int(v / 60).char('m')
	default:
		return w.int(v)
	}
}

func (w writer) bandwidth(b *Bandwidth) writer {
	return w.str(b.Type).char(':').int(int64(b.Value))
}

func (w writer) key(k *Key) writer {
	if k.Value == "" {
		return w.str(k.Method)
	}
	return w.str(k.Method).char(':').str(k.Value)
}

func (w writer) connection(c *Connection) writer {
	w = w.transport(c.Network, c.Type, c.Address)
	if c.TTL > 0 {
		w = w.char('/').int(int64(c.TTL))
	}
	if c.AddressNum > 1 {
		w = w.char('/').int(int64(c.AddressNum))
	}
	return w
}

func (w writer) transport(network, typ, addr string) writer {
	return w.fields(strdef(network, NetworkInternet), strdef(typ, TypeIPv4), strdef(addr, "127.0.0.1"))
}

func strdef(v, def string) string {
	if v == "" {
		return def
	}
	return v
}

func (w writer) fields(v ...string) writer {
	for i, s := range v {
		if i > 0 {
			w = append(w, ' ')
		}
		w = append(w, s...)
	}
	return w
}
<commit_msg>Fix timing order<commit_after>package sdp

import (
	"io"
	"strconv"
	"time"
)

// An Encoder writes a session description to a buffer.
type Encoder struct {
	w io.Writer
	b writer
}

// NewEncoder returns a new encoder that writes to w.
func NewEncoder(w io.Writer) *Encoder {
	return NewEncoderSize(w, 1024)
}

// NewEncoderSize returns a new encoder that writes to w.
func NewEncoderSize(w io.Writer, size int) *Encoder {
	return &Encoder{w: w, b: make([]byte, 0, size)}
}

// Encode encodes the session description.
func (e *Encoder) Encode(s *Session) error {
	e.Reset()
	e.b = e.b.session(s)
	if e.w != nil {
		return e.Flush()
	}
	return nil
}

// Flush writes encoded bytes to w.
func (e *Encoder) Flush() error {
	if b := e.b; len(b) > 0 {
		n, err := e.w.Write(b)
		e.b = b[:copy(b, b[n:])]
		return err
	}
	return nil
}

// Reset resets encoder state to be empty.
func (e *Encoder) Reset() {
	e.b = e.b[:0]
	//	e.pos, e.newline = 0, false
}

// Bytes returns encoded bytes of the last session description.
// The bytes stop being valid at the next encoder call.
func (e *Encoder) Bytes() []byte {
	return e.b
}

// Bytes returns the encoded session description as str.
func (e *Encoder) String() string {
	return string(e.b)
}

type writer []byte

func (w writer) char(v byte) writer {
	return append(w, v)
}

func (w writer) add(v byte) writer {
	return append(w, '\r', '\n', v, '=')
}

func (w writer) str(v string) writer {
	if v == "" {
		return append(w, '-')
	}
	return append(w, v...)
}

func (w writer) int(v int64) writer {
	return strconv.AppendInt(w, v, 10)
}

func (w writer) sp() writer {
	return append(w, ' ')
}

func (w writer) crlf() writer {
	return append(w, '\r', '\n')
}

func (w writer) session(s *Session) writer {
	w = w.str("v=").int(int64(s.Version))
	if s.Origin != nil {
		w = w.add('o').origin(s.Origin)
	}
	w = w.add('s').str(s.Name)
	if s.Information != "" {
		w = w.add('i').str(s.Information)
	}
	if s.URI != "" {
		w = w.add('u').str(s.URI)
	}
	for _, it := range s.Email {
		w = w.add('e').str(it)
	}
	for _, it := range s.Phone {
		w = w.add('p').str(it)
	}
	if s.Connection != nil {
		w = w.add('c').connection(s.Connection)
	}
	for _, b := range s.Bandwidth {
		w = w.add('b').bandwidth(b)
	}
	if len(s.TimeZone) > 0 {
		w = w.add('z').timezone(s.TimeZone)
	}
	for _, it := range s.Key {
		w = w.add('k').key(it)
	}
	w = w.add('t').timing(s.Timing)
	for _, it := range s.Repeat {
		w = w.add('r').repeat(it)
	}
	if s.Mode != "" {
		w = w.add('a').str(s.Mode)
	}
	for _, it := range s.Attributes {
		w = w.add('a').attr(it)
	}
	for _, it := range s.Media {
		w = w.media(it)
	}
	return w.crlf()
}

func (w writer) origin(o *Origin) writer {
	return w.str(strdef(o.Username, "-")).sp().int(o.SessionID).sp().int(o.SessionVersion).sp().transport(o.Network, o.Type, o.Address)
}

func (w writer) media(m *Media) writer {
	w = w.add('m').str(m.Type).sp().int(int64(m.Port))
	if m.PortNum > 0 {
		w = w.char('/').int(int64(m.PortNum))
	}
	w = w.sp().str(m.Proto)
	if f := m.FormatDescr; f != "" {
		w = w.sp().str(f)
	} else {
		for _, it := range m.Format {
			w = w.sp().int(int64(it.Payload))
		}
	}
	if m.Information != "" {
		w = w.add('i').str(m.Information)
	}
	for _, it := range m.Connection {
		w = w.add('c').connection(it)
	}
	for _, b := range m.Bandwidth {
		w = w.add('b').bandwidth(b)
	}
	for _, it := range m.Key {
		w = w.add('k').key(it)
	}
	for _, it := range m.Format {
		w = w.format(it)
	}
	if m.Mode != "" {
		w = w.add('a').str(m.Mode)
	}
	for _, it := range m.Attributes {
		w = w.add('a').attr(it)
	}
	return w
}

func (w writer) format(f *Format) writer {
	p := int64(f.Payload)
	if f.Name != "" {
		w = w.add('a').str("rtpmap:").int(p).sp().str(f.Name).char('/').int(int64(f.ClockRate))
		if f.Channels > 0 {
			w = w.char('/').int(int64(f.Channels))
		}
	}
	for _, it := range f.Feedback {
		w = w.add('a').str("rtcp-fb:").int(p).sp().str(it)
	}
	for _, it := range f.Params {
		w = w.add('a').str("fmtp:").int(p).sp().str(it)
	}
	return w
}

func (w writer) attr(a *Attr) writer {
	if a.Value == "" {
		return w.str(a.Name)
	}
	return w.str(a.Name).char(':').str(a.Value)
}

func (w writer) timezone(z []*TimeZone) writer {
	for i, it := range z {
		if i > 0 {
			w = w.char(' ')
		}
		w = w.time(it.Time).sp().duration(it.Offset)
	}
	return w
}

func (w writer) timing(t *Timing) writer {
	if t == nil {
		return w.str("0 0")
	}
	return w.time(t.Start).sp().time(t.Stop)
}

func (w writer) repeat(r *Repeat) writer {
	w = w.duration(r.Interval).sp().duration(r.Duration)
	for _, it := range r.Offsets {
		w = w.sp().duration(it)
	}
	return w
}

func (w writer) time(t time.Time) writer {
	if t.IsZero() {
		return w.char('0')
	}
	return w.int(int64(t.Sub(epoch).Seconds()))
}

func (w writer) duration(d time.Duration) writer {
	v := int64(d.Seconds())
	switch {
	case v == 0:
		return w.char('0')
	case v%86400 == 0:
		return w.int(v / 86400).char('d')
	case v%3600 == 0:
		return w.int(v / 3600).char('h')
	case v%60 == 0:
		return w.int(v / 60).char('m')
	default:
		return w.int(v)
	}
}

func (w writer) bandwidth(b *Bandwidth) writer {
	return w.str(b.Type).char(':').int(int64(b.Value))
}

func (w writer) key(k *Key) writer {
	if k.Value == "" {
		return w.str(k.Method)
	}
	return w.str(k.Method).char(':').str(k.Value)
}

func (w writer) connection(c *Connection) writer {
	w = w.transport(c.Network, c.Type, c.Address)
	if c.TTL > 0 {
		w = w.char('/').int(int64(c.TTL))
	}
	if c.AddressNum > 1 {
		w = w.char('/').int(int64(c.AddressNum))
	}
	return w
}

func (w writer) transport(network, typ, addr string) writer {
	return w.fields(strdef(network, NetworkInternet), strdef(typ, TypeIPv4), strdef(addr, "127.0.0.1"))
}

func strdef(v, def string) string {
	if v == "" {
		return def
	}
	return v
}

func (w writer) fields(v ...string) writer {
	for i, s := range v {
		if i > 0 {
			w = append(w, ' ')
		}
		w = append(w, s...)
	}
	return w
}
<|endoftext|>


<commit_before>package accessor

import (
	"encoding/json"

	"github.com/concourse/concourse/atc/db"
	"github.com/golang/groupcache/lru"
)

type claimsCacheEntry struct {
	claims db.Claims
	size   int
}

type claimsCacher struct {
	accessTokenFetcher AccessTokenFetcher
	maxCacheSizeBytes  int

	cache          *lru.Cache
	cacheSizeBytes int
}

func NewClaimsCacher(
	accessTokenFetcher AccessTokenFetcher,
	maxCacheSizeBytes int,
) *claimsCacher {
	c := &claimsCacher{
		accessTokenFetcher: accessTokenFetcher,
		maxCacheSizeBytes:  maxCacheSizeBytes,
		cache:              lru.New(0),
	}
	c.cache.OnEvicted = func(_ lru.Key, value interface{}) {
		entry, _ := value.(claimsCacheEntry)
		c.cacheSizeBytes -= entry.size
	}

	return c
}

func (c *claimsCacher) GetAccessToken(rawToken string) (db.AccessToken, bool, error) {
	claims, found := c.cache.Get(rawToken)
	if found {
		entry, _ := claims.(claimsCacheEntry)
		return db.AccessToken{Token: rawToken, Claims: entry.claims}, true, nil
	}

	token, found, err := c.accessTokenFetcher.GetAccessToken(rawToken)
	if err != nil {
		return db.AccessToken{}, false, err
	}
	payload, err := json.Marshal(token.Claims)
	if err != nil {
		return db.AccessToken{}, false, err
	}
	entry := claimsCacheEntry{claims: token.Claims, size: len(payload)}
	c.cache.Add(rawToken, entry)
	c.cacheSizeBytes += entry.size

	for c.cacheSizeBytes > c.maxCacheSizeBytes && c.cache.Len() > 0 {
		c.cache.RemoveOldest()
	}

	return token, true, nil
}
<commit_msg>make claims cacher safe for concurrent use<commit_after>package accessor

import (
	"encoding/json"
	"sync"

	"github.com/concourse/concourse/atc/db"
	"github.com/golang/groupcache/lru"
)

type claimsCacheEntry struct {
	claims db.Claims
	size   int
}

type claimsCacher struct {
	accessTokenFetcher AccessTokenFetcher
	maxCacheSizeBytes  int

	cache          *lru.Cache
	cacheSizeBytes int
	mu             sync.Mutex // lru.Cache is not safe for concurrent access
}

func NewClaimsCacher(
	accessTokenFetcher AccessTokenFetcher,
	maxCacheSizeBytes int,
) *claimsCacher {
	c := &claimsCacher{
		accessTokenFetcher: accessTokenFetcher,
		maxCacheSizeBytes:  maxCacheSizeBytes,
		cache:              lru.New(0),
	}
	c.cache.OnEvicted = func(_ lru.Key, value interface{}) {
		entry, _ := value.(claimsCacheEntry)
		c.cacheSizeBytes -= entry.size
	}

	return c
}

func (c *claimsCacher) GetAccessToken(rawToken string) (db.AccessToken, bool, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	claims, found := c.cache.Get(rawToken)
	if found {
		entry, _ := claims.(claimsCacheEntry)
		return db.AccessToken{Token: rawToken, Claims: entry.claims}, true, nil
	}

	token, found, err := c.accessTokenFetcher.GetAccessToken(rawToken)
	if err != nil {
		return db.AccessToken{}, false, err
	}
	payload, err := json.Marshal(token.Claims)
	if err != nil {
		return db.AccessToken{}, false, err
	}
	entry := claimsCacheEntry{claims: token.Claims, size: len(payload)}
	c.cache.Add(rawToken, entry)
	c.cacheSizeBytes += entry.size

	for c.cacheSizeBytes > c.maxCacheSizeBytes && c.cache.Len() > 0 {
		c.cache.RemoveOldest()
	}

	return token, true, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 RedHat, Inc.
// Copyright 2015 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 sdjournal

import (
	"errors"
	"fmt"
	"io"
	"log"
	"strings"
	"time"

	"golang.org/x/net/context"
)

var (
	ErrExpired = errors.New("Timeout expired")
)

// JournalReaderConfig represents options to drive the behavior of a JournalReader.
type JournalReaderConfig struct {
	// The Since and NumFromTail options are mutually exclusive and determine
	// where the reading begins within the journal.
	Since       time.Duration // start relative to a Duration from now
	NumFromTail uint64        // start relative to the tail

	// Show only journal entries whose fields match the supplied values. If
	// the array is empty, entries will not be filtered.
	Matches []Match

	// If not empty, the journal instance will point to a journal residing
	// in this directory. The supplied path may be relative or absolute.
	Path string
}

// JournalReader is an io.ReadCloser which provides a simple interface for iterating through the
// systemd journal.
type JournalReader struct {
	journal *Journal
}

// FollowFilter is a function which you pass into Follow to determine if you to
// filter key value pair from a journal entry.
type FollowFilter func(key, value string) bool

// AllKeys is a FollowFilter that allows all keys of a given journal entry through the filter
func AllKeys(key, value string) bool {
	return true
}

func OnlyMessages(key, value string) bool {
	if key == "MESSAGE" {
		return true
	}
	return false
}

// NewJournalReader creates a new JournalReader with configuration options that are similar to the
// systemd journalctl tool's iteration and filtering features.
func NewJournalReader(config JournalReaderConfig) (*JournalReader, error) {
	r := &JournalReader{}

	// Open the journal
	var err error
	if config.Path != "" {
		r.journal, err = NewJournalFromDir(config.Path)
	} else {
		r.journal, err = NewJournal()
	}
	if err != nil {
		return nil, err
	}

	// Add any supplied matches
	for _, m := range config.Matches {
		r.journal.AddMatch(m.String())
	}

	// Set the start position based on options
	if config.Since != 0 {
		// Start based on a relative time
		start := time.Now().Add(config.Since)
		if err := r.journal.SeekRealtimeUsec(uint64(start.UnixNano() / 1000)); err != nil {
			return nil, err
		}
	} else if config.NumFromTail != 0 {
		// Start based on a number of lines before the tail
		if err := r.journal.SeekTail(); err != nil {
			return nil, err
		}

		// Move the read pointer into position near the tail. Go one further than
		// the option so that the initial cursor advancement positions us at the
		// correct starting point.
		if _, err := r.journal.PreviousSkip(config.NumFromTail + 1); err != nil {
			return nil, err
		}
	}

	return r, nil
}

func (r *JournalReader) Read(b []byte) (int, error) {
	var err error
	var c int

	// Advance the journal cursor
	c, err = r.journal.Next()

	// An unexpected error
	if err != nil {
		return 0, err
	}

	// EOF detection
	if c == 0 {
		return 0, io.EOF
	}

	// Build a message
	var msg string
	msg, err = r.buildMessage()

	if err != nil {
		return 0, err
	}

	// Copy and return the message
	copy(b, []byte(msg))

	return len(msg), nil
}

func (r *JournalReader) Close() error {
	return r.journal.Close()
}

// Follow asynchronously follows the JournalReader it takes in a context to stop the following the Journal, it returns a
// buffered channel of errors and will stop following the journal on the first given error
func (r *JournalReader) Follow(ctx context.Context, msgs chan<- map[string]interface{}, filter FollowFilter) <-chan error {
	errChan := make(chan error, 1)
	go func() {
		for {
			kvMap := make(map[string]interface{})
			select {
			case <-ctx.Done():
				errChan <- ctx.Err()
				return
			default:
				var err error
				var c int

				// Advance the journal cursor
				c, err = r.journal.Next()

				// An unexpected error
				if err != nil {
					errChan <- err
					return
				}

				// We have a new journal entry go over the fields
				// get the data for what we care about and return
				r.journal.RestartData()
				if c > 0 {
				fields:
					for {
						s, err := r.journal.EnumerateData()
						if err != nil || len(s) == 0 {
							break fields
						}
						s = s[:len(s)]
						arr := strings.SplitN(s, "=", 2)
						// if we want the pair,
						// add it to the map
						if filter(arr[0], arr[1]) {
							kvMap[arr[0]] = arr[1]
						}
					}
					msgs <- kvMap
				}
			}

			// we're at the tail, so wait for new events or time out.
			// holds journal events to process. tightly bounded for now unless there's a
			// reason to unblock the journal watch routine more quickly
			events := make(chan int, 1)
			pollDone := make(chan bool, 1)
			go func() {
				for {
					select {
					case <-pollDone:
						return
					default:
						events <- r.journal.Wait(time.Duration(1) * time.Second)
					}
				}
			}()

			select {
			case <-ctx.Done():
				errChan <- ctx.Err()
				pollDone <- true
				return
			case e := <-events:
				pollDone <- true
				switch e {
				case SD_JOURNAL_NOP, SD_JOURNAL_APPEND, SD_JOURNAL_INVALIDATE:
					// TODO: need to account for any of these?
				default:
					log.Printf("Received unknown event: %d\n", e)
				}
			}
		}
	}()
	return errChan
}

// buildMessage returns a string representing the current journal entry in a simple format which
// includes the entry timestamp and MESSAGE field.
func (r *JournalReader) buildMessage() (string, error) {
	var msg string
	var usec uint64
	var err error

	if msg, err = r.journal.GetData("MESSAGE"); err != nil {
		return "", err
	}

	if usec, err = r.journal.GetRealtimeUsec(); err != nil {
		return "", err
	}

	timestamp := time.Unix(0, int64(usec)*int64(time.Microsecond))

	return fmt.Sprintf("%s %s\n", timestamp, msg), nil
}
<commit_msg>change wait time<commit_after>// Copyright 2015 RedHat, Inc.
// Copyright 2015 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 sdjournal

import (
	"errors"
	"fmt"
	"io"
	"log"
	"strings"
	"time"

	"golang.org/x/net/context"
)

var (
	ErrExpired = errors.New("Timeout expired")
)

// JournalReaderConfig represents options to drive the behavior of a JournalReader.
type JournalReaderConfig struct {
	// The Since and NumFromTail options are mutually exclusive and determine
	// where the reading begins within the journal.
	Since       time.Duration // start relative to a Duration from now
	NumFromTail uint64        // start relative to the tail

	// Show only journal entries whose fields match the supplied values. If
	// the array is empty, entries will not be filtered.
	Matches []Match

	// If not empty, the journal instance will point to a journal residing
	// in this directory. The supplied path may be relative or absolute.
	Path string
}

// JournalReader is an io.ReadCloser which provides a simple interface for iterating through the
// systemd journal.
type JournalReader struct {
	journal *Journal
}

// FollowFilter is a function which you pass into Follow to determine if you to
// filter key value pair from a journal entry.
type FollowFilter func(key, value string) bool

// AllKeys is a FollowFilter that allows all keys of a given journal entry through the filter
func AllKeys(key, value string) bool {
	return true
}

func OnlyMessages(key, value string) bool {
	if key == "MESSAGE" {
		return true
	}
	return false
}

// NewJournalReader creates a new JournalReader with configuration options that are similar to the
// systemd journalctl tool's iteration and filtering features.
func NewJournalReader(config JournalReaderConfig) (*JournalReader, error) {
	r := &JournalReader{}

	// Open the journal
	var err error
	if config.Path != "" {
		r.journal, err = NewJournalFromDir(config.Path)
	} else {
		r.journal, err = NewJournal()
	}
	if err != nil {
		return nil, err
	}

	// Add any supplied matches
	for _, m := range config.Matches {
		r.journal.AddMatch(m.String())
	}

	// Set the start position based on options
	if config.Since != 0 {
		// Start based on a relative time
		start := time.Now().Add(config.Since)
		if err := r.journal.SeekRealtimeUsec(uint64(start.UnixNano() / 1000)); err != nil {
			return nil, err
		}
	} else if config.NumFromTail != 0 {
		// Start based on a number of lines before the tail
		if err := r.journal.SeekTail(); err != nil {
			return nil, err
		}

		// Move the read pointer into position near the tail. Go one further than
		// the option so that the initial cursor advancement positions us at the
		// correct starting point.
		if _, err := r.journal.PreviousSkip(config.NumFromTail + 1); err != nil {
			return nil, err
		}
	}

	return r, nil
}

func (r *JournalReader) Read(b []byte) (int, error) {
	var err error
	var c int

	// Advance the journal cursor
	c, err = r.journal.Next()

	// An unexpected error
	if err != nil {
		return 0, err
	}

	// EOF detection
	if c == 0 {
		return 0, io.EOF
	}

	// Build a message
	var msg string
	msg, err = r.buildMessage()

	if err != nil {
		return 0, err
	}

	// Copy and return the message
	copy(b, []byte(msg))

	return len(msg), nil
}

func (r *JournalReader) Close() error {
	return r.journal.Close()
}

// Follow asynchronously follows the JournalReader it takes in a context to stop the following the Journal, it returns a
// buffered channel of errors and will stop following the journal on the first given error
func (r *JournalReader) Follow(ctx context.Context, msgs chan<- map[string]interface{}, filter FollowFilter) <-chan error {
	errChan := make(chan error, 1)
	go func() {
		for {
			kvMap := make(map[string]interface{})
			select {
			case <-ctx.Done():
				errChan <- ctx.Err()
				return
			default:
				var err error
				var c int

				// Advance the journal cursor
				c, err = r.journal.Next()

				// An unexpected error
				if err != nil {
					errChan <- err
					return
				}

				// We have a new journal entry go over the fields
				// get the data for what we care about and return
				r.journal.RestartData()
				if c > 0 {
				fields:
					for {
						s, err := r.journal.EnumerateData()
						if err != nil || len(s) == 0 {
							break fields
						}
						s = s[:len(s)]
						arr := strings.SplitN(s, "=", 2)
						// if we want the pair,
						// add it to the map
						if filter(arr[0], arr[1]) {
							kvMap[arr[0]] = arr[1]
						}
					}
					msgs <- kvMap
				}
			}

			// we're at the tail, so wait for new events or time out.
			// holds journal events to process. tightly bounded for now unless there's a
			// reason to unblock the journal watch routine more quickly
			events := make(chan int, 1)
			pollDone := make(chan bool, 1)
			go func() {
				for {
					select {
					case <-pollDone:
						return
					default:
						events <- r.journal.Wait(time.Duration(300) * time.Second)
					}
				}
			}()

			select {
			case <-ctx.Done():
				errChan <- ctx.Err()
				pollDone <- true
				return
			case e := <-events:
				pollDone <- true
				switch e {
				case SD_JOURNAL_NOP, SD_JOURNAL_APPEND, SD_JOURNAL_INVALIDATE:
					// TODO: need to account for any of these?
				default:
					log.Printf("Received unknown event: %d\n", e)
				}
			}
		}
	}()
	return errChan
}

// buildMessage returns a string representing the current journal entry in a simple format which
// includes the entry timestamp and MESSAGE field.
func (r *JournalReader) buildMessage() (string, error) {
	var msg string
	var usec uint64
	var err error

	if msg, err = r.journal.GetData("MESSAGE"); err != nil {
		return "", err
	}

	if usec, err = r.journal.GetRealtimeUsec(); err != nil {
		return "", err
	}

	timestamp := time.Unix(0, int64(usec)*int64(time.Microsecond))

	return fmt.Sprintf("%s %s\n", timestamp, msg), nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Kubernetes Authors.

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 mockcompute

import (
	"context"
	"fmt"
	"sync"

	compute "google.golang.org/api/compute/v1"
	"k8s.io/kops/upup/pkg/fi/cloudup/gce"
)

type targetPoolClient struct {
	// targetPools are targetPools keyed by project, region, and targetPool name.
	targetPools map[string]map[string]map[string]*compute.TargetPool
	sync.Mutex
}

var _ gce.TargetPoolClient = &targetPoolClient{}

func newTargetPoolClient() *targetPoolClient {
	return &targetPoolClient{
		targetPools: map[string]map[string]map[string]*compute.TargetPool{},
	}
}

func (c *targetPoolClient) All() map[string]interface{} {
	c.Lock()
	defer c.Unlock()
	m := map[string]interface{}{}
	for _, regions := range c.targetPools {
		for _, tps := range regions {
			for n, tp := range tps {
				m[n] = tp
			}
		}
	}
	return m
}

func (c *targetPoolClient) Insert(project, region string, tp *compute.TargetPool) (*compute.Operation, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		regions = map[string]map[string]*compute.TargetPool{}
		c.targetPools[project] = regions
	}
	tps, ok := regions[region]
	if !ok {
		tps = map[string]*compute.TargetPool{}
		regions[region] = tps
	}
	tp.SelfLink = fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/regions/%s/targetPools/%s", project, region, tp.Name)
	tps[tp.Name] = tp
	return doneOperation(), nil
}

func (c *targetPoolClient) Delete(project, region, name string) (*compute.Operation, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		return nil, notFoundError()
	}
	tps, ok := regions[region]
	if !ok {
		return nil, notFoundError()
	}
	if _, ok := tps[name]; !ok {
		return nil, notFoundError()
	}
	delete(tps, name)
	return doneOperation(), nil
}

func (c *targetPoolClient) Get(project, region, name string) (*compute.TargetPool, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		return nil, notFoundError()
	}
	tps, ok := regions[region]
	if !ok {
		return nil, notFoundError()
	}
	tp, ok := tps[name]
	if !ok {
		return nil, notFoundError()
	}
	return tp, nil
}

func (c *targetPoolClient) List(ctx context.Context, project, region string) ([]*compute.TargetPool, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		return nil, nil
	}
	tps, ok := regions[region]
	if !ok {
		return nil, nil
	}
	var l []*compute.TargetPool
	for _, tp := range tps {
		l = append(l, tp)
	}
	return l, nil
}

func (c *targetPoolClient) List(ctx context.Context, project, region string) ([]*compute.TargetPool, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		return nil, nil
	}
	tps, ok := regions[region]
	if !ok {
		return nil, nil
	}
	var l []*compute.TargetPool
	for _, tp := range tps {
		l = append(l, tp)
	}
	return l, nil
}

func (c *targetPoolClient) AddHealthCheck(project, region, name string, req *compute.TargetPoolsAddHealthCheckRequest) (*compute.Operation, error) {
	// TODO: AddHealthCheck test
	return doneOperation(), nil
}
<commit_msg>remove list<commit_after>/*
Copyright 2021 The Kubernetes Authors.

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 mockcompute

import (
	"context"
	"fmt"
	"sync"

	compute "google.golang.org/api/compute/v1"
	"k8s.io/kops/upup/pkg/fi/cloudup/gce"
)

type targetPoolClient struct {
	// targetPools are targetPools keyed by project, region, and targetPool name.
	targetPools map[string]map[string]map[string]*compute.TargetPool
	sync.Mutex
}

var _ gce.TargetPoolClient = &targetPoolClient{}

func newTargetPoolClient() *targetPoolClient {
	return &targetPoolClient{
		targetPools: map[string]map[string]map[string]*compute.TargetPool{},
	}
}

func (c *targetPoolClient) All() map[string]interface{} {
	c.Lock()
	defer c.Unlock()
	m := map[string]interface{}{}
	for _, regions := range c.targetPools {
		for _, tps := range regions {
			for n, tp := range tps {
				m[n] = tp
			}
		}
	}
	return m
}

func (c *targetPoolClient) Insert(project, region string, tp *compute.TargetPool) (*compute.Operation, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		regions = map[string]map[string]*compute.TargetPool{}
		c.targetPools[project] = regions
	}
	tps, ok := regions[region]
	if !ok {
		tps = map[string]*compute.TargetPool{}
		regions[region] = tps
	}
	tp.SelfLink = fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/regions/%s/targetPools/%s", project, region, tp.Name)
	tps[tp.Name] = tp
	return doneOperation(), nil
}

func (c *targetPoolClient) Delete(project, region, name string) (*compute.Operation, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		return nil, notFoundError()
	}
	tps, ok := regions[region]
	if !ok {
		return nil, notFoundError()
	}
	if _, ok := tps[name]; !ok {
		return nil, notFoundError()
	}
	delete(tps, name)
	return doneOperation(), nil
}

func (c *targetPoolClient) Get(project, region, name string) (*compute.TargetPool, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		return nil, notFoundError()
	}
	tps, ok := regions[region]
	if !ok {
		return nil, notFoundError()
	}
	tp, ok := tps[name]
	if !ok {
		return nil, notFoundError()
	}
	return tp, nil
}

func (c *targetPoolClient) List(ctx context.Context, project, region string) ([]*compute.TargetPool, error) {
	c.Lock()
	defer c.Unlock()
	regions, ok := c.targetPools[project]
	if !ok {
		return nil, nil
	}
	tps, ok := regions[region]
	if !ok {
		return nil, nil
	}
	var l []*compute.TargetPool
	for _, tp := range tps {
		l = append(l, tp)
	}
	return l, nil
}

func (c *targetPoolClient) AddHealthCheck(project, region, name string, req *compute.TargetPoolsAddHealthCheckRequest) (*compute.Operation, error) {
	// TODO: AddHealthCheck test
	return doneOperation(), nil
}
<|endoftext|>


<commit_before>package lib

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"strings"
	"text/template"

	"github.com/Masterminds/sprig"
	"github.com/fatih/color"
)

// TLSDescription has the basic information about a TLS connection
type TLSDescription struct {
	Version string `json:"version"`
	Cipher  string `json:"cipher"`
}

// CertificateRequestInfo has the basic information about requested client certificates
type CertificateRequestInfo struct {
	AcceptableCAs    []simplePKIXName `json:"acceptable_issuers,omitempty"`
	SignatureSchemes []string         `json:"signature_schemes,omitempty"`
}

type tlsInfoContext struct {
	Conn *TLSDescription
	CRI  *CertificateRequestInfo
}

var tlsLayout = `** TLS Connection **
Version: {{.Conn.Version}}
Cipher Suite: {{.Conn.Cipher}}
{{- if .CRI}}

{{"Server has requested a client certificate:" | greenify}}
Acceptable issuers:
{{- range .CRI.AcceptableCAs}}
	{{.Name | printShortName}}{{end}}
Supported Signature Schemes:
{{- range .CRI.SignatureSchemes}}
	{{.}}{{end}}
{{- end}}`

func tlscolor(d description) string {
	c, ok := qualityColors[d.Quality]
	if !ok {
		return d.Name
	}
	return c.SprintFunc()(d.Name)
}

// EncodeTLSInfoToText returns a human readable string, suitable for certigo console output.
func EncodeTLSInfoToText(tcs *tls.ConnectionState, cri *tls.CertificateRequestInfo) string {
	version := lookup(tlsVersions, tcs.Version)
	cipher := lookup(cipherSuites, tcs.CipherSuite)
	description := TLSDescription{
		Version: tlscolor(version),
		Cipher:  tlscolor(explainCipher(cipher)),
	}
	tlsInfoContext := tlsInfoContext{
		Conn: &description,
	}
	if cri != nil {
		criDesc, err := EncodeCRIToObject(cri)
		if err == nil {
			tlsInfoContext.CRI = criDesc.(*CertificateRequestInfo)
		}
	}

	funcMap := sprig.TxtFuncMap()
	extras := template.FuncMap{
		"printCommonName": PrintCommonName,
		"printShortName":  PrintShortName,
		"greenify":        greenify,
	}
	for k, v := range extras {
		funcMap[k] = v
	}

	t := template.New("TLS template").Funcs(funcMap)
	t, err := t.Parse(tlsLayout)
	if err != nil {
		// Should never happen
		panic(err)
	}
	var buffer bytes.Buffer
	w := bufio.NewWriter(&buffer)
	err = t.Execute(w, tlsInfoContext)
	if err != nil {
		// Should never happen
		panic(err)
	}
	w.Flush()
	return string(buffer.Bytes())
}

// EncodeTLSToObject returns a JSON-marshallable description of a TLS connection
func EncodeTLSToObject(t *tls.ConnectionState) interface{} {
	version := lookup(tlsVersions, t.Version)
	cipher := lookup(cipherSuites, t.CipherSuite)
	return &TLSDescription{
		version.Slug,
		cipher.Slug,
	}
}

// EncodeCRIToObject returns a JSON-marshallable representation of a CertificateRequestInfo object.
func EncodeCRIToObject(cri *tls.CertificateRequestInfo) (interface{}, error) {
	out := &CertificateRequestInfo{}
	for _, ca := range cri.AcceptableCAs {
		subject, err := parseRawSubject(ca)
		if err != nil {
			return nil, err
		}
		out.AcceptableCAs = append(out.AcceptableCAs, simplePKIXName{subject, nil})
	}
	for _, scheme := range cri.SignatureSchemes {
		desc, ok := signatureSchemeStrings[scheme]
		if !ok {
			desc = fmt.Sprintf("Unknown(0x%x)", scheme)
		}
		out.SignatureSchemes = append(out.SignatureSchemes, desc)
	}
	return out, nil
}

// Just a map lookup with a default
func lookup(descriptions map[uint16]description, what uint16) description {
	v, ok := descriptions[what]
	if !ok {
		unknown := fmt.Sprintf("UNKNOWN_%x", what)
		return description{unknown, unknown, 0}
	}
	return v
}

const (
	insecure = iota
	ok       = iota
	good     = iota
)

type description struct {
	Name    string // a human-friendly string
	Slug    string // a machine-friendly string
	Quality uint8  // insecure, ok, good
}

var qualityColors = map[uint8]*color.Color{
	insecure: red,
	ok:       yellow,
	good:     green,
}

var tlsVersions = map[uint16]description{
	tls.VersionSSL30: {"SSL 3.0", "ssl_3_0", insecure},
	tls.VersionTLS10: {"TLS 1.0", "tls_1_0", insecure},
	tls.VersionTLS11: {"TLS 1.1", "tls_1_1", ok},
	tls.VersionTLS12: {"TLS 1.2", "tls_1_2", good},
}

func parseRawSubject(subject []byte) (pkix.Name, error) {
	name := pkix.Name{}

	var seq pkix.RDNSequence
	_, err := asn1.Unmarshal(subject, &seq)
	if err != nil {
		return name, err
	}

	name.FillFromRDNSequence(&seq)
	return name, nil
}

// Fill in a human readable name, extracted from the slug
func explainCipher(d description) description {
	kexAndCipher := strings.Split(d.Slug, "_WITH_")
	d.Name = fmt.Sprintf("%s key exchange, %s cipher", kexAndCipher[0][len("TLS_"):], kexAndCipher[1])
	return d
}

var cipherSuites = map[uint16]description{
	tls.TLS_RSA_WITH_RC4_128_SHA:                {"", "TLS_RSA_WITH_RC4_128_SHA", insecure},
	tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA:           {"", "TLS_RSA_WITH_3DES_EDE_CBC_SHA", insecure},
	tls.TLS_RSA_WITH_AES_128_CBC_SHA:            {"", "TLS_RSA_WITH_AES_128_CBC_SHA", ok},
	tls.TLS_RSA_WITH_AES_256_CBC_SHA:            {"", "TLS_RSA_WITH_AES_256_CBC_SHA", ok},
	tls.TLS_RSA_WITH_AES_128_CBC_SHA256:         {"", "TLS_RSA_WITH_AES_128_CBC_SHA256", ok},
	tls.TLS_RSA_WITH_AES_128_GCM_SHA256:         {"", "TLS_RSA_WITH_AES_128_GCM_SHA256", ok},
	tls.TLS_RSA_WITH_AES_256_GCM_SHA384:         {"", "TLS_RSA_WITH_AES_256_GCM_SHA384", ok},
	tls.TLS_ECDHE_ECDSA_WITH_RC4_128_SHA:        {"", "TLS_ECDHE_ECDSA_WITH_RC4_128_SHA", insecure},
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA:    {"", "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA", ok},
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA:    {"", "TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA", ok},
	tls.TLS_ECDHE_RSA_WITH_RC4_128_SHA:          {"", "TLS_ECDHE_RSA_WITH_RC4_128_SHA", insecure},
	tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA:     {"", "TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA", insecure},
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA:      {"", "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA", ok},
	tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA:      {"", "TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA", ok},
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256: {"", "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256", ok},
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256:   {"", "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256", ok},
	tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:   {"", "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256", good},
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256: {"", "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256", good},
	tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384:   {"", "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384", good},
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384: {"", "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384", good},
	tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305:    {"", "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305", good},
	tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305:  {"", "TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305", good},
}
<commit_msg>Fix panic in explainCipher for TLS 1.3 ciphers<commit_after>package lib

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"crypto/x509/pkix"
	"encoding/asn1"
	"fmt"
	"strings"
	"text/template"

	"github.com/Masterminds/sprig"
	"github.com/fatih/color"
)

// TLSDescription has the basic information about a TLS connection
type TLSDescription struct {
	Version string `json:"version"`
	Cipher  string `json:"cipher"`
}

// CertificateRequestInfo has the basic information about requested client certificates
type CertificateRequestInfo struct {
	AcceptableCAs    []simplePKIXName `json:"acceptable_issuers,omitempty"`
	SignatureSchemes []string         `json:"signature_schemes,omitempty"`
}

type tlsInfoContext struct {
	Conn *TLSDescription
	CRI  *CertificateRequestInfo
}

var tlsLayout = `** TLS Connection **
Version: {{.Conn.Version}}
Cipher Suite: {{.Conn.Cipher}}
{{- if .CRI}}

{{"Server has requested a client certificate:" | greenify}}
Acceptable issuers:
{{- range .CRI.AcceptableCAs}}
	{{.Name | printShortName}}{{end}}
Supported Signature Schemes:
{{- range .CRI.SignatureSchemes}}
	{{.}}{{end}}
{{- end}}`

func tlscolor(d description) string {
	c, ok := qualityColors[d.Quality]
	if !ok {
		return d.Name
	}
	return c.SprintFunc()(d.Name)
}

// EncodeTLSInfoToText returns a human readable string, suitable for certigo console output.
func EncodeTLSInfoToText(tcs *tls.ConnectionState, cri *tls.CertificateRequestInfo) string {
	version := lookup(tlsVersions, tcs.Version)
	cipher := lookup(cipherSuites, tcs.CipherSuite)
	description := TLSDescription{
		Version: tlscolor(version),
		Cipher:  tlscolor(explainCipher(cipher)),
	}
	tlsInfoContext := tlsInfoContext{
		Conn: &description,
	}
	if cri != nil {
		criDesc, err := EncodeCRIToObject(cri)
		if err == nil {
			tlsInfoContext.CRI = criDesc.(*CertificateRequestInfo)
		}
	}

	funcMap := sprig.TxtFuncMap()
	extras := template.FuncMap{
		"printCommonName": PrintCommonName,
		"printShortName":  PrintShortName,
		"greenify":        greenify,
	}
	for k, v := range extras {
		funcMap[k] = v
	}

	t := template.New("TLS template").Funcs(funcMap)
	t, err := t.Parse(tlsLayout)
	if err != nil {
		// Should never happen
		panic(err)
	}
	var buffer bytes.Buffer
	w := bufio.NewWriter(&buffer)
	err = t.Execute(w, tlsInfoContext)
	if err != nil {
		// Should never happen
		panic(err)
	}
	w.Flush()
	return string(buffer.Bytes())
}

// EncodeTLSToObject returns a JSON-marshallable description of a TLS connection
func EncodeTLSToObject(t *tls.ConnectionState) interface{} {
	version := lookup(tlsVersions, t.Version)
	cipher := lookup(cipherSuites, t.CipherSuite)
	return &TLSDescription{
		version.Slug,
		cipher.Slug,
	}
}

// EncodeCRIToObject returns a JSON-marshallable representation of a CertificateRequestInfo object.
func EncodeCRIToObject(cri *tls.CertificateRequestInfo) (interface{}, error) {
	out := &CertificateRequestInfo{}
	for _, ca := range cri.AcceptableCAs {
		subject, err := parseRawSubject(ca)
		if err != nil {
			return nil, err
		}
		out.AcceptableCAs = append(out.AcceptableCAs, simplePKIXName{subject, nil})
	}
	for _, scheme := range cri.SignatureSchemes {
		desc, ok := signatureSchemeStrings[scheme]
		if !ok {
			desc = fmt.Sprintf("Unknown(0x%x)", scheme)
		}
		out.SignatureSchemes = append(out.SignatureSchemes, desc)
	}
	return out, nil
}

// Just a map lookup with a default
func lookup(descriptions map[uint16]description, what uint16) description {
	v, ok := descriptions[what]
	if !ok {
		unknown := fmt.Sprintf("UNKNOWN_%x", what)
		return description{unknown, unknown, 0}
	}
	return v
}

const (
	insecure = iota
	ok       = iota
	good     = iota
)

type description struct {
	Name    string // a human-friendly string
	Slug    string // a machine-friendly string
	Quality uint8  // insecure, ok, good
}

var qualityColors = map[uint8]*color.Color{
	insecure: red,
	ok:       yellow,
	good:     green,
}

var tlsVersions = map[uint16]description{
	tls.VersionSSL30: {"SSL 3.0", "ssl_3_0", insecure},
	tls.VersionTLS10: {"TLS 1.0", "tls_1_0", insecure},
	tls.VersionTLS11: {"TLS 1.1", "tls_1_1", ok},
	tls.VersionTLS12: {"TLS 1.2", "tls_1_2", good},
	tls.VersionTLS13: {"TLS 1.3", "tls_1_3", good},
}

func parseRawSubject(subject []byte) (pkix.Name, error) {
	name := pkix.Name{}

	var seq pkix.RDNSequence
	_, err := asn1.Unmarshal(subject, &seq)
	if err != nil {
		return name, err
	}

	name.FillFromRDNSequence(&seq)
	return name, nil
}

// Fill in a human readable name, extracted from the slug
func explainCipher(d description) description {
	kexAndCipher := strings.Split(d.Slug, "_WITH_")
	if len(kexAndCipher) < 2 {
		return d
	}
	d.Name = fmt.Sprintf("%s key exchange, %s cipher", kexAndCipher[0][len("TLS_"):], kexAndCipher[1])
	return d
}

var cipherSuites = map[uint16]description{
	tls.TLS_RSA_WITH_RC4_128_SHA:                {"", "TLS_RSA_WITH_RC4_128_SHA", insecure},
	tls.TLS_RSA_WITH_3DES_EDE_CBC_SHA:           {"", "TLS_RSA_WITH_3DES_EDE_CBC_SHA", insecure},
	tls.TLS_RSA_WITH_AES_128_CBC_SHA:            {"", "TLS_RSA_WITH_AES_128_CBC_SHA", ok},
	tls.TLS_RSA_WITH_AES_256_CBC_SHA:            {"", "TLS_RSA_WITH_AES_256_CBC_SHA", ok},
	tls.TLS_RSA_WITH_AES_128_CBC_SHA256:         {"", "TLS_RSA_WITH_AES_128_CBC_SHA256", ok},
	tls.TLS_RSA_WITH_AES_128_GCM_SHA256:         {"", "TLS_RSA_WITH_AES_128_GCM_SHA256", ok},
	tls.TLS_RSA_WITH_AES_256_GCM_SHA384:         {"", "TLS_RSA_WITH_AES_256_GCM_SHA384", ok},
	tls.TLS_ECDHE_ECDSA_WITH_RC4_128_SHA:        {"", "TLS_ECDHE_ECDSA_WITH_RC4_128_SHA", insecure},
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA:    {"", "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA", ok},
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA:    {"", "TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA", ok},
	tls.TLS_ECDHE_RSA_WITH_RC4_128_SHA:          {"", "TLS_ECDHE_RSA_WITH_RC4_128_SHA", insecure},
	tls.TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA:     {"", "TLS_ECDHE_RSA_WITH_3DES_EDE_CBC_SHA", insecure},
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA:      {"", "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA", ok},
	tls.TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA:      {"", "TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA", ok},
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256: {"", "TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA256", ok},
	tls.TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256:   {"", "TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA256", ok},
	tls.TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:   {"", "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256", good},
	tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256: {"", "TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256", good},
	tls.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384:   {"", "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384", good},
	tls.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384: {"", "TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384", good},
	tls.TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305:    {"", "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305", good},
	tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305:  {"", "TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305", good},
}
<|endoftext|>


<commit_before>package web

import(
  "fmt"
  "net/http"
  "net/url"
  "os"
  "os/signal"
  "io"
  "io/ioutil"
	"strings"
	"syscall"
  "crypto/tls"
  "strconv"
  "../arch"
  "../file"
)

var client = &http.Client{}
var nodeBaseAddress = "https://nodejs.org/dist/"
var npmBaseAddress = "https://github.com/npm/cli/archive/"
// var oldNpmBaseAddress = "https://github.com/npm/npm/archive/"

func SetProxy(p string, verifyssl bool){
  if p != "" && p != "none" {
    proxyUrl, _ := url.Parse(p)
    client = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyUrl), TLSClientConfig: &tls.Config{InsecureSkipVerify: verifyssl}}}
  } else {
    client = &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: verifyssl}}}
  }
}

func SetMirrors(node_mirror string, npm_mirror string){
  if node_mirror != "" && node_mirror != "none"{
    nodeBaseAddress = node_mirror;
    if strings.ToLower(nodeBaseAddress[0:4]) != "http" {
      nodeBaseAddress = "http://"+nodeBaseAddress
    }
  }
  if npm_mirror != "" && npm_mirror != "none"{
    npmBaseAddress = npm_mirror;
    if strings.ToLower(npmBaseAddress[0:4]) != "http" {
      npmBaseAddress = "http://"+npmBaseAddress
    }
  }
}

func GetFullNodeUrl(path string) string{
  return nodeBaseAddress+ path;
}

func  GetFullNpmUrl(path string) string{
  return npmBaseAddress + path;
}

func Download(url string, target string, version string) bool {

  output, err := os.Create(target)
  if err != nil {
    fmt.Println("Error while creating", target, "-", err)
  }
  defer output.Close()

  response, err := client.Get(url)
  if err != nil {
    fmt.Println("Error while downloading", url, "-", err)
  }
  defer response.Body.Close()
  c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		fmt.Println("Download interrupted.Rolling back...")
		output.Close()
		response.Body.Close()
		var err error
		if strings.Contains(target, "node") {
			err = os.RemoveAll(os.Getenv("NVM_HOME") + "\\v" + version)
		} else {
			err = os.Remove(target)
		}
		if err != nil {
			fmt.Println("Error while rolling back", err)
		}
		os.Exit(1)
	}()
  _, err = io.Copy(output, response.Body)
  if err != nil {
    fmt.Println("Error while downloading", url, "-", err)
  }
  if response.Status[0:3] != "200" {
    fmt.Println("Download failed. Rolling Back.")
    err := os.Remove(target)
    if err != nil {
      fmt.Println("Rollback failed.",err)
    }
    return false
  }

  return true
}

func GetNodeJS(root string, v string, a string) bool {

  a = arch.Validate(a)

  vpre := ""
  vers := strings.Fields(strings.Replace(v,"."," ",-1))
  main, _ := strconv.ParseInt(vers[0],0,0)

  if a == "32" {
    if main > 0 {
      vpre = "win-x86/"
    } else {
      vpre = ""
    }
  } else if a == "64" {
    if main > 0 {
      vpre = "win-x64/"
    } else {
      vpre = "x64/"
    }
  }

  url := getNodeUrl ( v, vpre );

  if url == "" {
    //No url should mean this version/arch isn't available
    fmt.Println("Node.js v"+v+" " + a + "bit isn't available right now.")
  } else {
   fileName := root+"\\v"+v+"\\node"+a+".exe"

    fmt.Println("Downloading node.js version "+v+" ("+a+"-bit)... ")

    if Download(url,fileName,v) {
      fmt.Printf("Complete\n")
      return true
    } else {
      return false
    }
  }
  return false

}

func GetNpm(root string, v string) bool {
  url := GetFullNpmUrl("v"+v+".zip")
  // temp directory to download the .zip file
  tempDir := root+"\\temp"

  // if the temp directory doesn't exist, create it
  if (!file.Exists(tempDir)) {
    fmt.Println("Creating "+tempDir+"\n")
    err := os.Mkdir(tempDir, os.ModePerm)
    if err != nil {
      fmt.Println(err)
      os.Exit(1)
    }
  }
  fileName := tempDir+"\\"+"npm-v"+v+".zip"

  fmt.Printf("Downloading npm version "+v+"... ")
  if Download(url,fileName,v) {
    fmt.Printf("Complete\n")
    return true
  } else {
    return false
  }
}

func GetRemoteTextFile(url string) string {
  response, httperr := client.Get(url)
  if httperr != nil {
    fmt.Println("\nCould not retrieve "+url+".\n\n")
    fmt.Printf("%s", httperr)
    os.Exit(1)
  } else {
    defer response.Body.Close()
    contents, readerr := ioutil.ReadAll(response.Body)
    if readerr != nil {
      fmt.Printf("%s", readerr)
      os.Exit(1)
    }
    return string(contents)
  }
  os.Exit(1)
  return ""
}

func IsNode64bitAvailable(v string) bool {
  if v == "latest" {
    return true
  }

  // Anything below version 8 doesn't have a 64 bit version
  vers := strings.Fields(strings.Replace(v,"."," ",-1))
  main, _ := strconv.ParseInt(vers[0],0,0)
  minor, _ := strconv.ParseInt(vers[1],0,0)
  if main == 0 && minor < 8 {
    return false
  }

  // TODO: fixme. Assume a 64 bit version exists
  return true
}

func getNodeUrl (v string,  vpre string) string {
  //url := "http://nodejs.org/dist/v"+v+"/" + vpre + "/node.exe"
  url := GetFullNodeUrl("v"+v+"/" + vpre + "/node.exe")
  // Check online to see if a 64 bit version exists
  _, err := client.Head( url )
  if err != nil {
    return ""
  }
  return url;
}
<commit_msg>web.download: return false ASAP if meet the error<commit_after>package web

import(
  "fmt"
  "net/http"
  "net/url"
  "os"
  "os/signal"
  "io"
  "io/ioutil"
	"strings"
	"syscall"
  "crypto/tls"
  "strconv"
  "../arch"
  "../file"
)

var client = &http.Client{}
var nodeBaseAddress = "https://nodejs.org/dist/"
var npmBaseAddress = "https://github.com/npm/cli/archive/"
// var oldNpmBaseAddress = "https://github.com/npm/npm/archive/"

func SetProxy(p string, verifyssl bool){
  if p != "" && p != "none" {
    proxyUrl, _ := url.Parse(p)
    client = &http.Client{Transport: &http.Transport{Proxy: http.ProxyURL(proxyUrl), TLSClientConfig: &tls.Config{InsecureSkipVerify: verifyssl}}}
  } else {
    client = &http.Client{Transport: &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: verifyssl}}}
  }
}

func SetMirrors(node_mirror string, npm_mirror string){
  if node_mirror != "" && node_mirror != "none"{
    nodeBaseAddress = node_mirror;
    if strings.ToLower(nodeBaseAddress[0:4]) != "http" {
      nodeBaseAddress = "http://"+nodeBaseAddress
    }
  }
  if npm_mirror != "" && npm_mirror != "none"{
    npmBaseAddress = npm_mirror;
    if strings.ToLower(npmBaseAddress[0:4]) != "http" {
      npmBaseAddress = "http://"+npmBaseAddress
    }
  }
}

func GetFullNodeUrl(path string) string{
  return nodeBaseAddress+ path;
}

func  GetFullNpmUrl(path string) string{
  return npmBaseAddress + path;
}

func Download(url string, target string, version string) bool {

  output, err := os.Create(target)
  if err != nil {
    fmt.Println("Error while creating", target, "-", err)
    return false
  }
  defer output.Close()

  response, err := client.Get(url)
  if err != nil {
    fmt.Println("Error while downloading", url, "-", err)
    return false
  }
  defer response.Body.Close()
  c := make(chan os.Signal, 2)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		fmt.Println("Download interrupted.Rolling back...")
		output.Close()
		response.Body.Close()
		var err error
		if strings.Contains(target, "node") {
			err = os.RemoveAll(os.Getenv("NVM_HOME") + "\\v" + version)
		} else {
			err = os.Remove(target)
		}
		if err != nil {
			fmt.Println("Error while rolling back", err)
		}
		os.Exit(1)
	}()
  _, err = io.Copy(output, response.Body)
  if err != nil {
    fmt.Println("Error while downloading", url, "-", err)
    return false
  }
  if response.Status[0:3] != "200" {
    fmt.Println("Download failed. Rolling Back.")
    err := os.Remove(target)
    if err != nil {
      fmt.Println("Rollback failed.",err)
      return false
    }
    return false
  }

  return true
}

func GetNodeJS(root string, v string, a string) bool {

  a = arch.Validate(a)

  vpre := ""
  vers := strings.Fields(strings.Replace(v,"."," ",-1))
  main, _ := strconv.ParseInt(vers[0],0,0)

  if a == "32" {
    if main > 0 {
      vpre = "win-x86/"
    } else {
      vpre = ""
    }
  } else if a == "64" {
    if main > 0 {
      vpre = "win-x64/"
    } else {
      vpre = "x64/"
    }
  }

  url := getNodeUrl ( v, vpre );

  if url == "" {
    //No url should mean this version/arch isn't available
    fmt.Println("Node.js v"+v+" " + a + "bit isn't available right now.")
  } else {
   fileName := root+"\\v"+v+"\\node"+a+".exe"

    fmt.Println("Downloading node.js version "+v+" ("+a+"-bit)... ")

    if Download(url,fileName,v) {
      fmt.Printf("Complete\n")
      return true
    } else {
      return false
    }
  }
  return false

}

func GetNpm(root string, v string) bool {
  url := GetFullNpmUrl("v"+v+".zip")
  // temp directory to download the .zip file
  tempDir := root+"\\temp"

  // if the temp directory doesn't exist, create it
  if (!file.Exists(tempDir)) {
    fmt.Println("Creating "+tempDir+"\n")
    err := os.Mkdir(tempDir, os.ModePerm)
    if err != nil {
      fmt.Println(err)
      os.Exit(1)
    }
  }
  fileName := tempDir+"\\"+"npm-v"+v+".zip"

  fmt.Printf("Downloading npm version "+v+"... ")
  if Download(url,fileName,v) {
    fmt.Printf("Complete\n")
    return true
  } else {
    return false
  }
}

func GetRemoteTextFile(url string) string {
  response, httperr := client.Get(url)
  if httperr != nil {
    fmt.Println("\nCould not retrieve "+url+".\n\n")
    fmt.Printf("%s", httperr)
    os.Exit(1)
  } else {
    defer response.Body.Close()
    contents, readerr := ioutil.ReadAll(response.Body)
    if readerr != nil {
      fmt.Printf("%s", readerr)
      os.Exit(1)
    }
    return string(contents)
  }
  os.Exit(1)
  return ""
}

func IsNode64bitAvailable(v string) bool {
  if v == "latest" {
    return true
  }

  // Anything below version 8 doesn't have a 64 bit version
  vers := strings.Fields(strings.Replace(v,"."," ",-1))
  main, _ := strconv.ParseInt(vers[0],0,0)
  minor, _ := strconv.ParseInt(vers[1],0,0)
  if main == 0 && minor < 8 {
    return false
  }

  // TODO: fixme. Assume a 64 bit version exists
  return true
}

func getNodeUrl (v string,  vpre string) string {
  //url := "http://nodejs.org/dist/v"+v+"/" + vpre + "/node.exe"
  url := GetFullNodeUrl("v"+v+"/" + vpre + "/node.exe")
  // Check online to see if a 64 bit version exists
  _, err := client.Head( url )
  if err != nil {
    return ""
  }
  return url;
}
<|endoftext|>


<commit_before><commit_msg>Update Connection Close() to return error<commit_after><|endoftext|>


<commit_before>package elasticsearch

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/maliceio/malice/malice/database"
	"github.com/maliceio/malice/malice/persist"
	"github.com/maliceio/malice/utils"
	elastic "gopkg.in/olivere/elastic.v3"
)

// PluginResults a malice plugin results object
type PluginResults struct {
	ID   string `json:"id" gorethink:"id,omitempty"`
	Data map[string]interface{}
}

// ElasticAddr ElasticSearch address to user for connections
var ElasticAddr string

// InitElasticSearch initalizes ElasticSearch for use with malice
func InitElasticSearch() error {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	exists, err := client.IndexExists("malice").Do()
	utils.Assert(err)

	if !exists {
		// Index does not exist yet.
		createIndex, err := client.CreateIndex("malice").BodyString(mapping).Do()
		utils.Assert(err)
		if !createIndex.Acknowledged {
			// Not acknowledged
			log.Error("Couldn't create Index.")
		} else {
			log.Info("Created Index: ", "malice")
		}
	} else {
		log.Info("Index malice already exists.")
	}

	return err
}

// TestConnection tests the ElasticSearch connection
func TestConnection(addr string) error {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	// connect to ElasticSearch where --link elastic was using via malice in Docker
	log.Debugf("Attempting to connect to: %s", ElasticAddr)
	_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

	// Ping the Elasticsearch server to get e.g. the version number
	// info, code, err := client.Ping(ElasticAddr).Do()
	// utils.Assert(err)
	// fmt.Printf("Elasticsearch returned with code %d and version %s", code, info.Version.Number)

	if err != nil {
		// connect to ElasticSearch via malice in Docker
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", addr))
		log.Debugf("Attempting to connect to: %s", ElasticAddr)
		_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

		if err != nil {
			// connect to ElasticSearch using Docker for Mac
			ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "localhost"))
			log.Debugf("Attempting to connect to: %s", ElasticAddr)
			_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
			return err
		}
		return err
	}
	return err
}

// WriteFileToDatabase inserts sample into Database
func WriteFileToDatabase(sample persist.File) elastic.IndexResponse {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file":      sample,
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	update, err := client.Update().Index("malice").Type("samples").Id(newScan.Id).
		Doc(map[string]interface{}{
			"plugins": map[string]interface{}{
				"intel": map[string]interface{}{
					"nsrl": "UPDATED",
				},
			},
		}).
		Do()
	utils.Assert(err)
	log.Debugf("New version of sample %q is now %d\n", update.Id, update.Version)

	return *newScan
}

// WriteHashToDatabase inserts sample into Database
func WriteHashToDatabase(hash string) elastic.IndexResponse {

	hashType, err := util.GetHashType(hash)
	utils.Assert(err)

	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file": map[string]interface{}{
			hashType: hash,
		},
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	return *newScan
}

// WritePluginResultsToDatabase upserts plugin results into Database
func WritePluginResultsToDatabase(results PluginResults) {

	// scanID := utils.Getopt("MALICE_SCANID", "")
	ElasticAddr := fmt.Sprintf("%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	getSample, err := client.Get().
		Index("malice").
		Type("samples").
		Id(results.ID).
		Do()

	fmt.Println(getSample)
	fmt.Println(err)
	if err != nil {

	}

	if getSample.Found {
		fmt.Printf("Got document %s in version %d from index %s, type %s\n", getSample.Id, getSample.Version, getSample.Index, getSample.Type)
		updateScan := map[string]interface{}{
			"scan_date": time.Now().Format(time.RFC3339Nano),
			"plugins": map[string]interface{}{
				category: map[string]interface{}{
					name: results.Data,
				},
			},
		}
		update, err := client.Update().Index("malice").Type("samples").Id(getSample.Id).
			Doc(updateScan).
			Do()
		utils.Assert(err)

		log.Debugf("New version of sample %q is now %d\n", update.Id, update.Version)
		// return *update

	} else {

		scan := map[string]interface{}{
			// "id":      sample.SHA256,
			// "file":      sample,
			"plugins": map[string]interface{}{
				category: map[string]interface{}{
					name: results.Data,
				},
			},
			"scan_date": time.Now().Format(time.RFC3339Nano),
		}

		newScan, err := client.Index().
			Index("malice").
			Type("samples").
			OpType("create").
			// Id("1").
			BodyJson(scan).
			Do()
		utils.Assert(err)

		log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)
		// return *newScan
	}
}
<commit_msg>fix database/elasticsearch<commit_after>package elasticsearch

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"

	"github.com/maliceio/malice/malice/database"
	"github.com/maliceio/malice/malice/persist"
	"github.com/maliceio/malice/utils"
	elastic "gopkg.in/olivere/elastic.v3"
)

// PluginResults a malice plugin results object
type PluginResults struct {
	ID       string `json:"id"`
	Name     string
	Category string
	Data     map[string]interface{}
}

// ElasticAddr ElasticSearch address to user for connections
var ElasticAddr string

// InitElasticSearch initalizes ElasticSearch for use with malice
func InitElasticSearch() error {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	exists, err := client.IndexExists("malice").Do()
	utils.Assert(err)

	if !exists {
		// Index does not exist yet.
		createIndex, err := client.CreateIndex("malice").BodyString(mapping).Do()
		utils.Assert(err)
		if !createIndex.Acknowledged {
			// Not acknowledged
			log.Error("Couldn't create Index.")
		} else {
			log.Info("Created Index: ", "malice")
		}
	} else {
		log.Info("Index malice already exists.")
	}

	return err
}

// TestConnection tests the ElasticSearch connection
func TestConnection(addr string) error {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	// connect to ElasticSearch where --link elastic was using via malice in Docker
	log.Debugf("Attempting to connect to: %s", ElasticAddr)
	_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

	// Ping the Elasticsearch server to get e.g. the version number
	// info, code, err := client.Ping(ElasticAddr).Do()
	// utils.Assert(err)
	// fmt.Printf("Elasticsearch returned with code %d and version %s", code, info.Version.Number)

	if err != nil {
		// connect to ElasticSearch via malice in Docker
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", addr))
		log.Debugf("Attempting to connect to: %s", ElasticAddr)
		_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))

		if err != nil {
			// connect to ElasticSearch using Docker for Mac
			ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "localhost"))
			log.Debugf("Attempting to connect to: %s", ElasticAddr)
			_, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
			return err
		}
		return err
	}
	return err
}

// WriteFileToDatabase inserts sample into Database
func WriteFileToDatabase(sample persist.File) elastic.IndexResponse {

	if ElasticAddr == "" {
		ElasticAddr = fmt.Sprintf("http://%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	}

	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file":      sample,
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	update, err := client.Update().Index("malice").Type("samples").Id(newScan.Id).
		Doc(map[string]interface{}{
			"plugins": map[string]interface{}{
				"intel": map[string]interface{}{
					"nsrl": "UPDATED",
				},
			},
		}).
		Do()
	utils.Assert(err)
	log.Debugf("New version of sample %q is now %d\n", update.Id, update.Version)

	return *newScan
}

// WriteHashToDatabase inserts sample into Database
func WriteHashToDatabase(hash string) elastic.IndexResponse {

	hashType, err := utils.GetHashType(hash)
	utils.Assert(err)

	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	scan := map[string]interface{}{
		// "id":      sample.SHA256,
		"file": map[string]interface{}{
			hashType: hash,
		},
		"plugins":   database.GetPluginsByCategory(),
		"scan_date": time.Now().Format(time.RFC3339Nano),
	}

	newScan, err := client.Index().
		Index("malice").
		Type("samples").
		OpType("create").
		// Id("1").
		BodyJson(scan).
		Do()
	utils.Assert(err)
	log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)

	return *newScan
}

// WritePluginResultsToDatabase upserts plugin results into Database
func WritePluginResultsToDatabase(results PluginResults) {

	// scanID := utils.Getopt("MALICE_SCANID", "")
	ElasticAddr := fmt.Sprintf("%s:9200", utils.Getopt("MALICE_ELASTICSEARCH", "elastic"))
	client, err := elastic.NewSimpleClient(elastic.SetURL(ElasticAddr))
	utils.Assert(err)

	getSample, err := client.Get().
		Index("malice").
		Type("samples").
		Id(results.ID).
		Do()

	fmt.Println(getSample)
	fmt.Println(err)
	if err != nil {

	}

	if getSample.Found {
		fmt.Printf("Got document %s in version %d from index %s, type %s\n", getSample.Id, getSample.Version, getSample.Index, getSample.Type)
		updateScan := map[string]interface{}{
			"scan_date": time.Now().Format(time.RFC3339Nano),
			"plugins": map[string]interface{}{
				results.Category: map[string]interface{}{
					results.Name: results.Data,
				},
			},
		}
		update, err := client.Update().Index("malice").Type("samples").Id(getSample.Id).
			Doc(updateScan).
			Do()
		utils.Assert(err)

		log.Debugf("New version of sample %q is now %d\n", update.Id, update.Version)
		// return *update

	} else {

		scan := map[string]interface{}{
			// "id":      sample.SHA256,
			// "file":      sample,
			"plugins": map[string]interface{}{
				results.Category: map[string]interface{}{
					results.Name: results.Data,
				},
			},
			"scan_date": time.Now().Format(time.RFC3339Nano),
		}

		newScan, err := client.Index().
			Index("malice").
			Type("samples").
			OpType("create").
			// Id("1").
			BodyJson(scan).
			Do()
		utils.Assert(err)

		log.Debugf("Indexed sample %s to index %s, type %s\n", newScan.Id, newScan.Index, newScan.Type)
		// return *newScan
	}
}
<|endoftext|>


<commit_before>package fptp_test

import (
	"testing"
	"time"

	"io"

	"context"

	"github.com/jabley/fptp"
)

func TestUnusedCounterHasZeroUnclosed(t *testing.T) {
	counter := &CounterCloser{}
	if counter.Unclosed() != 0 {
		t.Fatalf("Unused counter should have zero unclosed")
	}
}

func TestUnclosedSearcherIsCounted(t *testing.T) {
	counter := &CounterCloser{}
	searcher := NewCountingSearcher(counter)

	closer, _ := searcher.Search(fptp.NewSearchRequest())
	defer closer.Close()
	if counter.Unclosed() != 1 {
		t.Fatalf("Unclosed Closer should have been counted")
	}
}

func TestCompositeWillTimeout(t *testing.T) {
	counter := &CounterCloser{}
	searcher := &delayedSearcher{
		delegate: NewCountingSearcher(counter),
		delay:    20 * time.Millisecond,
	}

	composite := fptp.NewCompositeSearcher([]fptp.Searcher{searcher}, 10*time.Millisecond)

	result, err := composite.Search(fptp.NewSearchRequest())

	if result != nil {
		defer result.Close()
		t.Fatalf("Expected nil result, but got %v", result)
	}

	if err == nil {
		t.Fatalf("Expected to get a timeout error but got nil")
	}

	if err != context.DeadlineExceeded {
		t.Fatalf("Got error %v", err)
	}
}

func TestWillWaitUntilSuccessIsReturned(t *testing.T) {
	counter := &CounterCloser{}
	searcher := &delayedSearcher{
		delegate: NewCountingSearcher(counter),
		delay:    10 * time.Millisecond,
	}

	// Searchers such that the erroring one should complete faster than the other one
	searchers := []fptp.Searcher{
		&erroringSearcher{
			err: io.EOF,
		},
		searcher,
	}
	composite := fptp.NewCompositeSearcher(searchers, 5*time.Second)

	result, err := composite.Search(fptp.NewSearchRequest())

	if err != nil {
		t.Fatalf("Expected that err should be nil, but got %v", err)
	}

	if result == nil {
		t.Fatalf("Expected non-nil result")
	}

	result.Close()
}

func TestWillTryAllSearchersAndReturnTheFailure(t *testing.T) {
	var searchers []fptp.Searcher
	for i := 0; i < 20; i++ {
		searchers = append(searchers, &erroringSearcher{err: io.EOF})
	}

	composite := fptp.NewCompositeSearcher(searchers, 5*time.Second)

	result, err := composite.Search(fptp.NewSearchRequest())

	if result != nil {
		defer result.Close()
		t.Fatalf("Expected nil result, but got %v", result)
	}

	if err == nil {
		t.Fatalf("Expected to get an error but got nil")
	}

	unusedSearchers := 0
	for i := 0; i < len(searchers); i++ {
		es := searchers[i].(*erroringSearcher)
		if !es.used {
			unusedSearchers++
		}
	}

	if unusedSearchers != 0 {
		t.Fatalf("Expected that all of the searchers had been used, but %d were not", unusedSearchers)
	}
}

func TestClosersAreNotLeaked(t *testing.T) {
	counter := &CounterCloser{}
	var searchers []fptp.Searcher
	for i := 0; i < 20; i++ {
		searchers = append(searchers, NewCountingSearcher(counter))
	}

	composite := fptp.NewCompositeSearcher(searchers, 5*time.Second)

	for i := 0; i < 100000; i++ {
		winner, _ := composite.Search(fptp.NewSearchRequest())
		winner.Close()
		assertClosersAllClosed(t, counter)
	}
}

func assertClosersAllClosed(t *testing.T, counter *CounterCloser) {
	backoff := 1 * time.Microsecond
	for {
		unclosed := counter.Unclosed()
		if unclosed == 0 {
			break
		}
		// fail if we've exceeded our time budgete - we don't want to spin forever
		if backoff > 10*time.Millisecond {
			t.Fatalf("Expected all closers to have been closed, but had %d still open after backing off to %v", unclosed, backoff)
		}

		// Give it a little while for all of the laggards to be closed
		time.Sleep(backoff)
		backoff = backoff * 2
	}
}

type delayedSearcher struct {
	delay    time.Duration
	delegate fptp.Searcher
}

func (ds *delayedSearcher) Search(req *fptp.SearchRequest) (io.Closer, error) {
	time.Sleep(ds.delay)
	return ds.delegate.Search(req)
}

type erroringSearcher struct {
	used bool
	err  error
}

func (es *erroringSearcher) Search(req *fptp.SearchRequest) (io.Closer, error) {
	es.used = true
	return nil, es.err
}
<commit_msg>Nicer testing API<commit_after>package fptp_test

import (
	"fmt"
	"reflect"
	"testing"
	"time"

	"io"

	"context"

	"github.com/jabley/fptp"
)

func TestUnusedCounterHasZeroUnclosed(t *testing.T) {
	counter := &CounterCloser{}
	assertEqual(t, 0, counter.Unclosed(), "Unused counter should have zero unclosed")
}

func TestUnclosedSearcherIsCounted(t *testing.T) {
	counter := &CounterCloser{}
	searcher := NewCountingSearcher(counter)

	closer, _ := searcher.Search(fptp.NewSearchRequest())
	defer closer.Close()
	assertEqual(t, 1, counter.Unclosed(), "Unclosed Closer should have been counted")
}

func TestCompositeWillTimeout(t *testing.T) {
	counter := &CounterCloser{}
	searcher := &delayedSearcher{
		delegate: NewCountingSearcher(counter),
		delay:    20 * time.Millisecond,
	}

	composite := fptp.NewCompositeSearcher([]fptp.Searcher{searcher}, 10*time.Millisecond)

	result, err := composite.Search(fptp.NewSearchRequest())

	if result != nil {
		defer result.Close()
		assertEqual(t, nil, result, "Expected result to be nil")
	}

	assertEqual(t, context.DeadlineExceeded, err, "Expected to get a timeout error")
}

func TestWillWaitUntilSuccessIsReturned(t *testing.T) {
	counter := &CounterCloser{}
	searcher := &delayedSearcher{
		delegate: NewCountingSearcher(counter),
		delay:    10 * time.Millisecond,
	}

	// Searchers such that the erroring one should complete faster than the other one
	searchers := []fptp.Searcher{
		&erroringSearcher{
			err: io.EOF,
		},
		searcher,
	}
	composite := fptp.NewCompositeSearcher(searchers, 5*time.Second)

	result, err := composite.Search(fptp.NewSearchRequest())

	if result != nil {
		defer result.Close()
	}
	assertEqual(t, nil, err, "Expected err to be nil")
}

func TestWillTryAllSearchersAndReturnTheFailure(t *testing.T) {
	var searchers []fptp.Searcher
	for i := 0; i < 20; i++ {
		searchers = append(searchers, &erroringSearcher{err: io.EOF})
	}

	composite := fptp.NewCompositeSearcher(searchers, 5*time.Second)

	result, err := composite.Search(fptp.NewSearchRequest())

	if result != nil {
		defer result.Close()
		assertEqual(t, nil, result, "Expected result to be nil")
	}

	if err == nil {
		t.Fatalf("Expected to get an error but got nil")
	}

	unusedSearchers := 0
	for i := 0; i < len(searchers); i++ {
		es := searchers[i].(*erroringSearcher)
		if !es.used {
			unusedSearchers++
		}
	}

	assertEqual(t, 0, unusedSearchers, "Expected that all of the searchers had been used, but %d were not", unusedSearchers)
}

func TestClosersAreNotLeaked(t *testing.T) {
	counter := &CounterCloser{}
	var searchers []fptp.Searcher
	for i := 0; i < 20; i++ {
		searchers = append(searchers, NewCountingSearcher(counter))
	}

	composite := fptp.NewCompositeSearcher(searchers, 5*time.Second)

	for i := 0; i < 1000000; i++ {
		winner, _ := composite.Search(fptp.NewSearchRequest())
		winner.Close()
		assertClosersAllClosed(t, counter)
	}
}

func assertEqual(t *testing.T, expected interface{}, actual interface{}, msgAndArgs ...interface{}) {
	if expected == nil || actual == nil {
		if actual == expected {
			return
		}
		fail(t, expected, actual, msgAndArgs...)
	}

	if !reflect.DeepEqual(expected, actual) {
		fail(t, expected, actual, msgAndArgs...)
	}
}

func fail(t *testing.T, expected, actual interface{}, msgAndArgs ...interface{}) {
	t.Fatalf("\nExpected:\t%#v\nActual:\t\t%#v\n%s", expected, actual, messageFromMsgAndArgs(msgAndArgs...))
}

func messageFromMsgAndArgs(msgAndArgs ...interface{}) string {
	if len(msgAndArgs) == 0 || msgAndArgs == nil {
		return ""
	}
	if len(msgAndArgs) == 1 {
		return msgAndArgs[0].(string)
	}
	if len(msgAndArgs) > 1 {
		return fmt.Sprintf(msgAndArgs[0].(string), msgAndArgs[1:]...)
	}
	return ""
}

func assertClosersAllClosed(t *testing.T, counter *CounterCloser) {
	backoff := 1 * time.Nanosecond
	for {
		unclosed := counter.Unclosed()
		if unclosed == 0 {
			break
		}
		// fail if we've exceeded our time budgete - we don't want to spin forever
		if backoff > 10*time.Millisecond {
			assertEqual(t, 0, unclosed, "Expected all closers to have been closed, but had %d still open after backing off to %v", unclosed, backoff)
		}

		// Give it a little while for all of the laggards to be closed
		time.Sleep(backoff)
		backoff = backoff * 2
	}
}

type delayedSearcher struct {
	delay    time.Duration
	delegate fptp.Searcher
}

func (ds *delayedSearcher) Search(req *fptp.SearchRequest) (io.Closer, error) {
	time.Sleep(ds.delay)
	return ds.delegate.Search(req)
}

type erroringSearcher struct {
	used bool
	err  error
}

func (es *erroringSearcher) Search(req *fptp.SearchRequest) (io.Closer, error) {
	es.used = true
	return nil, es.err
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package command

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
	"k8s.io/client-go/util/jsonpath"
)

var (
	outputOpt string
	re        = regexp.MustCompile(`^jsonpath\=(.*)`)
)

// OutputJSON returns true if the JSON output option was specified
func OutputJSON() bool {
	return len(outputOpt) > 0
}

//AddJSONOutput adds the -o|--output option to any cmd to export to json
func AddJSONOutput(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&outputOpt, "output", "o", "", "json| jsonpath='{}'")
}

//ForceJSON sets output mode to JSON (for unit tests)
func ForceJSON() {
	outputOpt = "json"
}

//PrintOutput receives an interface and dump the data using the --output flag.
//ATM only json or jsonpath. In the future yaml
func PrintOutput(data interface{}) error {
	return PrintOutputWithType(data, outputOpt)
}

//PrintOutputWithPatch merges data with patch and dump the data using the --output flag.
func PrintOutputWithPatch(data interface{}, patch interface{}) error {
	mergedInterface, err := mergeInterfaces(data, patch)
	if err != nil {
		return fmt.Errorf("Unable to merge Interfaces:%v", err)
	}
	return PrintOutputWithType(mergedInterface, outputOpt)
}

func mergeInterfaces(data, patch interface{}) (interface{}, error) {
	var i1, i2 interface{}

	data1, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}
	data2, err := json.Marshal(patch)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(data1, &i1)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(data2, &i2)
	if err != nil {
		return nil, err
	}
	return recursiveMerge(i1, i2), nil
}

func recursiveMerge(i1, i2 interface{}) interface{} {
	switch i1 := i1.(type) {
	case map[string]interface{}:
		i2, ok := i2.(map[string]interface{})
		if !ok {
			return i1
		}
		for k, v2 := range i2 {
			if v1, ok := i1[k]; ok {
				i1[k] = recursiveMerge(v1, v2)
			} else {
				i1[k] = v2
			}
		}
	case nil:
		i2, ok := i2.(map[string]interface{})
		if ok {
			return i2
		}
	}
	return i1
}

//PrintOutputWithType receives an interface and dump the data using the --output flag.
//ATM only json or jsonpath. In the future yaml
func PrintOutputWithType(data interface{}, outputType string) error {
	if outputType == "json" {
		return dumpJSON(data, "")
	}

	if outputType == "yaml" {
		return dumpYAML(data)
	}

	if re.MatchString(outputType) {
		return dumpJSON(data, re.ReplaceAllString(outputType, "$1"))
	}

	return fmt.Errorf("couldn't found output printer")
}

// DumpJSONToString dumps the contents of data into a string. If jsonpath is
// non-empty, will attempt to do jsonpath filtering using said string. Returns a
// string containing the JSON in data, or an error if any JSON marshaling,
// parsing operations fail.
func DumpJSONToString(data interface{}, jsonPath string) (string, error) {
	if len(jsonPath) == 0 {
		result, err := json.MarshalIndent(data, "", "  ")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't marshal to json: '%s'\n", err)
			return "", err
		}
		fmt.Println(string(result))
		return "", nil
	}

	parser := jsonpath.New("").AllowMissingKeys(true)
	if err := parser.Parse(jsonPath); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return "", err
	}

	var sb strings.Builder
	if err := parser.Execute(&sb, data); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return "", err

	}
	return sb.String(), nil
}

// dumpJSON dumps the data variable to the stdout as json.
// If something fails, it returns an error
// If jsonPath is passed, it runs the json query over data var.
func dumpJSON(data interface{}, jsonPath string) error {
	jsonStr, err := DumpJSONToString(data, jsonPath)
	if err != nil {
		return err
	}
	fmt.Println(jsonStr)
	return nil
}

// dumpYAML dumps the data variable to the stdout as yaml.
// If something fails, it returns an error
func dumpYAML(data interface{}) error {
	result, err := yaml.Marshal(data)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't marshal to yaml: '%s'\n", err)
		return err
	}
	fmt.Println(string(result))
	return nil
}
<commit_msg>fix(command): typo in error message<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package command

import (
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"strings"

	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
	"k8s.io/client-go/util/jsonpath"
)

var (
	outputOpt string
	re        = regexp.MustCompile(`^jsonpath\=(.*)`)
)

// OutputJSON returns true if the JSON output option was specified
func OutputJSON() bool {
	return len(outputOpt) > 0
}

//AddJSONOutput adds the -o|--output option to any cmd to export to json
func AddJSONOutput(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&outputOpt, "output", "o", "", "json| jsonpath='{}'")
}

//ForceJSON sets output mode to JSON (for unit tests)
func ForceJSON() {
	outputOpt = "json"
}

//PrintOutput receives an interface and dump the data using the --output flag.
//ATM only json or jsonpath. In the future yaml
func PrintOutput(data interface{}) error {
	return PrintOutputWithType(data, outputOpt)
}

//PrintOutputWithPatch merges data with patch and dump the data using the --output flag.
func PrintOutputWithPatch(data interface{}, patch interface{}) error {
	mergedInterface, err := mergeInterfaces(data, patch)
	if err != nil {
		return fmt.Errorf("Unable to merge Interfaces:%v", err)
	}
	return PrintOutputWithType(mergedInterface, outputOpt)
}

func mergeInterfaces(data, patch interface{}) (interface{}, error) {
	var i1, i2 interface{}

	data1, err := json.Marshal(data)
	if err != nil {
		return nil, err
	}
	data2, err := json.Marshal(patch)
	if err != nil {
		return nil, err
	}
	err = json.Unmarshal(data1, &i1)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(data2, &i2)
	if err != nil {
		return nil, err
	}
	return recursiveMerge(i1, i2), nil
}

func recursiveMerge(i1, i2 interface{}) interface{} {
	switch i1 := i1.(type) {
	case map[string]interface{}:
		i2, ok := i2.(map[string]interface{})
		if !ok {
			return i1
		}
		for k, v2 := range i2 {
			if v1, ok := i1[k]; ok {
				i1[k] = recursiveMerge(v1, v2)
			} else {
				i1[k] = v2
			}
		}
	case nil:
		i2, ok := i2.(map[string]interface{})
		if ok {
			return i2
		}
	}
	return i1
}

//PrintOutputWithType receives an interface and dump the data using the --output flag.
//ATM only json or jsonpath. In the future yaml
func PrintOutputWithType(data interface{}, outputType string) error {
	if outputType == "json" {
		return dumpJSON(data, "")
	}

	if outputType == "yaml" {
		return dumpYAML(data)
	}

	if re.MatchString(outputType) {
		return dumpJSON(data, re.ReplaceAllString(outputType, "$1"))
	}

	return fmt.Errorf("couldn't find output printer")
}

// DumpJSONToString dumps the contents of data into a string. If jsonpath is
// non-empty, will attempt to do jsonpath filtering using said string. Returns a
// string containing the JSON in data, or an error if any JSON marshaling,
// parsing operations fail.
func DumpJSONToString(data interface{}, jsonPath string) (string, error) {
	if len(jsonPath) == 0 {
		result, err := json.MarshalIndent(data, "", "  ")
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't marshal to json: '%s'\n", err)
			return "", err
		}
		fmt.Println(string(result))
		return "", nil
	}

	parser := jsonpath.New("").AllowMissingKeys(true)
	if err := parser.Parse(jsonPath); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return "", err
	}

	var sb strings.Builder
	if err := parser.Execute(&sb, data); err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't parse jsonpath expression: '%s'\n", err)
		return "", err

	}
	return sb.String(), nil
}

// dumpJSON dumps the data variable to the stdout as json.
// If something fails, it returns an error
// If jsonPath is passed, it runs the json query over data var.
func dumpJSON(data interface{}, jsonPath string) error {
	jsonStr, err := DumpJSONToString(data, jsonPath)
	if err != nil {
		return err
	}
	fmt.Println(jsonStr)
	return nil
}

// dumpYAML dumps the data variable to the stdout as yaml.
// If something fails, it returns an error
func dumpYAML(data interface{}) error {
	result, err := yaml.Marshal(data)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Couldn't marshal to yaml: '%s'\n", err)
		return err
	}
	fmt.Println(string(result))
	return nil
}
<|endoftext|>