<commit_before>package manta

import (
	"bytes"
	"io/ioutil"
	"os"

	"github.com/dotabuff/manta/dota"
	"github.com/golang/snappy/snappy"
)

// The first 8 bytes of a replay for Source 1 and Source 2
var magicSource1 = []byte{'P', 'U', 'F', 'D', 'E', 'M', 'S', '\000'}
var magicSource2 = []byte{'P', 'B', 'D', 'E', 'M', 'S', '2', '\000'}

// A replay parser capable of parsing Source 2 replays
type Parser struct {
	// Callbacks provide a mechanism for receiving notification
	// when a specific message has been received and decoded.
	Callbacks *Callbacks

	// GameEvents provides a mechanism for receiving notification
	// when a specific game event message has been received and decoded.
	GameEvents *GameEvents

	// Contains the game tick associated with the last message processed.
	Tick uint32

	hasClassInfo   bool
	classInfo      map[int32]string
	classIdSize    int
	classBaseline  map[int32]map[string]interface{}
	packetEntities map[int32]*packetEntity
	sendTables     *sendTables
	stringTables   *stringTables

	reader     *reader
	isStopping bool
}

// Create a new Parser from a file on disk.
func NewParserFromFile(path string) (*Parser, error) {
	fd, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	buf, err := ioutil.ReadAll(fd)
	if err != nil {
		return nil, err
	}

	return NewParser(buf)
}

// Create a new parser from a byte slice.
func NewParser(buf []byte) (*Parser, error) {
	// Create a new parser with an internal reader for the given buffer.
	parser := &Parser{
		Callbacks:  &Callbacks{},
		GameEvents: &GameEvents{},

		Tick: 0,

		reader:     newReader(buf),
		isStopping: false,

		classInfo:      make(map[int32]string),
		classBaseline:  make(map[int32]map[string]interface{}),
		packetEntities: make(map[int32]*packetEntity),
		stringTables:   newStringTables(),
	}

	// Parse out the header, ensuring that it's valid.
	if magic := parser.reader.readBytes(8); !bytes.Equal(magic, magicSource2) {
		return nil, _errorf("unexpected magic: expected %s, got %s", magicSource2, magic)
	}

	// Skip the next 8 bytes, which appear to be two int32s related to the size
	// of the demo file. We may need them in the future, but not so far.
	parser.reader.seekBytes(8)

	// Internal handlers
	parser.Callbacks.OnCDemoPacket(parser.onCDemoPacket)
	parser.Callbacks.OnCDemoSignonPacket(parser.onCDemoPacket)
	parser.Callbacks.OnCDemoFullPacket(parser.onCDemoFullPacket)
	parser.Callbacks.OnCDemoClassInfo(parser.onCDemoClassInfo)
	parser.Callbacks.OnCDemoSendTables(parser.onCDemoSendTables)
	parser.Callbacks.OnCSVCMsg_CreateStringTable(parser.onCSVCMsg_CreateStringTable)
	parser.Callbacks.OnCSVCMsg_PacketEntities(parser.onCSVCMsg_PacketEntities)
	parser.Callbacks.OnCSVCMsg_SendTable(parser.onCSVCMsg_SendTable)
	parser.Callbacks.OnCSVCMsg_UpdateStringTable(parser.onCSVCMsg_UpdateStringTable)
	parser.Callbacks.OnCSVCMsg_ServerInfo(parser.onCSVCMsg_ServerInfo)
	parser.Callbacks.OnCMsgSource1LegacyGameEvent(parser.GameEvents.onCMsgSource1LegacyGameEvent)

	// Panic if we see any of these
	parser.Callbacks.OnCSVCMsg_GameEvent(func(m *dota.CSVCMsg_GameEvent) error {
		_panicf("unexpected: saw a CSVCMsg_GameEvent")
		return nil
	})

	// Maintains the value of parser.Tick
	parser.Callbacks.OnCNETMsg_Tick(func(m *dota.CNETMsg_Tick) error {
		parser.Tick = m.GetTick()
		return nil
	})

	// Stops parsing when we reach the end of the replay.
	parser.Callbacks.OnCDemoStop(func(m *dota.CDemoStop) error {
		parser.Stop()
		return nil
	})

	return parser, nil
}

// Start parsing the replay. Will stop processing new events after Stop() is called.
func (p *Parser) Start() error {
	var msg *outerMessage
	var err error

	// Loop through all outer messages until we're signaled to stop. Stopping
	// happens when either the OnCDemoStop message is encountered or
	// parser.Stop() is called programatically.
	for !p.isStopping {
		// Read the next outer message.
		if msg, err = p.readOuterMessage(); err != nil {
			return err
		}

		// Invoke callbacks for the given message type.
		if err = p.CallByDemoType(msg.typeId, msg.data); err != nil {
			return err
		}
	}

	return nil
}

// Stop parsing the replay, causing the parser to stop processing new events.
func (p *Parser) Stop() {
	p.isStopping = true
}

// Describes a demo message parsed from the replay.
type outerMessage struct {
	tick   uint32
	typeId int32
	data   []byte
}

// Read the next outer message from the buffer.
func (p *Parser) readOuterMessage() (*outerMessage, error) {
	// Read a command header, which includes both the message type
	// well as a flag to determine whether or not whether or not the
	// message is compressed with snappy.
	command := dota.EDemoCommands(p.reader.readVarUint32())

	// Extract the type and compressed flag out of the command
	msgType := int32(command & ^dota.EDemoCommands_DEM_IsCompressed)
	msgCompressed := (command & dota.EDemoCommands_DEM_IsCompressed) == dota.EDemoCommands_DEM_IsCompressed

	// Read the tick that the message corresponds with.
	tick := p.reader.readVarUint32()

	// Read the size and following buffer.
	size := int(p.reader.readVarUint32())
	buf := p.reader.readBytes(size)

	// If the buffer is compressed, decompress it with snappy.
	if msgCompressed {
		var err error
		if buf, err = snappy.Decode(nil, buf); err != nil {
			return nil, err
		}
	}

	// Return the message
	msg := &outerMessage{
		tick:   tick,
		typeId: msgType,
		data:   buf,
	}
	return msg, nil
}
<commit_msg>Separate demo message tick and net tick<commit_after>package manta

import (
	"bytes"
	"io/ioutil"
	"os"

	"github.com/dotabuff/manta/dota"
	"github.com/golang/snappy/snappy"
)

// The first 8 bytes of a replay for Source 1 and Source 2
var magicSource1 = []byte{'P', 'U', 'F', 'D', 'E', 'M', 'S', '\000'}
var magicSource2 = []byte{'P', 'B', 'D', 'E', 'M', 'S', '2', '\000'}

// A replay parser capable of parsing Source 2 replays
type Parser struct {
	// Callbacks provide a mechanism for receiving notification
	// when a specific message has been received and decoded.
	Callbacks *Callbacks

	// GameEvents provides a mechanism for receiving notification
	// when a specific game event message has been received and decoded.
	GameEvents *GameEvents

	// Contains the game tick associated with the last message processed.
	Tick uint32

	// Contains the net tick associated with the last net message processed.
	NetTick uint32

	hasClassInfo   bool
	classInfo      map[int32]string
	classIdSize    int
	classBaseline  map[int32]map[string]interface{}
	packetEntities map[int32]*packetEntity
	sendTables     *sendTables
	stringTables   *stringTables

	reader     *reader
	isStopping bool
}

// Create a new Parser from a file on disk.
func NewParserFromFile(path string) (*Parser, error) {
	fd, err := os.Open(path)
	if err != nil {
		return nil, err
	}

	buf, err := ioutil.ReadAll(fd)
	if err != nil {
		return nil, err
	}

	return NewParser(buf)
}

// Create a new parser from a byte slice.
func NewParser(buf []byte) (*Parser, error) {
	// Create a new parser with an internal reader for the given buffer.
	parser := &Parser{
		Callbacks:  &Callbacks{},
		GameEvents: &GameEvents{},

		Tick:    0,
		NetTick: 0,

		reader:     newReader(buf),
		isStopping: false,

		classInfo:      make(map[int32]string),
		classBaseline:  make(map[int32]map[string]interface{}),
		packetEntities: make(map[int32]*packetEntity),
		stringTables:   newStringTables(),
	}

	// Parse out the header, ensuring that it's valid.
	if magic := parser.reader.readBytes(8); !bytes.Equal(magic, magicSource2) {
		return nil, _errorf("unexpected magic: expected %s, got %s", magicSource2, magic)
	}

	// Skip the next 8 bytes, which appear to be two int32s related to the size
	// of the demo file. We may need them in the future, but not so far.
	parser.reader.seekBytes(8)

	// Internal handlers
	parser.Callbacks.OnCDemoPacket(parser.onCDemoPacket)
	parser.Callbacks.OnCDemoSignonPacket(parser.onCDemoPacket)
	parser.Callbacks.OnCDemoFullPacket(parser.onCDemoFullPacket)
	parser.Callbacks.OnCDemoClassInfo(parser.onCDemoClassInfo)
	parser.Callbacks.OnCDemoSendTables(parser.onCDemoSendTables)
	parser.Callbacks.OnCSVCMsg_CreateStringTable(parser.onCSVCMsg_CreateStringTable)
	parser.Callbacks.OnCSVCMsg_PacketEntities(parser.onCSVCMsg_PacketEntities)
	parser.Callbacks.OnCSVCMsg_SendTable(parser.onCSVCMsg_SendTable)
	parser.Callbacks.OnCSVCMsg_UpdateStringTable(parser.onCSVCMsg_UpdateStringTable)
	parser.Callbacks.OnCSVCMsg_ServerInfo(parser.onCSVCMsg_ServerInfo)
	parser.Callbacks.OnCMsgSource1LegacyGameEvent(parser.GameEvents.onCMsgSource1LegacyGameEvent)

	// Panic if we see any of these
	parser.Callbacks.OnCSVCMsg_GameEvent(func(m *dota.CSVCMsg_GameEvent) error {
		_panicf("unexpected: saw a CSVCMsg_GameEvent")
		return nil
	})

	// Maintains the value of parser.Tick
	parser.Callbacks.OnCNETMsg_Tick(func(m *dota.CNETMsg_Tick) error {
		parser.NetTick = m.GetTick()
		return nil
	})

	// Stops parsing when we reach the end of the replay.
	parser.Callbacks.OnCDemoStop(func(m *dota.CDemoStop) error {
		parser.Stop()
		return nil
	})

	return parser, nil
}

// Start parsing the replay. Will stop processing new events after Stop() is called.
func (p *Parser) Start() error {
	var msg *outerMessage
	var err error

	// Loop through all outer messages until we're signaled to stop. Stopping
	// happens when either the OnCDemoStop message is encountered or
	// parser.Stop() is called programatically.
	for !p.isStopping {
		// Read the next outer message.
		if msg, err = p.readOuterMessage(); err != nil {
			return err
		}

		// Update the parser tick
		p.Tick = msg.tick

		// Invoke callbacks for the given message type.
		if err = p.CallByDemoType(msg.typeId, msg.data); err != nil {
			return err
		}
	}

	return nil
}

// Stop parsing the replay, causing the parser to stop processing new events.
func (p *Parser) Stop() {
	p.isStopping = true
}

// Describes a demo message parsed from the replay.
type outerMessage struct {
	tick   uint32
	typeId int32
	data   []byte
}

// Read the next outer message from the buffer.
func (p *Parser) readOuterMessage() (*outerMessage, error) {
	// Read a command header, which includes both the message type
	// well as a flag to determine whether or not whether or not the
	// message is compressed with snappy.
	command := dota.EDemoCommands(p.reader.readVarUint32())

	// Extract the type and compressed flag out of the command
	msgType := int32(command & ^dota.EDemoCommands_DEM_IsCompressed)
	msgCompressed := (command & dota.EDemoCommands_DEM_IsCompressed) == dota.EDemoCommands_DEM_IsCompressed

	// Read the tick that the message corresponds with.
	tick := p.reader.readVarUint32()

	// This appears to actually be an int32, where a -1 means pre-game.
	if tick == 4294967295 {
		tick = 0
	}

	// Read the size and following buffer.
	size := int(p.reader.readVarUint32())
	buf := p.reader.readBytes(size)

	// If the buffer is compressed, decompress it with snappy.
	if msgCompressed {
		var err error
		if buf, err = snappy.Decode(nil, buf); err != nil {
			return nil, err
		}
	}

	// Return the message
	msg := &outerMessage{
		tick:   tick,
		typeId: msgType,
		data:   buf,
	}
	return msg, nil
}
<|endoftext|>


<commit_before>package model

import (
	"fmt"
	"time"
)

const (
	TimestampFormat = time.RFC822
	Indent          = "  "
)

func Age(timestamp time.Time) string {
	if timestamp.Equal(time.Unix(0, 0)) {
		return "unknown"
	}
	age := time.Now().Sub(timestamp)
	ageString := age.String()
	const year = 365 * 24
	switch {
	case age.Hours() > year:
		years := uint64(age.Hours()) / year
		ageString = fmt.Sprintf("%dy", years)
	case age.Hours() > 24 && age.Hours() < year:
		days := uint64(age.Hours()) / 24
		ageString = fmt.Sprintf("%dd", days)
	case age.Hours() <= 24 && age.Hours() > 1:
		hours := uint64(age.Hours())
		ageString = fmt.Sprintf("%dh", hours)
	case age.Hours() < 1 && age.Minutes() > 1:
		minutes := uint64(age.Minutes())
		ageString = fmt.Sprintf("%dm", minutes)
	default:
		seconds := uint64(age.Seconds())
		ageString = fmt.Sprintf("%ds", seconds)
	}
	return ageString
}
<commit_msg>minor fixes<commit_after>package model

import (
	"fmt"
	"time"
)

const (
	TimestampFormat = time.RFC3339
	Indent          = "  "
)

func Age(timestamp time.Time) string {
	if timestamp.Equal(time.Unix(0, 0)) {
		return "unknown"
	}
	age := time.Now().Sub(timestamp)
	var ageString string
	const year = 365 * 24
	switch {
	case age.Hours() > year:
		years := uint64(age.Hours()) / year
		ageString = fmt.Sprintf("%dy", years)
	case age.Hours() > 24 && age.Hours() < year:
		days := uint64(age.Hours()) / 24
		ageString = fmt.Sprintf("%dd", days)
	case age.Hours() <= 24 && age.Hours() > 1:
		hours := uint64(age.Hours())
		ageString = fmt.Sprintf("%dh", hours)
	case age.Hours() < 1 && age.Minutes() > 1:
		minutes := uint64(age.Minutes())
		ageString = fmt.Sprintf("%dm", minutes)
	default:
		seconds := uint64(age.Seconds())
		ageString = fmt.Sprintf("%ds", seconds)
	}
	return ageString
}
<|endoftext|>


<commit_before>/*
Package intern maps strings to integers for fast comparisons.

Description

Consider the string "This is string one" stored in variable x and the string
"This is string two" stored in variable y.  When a program tests if x == y, it
compares their first characters and finds that both are "T".  It then compares
their second characters and finds that both are "h".  This continues until the
16th character, where it finds that "o" is different from "t" and concludes
that the two strings are different.

The intern package provides two "symbol" data types to speed up such
comparisons.  The first, Eq, is an integer that is associated with particular
string contents.  For example, "This is string one" may be mapped to 52413
while "This is string two" may be mapped to 29488.  Comparing 52413 == 29488
fails in a single step, which can save a lot of time if such comparisons are
performed frequently.

The second data type, LGE, supports less than, greater than, and equal to
comparisons (and variations such as <=, >=, and !=).  For example, if x is "But
I have promises to keep", y is "And miles to go before I sleep", and z is also
"And miles to go before I sleep", then these might map to symbols 48712, 29519,
and 29519, respectively.  Let's call these sx, sy, and sz.  Then, x > y and sx
> sy; y == z and sy == sz; z < x and sz < sx; and any other such comparison of
x, y, and z, also holds for sx, sy, and sz.

A String method is defined for both Eq and LGE.  This maps a symbol back to its
original string.  Ergo, no information is lost when mapping strings to symbols.

Usage

NewEq maps a string to an Eq symbol, and NewLGE maps a string to an LGE symbol.
The former is faster and always succeeds (until the set of 64-bit integers is
exhausted).  NewLGE, in addition to being slower, can fail if earlier
assignments of integers to strings preclude a new string from being mapped to
an integer that respects comparisons with all existing symbols.  (In the
current implementation, a worst-case sequence of NewLGE calls will fail on the
65th call.)  As a workaround, the package provides a PreLGE function that
indicates an intention to invoke NewLGE on a particular string but without
actually assigning an integer.

Best practice is to pre-allocate as many LGE symbols as possible before calling
NewLGE.  When NewLGE is called, all strings previously passed to PreLGE are
assigned integers in an order that helps ensure that the desired relations are
preserved.  The process of pre-allocating LGE symbols with PreLGE and later
allocating them with NewLGE can be repeated as many times as necessary but with
increasing likelihood of failure with each repetition.  If NewLGE does fail,
the RemapAllLGEs function can be called to completely redo the mapping from
strings to LGE symbols.  The program will need to update any live LGE symbols
it has stored in data structures.

All functions in this package are thread-safe.
*/
package intern

import (
	"fmt"
	"sync"
)

// symbol represents either package symbol type (Eq or LGE).
type symbol uint64

// symbolList is a list of symbols.  It implements sort.Interface.
type symbolList []symbol

// Len returns the length of a symbolList.
func (sl symbolList) Len() int { return len(sl) }

// Less says is one symbolList element is less than another.
func (sl symbolList) Less(i, j int) bool { return sl[i] < sl[j] }

// Swap swaps two elements of a symbolList.
func (sl symbolList) Swap(i, j int) { sl[i], sl[j] = sl[j], sl[i] }

// state includes all the state needed to manipulate all interned-string types.
type state struct {
	symToStr     map[symbol]string // Mapping from symbols to strings
	strToSym     map[string]symbol // Mapping from strings to symbols
	tree         *tree             // Tree for maintaining symbols assignments
	pending      []string          // Strings not yet mapped to symbols
	sync.RWMutex                   // Mutex protecting all of the above
}

// forgetAll discards all extant string/symbol mappings and resets the
// assignment tables to their initial state.
func (st *state) forgetAll() {
	st.symToStr = make(map[symbol]string)
	st.strToSym = make(map[string]symbol)
	st.tree = nil
	st.pending = make([]string, 0, 100)
}

// toString converts a symbol back to a string.  It panics if given a symbol
// that was not created using New*.
func (st *state) toString(s symbol, ty string) string {
	st.RLock()
	defer st.RUnlock()
	if str, ok := st.symToStr[s]; ok {
		return str
	}
	panic(fmt.Sprintf("%d is not a valid intern.%s", s, ty))
}

// flushPending flushes all pending symbols, converting strings to symbols.
// The function returns an error status.
func (st *state) flushPending() error {
	var err error
	if len(st.pending) > 0 {
		var sMap map[string]symbol
		st.tree, sMap, err = st.tree.insertMany(st.pending)
		if err != nil {
			return err
		}
		st.pending = st.pending[:0]
		for k, v := range sMap {
			st.strToSym[k] = v
			st.symToStr[v] = k
		}
	}
	return nil
}

// assignSymbol assigns a new symbol to a string.  It takes as arguments the
// string to assign and a Boolean that indicates whether to use a tree to
// preserve order.  This method returns the assigned symbol and an error value.
func (st *state) assignSymbol(s string, useTree bool) (symbol, error) {
	// Check if the string was already assigned a symbol.
	sym, ok := st.strToSym[s]
	if ok {
		// The string was already assigned a symbol.
		return sym, nil
	}

	// We haven't seen this string before.  Find a symbol for it.
	if useTree {
		// Assign the symbol using a tree to maintain order.
		var err error
		st.tree, sym, err = st.tree.insert(s)
		if err != nil {
			return 0, err
		}
	} else {
		// Assign the next available number, starting at 1 to ensure
		// that an uninitialized symbol is treated as invalid.
		sym = symbol(len(st.symToStr)) + 1
	}
	st.symToStr[sym] = s
	st.strToSym[s] = sym
	return sym, nil
}
<commit_msg>Discuss performance in the package-level documentation<commit_after>/*
Package intern maps strings to integers for fast comparisons.

Description

Consider the string "This is string one" stored in variable x and the string
"This is string two" stored in variable y.  When a program tests if x == y, it
compares their first characters and finds that both are "T".  It then compares
their second characters and finds that both are "h".  This continues until the
16th character, where it finds that "o" is different from "t" and concludes
that the two strings are different.

The intern package provides two "symbol" data types to speed up such
comparisons.  The first, Eq, is an integer that is associated with particular
string contents.  For example, "This is string one" may be mapped to 52413
while "This is string two" may be mapped to 29488.  Comparing 52413 == 29488
fails in a single step, which can save a lot of time if such comparisons are
performed frequently.

The second data type, LGE, supports less than, greater than, and equal to
comparisons (and variations such as <=, >=, and !=).  For example, if x is "But
I have promises to keep", y is "And miles to go before I sleep", and z is also
"And miles to go before I sleep", then these might map to symbols 48712, 29519,
and 29519, respectively.  Let's call these sx, sy, and sz.  Then, x > y and sx
> sy; y == z and sy == sz; z < x and sz < sx; and any other such comparison of
x, y, and z, also holds for sx, sy, and sz.

A String method is defined for both Eq and LGE.  This maps a symbol back to its
original string.  Ergo, no information is lost when mapping strings to symbols.

Usage

NewEq maps a string to an Eq symbol, and NewLGE maps a string to an LGE symbol.
The former is faster and always succeeds (until the set of 64-bit integers is
exhausted).  NewLGE, in addition to being slower, can fail if earlier
assignments of integers to strings preclude a new string from being mapped to
an integer that respects comparisons with all existing symbols.  (In the
current implementation, a worst-case sequence of NewLGE calls will fail on the
65th call.)  As a workaround, the package provides a PreLGE function that
indicates an intention to invoke NewLGE on a particular string but without
actually assigning an integer.

Best practice is to pre-allocate as many LGE symbols as possible before calling
NewLGE.  When NewLGE is called, all strings previously passed to PreLGE are
assigned integers in an order that helps ensure that the desired relations are
preserved.  The process of pre-allocating LGE symbols with PreLGE and later
allocating them with NewLGE can be repeated as many times as necessary but with
increasing likelihood of failure with each repetition.  If NewLGE does fail,
the RemapAllLGEs function can be called to completely redo the mapping from
strings to LGE symbols.  The program will need to update any live LGE symbols
it has stored in data structures.

All functions in this package are thread-safe.

Performance

It's tricky to discuss the speed of Eq and LGE symbol comparisons relative to
string comparisons.  First, the time needed to compare two strings is a
function of the length of the common prefix.  For example, comparing "Jonas
Grumby" to "Roy Hinkley" is faster than comparing
"antidisestablishmentarianism" to "antidisestablishmentarianist".  In contrast,
symbol comparisons take the same amount of time regardless of the strings they
represent.  Second, symbols must be allocated, which takes time (more for an
LGE than an Eq).  Consequently, if few comparisons are performed per
allocation, it may be faster to use strings than symbols.  As a further twist,
the NewEqMulti and PreLGEMulti/NewLGEMulti functions help amortize some of the
allocation costs when allocating multiple symbols at once, but not all programs
have multiple strings they need to intern at once.  Third, there is a memory
cost associated with maintaining a bidirectional mapping between strings and
symbols.  If this extra memory causes a program's working set to expand beyond
the size of a cache, it may be faster to use strings than symbols.

The intern package includes a number of benchmarks to help programmers
determine if it may be beneficial to use the package.  Run them in
the usual manner:

    go test --bench=. --run=None

On my computer, the results show, very roughly, that

1) for random strings (i.e., very short common prefixes), symbols are faster
than strings if the program performs more than twice as many comparisons as Eq
allocations or more than 13 times as many comparisons as LGE allocations, and
that

2) for very long, very similar strings (well over 100 initial characters in
common), both Eq and LGE symbols are faster than strings if the program
performs at least as many comparisons as symbol allocations.

*/
package intern

import (
	"fmt"
	"sync"
)

// symbol represents either package symbol type (Eq or LGE).
type symbol uint64

// symbolList is a list of symbols.  It implements sort.Interface.
type symbolList []symbol

// Len returns the length of a symbolList.
func (sl symbolList) Len() int { return len(sl) }

// Less says is one symbolList element is less than another.
func (sl symbolList) Less(i, j int) bool { return sl[i] < sl[j] }

// Swap swaps two elements of a symbolList.
func (sl symbolList) Swap(i, j int) { sl[i], sl[j] = sl[j], sl[i] }

// state includes all the state needed to manipulate all interned-string types.
type state struct {
	symToStr     map[symbol]string // Mapping from symbols to strings
	strToSym     map[string]symbol // Mapping from strings to symbols
	tree         *tree             // Tree for maintaining symbols assignments
	pending      []string          // Strings not yet mapped to symbols
	sync.RWMutex                   // Mutex protecting all of the above
}

// forgetAll discards all extant string/symbol mappings and resets the
// assignment tables to their initial state.
func (st *state) forgetAll() {
	st.symToStr = make(map[symbol]string)
	st.strToSym = make(map[string]symbol)
	st.tree = nil
	st.pending = make([]string, 0, 100)
}

// toString converts a symbol back to a string.  It panics if given a symbol
// that was not created using New*.
func (st *state) toString(s symbol, ty string) string {
	st.RLock()
	defer st.RUnlock()
	if str, ok := st.symToStr[s]; ok {
		return str
	}
	panic(fmt.Sprintf("%d is not a valid intern.%s", s, ty))
}

// flushPending flushes all pending symbols, converting strings to symbols.
// The function returns an error status.
func (st *state) flushPending() error {
	var err error
	if len(st.pending) > 0 {
		var sMap map[string]symbol
		st.tree, sMap, err = st.tree.insertMany(st.pending)
		if err != nil {
			return err
		}
		st.pending = st.pending[:0]
		for k, v := range sMap {
			st.strToSym[k] = v
			st.symToStr[v] = k
		}
	}
	return nil
}

// assignSymbol assigns a new symbol to a string.  It takes as arguments the
// string to assign and a Boolean that indicates whether to use a tree to
// preserve order.  This method returns the assigned symbol and an error value.
func (st *state) assignSymbol(s string, useTree bool) (symbol, error) {
	// Check if the string was already assigned a symbol.
	sym, ok := st.strToSym[s]
	if ok {
		// The string was already assigned a symbol.
		return sym, nil
	}

	// We haven't seen this string before.  Find a symbol for it.
	if useTree {
		// Assign the symbol using a tree to maintain order.
		var err error
		st.tree, sym, err = st.tree.insert(s)
		if err != nil {
			return 0, err
		}
	} else {
		// Assign the next available number, starting at 1 to ensure
		// that an uninitialized symbol is treated as invalid.
		sym = symbol(len(st.symToStr)) + 1
	}
	st.symToStr[sym] = s
	st.strToSym[s] = sym
	return sym, nil
}
<|endoftext|>


<commit_before>package template

import (
	"fmt"
	"os"
	"strconv"
)

type parser struct {
	l   *lexer
	tok token
	lit []byte
	s   *scope
}

func (p *parser) Error(format string, args ...interface{}) {
	panic(fmt.Sprintf(format, args...))
}

func (p *parser) next() {
	p.tok, p.lit = p.l.scan()
}

func (p *parser) Expect(tok token) string {
	if p.tok != tok {
		p.Error("expected %s, got %s", tokStrings[tok], tokStrings[p.tok])
	}
	lit := p.lit
	p.next()
	return string(lit)
}

func (p *parser) ExpectWord(word string) {
	if p.tok != tokIdent || string(p.lit) != word {
		p.Error("expected ident %s, got token %s, %s", word, tokStrings[p.tok], p.lit)
	}
	p.next()
}

// parse until one of the following tags
func (p *parser) ParseUntil(tags ...string) (string, RenderList) {
	r := make(RenderList, 0, 10)
	for p.tok != tokEof {
		switch p.tok {
		case tokText:
			r = append(r, &printLit{p.lit})
			p.next()
		case tokBlockTagStart:
			p.next()
			lit := string(p.lit)
			for _, t := range tags {
				if t == lit {
					p.next()
					return t, r
				}
			}
			r = append(r, p.parseBlockTag())
		case tokVarTagStart:
			r = append(r, p.parseVarTag())
		default:
			p.Error("unexpected token %s", tokStrings[p.tok])
		}
	}
	return "", r
}

func (p *parser) parseBlockTag() Renderer {
	if tag, ok := tags[p.Expect(tokIdent)]; ok {
		return tag(p)
	}
	p.Error("tag isn't registered")
	return nil
}

func (p *parser) parseVarTag() Renderer {
	p.Expect(tokVarTagStart)
	v := p.parseVar()
	f := p.parseFilters()
	p.Expect(tokVarTagEnd)
	return &printVar{&expr{v, f}}
}

func (p *parser) parseVar() valuer {
	var ret valuer
	switch p.tok {
	case tokInt:
		i, err := strconv.Atoi64(string(p.lit))
		if err != nil {
			p.Error("internal int error: %s", err)
		}
		ret = intLit(i)
		p.next()
	case tokFloat:
		f, err := strconv.Atof64(string(p.lit))
		if err != nil {
			p.Error("Internal float error: %s", err)
		}
		ret = floatLit(f)
		p.next()
	case tokString:
		ret = stringLit(p.lit)
		p.next()
	case tokIdent:
		ret = p.parseAttrVar()
	default:
		p.Error("Unexpected token %s", tokStrings[p.tok])
	}
	return ret
}

func (p *parser) parseAttrVar() *variable {
	var v variable
	v.v = p.s.Lookup(p.Expect(tokIdent))
	for p.tok == tokDot {
		p.Expect(tokDot)
		v.attrs = append(v.attrs, string(p.lit))
		if p.tok == tokInt {
			p.Expect(tokInt)
		} else {
			p.Expect(tokIdent)
		}
	}
	return &v
}

func (p *parser) parseFilters() []*filter {
	f := make([]*filter, 0, 2)
	for p.tok == tokFilter {
		p.next()
		rf, ok := filters[string(p.lit)]
		if !ok {
			p.Error("filter does not exist")
		}
		p.Expect(tokIdent)
		var val valuer
		args := false
		switch rf.arg {
		case ReqArg:
			args = true
		case OptArg:
			args = p.tok == tokArgument
		case NoArg:
			if p.tok == tokArgument {
				p.Error("filter accepts no arguments")
			}
		}
		if args {
			p.Expect(tokArgument)
			val = p.parseVar()
		}
		f = append(f, &filter{rf.f, val})
	}
	return f
}

func Parse(s []byte) (*Template, os.Error) {
	t := new(Template)
	l := &lexer{src: s}
	l.init()
	p := &parser{l: l, s: newScope()}

	p.next()
	_, t.nodes = p.ParseUntil()
	t.scope = p.s

	return t, nil
}

func ParseString(s string) (*Template, os.Error) {
	return Parse([]byte(s))
}
<commit_msg>Add a way to parse files<commit_after>package template

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
)

type parser struct {
	l   *lexer
	tok token
	lit []byte
	s   *scope
}

func (p *parser) Error(format string, args ...interface{}) {
	panic(fmt.Sprintf(format, args...))
}

func (p *parser) next() {
	p.tok, p.lit = p.l.scan()
}

func (p *parser) Expect(tok token) string {
	if p.tok != tok {
		p.Error("expected %s, got %s", tokStrings[tok], tokStrings[p.tok])
	}
	lit := p.lit
	p.next()
	return string(lit)
}

func (p *parser) ExpectWord(word string) {
	if p.tok != tokIdent || string(p.lit) != word {
		p.Error("expected ident %s, got token %s, %s", word, tokStrings[p.tok], p.lit)
	}
	p.next()
}

// parse until one of the following tags
func (p *parser) ParseUntil(tags ...string) (string, RenderList) {
	r := make(RenderList, 0, 10)
	for p.tok != tokEof {
		switch p.tok {
		case tokText:
			r = append(r, &printLit{p.lit})
			p.next()
		case tokBlockTagStart:
			p.next()
			lit := string(p.lit)
			for _, t := range tags {
				if t == lit {
					p.next()
					return t, r
				}
			}
			r = append(r, p.parseBlockTag())
		case tokVarTagStart:
			r = append(r, p.parseVarTag())
		default:
			p.Error("unexpected token %s", tokStrings[p.tok])
		}
	}
	return "", r
}

func (p *parser) parseBlockTag() Renderer {
	if tag, ok := tags[p.Expect(tokIdent)]; ok {
		return tag(p)
	}
	p.Error("tag isn't registered")
	return nil
}

func (p *parser) parseVarTag() Renderer {
	p.Expect(tokVarTagStart)
	v := p.parseVar()
	f := p.parseFilters()
	p.Expect(tokVarTagEnd)
	return &printVar{&expr{v, f}}
}

func (p *parser) parseVar() valuer {
	var ret valuer
	switch p.tok {
	case tokInt:
		i, err := strconv.Atoi64(string(p.lit))
		if err != nil {
			p.Error("internal int error: %s", err)
		}
		ret = intLit(i)
		p.next()
	case tokFloat:
		f, err := strconv.Atof64(string(p.lit))
		if err != nil {
			p.Error("Internal float error: %s", err)
		}
		ret = floatLit(f)
		p.next()
	case tokString:
		ret = stringLit(p.lit)
		p.next()
	case tokIdent:
		ret = p.parseAttrVar()
	default:
		p.Error("Unexpected token %s", tokStrings[p.tok])
	}
	return ret
}

func (p *parser) parseAttrVar() *variable {
	var v variable
	v.v = p.s.Lookup(p.Expect(tokIdent))
	for p.tok == tokDot {
		p.Expect(tokDot)
		v.attrs = append(v.attrs, string(p.lit))
		if p.tok == tokInt {
			p.Expect(tokInt)
		} else {
			p.Expect(tokIdent)
		}
	}
	return &v
}

func (p *parser) parseFilters() []*filter {
	f := make([]*filter, 0, 2)
	for p.tok == tokFilter {
		p.next()
		rf, ok := filters[string(p.lit)]
		if !ok {
			p.Error("filter does not exist")
		}
		p.Expect(tokIdent)
		var val valuer
		args := false
		switch rf.arg {
		case ReqArg:
			args = true
		case OptArg:
			args = p.tok == tokArgument
		case NoArg:
			if p.tok == tokArgument {
				p.Error("filter accepts no arguments")
			}
		}
		if args {
			p.Expect(tokArgument)
			val = p.parseVar()
		}
		f = append(f, &filter{rf.f, val})
	}
	return f
}

func Parse(s []byte) (*Template, os.Error) {
	t := new(Template)
	l := &lexer{src: s}
	l.init()
	p := &parser{l: l, s: newScope()}

	p.next()
	_, t.nodes = p.ParseUntil()
	t.scope = p.s

	return t, nil
}

func ParseString(s string) (*Template, os.Error) {
	return Parse([]byte(s))
}

func ParseFile(name string) (*Template, os.Error) {
	b, err := ioutil.ReadFile(name)
	if err != nil {
		return nil, err
	}
	return Parse(b)
}
<|endoftext|>


<commit_before>/*
   Copyright The containerd 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 progress

import (
	"bytes"
	"fmt"
)

// TODO(stevvooe): We may want to support more interesting parameterization of
// the bar. For now, it is very simple.

// Bar provides a very simple progress bar implementation.
//
// Use with fmt.Printf and "r" to format the progress bar. A "-" flag makes it
// progress from right to left.
type Bar float64

var _ fmt.Formatter = Bar(1.0)

// Format the progress bar as output
func (h Bar) Format(state fmt.State, r rune) {
	switch r {
	case 'r':
	default:
		panic(fmt.Sprintf("%v: unexpected format character", float64(h)))
	}

	if h > 1.0 {
		h = 1.0
	}

	if h < 0.0 {
		h = 0.0
	}

	if state.Flag('-') {
		h = 1.0 - h
	}

	width, ok := state.Width()
	if !ok {
		// default width of 40
		width = 40
	}

	var pad int

	extra := len([]byte(green)) + len([]byte(reset))

	p := make([]byte, width+extra)
	p[0], p[len(p)-1] = '|', '|'
	pad += 2

	positive := int(Bar(width-pad) * h)
	negative := width - pad - positive

	n := 1
	n += copy(p[n:], []byte(green))
	n += copy(p[n:], bytes.Repeat([]byte("+"), positive))
	n += copy(p[n:], []byte(reset))

	if negative > 0 {
		copy(p[n:len(p)-1], bytes.Repeat([]byte("-"), negative))
	}

	state.Write(p)
}
<commit_msg>remove excessive []byte(s) conversion <commit_after>/*
   Copyright The containerd 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 progress

import (
	"bytes"
	"fmt"
)

// TODO(stevvooe): We may want to support more interesting parameterization of
// the bar. For now, it is very simple.

// Bar provides a very simple progress bar implementation.
//
// Use with fmt.Printf and "r" to format the progress bar. A "-" flag makes it
// progress from right to left.
type Bar float64

var _ fmt.Formatter = Bar(1.0)

// Format the progress bar as output
func (h Bar) Format(state fmt.State, r rune) {
	switch r {
	case 'r':
	default:
		panic(fmt.Sprintf("%v: unexpected format character", float64(h)))
	}

	if h > 1.0 {
		h = 1.0
	}

	if h < 0.0 {
		h = 0.0
	}

	if state.Flag('-') {
		h = 1.0 - h
	}

	width, ok := state.Width()
	if !ok {
		// default width of 40
		width = 40
	}

	var pad int

	extra := len([]byte(green)) + len([]byte(reset))

	p := make([]byte, width+extra)
	p[0], p[len(p)-1] = '|', '|'
	pad += 2

	positive := int(Bar(width-pad) * h)
	negative := width - pad - positive

	n := 1
	n += copy(p[n:], green)
	n += copy(p[n:], bytes.Repeat([]byte("+"), positive))
	n += copy(p[n:], reset)

	if negative > 0 {
		copy(p[n:len(p)-1], bytes.Repeat([]byte("-"), negative))
	}

	state.Write(p)
}
<|endoftext|>


<commit_before>package ipfix

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"runtime"
)

// The version field in IPFIX messages should always have the value 10. If it
// does not, you get this error. It's probably a sign of a bug in the parser or
// the exporter and that we have lost synchronization with the data stream.
// Reestablishing the session is the only way forward at this point.
var ErrVersion = errors.New("incorrect version field in message header - out of sync?")

// A Message is the top level construct representing an IPFIX message. A well
// formed message contains one or more sets of data or template information.
type Message struct {
	Header       MessageHeader
	DataSets     []DataSet
	TemplateSets []TemplateSet
}

// The MessageHeader provides metadata for the entire Message. The sequence
// number and domain ID can be used to gain knowledge of messages lost on an
// unreliable transport such as UDP.
type MessageHeader struct {
	Version        uint16 // Always 0x0a
	Length         uint16
	ExportTime     uint32 // Epoch seconds
	SequenceNumber uint32
	DomainId       uint32
}

type setHeader struct {
	SetId  uint16
	Length uint16
}

type templateHeader struct {
	TemplateId uint16
	FieldCount uint16
}

// The DataSet represents a single exported flow. The Records each describe
// different aspects of the flow (source and destination address, counters,
// service, etc.).
type DataSet struct {
	TemplateId uint16
	Records    [][]byte
}

// The TemplateSet describes a data template, as used by DataSets.
type TemplateSet struct {
	TemplateId uint16
	Records    []TemplateRecord
}

// The TemplateRecord describes the ID and size of the corresponding Records in a DataSet.
type TemplateRecord struct {
	EnterpriseId uint32
	FieldId      uint16
	Length       uint16
}

// The Session is the context for IPFIX messages.
type Session struct {
	templates  [][]TemplateRecord
	reader     *bufio.Reader
	minRecord  []uint16
	dictionary dictionary
}

// NewSession initializes a new Session based on the provided io.Reader.
func NewSession(reader io.Reader) *Session {
	s := Session{}
	s.templates = make([][]TemplateRecord, 65536)
	s.reader = bufio.NewReaderSize(reader, 32768)
	s.minRecord = make([]uint16, 65536)
	s.dictionary = builtinDictionary
	return &s
}

var msgHeaderLength = binary.Size(MessageHeader{})
var setHeaderLength = binary.Size(setHeader{})
var templateHeaderLength = binary.Size(templateHeader{})

// ReadMessage extracts and returns one message from the IPFIX stream. As long
// as err is nil, further messages can be read from the stream. Errors are not
// recoverable -- once an error has been returned, ReadMessage should not be
// called again on the same session.
func (s *Session) ReadMessage() (msg *Message, err error) {
	defer func() {
		if r := recover(); r != nil {
			if _, ok := r.(runtime.Error); ok {
				panic(r)
			}

			msg = nil
			err = r.(error)
		}
	}()

	msg = &Message{}
	msg.DataSets = make([]DataSet, 0)
	msg.TemplateSets = make([]TemplateSet, 0)

	msgHdr := MessageHeader{}
	err = binary.Read(s.reader, binary.BigEndian, &msgHdr)
	s.errorIf(err)
	if msgHdr.Version != 10 {
		s.errorIf(ErrVersion)
	}
	read := msgHeaderLength
	msg.Header = msgHdr

	for read < int(msgHdr.Length) {
		tsets, dsets, tr := s.readSet()
		read += tr
		msg.TemplateSets = append(msg.TemplateSets, tsets...)
		msg.DataSets = append(msg.DataSets, dsets...)
	}
	return
}

func (s *Session) errorIf(err error) {
	if err != nil {
		panic(err)
	}
}

func (s *Session) readSet() (tsets []TemplateSet, dsets []DataSet, read int) {
	tsets = make([]TemplateSet, 0)
	dsets = make([]DataSet, 0)

	setHdr := setHeader{}
	err := binary.Read(s.reader, binary.BigEndian, &setHdr)
	s.errorIf(err)
	read += setHeaderLength

	end := int(setHdr.Length)
	for read < end {
		if end-read < int(s.minRecord[setHdr.SetId]) {
			// Padding
			_, err = io.ReadFull(s.reader, make([]byte, end-read))
			s.errorIf(err)
			read += end - read
		} else if setHdr.SetId == 2 {
			// Template Set
			ts, tsRead := s.readTemplateSet()
			read += tsRead
			tsets = append(tsets, *ts)

			// Update the template cache
			tid := ts.TemplateId
			if len(ts.Records) == 0 {
				// Set was withdrawn
				s.templates[tid] = nil
			} else {
				s.templates[tid] = ts.Records
			}

			// Update the minimum record length cache
			var minLength uint16
			for i := range ts.Records {
				if ts.Records[i].Length == 65535 {
					minLength += 1
				} else {
					minLength += ts.Records[i].Length
				}
			}
			s.minRecord[tid] = minLength
		} else if setHdr.SetId == 3 {
			// Options Template Set, not handled
			_, err = io.ReadFull(s.reader, make([]byte, end-read))
			s.errorIf(err)
			read += end - read
		} else {
			if tpl := s.templates[setHdr.SetId]; tpl != nil {
				// Data set
				ds, dsRead := s.readDataSet(tpl)
				read += dsRead
				ds.TemplateId = setHdr.SetId
				dsets = append(dsets, *ds)
			} else {
				// Data set with unknown template
				_, err := io.ReadFull(s.reader, make([]byte, end-read))
				s.errorIf(err)
				read += end - read
			}
		}
	}
	return
}

func (s *Session) readDataSet(tpl []TemplateRecord) (ds *DataSet, read int) {
	ds = &DataSet{}

	ds.Records = make([][]byte, len(tpl))
	for i := range tpl {
		var bs []byte
		if tpl[i].Length == 65535 {
			var bsRead int
			bs, bsRead = s.readVariableLength()
			read += bsRead
		} else {
			bs = make([]byte, tpl[i].Length)
			_, err := io.ReadFull(s.reader, bs)
			s.errorIf(err)
			read += len(bs)
		}
		ds.Records[i] = bs
	}

	return
}

func (s *Session) readTemplateSet() (ts *TemplateSet, read int) {
	ts = &TemplateSet{}
	th := templateHeader{}
	err := binary.Read(s.reader, binary.BigEndian, &th)
	s.errorIf(err)
	read = templateHeaderLength

	ts.TemplateId = th.TemplateId
	ts.Records = make([]TemplateRecord, th.FieldCount)
	for i := 0; i < int(th.FieldCount); i++ {
		f := TemplateRecord{}
		err = binary.Read(s.reader, binary.BigEndian, &f.FieldId)
		s.errorIf(err)
		read += 2
		err = binary.Read(s.reader, binary.BigEndian, &f.Length)
		s.errorIf(err)
		read += 2
		if f.FieldId >= 0x8000 {
			f.FieldId -= 0x8000
			err = binary.Read(s.reader, binary.BigEndian, &f.EnterpriseId)
			s.errorIf(err)
			read += 4
		}
		ts.Records[i] = f
	}

	return
}

// Reads a variable length information element. Returns a slice with the data
// in it, the total number of bytes read.
func (s *Session) readVariableLength() (bs []byte, r int) {
	var l int

	var l0 uint8
	err := binary.Read(s.reader, binary.BigEndian, &l0)
	s.errorIf(err)
	r += 1

	if l0 < 255 {
		l = int(l0)
	} else {
		var l1 uint16
		err = binary.Read(s.reader, binary.BigEndian, &l1)
		s.errorIf(err)
		r += 2
		l = int(l1)
	}

	bs = make([]byte, l)
	_, err = io.ReadFull(s.reader, bs)
	s.errorIf(err)
	r += len(bs)

	return
}
<commit_msg>Refactor out bytesRead<commit_after>package ipfix

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
	"runtime"
)

// The version field in IPFIX messages should always have the value 10. If it
// does not, you get this error. It's probably a sign of a bug in the parser or
// the exporter and that we have lost synchronization with the data stream.
// Reestablishing the session is the only way forward at this point.
var ErrVersion = errors.New("incorrect version field in message header - out of sync?")

// A Message is the top level construct representing an IPFIX message. A well
// formed message contains one or more sets of data or template information.
type Message struct {
	Header       MessageHeader
	DataSets     []DataSet
	TemplateSets []TemplateSet
}

// The MessageHeader provides metadata for the entire Message. The sequence
// number and domain ID can be used to gain knowledge of messages lost on an
// unreliable transport such as UDP.
type MessageHeader struct {
	Version        uint16 // Always 0x0a
	Length         uint16
	ExportTime     uint32 // Epoch seconds
	SequenceNumber uint32
	DomainId       uint32
}

type setHeader struct {
	SetId  uint16
	Length uint16
}

type templateHeader struct {
	TemplateId uint16
	FieldCount uint16
}

// The DataSet represents a single exported flow. The Records each describe
// different aspects of the flow (source and destination address, counters,
// service, etc.).
type DataSet struct {
	TemplateId uint16
	Records    [][]byte
}

// The TemplateSet describes a data template, as used by DataSets.
type TemplateSet struct {
	TemplateId uint16
	Records    []TemplateRecord
}

// The TemplateRecord describes the ID and size of the corresponding Records in a DataSet.
type TemplateRecord struct {
	EnterpriseId uint32
	FieldId      uint16
	Length       uint16
}

// The Session is the context for IPFIX messages.
type Session struct {
	templates  [][]TemplateRecord
	reader     *bufio.Reader
	minRecord  []uint16
	dictionary dictionary
	bytesRead  int
}

// NewSession initializes a new Session based on the provided io.Reader.
func NewSession(reader io.Reader) *Session {
	s := Session{}
	s.templates = make([][]TemplateRecord, 65536)
	s.reader = bufio.NewReaderSize(reader, 32768)
	s.minRecord = make([]uint16, 65536)
	s.dictionary = builtinDictionary
	return &s
}

var msgHeaderLength = binary.Size(MessageHeader{})
var setHeaderLength = binary.Size(setHeader{})
var templateHeaderLength = binary.Size(templateHeader{})

// ReadMessage extracts and returns one message from the IPFIX stream. As long
// as err is nil, further messages can be read from the stream. Errors are not
// recoverable -- once an error has been returned, ReadMessage should not be
// called again on the same session.
func (s *Session) ReadMessage() (msg *Message, err error) {
	defer func() {
		if r := recover(); r != nil {
			if _, ok := r.(runtime.Error); ok {
				panic(r)
			}

			msg = nil
			err = r.(error)
		}
	}()

	s.bytesRead = 0
	msg = &Message{}
	msg.DataSets = make([]DataSet, 0)
	msg.TemplateSets = make([]TemplateSet, 0)

	msgHdr := MessageHeader{}
	err = binary.Read(s.reader, binary.BigEndian, &msgHdr)
	s.errorIf(err)
	if msgHdr.Version != 10 {
		s.errorIf(ErrVersion)
	}
	s.bytesRead += msgHeaderLength
	msg.Header = msgHdr

	for s.bytesRead < int(msgHdr.Length) {
		tsets, dsets := s.readSet()
		msg.TemplateSets = append(msg.TemplateSets, tsets...)
		msg.DataSets = append(msg.DataSets, dsets...)
	}

	return
}

func (s *Session) errorIf(err error) {
	if err != nil {
		panic(err)
	}
}

func (s *Session) readSet() ([]TemplateSet, []DataSet) {
	tsets := make([]TemplateSet, 0)
	dsets := make([]DataSet, 0)

	setHdr := setHeader{}
	err := binary.Read(s.reader, binary.BigEndian, &setHdr)
	s.errorIf(err)

	end := s.bytesRead + int(setHdr.Length)
	s.bytesRead += setHeaderLength

	for s.bytesRead < end {
		if end-s.bytesRead < int(s.minRecord[setHdr.SetId]) {
			// Padding
			_, err = io.ReadFull(s.reader, make([]byte, end-s.bytesRead))
			s.errorIf(err)
			s.bytesRead = end
		} else if setHdr.SetId == 2 {
			// Template Set
			ts := s.readTemplateSet()
			tsets = append(tsets, *ts)

			// Update the template cache
			tid := ts.TemplateId
			if len(ts.Records) == 0 {
				// Set was withdrawn
				s.templates[tid] = nil
			} else {
				s.templates[tid] = ts.Records
			}

			// Update the minimum record length cache
			var minLength uint16
			for i := range ts.Records {
				if ts.Records[i].Length == 65535 {
					minLength += 1
				} else {
					minLength += ts.Records[i].Length
				}
			}
			s.minRecord[tid] = minLength
		} else if setHdr.SetId == 3 {
			// Options Template Set, not handled
			s.readFull(end - s.bytesRead)
		} else {
			if tpl := s.templates[setHdr.SetId]; tpl != nil {
				// Data set
				ds := s.readDataSet(tpl)
				ds.TemplateId = setHdr.SetId
				dsets = append(dsets, *ds)
			} else {
				// Data set with unknown template
				s.readFull(end - s.bytesRead)
			}
		}
	}

	return tsets, dsets
}

func (s *Session) readFull(n int) []byte {
	bs := make([]byte, n)
	_, err := io.ReadFull(s.reader, bs)
	s.errorIf(err)
	s.bytesRead += len(bs)
	return bs
}

func (s *Session) readDataSet(tpl []TemplateRecord) *DataSet {
	ds := DataSet{}
	ds.Records = make([][]byte, len(tpl))

	for i := range tpl {
		var bs []byte
		if tpl[i].Length == 65535 {
			bs = s.readVariableLength()
		} else {
			bs = s.readFull(int(tpl[i].Length))
		}
		ds.Records[i] = bs
	}

	return &ds
}

func (s *Session) readTemplateSet() *TemplateSet {
	ts := TemplateSet{}
	th := templateHeader{}
	err := binary.Read(s.reader, binary.BigEndian, &th)
	s.errorIf(err)
	s.bytesRead += templateHeaderLength

	ts.TemplateId = th.TemplateId
	ts.Records = make([]TemplateRecord, th.FieldCount)
	for i := 0; i < int(th.FieldCount); i++ {
		f := TemplateRecord{}
		err = binary.Read(s.reader, binary.BigEndian, &f.FieldId)
		s.errorIf(err)
		s.bytesRead += 2
		err = binary.Read(s.reader, binary.BigEndian, &f.Length)
		s.errorIf(err)
		s.bytesRead += 2
		if f.FieldId >= 0x8000 {
			f.FieldId -= 0x8000
			err = binary.Read(s.reader, binary.BigEndian, &f.EnterpriseId)
			s.errorIf(err)
			s.bytesRead += 4
		}
		ts.Records[i] = f
	}

	return &ts
}

func (s *Session) readVariableLength() []byte {
	var l int

	var l0 uint8
	err := binary.Read(s.reader, binary.BigEndian, &l0)
	s.errorIf(err)
	s.bytesRead += 1

	if l0 < 255 {
		l = int(l0)
	} else {
		var l1 uint16
		err = binary.Read(s.reader, binary.BigEndian, &l1)
		s.errorIf(err)
		s.bytesRead += 2
		l = int(l1)
	}

	return s.readFull(l)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package version

// Base version information.
//
// This is the fallback data used when version information from git is not
// provided via go ldflags. It provides an approximation of the Kubernetes
// version for ad-hoc builds (e.g. `go build`) that cannot get the version
// information from git.
//
// The "-dev" suffix in the version info indicates that fact, and it means the
// current build is from a version greater that version. For example, v0.7-dev
// means version > 0.7 and < 0.8. (There's exceptions to this rule, see
// docs/releasing.md for more details.)
//
// When releasing a new Kubernetes version, this file should be updated to
// reflect the new version, and then a git annotated tag (using format vX.Y
// where X == Major version and Y == Minor version) should be created to point
// to the commit that updates pkg/version/base.go

var (
	// TODO: Deprecate gitMajor and gitMinor, use only gitVersion instead.
	gitMajor     string = "0"              // major version, always numeric
	gitMinor     string = "4.1"            // minor version, numeric possibly followed by "+"
	gitVersion   string = "v0.4.1"         // version from git, output of $(git describe)
	gitCommit    string = ""               // sha1 from git, output of $(git rev-parse HEAD)
	gitTreeState string = "not a git tree" // state of git tree, either "clean" or "dirty"
)
<commit_msg>Cut v0.4.1-dev<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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

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

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

package version

// Base version information.
//
// This is the fallback data used when version information from git is not
// provided via go ldflags. It provides an approximation of the Kubernetes
// version for ad-hoc builds (e.g. `go build`) that cannot get the version
// information from git.
//
// The "-dev" suffix in the version info indicates that fact, and it means the
// current build is from a version greater that version. For example, v0.7-dev
// means version > 0.7 and < 0.8. (There's exceptions to this rule, see
// docs/releasing.md for more details.)
//
// When releasing a new Kubernetes version, this file should be updated to
// reflect the new version, and then a git annotated tag (using format vX.Y
// where X == Major version and Y == Minor version) should be created to point
// to the commit that updates pkg/version/base.go

var (
	// TODO: Deprecate gitMajor and gitMinor, use only gitVersion instead.
	gitMajor     string = "0"              // major version, always numeric
	gitMinor     string = "4.1+"           // minor version, numeric possibly followed by "+"
	gitVersion   string = "v0.4.1-dev"     // version from git, output of $(git describe)
	gitCommit    string = ""               // sha1 from git, output of $(git rev-parse HEAD)
	gitTreeState string = "not a git tree" // state of git tree, either "clean" or "dirty"
)
<|endoftext|>


<commit_before>/*
Copyright 2014 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 version

// Base version information.
//
// This is the fallback data used when version information from git is not
// provided via go ldflags. It provides an approximation of the Kubernetes
// version for ad-hoc builds (e.g. `go build`) that cannot get the version
// information from git.
//
// If you are looking at these fields in the git tree, they look
// strange. They are modified on the fly by the build process. The
// in-tree values are dummy values used for "git archive", which also
// works for GitHub tar downloads.
//
// When releasing a new Kubernetes version, this file is updated by
// build/mark_new_version.sh to reflect the new version, and then a
// git annotated tag (using format vX.Y where X == Major version and Y
// == Minor version) is created to point to the commit that updates
// pkg/version/base.go
var (
	// TODO: Deprecate gitMajor and gitMinor, use only gitVersion
	// instead. First step in deprecation, keep the fields but make
	// them irrelevant. (Next we'll take it out, which may muck with
	// scripts consuming the kubectl version output - but most of
	// these should be looking at gitVersion already anyways.)
	gitMajor string = "1"  // major version, always numeric
	gitMinor string = "5+" // minor version, numeric possibly followed by "+"

	// semantic version, derived by build scripts (see
	// https://github.com/kubernetes/kubernetes/blob/master/docs/design/versioning.md
	// for a detailed discussion of this field)
	//
	// TODO: This field is still called "gitVersion" for legacy
	// reasons. For prerelease versions, the build metadata on the
	// semantic version is a git hash, but the version itself is no
	// longer the direct output of "git describe", but a slight
	// translation to be semver compliant.
	gitVersion   string = "v1.5.2-beta.0+$Format:%h$"
	gitCommit    string = "$Format:%H$"    // sha1 from git, output of $(git rev-parse HEAD)
	gitTreeState string = "not a git tree" // state of git tree, either "clean" or "dirty"

	buildDate string = "1970-01-01T00:00:00Z" // build date in ISO8601 format, output of $(date -u +'%Y-%m-%dT%H:%M:%SZ')
)
<commit_msg>Kubernetes version v1.5.2<commit_after>/*
Copyright 2014 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 version

// Base version information.
//
// This is the fallback data used when version information from git is not
// provided via go ldflags. It provides an approximation of the Kubernetes
// version for ad-hoc builds (e.g. `go build`) that cannot get the version
// information from git.
//
// If you are looking at these fields in the git tree, they look
// strange. They are modified on the fly by the build process. The
// in-tree values are dummy values used for "git archive", which also
// works for GitHub tar downloads.
//
// When releasing a new Kubernetes version, this file is updated by
// build/mark_new_version.sh to reflect the new version, and then a
// git annotated tag (using format vX.Y where X == Major version and Y
// == Minor version) is created to point to the commit that updates
// pkg/version/base.go
var (
	// TODO: Deprecate gitMajor and gitMinor, use only gitVersion
	// instead. First step in deprecation, keep the fields but make
	// them irrelevant. (Next we'll take it out, which may muck with
	// scripts consuming the kubectl version output - but most of
	// these should be looking at gitVersion already anyways.)
	gitMajor string = "1" // major version, always numeric
	gitMinor string = "5" // minor version, numeric possibly followed by "+"

	// semantic version, derived by build scripts (see
	// https://github.com/kubernetes/kubernetes/blob/master/docs/design/versioning.md
	// for a detailed discussion of this field)
	//
	// TODO: This field is still called "gitVersion" for legacy
	// reasons. For prerelease versions, the build metadata on the
	// semantic version is a git hash, but the version itself is no
	// longer the direct output of "git describe", but a slight
	// translation to be semver compliant.
	gitVersion   string = "v1.5.2+$Format:%h$"
	gitCommit    string = "$Format:%H$"    // sha1 from git, output of $(git rev-parse HEAD)
	gitTreeState string = "not a git tree" // state of git tree, either "clean" or "dirty"

	buildDate string = "1970-01-01T00:00:00Z" // build date in ISO8601 format, output of $(date -u +'%Y-%m-%dT%H:%M:%SZ')
)
<|endoftext|>


<commit_before>package idm

import (
	"fmt"
	"io"
)

// Parser represents a parser.
type Parser struct {
	s   *Scanner
	buf struct {
		t   Token  // last read token
		lit string // last read literal
		n   int    // buffer size (max=1)
	}
}

// NewParser returns a new instance of Parser.
func NewParser(r io.Reader) *Parser {
	return &Parser{s: NewScanner(r)}
}

// scan returns the next token from the underlying scanner.
// if a token has been unscanned then read that instead.
func (p *Parser) scan() (t Token, lit string) {

	if p.buf.n != 0 {
		p.buf.n = 0
		return p.buf.t, p.buf.lit
	}

	t, lit = p.s.Scan()

	p.buf.t, p.buf.lit = t, lit
	return
}

// unscan pushes the previously read token back onto the buffer.
func (p *Parser) unscan() { p.buf.n = 1 }

// scanIgnoreWhitespace scans the next non-whitespace token.
func (p *Parser) scanIgnoreWhitespace() (t Token, lit string) {
	t, lit = p.scan()
	if t == Space {
		t, lit = p.scan()
	}
	return
}

// Statement represents a code statement a = 2.
type Statement struct {
	Left  string
	Rigth string
}

// Parse parse a assign statement a = b
func (p *Parser) Parse() (*Statement, error) {
	stmt := &Statement{}

	tok, lit := p.scanIgnoreWhitespace()
	if tok != Identifier {
		return nil, fmt.Errorf("found %q, expected left", lit)
	}
	stmt.Left = lit

	if tok, lit := p.scanIgnoreWhitespace(); tok != Assign {
		return nil, fmt.Errorf("found %q, expected '=' with tok: %v, expected %v", lit, tok, Assign)
	}

	tok, lit = p.scanIgnoreWhitespace()
	if tok != Identifier {
		return nil, fmt.Errorf("found %q, expected identifier name", lit)
	}
	stmt.Rigth = lit

	return stmt, nil
}
<commit_msg>fix typo<commit_after>package idm

import (
	"fmt"
	"io"
)

// Parser represents a parser.
type Parser struct {
	s   *Scanner
	buf struct {
		t   Token  // last read token
		lit string // last read literal
		n   int    // buffer size (max=1)
	}
}

// NewParser returns a new instance of Parser.
func NewParser(r io.Reader) *Parser {
	return &Parser{s: NewScanner(r)}
}

// scan returns the next token from the underlying scanner.
// if a token has been unscanned then read that instead.
func (p *Parser) scan() (t Token, lit string) {

	if p.buf.n != 0 {
		p.buf.n = 0
		return p.buf.t, p.buf.lit
	}

	t, lit = p.s.Scan()

	p.buf.t, p.buf.lit = t, lit
	return
}

// unscan pushes the previously read token back onto the buffer.
func (p *Parser) unscan() { p.buf.n = 1 }

// scanIgnoreWhitespace scans the next non-whitespace token.
func (p *Parser) scanIgnoreWhitespace() (t Token, lit string) {
	t, lit = p.scan()
	if t == Space {
		t, lit = p.scan()
	}
	return
}

// Statement represents a code statement a = 2.
type Statement struct {
	Left  string
	Right string
}

// Parse parse a assign statement a = b
func (p *Parser) Parse() (*Statement, error) {
	stmt := &Statement{}

	tok, lit := p.scanIgnoreWhitespace()
	if tok != Identifier {
		return nil, fmt.Errorf("found %q, expected left", lit)
	}
	stmt.Left = lit

	if tok, lit := p.scanIgnoreWhitespace(); tok != Assign {
		return nil, fmt.Errorf("found %q, expected '=' with tok: %v, expected %v", lit, tok, Assign)
	}

	tok, lit = p.scanIgnoreWhitespace()
	if tok != Identifier {
		return nil, fmt.Errorf("found %q, expected identifier name", lit)
	}
	stmt.Right = lit

	return stmt, nil
}
<|endoftext|>


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

package cmd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"regexp"
	"strings"

	"github.com/spf13/cobra"

	"github.com/plusvic/go-ansi"

	"github.com/fatih/color"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"

	"github.com/VirusTotal/vt-cli/utils"
	"github.com/VirusTotal/vt-cli/yaml"
	"github.com/VirusTotal/vt-go/vt"
)

var colorScheme = yaml.Colors{
	KeyColor:     color.New(color.FgYellow),
	ValueColor:   color.New(color.FgHiGreen),
	CommentColor: color.New(color.Faint)}

func addAPIKeyFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"apikey", "k", "",
		"api key")
}

func addHostFlag(flags *pflag.FlagSet) {
	flags.String(
		"host", "www.virustotal.com",
		"API host name")
	flags.MarkHidden("host")
}

func addIncludeExcludeFlags(flags *pflag.FlagSet) {
	flags.StringSliceP(
		"include", "i", []string{"**"},
		"include fields matching the provided pattern")

	flags.StringSliceP(
		"exclude", "x", []string{},
		"exclude fields matching the provided pattern")
}

func addThreadsFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"threads", "t", 5,
		"number of threads working in parallel")
}

func addIDOnlyFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"identifiers-only", "I", false,
		"print identifiers only")
}

func addLimitFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"limit", "n", 10,
		"maximum number of results")
}

func addCursorFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"cursor", "c", "",
		"cursor")
}

func addOutputFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"output", "o", ".",
		"directory where downloaded files are put")
}

func addFilterFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"filter", "f", "",
		"filter")
}

func addVerboseFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"verbose", "v", false,
		"verbose output")
}

func addYAMLFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"yaml", "y", false,
		"output in YAML format")
}

// ReadFile reads the specified file and returns its content. If filename is "-"
// the data is read from stdin.
func ReadFile(filename string) ([]byte, error) {
	if filename == "-" {
		return ioutil.ReadAll(os.Stdin)
	}
	return ioutil.ReadFile(filename)
}

// PrintCommandLineWithCursor prints the same command-line that was used for
// executing the program but adding or replacing the --cursor flag with
// the current cursor for the given iterator.
func PrintCommandLineWithCursor(cmd *cobra.Command, it *vt.Iterator) {
	if cursor := it.Cursor(); cursor != "" {
		args := cmd.Flags().Args()
		cursorFound := false
		for i, arg := range args {
			if arg == "-c" {
				args[i+1] = cursor
				cursorFound = true
				break
			} else if strings.HasPrefix(arg, "--cursor=") {
				args[i] = fmt.Sprintf("--cursor=%s", cursor)
				cursorFound = true
				break
			}
			args[i] = fmt.Sprintf("'%s'", arg)
		}
		if !cursorFound {
			args = append(args, fmt.Sprintf("--cursor=%s", cursor))
		}
		color.New(color.Faint).Fprintf(
			ansi.NewAnsiStderr(), "\nMORE WITH:\n%s %s\n",
			cmd.CommandPath(), strings.Join(args, " "))
	}
}

// NewAPIClient returns a new utils.APIClient with the API key specified via
// command-line argument or config file.
func NewAPIClient() (*utils.APIClient, error) {
	apikey := viper.GetString("apikey")
	if apikey == "" {
		return nil, errors.New("An API key is needed. Either use the --apikey flag or run \"vt init\" to set up your API key")
	}
	return utils.NewAPIClient(apikey, fmt.Sprintf("vt-cli %s", Version))
}

// ObjectPrinter ...
type ObjectPrinter struct {
	client *utils.APIClient
	cmd    *cobra.Command
}

// NewObjectPrinter ...
func NewObjectPrinter(cmd *cobra.Command) (*ObjectPrinter, error) {
	client, err := NewAPIClient()
	if err != nil {
		return nil, err
	}
	return &ObjectPrinter{client: client, cmd: cmd}, nil
}

// Print ...
func (p *ObjectPrinter) Print(objType string, args []string, argRe *regexp.Regexp) error {

	var r utils.StringReader

	if len(args) == 1 && args[0] == "-" {
		r = utils.NewStringIOReader(os.Stdin)
	} else {
		r = utils.NewStringArrayReader(args)
	}

	if argRe != nil {
		r = utils.NewFilteredStringReader(r, argRe)
	}

	filteredArgs := make([]string, 0)
	for s, err := r.ReadString(); s != "" || err == nil; s, err = r.ReadString() {
		filteredArgs = append(filteredArgs, s)
	}

	objectsCh := make(chan *vt.Object)
	errorsCh := make(chan error, len(filteredArgs))

	go p.client.RetrieveObjects(objType, filteredArgs, objectsCh, errorsCh)

	objs := make([]*vt.Object, 0)

	for obj := range objectsCh {
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	for err := range errorsCh {
		fmt.Fprintln(os.Stderr, err)
	}

	return nil
}

func (p *ObjectPrinter) PrintCollection(collection *url.URL) error {
	it, err := p.client.Iterator(collection,
		vt.IteratorOptions{
			Limit:  viper.GetInt("limit"),
			Cursor: viper.GetString("cursor"),
			Filter: viper.GetString("filter"),
		})
	if err != nil {
		return err
	}
	return p.PrintIter(it)
}

func (p *ObjectPrinter) PrintIter(it *vt.Iterator) error {

	objs := make([]*vt.Object, 0)
	for it.Next() {
		obj := it.Get()
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if err := it.Error(); err != nil {
		return err
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	PrintCommandLineWithCursor(p.cmd, it)
	return nil
}

func (p *ObjectPrinter) PrintObject(obj *vt.Object) error {
	objs := make([]*vt.Object, 1)
	objs[0] = obj
	return p.PrintObjects(objs)
}

func (p *ObjectPrinter) PrintObjects(objs []*vt.Object) error {

	list := make([]map[string]interface{}, 0)

	for _, obj := range objs {
		m := obj.Attributes
		if viper.IsSet("include") && viper.IsSet("exclude") {
			m = utils.FilterMap(
				m, viper.GetStringSlice("include"), viper.GetStringSlice("exclude"))
		}
		for name, r := range obj.Relationships {
			if r.IsOneToOne && len(r.RelatedObjects) > 0 {
				m[name] = r.RelatedObjects[0].ID
			} else {
				l := make([]string, 0)
				for _, obj := range r.RelatedObjects {
					l = append(l, obj.ID)
				}
				m[name] = l
			}
		}
		key := fmt.Sprintf("%s <%s>", obj.Type, obj.ID)
		list = append(list, map[string]interface{}{key: m})
	}

	return yaml.NewColorEncoder(ansi.NewAnsiStdout(), colorScheme).Encode(list)
}
<commit_msg>Fix issues with generation of command-line with cursor for retrieving next set of results.<commit_after>// Copyright © 2017 The VirusTotal CLI authors. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"regexp"
	"strings"

	"github.com/spf13/cobra"

	"github.com/plusvic/go-ansi"

	"github.com/fatih/color"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"

	"github.com/VirusTotal/vt-cli/utils"
	"github.com/VirusTotal/vt-cli/yaml"
	"github.com/VirusTotal/vt-go/vt"
)

var colorScheme = yaml.Colors{
	KeyColor:     color.New(color.FgYellow),
	ValueColor:   color.New(color.FgHiGreen),
	CommentColor: color.New(color.Faint)}

func addAPIKeyFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"apikey", "k", "",
		"api key")
}

func addHostFlag(flags *pflag.FlagSet) {
	flags.String(
		"host", "www.virustotal.com",
		"API host name")
	flags.MarkHidden("host")
}

func addIncludeExcludeFlags(flags *pflag.FlagSet) {
	flags.StringSliceP(
		"include", "i", []string{"**"},
		"include fields matching the provided pattern")

	flags.StringSliceP(
		"exclude", "x", []string{},
		"exclude fields matching the provided pattern")
}

func addThreadsFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"threads", "t", 5,
		"number of threads working in parallel")
}

func addIDOnlyFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"identifiers-only", "I", false,
		"print identifiers only")
}

func addLimitFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"limit", "n", 10,
		"maximum number of results")
}

func addCursorFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"cursor", "c", "",
		"cursor")
}

func addOutputFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"output", "o", ".",
		"directory where downloaded files are put")
}

func addFilterFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"filter", "f", "",
		"filter")
}

func addVerboseFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"verbose", "v", false,
		"verbose output")
}

func addYAMLFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"yaml", "y", false,
		"output in YAML format")
}

// ReadFile reads the specified file and returns its content. If filename is "-"
// the data is read from stdin.
func ReadFile(filename string) ([]byte, error) {
	if filename == "-" {
		return ioutil.ReadAll(os.Stdin)
	}
	return ioutil.ReadFile(filename)
}

// PrintCommandLineWithCursor prints the same command-line that was used for
// executing the program but adding or replacing the --cursor flag with
// the current cursor for the given iterator.
func PrintCommandLineWithCursor(cmd *cobra.Command, it *vt.Iterator) {
	if cursor := it.Cursor(); cursor != "" {
		args := cmd.Flags().Args()
		for i, arg := range args {
			args[i] = fmt.Sprintf("'%s'", arg)
		}
		flags := make([]string, 0)
		cmd.Flags().Visit(func(flag *pflag.Flag) {
			if flag.Name != "cursor" {
				flags = append(flags, fmt.Sprintf("--%s=%v", flag.Name, flag.Value.String()))
			}
		})
		flags = append(flags, fmt.Sprintf("--cursor=%s", cursor))
		color.New(color.Faint).Fprintf(
			ansi.NewAnsiStderr(), "\nMORE WITH:\n%s %s %s\n",
			cmd.CommandPath(), strings.Join(args, " "), strings.Join(flags, " "))
	}
}

// NewAPIClient returns a new utils.APIClient with the API key specified via
// command-line argument or config file.
func NewAPIClient() (*utils.APIClient, error) {
	apikey := viper.GetString("apikey")
	if apikey == "" {
		return nil, errors.New("An API key is needed. Either use the --apikey flag or run \"vt init\" to set up your API key")
	}
	return utils.NewAPIClient(apikey, fmt.Sprintf("vt-cli %s", Version))
}

// ObjectPrinter ...
type ObjectPrinter struct {
	client *utils.APIClient
	cmd    *cobra.Command
}

// NewObjectPrinter ...
func NewObjectPrinter(cmd *cobra.Command) (*ObjectPrinter, error) {
	client, err := NewAPIClient()
	if err != nil {
		return nil, err
	}
	return &ObjectPrinter{client: client, cmd: cmd}, nil
}

// Print ...
func (p *ObjectPrinter) Print(objType string, args []string, argRe *regexp.Regexp) error {

	var r utils.StringReader

	if len(args) == 1 && args[0] == "-" {
		r = utils.NewStringIOReader(os.Stdin)
	} else {
		r = utils.NewStringArrayReader(args)
	}

	if argRe != nil {
		r = utils.NewFilteredStringReader(r, argRe)
	}

	filteredArgs := make([]string, 0)
	for s, err := r.ReadString(); s != "" || err == nil; s, err = r.ReadString() {
		filteredArgs = append(filteredArgs, s)
	}

	objectsCh := make(chan *vt.Object)
	errorsCh := make(chan error, len(filteredArgs))

	go p.client.RetrieveObjects(objType, filteredArgs, objectsCh, errorsCh)

	objs := make([]*vt.Object, 0)

	for obj := range objectsCh {
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	for err := range errorsCh {
		fmt.Fprintln(os.Stderr, err)
	}

	return nil
}

func (p *ObjectPrinter) PrintCollection(collection *url.URL) error {
	it, err := p.client.Iterator(collection,
		vt.IteratorOptions{
			Limit:  viper.GetInt("limit"),
			Cursor: viper.GetString("cursor"),
			Filter: viper.GetString("filter"),
		})
	if err != nil {
		return err
	}
	return p.PrintIter(it)
}

func (p *ObjectPrinter) PrintIter(it *vt.Iterator) error {

	objs := make([]*vt.Object, 0)
	for it.Next() {
		obj := it.Get()
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if err := it.Error(); err != nil {
		return err
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	PrintCommandLineWithCursor(p.cmd, it)
	return nil
}

func (p *ObjectPrinter) PrintObject(obj *vt.Object) error {
	objs := make([]*vt.Object, 1)
	objs[0] = obj
	return p.PrintObjects(objs)
}

func (p *ObjectPrinter) PrintObjects(objs []*vt.Object) error {

	list := make([]map[string]interface{}, 0)

	for _, obj := range objs {
		m := obj.Attributes
		if viper.IsSet("include") && viper.IsSet("exclude") {
			m = utils.FilterMap(
				m, viper.GetStringSlice("include"), viper.GetStringSlice("exclude"))
		}
		for name, r := range obj.Relationships {
			if r.IsOneToOne && len(r.RelatedObjects) > 0 {
				m[name] = r.RelatedObjects[0].ID
			} else {
				l := make([]string, 0)
				for _, obj := range r.RelatedObjects {
					l = append(l, obj.ID)
				}
				m[name] = l
			}
		}
		key := fmt.Sprintf("%s <%s>", obj.Type, obj.ID)
		list = append(list, map[string]interface{}{key: m})
	}

	return yaml.NewColorEncoder(ansi.NewAnsiStdout(), colorScheme).Encode(list)
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"syscall"

	"github.com/github/hub/ui"
	"github.com/github/hub/utils"
	"github.com/kballard/go-shellquote"
)

type Cmd struct {
	Name   string
	Args   []string
	Stdin  *os.File
	Stdout *os.File
	Stderr *os.File
}

func (cmd Cmd) String() string {
	return fmt.Sprintf("%s %s", cmd.Name, strings.Join(cmd.Args, " "))
}

func (cmd *Cmd) WithArg(arg string) *Cmd {
	cmd.Args = append(cmd.Args, arg)

	return cmd
}

func (cmd *Cmd) WithArgs(args ...string) *Cmd {
	for _, arg := range args {
		cmd.WithArg(arg)
	}

	return cmd
}

func (cmd *Cmd) CombinedOutput() (string, error) {
	verboseLog(cmd)
	output, err := exec.Command(cmd.Name, cmd.Args...).CombinedOutput()

	return string(output), err
}

func (cmd *Cmd) Success() bool {
	verboseLog(cmd)
	err := exec.Command(cmd.Name, cmd.Args...).Run()
	return err == nil
}

// Run runs command with `Exec` on platforms except Windows
// which only supports `Spawn`
func (cmd *Cmd) Run() error {
	if runtime.GOOS == "windows" {
		return cmd.Spawn()
	} else {
		return cmd.Exec()
	}
}

// Spawn runs command with spawn(3)
func (cmd *Cmd) Spawn() error {
	verboseLog(cmd)
	c := exec.Command(cmd.Name, cmd.Args...)
	c.Stdin = cmd.Stdin
	c.Stdout = cmd.Stdout
	c.Stderr = cmd.Stderr

	return c.Run()
}

// Exec runs command with exec(3)
// Note that Windows doesn't support exec(3): http://golang.org/src/pkg/syscall/exec_windows.go#L339
func (cmd *Cmd) Exec() error {
	verboseLog(cmd)

	binary, err := exec.LookPath(cmd.Name)
	if err != nil {
		return &exec.Error{
			Name: cmd.Name,
			Err:  fmt.Errorf("command not found"),
		}
	}

	args := []string{binary}
	args = append(args, cmd.Args...)

	return syscall.Exec(binary, args, os.Environ())
}

func New(cmd string) *Cmd {
	cmds, err := shellquote.Split(cmd)
	utils.Check(err)

	name := cmds[0]
	args := make([]string, len(cmds)-1)
	for _, arg := range cmds[1:] {
		args = append(args, arg)
	}
	return &Cmd{Name: name, Args: args, Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr}
}

func NewWithArray(cmd []string) *Cmd {
	return &Cmd{Name: cmd[0], Args: cmd[1:], Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr}
}

func verboseLog(cmd *Cmd) {
	if os.Getenv("HUB_VERBOSE") != "" {
		msg := fmt.Sprintf("$ %s %s", cmd.Name, strings.Join(cmd.Args, " "))
		if ui.IsTerminal(os.Stderr) {
			msg = fmt.Sprintf("\033[35m%s\033[0m", msg)
		}
		ui.Errorln(msg)
	}
}
<commit_msg>Use copy() to copy args values<commit_after>package cmd

import (
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"syscall"

	"github.com/github/hub/ui"
	"github.com/github/hub/utils"
	"github.com/kballard/go-shellquote"
)

type Cmd struct {
	Name   string
	Args   []string
	Stdin  *os.File
	Stdout *os.File
	Stderr *os.File
}

func (cmd Cmd) String() string {
	return fmt.Sprintf("%s %s", cmd.Name, strings.Join(cmd.Args, " "))
}

func (cmd *Cmd) WithArg(arg string) *Cmd {
	cmd.Args = append(cmd.Args, arg)

	return cmd
}

func (cmd *Cmd) WithArgs(args ...string) *Cmd {
	for _, arg := range args {
		cmd.WithArg(arg)
	}

	return cmd
}

func (cmd *Cmd) CombinedOutput() (string, error) {
	verboseLog(cmd)
	output, err := exec.Command(cmd.Name, cmd.Args...).CombinedOutput()

	return string(output), err
}

func (cmd *Cmd) Success() bool {
	verboseLog(cmd)
	err := exec.Command(cmd.Name, cmd.Args...).Run()
	return err == nil
}

// Run runs command with `Exec` on platforms except Windows
// which only supports `Spawn`
func (cmd *Cmd) Run() error {
	if runtime.GOOS == "windows" {
		return cmd.Spawn()
	} else {
		return cmd.Exec()
	}
}

// Spawn runs command with spawn(3)
func (cmd *Cmd) Spawn() error {
	verboseLog(cmd)
	c := exec.Command(cmd.Name, cmd.Args...)
	c.Stdin = cmd.Stdin
	c.Stdout = cmd.Stdout
	c.Stderr = cmd.Stderr

	return c.Run()
}

// Exec runs command with exec(3)
// Note that Windows doesn't support exec(3): http://golang.org/src/pkg/syscall/exec_windows.go#L339
func (cmd *Cmd) Exec() error {
	verboseLog(cmd)

	binary, err := exec.LookPath(cmd.Name)
	if err != nil {
		return &exec.Error{
			Name: cmd.Name,
			Err:  fmt.Errorf("command not found"),
		}
	}

	args := []string{binary}
	args = append(args, cmd.Args...)

	return syscall.Exec(binary, args, os.Environ())
}

func New(cmd string) *Cmd {
	cmds, err := shellquote.Split(cmd)
	utils.Check(err)

	name := cmds[0]
	args := make([]string, len(cmds)-1)
	copy(args, cmds[1:])

	return &Cmd{Name: name, Args: args, Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr}
}

func NewWithArray(cmd []string) *Cmd {
	return &Cmd{Name: cmd[0], Args: cmd[1:], Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr}
}

func verboseLog(cmd *Cmd) {
	if os.Getenv("HUB_VERBOSE") != "" {
		msg := fmt.Sprintf("$ %s %s", cmd.Name, strings.Join(cmd.Args, " "))
		if ui.IsTerminal(os.Stderr) {
			msg = fmt.Sprintf("\033[35m%s\033[0m", msg)
		}
		ui.Errorln(msg)
	}
}
<|endoftext|>


<commit_before>package crawler

import (
	"fmt"
	"github.com/AaronGoldman/ccfs/objects"
	"github.com/AaronGoldman/ccfs/services"
	"html/template"
	"log"
	"net/http"
	"net/url"
)

func webSearchHandler(w http.ResponseWriter, r *http.Request) {

	ResponseStruct := struct {
		Query                  string
		NameSegmentInfo        map[nameSegmentIndexEntry]int
		NameSegmentInfoPresent bool
		BlobInfo               blobIndexEntry
		BlobInfoPresent        bool
		CommitInfo             commitIndexEntry
		CommitInfoPresent      bool
		TagInfo                tagIndexEntry
		TagInfoPresent         bool
	}{}

	parsedURL, err := url.Parse(r.RequestURI)
	if err != nil {
		log.Println(err)
	}
	ResponseStruct.Query = parsedURL.Query().Get("q")
	log.Println(ResponseStruct.Query)
	log.Println(len(ResponseStruct.Query))
	hcid, err := objects.HcidFromHex(ResponseStruct.Query)
	if err == nil {
		ResponseStruct.BlobInfo, ResponseStruct.BlobInfoPresent =
			blobIndex[ResponseStruct.Query]
		if !ResponseStruct.BlobInfoPresent {
			log.Println("HID is not present")
		} else {
			switch ResponseStruct.BlobInfo.TypeString {
			case "list":
			case "commit":
				commit, err := services.GetCommitForHcid(hcid)
				if err == nil {
					ResponseStruct.CommitInfo,
						ResponseStruct.CommitInfoPresent =
						commitIndex[commit.Hkid.Hex()]
				}
			case "tag":
				tag, err := services.GetTagForHcid(hcid)
				if err == nil {
					ResponseStruct.TagInfo, ResponseStruct.TagInfoPresent =
						tagIndex[tag.Hkid.Hex()]
				}
			case "Repository":
				ResponseStruct.CommitInfo, ResponseStruct.CommitInfoPresent =
					commitIndex[ResponseStruct.Query]
			case "Domain":
				ResponseStruct.TagInfo, ResponseStruct.TagInfoPresent =
					tagIndex[ResponseStruct.Query]
			default:
				log.Printf("Unrecognized Type",
					ResponseStruct.BlobInfo.TypeString,
				)
			}
		}
		log.Println(err)
	}

	ResponseStruct.NameSegmentInfo,
		ResponseStruct.NameSegmentInfoPresent =
		nameSegmentIndex[ResponseStruct.Query]
	if !ResponseStruct.NameSegmentInfoPresent {
		log.Println("Name Segment is not present")
	}

	t, err := template.New("WebSearch template").Parse(`
	{{define "NameSegTemp"}}

			{{range $key, $value := .}}

 {{$key}}:

				{{range $key1:= $value}}

 {{$key1}}

				{{end}}
			{{end}}

	{{end}}
	{{define "VersionTemp"}}

			{{range $key, $value := .}}

					 {{$key}}:
					 {{$value}}

			{{end}}

	{{end}}

			<title>
				Search - CCFS
			</title>

			<form action = "./" method="get">
				<input type = "text" name = "q" value="{{.Query}}">
				<input type = "submit" value="Search">
			</form>

			{{with .Query}}
				Search results for: {{.}}
			{{end}}

				{{if .NameSegmentInfoPresent}}

						{{range $key, $value:= .NameSegmentInfo}}

								{{$key.TypeString}}:
								{{$key.Hash}}

						{{end}}

				{{end}}
				{{if .BlobInfoPresent}}

					{{.BlobInfo.TypeString}}[{{.BlobInfo.Size}}]:
					{{with .BlobInfo.Collection}}
					{{.}}:
					{{end}}
					{{template "NameSegTemp" .BlobInfo.NameSeg}}
					{{template "VersionTemp" .BlobInfo.Descendants}}
				{{end}}
				{{if .CommitInfoPresent}}
					{{template "NameSegTemp" .CommitInfo.NameSeg}}
					{{template "VersionTemp" .CommitInfo.Version}}
				{{end}}
				{{if .TagInfoPresent}}
					{{.TagInfo}}
				{{end}}

`)
	if err != nil {
		log.Println(err)
		http.Error(
			w,
			fmt.Sprintf("HTTP Error 500 Internal Search server error\n%s\n",
				err,
			),
			500,
		)
	} else {
		t.Execute(w, ResponseStruct)
	}
}
<commit_msg>Further Format of the search webpage (mainly Tags)<commit_after>package crawler

import (
	"fmt"
	"github.com/AaronGoldman/ccfs/objects"
	"github.com/AaronGoldman/ccfs/services"
	"html/template"
	"log"
	"net/http"
	"net/url"
)

func webSearchHandler(w http.ResponseWriter, r *http.Request) {

	ResponseStruct := struct {
		Query                  string
		NameSegmentInfo        map[nameSegmentIndexEntry]int
		NameSegmentInfoPresent bool
		BlobInfo               blobIndexEntry
		BlobInfoPresent        bool
		CommitInfo             commitIndexEntry
		CommitInfoPresent      bool
		TagInfo                tagIndexEntry
		TagInfoPresent         bool
	}{}

	parsedURL, err := url.Parse(r.RequestURI)
	if err != nil {
		log.Println(err)
	}
	ResponseStruct.Query = parsedURL.Query().Get("q")
	log.Println(ResponseStruct.Query)
	log.Println(len(ResponseStruct.Query))
	hcid, err := objects.HcidFromHex(ResponseStruct.Query)
	if err == nil {
		ResponseStruct.BlobInfo, ResponseStruct.BlobInfoPresent =
			blobIndex[ResponseStruct.Query]
		if !ResponseStruct.BlobInfoPresent {
			log.Println("HID is not present")
		} else {
			switch ResponseStruct.BlobInfo.TypeString {
			case "list":
			case "commit":
				commit, err := services.GetCommitForHcid(hcid)
				if err == nil {
					ResponseStruct.CommitInfo,
						ResponseStruct.CommitInfoPresent =
						commitIndex[commit.Hkid.Hex()]
				}
			case "tag":
				tag, err := services.GetTagForHcid(hcid)
				if err == nil {
					ResponseStruct.TagInfo, ResponseStruct.TagInfoPresent =
						tagIndex[tag.Hkid.Hex()]
				}
			case "Repository":
				ResponseStruct.CommitInfo, ResponseStruct.CommitInfoPresent =
					commitIndex[ResponseStruct.Query]
			case "Domain":
				ResponseStruct.TagInfo, ResponseStruct.TagInfoPresent =
					tagIndex[ResponseStruct.Query]
			default:
				log.Printf("Unrecognized Type",
					ResponseStruct.BlobInfo.TypeString,
				)
			}
		}
		log.Println(err)
	}

	ResponseStruct.NameSegmentInfo,
		ResponseStruct.NameSegmentInfoPresent =
		nameSegmentIndex[ResponseStruct.Query]
	if !ResponseStruct.NameSegmentInfoPresent {
		log.Println("Name Segment is not present")
	}

	t, err := template.New("WebSearch template").Parse(`
	{{define "NameSegTemp"}}

			{{range $key, $value := .}}

 {{$key}}:

				{{range $key1:= $value}}

 {{$key1}}

				{{end}}
			{{end}}

	{{end}}
	{{define "VersionTemp"}}

			{{range $key, $value := .}}

					 {{$key}}:
					 {{$value}}

			{{end}}

	{{end}}

			<title>
				Search - CCFS
			</title>

			<form action = "./" method="get">
				<input type = "text" name = "q" value="{{.Query}}">
				<input type = "submit" value="Search">
			</form>

			{{with .Query}}
				Search results for: {{.}}
			{{end}}

				{{if .NameSegmentInfoPresent}}

						{{range $key, $value:= .NameSegmentInfo}}

								{{$key.TypeString}}:
								{{$key.Hash}}

						{{end}}

				{{end}}
				{{if .BlobInfoPresent}}

					{{.BlobInfo.TypeString}}[{{.BlobInfo.Size}}]:
					{{with .BlobInfo.Collection}}
					{{.}}:
					{{end}}
					{{template "NameSegTemp" .BlobInfo.NameSeg}}
					{{template "VersionTemp" .BlobInfo.Descendants}}
				{{end}}
				{{if .CommitInfoPresent}}
					{{template "NameSegTemp" .CommitInfo.NameSeg}}
					{{template "VersionTemp" .CommitInfo.Version}}
				{{end}}
				{{if .TagInfoPresent}}

Aliases:

{{template "NameSegTemp" .TagInfo.NameSeg}}

Sub Domains:

								{{range $key, $value := .TagInfo.Version}}

{{$key}}:

{{template "VersionTemp" $value}}

								{{end}}

				{{end}}

`)
	if err != nil {
		log.Println(err)
		http.Error(
			w,
			fmt.Sprintf("HTTP Error 500 Internal Search server error\n%s\n",
				err,
			),
			500,
		)
	} else {
		t.Execute(w, ResponseStruct)
	}
}
<|endoftext|>