<commit_before>package rbxmk

import (
	"strings"

	"github.com/robloxapi/rbxfile"
	"github.com/yuin/gopher-lua"
)

// Value is a value.
type Value interface{}

// Copier is a Value that can be copied.
type Copier interface {
	Value
	Copy() Value
}

// TValue is a value annotated with a type.
type TValue struct {
	Type string
	Value
}

type Type struct {
	// Name is the name of the type.
	Name string

	// Count indicates the number of Lua values that the type can reflect to and
	// from. A Count of 0 is the same as 1. Less than 0 indicates a variable
	// amount.
	Count int

	// ReflectTo converts v to a number of Lua values. l must be used only for
	// the conversion of values as needed. If err is nil, then lvs must have a
	// length of 1 or greater.
	ReflectTo func(s State, t Type, v Value) (lvs []lua.LValue, err error)

	// ReflectFrom converts a Lua value to v. l must be used only for the
	// conversion of values as needed. lvs must have a length of 1 or greater.
	ReflectFrom func(s State, t Type, lvs ...lua.LValue) (v Value, err error)

	// Serialize serializes a value.
	Serialize func(s State, v Value) (sv rbxfile.Value, err error)

	// Deserialize deserializes a value.
	Deserialize func(s State, sv rbxfile.Value) (v Value, err error)

	// Metatable defines the metamethods of a custom type. If Metatable is
	// non-nil, then a metatable is constructed and registered as a type
	// metatable under Name.
	Metatable Metatable

	// Members defines the members of a custom type. If the __index and
	// __newindex metamethods are not defined by Metatable, then Members defines
	// them according to the given members.
	Members Members

	// Constructors defines functions that construct the type. If non-nil, a
	// table containing each constructor is created and set as a global
	// referenced by Name.
	Constructors Constructors
}

// ValueCount returns the normalized number of Lua values that the type reflects
// between. Less than 0 means the amount is variable.
func (t Type) ValueCount() int {
	switch {
	case t.Count == 0, t.Count == 1:
		return 1
	case t.Count < 0:
		return -1
	}
	return t.Count
}

// Metatable defines the metamethods of a custom type.
type Metatable map[string]Metamethod

// Metamethod is called when a metamethod is invoked. v is the receiver, which
// was successfully reflected from the first argument to the metamethod.
type Metamethod func(s State, v Value) int

// Members is a set of members keyed by name.
type Members map[string]Member

// Member defines a member of a custom type. There are several kinds of members:
//
//     - Property: Method is false, Get must be defined, Set is optionally
//       defined. If Set is not defined, then the member is marked as read-only.
//     - Method: Method is true, Get must be defined, Set is ignored. Get is the
//       method itself.
type Member struct {
	// Get gets the value of a member from v and pushes it onto l. The index is
	// the 2nd argument in s.L. If Method is true, the function is pushed as a
	// method. The first argument will be the same value as v.
	Get func(s State, v Value) int
	// Set gets a value from l and sets it to a member of v. The index and value
	// are the 2nd and 3rd arguments in s.L.
	Set func(s State, v Value)
	// Method marks the member as being a method.
	Method bool
}

// Constructors is a set of constructor functions keyed by name.
type Constructors map[string]Constructor

// Constructor creates a new value of a Type. The function can receive arguments
// from s.L, and must push a new value to s.L.
type Constructor func(s State) int

// State contains references to an environment surrounding a value.
type State struct {
	*World

	L *lua.LState

	// Cycle is used to mark a table as having been traversed. This is non-nil
	// only for types that can contain other types.
	Cycle *Cycle
}

// Count returns the number of arguments in the stack frame.
func (s State) Count() int {
	return s.L.GetTop()
}

// Push reflects v according to type t registered with s.World, then pushes the
// results to s.L.
func (s State) Push(t string, v Value) int {
	typ := s.Type(t)
	lvs, err := typ.ReflectTo(s, typ, v)
	if err != nil {
		s.L.RaiseError(err.Error())
		return 0
	}
	for _, lv := range lvs {
		s.L.Push(lv)
	}
	return len(lvs)
}

// Pull gets from s.L the values starting from n, and reflects a value from them
// according to type t registered with s.World.
func (s State) Pull(n int, t string) Value {
	typ := s.Type(t)
	var v Value
	var err error
	if typ.Count < 0 {
		lvs := make([]lua.LValue, 0, 4)
		for i := n; i <= s.L.GetTop(); i++ {
			lvs = append(lvs, s.L.Get(i))
		}
		v, err = typ.ReflectFrom(s, typ, lvs...)
	} else if typ.Count > 1 {
		lvs := make([]lua.LValue, 0, 4)
		for i := n; i <= typ.Count; i++ {
			lvs = append(lvs, s.L.CheckAny(i))
		}
		v, err = typ.ReflectFrom(s, typ, lvs...)
	} else {
		v, err = typ.ReflectFrom(s, typ, s.L.CheckAny(n))
	}
	if err != nil {
		s.L.ArgError(n, err.Error())
		return nil
	}
	return v
}

// PullOpt gets from s.L the value at n, and reflects a value from it according
// to type t registered with s.World. If the value is nil, d is returned
// instead.
func (s State) PullOpt(n int, t string, d Value) Value {
	typ := s.Type(t)
	if typ.Count < 0 {
		panic("PullOpt cannot pull variable types")
	} else if typ.Count > 1 {
		panic("PullOpt cannot pull multi-value types")
	}
	lv := s.L.Get(n)
	if lv == lua.LNil {
		return d
	}
	v, err := typ.ReflectFrom(s, typ, lv)
	if err != nil {
		s.L.ArgError(n, err.Error())
		return d
	}
	return v
}

func listTypes(types []string) string {
	switch len(types) {
	case 0:
		return ""
	case 1:
		return types[0]
	case 2:
		return types[0] + " or " + types[1]
	}
	return strings.Join(types[:len(types)-2], ", ") + ", or " + types[len(types)-1]
}

// PullAnyOf gets from s.L the values starting from n, and reflects a value from
// them according to any of the types in t registered with s.World. Returns the
// first successful reflection among the types in t. If no types succeeded, then
// a type error is thrown.
func (s State) PullAnyOf(n int, t ...string) Value {
	if n > s.L.GetTop() {
		// Every type must reflect at least one value, so no values is an
		// immediate error.
		s.L.ArgError(n, "value expected")
		return nil
	}
	// Find the maximum count among the given types. 0 is treated the same as 1.
	// <0 indicates an arbitrary number of values.
	max := 1
	types := make([]Type, 0, 4)
	for _, t := range t {
		typ := s.Type(t)
		types = append(types, typ)
		if typ.Count > 1 {
			max = typ.Count
		} else if typ.Count < 0 {
			max = -1
			break
		}
	}
	switch max {
	case 1: // All types have 1 value.
		v := s.L.CheckAny(n)
		for _, t := range types {
			if v, err := t.ReflectFrom(s, t, v); err == nil {
				return v
			}
		}
	case -1: // At least one type has arbitrary values.
		lvs := make([]lua.LValue, 0, 4)
		for _, t := range types {
			lvs = lvs[:0]
			var v Value
			var err error
			if t.Count < 0 {
				// Append all values.
				for i := n; i <= s.L.GetTop(); i++ {
					lvs = append(lvs, s.L.Get(i))
				}
				v, err = t.ReflectFrom(s, t, lvs...)
			} else if t.Count > 1 {
				// Append up to type count.
				for i := n; i <= t.Count; i++ {
					lvs = append(lvs, s.L.CheckAny(i))
				}
				v, err = t.ReflectFrom(s, t, lvs...)
			} else {
				// Append single value.
				v, err = t.ReflectFrom(s, t, s.L.CheckAny(n))
			}
			if err != nil {
				continue
			}
			return v
		}
	default: // Constant maximum.
		lvs := make([]lua.LValue, 0, 4)
		for _, t := range types {
			lvs = lvs[:0]
			n := t.Count
			if n == 0 {
				n = 1
			}
			for i := n; i <= t.Count; i++ {
				lvs = append(lvs, s.L.CheckAny(i))
			}
			v, err := t.ReflectFrom(s, t, lvs...)
			if err != nil {
				continue
			}
			return v
		}
	}
	TypeError(s.L, n, listTypes(t))
	return nil
}

// Cycle is used to detect cyclic references by containing values that have
// already been traversed.
type Cycle struct {
	m map[interface{}]struct{}
}

// Has returns whether a value has already been visited.
func (c *Cycle) Has(t interface{}) bool {
	if c == nil {
		return false
	}
	_, ok := c.m[t]
	return ok
}

// Put marks a value has having been visited.
func (c *Cycle) Put(t interface{}) {
	if c == nil {
		return
	}
	if c.m == nil {
		c.m = make(map[interface{}]struct{}, 4)
	}
	c.m[t] = struct{}{}
}

// ReflectTypeTo is a Type.ReflectTo that converts v to a userdata set with a
// type metatable registered as t.Name.
func ReflectTypeTo(s State, t Type, v Value) (lvs []lua.LValue, err error) {
	u := s.L.NewUserData()
	u.Value = v
	s.L.SetMetatable(u, s.L.GetTypeMetatable(t.Name))
	return append(lvs, u), nil
}

// ReflectTypeFrom is a Type.ReflectFrom that converts v from a userdata set
// with a type metatable registered as t.Name.
func ReflectTypeFrom(s State, t Type, lvs ...lua.LValue) (v Value, err error) {
	u, ok := lvs[0].(*lua.LUserData)
	if ok {
		return nil, TypeError(nil, 0, t.Name)
	}
	if u.Metatable != s.L.GetTypeMetatable(t.Name) {
		return nil, TypeError(nil, 0, t.Name)
	}
	if v, ok = u.Value.(Value); !ok {
		return nil, TypeError(nil, 0, t.Name)
	}
	return v, nil
}

type typeError struct {
	expected string
	got      string
}

func (err typeError) Error() string {
	if err.got == "" {
		return err.expected + " expected"
	}
	return err.expected + " expected, got " + err.got
}

// TypeError raises an argument error indicating that a given type was expected.
func TypeError(l *lua.LState, n int, typ string) (err error) {
	if l != nil {
		err = typeError{expected: typ, got: l.Get(n).Type().String()}
		l.ArgError(n, err.Error())
	} else {
		err = typeError{expected: typ}
	}
	return err
}

// CheckType returns the underlying value if the value at n is a userdata that
// has type metatable corresponding to typ.
func CheckType(l *lua.LState, n int, typ string) interface{} {
	if u, ok := l.Get(n).(*lua.LUserData); ok {
		if u.Metatable == l.GetTypeMetatable(typ) {
			return u.Value
		}
	}
	TypeError(l, n, typ)
	return nil
}

// OptType is like CheckType, but returns nil if the value is nil.
func OptType(l *lua.LState, n int, typ string) interface{} {
	v := l.Get(n)
	if v == lua.LNil {
		return nil
	}
	if u, ok := v.(*lua.LUserData); ok {
		if u.Metatable == l.GetTypeMetatable(typ) {
			return u.Value
		}
	}
	TypeError(l, n, typ)
	return nil
}
<commit_msg>Panic when pushing an unregistered type.<commit_after>package rbxmk

import (
	"strings"

	"github.com/robloxapi/rbxfile"
	"github.com/yuin/gopher-lua"
)

// Value is a value.
type Value interface{}

// Copier is a Value that can be copied.
type Copier interface {
	Value
	Copy() Value
}

// TValue is a value annotated with a type.
type TValue struct {
	Type string
	Value
}

type Type struct {
	// Name is the name of the type.
	Name string

	// Count indicates the number of Lua values that the type can reflect to and
	// from. A Count of 0 is the same as 1. Less than 0 indicates a variable
	// amount.
	Count int

	// ReflectTo converts v to a number of Lua values. l must be used only for
	// the conversion of values as needed. If err is nil, then lvs must have a
	// length of 1 or greater.
	ReflectTo func(s State, t Type, v Value) (lvs []lua.LValue, err error)

	// ReflectFrom converts a Lua value to v. l must be used only for the
	// conversion of values as needed. lvs must have a length of 1 or greater.
	ReflectFrom func(s State, t Type, lvs ...lua.LValue) (v Value, err error)

	// Serialize serializes a value.
	Serialize func(s State, v Value) (sv rbxfile.Value, err error)

	// Deserialize deserializes a value.
	Deserialize func(s State, sv rbxfile.Value) (v Value, err error)

	// Metatable defines the metamethods of a custom type. If Metatable is
	// non-nil, then a metatable is constructed and registered as a type
	// metatable under Name.
	Metatable Metatable

	// Members defines the members of a custom type. If the __index and
	// __newindex metamethods are not defined by Metatable, then Members defines
	// them according to the given members.
	Members Members

	// Constructors defines functions that construct the type. If non-nil, a
	// table containing each constructor is created and set as a global
	// referenced by Name.
	Constructors Constructors
}

// ValueCount returns the normalized number of Lua values that the type reflects
// between. Less than 0 means the amount is variable.
func (t Type) ValueCount() int {
	switch {
	case t.Count == 0, t.Count == 1:
		return 1
	case t.Count < 0:
		return -1
	}
	return t.Count
}

// Metatable defines the metamethods of a custom type.
type Metatable map[string]Metamethod

// Metamethod is called when a metamethod is invoked. v is the receiver, which
// was successfully reflected from the first argument to the metamethod.
type Metamethod func(s State, v Value) int

// Members is a set of members keyed by name.
type Members map[string]Member

// Member defines a member of a custom type. There are several kinds of members:
//
//     - Property: Method is false, Get must be defined, Set is optionally
//       defined. If Set is not defined, then the member is marked as read-only.
//     - Method: Method is true, Get must be defined, Set is ignored. Get is the
//       method itself.
type Member struct {
	// Get gets the value of a member from v and pushes it onto l. The index is
	// the 2nd argument in s.L. If Method is true, the function is pushed as a
	// method. The first argument will be the same value as v.
	Get func(s State, v Value) int
	// Set gets a value from l and sets it to a member of v. The index and value
	// are the 2nd and 3rd arguments in s.L.
	Set func(s State, v Value)
	// Method marks the member as being a method.
	Method bool
}

// Constructors is a set of constructor functions keyed by name.
type Constructors map[string]Constructor

// Constructor creates a new value of a Type. The function can receive arguments
// from s.L, and must push a new value to s.L.
type Constructor func(s State) int

// State contains references to an environment surrounding a value.
type State struct {
	*World

	L *lua.LState

	// Cycle is used to mark a table as having been traversed. This is non-nil
	// only for types that can contain other types.
	Cycle *Cycle
}

// Count returns the number of arguments in the stack frame.
func (s State) Count() int {
	return s.L.GetTop()
}

// Push reflects v according to type t registered with s.World, then pushes the
// results to s.L.
func (s State) Push(t string, v Value) int {
	typ := s.Type(t)
	if typ.Name == "" {
		panic("unregistered type " + t)
	}
	lvs, err := typ.ReflectTo(s, typ, v)
	if err != nil {
		s.L.RaiseError(err.Error())
		return 0
	}
	for _, lv := range lvs {
		s.L.Push(lv)
	}
	return len(lvs)
}

// Pull gets from s.L the values starting from n, and reflects a value from them
// according to type t registered with s.World.
func (s State) Pull(n int, t string) Value {
	typ := s.Type(t)
	var v Value
	var err error
	if typ.Count < 0 {
		lvs := make([]lua.LValue, 0, 4)
		for i := n; i <= s.L.GetTop(); i++ {
			lvs = append(lvs, s.L.Get(i))
		}
		v, err = typ.ReflectFrom(s, typ, lvs...)
	} else if typ.Count > 1 {
		lvs := make([]lua.LValue, 0, 4)
		for i := n; i <= typ.Count; i++ {
			lvs = append(lvs, s.L.CheckAny(i))
		}
		v, err = typ.ReflectFrom(s, typ, lvs...)
	} else {
		v, err = typ.ReflectFrom(s, typ, s.L.CheckAny(n))
	}
	if err != nil {
		s.L.ArgError(n, err.Error())
		return nil
	}
	return v
}

// PullOpt gets from s.L the value at n, and reflects a value from it according
// to type t registered with s.World. If the value is nil, d is returned
// instead.
func (s State) PullOpt(n int, t string, d Value) Value {
	typ := s.Type(t)
	if typ.Count < 0 {
		panic("PullOpt cannot pull variable types")
	} else if typ.Count > 1 {
		panic("PullOpt cannot pull multi-value types")
	}
	lv := s.L.Get(n)
	if lv == lua.LNil {
		return d
	}
	v, err := typ.ReflectFrom(s, typ, lv)
	if err != nil {
		s.L.ArgError(n, err.Error())
		return d
	}
	return v
}

func listTypes(types []string) string {
	switch len(types) {
	case 0:
		return ""
	case 1:
		return types[0]
	case 2:
		return types[0] + " or " + types[1]
	}
	return strings.Join(types[:len(types)-2], ", ") + ", or " + types[len(types)-1]
}

// PullAnyOf gets from s.L the values starting from n, and reflects a value from
// them according to any of the types in t registered with s.World. Returns the
// first successful reflection among the types in t. If no types succeeded, then
// a type error is thrown.
func (s State) PullAnyOf(n int, t ...string) Value {
	if n > s.L.GetTop() {
		// Every type must reflect at least one value, so no values is an
		// immediate error.
		s.L.ArgError(n, "value expected")
		return nil
	}
	// Find the maximum count among the given types. 0 is treated the same as 1.
	// <0 indicates an arbitrary number of values.
	max := 1
	types := make([]Type, 0, 4)
	for _, t := range t {
		typ := s.Type(t)
		types = append(types, typ)
		if typ.Count > 1 {
			max = typ.Count
		} else if typ.Count < 0 {
			max = -1
			break
		}
	}
	switch max {
	case 1: // All types have 1 value.
		v := s.L.CheckAny(n)
		for _, t := range types {
			if v, err := t.ReflectFrom(s, t, v); err == nil {
				return v
			}
		}
	case -1: // At least one type has arbitrary values.
		lvs := make([]lua.LValue, 0, 4)
		for _, t := range types {
			lvs = lvs[:0]
			var v Value
			var err error
			if t.Count < 0 {
				// Append all values.
				for i := n; i <= s.L.GetTop(); i++ {
					lvs = append(lvs, s.L.Get(i))
				}
				v, err = t.ReflectFrom(s, t, lvs...)
			} else if t.Count > 1 {
				// Append up to type count.
				for i := n; i <= t.Count; i++ {
					lvs = append(lvs, s.L.CheckAny(i))
				}
				v, err = t.ReflectFrom(s, t, lvs...)
			} else {
				// Append single value.
				v, err = t.ReflectFrom(s, t, s.L.CheckAny(n))
			}
			if err != nil {
				continue
			}
			return v
		}
	default: // Constant maximum.
		lvs := make([]lua.LValue, 0, 4)
		for _, t := range types {
			lvs = lvs[:0]
			n := t.Count
			if n == 0 {
				n = 1
			}
			for i := n; i <= t.Count; i++ {
				lvs = append(lvs, s.L.CheckAny(i))
			}
			v, err := t.ReflectFrom(s, t, lvs...)
			if err != nil {
				continue
			}
			return v
		}
	}
	TypeError(s.L, n, listTypes(t))
	return nil
}

// Cycle is used to detect cyclic references by containing values that have
// already been traversed.
type Cycle struct {
	m map[interface{}]struct{}
}

// Has returns whether a value has already been visited.
func (c *Cycle) Has(t interface{}) bool {
	if c == nil {
		return false
	}
	_, ok := c.m[t]
	return ok
}

// Put marks a value has having been visited.
func (c *Cycle) Put(t interface{}) {
	if c == nil {
		return
	}
	if c.m == nil {
		c.m = make(map[interface{}]struct{}, 4)
	}
	c.m[t] = struct{}{}
}

// ReflectTypeTo is a Type.ReflectTo that converts v to a userdata set with a
// type metatable registered as t.Name.
func ReflectTypeTo(s State, t Type, v Value) (lvs []lua.LValue, err error) {
	u := s.L.NewUserData()
	u.Value = v
	s.L.SetMetatable(u, s.L.GetTypeMetatable(t.Name))
	return append(lvs, u), nil
}

// ReflectTypeFrom is a Type.ReflectFrom that converts v from a userdata set
// with a type metatable registered as t.Name.
func ReflectTypeFrom(s State, t Type, lvs ...lua.LValue) (v Value, err error) {
	u, ok := lvs[0].(*lua.LUserData)
	if ok {
		return nil, TypeError(nil, 0, t.Name)
	}
	if u.Metatable != s.L.GetTypeMetatable(t.Name) {
		return nil, TypeError(nil, 0, t.Name)
	}
	if v, ok = u.Value.(Value); !ok {
		return nil, TypeError(nil, 0, t.Name)
	}
	return v, nil
}

type typeError struct {
	expected string
	got      string
}

func (err typeError) Error() string {
	if err.got == "" {
		return err.expected + " expected"
	}
	return err.expected + " expected, got " + err.got
}

// TypeError raises an argument error indicating that a given type was expected.
func TypeError(l *lua.LState, n int, typ string) (err error) {
	if l != nil {
		err = typeError{expected: typ, got: l.Get(n).Type().String()}
		l.ArgError(n, err.Error())
	} else {
		err = typeError{expected: typ}
	}
	return err
}

// CheckType returns the underlying value if the value at n is a userdata that
// has type metatable corresponding to typ.
func CheckType(l *lua.LState, n int, typ string) interface{} {
	if u, ok := l.Get(n).(*lua.LUserData); ok {
		if u.Metatable == l.GetTypeMetatable(typ) {
			return u.Value
		}
	}
	TypeError(l, n, typ)
	return nil
}

// OptType is like CheckType, but returns nil if the value is nil.
func OptType(l *lua.LState, n int, typ string) interface{} {
	v := l.Get(n)
	if v == lua.LNil {
		return nil
	}
	if u, ok := v.(*lua.LUserData); ok {
		if u.Metatable == l.GetTypeMetatable(typ) {
			return u.Value
		}
	}
	TypeError(l, n, typ)
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	_ "net/http/httputil"
	"os"

	"github.com/olivere/elastic"
)

var cmdReindex = &Command{
	Run:   runReindex,
	Usage: "reindex [-v] [-bulk=<n>] [-shards=<n>] [-replicas=<n>] <source> <target>",
	Short: "reindex one index to another index",
	Long: `
The reindex command takes the documents from the source index
and bulk imports them to the specified target index.

This is quite handy if you want to change the settings of an
index and won't lose any data.

You can also copy indices from one cluster to another by using the -source
and -target options.

Use the -v flag to print progress.

Example:

  $ es reindex -v twitter twitter-snapshot
  $ es reindex -v -source=http://cluster1:9200 -target=http://cluster2:9200 twitter twitter-snapshot
`,
	ApiUrl: "http://www.elasticsearch.org/guide/reference/api/bulk.html",
}

var (
	sourceURL, targetURL string
	bulkSize             int
	shards, replicas     int
)

func init() {
	cmdReindex.Flag.BoolVar(&verbose, "v", false, "verbose")
	cmdReindex.Flag.StringVar(&sourceURL, "source", "", "URL of source cluster")
	cmdReindex.Flag.StringVar(&targetURL, "target", "", "URL of target cluster")
	cmdReindex.Flag.IntVar(&bulkSize, "bulk", 1000, "bulk size")
	cmdReindex.Flag.IntVar(&shards, "shards", -1, "number of shards for target index")
	cmdReindex.Flag.IntVar(&replicas, "replicas", -1, "number of replicas for target index")
}

func runReindex(cmd *Command, args []string) {
	if len(args) < 2 {
		cmd.printUsage()
		os.Exit(1)
	}

	sourceIndex := args[0]
	targetIndex := args[1]

	if sourceURL == "" {
		sourceURL = esUrl
	}
	if targetURL == "" {
		targetURL = esUrl
	}

	// Get a client
	sourceClient, err := elastic.NewClient(elastic.SetURL(sourceURL))
	if err != nil {
		log.Fatalf("%v", err)
	}
	targetClient := sourceClient
	if sourceURL != targetURL {
		targetClient, err = elastic.NewClient(elastic.SetURL(targetURL))
		if err != nil {
			log.Fatalf("%v", err)
		}
	}

	// Check if source index exists. Stop if it doesn't.
	exists, err := sourceClient.IndexExists(sourceIndex).Do()
	if err != nil {
		log.Fatalf("%v", err)
	}
	if !exists {
		os.Exit(0)
	}

	// Check if the target index exists. If it doesn't, create it with
	// the number of shards/replicas specified.
	exists, err = targetClient.IndexExists(targetIndex).Do()
	if err != nil {
		log.Fatalf("%v", err)
	}
	if !exists {
		settings := make(map[string]interface{})
		if shards > 0 || replicas >= 0 {
			ixs := make(map[string]interface{})
			if shards > 0 {
				ixs["number_of_shards"] = shards
			}
			if replicas >= 0 {
				ixs["number_of_replicas"] = replicas
			}
			settings["index"] = ixs
		}
		_, err = targetClient.CreateIndex(targetIndex).BodyJson(settings).Do()
		if err != nil {
			log.Fatalf("%v", err)
		}
	} else if shards > 0 || replicas >= 0 {
		fmt.Fprint(os.Stderr, "Shards and/or replicas will not be changed on an existing index\n")
	}

	// Progress callback
	progress := func(current, total int64) {
		if verbose {
			var percent int64
			if total > 0 {
				percent = 100 * current / total
			}
			fmt.Fprintf(os.Stderr, "Reindexing %9d of %9d (%3d%%)\r", current, total, percent)
		}
	}

	// Use the Elastic Reindexer
	ix := elastic.NewReindexer(sourceClient, sourceIndex, elastic.CopyToTargetIndex(targetIndex))
	ix = ix.TargetClient(targetClient)
	if bulkSize > 0 {
		ix = ix.BulkSize(bulkSize)
	}
	ix = ix.Scroll("5m")
	ix = ix.Progress(progress)
	ix = ix.StatsOnly(true)
	res, err := ix.Do()
	if err != nil {
		log.Fatalf("%v", err)
	}

	fmt.Fprintf(os.Stderr, "%d successful and %d failed request(s)\n", res.Success, res.Failed)
}
<commit_msg>Clarify shards/replicas in reindex command<commit_after>package main

import (
	"fmt"
	"log"
	_ "net/http/httputil"
	"os"

	"github.com/olivere/elastic"
)

var cmdReindex = &Command{
	Run:   runReindex,
	Usage: "reindex [-v] [-bulk=<n>] [-shards=<n>] [-replicas=<n>] <source> <target>",
	Short: "reindex one index to another index",
	Long: `
The reindex command takes the documents from the source index
and bulk imports them to the specified target index.

This is quite handy if you want to change the settings of an
index and won't lose any data.

You can also copy indices from one cluster to another by using the -source
and -target options.

Specifying -shards and/or -replicas is also possible; however, the target
index must not exist. Otherwise, -shards/-replicas will not be changed in
the target index.

Use the -v flag to print progress.

Example:

  $ es reindex -v twitter twitter-snapshot
  $ es reindex -v -source=http://cluster1:9200 -target=http://cluster2:9200 twitter twitter-snapshot
`,
	ApiUrl: "http://www.elasticsearch.org/guide/reference/api/bulk.html",
}

var (
	sourceURL, targetURL string
	bulkSize             int
	shards, replicas     int
)

func init() {
	cmdReindex.Flag.BoolVar(&verbose, "v", false, "verbose")
	cmdReindex.Flag.StringVar(&sourceURL, "source", "", "URL of source cluster")
	cmdReindex.Flag.StringVar(&targetURL, "target", "", "URL of target cluster")
	cmdReindex.Flag.IntVar(&bulkSize, "bulk", 1000, "bulk size")
	cmdReindex.Flag.IntVar(&shards, "shards", -1, "number of shards for target index")
	cmdReindex.Flag.IntVar(&replicas, "replicas", -1, "number of replicas for target index")
}

func runReindex(cmd *Command, args []string) {
	if len(args) < 2 {
		cmd.printUsage()
		os.Exit(1)
	}

	sourceIndex := args[0]
	targetIndex := args[1]

	if sourceURL == "" {
		sourceURL = esUrl
	}
	if targetURL == "" {
		targetURL = esUrl
	}

	// Get a client
	sourceClient, err := elastic.NewClient(elastic.SetURL(sourceURL))
	if err != nil {
		log.Fatalf("%v", err)
	}
	targetClient := sourceClient
	if sourceURL != targetURL {
		targetClient, err = elastic.NewClient(elastic.SetURL(targetURL))
		if err != nil {
			log.Fatalf("%v", err)
		}
	}

	// Check if source index exists. Stop if it doesn't.
	exists, err := sourceClient.IndexExists(sourceIndex).Do()
	if err != nil {
		log.Fatalf("%v", err)
	}
	if !exists {
		os.Exit(0)
	}

	// Check if the target index exists. If it doesn't, create it with
	// the number of shards/replicas specified.
	exists, err = targetClient.IndexExists(targetIndex).Do()
	if err != nil {
		log.Fatalf("%v", err)
	}
	if !exists {
		settings := make(map[string]interface{})
		if shards > 0 || replicas >= 0 {
			ixs := make(map[string]interface{})
			if shards > 0 {
				ixs["number_of_shards"] = shards
			}
			if replicas >= 0 {
				ixs["number_of_replicas"] = replicas
			}
			settings["index"] = ixs
		}
		_, err = targetClient.CreateIndex(targetIndex).BodyJson(settings).Do()
		if err != nil {
			log.Fatalf("%v", err)
		}
	} else if shards > 0 || replicas >= 0 {
		fmt.Fprint(os.Stderr, "Shards and/or replicas will not be changed on an existing index\n")
	}

	// Progress callback
	progress := func(current, total int64) {
		if verbose {
			var percent int64
			if total > 0 {
				percent = 100 * current / total
			}
			fmt.Fprintf(os.Stderr, "Reindexing %9d of %9d (%3d%%)\r", current, total, percent)
		}
	}

	// Use the Elastic Reindexer
	ix := elastic.NewReindexer(sourceClient, sourceIndex, elastic.CopyToTargetIndex(targetIndex))
	ix = ix.TargetClient(targetClient)
	if bulkSize > 0 {
		ix = ix.BulkSize(bulkSize)
	}
	ix = ix.Scroll("5m")
	ix = ix.Progress(progress)
	ix = ix.StatsOnly(true)
	res, err := ix.Do()
	if err != nil {
		log.Fatalf("%v", err)
	}

	fmt.Fprintf(os.Stderr, "%d successful and %d failed request(s)\n", res.Success, res.Failed)
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
)

type Renamer struct {
	in     string
	out    string
	rel    string
	needle string
}

func (r Renamer) Run() error {
	contents, err := ioutil.ReadFile(r.in)
	// println(string(contents))
	// println(r.out)
	err = ioutil.WriteFile(r.out, contents, 0644)
	return err
}
<commit_msg>stub renamer initial loop over json file contents<commit_after>package main

import (
	"encoding/json"
	"io/ioutil"
)

type Renamer struct {
	in     string
	out    string
	rel    string
	needle string
}

func (r Renamer) Run() error {

	var c map[string]interface{}
	contents, err := ioutil.ReadFile(r.in)
	json.Unmarshal(contents, &c)

	for _, v := range c {
		switch v.(type) {
		case string:
			println(v.(string))
		}
	}

	// println(string(contents))
	// println(r.out)
	err = ioutil.WriteFile(r.out, contents, 0644)
	return err
}
<|endoftext|>


<commit_before>package req

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/hashicorp/go-multierror"
	"io"
	"io/ioutil"
	"net/http"
	urlpkg "net/url"
	"strings"
)

// Request is the http request
type Request struct {
	error       error
	client      *Client
	httpRequest *http.Request
}

// New create a new request using the global default client.
func New() *Request {
	return defaultClient.R()
}

func (r *Request) appendError(err error) {
	r.error = multierror.Append(r.error, err)
}

// Error return the underlying error, not nil if some error
// happend when constructing the request.
func (r *Request) Error() error {
	return r.error
}

// Method set the http request method.
func (r *Request) Method(method string) *Request {
	if method == "" {
		// We document that "" means "GET" for Request.Method, and people have
		// relied on that from NewRequest, so keep that working.
		// We still enforce validMethod for non-empty methods.
		method = "GET"
	}
	if !validMethod(method) {
		err := fmt.Errorf("net/http: invalid method %q", method)
		if err != nil {
			r.appendError(err)
		}
	}
	r.httpRequest.Method = method
	r.httpRequest = r.httpRequest.WithContext(context.Background())
	return r
}

// URL set the http request url.
func (r *Request) URL(url string) *Request {
	u, err := urlpkg.Parse(url)
	if err != nil {
		r.appendError(err)
		return r
	}
	// The host's colon:port should be normalized. See Issue 14836.
	u.Host = removeEmptyPort(u.Host)
	r.httpRequest.URL = u
	r.httpRequest.Host = u.Host
	return r
}

func (r *Request) Send(method, url string) (*Response, error) {
	r.httpRequest.Method = method
	u, err := urlpkg.Parse(url)
	if err != nil {
		return nil, err
	}
	// The host's colon:port should be normalized. See Issue 14836.
	u.Host = removeEmptyPort(u.Host)
	r.httpRequest.URL = u
	r.httpRequest.Host = u.Host
	return r.execute()
}

// MustGet like Get, panic if error happens.
func (r *Request) MustGet(url string) *Response {
	resp, err := r.Get(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Get Send the request with GET method and specified url.
func (r *Request) Get(url string) (*Response, error) {
	return r.Send(http.MethodGet, url)
}

// MustPost like Post, panic if error happens.
func (r *Request) MustPost(url string) *Response {
	resp, err := r.Post(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Post Send the request with POST method and specified url.
func (r *Request) Post(url string) (*Response, error) {
	return r.Send(http.MethodPost, url)
}

// MustPut like Put, panic if error happens.
func (r *Request) MustPut(url string) *Response {
	resp, err := r.Put(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Put Send the request with Put method and specified url.
func (r *Request) Put(url string) (*Response, error) {
	return r.Send(http.MethodPut, url)
}

// MustPatch like Patch, panic if error happens.
func (r *Request) MustPatch(url string) *Response {
	resp, err := r.Patch(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Patch Send the request with PATCH method and specified url.
func (r *Request) Patch(url string) (*Response, error) {
	return r.Send(http.MethodPatch, url)
}

// MustDelete like Delete, panic if error happens.
func (r *Request) MustDelete(url string) *Response {
	resp, err := r.Delete(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Delete Send the request with DELETE method and specified url.
func (r *Request) Delete(url string) (*Response, error) {
	return r.Send(http.MethodDelete, url)
}

// MustOptions like Options, panic if error happens.
func (r *Request) MustOptions(url string) *Response {
	resp, err := r.Options(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Options Send the request with OPTIONS method and specified url.
func (r *Request) Options(url string) (*Response, error) {
	return r.Send(http.MethodOptions, url)
}

// MustHead like Head, panic if error happens.
func (r *Request) MustHead(url string) *Response {
	resp, err := r.Send(http.MethodHead, url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Head Send the request with HEAD method and specified url.
func (r *Request) Head(url string) (*Response, error) {
	return r.Send(http.MethodHead, url)
}

// Body set the request body.
func (r *Request) Body(body interface{}) *Request {
	if body == nil {
		return r
	}
	switch b := body.(type) {
	case io.ReadCloser:
		r.httpRequest.Body = b
	case io.Reader:
		r.httpRequest.Body = ioutil.NopCloser(b)
	case []byte:
		r.BodyBytes(b)
	case string:
		r.BodyString(b)
	}
	return r
}

// BodyBytes set the request body as []byte.
func (r *Request) BodyBytes(body []byte) *Request {
	r.httpRequest.Body = ioutil.NopCloser(bytes.NewReader(body))
	return r
}

// BodyString set the request body as string.
func (r *Request) BodyString(body string) *Request {
	r.httpRequest.Body = ioutil.NopCloser(strings.NewReader(body))
	return r
}

// BodyJsonString set the request body as string and set Content-Type header
// as "application/json; charset=UTF-8"
func (r *Request) BodyJsonString(body string) *Request {
	r.httpRequest.Body = ioutil.NopCloser(strings.NewReader(body))
	r.setContentType(CONTENT_TYPE_APPLICATION_JSON_UTF8)
	return r
}

// BodyJsonBytes set the request body as []byte and set Content-Type header
// as "application/json; charset=UTF-8"
func (r *Request) BodyJsonBytes(body []byte) *Request {
	r.httpRequest.Body = ioutil.NopCloser(bytes.NewReader(body))
	r.setContentType(CONTENT_TYPE_APPLICATION_JSON_UTF8)
	return r
}

// BodyJsonMarshal set the request body that marshaled from object, and
// set Content-Type header as "application/json; charset=UTF-8"
func (r *Request) BodyJsonMarshal(v interface{}) *Request {
	b, err := json.Marshal(v)
	if err != nil {
		r.appendError(err)
		return r
	}
	return r.BodyBytes(b)
}

func (r *Request) setContentType(contentType string) *Request {
	r.httpRequest.Header.Set("Content-Type", contentType)
	return r
}

func (r *Request) execute() (resp *Response, err error) {
	if r.error != nil {
		return nil, r.error
	}
	for k, v := range r.client.commonHeader {
		if r.httpRequest.Header.Get(k) == "" {
			r.httpRequest.Header.Set(k, v)
		}
	}
	logf(r.client.log, "%s %s", r.httpRequest.Method, r.httpRequest.URL.String())
	httpResponse, err := r.client.httpClient.Do(r.httpRequest)
	if err != nil {
		return
	}
	resp = &Response{
		request:  r,
		Response: httpResponse,
	}
	if r.client.t.ResponseOptions != nil && r.client.t.ResponseOptions.AutoDiscard {
		err = resp.Discard()
	}
	return
}
<commit_msg>support path params<commit_after>package req

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"github.com/hashicorp/go-multierror"
	"io"
	"io/ioutil"
	"net/http"
	urlpkg "net/url"
	"strings"
)

// Request is the http request
type Request struct {
	pathParams  map[string]string
	error       error
	client      *Client
	httpRequest *http.Request
}

// New create a new request using the global default client.
func New() *Request {
	return defaultClient.R()
}

func (r *Request) PathParams(params map[string]string) *Request {
	for key, value := range params {
		r.PathParam(key, value)
	}
	return r
}

func (r *Request) PathParam(key, value string) *Request {
	if r.pathParams == nil {
		r.pathParams = make(map[string]string)
	}
	r.pathParams[key] = value
	return r
}

func (r *Request) appendError(err error) {
	r.error = multierror.Append(r.error, err)
}

// Error return the underlying error, not nil if some error
// happend when constructing the request.
func (r *Request) Error() error {
	return r.error
}

// Method set the http request method.
func (r *Request) Method(method string) *Request {
	if method == "" {
		// We document that "" means "GET" for Request.Method, and people have
		// relied on that from NewRequest, so keep that working.
		// We still enforce validMethod for non-empty methods.
		method = "GET"
	}
	if !validMethod(method) {
		err := fmt.Errorf("net/http: invalid method %q", method)
		if err != nil {
			r.appendError(err)
		}
	}
	r.httpRequest.Method = method
	r.httpRequest = r.httpRequest.WithContext(context.Background())
	return r
}

// URL set the http request url.
func (r *Request) URL(url string) *Request {
	u, err := urlpkg.Parse(url)
	if err != nil {
		r.appendError(err)
		return r
	}
	// The host's colon:port should be normalized. See Issue 14836.
	u.Host = removeEmptyPort(u.Host)
	r.httpRequest.URL = u
	r.httpRequest.Host = u.Host
	return r
}

func (r *Request) Send(method, url string) (*Response, error) {
	if r.error != nil {
		return nil, r.error
	}

	r.httpRequest.Method = method

	// handle path params
	if len(r.pathParams) > 0 {
		for k, v := range r.pathParams {
			url = strings.Replace(url, "{"+k+"}", urlpkg.PathEscape(v), -1)
		}
	}

	u, err := urlpkg.Parse(url)
	if err != nil {
		return nil, err
	}
	// The host's colon:port should be normalized. See Issue 14836.
	u.Host = removeEmptyPort(u.Host)
	r.httpRequest.URL = u
	r.httpRequest.Host = u.Host
	return r.execute()
}

// MustGet like Get, panic if error happens.
func (r *Request) MustGet(url string) *Response {
	resp, err := r.Get(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Get Send the request with GET method and specified url.
func (r *Request) Get(url string) (*Response, error) {
	return r.Send(http.MethodGet, url)
}

// MustPost like Post, panic if error happens.
func (r *Request) MustPost(url string) *Response {
	resp, err := r.Post(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Post Send the request with POST method and specified url.
func (r *Request) Post(url string) (*Response, error) {
	return r.Send(http.MethodPost, url)
}

// MustPut like Put, panic if error happens.
func (r *Request) MustPut(url string) *Response {
	resp, err := r.Put(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Put Send the request with Put method and specified url.
func (r *Request) Put(url string) (*Response, error) {
	return r.Send(http.MethodPut, url)
}

// MustPatch like Patch, panic if error happens.
func (r *Request) MustPatch(url string) *Response {
	resp, err := r.Patch(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Patch Send the request with PATCH method and specified url.
func (r *Request) Patch(url string) (*Response, error) {
	return r.Send(http.MethodPatch, url)
}

// MustDelete like Delete, panic if error happens.
func (r *Request) MustDelete(url string) *Response {
	resp, err := r.Delete(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Delete Send the request with DELETE method and specified url.
func (r *Request) Delete(url string) (*Response, error) {
	return r.Send(http.MethodDelete, url)
}

// MustOptions like Options, panic if error happens.
func (r *Request) MustOptions(url string) *Response {
	resp, err := r.Options(url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Options Send the request with OPTIONS method and specified url.
func (r *Request) Options(url string) (*Response, error) {
	return r.Send(http.MethodOptions, url)
}

// MustHead like Head, panic if error happens.
func (r *Request) MustHead(url string) *Response {
	resp, err := r.Send(http.MethodHead, url)
	if err != nil {
		panic(err)
	}
	return resp
}

// Head Send the request with HEAD method and specified url.
func (r *Request) Head(url string) (*Response, error) {
	return r.Send(http.MethodHead, url)
}

// Body set the request body.
func (r *Request) Body(body interface{}) *Request {
	if body == nil {
		return r
	}
	switch b := body.(type) {
	case io.ReadCloser:
		r.httpRequest.Body = b
	case io.Reader:
		r.httpRequest.Body = ioutil.NopCloser(b)
	case []byte:
		r.BodyBytes(b)
	case string:
		r.BodyString(b)
	}
	return r
}

// BodyBytes set the request body as []byte.
func (r *Request) BodyBytes(body []byte) *Request {
	r.httpRequest.Body = ioutil.NopCloser(bytes.NewReader(body))
	return r
}

// BodyString set the request body as string.
func (r *Request) BodyString(body string) *Request {
	r.httpRequest.Body = ioutil.NopCloser(strings.NewReader(body))
	return r
}

// BodyJsonString set the request body as string and set Content-Type header
// as "application/json; charset=UTF-8"
func (r *Request) BodyJsonString(body string) *Request {
	r.httpRequest.Body = ioutil.NopCloser(strings.NewReader(body))
	r.setContentType(CONTENT_TYPE_APPLICATION_JSON_UTF8)
	return r
}

// BodyJsonBytes set the request body as []byte and set Content-Type header
// as "application/json; charset=UTF-8"
func (r *Request) BodyJsonBytes(body []byte) *Request {
	r.httpRequest.Body = ioutil.NopCloser(bytes.NewReader(body))
	r.setContentType(CONTENT_TYPE_APPLICATION_JSON_UTF8)
	return r
}

// BodyJsonMarshal set the request body that marshaled from object, and
// set Content-Type header as "application/json; charset=UTF-8"
func (r *Request) BodyJsonMarshal(v interface{}) *Request {
	b, err := json.Marshal(v)
	if err != nil {
		r.appendError(err)
		return r
	}
	return r.BodyBytes(b)
}

func (r *Request) setContentType(contentType string) *Request {
	r.httpRequest.Header.Set("Content-Type", contentType)
	return r
}

func (r *Request) execute() (resp *Response, err error) {
	if r.error != nil {
		return nil, r.error
	}
	for k, v := range r.client.commonHeader {
		if r.httpRequest.Header.Get(k) == "" {
			r.httpRequest.Header.Set(k, v)
		}
	}
	logf(r.client.log, "%s %s", r.httpRequest.Method, r.httpRequest.URL.String())
	httpResponse, err := r.client.httpClient.Do(r.httpRequest)
	if err != nil {
		return
	}
	resp = &Response{
		request:  r,
		Response: httpResponse,
	}
	if r.client.t.ResponseOptions != nil && r.client.t.ResponseOptions.AutoDiscard {
		err = resp.Discard()
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"time"

	"github.com/dickeyxxx/golock"
)

var analyticsPath = filepath.Join(HomeDir, ".heroku", "analytics.json")
var currentAnalyticsCommand = &AnalyticsCommand{
	Timestamp:  time.Now().Unix(),
	Version:    version(),
	OS:         runtime.GOOS,
	Arch:       runtime.GOARCH,
	Language:   runtime.Version(),
	CLIVersion: version(),
}

// AnalyticsCommand represents an analytics command
type AnalyticsCommand struct {
	Command    string `json:"command"`
	Plugin     string `json:"plugin,omitempty"`
	Timestamp  int64  `json:"timestamp"`
	CLIVersion string `json:"cli_version"`
	Version    string `json:"version"`
	OS         string `json:"os"`
	Arch       string `json:"arch"`
	Language   string `json:"language"`
	Status     int    `json:"status"`
	Runtime    int64  `json:"runtime"`
	Valid      bool   `json:"valid"`
	start      time.Time
}

// RecordStart marks when a command was started (for tracking runtime)
func (c *AnalyticsCommand) RecordStart() {
	c.start = time.Now()
}

// RecordEnd marks when a command was completed
// and records it to the analytics file
func (c *AnalyticsCommand) RecordEnd(status int) {
	if c == nil || skipAnalytics() || len(os.Args) < 2 {
		return
	}
	c.Command = os.Args[1]
	c.Status = status
	if !c.start.IsZero() {
		c.Valid = true
		c.Runtime = (time.Now().UnixNano() - c.start.UnixNano()) / 1000000
	}
	commands := readAnalyticsFile()
	commands = append(commands, *c)
	LogIfError(writeAnalyticsFile(commands))
}

func readAnalyticsFile() (commands []AnalyticsCommand) {
	f, err := os.Open(analyticsPath)
	if err != nil {
		return
	}
	json.NewDecoder(f).Decode(&commands)
	return commands
}

func writeAnalyticsFile(commands []AnalyticsCommand) error {
	data, err := json.MarshalIndent(commands, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(analyticsPath, data, 0644)
}

// SubmitAnalytics sends the analytics info to the analytics service
func SubmitAnalytics() {
	wg.Add(1)
	go func() {
		defer wg.Done()
		if skipAnalytics() {
			return
		}
		commands := readAnalyticsFile()
		if len(commands) < 10 {
			// do not record if less than 10 commands
			return
		}
		lockfile := filepath.Join(AppDir(), "analytics.lock")
		if locked, _ := golock.IsLocked(lockfile); locked {
			// skip if already updating
			return
		}
		golock.Lock(lockfile)
		defer golock.Unlock(lockfile)
		plugins := func() map[string]string {
			plugins := make(map[string]string)
			for _, plugin := range GetPlugins() {
				plugins[plugin.Name] = plugin.Version
			}
			dirs, _ := ioutil.ReadDir(filepath.Join(HomeDir, ".heroku", "plugins"))
			for _, dir := range dirs {
				plugins[dir.Name()] = "ruby"
			}
			return plugins
		}

		req := apiRequestBase("")
		host := os.Getenv("HEROKU_ANALYTICS_HOST")
		if host == "" {
			host = "https://cli-analytics.heroku.com"
		}
		req.Uri = host + "/record"
		req.Method = "POST"
		req.Body = struct {
			Version  string             `json:"version"`
			Commands []AnalyticsCommand `json:"commands"`
			User     string             `json:"user"`
			Plugins  map[string]string  `json:"plugins"`
		}{version(), commands, netrcLogin(), plugins()}
		resp, err := req.Do()
		if err != nil {
			LogIfError(err)
			return
		}
		if resp.StatusCode != 201 {
			Logln("analytics: HTTP " + resp.Status)
		}
		os.Truncate(analyticsPath, 0)
	}()
}

func skipAnalytics() bool {
	skip, err := config.GetBool("skip_analytics")
	if err != nil {
		Logln(err)
		return true
	}
	return skip || netrcLogin() == ""
}
<commit_msg>change golang version<commit_after>package main

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"github.com/dickeyxxx/golock"
)

var analyticsPath = filepath.Join(HomeDir, ".heroku", "analytics.json")
var currentAnalyticsCommand = &AnalyticsCommand{
	Timestamp:  time.Now().Unix(),
	Version:    version(),
	OS:         runtime.GOOS,
	Arch:       runtime.GOARCH,
	Language:   "go/" + strings.TrimPrefix(runtime.Version(), "go"),
	CLIVersion: version(),
}

// AnalyticsCommand represents an analytics command
type AnalyticsCommand struct {
	Command    string `json:"command"`
	Plugin     string `json:"plugin,omitempty"`
	Timestamp  int64  `json:"timestamp"`
	CLIVersion string `json:"cli_version"`
	Version    string `json:"version"`
	OS         string `json:"os"`
	Arch       string `json:"arch"`
	Language   string `json:"language"`
	Status     int    `json:"status"`
	Runtime    int64  `json:"runtime"`
	Valid      bool   `json:"valid"`
	start      time.Time
}

// RecordStart marks when a command was started (for tracking runtime)
func (c *AnalyticsCommand) RecordStart() {
	c.start = time.Now()
}

// RecordEnd marks when a command was completed
// and records it to the analytics file
func (c *AnalyticsCommand) RecordEnd(status int) {
	if c == nil || skipAnalytics() || len(os.Args) < 2 {
		return
	}
	c.Command = os.Args[1]
	c.Status = status
	if !c.start.IsZero() {
		c.Valid = true
		c.Runtime = (time.Now().UnixNano() - c.start.UnixNano()) / 1000000
	}
	commands := readAnalyticsFile()
	commands = append(commands, *c)
	LogIfError(writeAnalyticsFile(commands))
}

func readAnalyticsFile() (commands []AnalyticsCommand) {
	f, err := os.Open(analyticsPath)
	if err != nil {
		return
	}
	json.NewDecoder(f).Decode(&commands)
	return commands
}

func writeAnalyticsFile(commands []AnalyticsCommand) error {
	data, err := json.MarshalIndent(commands, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(analyticsPath, data, 0644)
}

// SubmitAnalytics sends the analytics info to the analytics service
func SubmitAnalytics() {
	wg.Add(1)
	go func() {
		defer wg.Done()
		if skipAnalytics() {
			return
		}
		commands := readAnalyticsFile()
		if len(commands) < 10 {
			// do not record if less than 10 commands
			return
		}
		lockfile := filepath.Join(AppDir(), "analytics.lock")
		if locked, _ := golock.IsLocked(lockfile); locked {
			// skip if already updating
			return
		}
		golock.Lock(lockfile)
		defer golock.Unlock(lockfile)
		plugins := func() map[string]string {
			plugins := make(map[string]string)
			for _, plugin := range GetPlugins() {
				plugins[plugin.Name] = plugin.Version
			}
			dirs, _ := ioutil.ReadDir(filepath.Join(HomeDir, ".heroku", "plugins"))
			for _, dir := range dirs {
				plugins[dir.Name()] = "ruby"
			}
			return plugins
		}

		req := apiRequestBase("")
		host := os.Getenv("HEROKU_ANALYTICS_HOST")
		if host == "" {
			host = "https://cli-analytics.heroku.com"
		}
		req.Uri = host + "/record"
		req.Method = "POST"
		req.Body = struct {
			Version  string             `json:"version"`
			Commands []AnalyticsCommand `json:"commands"`
			User     string             `json:"user"`
			Plugins  map[string]string  `json:"plugins"`
		}{version(), commands, netrcLogin(), plugins()}
		resp, err := req.Do()
		if err != nil {
			LogIfError(err)
			return
		}
		if resp.StatusCode != 201 {
			Logln("analytics: HTTP " + resp.Status)
		}
		os.Truncate(analyticsPath, 0)
	}()
}

func skipAnalytics() bool {
	skip, err := config.GetBool("skip_analytics")
	if err != nil {
		Logln(err)
		return true
	}
	return skip || netrcLogin() == ""
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "net/http"
    "io"
    "github.com/nfnt/resize"
    "image"
    "image/jpeg"
    "image/png"
    "strconv"
)

type Size struct {
    Width uint
    Height uint
}

// Return a given error in JSON format to the ResponseWriter
func format_error(err error, w http.ResponseWriter) {
    w.Header().Set("Content-Type", "application/json")
    io.WriteString(w, fmt.Sprintf("{ \"error\": \"%s\"}", err))
    return
}

// Parse a given string into a uint value
func parseInteger(value string) (uint, error) {
    integer, err := strconv.Atoi(value)
    return uint(integer), err
}

// Resizing endpoint.
func resizing(w http.ResponseWriter, r *http.Request) {
    size := Size{}

    // Get parameters
    imageUrl := r.FormValue("image")
    size.Width, _ = parseInteger(r.FormValue("width"))
    size.Height, _ = parseInteger(r.FormValue("height"))

    // Download the image
    imageBuffer, err := http.Get(imageUrl)
    if err != nil {
        format_error(err, w)
    }

    finalImage, _, _ := image.Decode(imageBuffer.Body)

    r.Body.Close()

    imageResized := resize.Resize(size.Width, size.Height, finalImage, resize.Lanczos3)

    if imageBuffer.Header.Get("Content-Type") == "image/png" {
        png.Encode(w, imageResized)
    }

    if imageBuffer.Header.Get("Content-Type") == "image/jpeg" {
        jpeg.Encode(w, imageResized, nil)
    }

    if imageBuffer.Header.Get("Content-Type") == "binary/octet-stream" {
        jpeg.Encode(w, imageResized, nil)
    }
}

func main() {
    http.HandleFunc("/resize", resizing)
    http.ListenAndServe(":8080", nil)
}
<commit_msg>added basic validation<commit_after>package main

import (
    "fmt"
    "net/http"
    "github.com/nfnt/resize"
    "image"
    "image/jpeg"
    "image/png"
    "strconv"
)

type Size struct {
    Width uint
    Height uint
}

func validateSize(s *Size) error {
    if s.Height >= 1000 {
        return error(fmt.Errorf("Height cannot be higher than 1000"))
    }

    if s.Width >= 1000 {
        return error(fmt.Errorf("Width cannot be higher than 1000"))
    }

    return nil
}

// Return a given error in JSON format to the ResponseWriter
func formatError(err error, w http.ResponseWriter)  {
    http.Error(w, fmt.Sprintf("{ \"error\": \"%s\"}", err), 400)
}

// Parse a given string into a uint value
func parseInteger(value string) (uint, error) {
    integer, err := strconv.Atoi(value)
    return uint(integer), err
}

// Resizing endpoint.
func resizing(w http.ResponseWriter, r *http.Request) {
    size := Size{}

    // Get parameters
    imageUrl := r.FormValue("image")
    size.Width, _ = parseInteger(r.FormValue("width"))
    size.Height, _ = parseInteger(r.FormValue("height"))

    if err := validateSize(&size); err != nil {
        formatError(err, w)
        return
    }

    // Download the image
    imageBuffer, err := http.Get(imageUrl)
    if err != nil {
        formatError(err, w)
        return
    }

    finalImage, _, _ := image.Decode(imageBuffer.Body)

    r.Body.Close()

    imageResized := resize.Resize(size.Width, size.Height, finalImage, resize.Lanczos3)

    if imageBuffer.Header.Get("Content-Type") == "image/png" {
        png.Encode(w, imageResized)
    }

    if imageBuffer.Header.Get("Content-Type") == "image/jpeg" {
        jpeg.Encode(w, imageResized, nil)
    }

    if imageBuffer.Header.Get("Content-Type") == "binary/octet-stream" {
        jpeg.Encode(w, imageResized, nil)
    }
}

func main() {
    http.HandleFunc("/resize", resizing)
    http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>package bagman

import (
	"fmt"
	"github.com/APTrust/bagins"
	"github.com/APTrust/bagman/fluctus/models"
	"github.com/diamondap/goamz/aws"
	"os"
	"path/filepath"
	"strings"
)

const (
	// BagSizeLimit is 250GB.
	BagSizeLimit = int64(268435456000)

	// Allox approx 1MB padding for tag files,
	// manifest, and tar file headers.
	BagPadding   = int64(1000000)

	// The max number of bytes files can occupy
	// in a bag.
	FileSetSizeLimit = BagSizeLimit - BagPadding

	S3UriPrefix = "https://s3.amazonaws.com/"
)

// (nsq worker) Make sure we have enough disk space.
// Create the bag (baggins).
// Create tag file with title and description.
// Create subdirectories.
// Download files from S3 to tmp area.
// Copy files into subdirectories.
// Save bag.
// (nsq worker) Tar the bag.
// (nsq worker) Copy the tar file to S3 restoration bucket.
// (nsq worker) Log results.
// (nsq worker) Clean up all temp files.

type FileSet struct {
	Files []*models.GenericFile
}

type BagRestorer struct {
	// The intellectual object we'll be restoring.
	IntellectualObject *models.IntellectualObject
	s3Client           *S3Client
	workingDir         string
	errorMessage       string
	tarFiles           []string
	fileSets           []*FileSet
	bags               []*bagins.Bag
}

// Creates a new bag restorer from the intellectual object.
// Param working dir is the path to the directory into which
// files should be downloaded and the bag should be built.
func NewBagRestorer(intelObj *models.IntellectualObject, workingDir string) (*BagRestorer, error) {
	if intelObj == nil {
		return nil, fmt.Errorf("IntellectualObject cannot be nil")
	}
	absWorkingDir, err := filepath.Abs(workingDir)
	if err != nil {
		return nil, err
	}
	// if _, err := os.Stat(absWorkingDir); os.IsNotExist(err) {
	// 	return nil, fmt.Errorf("The bag restoration working directory, '%s', does not exist", absWorkingDir)
	// }
	s3Client, err := NewS3Client(aws.USEast)
	if err != nil {
		return nil, err
	}
	// Specify the location & create a new empty bag.
	restorer := BagRestorer {
		IntellectualObject: intelObj,
		s3Client: s3Client,
		workingDir: absWorkingDir,
	}
	return &restorer, nil
}

// Fills restorer.fileSets with lists of files that can be packaged
// into individual bags.
func (restorer *BagRestorer) buildFileSets() {
	bytesInSet := int64(0)
	fileSet := &FileSet{}
	for _, gf := range restorer.IntellectualObject.GenericFiles {
		if bytesInSet + gf.Size > FileSetSizeLimit {
			restorer.fileSets = append(restorer.fileSets, fileSet)
			fileSet = &FileSet{}
			bytesInSet = 0
		}
		fileSet.Files = append(fileSet.Files, gf)
		bytesInSet += gf.Size
	}
	if bytesInSet > 0 {
		restorer.fileSets = append(restorer.fileSets, fileSet)
	}
}

// Restores an IntellectualObject by downloading all of its files
// and assembling them into one or more bags. Returns a slice of
// strings, each of which is the path to a bag.
func (restorer *BagRestorer) Restore() ([]string, error) {
	restorer.buildFileSets()
	paths := make([]string, len(restorer.fileSets))
	for i := range(restorer.fileSets) {
		bag, err := restorer.buildBag(i)
		if err != nil {
			return nil, err
		}
		paths[i] = bag.Path()
	}
	return paths, nil
}

// Creates a single bag and returns a reference to the bag object.
func (restorer *BagRestorer) buildBag(setNumber int) (*bagins.Bag, error) {
	bagName := restorer.bagName(setNumber)
	err := restorer.makeDirectory(bagName)
	if err != nil {
		return nil, err
	}
	bag, err := bagins.NewBag(restorer.workingDir, bagName, "md5")
	if err != nil {
	 	return nil, err
	}
	// Add tag files. See https://github.com/APTrust/bagins/blob/develop/bag.go#L88
	err = restorer.writeAPTrustTagFile(bag)
	if err != nil {
	 	return nil, err
	}

	// Fetch the generic files
	filesFetched, err := restorer.fetchAllFiles(setNumber)
	if err != nil {
	 	return nil, err
	}

	// Add the fetched files to the bag
	for _, fileName := range filesFetched {
		err = bag.AddFile(fileName, filepath.Base(fileName))
		if err != nil {
			return nil, err
		}
	}

	// Call save to make sure the manifest and tag files
	// are all written to disk.
	errs := bag.Save()
	if errs != nil {
		errMsg := ""
		for i := range(errs) {
			errMsg += fmt.Sprintf("%v | ", errs[i])
		}
		return nil, fmt.Errorf(errMsg)
	}

	return bag, nil
}

// Writes the aptrust-info.txt tag file.
func (restorer *BagRestorer) writeAPTrustTagFile(bag *bagins.Bag) (error) {
	if err := bag.AddTagfile("aptrust-info.txt"); err != nil {
		return err
	}
	tagFile, err := bag.TagFile("aptrust-info.txt")
	if err != nil {
		return err
	}
	tagFile.Data.AddField(*bagins.NewTagField("Title", restorer.IntellectualObject.Title))
	tagFile.Data.AddField(*bagins.NewTagField("Access", restorer.IntellectualObject.Access))
	if restorer.IntellectualObject.Description != "" {
		tagFile.Data.AddField(*bagins.NewTagField("Access", restorer.IntellectualObject.Access))
	}
	return nil
}

// Fetches all of the data files for a bag.
func (restorer *BagRestorer) fetchAllFiles(setNumber int) ([]string, error) {
	fileSet := restorer.fileSets[setNumber]
	localFilePaths := make([]string, len(fileSet.Files))
	for i, gf := range fileSet.Files {
		// TODO: Use go-routines to fetch multiple files at once?
		// If we are restoring many bags simultaneously, we could
		// wind up with too many open connections/file handles.
		fetchResult := restorer.fetchFile(gf)
		if fetchResult.ErrorMessage != "" {
			restorer.cleanup(setNumber)
			err := fmt.Errorf("Error fetching file %s from %s: %s",
				gf.Identifier, gf.URI, fetchResult.ErrorMessage)
			return nil, err
		}
		localFilePaths[i] = fetchResult.LocalTarFile
	}
	return localFilePaths, nil
}

// Creates the directories necessary to restore a bag.
func (restorer *BagRestorer) makeDirectory(bagName string) (error){
	localPath := filepath.Join(restorer.workingDir, bagName)
	localDir := filepath.Dir(localPath)
	if _, err := os.Stat(localDir); os.IsNotExist(err) {
		err = os.MkdirAll(localDir, 0755)
		if err != nil {
			return err
		}
	}
	return nil
}

// Fetches the requested file from S3 and returns a FetchResult.
func (restorer *BagRestorer) fetchFile(gf *models.GenericFile) (*FetchResult) {
	localPath := filepath.Join(restorer.workingDir, gf.Identifier)
	bucketName, key := bucketNameAndKey(gf.URI)
	s3Key, err := restorer.s3Client.GetKey(bucketName, key)
	if err != nil {
		errMsg := fmt.Sprintf("Could not get key info for %s: %v", gf.URI, err)
		return &FetchResult {
			ErrorMessage: errMsg,
		}
	}
	return restorer.s3Client.FetchToFile(bucketName, *s3Key, localPath)
}

// Given an S3 URI, returns the bucket name and key.
func bucketNameAndKey(uri string) (string, string) {
	relativeUri := strings.Replace(uri, S3UriPrefix, "", 1)
	parts := strings.SplitN(relativeUri, "/", 2)
	return parts[0], parts[1]
}

// Deletes a single bag created by Restore()
func (restorer *BagRestorer) cleanup(setNumber int) {
	bagDir := filepath.Join(restorer.workingDir, restorer.bagName(setNumber))
	_ = os.RemoveAll(bagDir)
}

// Deletes all of the bags created by Restore()
func (restorer *BagRestorer) Cleanup() {
	for i := range restorer.fileSets {
		restorer.cleanup(i)
	}
}

// BagName returns the IntelObj identifier, minus the institution name prefix,
// plus a suffix like .b001.of125, if necessary. Param setNumber is the
// index of the fileset whose files should go into the bag.
func (restorer *BagRestorer) bagName(setNumber int) (string) {
	instPrefix := fmt.Sprintf("%s/", restorer.IntellectualObject.InstitutionId)
	bagName := strings.Replace(restorer.IntellectualObject.Identifier, instPrefix, "", 1)
	if len(restorer.fileSets) > 1 {
		partNumber := setNumber + 1
		return fmt.Sprintf("%s.b%04d.of%04d", bagName, partNumber, len(restorer.fileSets))
	}
	return bagName
}

// Returns the path to the tar file. This is the file that should
// be copied to the S3 restoration bucket.
func (bagRestorer *BagRestorer) TarFiles() ([]string) {
	return bagRestorer.tarFiles
}
<commit_msg>Added optional logging to BagRestorer<commit_after>package bagman

import (
	"fmt"
	"github.com/APTrust/bagins"
	"github.com/APTrust/bagman/fluctus/models"
	"github.com/diamondap/goamz/aws"
	"github.com/op/go-logging"
	"os"
	"path/filepath"
	"strings"
)

const (
	// BagSizeLimit is 250GB.
	BagSizeLimit = int64(268435456000)

	// Allox approx 1MB padding for tag files,
	// manifest, and tar file headers.
	BagPadding   = int64(1000000)

	// The max number of bytes files can occupy
	// in a bag.
	FileSetSizeLimit = BagSizeLimit - BagPadding

	S3UriPrefix = "https://s3.amazonaws.com/"
)

// (nsq worker) Make sure we have enough disk space.
// Create the bag (baggins).
// Create tag file with title and description.
// Create subdirectories.
// Download files from S3 to tmp area.
// Copy files into subdirectories.
// Save bag.
// (nsq worker) Tar the bag.
// (nsq worker) Copy the tar file to S3 restoration bucket.
// (nsq worker) Log results.
// (nsq worker) Clean up all temp files.

type FileSet struct {
	Files []*models.GenericFile
}

type BagRestorer struct {
	// The intellectual object we'll be restoring.
	IntellectualObject *models.IntellectualObject
	s3Client           *S3Client
	workingDir         string
	errorMessage       string
	tarFiles           []string
	fileSets           []*FileSet
	bags               []*bagins.Bag
	logger             *logging.Logger
}

// Creates a new bag restorer from the intellectual object.
// Param working dir is the path to the directory into which
// files should be downloaded and the bag should be built.
func NewBagRestorer(intelObj *models.IntellectualObject, workingDir string) (*BagRestorer, error) {
	if intelObj == nil {
		return nil, fmt.Errorf("IntellectualObject cannot be nil")
	}
	absWorkingDir, err := filepath.Abs(workingDir)
	if err != nil {
		return nil, err
	}
	s3Client, err := NewS3Client(aws.USEast)
	if err != nil {
		return nil, err
	}
	// Specify the location & create a new empty bag.
	restorer := BagRestorer {
		IntellectualObject: intelObj,
		s3Client: s3Client,
		workingDir: absWorkingDir,
	}
	return &restorer, nil
}

// Sets a logger to which the BagRestorer will print debug messages.
// This is optional.
func (restorer *BagRestorer) SetLogger(logger *logging.Logger) {
	restorer.logger = logger
}

func (restorer *BagRestorer) debug (message string) {
	if restorer.logger != nil {
		restorer.logger.Debug(message)
	}
}

// Fills restorer.fileSets with lists of files that can be packaged
// into individual bags.
func (restorer *BagRestorer) buildFileSets() {
	bytesInSet := int64(0)
	fileSet := &FileSet{}
	for _, gf := range restorer.IntellectualObject.GenericFiles {
		if bytesInSet + gf.Size > FileSetSizeLimit {
			restorer.fileSets = append(restorer.fileSets, fileSet)
			fileSet = &FileSet{}
			bytesInSet = 0
		}
		fileSet.Files = append(fileSet.Files, gf)
		bytesInSet += gf.Size
		restorer.debug(fmt.Sprintf("Added %s to fileset %d", gf.Identifier, len(restorer.fileSets) + 1))
	}
	if bytesInSet > 0 {
		restorer.fileSets = append(restorer.fileSets, fileSet)
	}
}

// Restores an IntellectualObject by downloading all of its files
// and assembling them into one or more bags. Returns a slice of
// strings, each of which is the path to a bag.
func (restorer *BagRestorer) Restore() ([]string, error) {
	restorer.buildFileSets()
	paths := make([]string, len(restorer.fileSets))
	for i := range(restorer.fileSets) {
		bag, err := restorer.buildBag(i)
		if err != nil {
			return nil, err
		}
		paths[i] = bag.Path()
		restorer.debug(fmt.Sprintf("Finished bag %s", bag.Path()))
	}
	return paths, nil
}

// Creates a single bag and returns a reference to the bag object.
func (restorer *BagRestorer) buildBag(setNumber int) (*bagins.Bag, error) {
	bagName := restorer.bagName(setNumber)
	err := restorer.makeDirectory(bagName)
	if err != nil {
		return nil, err
	}
	restorer.debug(fmt.Sprintf("Creating bag %s", bagName))
	bag, err := bagins.NewBag(restorer.workingDir, bagName, "md5")
	if err != nil {
	 	return nil, err
	}
	// Add tag files. See https://github.com/APTrust/bagins/blob/develop/bag.go#L88
	err = restorer.writeAPTrustTagFile(bag)
	if err != nil {
	 	return nil, err
	}

	// Fetch the generic files
	filesFetched, err := restorer.fetchAllFiles(setNumber)
	if err != nil {
	 	return nil, err
	}

	// Add the fetched files to the bag
	for _, fileName := range filesFetched {
		err = bag.AddFile(fileName, filepath.Base(fileName))
		if err != nil {
			return nil, err
		}
	}

	// Call save to make sure the manifest and tag files
	// are all written to disk.
	errs := bag.Save()
	if errs != nil {
		errMsg := ""
		for i := range(errs) {
			errMsg += fmt.Sprintf("%v | ", errs[i])
		}
		return nil, fmt.Errorf(errMsg)
	}

	return bag, nil
}

// Writes the aptrust-info.txt tag file.
func (restorer *BagRestorer) writeAPTrustTagFile(bag *bagins.Bag) (error) {
	restorer.debug(fmt.Sprintf("Creating aptrust-info.txt"))
	if err := bag.AddTagfile("aptrust-info.txt"); err != nil {
		return err
	}
	tagFile, err := bag.TagFile("aptrust-info.txt")
	if err != nil {
		return err
	}
	tagFile.Data.AddField(*bagins.NewTagField("Title", restorer.IntellectualObject.Title))
	tagFile.Data.AddField(*bagins.NewTagField("Access", restorer.IntellectualObject.Access))
	if restorer.IntellectualObject.Description != "" {
		tagFile.Data.AddField(*bagins.NewTagField("Access", restorer.IntellectualObject.Access))
	}
	return nil
}

// Fetches all of the data files for a bag.
func (restorer *BagRestorer) fetchAllFiles(setNumber int) ([]string, error) {
	fileSet := restorer.fileSets[setNumber]
	localFilePaths := make([]string, len(fileSet.Files))
	for i, gf := range fileSet.Files {
		// TODO: Use go-routines to fetch multiple files at once?
		// If we are restoring many bags simultaneously, we could
		// wind up with too many open connections/file handles.
		fetchResult := restorer.fetchFile(gf)
		if fetchResult.ErrorMessage != "" {
			restorer.cleanup(setNumber)
			err := fmt.Errorf("Error fetching file %s from %s: %s",
				gf.Identifier, gf.URI, fetchResult.ErrorMessage)
			return nil, err
		}
		localFilePaths[i] = fetchResult.LocalTarFile
	}
	return localFilePaths, nil
}

// Creates the directories necessary to restore a bag.
func (restorer *BagRestorer) makeDirectory(bagName string) (error){
	localPath := filepath.Join(restorer.workingDir, bagName)
	localDir := filepath.Dir(localPath)
	if _, err := os.Stat(localDir); os.IsNotExist(err) {
		restorer.debug(fmt.Sprintf("Creating directory %s", localDir))
		err = os.MkdirAll(localDir, 0755)
		if err != nil {
			return err
		}
	}
	return nil
}

// Fetches the requested file from S3 and returns a FetchResult.
func (restorer *BagRestorer) fetchFile(gf *models.GenericFile) (*FetchResult) {
	localPath := filepath.Join(restorer.workingDir, gf.Identifier)
	bucketName, key := bucketNameAndKey(gf.URI)
	restorer.debug(fmt.Sprintf("Fetching key %s from bucket %s for file %s",
		key, bucketName, gf.Identifier))
	s3Key, err := restorer.s3Client.GetKey(bucketName, key)
	if err != nil {
		errMsg := fmt.Sprintf("Could not get key info for %s: %v", gf.URI, err)
		return &FetchResult {
			ErrorMessage: errMsg,
		}
	}
	return restorer.s3Client.FetchToFile(bucketName, *s3Key, localPath)
}

// Given an S3 URI, returns the bucket name and key.
func bucketNameAndKey(uri string) (string, string) {
	relativeUri := strings.Replace(uri, S3UriPrefix, "", 1)
	parts := strings.SplitN(relativeUri, "/", 2)
	return parts[0], parts[1]
}

// Deletes a single bag created by Restore()
func (restorer *BagRestorer) cleanup(setNumber int) {
	bagDir := filepath.Join(restorer.workingDir, restorer.bagName(setNumber))
	restorer.debug(fmt.Sprintf("Cleaning up %s", bagDir))
	_ = os.RemoveAll(bagDir)
}

// Deletes all of the bags created by Restore()
func (restorer *BagRestorer) Cleanup() {
	for i := range restorer.fileSets {
		restorer.cleanup(i)
	}
}

// BagName returns the IntelObj identifier, minus the institution name prefix,
// plus a suffix like .b001.of125, if necessary. Param setNumber is the
// index of the fileset whose files should go into the bag.
func (restorer *BagRestorer) bagName(setNumber int) (string) {
	instPrefix := fmt.Sprintf("%s/", restorer.IntellectualObject.InstitutionId)
	bagName := strings.Replace(restorer.IntellectualObject.Identifier, instPrefix, "", 1)
	if len(restorer.fileSets) > 1 {
		partNumber := setNumber + 1
		return fmt.Sprintf("%s.b%04d.of%04d", bagName, partNumber, len(restorer.fileSets))
	}
	return bagName
}
<|endoftext|>


<commit_before>package follower

import (
	"time"

	"github.com/bahadley/mgc/log"
)

const (
	// Length of window.
	bufSz uint32 = 4
)

var (
	// Invariant:  Heartbeats are in descending order by event.seqNo.
	hbWindow []*event

	// Used to calculate next freshness points.  Defined in freshnessPoint.go
	fpCalc deadline
)

func runObservations() {
	for {
		switch event := <-eventChan; event.eventType {
		case heartbeatEvent:
			if !insert(event) {
				log.Warning.Printf("Heartbeat from %s with seqNo %d not inserted",
					event.src, event.seqNo)
			}
			outputChan <- event
		case queryEvent:
			reportChan <- &report{freshnessPoint: event.eventTime.Add(time.Millisecond * 500)}
		case freshnessEvent:
			reportChan <- &report{suspect: false}
		default:
			log.Error.Println("Invalid event type encountered")
		}
	}
}

func insert(tmp *event) bool {
	inserted := false

	for idx, hb := range hbWindow {
		if inserted ||
			(!inserted && hb != nil && tmp.seqNo > hb.seqNo) {
			// Insert the new heartbeat and shift the subsequent heartbeats towards
			// the back of the window.  The last heartbeat will fall off if the
			// window is full.
			hbWindow[idx] = tmp
			tmp = hb
			inserted = true
		} else if !inserted && hb == nil {
			// Window is currently empty and this is the first arriving heartbeat, or ...
			// Out of order arrival and there is room at the back of the window.
			hbWindow[idx] = tmp
			inserted = true
			break
		}
	}

	return inserted
}

func init() {
	hbWindow = make([]*event, bufSz)

	fpCalc = &last{}
}
<commit_msg>Allow insert function to work with generic windows<commit_after>package follower

import (
	"time"

	"github.com/bahadley/mgc/log"
)

const (
	// Length of window.
	bufSz uint32 = 4
)

var (
	// Invariant:  Heartbeats are in descending order by event.seqNo.
	hbWindow []*event

	// Used to calculate next freshness points.  Defined in freshnessPoint.go
	fpCalc deadline
)

func runObservations() {
	for {
		switch event := <-eventChan; event.eventType {
		case heartbeatEvent:
			if !insert(event, hbWindow) {
				log.Warning.Printf("Heartbeat from %s with seqNo %d not inserted",
					event.src, event.seqNo)
			}
			outputChan <- event
		case queryEvent:
			reportChan <- &report{freshnessPoint: event.eventTime.Add(time.Millisecond * 500)}
		case freshnessEvent:
			reportChan <- &report{suspect: false}
		default:
			log.Error.Println("Invalid event type encountered")
		}
	}
}

func insert(tmp *event, window []*event) bool {
	inserted := false

	for idx, hb := range window {
		if inserted ||
			(!inserted && hb != nil && tmp.seqNo > hb.seqNo) {
			// Insert the new heartbeat and shift the subsequent heartbeats towards
			// the back of the window.  The last heartbeat will fall off if the
			// window is full.
			window[idx] = tmp
			tmp = hb
			inserted = true
		} else if !inserted && hb == nil {
			// Window is currently empty and this is the first arriving heartbeat, or ...
			// Out of order arrival and there is room at the back of the window.
			window[idx] = tmp
			inserted = true
			break
		}
	}

	return inserted
}

func init() {
	hbWindow = make([]*event, bufSz)

	fpCalc = &last{}
}
<|endoftext|>


<commit_before>// Package ns provides name server methods for selected name server(s)
package ns

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/mehrdadrad/mylg/data"
	"github.com/miekg/dns"
)

const (
	publicDNSHost      = "http://public-dns.info"
	publicDNSNodesPath = "/nameservers.csv"
)

// A Host represents a name server host
type Host struct {
	IP      string
	Alpha2  string
	Country string
	City    string
}

// A Request represents a name server request
type Request struct {
	Country string
	City    string
	Host    string
	Hosts   []Host
}

// NewRequest creates a new dns request object
func NewRequest() *Request {
	return &Request{Host: ""}
}

// Init configure dns command and fetch name servers
func (d *Request) Init() {
	if !d.cache("validate") {
		d.Hosts = fetchNSHosts()
		d.cache("write")
	} else {
		d.cache("read")
	}
}

// CountryList init the connect contry items
func (d *Request) CountryList() []string {
	var countries []string
	for _, host := range d.Hosts {
		countries = append(countries, host.Country)
	}
	countries = uniqStrSlice(countries)
	sort.Strings(countries)
	return countries
}

// NodeList gets the node city items
func (d *Request) NodeList() []string {
	var node []string
	for _, host := range d.Hosts {
		if host.Country == d.Country {
			node = append(node, host.City)
		}
	}
	sort.Strings(node)
	return node
}

// ChkCountry validates and set requested country
func (d *Request) ChkCountry(country string) bool {
	country = strings.ToLower(country)
	for _, h := range d.Hosts {
		if country == h.Country {
			d.Country = country
			return true
		}
	}

	return false
}

// ChkNode set requested country
func (d *Request) ChkNode(city string) bool {
	city = strings.ToLower(city)
	for _, h := range d.Hosts {
		if d.Country == h.Country && city == h.City {
			d.Host = h.IP
			d.City = h.City
			return true
		}
	}
	return false
}

// Local set host to nothing means local
func (d *Request) Local() {
	d.Host = ""
	d.Country = ""
}

// Dig look up name server
func (d *Request) Dig(args string) {
	c := new(dns.Client)
	m := new(dns.Msg)

	m.SetQuestion(dns.Fqdn(args), dns.TypeANY)
	m.RecursionDesired = true

	if d.Host == "" {
		config, _ := dns.ClientConfigFromFile("/etc/resolv.conf")
		d.Host = config.Servers[0]
		d.City = "your local dns server"
	}

	println("Trying to query server:", d.Host, d.Country, d.City)

	t := time.Now()
	r, _, err := c.Exchange(m, d.Host+":53")
	elapsed := time.Now().Sub(t)
	if err != nil {
		println(err.Error())
		return
	}
	fmt.Printf("Query time: %.4f ms\n", elapsed.Seconds())
	for _, a := range r.Answer {
		fmt.Println(a)
	}
}

// cache provides caching for name servers
func (d *Request) cache(r string) bool {
	switch r {
	case "read":
		b, err := ioutil.ReadFile("/tmp/mylg.ns")
		if err != nil {
			panic(err.Error())
		}
		d.Hosts = d.Hosts[:0]
		r := bytes.NewBuffer(b)
		s := bufio.NewScanner(r)
		for s.Scan() {
			csv := strings.Split(s.Text(), ";")
			if len(csv) != 4 {
				continue
			}
			d.Hosts = append(d.Hosts, Host{Alpha2: csv[0], Country: csv[1], City: csv[2], IP: csv[3]})
		}
	case "write":
		var data []string
		for _, h := range d.Hosts {
			data = append(data, fmt.Sprintf("%s;%s;%s;%s", h.Alpha2, h.Country, h.City, h.IP))
		}
		err := ioutil.WriteFile("/tmp/mylg.ns", []byte(strings.Join(data, "\n")), 0644)
		if err != nil {
			panic(err.Error())
		}
	case "validate":
		f, err := os.Stat("/tmp/mylg.ns")
		if err != nil {
			return false
		}
		d := time.Since(f.ModTime())
		if d.Hours() > 48 {
			return false
		}
	}
	return true
}

// Fetch name servers from public-dns.info
func fetchNSHosts() []Host {
	var (
		hosts   []Host
		city    string
		counter = make(map[string]int)
		chkDup  = make(map[string]int)
	)
	resp, err := http.Get(publicDNSHost + publicDNSNodesPath)
	if err != nil {
		println(err.Error())
		return []Host{}
	}
	if resp.StatusCode != 200 {
		println("error: public dns is not available")
		return []Host{}
	}
	defer resp.Body.Close()
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		csv := strings.Split(scanner.Text(), ",")
		if csv[3] != "\"\"" {
			if name, ok := data.Country[csv[2]]; ok && counter[csv[2]] < 5 {
				name = strings.ToLower(name)
				city = strings.ToLower(csv[3])
				chkDup[name+city] += 1
				if chkDup[name+city] > 1 {
					city = fmt.Sprintf("%s0%d", city, chkDup[name+city]-1)
				}

				hosts = append(hosts, Host{IP: csv[0], Alpha2: csv[2], Country: name, City: city})
				counter[csv[2]]++
			}
		}
	}
	return hosts
}

// uniqStrSlice return unique slice
func uniqStrSlice(src []string) []string {
	var rst []string
	tmp := make(map[string]struct{})
	for _, s := range src {
		tmp[s] = struct{}{}
	}
	for s := range tmp {
		rst = append(rst, s)
	}
	return rst
}
<commit_msg>added @server feature/type option/CHAOS, +trace added but not completed<commit_after>// Package ns provides name server methods for selected name server(s)
package ns

import (
	"bufio"
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/miekg/dns"

	"github.com/mehrdadrad/mylg/cli"
	"github.com/mehrdadrad/mylg/data"
)

const (
	publicDNSHost      = "http://public-dns.info"
	publicDNSNodesPath = "/nameservers.csv"
)

// A Host represents a name server host
type Host struct {
	IP      string
	Alpha2  string
	Country string
	City    string
}

// A Request represents a name server request
type Request struct {
	Country      string
	City         string
	Target       string
	Type         uint16
	Host         string
	Hosts        []Host
	TraceEnabled bool
}

// NewRequest creates a new dns request object
func NewRequest() *Request {
	return &Request{Host: ""}
}

// SetOptions passes arguments to appropriate variable
func (d *Request) SetOptions(args, prompt string) {
	nArgs, _ := cli.Flag(args)
	d.Host = ""
	d.TraceEnabled = false
	d.Type = dns.TypeANY

	for _, a := range strings.Fields(nArgs) {
		if a[0] == '@' {
			d.Host = a[1:]
			d.City = ""
			continue
		}
		if t, ok := dns.StringToType[strings.ToUpper(a)]; ok {
			d.Type = t
			continue
		}
		if a == "+trace" {
			d.TraceEnabled = true
			continue
		}
		d.Target = a
	}

	p := strings.Split(prompt, "/")

	if d.Host == "" {
		if p[0] == "local" || len(p) < 3 {
			config, _ := dns.ClientConfigFromFile("/etc/resolv.conf")
			d.Host = config.Servers[0]
			d.City = "your local dns server"
		} else {
			d.ChkNode(p[2])
		}
	}
}

// Init configure dns command and fetch name servers
func (d *Request) Init() {
	if !d.cache("validate") {
		d.Hosts = fetchNSHosts()
		d.cache("write")
	} else {
		d.cache("read")
	}
}

// CountryList init the connect contry items
func (d *Request) CountryList() []string {
	var countries []string
	for _, host := range d.Hosts {
		countries = append(countries, host.Country)
	}
	countries = uniqStrSlice(countries)
	sort.Strings(countries)
	return countries
}

// NodeList gets the node city items
func (d *Request) NodeList() []string {
	var node []string
	for _, host := range d.Hosts {
		if host.Country == d.Country {
			node = append(node, host.City)
		}
	}
	sort.Strings(node)
	return node
}

// ChkCountry validates and set requested country
func (d *Request) ChkCountry(country string) bool {
	country = strings.ToLower(country)
	for _, h := range d.Hosts {
		if country == h.Country {
			d.Country = country
			return true
		}
	}

	return false
}

// ChkNode set requested country
func (d *Request) ChkNode(city string) bool {
	city = strings.ToLower(city)
	for _, h := range d.Hosts {
		if d.Country == h.Country && city == h.City {
			d.Host = h.IP
			d.City = h.City
			return true
		}
	}
	return false
}

// Local set host to nothing means local
func (d *Request) Local() {
	d.Host = ""
	d.Country = ""
}

// Dig looks up name server w/ trace feature
func (d *Request) Dig() {
	if !d.TraceEnabled {
		d.RunDig()
	} else {
		d.RunDigTrace()
	}
}

// RunDig looks up name server
func (d *Request) RunDig() {
	c := new(dns.Client)
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(d.Target), d.Type)
	m.RecursionDesired = true
	m.RecursionAvailable = true

	println("Trying to query server:", d.Host, d.Country, d.City)

	r, rtt, err := c.Exchange(m, d.Host+":53")
	if err != nil {
		println(err.Error())
		return
	}
	fmt.Printf("Query time: %d ms\n", rtt/1e6)
	for _, a := range r.Answer {
		fmt.Println(a)
	}
	// CHAOS info
	c.Timeout = ((rtt / 1e6) + 100) * time.Millisecond
	fmt.Printf("\n=== CHAOS class BIND information ===\n")
	for _, q := range []string{"version.bind.", "hostname.bind."} {
		m.Question[0] = dns.Question{q, dns.TypeTXT, dns.ClassCHAOS}
		r, _, err = c.Exchange(m, d.Host+":53")
		if err != nil {
			continue
		}
		for _, a := range r.Answer {
			fmt.Println(a)
		}
	}
}

// RunDigTrace handles dig trace
func (d *Request) RunDigTrace() {
	// TODO
	var (
		rr   = make([]string, 30)
		host = d.Host
	)
	c := new(dns.Client)
	m := new(dns.Msg)
	m.RecursionDesired = true
	q := ""

	domain := strings.Split(dns.Fqdn(d.Target), ".")
	for i, _ := range domain {
		if i != 1 {
			q = domain[len(domain)-i-1] + "." + q
		} else {
			q = domain[len(domain)-i-1] + q
		}
		m.SetQuestion(q, dns.TypeNS)
		r, rtt, err := c.Exchange(m, host+":53")
		if err != nil {
			println(err.Error())
		}
		for _, a := range r.Answer {
			fmt.Printf("%s\n", a)
			rr = strings.Fields(a.String())
		}
		fmt.Printf("from: %s#53 in %d ms\n", host, rtt/1e6)
		ips, _ := net.LookupHost(rr[len(rr)-1])
		if len(ips) > 0 {
			host = ips[0]
		} else {
			println("can not resolve: %s", rr[len(rr)-1])
			break
		}
	}

}

// cache provides caching for name servers
func (d *Request) cache(r string) bool {
	switch r {
	case "read":
		b, err := ioutil.ReadFile("/tmp/mylg.ns")
		if err != nil {
			panic(err.Error())
		}
		d.Hosts = d.Hosts[:0]
		r := bytes.NewBuffer(b)
		s := bufio.NewScanner(r)
		for s.Scan() {
			csv := strings.Split(s.Text(), ";")
			if len(csv) != 4 {
				continue
			}
			d.Hosts = append(d.Hosts, Host{Alpha2: csv[0], Country: csv[1], City: csv[2], IP: csv[3]})
		}
	case "write":
		var data []string
		for _, h := range d.Hosts {
			data = append(data, fmt.Sprintf("%s;%s;%s;%s", h.Alpha2, h.Country, h.City, h.IP))
		}
		err := ioutil.WriteFile("/tmp/mylg.ns", []byte(strings.Join(data, "\n")), 0644)
		if err != nil {
			panic(err.Error())
		}
	case "validate":
		f, err := os.Stat("/tmp/mylg.ns")
		if err != nil {
			return false
		}
		d := time.Since(f.ModTime())
		if d.Hours() > 48 {
			return false
		}
	}
	return true
}

// Fetch name servers from public-dns.info
func fetchNSHosts() []Host {
	var (
		hosts   []Host
		city    string
		counter = make(map[string]int)
		chkDup  = make(map[string]int)
	)
	resp, err := http.Get(publicDNSHost + publicDNSNodesPath)
	if err != nil {
		println(err.Error())
		return []Host{}
	}
	if resp.StatusCode != 200 {
		println("error: public dns is not available")
		return []Host{}
	}
	defer resp.Body.Close()
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		csv := strings.Split(scanner.Text(), ",")
		if csv[3] != "\"\"" {
			if name, ok := data.Country[csv[2]]; ok && counter[csv[2]] < 5 {
				name = strings.ToLower(name)
				city = strings.ToLower(csv[3])
				chkDup[name+city] += 1
				if chkDup[name+city] > 1 {
					city = fmt.Sprintf("%s0%d", city, chkDup[name+city]-1)
				}

				hosts = append(hosts, Host{IP: csv[0], Alpha2: csv[2], Country: name, City: city})
				counter[csv[2]]++
			}
		}
	}
	return hosts
}

// uniqStrSlice return unique slice
func uniqStrSlice(src []string) []string {
	var rst []string
	tmp := make(map[string]struct{})
	for _, s := range src {
		tmp[s] = struct{}{}
	}
	for s := range tmp {
		rst = append(rst, s)
	}
	return rst
}
<|endoftext|>


<commit_before>package chroot

import (
	"context"
	"errors"
	"log"
	"strings"
	"time"

	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
	"github.com/Azure/go-autorest/autorest/azure"
	"github.com/Azure/go-autorest/autorest/to"
	"github.com/hashicorp/packer/builder/azure/common/client"
)

type DiskAttacher interface {
	AttachDisk(ctx context.Context, disk string) (lun int32, err error)
	WaitForDevice(ctx context.Context, i int32) (device string, err error)
	DetachDisk(ctx context.Context, disk string) (err error)
	WaitForDetach(ctx context.Context, diskID string) error
}

var NewDiskAttacher = func(azureClient client.AzureClientSet) DiskAttacher {
	return &diskAttacher{
		azcli: azureClient,
	}
}

type diskAttacher struct {
	azcli client.AzureClientSet

	vm *client.ComputeInfo // store info about this VM so that we don't have to ask metadata service on every call
}

var DiskNotFoundError = errors.New("Disk not found")

func (da *diskAttacher) AttachDisk(ctx context.Context, diskID string) (int32, error) {
	dataDisks, err := da.getDisks(ctx)
	if err != nil {
		return -1, err
	}

	// check to see if disk is already attached, remember lun if found
	if disk := findDiskInList(dataDisks, diskID); disk != nil {
		// disk is already attached, just take this lun
		if disk.Lun == nil {
			return -1, errors.New("disk is attached, but lun was not set in VM model (possibly an error in the Azure APIs)")
		}
		return to.Int32(disk.Lun), nil
	}

	// disk was not found on VM, go and actually attach it

	var lun int32 = -1
findFreeLun:
	for lun = 0; lun < 64; lun++ {
		for _, v := range dataDisks {
			if to.Int32(v.Lun) == lun {
				continue findFreeLun
			}
		}
		// no datadisk is using this lun
		break
	}

	// append new data disk to collection
	dataDisks = append(dataDisks, compute.DataDisk{
		CreateOption: compute.DiskCreateOptionTypesAttach,
		ManagedDisk: &compute.ManagedDiskParameters{
			ID: to.StringPtr(diskID),
		},
		Lun: to.Int32Ptr(lun),
	})

	// prepare resource object for update operation
	err = da.setDisks(ctx, dataDisks)
	if err != nil {
		return -1, err
	}

	return lun, nil
}

func (da *diskAttacher) DetachDisk(ctx context.Context, diskID string) error {
	log.Println("Fetching list of disks currently attached to VM")
	currentDisks, err := da.getDisks(ctx)
	if err != nil {
		return err
	}

	log.Printf("Removing %q from list of disks currently attached to VM", diskID)
	newDisks := []compute.DataDisk{}
	for _, disk := range currentDisks {
		if disk.ManagedDisk != nil &&
			!strings.EqualFold(to.String(disk.ManagedDisk.ID), diskID) {
			newDisks = append(newDisks, disk)
		}
	}
	if len(currentDisks) == len(newDisks) {
		return DiskNotFoundError
	}

	log.Println("Updating new list of disks attached to VM")
	err = da.setDisks(ctx, newDisks)
	if err != nil {
		return err
	}

	return nil
}

func (da *diskAttacher) WaitForDetach(ctx context.Context, diskID string) error {
	for { // loop until disk is not attached, timeout or error
		list, err := da.getDisks(ctx)
		if err != nil {
			return err
		}
		if findDiskInList(list, diskID) == nil {
			log.Println("Disk is no longer in VM model, assuming detached")
			return nil
		}

		select {
		case <-time.After(time.Second): //continue
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (da *diskAttacher) getThisVM(ctx context.Context) (compute.VirtualMachine, error) {
	// getting resource info for this VM
	if da.vm == nil {
		vm, err := da.azcli.MetadataClient().GetComputeInfo()
		if err != nil {
			return compute.VirtualMachine{}, err
		}
		da.vm = vm
	}

	// retrieve actual VM
	vmResource, err := da.azcli.VirtualMachinesClient().Get(ctx, da.vm.ResourceGroupName, da.vm.Name, "")
	if err != nil {
		return compute.VirtualMachine{}, err
	}
	if vmResource.StorageProfile == nil {
		return compute.VirtualMachine{}, errors.New("properties.storageProfile is not set on VM, this is unexpected")
	}

	return vmResource, nil
}

func (da diskAttacher) getDisks(ctx context.Context) ([]compute.DataDisk, error) {
	vmResource, err := da.getThisVM(ctx)
	if err != nil {
		return []compute.DataDisk{}, err
	}

	return *vmResource.StorageProfile.DataDisks, nil
}

func (da diskAttacher) setDisks(ctx context.Context, disks []compute.DataDisk) error {
	vmResource, err := da.getThisVM(ctx)
	if err != nil {
		return err
	}

	id, err := azure.ParseResourceID(to.String(vmResource.ID))
	if err != nil {
		return err
	}

	vmResource.StorageProfile.DataDisks = &disks
	vmResource.Resources = nil

	// update the VM resource, attach disk
	_, err = da.azcli.VirtualMachinesClient().CreateOrUpdate(ctx, id.ResourceGroup, id.ResourceName, vmResource)

	return err
}

func findDiskInList(list []compute.DataDisk, diskID string) *compute.DataDisk {
	for _, disk := range list {
		if disk.ManagedDisk != nil &&
			strings.EqualFold(to.String(disk.ManagedDisk.ID), diskID) {
			return &disk
		}
	}
	return nil
}
<commit_msg>Fix ordering to match master for better diffs.<commit_after>package chroot

import (
	"context"
	"errors"
	"log"
	"strings"
	"time"

	"github.com/Azure/azure-sdk-for-go/profiles/latest/compute/mgmt/compute"
	"github.com/Azure/go-autorest/autorest/azure"
	"github.com/Azure/go-autorest/autorest/to"
	"github.com/hashicorp/packer/builder/azure/common/client"
)

type DiskAttacher interface {
	AttachDisk(ctx context.Context, disk string) (lun int32, err error)
	WaitForDevice(ctx context.Context, i int32) (device string, err error)
	DetachDisk(ctx context.Context, disk string) (err error)
	WaitForDetach(ctx context.Context, diskID string) error
}

var NewDiskAttacher = func(azureClient client.AzureClientSet) DiskAttacher {
	return &diskAttacher{
		azcli: azureClient,
	}
}

type diskAttacher struct {
	azcli client.AzureClientSet

	vm *client.ComputeInfo // store info about this VM so that we don't have to ask metadata service on every call
}

var DiskNotFoundError = errors.New("Disk not found")

func (da *diskAttacher) DetachDisk(ctx context.Context, diskID string) error {
	log.Println("Fetching list of disks currently attached to VM")
	currentDisks, err := da.getDisks(ctx)
	if err != nil {
		return err
	}

	log.Printf("Removing %q from list of disks currently attached to VM", diskID)
	newDisks := []compute.DataDisk{}
	for _, disk := range currentDisks {
		if disk.ManagedDisk != nil &&
			!strings.EqualFold(to.String(disk.ManagedDisk.ID), diskID) {
			newDisks = append(newDisks, disk)
		}
	}
	if len(currentDisks) == len(newDisks) {
		return DiskNotFoundError
	}

	log.Println("Updating new list of disks attached to VM")
	err = da.setDisks(ctx, newDisks)
	if err != nil {
		return err
	}

	return nil
}

func (da *diskAttacher) WaitForDetach(ctx context.Context, diskID string) error {
	for { // loop until disk is not attached, timeout or error
		list, err := da.getDisks(ctx)
		if err != nil {
			return err
		}
		if findDiskInList(list, diskID) == nil {
			log.Println("Disk is no longer in VM model, assuming detached")
			return nil
		}

		select {
		case <-time.After(time.Second): //continue
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (da *diskAttacher) AttachDisk(ctx context.Context, diskID string) (int32, error) {
	dataDisks, err := da.getDisks(ctx)
	if err != nil {
		return -1, err
	}

	// check to see if disk is already attached, remember lun if found
	if disk := findDiskInList(dataDisks, diskID); disk != nil {
		// disk is already attached, just take this lun
		if disk.Lun == nil {
			return -1, errors.New("disk is attached, but lun was not set in VM model (possibly an error in the Azure APIs)")
		}
		return to.Int32(disk.Lun), nil
	}

	// disk was not found on VM, go and actually attach it

	var lun int32 = -1
findFreeLun:
	for lun = 0; lun < 64; lun++ {
		for _, v := range dataDisks {
			if to.Int32(v.Lun) == lun {
				continue findFreeLun
			}
		}
		// no datadisk is using this lun
		break
	}

	// append new data disk to collection
	dataDisks = append(dataDisks, compute.DataDisk{
		CreateOption: compute.DiskCreateOptionTypesAttach,
		ManagedDisk: &compute.ManagedDiskParameters{
			ID: to.StringPtr(diskID),
		},
		Lun: to.Int32Ptr(lun),
	})

	// prepare resource object for update operation
	err = da.setDisks(ctx, dataDisks)
	if err != nil {
		return -1, err
	}

	return lun, nil
}

func (da *diskAttacher) getThisVM(ctx context.Context) (compute.VirtualMachine, error) {
	// getting resource info for this VM
	if da.vm == nil {
		vm, err := da.azcli.MetadataClient().GetComputeInfo()
		if err != nil {
			return compute.VirtualMachine{}, err
		}
		da.vm = vm
	}

	// retrieve actual VM
	vmResource, err := da.azcli.VirtualMachinesClient().Get(ctx, da.vm.ResourceGroupName, da.vm.Name, "")
	if err != nil {
		return compute.VirtualMachine{}, err
	}
	if vmResource.StorageProfile == nil {
		return compute.VirtualMachine{}, errors.New("properties.storageProfile is not set on VM, this is unexpected")
	}

	return vmResource, nil
}

func (da diskAttacher) getDisks(ctx context.Context) ([]compute.DataDisk, error) {
	vmResource, err := da.getThisVM(ctx)
	if err != nil {
		return []compute.DataDisk{}, err
	}

	return *vmResource.StorageProfile.DataDisks, nil
}

func (da diskAttacher) setDisks(ctx context.Context, disks []compute.DataDisk) error {
	vmResource, err := da.getThisVM(ctx)
	if err != nil {
		return err
	}

	id, err := azure.ParseResourceID(to.String(vmResource.ID))
	if err != nil {
		return err
	}

	vmResource.StorageProfile.DataDisks = &disks
	vmResource.Resources = nil

	// update the VM resource, attach disk
	_, err = da.azcli.VirtualMachinesClient().CreateOrUpdate(ctx, id.ResourceGroup, id.ResourceName, vmResource)

	return err
}

func findDiskInList(list []compute.DataDisk, diskID string) *compute.DataDisk {
	for _, disk := range list {
		if disk.ManagedDisk != nil &&
			strings.EqualFold(to.String(disk.ManagedDisk.ID), diskID) {
			return &disk
		}
	}
	return nil
}
<|endoftext|>