<commit_before>/*
This package reads pickled data written from python using the "pickle" module.
Protocols 0,1,2 are implemented. These are the versions written by the Python
2.x series. Python 3 defines newer protocol versions, but can write the older
protocol versions so they are readable by this package.

TLDR:

Read a pickled string or unicode object
	var somePickledData io.Reader
	mystring, err := stalecucumber.String(stalecucumber.Unpickle(somePickledData))

Read a pickled integer
	var somePickledData io.Reader
	myint64, err := stalecucumber.Int(stalecucumber.Unpickle(somePickledData))

Read a python dictionary into a structure
	var somePickledData io.Reader
	mystruct := struct{
		Apple int
		Banana uint
		Cat string
		Dog float32}{}

	err := stalecucumber.UnpackInto(&mystruct).From(stalecucumber.From(somePickledData))

Reading Data

The stalecucumber.Unpickle function takes a reader and attempts to read
a complete pickle program from it. This is normally the output of the function
like "pickle.dump" from Python.

The returned type is interface{} because unpickling can generate any type. Use
a helper function to convert to another type without an additional type check.

An error is returned if the underlying reader fails, the program
is invalid, or unsupported opcodes are encountered. See below for the details
of unsupported opcodes.

If the version of python you are using supports protocol version 1 or 2,
you should always specify that protocol version. By default the "pickle"
and "cPickle" modules in Python write using protocol 0. Protocol 0
requires much more space to represent the same values and is much
slower to parse.

Type Conversions

Types conversion Python types to Go types is performed as followed
	int -> int64
	string -> string
	unicode -> string
	float -> float64
	long -> big.Int from the "math/big" package
	lists -> []interface{}
	tuples -> []interface{}
	dict -> map[interface{}]interface{}

The following values are converted from Python to the Go types
	True & False -> bool
	None -> stalecucumber.PickleNone

Helper Functions

The following helper functions were inspired by the github.com/garyburd/redigo
package. Each function takes the result of Unpickle as its arguments. If unpickle
fails it does nothing and returns that error. Otherwise it attempts to
convert to the appropriate type. If type conversion fails it returns an error

	String - string from Python string or unicode
	Int - int64 from Python int or long
	Bool - bool from Python True or False
	Big - *big.Int from Python long
	ListOrTuple - []interface{} from Python Tuple or List
	Float - float64 from Python float
	Dict - map[interface{}]interface{} from Python dictionary
	DictString -
		map[string]interface{} from Python dictionary.
		Keys must all be of type unicode or string.

Unpacking into structures

If the pickled object is a python dictionary that has only unicode and string
objects for keys, that object can be unpickled into a struct in Go by using
the "UnpackInto" function. The "From" receiver on the return value accepts
the result of "Unpickle" as its actual parameters.

The keys of the python dictionary are assigned to fields in a structure. If
the first character of the key is not uppercase, it is uppercased. If a field
matching that name is found, the value in the python dictionary is unpacked
into the value of the field within the structure.

A nested python dictionary is unpickled into nested structures in Go. If a
field is of type map[interface{}]interface{} is of course unpacked into that
as well.

By default UnpackInto skips any missing fields and fails if a field's
type is not compatible with the object's type.

This behavior can be changed by setting "AllowMissingFields" and
"AllowMismatchedFields" on the return value of UnpackInto before calling
From.

Recursive objects

You can pickle recursive objects like so

	a = {}
	a["self"] = a
	pickle.dumps(a)

Python's pickler is intelligent enough not to emit an infinite data structure
when a recursive object is pickled.

I reccomend against pickling recursive objects in the first place, but this
library handles unpickling them without a problem. The result of unpickling
the above is map[interface{}]interface{} with a key "a" that contains
a reference to itself.

Attempting to unpack the result of the above python code into a structure
with UnpackInto would either fail or recurse forever.

Unsupported Opcodes

The pickle format is incredibly flexible and as a result has some
features that are impractical or unimportant when implementing a reader in
another language.

Each set of opcodes is listed below by protocol version with the impact.

Protocol 0

	GLOBAL

This opcode is equivalent to calling "import foo; foo.bar" in python. It is
generated whenever an object instance, class definition, or method definition
is serialized. As long as the pickled data does not contain an instance
of a python class or a reference to a python callable this opcode is not
emitted by the "pickle" module.

A few examples of what will definitely cause this opcode to be emitted

	pickle.dumps(range) #Pickling the range function
	pickle.dumps(Exception()) #Pickling an instance of a python class

This opcodes will be partially supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	REDUCE
	BUILD
	INST

These opcodes are used in recreating pickled python objects. That is currently
not supported by this package.

These opcodes will be supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	PERSID

This opcode is used to reference concrete definitions of objects between
a pickler and an unpickler by an ID#. The pickle protocol doesn't define
what a persistent ID means.

This opcode is unlikely to ever be supported by this package.

Protocol 1

	OBJ

This opcodes is used in recreating pickled python objects. That is currently
not supported by this package.

This opcode will supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	BINPERSID

This opcode is equivalent to PERSID in protocol 0 and won't be supported
for the same reason.

Protocol 2

	NEWOBJ

This opcodes is used in recreating pickled python objects. That is currently
not supported by this package.

This opcode will supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	EXT1
	EXT2
	EXT4

These opcodes allow using a registry
of popular objects that are pickled by name, typically classes.
It is envisioned that through a global negotiation and
registration process, third parties can set up a mapping between
ints and object names.

These opcodes are unlikely to ever be supported by this package.

*/
package stalecucumber

import "errors"
import "io"
import "bytes"
import "encoding/binary"
import "fmt"

var ErrOpcodeStopped = errors.New("STOP opcode found")
var ErrStackTooSmall = errors.New("Stack is too small to perform requested operation")
var ErrInputTruncated = errors.New("Input to the pickle machine was truncated")
var ErrOpcodeNotImplemented = errors.New("Input encountered opcode that is not implemented")
var ErrNoResult = errors.New("Input did not place a value onto the stack")
var ErrMarkNotFound = errors.New("Mark could not be found on the stack")

/*
Unpickle a value from a reader. This function returns an error if
the reader fails, the pickled data is invalid, or if the pickled data contains
an unsupported opcode.

See unsupported opcodes in the documentation of this package for more
information.

This function is generally not used directly, but with one of the helpers
such as string.
*/
func Unpickle(reader io.Reader) (interface{}, error) {
	var pm PickleMachine
	pm.buf = &bytes.Buffer{}
	pm.Reader = reader

	err := (&pm).execute()
	if err != nil {
		return nil, err
	}

	if len(pm.Stack) == 0 {
		return nil, ErrNoResult
	}

	return pm.Stack[0], nil

}

var jumpList = buildEmptyJumpList()

func init() {
	populateJumpList(&jumpList)
}

/*
This type is returned whenever Unpickle encounters an error in pickled data.
*/
type PickleMachineError struct {
	Err       error
	StackSize int
	MemoSize  int
	Opcode    uint8
}

/*
This struct is current exposed but not useful. It is likely to be hidden
in the near future.
*/
type PickleMachine struct {
	Stack  []interface{}
	Memo   []interface{}
	Reader io.Reader

	currentOpcode uint8
	buf           *bytes.Buffer
}

func (pme PickleMachineError) Error() string {
	return fmt.Sprintf("Pickle Machine failed on opcode:0x%x. Stack size:%d. Memo size:%d. Cause:%v", pme.Opcode, pme.StackSize, pme.MemoSize, pme.Err)
}

func (pm *PickleMachine) error(src error) error {
	return PickleMachineError{
		StackSize: len(pm.Stack),
		MemoSize:  len(pm.Memo),
		Err:       src,
		Opcode:    pm.currentOpcode,
	}
}

func (pm *PickleMachine) execute() error {
	for {
		var opcode uint8
		err := binary.Read(pm.Reader, binary.BigEndian, &opcode)
		if err != nil {
			return err
		}

		pm.currentOpcode = opcode
		err = jumpList[int(opcode)](pm)
		if err == ErrOpcodeStopped {
			return nil
		} else if err == ErrOpcodeNotImplemented {
			return pm.error(ErrOpcodeNotImplemented)
		} else if err != nil {
			return pm.error(err)
		}
	}
}

func (pm *PickleMachine) storeMemo(index int64, v interface{}) error {
	if index < 0 {
		return fmt.Errorf("Requested to write to invalid memo index:%v", index)
	}

	if int64(len(pm.Memo)) <= index {
		replacement := make([]interface{}, index+1)
		copy(replacement, pm.Memo)
		pm.Memo = replacement
	}

	pm.Memo[index] = v

	return nil
}

func (pm *PickleMachine) readFromMemo(index int64) (interface{}, error) {
	if index < 0 || index >= int64(len(pm.Memo)) {
		return nil, fmt.Errorf("Requested to read from invalid memo index %d", index)
	}

	return pm.Memo[index], nil
}

func (pm *PickleMachine) push(v interface{}) {
	pm.Stack = append(pm.Stack, v)
}

func (pm *PickleMachine) pop() (interface{}, error) {
	if len(pm.Stack) == 0 {
		return nil, ErrStackTooSmall
	}

	lastIndex := len(pm.Stack) - 1
	top := pm.Stack[lastIndex]

	pm.Stack = pm.Stack[:lastIndex]
	return top, nil
}

func (pm *PickleMachine) readFromStack(offset int) (interface{}, error) {
	return pm.readFromStackAt(len(pm.Stack) - 1 - offset)
}

func (pm *PickleMachine) readFromStackAt(position int) (interface{}, error) {

	if position < 0 {
		return nil, fmt.Errorf("Request to read from invalid stack position %d", position)
	}

	return pm.Stack[position], nil

}

func (pm *PickleMachine) readIntFromStack(offset int) (int64, error) {
	v, err := pm.readFromStack(offset)
	if err != nil {
		return 0, err
	}

	vi, ok := v.(int64)
	if !ok {
		return 0, fmt.Errorf("Type %T was requested from stack but found %v(%T)", vi, v, v)
	}

	return vi, nil
}

func (pm *PickleMachine) popAfterIndex(index int) error {
	if len(pm.Stack)-1 < index {
		return ErrStackTooSmall
	}

	pm.Stack = pm.Stack[0:index]
	return nil
}

func (pm *PickleMachine) putMemo(index int, v interface{}) {
	for len(pm.Memo) <= index {
		pm.Memo = append(pm.Memo, nil)
	}

	pm.Memo[index] = v
}

func (pm *PickleMachine) findMark() (int, error) {
	for i := len(pm.Stack) - 1; i != -1; i-- {
		if _, ok := pm.Stack[i].(PickleMark); ok {
			return i, nil
		}
	}
	return -1, ErrMarkNotFound
}

func (pm *PickleMachine) readFixedLengthRaw(l int64) ([]byte, error) {

	pm.buf.Reset()
	_, err := io.CopyN(pm.buf, pm.Reader, l)
	if err != nil {
		return nil, err
	}

	return pm.buf.Bytes(), nil
}

func (pm *PickleMachine) readFixedLengthString(l int64) (string, error) {

	//Avoid getting "<nil>"
	if l == 0 {
		return "", nil
	}

	pm.buf.Reset()
	_, err := io.CopyN(pm.buf, pm.Reader, l)
	if err != nil {
		return "", err
	}
	return pm.buf.String(), nil
}

func (pm *PickleMachine) readString() (string, error) {
	pm.buf.Reset()
	for {
		var v [1]byte
		n, err := pm.Reader.Read(v[:])
		if n != 1 {
			return "", ErrInputTruncated
		}
		if err != nil {
			return "", err
		}

		if v[0] == '\n' {
			break
		}
		pm.buf.WriteByte(v[0])
	}

	//Avoid getting "<nil>"
	if pm.buf.Len() == 0 {
		return "", nil
	}
	return pm.buf.String(), nil
}

func (pm *PickleMachine) readBinaryInto(dst interface{}, bigEndian bool) error {
	var bo binary.ByteOrder
	if bigEndian {
		bo = binary.BigEndian
	} else {
		bo = binary.LittleEndian
	}
	return binary.Read(pm.Reader, bo, dst)
}
<commit_msg>correct docs<commit_after>/*
This package reads pickled data written from python using the "pickle" module.
Protocols 0,1,2 are implemented. These are the versions written by the Python
2.x series. Python 3 defines newer protocol versions, but can write the older
protocol versions so they are readable by this package.

TLDR:

Read a pickled string or unicode object
	var somePickledData io.Reader
	mystring, err := stalecucumber.String(stalecucumber.Unpickle(somePickledData))

Read a pickled integer
	var somePickledData io.Reader
	myint64, err := stalecucumber.Int(stalecucumber.Unpickle(somePickledData))

Read a python dictionary into a structure
	var somePickledData io.Reader
	mystruct := struct{
		Apple int
		Banana uint
		Cat string
		Dog float32}{}

	err := stalecucumber.UnpackInto(&mystruct).From(stalecucumber.From(somePickledData))

Reading Data

The stalecucumber.Unpickle function takes a reader and attempts to read
a complete pickle program from it. This is normally the output of the function
like "pickle.dump" from Python.

The returned type is interface{} because unpickling can generate any type. Use
a helper function to convert to another type without an additional type check.

An error is returned if the underlying reader fails, the program
is invalid, or unsupported opcodes are encountered. See below for the details
of unsupported opcodes.

If the version of python you are using supports protocol version 1 or 2,
you should always specify that protocol version. By default the "pickle"
and "cPickle" modules in Python write using protocol 0. Protocol 0
requires much more space to represent the same values and is much
slower to parse.

Type Conversions

Types conversion Python types to Go types is performed as followed
	int -> int64
	string -> string
	unicode -> string
	float -> float64
	long -> big.Int from the "math/big" package
	lists -> []interface{}
	tuples -> []interface{}
	dict -> map[interface{}]interface{}

The following values are converted from Python to the Go types
	True & False -> bool
	None -> stalecucumber.PickleNone

Helper Functions

The following helper functions were inspired by the github.com/garyburd/redigo
package. Each function takes the result of Unpickle as its arguments. If unpickle
fails it does nothing and returns that error. Otherwise it attempts to
convert to the appropriate type. If type conversion fails it returns an error

	String - string from Python string or unicode
	Int - int64 from Python int or long
	Bool - bool from Python True or False
	Big - *big.Int from Python long
	ListOrTuple - []interface{} from Python Tuple or List
	Float - float64 from Python float
	Dict - map[interface{}]interface{} from Python dictionary
	DictString -
		map[string]interface{} from Python dictionary.
		Keys must all be of type unicode or string.

Unpacking into structures

If the pickled object is a python dictionary that has only unicode and string
objects for keys, that object can be unpickled into a struct in Go by using
the "UnpackInto" function. The "From" receiver on the return value accepts
the result of "Unpickle" as its actual parameters.

The keys of the python dictionary are assigned to fields in a structure. If
the first character of the key is not uppercase, it is uppercased. If a field
matching that name is found, the value in the python dictionary is unpacked
into the value of the field within the structure.

A nested python dictionary is unpickled into nested structures in Go. If a
field is of type map[interface{}]interface{} is of course unpacked into that
as well.

By default UnpackInto skips any missing fields and fails if a field's
type is not compatible with the object's type.

This behavior can be changed by setting "AllowMissingFields" and
"AllowMismatchedFields" on the return value of UnpackInto before calling
From.

Recursive objects

You can pickle recursive objects like so

	a = {}
	a["self"] = a
	pickle.dumps(a)

Python's pickler is intelligent enough not to emit an infinite data structure
when a recursive object is pickled.

I reccomend against pickling recursive objects in the first place, but this
library handles unpickling them without a problem. The result of unpickling
the above is map[interface{}]interface{} with a key "a" that contains
a reference to itself.

Attempting to unpack the result of the above python code into a structure
with UnpackInto would either fail or recurse forever.

Unsupported Opcodes

The pickle format is incredibly flexible and as a result has some
features that are impractical or unimportant when implementing a reader in
another language.

Each set of opcodes is listed below by protocol version with the impact.

Protocol 0

	GLOBAL

This opcode is equivalent to calling "import foo; foo.bar" in python. It is
generated whenever an object instance, class definition, or method definition
is serialized. As long as the pickled data does not contain an instance
of a python class or a reference to a python callable this opcode is not
emitted by the "pickle" module.

A few examples of what will definitely cause this opcode to be emitted

	pickle.dumps(range) #Pickling the range function
	pickle.dumps(Exception()) #Pickling an instance of a python class

This opcode will be partially supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	REDUCE
	BUILD
	INST

These opcodes are used in recreating pickled python objects. That is currently
not supported by this package.

These opcodes will be supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	PERSID

This opcode is used to reference concrete definitions of objects between
a pickler and an unpickler by an ID number. The pickle protocol doesn't define
what a persistent ID means.

This opcode is unlikely to ever be supported by this package.

Protocol 1

	OBJ

This opcodes is used in recreating pickled python objects. That is currently
not supported by this package.

This opcode will supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	BINPERSID

This opcode is equivalent to PERSID in protocol 0 and won't be supported
for the same reason.

Protocol 2

	NEWOBJ

This opcodes is used in recreating pickled python objects. That is currently
not supported by this package.

This opcode will supported in a future revision to this package
that allows the unpickling of instances of Python classes.

	EXT1
	EXT2
	EXT4

These opcodes allow using a registry
of popular objects that are pickled by name, typically classes.
It is envisioned that through a global negotiation and
registration process, third parties can set up a mapping between
ints and object names.

These opcodes are unlikely to ever be supported by this package.

*/
package stalecucumber

import "errors"
import "io"
import "bytes"
import "encoding/binary"
import "fmt"

var ErrOpcodeStopped = errors.New("STOP opcode found")
var ErrStackTooSmall = errors.New("Stack is too small to perform requested operation")
var ErrInputTruncated = errors.New("Input to the pickle machine was truncated")
var ErrOpcodeNotImplemented = errors.New("Input encountered opcode that is not implemented")
var ErrNoResult = errors.New("Input did not place a value onto the stack")
var ErrMarkNotFound = errors.New("Mark could not be found on the stack")

/*
Unpickle a value from a reader. This function returns an error if
the reader fails, the pickled data is invalid, or if the pickled data contains
an unsupported opcode.

See unsupported opcodes in the documentation of this package for more
information.

This function is generally not used directly, but with one of the helpers
such as string.
*/
func Unpickle(reader io.Reader) (interface{}, error) {
	var pm PickleMachine
	pm.buf = &bytes.Buffer{}
	pm.Reader = reader

	err := (&pm).execute()
	if err != nil {
		return nil, err
	}

	if len(pm.Stack) == 0 {
		return nil, ErrNoResult
	}

	return pm.Stack[0], nil

}

var jumpList = buildEmptyJumpList()

func init() {
	populateJumpList(&jumpList)
}

/*
This type is returned whenever Unpickle encounters an error in pickled data.
*/
type PickleMachineError struct {
	Err       error
	StackSize int
	MemoSize  int
	Opcode    uint8
}

/*
This struct is current exposed but not useful. It is likely to be hidden
in the near future.
*/
type PickleMachine struct {
	Stack  []interface{}
	Memo   []interface{}
	Reader io.Reader

	currentOpcode uint8
	buf           *bytes.Buffer
}

func (pme PickleMachineError) Error() string {
	return fmt.Sprintf("Pickle Machine failed on opcode:0x%x. Stack size:%d. Memo size:%d. Cause:%v", pme.Opcode, pme.StackSize, pme.MemoSize, pme.Err)
}

func (pm *PickleMachine) error(src error) error {
	return PickleMachineError{
		StackSize: len(pm.Stack),
		MemoSize:  len(pm.Memo),
		Err:       src,
		Opcode:    pm.currentOpcode,
	}
}

func (pm *PickleMachine) execute() error {
	for {
		var opcode uint8
		err := binary.Read(pm.Reader, binary.BigEndian, &opcode)
		if err != nil {
			return err
		}

		pm.currentOpcode = opcode
		err = jumpList[int(opcode)](pm)
		if err == ErrOpcodeStopped {
			return nil
		} else if err == ErrOpcodeNotImplemented {
			return pm.error(ErrOpcodeNotImplemented)
		} else if err != nil {
			return pm.error(err)
		}
	}
}

func (pm *PickleMachine) storeMemo(index int64, v interface{}) error {
	if index < 0 {
		return fmt.Errorf("Requested to write to invalid memo index:%v", index)
	}

	if int64(len(pm.Memo)) <= index {
		replacement := make([]interface{}, index+1)
		copy(replacement, pm.Memo)
		pm.Memo = replacement
	}

	pm.Memo[index] = v

	return nil
}

func (pm *PickleMachine) readFromMemo(index int64) (interface{}, error) {
	if index < 0 || index >= int64(len(pm.Memo)) {
		return nil, fmt.Errorf("Requested to read from invalid memo index %d", index)
	}

	return pm.Memo[index], nil
}

func (pm *PickleMachine) push(v interface{}) {
	pm.Stack = append(pm.Stack, v)
}

func (pm *PickleMachine) pop() (interface{}, error) {
	if len(pm.Stack) == 0 {
		return nil, ErrStackTooSmall
	}

	lastIndex := len(pm.Stack) - 1
	top := pm.Stack[lastIndex]

	pm.Stack = pm.Stack[:lastIndex]
	return top, nil
}

func (pm *PickleMachine) readFromStack(offset int) (interface{}, error) {
	return pm.readFromStackAt(len(pm.Stack) - 1 - offset)
}

func (pm *PickleMachine) readFromStackAt(position int) (interface{}, error) {

	if position < 0 {
		return nil, fmt.Errorf("Request to read from invalid stack position %d", position)
	}

	return pm.Stack[position], nil

}

func (pm *PickleMachine) readIntFromStack(offset int) (int64, error) {
	v, err := pm.readFromStack(offset)
	if err != nil {
		return 0, err
	}

	vi, ok := v.(int64)
	if !ok {
		return 0, fmt.Errorf("Type %T was requested from stack but found %v(%T)", vi, v, v)
	}

	return vi, nil
}

func (pm *PickleMachine) popAfterIndex(index int) error {
	if len(pm.Stack)-1 < index {
		return ErrStackTooSmall
	}

	pm.Stack = pm.Stack[0:index]
	return nil
}

func (pm *PickleMachine) putMemo(index int, v interface{}) {
	for len(pm.Memo) <= index {
		pm.Memo = append(pm.Memo, nil)
	}

	pm.Memo[index] = v
}

func (pm *PickleMachine) findMark() (int, error) {
	for i := len(pm.Stack) - 1; i != -1; i-- {
		if _, ok := pm.Stack[i].(PickleMark); ok {
			return i, nil
		}
	}
	return -1, ErrMarkNotFound
}

func (pm *PickleMachine) readFixedLengthRaw(l int64) ([]byte, error) {

	pm.buf.Reset()
	_, err := io.CopyN(pm.buf, pm.Reader, l)
	if err != nil {
		return nil, err
	}

	return pm.buf.Bytes(), nil
}

func (pm *PickleMachine) readFixedLengthString(l int64) (string, error) {

	//Avoid getting "<nil>"
	if l == 0 {
		return "", nil
	}

	pm.buf.Reset()
	_, err := io.CopyN(pm.buf, pm.Reader, l)
	if err != nil {
		return "", err
	}
	return pm.buf.String(), nil
}

func (pm *PickleMachine) readString() (string, error) {
	pm.buf.Reset()
	for {
		var v [1]byte
		n, err := pm.Reader.Read(v[:])
		if n != 1 {
			return "", ErrInputTruncated
		}
		if err != nil {
			return "", err
		}

		if v[0] == '\n' {
			break
		}
		pm.buf.WriteByte(v[0])
	}

	//Avoid getting "<nil>"
	if pm.buf.Len() == 0 {
		return "", nil
	}
	return pm.buf.String(), nil
}

func (pm *PickleMachine) readBinaryInto(dst interface{}, bigEndian bool) error {
	var bo binary.ByteOrder
	if bigEndian {
		bo = binary.BigEndian
	} else {
		bo = binary.LittleEndian
	}
	return binary.Read(pm.Reader, bo, dst)
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go-Commander Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Based on the original work by The Go Authors:
// Copyright 2011 The Go Authors.  All rights reserved.

// commander helps creating command line programs whose arguments are flags,
// commands and subcommands.
package commander

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sort"
	"strings"
	"text/template"

	"github.com/gonuts/flag"
)

// UsageSection differentiates between sections in the usage text.
type Listing int

const (
	CommandsList = iota
	HelpTopicsList
	Unlisted
)

// A Command is an implementation of a subcommand.
type Command struct {

	// UsageLine is the short usage message.
	// The first word in the line is taken to be the command name.
	UsageLine string

	// Short is the short description line shown in command lists.
	Short string

	// Long is the long description shown in the 'help <this-command>' output.
	Long string

	// List reports which list to show this command in Usage and Help.
	// Choose between {CommandsList (default), HelpTopicsList, Unlisted}
	List Listing

	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string) error

	// Flag is a set of flags specific to this command.
	Flag flag.FlagSet

	// CustomFlags indicates that the command will do its own
	// flag parsing.
	CustomFlags bool

	// Subcommands are dispatched from this command
	Subcommands []*Command

	// Parent command, nil for root.
	Parent *Command

	// UsageTemplate formats the usage (short) information displayed to the user
	// (leave empty for default)
	UsageTemplate string

	// HelpTemplate formats the help (long) information displayed to the user
	// (leave empty for default)
	HelpTemplate string

	// Stdout and Stderr by default are os.Stdout and os.Stderr, but you can
	// point them at any io.Writer
	Stdout io.Writer
	Stderr io.Writer
}

// Name returns the command's name: the first word in the usage line.
func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Usage prints the usage details to the standard error output.
func (c *Command) Usage() {
	c.usage()
}

// FlagOptions returns the flag's options as a string
func (c *Command) FlagOptions() string {
	var buf bytes.Buffer
	c.Flag.SetOutput(&buf)
	c.Flag.PrintDefaults()

	str := string(buf.Bytes())
	if len(str) > 0 {
		return fmt.Sprintf("\nOptions:\n%s", str)
	}
	return ""
}

// Runnable reports whether the command can be run; otherwise
// it is a documentation pseudo-command such as importpath.
func (c *Command) Runnable() bool {
	return c.Run != nil
}

// Type to allow us to use sort.Sort on a slice of Commands
type CommandSlice []*Command

func (c CommandSlice) Len() int {
	return len(c)
}

func (c CommandSlice) Less(i, j int) bool {
	return c[i].Name() < c[j].Name()
}

func (c CommandSlice) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}

// Sort the commands
func (c *Command) SortCommands() {
	sort.Sort(CommandSlice(c.Subcommands))
}

// Init the command
func (c *Command) init() {
	if c.Parent != nil {
		return // already initialized.
	}

	// setup strings
	if len(c.UsageLine) < 1 {
		c.UsageLine = Defaults.UsageLine
	}
	if len(c.UsageTemplate) < 1 {
		c.UsageTemplate = Defaults.UsageTemplate
	}
	if len(c.HelpTemplate) < 1 {
		c.HelpTemplate = Defaults.HelpTemplate
	}

	if c.Stderr == nil {
		c.Stderr = os.Stderr
	}
	if c.Stdout == nil {
		c.Stdout = os.Stdout
	}

	// init subcommands
	for _, cmd := range c.Subcommands {
		cmd.init()
	}

	// init hierarchy...
	for _, cmd := range c.Subcommands {
		cmd.Parent = c
	}
}

// Dispatch executes the command using the provided arguments.
// If a subcommand exists matching the first argument, it is dispatched.
// Otherwise, the command's Run function is called.
func (c *Command) Dispatch(args []string) error {
	if c == nil {
		return fmt.Errorf("Called Run() on a nil Command")
	}

	// Ensure command is initialized.
	c.init()

	// First, try a sub-command
	if len(args) > 0 {
		for _, cmd := range c.Subcommands {
			n := cmd.Name()
			if n == args[0] {
				return cmd.Dispatch(args[1:])
			}
		}

		// help is builtin (but after, to allow overriding)
		if args[0] == "help" {
			return c.help(args[1:])
		}

		// then, try out an external binary (git-style)
		bin, err := exec.LookPath(c.FullName() + "-" + args[0])
		if err == nil {
			cmd := exec.Command(bin, args[1:]...)
			cmd.Stdin = os.Stdin
			cmd.Stdout = c.Stdout
			cmd.Stderr = c.Stderr
			return cmd.Run()
		}
	}

	// then, try running this command
	if c.Runnable() {
		if !c.CustomFlags {
			var err = error(nil)
			c.Flag.Usage = func() {
				c.Usage()
				err = fmt.Errorf("Failed to parse flags.")
			}
			c.Flag.Parse(args)
			if err != nil {
				return err
			}
			args = c.Flag.Args()
		}
		return c.Run(c, args)
	}

	// TODO: try an alias
	//...

	// Last, print usage
	if err := c.usage(); err != nil {
		return err
	}
	return nil
}

func (c *Command) usage() error {
	// c.SortCommands()
	err := tmpl(c.Stderr, c.UsageTemplate, c)
	if err != nil {
		fmt.Println(err)
	}
	return err
}

// help implements the 'help' command.
func (c *Command) help(args []string) error {

	// help exactly for this command?
	if len(args) == 0 {
		if len(c.Long) > 0 {
			return tmpl(c.Stdout, c.HelpTemplate, c)
		} else {
			return c.usage()
		}
	}

	arg := args[0]

	// is this help for a subcommand?
	for _, cmd := range c.Subcommands {
		n := cmd.Name()
		// strip out "<parent>-"" name
		if strings.HasPrefix(n, c.Name()+"-") {
			n = n[len(c.Name()+"-"):]
		}
		if n == arg {
			return cmd.help(args[1:])
		}
	}

	return fmt.Errorf("Unknown help topic %#q.  Run '%v help'.\n", arg, c.Name())
}

func (c *Command) MaxLen() (res int) {
	res = 0
	for _, cmd := range c.Subcommands {
		i := len(cmd.Name())
		if i > res {
			res = i
		}
	}
	return
}

// ColFormat returns the column header size format for printing in the template
func (c *Command) ColFormat() string {
	sz := c.MaxLen()
	if sz < 11 {
		sz = 11
	}
	return fmt.Sprintf("%%-%ds", sz)
}

// FullName returns the full name of the command, prefixed with parent commands
func (c *Command) FullName() string {
	n := c.Name()
	if c.Parent != nil {
		n = c.Parent.FullName() + "-" + n
	}
	return n
}

// FullSpacedName returns the full name of the command, with ' ' instead of '-'
func (c *Command) FullSpacedName() string {
	n := c.Name()
	if c.Parent != nil {
		n = c.Parent.FullSpacedName() + " " + n
	}
	return n
}

func (c *Command) SubcommandList(list Listing) []*Command {
	var cmds []*Command
	for _, cmd := range c.Subcommands {
		if cmd.List == list {
			cmds = append(cmds, cmd)
		}
	}
	return cmds
}

var Defaults = Command{
	UsageTemplate: `{{if .Runnable}}Usage: {{if .Parent}}{{.Parent.FullSpacedName}}{{end}} {{.UsageLine}}

{{end}}{{.FullSpacedName}} - {{.Short}}

{{if commandList}}Commands:
{{range commandList}}
    {{.Name | printf (colfmt)}} {{.Short}}{{end}}

Use "{{.Name}} help <command>" for more information about a command.

{{end}}{{.FlagOptions}}{{if helpList}}
Additional help topics:
{{range helpList}}
    {{.Name | printf (colfmt)}} {{.Short}}{{end}}

Use "{{.Name}} help <topic>" for more information about that topic.

{{end}}`,

	HelpTemplate: `{{if .Runnable}}Usage: {{if .Parent}}{{.Parent.FullSpacedName}}{{end}} {{.UsageLine}}

{{end}}{{.Long | trim}}
{{.FlagOptions}}
`,
}

// tmpl executes the given template text on data, writing the result to w.
func tmpl(w io.Writer, text string, data interface{}) error {
	t := template.New("top")
	t.Funcs(template.FuncMap{
		"trim":        strings.TrimSpace,
		"colfmt":      func() string { return data.(*Command).ColFormat() },
		"commandList": func() []*Command { return data.(*Command).SubcommandList(CommandsList) },
		"helpList":    func() []*Command { return data.(*Command).SubcommandList(HelpTopicsList) },
	})
	template.Must(t.Parse(text))
	return t.Execute(w, data)
}
<commit_msg>different tmpl<commit_after>// Copyright 2012 The Go-Commander Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
//
// Based on the original work by The Go Authors:
// Copyright 2011 The Go Authors.  All rights reserved.

// commander helps creating command line programs whose arguments are flags,
// commands and subcommands.
package commander

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sort"
	"strings"
	"text/template"

	"github.com/gonuts/flag"
)

// UsageSection differentiates between sections in the usage text.
type Listing int

const (
	CommandsList = iota
	HelpTopicsList
	Unlisted
)

// A Command is an implementation of a subcommand.
type Command struct {

	// UsageLine is the short usage message.
	// The first word in the line is taken to be the command name.
	UsageLine string

	// Short is the short description line shown in command lists.
	Short string

	// Long is the long description shown in the 'help <this-command>' output.
	Long string

	// List reports which list to show this command in Usage and Help.
	// Choose between {CommandsList (default), HelpTopicsList, Unlisted}
	List Listing

	// Run runs the command.
	// The args are the arguments after the command name.
	Run func(cmd *Command, args []string) error

	// Flag is a set of flags specific to this command.
	Flag flag.FlagSet

	// CustomFlags indicates that the command will do its own
	// flag parsing.
	CustomFlags bool

	// Subcommands are dispatched from this command
	Subcommands []*Command

	// Parent command, nil for root.
	Parent *Command

	// UsageTemplate formats the usage (short) information displayed to the user
	// (leave empty for default)
	UsageTemplate string

	// HelpTemplate formats the help (long) information displayed to the user
	// (leave empty for default)
	HelpTemplate string

	// Stdout and Stderr by default are os.Stdout and os.Stderr, but you can
	// point them at any io.Writer
	Stdout io.Writer
	Stderr io.Writer
}

// Name returns the command's name: the first word in the usage line.
func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Usage prints the usage details to the standard error output.
func (c *Command) Usage() {
	c.usage()
}

// FlagOptions returns the flag's options as a string
func (c *Command) FlagOptions() string {
	var buf bytes.Buffer
	c.Flag.SetOutput(&buf)
	c.Flag.PrintDefaults()

	str := string(buf.Bytes())
	if len(str) > 0 {
		return fmt.Sprintf("\nOptions:\n%s", str)
	}
	return ""
}

// Runnable reports whether the command can be run; otherwise
// it is a documentation pseudo-command such as importpath.
func (c *Command) Runnable() bool {
	return c.Run != nil
}

// Type to allow us to use sort.Sort on a slice of Commands
type CommandSlice []*Command

func (c CommandSlice) Len() int {
	return len(c)
}

func (c CommandSlice) Less(i, j int) bool {
	return c[i].Name() < c[j].Name()
}

func (c CommandSlice) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}

// Sort the commands
func (c *Command) SortCommands() {
	sort.Sort(CommandSlice(c.Subcommands))
}

// Init the command
func (c *Command) init() {
	if c.Parent != nil {
		return // already initialized.
	}

	// setup strings
	if len(c.UsageLine) < 1 {
		c.UsageLine = Defaults.UsageLine
	}
	if len(c.UsageTemplate) < 1 {
		c.UsageTemplate = Defaults.UsageTemplate
	}
	if len(c.HelpTemplate) < 1 {
		c.HelpTemplate = Defaults.HelpTemplate
	}

	if c.Stderr == nil {
		c.Stderr = os.Stderr
	}
	if c.Stdout == nil {
		c.Stdout = os.Stdout
	}

	// init subcommands
	for _, cmd := range c.Subcommands {
		cmd.init()
	}

	// init hierarchy...
	for _, cmd := range c.Subcommands {
		cmd.Parent = c
	}
}

// Dispatch executes the command using the provided arguments.
// If a subcommand exists matching the first argument, it is dispatched.
// Otherwise, the command's Run function is called.
func (c *Command) Dispatch(args []string) error {
	if c == nil {
		return fmt.Errorf("Called Run() on a nil Command")
	}

	// Ensure command is initialized.
	c.init()

	// First, try a sub-command
	if len(args) > 0 {
		for _, cmd := range c.Subcommands {
			n := cmd.Name()
			if n == args[0] {
				return cmd.Dispatch(args[1:])
			}
		}

		// help is builtin (but after, to allow overriding)
		if args[0] == "help" {
			return c.help(args[1:])
		}

		// then, try out an external binary (git-style)
		bin, err := exec.LookPath(c.FullName() + "-" + args[0])
		if err == nil {
			cmd := exec.Command(bin, args[1:]...)
			cmd.Stdin = os.Stdin
			cmd.Stdout = c.Stdout
			cmd.Stderr = c.Stderr
			return cmd.Run()
		}
	}

	// then, try running this command
	if c.Runnable() {
		if !c.CustomFlags {
			var err = error(nil)
			c.Flag.Usage = func() {
				c.Usage()
				err = fmt.Errorf("Failed to parse flags.")
			}
			c.Flag.Parse(args)
			if err != nil {
				return err
			}
			args = c.Flag.Args()
		}
		return c.Run(c, args)
	}

	// TODO: try an alias
	//...

	// Last, print usage
	if err := c.usage(); err != nil {
		return err
	}
	return nil
}

func (c *Command) usage() error {
	// c.SortCommands()
	err := tmpl(c.Stderr, c.UsageTemplate, c)
	if err != nil {
		fmt.Println(err)
	}
	return err
}

// help implements the 'help' command.
func (c *Command) help(args []string) error {

	// help exactly for this command?
	if len(args) == 0 {
		if len(c.Long) > 0 {
			return tmpl(c.Stdout, c.HelpTemplate, c)
		} else {
			return c.usage()
		}
	}

	arg := args[0]

	// is this help for a subcommand?
	for _, cmd := range c.Subcommands {
		n := cmd.Name()
		// strip out "<parent>-"" name
		if strings.HasPrefix(n, c.Name()+"-") {
			n = n[len(c.Name()+"-"):]
		}
		if n == arg {
			return cmd.help(args[1:])
		}
	}

	return fmt.Errorf("Unknown help topic %#q.  Run '%v help'.\n", arg, c.Name())
}

func (c *Command) MaxLen() (res int) {
	res = 0
	for _, cmd := range c.Subcommands {
		i := len(cmd.Name())
		if i > res {
			res = i
		}
	}
	return
}

// ColFormat returns the column header size format for printing in the template
func (c *Command) ColFormat() string {
	sz := c.MaxLen()
	if sz < 11 {
		sz = 11
	}
	return fmt.Sprintf("%%-%ds", sz)
}

// FullName returns the full name of the command, prefixed with parent commands
func (c *Command) FullName() string {
	n := c.Name()
	if c.Parent != nil {
		n = c.Parent.FullName() + "-" + n
	}
	return n
}

// FullSpacedName returns the full name of the command, with ' ' instead of '-'
func (c *Command) FullSpacedName() string {
	n := c.Name()
	if c.Parent != nil {
		n = c.Parent.FullSpacedName() + " " + n
	}
	return n
}

func (c *Command) SubcommandList(list Listing) []*Command {
	var cmds []*Command
	for _, cmd := range c.Subcommands {
		if cmd.List == list {
			cmds = append(cmds, cmd)
		}
	}
	return cmds
}

var Defaults = Command{
	UsageTemplate: `{{if .Runnable}}Usage: {{if .Parent}}{{.Parent.FullSpacedName}}{{end}} {{.UsageLine}}

{{else}}{{.FullSpacedName}} - {{end}}{{.Short}}

{{if commandList}}Commands:
{{range commandList}}
    {{.Name | printf (colfmt)}} {{.Short}}{{end}}

Use "{{.Name}} help <command>" for more information about a command.

{{end}}{{.FlagOptions}}{{if helpList}}
Additional help topics:
{{range helpList}}
    {{.Name | printf (colfmt)}} {{.Short}}{{end}}

Use "{{.Name}} help <topic>" for more information about that topic.

{{end}}`,

	HelpTemplate: `{{if .Runnable}}Usage: {{if .Parent}}{{.Parent.FullSpacedName}}{{end}} {{.UsageLine}}

{{end}}{{.Long | trim}}
{{.FlagOptions}}
`,
}

// tmpl executes the given template text on data, writing the result to w.
func tmpl(w io.Writer, text string, data interface{}) error {
	t := template.New("top")
	t.Funcs(template.FuncMap{
		"trim":        strings.TrimSpace,
		"colfmt":      func() string { return data.(*Command).ColFormat() },
		"commandList": func() []*Command { return data.(*Command).SubcommandList(CommandsList) },
		"helpList":    func() []*Command { return data.(*Command).SubcommandList(HelpTopicsList) },
	})
	template.Must(t.Parse(text))
	return t.Execute(w, data)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"sort"
	"strings"
	"sync"

	"github.com/msoap/raphanus"
)

// Ctx - context for bot command function (users, command, args, ...)
type Ctx struct {
	appConfig      *Config           // configuration
	users          *Users            // all users
	commands       Commands          // all chat commands
	userID         int               // current user
	allowExec      bool              // is user authorized
	messageCmd     string            // command name
	messageArgs    string            // command arguments
	messageSignal  chan<- BotMessage // for send telegram messages
	chatID         int               // chat for send replay
	exitSignal     chan<- struct{}   // for signal for terminate bot
	cache          *raphanus.DB      // cache for commands output
	cacheTTL       int               // cache timeout
	oneThreadMutex *sync.Mutex       // mutex for run shell commands in one thread
}

// /auth and /authroot - authorize users
func cmdAuth(ctx Ctx) (replayMsg string) {
	forRoot := ctx.messageCmd == "/authroot"

	if ctx.messageArgs == "" {

		replayMsg = "See code in terminal with shell2telegram or ask code from root user and type:\n" + ctx.messageCmd + " code"
		authCode := ctx.users.DoLogin(ctx.userID, forRoot)

		rootRoleStr := ""
		if forRoot {
			rootRoleStr = "root "
		}
		secretCodeMsg := fmt.Sprintf("Request %saccess for %s. Code: %s\n", rootRoleStr, ctx.users.String(ctx.userID), authCode)
		fmt.Print(secretCodeMsg)
		ctx.users.BroadcastForRoots(ctx.messageSignal, secretCodeMsg, 0)

	} else {
		if ctx.users.IsValidCode(ctx.userID, ctx.messageArgs, forRoot) {
			ctx.users.SetAuthorized(ctx.userID, forRoot)
			if forRoot {
				replayMsg = fmt.Sprintf("You (%s) authorized as root.", ctx.users.String(ctx.userID))
				log.Print("root authorized: ", ctx.users.String(ctx.userID))
			} else {
				replayMsg = fmt.Sprintf("You (%s) authorized.", ctx.users.String(ctx.userID))
				log.Print("authorized: ", ctx.users.String(ctx.userID))
			}
		} else {
			replayMsg = fmt.Sprintf("Code is not valid.")
		}
	}

	return replayMsg
}

// /help
func cmdHelp(ctx Ctx) (replayMsg string) {
	helpMsg := []string{}

	if ctx.allowExec {
		for cmd, shellCmdRow := range ctx.commands {
			description := shellCmdRow.description
			if description == "" {
				description = shellCmdRow.shellCmd
			}
			helpMsg = append(helpMsg, cmd+" → "+description)
		}
	}
	sort.Strings(helpMsg)

	if !ctx.appConfig.isPublicBot {
		helpMsg = append(helpMsg,
			"/auth [code] → authorize user",
			"/authroot [code] → authorize user as root",
		)
	}

	if ctx.users.IsRoot(ctx.userID) {
		helpMsgForRoot := []string{
			"/shell2telegram ban <user_id|username> → ban user",
			"/shell2telegram broadcast_to_root <message> → send message to all root users in private chat",
			"/shell2telegram desc <bot description> → set bot description",
			"/shell2telegram message_to_user <user_id|username> <message> → send message to user in private chat",
			"/shell2telegram rm </command> → delete command",
			"/shell2telegram search <query> → search users by name/id",
			"/shell2telegram stat → get stat about users",
			"/shell2telegram version → show version",
		}
		if ctx.appConfig.addExit {
			helpMsgForRoot = append(helpMsgForRoot, "/shell2telegram exit → terminate bot")
		}
		sort.Strings(helpMsgForRoot)

		helpMsg = append(helpMsg, helpMsgForRoot...)
	}

	if ctx.appConfig.description != "" {
		replayMsg = ctx.appConfig.description
	} else {
		replayMsg = "This bot created with shell2telegram"
	}
	replayMsg += "\n\n" +
		"available commands:\n" +
		strings.Join(helpMsg, "\n")

	return replayMsg
}

// all commands from command-line
func cmdUser(ctx Ctx) {
	if cmd, found := ctx.commands[ctx.messageCmd]; found {
		go func() {
			if ctx.appConfig.oneThread {
				ctx.oneThreadMutex.Lock()
			}
			replayMsgRaw := execShell(
				cmd.shellCmd,
				ctx.messageArgs,
				ctx.commands[ctx.messageCmd].vars,
				ctx.userID,
				ctx.chatID,
				ctx.users.list[ctx.userID].UserName,
				ctx.users.list[ctx.userID].FirstName+" "+ctx.users.list[ctx.userID].LastName,
				ctx.cache,
				ctx.cacheTTL,
				ctx.appConfig,
			)
			if ctx.appConfig.oneThread {
				ctx.oneThreadMutex.Unlock()
			}

			sendMessage(ctx.messageSignal, ctx.chatID, replayMsgRaw, cmd.isMarkdown)
		}()
	}
}

// /shell2telegram stat
func cmdShell2telegramStat(ctx Ctx) (replayMsg string) {
	for userID := range ctx.users.list {
		replayMsg += ctx.users.StringVerbose(userID) + "\n"
	}

	return replayMsg
}

// /shell2telegram search
func cmdShell2telegramSearch(ctx Ctx) (replayMsg string) {
	query := ctx.messageArgs

	if query == "" {
		return "Please set query: /shell2telegram search <query>"
	}

	for _, userID := range ctx.users.Search(query) {
		replayMsg += ctx.users.StringVerbose(userID) + "\n"
	}

	return replayMsg
}

// /shell2telegram ban
func cmdShell2telegramBan(ctx Ctx) (replayMsg string) {
	userName := ctx.messageArgs

	if userName == "" {
		return "Please set user_id or login: /shell2telegram ban <user_id|username>"
	}

	userID := ctx.users.FindByIDOrUserName(userName)

	if userID > 0 && ctx.users.BanUser(userID) {
		replayMsg = fmt.Sprintf("User %s banned", ctx.users.String(userID))
	} else {
		replayMsg = "User not found"
	}

	return replayMsg
}

// set bot description
func cmdShell2telegramDesc(ctx Ctx) (replayMsg string) {
	description := ctx.messageArgs

	if description == "" {
		return "Please set description: /shell2telegram desc <bot description>"
	}

	ctx.appConfig.description = description
	replayMsg = "Bot description set to: " + description

	return replayMsg
}

// /shell2telegram rm "/command" - delete command
func cmdShell2telegramRm(ctx Ctx) (replayMsg string) {
	commandName := ctx.messageArgs

	if commandName == "" {
		return "Please set command for delete: /shell2telegram rm </command>"
	}
	if _, ok := ctx.commands[commandName]; ok {
		delete(ctx.commands, commandName)
		replayMsg = "Deleted command: " + commandName
	} else {
		replayMsg = fmt.Sprintf("Command %s not found", commandName)
	}

	return replayMsg
}

// /shell2telegram version - get version
func cmdShell2telegramVersion(ctx Ctx) (replayMsg string) {
	replayMsg = fmt.Sprintf("shell2telegram %s", Version)
	return replayMsg
}

// /shell2telegram exit - terminate bot
func cmdShell2telegramExit(ctx Ctx) (replayMsg string) {
	if ctx.appConfig.addExit {
		replayMsg = "bye..."
		go func() {
			ctx.exitSignal <- struct{}{}
		}()
	}
	return replayMsg
}

// /shell2telegram broadcast_to_root - broadcast message to root users in private chat
func cmdShell2telegramBroadcastToRoot(ctx Ctx) (replayMsg string) {
	message := ctx.messageArgs

	if message == "" {
		replayMsg = "Please set message: /shell2telegram broadcast_to_root <message>"
	} else {
		ctx.users.BroadcastForRoots(ctx.messageSignal,
			fmt.Sprintf("Message from %s:\n%s", ctx.users.String(ctx.userID), message),
			ctx.userID, // don't send self
		)
		replayMsg = "Message sent"
	}

	return replayMsg
}

// /shell2telegram message_to_user user_id|username "message" - send message to user in private chat
func cmdShell2telegramMessageToUser(ctx Ctx) (replayMsg string) {
	userName, message := splitStringHalfBySpace(ctx.messageArgs)

	if userName == "" || message == "" {
		replayMsg = "Please set user_name and message: /shell2telegram message_to_user <user_id|username> <message>"
	} else {
		userID := ctx.users.FindByIDOrUserName(userName)

		if userID > 0 {
			ctx.users.SendMessageToPrivate(ctx.messageSignal, userID, message)
			replayMsg = "Message sent"
		} else {
			replayMsg = "User not found"
		}
	}

	return replayMsg
}
<commit_msg>Removed unused params<commit_after>package main

import (
	"fmt"
	"log"
	"sort"
	"strings"
	"sync"

	"github.com/msoap/raphanus"
)

// Ctx - context for bot command function (users, command, args, ...)
type Ctx struct {
	appConfig      *Config           // configuration
	users          *Users            // all users
	commands       Commands          // all chat commands
	userID         int               // current user
	allowExec      bool              // is user authorized
	messageCmd     string            // command name
	messageArgs    string            // command arguments
	messageSignal  chan<- BotMessage // for send telegram messages
	chatID         int               // chat for send replay
	exitSignal     chan<- struct{}   // for signal for terminate bot
	cache          *raphanus.DB      // cache for commands output
	cacheTTL       int               // cache timeout
	oneThreadMutex *sync.Mutex       // mutex for run shell commands in one thread
}

// /auth and /authroot - authorize users
func cmdAuth(ctx Ctx) (replayMsg string) {
	forRoot := ctx.messageCmd == "/authroot"

	if ctx.messageArgs == "" {

		replayMsg = "See code in terminal with shell2telegram or ask code from root user and type:\n" + ctx.messageCmd + " code"
		authCode := ctx.users.DoLogin(ctx.userID, forRoot)

		rootRoleStr := ""
		if forRoot {
			rootRoleStr = "root "
		}
		secretCodeMsg := fmt.Sprintf("Request %saccess for %s. Code: %s\n", rootRoleStr, ctx.users.String(ctx.userID), authCode)
		fmt.Print(secretCodeMsg)
		ctx.users.BroadcastForRoots(ctx.messageSignal, secretCodeMsg, 0)

	} else {
		if ctx.users.IsValidCode(ctx.userID, ctx.messageArgs, forRoot) {
			ctx.users.SetAuthorized(ctx.userID, forRoot)
			if forRoot {
				replayMsg = fmt.Sprintf("You (%s) authorized as root.", ctx.users.String(ctx.userID))
				log.Print("root authorized: ", ctx.users.String(ctx.userID))
			} else {
				replayMsg = fmt.Sprintf("You (%s) authorized.", ctx.users.String(ctx.userID))
				log.Print("authorized: ", ctx.users.String(ctx.userID))
			}
		} else {
			replayMsg = fmt.Sprintf("Code is not valid.")
		}
	}

	return replayMsg
}

// /help
func cmdHelp(ctx Ctx) (replayMsg string) {
	helpMsg := []string{}

	if ctx.allowExec {
		for cmd, shellCmdRow := range ctx.commands {
			description := shellCmdRow.description
			if description == "" {
				description = shellCmdRow.shellCmd
			}
			helpMsg = append(helpMsg, cmd+" → "+description)
		}
	}
	sort.Strings(helpMsg)

	if !ctx.appConfig.isPublicBot {
		helpMsg = append(helpMsg,
			"/auth [code] → authorize user",
			"/authroot [code] → authorize user as root",
		)
	}

	if ctx.users.IsRoot(ctx.userID) {
		helpMsgForRoot := []string{
			"/shell2telegram ban <user_id|username> → ban user",
			"/shell2telegram broadcast_to_root <message> → send message to all root users in private chat",
			"/shell2telegram desc <bot description> → set bot description",
			"/shell2telegram message_to_user <user_id|username> <message> → send message to user in private chat",
			"/shell2telegram rm </command> → delete command",
			"/shell2telegram search <query> → search users by name/id",
			"/shell2telegram stat → get stat about users",
			"/shell2telegram version → show version",
		}
		if ctx.appConfig.addExit {
			helpMsgForRoot = append(helpMsgForRoot, "/shell2telegram exit → terminate bot")
		}
		sort.Strings(helpMsgForRoot)

		helpMsg = append(helpMsg, helpMsgForRoot...)
	}

	if ctx.appConfig.description != "" {
		replayMsg = ctx.appConfig.description
	} else {
		replayMsg = "This bot created with shell2telegram"
	}
	replayMsg += "\n\n" +
		"available commands:\n" +
		strings.Join(helpMsg, "\n")

	return replayMsg
}

// all commands from command-line
func cmdUser(ctx Ctx) {
	if cmd, found := ctx.commands[ctx.messageCmd]; found {
		go func() {
			if ctx.appConfig.oneThread {
				ctx.oneThreadMutex.Lock()
			}
			replayMsgRaw := execShell(
				cmd.shellCmd,
				ctx.messageArgs,
				ctx.commands[ctx.messageCmd].vars,
				ctx.userID,
				ctx.chatID,
				ctx.users.list[ctx.userID].UserName,
				ctx.users.list[ctx.userID].FirstName+" "+ctx.users.list[ctx.userID].LastName,
				ctx.cache,
				ctx.cacheTTL,
				ctx.appConfig,
			)
			if ctx.appConfig.oneThread {
				ctx.oneThreadMutex.Unlock()
			}

			sendMessage(ctx.messageSignal, ctx.chatID, replayMsgRaw, cmd.isMarkdown)
		}()
	}
}

// /shell2telegram stat
func cmdShell2telegramStat(ctx Ctx) (replayMsg string) {
	for userID := range ctx.users.list {
		replayMsg += ctx.users.StringVerbose(userID) + "\n"
	}

	return replayMsg
}

// /shell2telegram search
func cmdShell2telegramSearch(ctx Ctx) (replayMsg string) {
	query := ctx.messageArgs

	if query == "" {
		return "Please set query: /shell2telegram search <query>"
	}

	for _, userID := range ctx.users.Search(query) {
		replayMsg += ctx.users.StringVerbose(userID) + "\n"
	}

	return replayMsg
}

// /shell2telegram ban
func cmdShell2telegramBan(ctx Ctx) (replayMsg string) {
	userName := ctx.messageArgs

	if userName == "" {
		return "Please set user_id or login: /shell2telegram ban <user_id|username>"
	}

	userID := ctx.users.FindByIDOrUserName(userName)

	if userID > 0 && ctx.users.BanUser(userID) {
		replayMsg = fmt.Sprintf("User %s banned", ctx.users.String(userID))
	} else {
		replayMsg = "User not found"
	}

	return replayMsg
}

// set bot description
func cmdShell2telegramDesc(ctx Ctx) (replayMsg string) {
	description := ctx.messageArgs

	if description == "" {
		return "Please set description: /shell2telegram desc <bot description>"
	}

	ctx.appConfig.description = description
	replayMsg = "Bot description set to: " + description

	return replayMsg
}

// /shell2telegram rm "/command" - delete command
func cmdShell2telegramRm(ctx Ctx) (replayMsg string) {
	commandName := ctx.messageArgs

	if commandName == "" {
		return "Please set command for delete: /shell2telegram rm </command>"
	}
	if _, ok := ctx.commands[commandName]; ok {
		delete(ctx.commands, commandName)
		replayMsg = "Deleted command: " + commandName
	} else {
		replayMsg = fmt.Sprintf("Command %s not found", commandName)
	}

	return replayMsg
}

// /shell2telegram version - get version
func cmdShell2telegramVersion(_ Ctx) (replayMsg string) {
	replayMsg = fmt.Sprintf("shell2telegram %s", Version)
	return replayMsg
}

// /shell2telegram exit - terminate bot
func cmdShell2telegramExit(ctx Ctx) (replayMsg string) {
	if ctx.appConfig.addExit {
		replayMsg = "bye..."
		go func() {
			ctx.exitSignal <- struct{}{}
		}()
	}
	return replayMsg
}

// /shell2telegram broadcast_to_root - broadcast message to root users in private chat
func cmdShell2telegramBroadcastToRoot(ctx Ctx) (replayMsg string) {
	message := ctx.messageArgs

	if message == "" {
		replayMsg = "Please set message: /shell2telegram broadcast_to_root <message>"
	} else {
		ctx.users.BroadcastForRoots(ctx.messageSignal,
			fmt.Sprintf("Message from %s:\n%s", ctx.users.String(ctx.userID), message),
			ctx.userID, // don't send self
		)
		replayMsg = "Message sent"
	}

	return replayMsg
}

// /shell2telegram message_to_user user_id|username "message" - send message to user in private chat
func cmdShell2telegramMessageToUser(ctx Ctx) (replayMsg string) {
	userName, message := splitStringHalfBySpace(ctx.messageArgs)

	if userName == "" || message == "" {
		replayMsg = "Please set user_name and message: /shell2telegram message_to_user <user_id|username> <message>"
	} else {
		userID := ctx.users.FindByIDOrUserName(userName)

		if userID > 0 {
			ctx.users.SendMessageToPrivate(ctx.messageSignal, userID, message)
			replayMsg = "Message sent"
		} else {
			replayMsg = "User not found"
		}
	}

	return replayMsg
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"os/signal"

	"github.com/hashicorp/terraform/command"
	"github.com/mitchellh/cli"
)

// Commands is the mapping of all the available Terraform commands.
var Commands map[string]cli.CommandFactory

// Ui is the cli.Ui used for communicating to the outside world.
var Ui cli.Ui

const ErrorPrefix = "e:"
const OutputPrefix = "o:"

func init() {
	Ui = &cli.PrefixedUi{
		AskPrefix:    OutputPrefix,
		OutputPrefix: OutputPrefix,
		InfoPrefix:   OutputPrefix,
		ErrorPrefix:  ErrorPrefix,
		Ui:           &cli.BasicUi{Writer: os.Stdout},
	}

	meta := command.Meta{
		ContextOpts: &ContextOpts,
		Ui:          Ui,
	}

	Commands = map[string]cli.CommandFactory{
		"apply": func() (cli.Command, error) {
			return &command.ApplyCommand{
				Meta:       meta,
				ShutdownCh: makeShutdownCh(),
			}, nil
		},

		"graph": func() (cli.Command, error) {
			return &command.GraphCommand{
				Meta: meta,
			}, nil
		},

		"output": func() (cli.Command, error) {
			return &command.OutputCommand{
				Meta: meta,
			}, nil
		},

		"plan": func() (cli.Command, error) {
			return &command.PlanCommand{
				Meta: meta,
			}, nil
		},

		"refresh": func() (cli.Command, error) {
			return &command.RefreshCommand{
				Meta: meta,
			}, nil
		},

		"show": func() (cli.Command, error) {
			return &command.ShowCommand{
				Meta: meta,
			}, nil
		},

		"version": func() (cli.Command, error) {
			return &command.VersionCommand{
				Meta:              meta,
				Revision:          GitCommit,
				Version:           Version,
				VersionPrerelease: VersionPrerelease,
			}, nil
		},
	}
}

// makeShutdownCh creates an interrupt listener and returns a channel.
// A message will be sent on the channel for every interrupt received.
func makeShutdownCh() <-chan struct{} {
	resultCh := make(chan struct{})

	signalCh := make(chan os.Signal, 4)
	signal.Notify(signalCh, os.Interrupt)
	go func() {
		for {
			<-signalCh
			resultCh <- struct{}{}
		}
	}()

	return resultCh
}
<commit_msg>Always enable colors for now<commit_after>package main

import (
	"os"
	"os/signal"

	"github.com/hashicorp/terraform/command"
	"github.com/mitchellh/cli"
)

// Commands is the mapping of all the available Terraform commands.
var Commands map[string]cli.CommandFactory

// Ui is the cli.Ui used for communicating to the outside world.
var Ui cli.Ui

const ErrorPrefix = "e:"
const OutputPrefix = "o:"

func init() {
	Ui = &cli.PrefixedUi{
		AskPrefix:    OutputPrefix,
		OutputPrefix: OutputPrefix,
		InfoPrefix:   OutputPrefix,
		ErrorPrefix:  ErrorPrefix,
		Ui:           &cli.BasicUi{Writer: os.Stdout},
	}

	meta := command.Meta{
		Color:       true,
		ContextOpts: &ContextOpts,
		Ui:          Ui,
	}

	Commands = map[string]cli.CommandFactory{
		"apply": func() (cli.Command, error) {
			return &command.ApplyCommand{
				Meta:       meta,
				ShutdownCh: makeShutdownCh(),
			}, nil
		},

		"graph": func() (cli.Command, error) {
			return &command.GraphCommand{
				Meta: meta,
			}, nil
		},

		"output": func() (cli.Command, error) {
			return &command.OutputCommand{
				Meta: meta,
			}, nil
		},

		"plan": func() (cli.Command, error) {
			return &command.PlanCommand{
				Meta: meta,
			}, nil
		},

		"refresh": func() (cli.Command, error) {
			return &command.RefreshCommand{
				Meta: meta,
			}, nil
		},

		"show": func() (cli.Command, error) {
			return &command.ShowCommand{
				Meta: meta,
			}, nil
		},

		"version": func() (cli.Command, error) {
			return &command.VersionCommand{
				Meta:              meta,
				Revision:          GitCommit,
				Version:           Version,
				VersionPrerelease: VersionPrerelease,
			}, nil
		},
	}
}

// makeShutdownCh creates an interrupt listener and returns a channel.
// A message will be sent on the channel for every interrupt received.
func makeShutdownCh() <-chan struct{} {
	resultCh := make(chan struct{})

	signalCh := make(chan os.Signal, 4)
	signal.Notify(signalCh, os.Interrupt)
	go func() {
		for {
			<-signalCh
			resultCh <- struct{}{}
		}
	}()

	return resultCh
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"io"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"

	"github.com/yuuki1/grabeni/aws/model"
)

type ENIClient struct {
	svc ec2iface.EC2API
	logger *log.Logger
	logWriter io.Writer
}

type AttachENIParam struct {
	InterfaceID string
	InstanceID  string
	DeviceIndex int
}

type DetachENIParam struct {
	InterfaceID string
}

type GrabENIParam AttachENIParam

type WaiterParam struct {
	MaxAttempts int
	IntervalSec int
}

func validateWaitUntilParam(p *WaiterParam) error {
	if p == nil {
		return fmt.Errorf("WaitUntilParam require")
	}

	if p.MaxAttempts <= 0 {
		return fmt.Errorf("invalid max attempts (%d)", p.MaxAttempts)
	}
	if p.IntervalSec <= 0 {
		return fmt.Errorf("invalid interval (%d) seconds", p.IntervalSec)
	}

	return nil
}

func NewENIClient() *ENIClient {
	session := session.New()
	region := os.Getenv("AWS_REGION")
	if region == "" {
		region, _ = NewMetaDataClientFromSession(session).GetRegion()
	}
	svc := ec2.New(session, &aws.Config{Region: aws.String(region)})

	f, _ := os.Open(os.DevNull)
	l := log.New(f, "", 0)

	return &ENIClient{svc: svc, logger: l}
}

func (c *ENIClient) WithLogWriter(w io.Writer) *ENIClient {
	c.logger.SetOutput(w)
	c.logWriter = w
	return c
}

func (c *ENIClient) DescribeENIByID(InterfaceID string) (*model.ENI, error) {
	params := &ec2.DescribeNetworkInterfacesInput{
		NetworkInterfaceIds: []*string{
			aws.String(InterfaceID),
		},
	}
	resp, err := c.svc.DescribeNetworkInterfaces(params)
	if err != nil {
		return nil, err
	}

	if len(resp.NetworkInterfaces) < 1 {
		return nil, nil
	}

	eni := model.NewENI(resp.NetworkInterfaces[0])
	instance, err := c.DescribeInstanceByID(eni.AttachedInstanceID())
	if err != nil {
		return nil, err
	}
	eni.SetInstance(instance)

	return eni, nil
}

func (c *ENIClient) DescribeENIs() ([]*model.ENI, error) {
	resp, err := c.svc.DescribeNetworkInterfaces(nil)
	if err != nil {
		return nil, err
	}

	if len(resp.NetworkInterfaces) < 1 {
		return nil, nil
	}

	enis := make([]*model.ENI, 0)
	for _, iface := range resp.NetworkInterfaces {
		enis = append(enis, model.NewENI(iface))
	}

	instanceIDs := make([]string, 0)
	for _, eni := range enis {
		if id := eni.AttachedInstanceID(); id != "" {
			instanceIDs = append(instanceIDs, id)
		}
	}

	instances, err := c.DescribeInstancesByIDs(instanceIDs)
	if err != nil {
		return nil, err
	}
	//TODO make hashmap

	for _, eni := range enis {
		for _, instance := range instances {
			if eni.AttachedInstanceID() == instance.InstanceID() {
				eni.SetInstance(instance)
			}
		}
	}

	return enis, nil
}

func (c *ENIClient) AttachENI(param *AttachENIParam) (*model.ENI, error) {
	eni, err := c.DescribeENIByID(param.InterfaceID)
	if err != nil {
		return nil, err
	}

	// Do nothing if the target ENI already attached with the target instance
	if eni.AttachedInstanceID() == param.InstanceID {
		return nil, nil
	}

	input := &ec2.AttachNetworkInterfaceInput{
		NetworkInterfaceId: aws.String(param.InterfaceID),
		InstanceId:         aws.String(param.InstanceID),
		DeviceIndex:        aws.Int64(int64(param.DeviceIndex)),
	}
	_, err = c.svc.AttachNetworkInterface(input)
	if err != nil {
		return nil, err
	}

	return eni, nil
}

func (c *ENIClient) AttachENIWithWaiter(p *AttachENIParam, wp *WaiterParam) (*model.ENI, error) {
	if err := validateWaitUntilParam(wp); err != nil {
		return nil, err
	}

	if eni, err := c.AttachENI(p); eni == nil || err != nil {
		return nil, err
	}

	c.logger.Printf("--> Attaching: %15s\n", p.InterfaceID)

	// Wait until attach event completed or timeout
	for i := 0; i < wp.MaxAttempts; i++ {
		fmt.Fprint(c.logWriter, ".") // use fmt.Fprint because standard log package always newline

		eni, err := c.DescribeENIByID(p.InterfaceID)
		if err != nil {
			return nil, err
		}
		if eni.Status() == "in-use" && eni.AttachedStatus() == "attached" {
			c.logger.Println()
			c.logger.Printf("--> Attached: %15s\n", p.InterfaceID)
			return eni, nil // detach completed
		}

		time.Sleep(time.Duration(wp.IntervalSec) * time.Second)
	}

	return nil, fmt.Errorf("attach %s error: over %d polling attempts", p.InterfaceID, wp.MaxAttempts)
}

func (c *ENIClient) DetachENIByAttachmentID(attachmentID string) error {
	params := &ec2.DetachNetworkInterfaceInput{
		AttachmentId: aws.String(attachmentID),
		Force:        aws.Bool(false),
	}
	_, err := c.svc.DetachNetworkInterface(params)
	if err != nil {
		return err
	}

	return nil
}

func (c *ENIClient) DetachENI(param *DetachENIParam) (*model.ENI, error) {
	eni, err := c.DescribeENIByID(param.InterfaceID)
	if err != nil {
		return nil, err
	}

	if eni.Status() == "available" {
		// already detached
		return nil, nil
	}

	if err := c.DetachENIByAttachmentID(eni.AttachmentID()); err != nil {
		return nil, err
	}

	return eni, nil
}

func (c *ENIClient) DetachENIWithWaiter(p *DetachENIParam, wp *WaiterParam) (*model.ENI, error) {
	if err := validateWaitUntilParam(wp); err != nil {
		return nil, err
	}

	if eni, err := c.DetachENI(p); eni == nil || err != nil {
		return nil, err
	}

	c.logger.Printf("--> Detaching: %15s\n", p.InterfaceID)

	// Wait until detach event completed or timeout
	for i := 0; i < wp.MaxAttempts; i++ {
		fmt.Fprint(c.logWriter, ".") // use fmt.Fprint because standard log package always newline

		eni, err := c.DescribeENIByID(p.InterfaceID)
		if err != nil {
			return nil, err
		}
		if eni.Status() == "available" {
			c.logger.Println()
			c.logger.Printf("--> Detached: %15s\n", eni.InterfaceID())
			return eni, nil // detach completed
		}

		time.Sleep(time.Duration(wp.IntervalSec) * time.Second)
	}

	return nil, fmt.Errorf("detach %s error: over %d polling attempts", p.InterfaceID, wp.MaxAttempts)
}

func (c *ENIClient) GrabENI(p *GrabENIParam, wp *WaiterParam) (*model.ENI, error) {
	eni, err := c.DescribeENIByID(p.InterfaceID)
	if err != nil {
		return nil, err
	}

	// Skip detaching if the target ENI has still not attached with the other instance
	if eni.Status() == "in-use" && eni.AttachedStatus() == "attached" {
		// Do nothing if the target ENI already attached with the target instance
		if eni.AttachedInstanceID() == p.InstanceID {
			return nil, nil
		}

		if _, err := c.DetachENIWithWaiter(&DetachENIParam{InterfaceID: eni.InterfaceID()}, wp); err != nil {
			return nil, err
		}
	}

	param := AttachENIParam(*p)
	if eni, err = c.AttachENIWithWaiter(&param, wp); err != nil {
		return nil, err
	}

	return eni, nil
}

func (c *ENIClient) DescribeInstanceByID(instanceID string) (*model.Instance, error) {
	p := &ec2.DescribeInstancesInput{
		InstanceIds: []*string{ aws.String(instanceID) },
		MaxResults: aws.Int64(1),
	}
	resp, err := c.svc.DescribeInstances(p)
	if err != nil {
		return nil, err
	}

	if len(resp.Reservations) < 1 {
		return nil, nil // Not found
	}

	instances := resp.Reservations[0].Instances

	if len(instances) < 1 {
		return nil, nil // Not found
	}

	return model.NewInstance(instances[0]), nil
}

func (c *ENIClient) DescribeInstancesByIDs(instanceIDs []string) ([]*model.Instance, error) {
	p := &ec2.DescribeInstancesInput{
		InstanceIds: aws.StringSlice(instanceIDs),
	}
	resp, err := c.svc.DescribeInstances(p)
	if err != nil {
		return nil, err
	}

	if len(resp.Reservations) < 1 {
		return nil, nil // Not found
	}

	instances := make([]*model.Instance, 1)
	for _, r := range resp.Reservations {
		for _, i := range r.Instances {
			instances = append(instances, model.NewInstance(i))
		}
	}

	return instances, nil
}
<commit_msg>Fix invalid memory address bug<commit_after>package aws

import (
	"fmt"
	"log"
	"io"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"

	"github.com/yuuki1/grabeni/aws/model"
)

type ENIClient struct {
	svc ec2iface.EC2API
	logger *log.Logger
	logWriter io.Writer
}

type AttachENIParam struct {
	InterfaceID string
	InstanceID  string
	DeviceIndex int
}

type DetachENIParam struct {
	InterfaceID string
}

type GrabENIParam AttachENIParam

type WaiterParam struct {
	MaxAttempts int
	IntervalSec int
}

func validateWaitUntilParam(p *WaiterParam) error {
	if p == nil {
		return fmt.Errorf("WaitUntilParam require")
	}

	if p.MaxAttempts <= 0 {
		return fmt.Errorf("invalid max attempts (%d)", p.MaxAttempts)
	}
	if p.IntervalSec <= 0 {
		return fmt.Errorf("invalid interval (%d) seconds", p.IntervalSec)
	}

	return nil
}

func NewENIClient() *ENIClient {
	session := session.New()
	region := os.Getenv("AWS_REGION")
	if region == "" {
		region, _ = NewMetaDataClientFromSession(session).GetRegion()
	}
	svc := ec2.New(session, &aws.Config{Region: aws.String(region)})

	f, _ := os.Open(os.DevNull)
	l := log.New(f, "", 0)

	return &ENIClient{svc: svc, logger: l}
}

func (c *ENIClient) WithLogWriter(w io.Writer) *ENIClient {
	c.logger.SetOutput(w)
	c.logWriter = w
	return c
}

func (c *ENIClient) DescribeENIByID(InterfaceID string) (*model.ENI, error) {
	params := &ec2.DescribeNetworkInterfacesInput{
		NetworkInterfaceIds: []*string{
			aws.String(InterfaceID),
		},
	}
	resp, err := c.svc.DescribeNetworkInterfaces(params)
	if err != nil {
		return nil, err
	}

	if len(resp.NetworkInterfaces) < 1 {
		return nil, nil
	}

	eni := model.NewENI(resp.NetworkInterfaces[0])
	instance, err := c.DescribeInstanceByID(eni.AttachedInstanceID())
	if err != nil {
		return nil, err
	}
	eni.SetInstance(instance)

	return eni, nil
}

func (c *ENIClient) DescribeENIs() ([]*model.ENI, error) {
	resp, err := c.svc.DescribeNetworkInterfaces(nil)
	if err != nil {
		return nil, err
	}

	if len(resp.NetworkInterfaces) < 1 {
		return nil, nil
	}

	enis := make([]*model.ENI, 0)
	for _, iface := range resp.NetworkInterfaces {
		enis = append(enis, model.NewENI(iface))
	}

	instanceIDs := make([]string, 0)
	for _, eni := range enis {
		if id := eni.AttachedInstanceID(); id != "" {
			instanceIDs = append(instanceIDs, id)
		}
	}

	instances, err := c.DescribeInstancesByIDs(instanceIDs)
	if err != nil {
		return nil, err
	}
	if len(instances) < 1 {
		return enis, nil
	}
	//TODO make hashmap

	for _, eni := range enis {
		for _, instance := range instances {
			if eni.AttachedInstanceID() == instance.InstanceID() {
				eni.SetInstance(instance)
			}
		}
	}

	return enis, nil
}

func (c *ENIClient) AttachENI(param *AttachENIParam) (*model.ENI, error) {
	eni, err := c.DescribeENIByID(param.InterfaceID)
	if err != nil {
		return nil, err
	}

	// Do nothing if the target ENI already attached with the target instance
	if eni.AttachedInstanceID() == param.InstanceID {
		return nil, nil
	}

	input := &ec2.AttachNetworkInterfaceInput{
		NetworkInterfaceId: aws.String(param.InterfaceID),
		InstanceId:         aws.String(param.InstanceID),
		DeviceIndex:        aws.Int64(int64(param.DeviceIndex)),
	}
	_, err = c.svc.AttachNetworkInterface(input)
	if err != nil {
		return nil, err
	}

	return eni, nil
}

func (c *ENIClient) AttachENIWithWaiter(p *AttachENIParam, wp *WaiterParam) (*model.ENI, error) {
	if err := validateWaitUntilParam(wp); err != nil {
		return nil, err
	}

	if eni, err := c.AttachENI(p); eni == nil || err != nil {
		return nil, err
	}

	c.logger.Printf("--> Attaching: %15s\n", p.InterfaceID)

	// Wait until attach event completed or timeout
	for i := 0; i < wp.MaxAttempts; i++ {
		fmt.Fprint(c.logWriter, ".") // use fmt.Fprint because standard log package always newline

		eni, err := c.DescribeENIByID(p.InterfaceID)
		if err != nil {
			return nil, err
		}
		if eni.Status() == "in-use" && eni.AttachedStatus() == "attached" {
			c.logger.Println()
			c.logger.Printf("--> Attached: %15s\n", p.InterfaceID)
			return eni, nil // detach completed
		}

		time.Sleep(time.Duration(wp.IntervalSec) * time.Second)
	}

	return nil, fmt.Errorf("attach %s error: over %d polling attempts", p.InterfaceID, wp.MaxAttempts)
}

func (c *ENIClient) DetachENIByAttachmentID(attachmentID string) error {
	params := &ec2.DetachNetworkInterfaceInput{
		AttachmentId: aws.String(attachmentID),
		Force:        aws.Bool(false),
	}
	_, err := c.svc.DetachNetworkInterface(params)
	if err != nil {
		return err
	}

	return nil
}

func (c *ENIClient) DetachENI(param *DetachENIParam) (*model.ENI, error) {
	eni, err := c.DescribeENIByID(param.InterfaceID)
	if err != nil {
		return nil, err
	}

	if eni.Status() == "available" {
		// already detached
		return nil, nil
	}

	if err := c.DetachENIByAttachmentID(eni.AttachmentID()); err != nil {
		return nil, err
	}

	return eni, nil
}

func (c *ENIClient) DetachENIWithWaiter(p *DetachENIParam, wp *WaiterParam) (*model.ENI, error) {
	if err := validateWaitUntilParam(wp); err != nil {
		return nil, err
	}

	if eni, err := c.DetachENI(p); eni == nil || err != nil {
		return nil, err
	}

	c.logger.Printf("--> Detaching: %15s\n", p.InterfaceID)

	// Wait until detach event completed or timeout
	for i := 0; i < wp.MaxAttempts; i++ {
		fmt.Fprint(c.logWriter, ".") // use fmt.Fprint because standard log package always newline

		eni, err := c.DescribeENIByID(p.InterfaceID)
		if err != nil {
			return nil, err
		}
		if eni.Status() == "available" {
			c.logger.Println()
			c.logger.Printf("--> Detached: %15s\n", eni.InterfaceID())
			return eni, nil // detach completed
		}

		time.Sleep(time.Duration(wp.IntervalSec) * time.Second)
	}

	return nil, fmt.Errorf("detach %s error: over %d polling attempts", p.InterfaceID, wp.MaxAttempts)
}

func (c *ENIClient) GrabENI(p *GrabENIParam, wp *WaiterParam) (*model.ENI, error) {
	eni, err := c.DescribeENIByID(p.InterfaceID)
	if err != nil {
		return nil, err
	}

	// Skip detaching if the target ENI has still not attached with the other instance
	if eni.Status() == "in-use" && eni.AttachedStatus() == "attached" {
		// Do nothing if the target ENI already attached with the target instance
		if eni.AttachedInstanceID() == p.InstanceID {
			return nil, nil
		}

		if _, err := c.DetachENIWithWaiter(&DetachENIParam{InterfaceID: eni.InterfaceID()}, wp); err != nil {
			return nil, err
		}
	}

	param := AttachENIParam(*p)
	if eni, err = c.AttachENIWithWaiter(&param, wp); err != nil {
		return nil, err
	}

	return eni, nil
}

func (c *ENIClient) DescribeInstanceByID(instanceID string) (*model.Instance, error) {
	p := &ec2.DescribeInstancesInput{
		InstanceIds: []*string{ aws.String(instanceID) },
		MaxResults: aws.Int64(1),
	}
	resp, err := c.svc.DescribeInstances(p)
	if err != nil {
		return nil, err
	}

	if len(resp.Reservations) < 1 {
		return nil, nil // Not found
	}

	instances := resp.Reservations[0].Instances

	if len(instances) < 1 {
		return nil, nil // Not found
	}

	return model.NewInstance(instances[0]), nil
}

func (c *ENIClient) DescribeInstancesByIDs(instanceIDs []string) ([]*model.Instance, error) {
	p := &ec2.DescribeInstancesInput{
		InstanceIds: aws.StringSlice(instanceIDs),
	}
	resp, err := c.svc.DescribeInstances(p)
	if err != nil {
		return nil, err
	}

	instances := make([]*model.Instance, 0)

	for _, r := range resp.Reservations {
		for _, i := range r.Instances {
			instances = append(instances, model.NewInstance(i))
		}
	}

	return instances, nil
}
<|endoftext|>


<commit_before>// Package awsauth implements AWS request signing using Signed Signature Version 2,
// Signed Signature Version 4, and the S3 Custom HTTP Authentication Scheme.
package awsauth

import (
	"net/http"
	"net/url"
	"time"
)

// Keys stores the authentication credentials to be used when signing requests.
// You can set them manually or leave it to awsauth to use environment variables.
var Keys *Credentials

type Credentials struct {
	AccessKeyID     string
	SecretAccessKey string
	SecurityToken   string `json:"Token"`
	Expiration      string
}

func (k *Credentials) Expired() bool {
	if k.Expiration != "" {
		// Credentials with no expiration can't expire
		return false
	}
	const awsform = "2006-01-02T15:04:05Z"
	t, _ := time.Parse(awsform, k.Expiration)
	expireTime := t.Add(4 * time.Minute)
	// if t + 4 mins is after now, true
	if expireTime.After(time.Now()) {
		return true
	} else {
		return false
	}
}

// Sign signs a request bound for AWS. It automatically chooses the best
// authentication scheme based on the service the request is going to.
func Sign(req *http.Request) *http.Request {
	service, _ := serviceAndRegion(req.URL.Host)
	sigVersion := awsSignVersion[service]

	switch sigVersion {
	case 2:
		return Sign2(req)
	case 3:
		return Sign3(req)
	case 4:
		return Sign4(req)
	case -1:
		return SignS3(req)
	}

	return nil
}

// Sign4 signs a request with Signed Signature Version 4.
func Sign4(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV4(req)
	meta := new(metadata)

	// Task 1
	hashedCanonReq := hashedCanonicalRequestV4(req, meta)

	// Task 2
	stringToSign := stringToSignV4(req, hashedCanonReq, meta)

	// Task 3
	signingKey := signingKeyV4(Keys.SecretAccessKey, meta.date, meta.region, meta.service)
	signature := signatureV4(signingKey, stringToSign)

	req.Header.Set("Authorization", buildAuthHeaderV4(signature, meta))

	return req
}

// Sign3 signs a request with Signed Signature Version 3.
// If the service you're accessing supports Version 4, use that instead.
func Sign3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV3(req)

	// Task 1
	stringToSign := stringToSignV3(req)

	// Task 2
	signature := signatureV3(stringToSign)

	// Task 3
	req.Header.Set("X-Amzn-Authorization", buildAuthHeaderV3(signature))

	return req
}

// Sign2 signs a request with Signed Signature Version 2.
// If the service you're accessing supports Version 4, use that instead.
func Sign2(req *http.Request) *http.Request {
	checkKeys()

	// Add the SecurityToken parameter when using STS
	// This must be added before the signature is calculated
	if Keys.SecurityToken != "" {
		v := url.Values{}
		v.Set("SecurityToken", Keys.SecurityToken)
		augmentRequestQuery(req, v)

	}

	prepareRequestV2(req)

	stringToSign := stringToSignV2(req)
	signature := signatureV2(stringToSign)

	values := url.Values{}
	values.Set("Signature", signature)

	augmentRequestQuery(req, values)

	return req
}

// SignS3 signs a request bound for Amazon S3 using their custom
// HTTP authentication scheme.
func SignS3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestS3(req)

	stringToSign := stringToSignS3(req)
	signature := signatureS3(stringToSign)

	authHeader := "AWS " + Keys.AccessKeyID + ":" + signature
	req.Header.Set("Authorization", authHeader)

	return req
}

type metadata struct {
	algorithm       string
	credentialScope string
	signedHeaders   string
	date            string
	region          string
	service         string
}

const (
	envAccessKeyID     = "AWS_ACCESS_KEY_ID"
	envSecretAccessKey = "AWS_SECRET_ACCESS_KEY"
	envSecurityToken   = "AWS_SECURITY_TOKEN"
)

var awsSignVersion = map[string]int{
	"autoscaling":          4,
	"cloudfront":           4,
	"cloudformation":       4,
	"cloudsearch":          4,
	"monitoring":           4,
	"dynamodb":             4,
	"ec2":                  2,
	"elasticmapreduce":     4,
	"elastictranscoder":    4,
	"elasticache":          2,
	"glacier":              4,
	"kinesis":              4,
	"redshift":             4,
	"rds":                  4,
	"sdb":                  2,
	"sns":                  4,
	"sqs":                  4,
	"s3":                   -1, // custom
	"elasticbeanstalk":     4,
	"importexport":         2,
	"iam":                  4,
	"route53":              3,
	"elasticloadbalancing": 4,
	"email":                3, // Simple Email Service (SES)
}
<commit_msg>somehow confused =! with ==<commit_after>// Package awsauth implements AWS request signing using Signed Signature Version 2,
// Signed Signature Version 4, and the S3 Custom HTTP Authentication Scheme.
package awsauth

import (
	"net/http"
	"net/url"
	"time"
)

// Keys stores the authentication credentials to be used when signing requests.
// You can set them manually or leave it to awsauth to use environment variables.
var Keys *Credentials

type Credentials struct {
	AccessKeyID     string
	SecretAccessKey string
	SecurityToken   string `json:"Token"`
	Expiration      string
}

func (k *Credentials) Expired() bool {
	if k.Expiration == "" {
		// Credentials with no expiration can't expire
		return false
	}
	const awsform = "2006-01-02T15:04:05Z"
	t, _ := time.Parse(awsform, k.Expiration)
	expireTime := t.Add(4 * time.Minute)
	// if t + 4 mins is after now, true
	if expireTime.After(time.Now()) {
		return true
	} else {
		return false
	}
}

// Sign signs a request bound for AWS. It automatically chooses the best
// authentication scheme based on the service the request is going to.
func Sign(req *http.Request) *http.Request {
	service, _ := serviceAndRegion(req.URL.Host)
	sigVersion := awsSignVersion[service]

	switch sigVersion {
	case 2:
		return Sign2(req)
	case 3:
		return Sign3(req)
	case 4:
		return Sign4(req)
	case -1:
		return SignS3(req)
	}

	return nil
}

// Sign4 signs a request with Signed Signature Version 4.
func Sign4(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV4(req)
	meta := new(metadata)

	// Task 1
	hashedCanonReq := hashedCanonicalRequestV4(req, meta)

	// Task 2
	stringToSign := stringToSignV4(req, hashedCanonReq, meta)

	// Task 3
	signingKey := signingKeyV4(Keys.SecretAccessKey, meta.date, meta.region, meta.service)
	signature := signatureV4(signingKey, stringToSign)

	req.Header.Set("Authorization", buildAuthHeaderV4(signature, meta))

	return req
}

// Sign3 signs a request with Signed Signature Version 3.
// If the service you're accessing supports Version 4, use that instead.
func Sign3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestV3(req)

	// Task 1
	stringToSign := stringToSignV3(req)

	// Task 2
	signature := signatureV3(stringToSign)

	// Task 3
	req.Header.Set("X-Amzn-Authorization", buildAuthHeaderV3(signature))

	return req
}

// Sign2 signs a request with Signed Signature Version 2.
// If the service you're accessing supports Version 4, use that instead.
func Sign2(req *http.Request) *http.Request {
	checkKeys()

	// Add the SecurityToken parameter when using STS
	// This must be added before the signature is calculated
	if Keys.SecurityToken != "" {
		v := url.Values{}
		v.Set("SecurityToken", Keys.SecurityToken)
		augmentRequestQuery(req, v)

	}

	prepareRequestV2(req)

	stringToSign := stringToSignV2(req)
	signature := signatureV2(stringToSign)

	values := url.Values{}
	values.Set("Signature", signature)

	augmentRequestQuery(req, values)

	return req
}

// SignS3 signs a request bound for Amazon S3 using their custom
// HTTP authentication scheme.
func SignS3(req *http.Request) *http.Request {
	checkKeys()

	// Add the X-Amz-Security-Token header when using STS
	if Keys.SecurityToken != "" {
		req.Header.Set("X-Amz-Security-Token", Keys.SecurityToken)
	}

	prepareRequestS3(req)

	stringToSign := stringToSignS3(req)
	signature := signatureS3(stringToSign)

	authHeader := "AWS " + Keys.AccessKeyID + ":" + signature
	req.Header.Set("Authorization", authHeader)

	return req
}

type metadata struct {
	algorithm       string
	credentialScope string
	signedHeaders   string
	date            string
	region          string
	service         string
}

const (
	envAccessKeyID     = "AWS_ACCESS_KEY_ID"
	envSecretAccessKey = "AWS_SECRET_ACCESS_KEY"
	envSecurityToken   = "AWS_SECURITY_TOKEN"
)

var awsSignVersion = map[string]int{
	"autoscaling":          4,
	"cloudfront":           4,
	"cloudformation":       4,
	"cloudsearch":          4,
	"monitoring":           4,
	"dynamodb":             4,
	"ec2":                  2,
	"elasticmapreduce":     4,
	"elastictranscoder":    4,
	"elasticache":          2,
	"glacier":              4,
	"kinesis":              4,
	"redshift":             4,
	"rds":                  4,
	"sdb":                  2,
	"sns":                  4,
	"sqs":                  4,
	"s3":                   -1, // custom
	"elasticbeanstalk":     4,
	"importexport":         2,
	"iam":                  4,
	"route53":              3,
	"elasticloadbalancing": 4,
	"email":                3, // Simple Email Service (SES)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/taketin/takotsubo/gh"
)

var Commands = []cli.Command{
	commandGetteam,
	commandAdd,
	commandMembers,
}

var commandGetteam = cli.Command{
	Name:  "getteam",
	Usage: "getteam {repository}",
	Description: `
	Get teams information from repository.
`,
	Action: doGetteam,
}

var commandAdd = cli.Command{
	Name:  "add",
	Usage: "add {team_id} {member,member,member}",
	Description: `
	Adding member to team.
`,
	Action: doAdd,
}

var commandMembers = cli.Command{
	Name: "members",
	Usage: "members {team_id}",
	Description: `
	Show team members.
`,
	Action: doMembers,
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func assert(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func doGetteam(c *cli.Context) {
	repo := c.Args()[0]
	teams := gh.Teams(repo)

	for _, team := range teams {
		fmt.Printf("ID: %d NAME: %s URL: %s\n", team.Id, team.Name, team.Url)
	}
}

func doAdd(c *cli.Context) {
	teamId := c.Args()[0]
	membersCsv := c.Args()[1]

	members := strings.Split(membersCsv, ",")
	if gh.AddMembers(teamId, members) {
		fmt.Print("Succeeded !")
	} else {
		fmt.Print("Failed! Sorry... ><")
	}
}

func doMembers(c *cli.Context) {
	teamId := c.Args()[0]
	members := gh.Members(teamId)

	for _, member := range members {
		fmt.Printf("%s\n", member)
	}
}
<commit_msg>Go fmt<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/taketin/takotsubo/gh"
)

var Commands = []cli.Command{
	commandGetteam,
	commandAdd,
	commandMembers,
}

var commandGetteam = cli.Command{
	Name:  "getteam",
	Usage: "getteam {repository}",
	Description: `
	Get teams information from repository.
`,
	Action: doGetteam,
}

var commandAdd = cli.Command{
	Name:  "add",
	Usage: "add {team_id} {member,member,member}",
	Description: `
	Adding member to team.
`,
	Action: doAdd,
}

var commandMembers = cli.Command{
	Name:  "members",
	Usage: "members {team_id}",
	Description: `
	Show team members.
`,
	Action: doMembers,
}

func debug(v ...interface{}) {
	if os.Getenv("DEBUG") != "" {
		log.Println(v...)
	}
}

func assert(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

func doGetteam(c *cli.Context) {
	repo := c.Args()[0]
	teams := gh.Teams(repo)

	for _, team := range teams {
		fmt.Printf("ID: %d NAME: %s URL: %s\n", team.Id, team.Name, team.Url)
	}
}

func doAdd(c *cli.Context) {
	teamId := c.Args()[0]
	membersCsv := c.Args()[1]

	members := strings.Split(membersCsv, ",")
	if gh.AddMembers(teamId, members) {
		fmt.Print("Succeeded !")
	} else {
		fmt.Print("Failed! Sorry... ><")
	}
}

func doMembers(c *cli.Context) {
	teamId := c.Args()[0]
	members := gh.Members(teamId)

	for _, member := range members {
		fmt.Printf("%s\n", member)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package boot

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"

	"github.com/u-root/u-root/pkg/boot/kexec"
	"github.com/u-root/u-root/pkg/boot/util"
	"github.com/u-root/u-root/pkg/uio"
)

// LinuxImage implements OSImage for a Linux kernel + initramfs.
type LinuxImage struct {
	Name string

	Kernel  io.ReaderAt
	Initrd  io.ReaderAt
	Cmdline string
}

var _ OSImage = &LinuxImage{}

// named is satisifed by both *os.File and *vfile.File. Hack hack hack.
type named interface {
	Name() string
}

func stringer(mod io.ReaderAt) string {
	if s, ok := mod.(fmt.Stringer); ok {
		return s.String()
	}
	if f, ok := mod.(named); ok {
		return f.Name()
	}
	return fmt.Sprintf("%T", mod)
}

// Label returns either the Name or a short description.
func (li *LinuxImage) Label() string {
	if len(li.Name) > 0 {
		return li.Name
	}
	return fmt.Sprintf("Linux(kernel=%s initrd=%s)", stringer(li.Kernel), stringer(li.Initrd))
}

// String prints a human-readable version of this linux image.
func (li *LinuxImage) String() string {
	return fmt.Sprintf("LinuxImage(\n  Name: %s\n  Kernel: %s\n  Initrd: %s\n  Cmdline: %s\n)\n", li.Name, stringer(li.Kernel), stringer(li.Initrd), li.Cmdline)
}

func copyToFile(r io.Reader) (*os.File, error) {
	f, err := ioutil.TempFile("", "linux-image")
	if err != nil {
		return nil, err
	}
	defer f.Close()
	if _, err := io.Copy(f, r); err != nil {
		return nil, err
	}
	if err := f.Sync(); err != nil {
		return nil, err
	}

	readOnlyF, err := os.Open(f.Name())
	if err != nil {
		return nil, err
	}
	return readOnlyF, nil
}

// Edit the kernel command line.
func (li *LinuxImage) Edit(f func(cmdline string) string) {
	li.Cmdline = f(li.Cmdline)
}

// Load implements OSImage.Load and kexec_load's the kernel with its initramfs.
func (li *LinuxImage) Load(verbose bool) error {
	if li.Kernel == nil {
		return errors.New("LinuxImage.Kernel must be non-nil")
	}

	kernel, initrd := uio.Reader(util.TryGzipFilter(li.Kernel)), uio.Reader(li.Initrd)
	if verbose {
		// In verbose mode, print a dot every 5MiB. It is not pretty,
		// but it at least proves the files are still downloading.
		progress := func(r io.Reader) io.Reader {
			return &uio.ProgressReader{
				R:        r,
				Symbol:   ".",
				Interval: 5 * 1024 * 1024,
				W:        os.Stdout,
			}
		}
		kernel = progress(kernel)
		initrd = progress(initrd)
	}

	// It seams inefficient to always copy, in particular when the reader
	// is an io.File but that's not sufficient, os.File could be a socket,
	// a pipe or some other strange thing. Also kexec_file_load will fail
	// (similar to execve) if anything as the file opened for writing.
	// That's unfortunately something we can't guarantee here - unless we
	// make a copy of the file and dump it somewhere.
	k, err := copyToFile(kernel)
	if err != nil {
		return err
	}
	defer k.Close()

	var i *os.File
	if li.Initrd != nil {
		i, err = copyToFile(initrd)
		if err != nil {
			return err
		}
		defer i.Close()
	}

	if verbose {
		log.Printf("Kernel: %s", k.Name())
		if i != nil {
			log.Printf("Initrd: %s", i.Name())
		}
		log.Printf("Command line: %s", li.Cmdline)
	}
	return kexec.FileLoad(k, i, li.Cmdline)
}
<commit_msg>update name of tmp files<commit_after>// Copyright 2018 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package boot

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"

	"github.com/u-root/u-root/pkg/boot/kexec"
	"github.com/u-root/u-root/pkg/boot/util"
	"github.com/u-root/u-root/pkg/uio"
)

// LinuxImage implements OSImage for a Linux kernel + initramfs.
type LinuxImage struct {
	Name string

	Kernel  io.ReaderAt
	Initrd  io.ReaderAt
	Cmdline string
}

var _ OSImage = &LinuxImage{}

// named is satisifed by both *os.File and *vfile.File. Hack hack hack.
type named interface {
	Name() string
}

func stringer(mod io.ReaderAt) string {
	if s, ok := mod.(fmt.Stringer); ok {
		return s.String()
	}
	if f, ok := mod.(named); ok {
		return f.Name()
	}
	return fmt.Sprintf("%T", mod)
}

// Label returns either the Name or a short description.
func (li *LinuxImage) Label() string {
	if len(li.Name) > 0 {
		return li.Name
	}
	return fmt.Sprintf("Linux(kernel=%s initrd=%s)", stringer(li.Kernel), stringer(li.Initrd))
}

// String prints a human-readable version of this linux image.
func (li *LinuxImage) String() string {
	return fmt.Sprintf("LinuxImage(\n  Name: %s\n  Kernel: %s\n  Initrd: %s\n  Cmdline: %s\n)\n", li.Name, stringer(li.Kernel), stringer(li.Initrd), li.Cmdline)
}

func copyToFile(r io.Reader) (*os.File, error) {
	f, err := ioutil.TempFile("", "kexec-image")
	if err != nil {
		return nil, err
	}
	defer f.Close()
	if _, err := io.Copy(f, r); err != nil {
		return nil, err
	}
	if err := f.Sync(); err != nil {
		return nil, err
	}

	readOnlyF, err := os.Open(f.Name())
	if err != nil {
		return nil, err
	}
	return readOnlyF, nil
}

// Edit the kernel command line.
func (li *LinuxImage) Edit(f func(cmdline string) string) {
	li.Cmdline = f(li.Cmdline)
}

// Load implements OSImage.Load and kexec_load's the kernel with its initramfs.
func (li *LinuxImage) Load(verbose bool) error {
	if li.Kernel == nil {
		return errors.New("LinuxImage.Kernel must be non-nil")
	}

	kernel, initrd := uio.Reader(util.TryGzipFilter(li.Kernel)), uio.Reader(li.Initrd)
	if verbose {
		// In verbose mode, print a dot every 5MiB. It is not pretty,
		// but it at least proves the files are still downloading.
		progress := func(r io.Reader) io.Reader {
			return &uio.ProgressReader{
				R:        r,
				Symbol:   ".",
				Interval: 5 * 1024 * 1024,
				W:        os.Stdout,
			}
		}
		kernel = progress(kernel)
		initrd = progress(initrd)
	}

	// It seams inefficient to always copy, in particular when the reader
	// is an io.File but that's not sufficient, os.File could be a socket,
	// a pipe or some other strange thing. Also kexec_file_load will fail
	// (similar to execve) if anything as the file opened for writing.
	// That's unfortunately something we can't guarantee here - unless we
	// make a copy of the file and dump it somewhere.
	k, err := copyToFile(kernel)
	if err != nil {
		return err
	}
	defer k.Close()

	var i *os.File
	if li.Initrd != nil {
		i, err = copyToFile(initrd)
		if err != nil {
			return err
		}
		defer i.Close()
	}

	if verbose {
		log.Printf("Kernel: %s", k.Name())
		if i != nil {
			log.Printf("Initrd: %s", i.Name())
		}
		log.Printf("Command line: %s", li.Cmdline)
	}
	return kexec.FileLoad(k, i, li.Cmdline)
}
<|endoftext|>