<commit_before>package helpers

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/runner"
)

type Config struct {
	ApiEndpoint string `json:"api"`
	AppsDomain  string `json:"apps_domain"`
	UseHttp     bool   `json:"use_http"`

	AdminUser     string `json:"admin_user"`
	AdminPassword string `json:"admin_password"`

	UseExistingUser      bool   `json:"use_existing_user"`
	ShouldKeepUser       bool   `json:"keep_user_at_suite_end"`
	ExistingUser         string `json:"existing_user"`
	ExistingUserPassword string `json:"existing_user_password"`

	ConfigurableTestPassword string `json:"test_password"`

	PersistentAppHost      string `json:"persistent_app_host"`
	PersistentAppSpace     string `json:"persistent_app_space"`
	PersistentAppOrg       string `json:"persistent_app_org"`
	PersistentAppQuotaName string `json:"persistent_app_quota_name"`

	SkipSSLValidation                 bool   `json:"skip_ssl_validation"`
	Backend                           string `json:"backend"`
	IncludeDiegoDocker                bool   `json:"include_diego_docker"`
	IncludeTasks                      bool   `json:"include_tasks"`
	IncludePrivilegedContainerSupport bool   `json:"include_privileged_container_support"`
	IncludeSSO                        bool   `json:"include_sso"`

	ArtifactsDirectory string `json:"artifacts_directory"`

	DefaultTimeout     time.Duration `json:"default_timeout"`
	SleepTimeout       time.Duration `json:"sleep_timeout"`
	DetectTimeout      time.Duration `json:"detect_timeout"`
	CfPushTimeout      time.Duration `json:"cf_push_timeout"`
	LongCurlTimeout    time.Duration `json:"long_curl_timeout"`
	BrokerStartTimeout time.Duration `json:"broker_start_timeout"`

	TimeoutScale float64 `json:"timeout_scale"`

	SecureAddress string `json:"secure_address"`

	DockerExecutable      string   `json:"docker_executable"`
	DockerParameters      []string `json:"docker_parameters"`
	DockerRegistryAddress string   `json:"docker_registry_address"`
	DockerPrivateImage    string   `json:"docker_private_image"`
	DockerUser            string   `json:"docker_user"`
	DockerPassword        string   `json:"docker_password"`
	DockerEmail           string   `json:"docker_email"`

	StaticFileBuildpackName string `json:"staticfile_buildpack_name"`
	JavaBuildpackName       string `json:"java_buildpack_name"`
	RubyBuildpackName       string `json:"ruby_buildpack_name"`
	NodejsBuildpackName     string `json:"nodejs_buildpack_name"`
	GoBuildpackName         string `json:"go_buildpack_name"`
	PythonBuildpackName     string `json:"python_buildpack_name"`
	PhpBuildpackName        string `json:"php_buildpack_name"`
	BinaryBuildpackName     string `json:"binary_buildpack_name"`

	NamePrefix string `json:"name_prefix"`
}

var defaults = Config{
	PersistentAppHost:      "CATS-persistent-app",
	PersistentAppSpace:     "CATS-persistent-space",
	PersistentAppOrg:       "CATS-persistent-org",
	PersistentAppQuotaName: "CATS-persistent-quota",

	StaticFileBuildpackName: "staticfile_buildpack",
	JavaBuildpackName:       "java_buildpack",
	RubyBuildpackName:       "ruby_buildpack",
	NodejsBuildpackName:     "nodejs_buildpack",
	GoBuildpackName:         "go_buildpack",
	PythonBuildpackName:     "python_buildpack",
	PhpBuildpackName:        "php_buildpack",
	BinaryBuildpackName:     "binary_buildpack",

	ArtifactsDirectory: filepath.Join("..", "results"),

	NamePrefix: "CATS",
}

func (c Config) ScaledTimeout(timeout time.Duration) time.Duration {
	return time.Duration(float64(timeout) * c.TimeoutScale)
}

var loadedConfig *Config

func Load(path string, config interface{}) error {
	c, ok := config.(*Config)
	if !ok {
		val := reflect.ValueOf(config).Elem().FieldByName("Config").Addr()
		c = val.Interface().(*Config)
	}

	*c = defaults
	err := loadConfigFromPath(path, config)
	if err != nil {
		return err
	}

	if c.ApiEndpoint == "" {
		return fmt.Errorf("missing configuration 'api'")
	}

	if c.AdminUser == "" {
		return fmt.Errorf("missing configuration 'admin_user'")
	}

	if c.AdminPassword == "" {
		return fmt.Errorf("missing configuration 'admin_password'")
	}

	if c.TimeoutScale <= 0 {
		c.TimeoutScale = 1.0
	}

	runner.SkipSSLValidation = c.SkipSSLValidation
	return nil
}

func LoadConfig() Config {
	if loadedConfig != nil {
		return *loadedConfig
	}

	var config Config

	err := Load(ConfigPath(), &config)
	if err != nil {
		panic(err)
	}

	loadedConfig = &config
	return config
}

func (c Config) Protocol() string {
	if c.UseHttp {
		return "http://"
	} else {
		return "https://"
	}
}

func loadConfigFromPath(path string, config interface{}) error {
	configFile, err := os.Open(path)
	if err != nil {
		return err
	}
	defer configFile.Close()

	decoder := json.NewDecoder(configFile)
	return decoder.Decode(config)
}

func ConfigPath() string {
	path := os.Getenv("CONFIG")
	if path == "" {
		panic("Must set $CONFIG to point to an integration config .json file.")
	}

	return path
}
<commit_msg>Add new include suites to config params<commit_after>package helpers

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"reflect"
	"time"

	"github.com/cloudfoundry-incubator/cf-test-helpers/runner"
)

type Config struct {
	ApiEndpoint string `json:"api"`
	AppsDomain  string `json:"apps_domain"`
	UseHttp     bool   `json:"use_http"`

	AdminUser     string `json:"admin_user"`
	AdminPassword string `json:"admin_password"`

	UseExistingUser      bool   `json:"use_existing_user"`
	ShouldKeepUser       bool   `json:"keep_user_at_suite_end"`
	ExistingUser         string `json:"existing_user"`
	ExistingUserPassword string `json:"existing_user_password"`

	ConfigurableTestPassword string `json:"test_password"`

	PersistentAppHost      string `json:"persistent_app_host"`
	PersistentAppSpace     string `json:"persistent_app_space"`
	PersistentAppOrg       string `json:"persistent_app_org"`
	PersistentAppQuotaName string `json:"persistent_app_quota_name"`

	SkipSSLValidation                 bool   `json:"skip_ssl_validation"`
	Backend                           string `json:"backend"`
	IncludeDiegoDocker                bool   `json:"include_diego_docker"`
	IncludeTasks                      bool   `json:"include_tasks"`
	IncludePrivilegedContainerSupport bool   `json:"include_privileged_container_support"`
	IncludeSSO                        bool   `json:"include_sso"`

	ArtifactsDirectory string `json:"artifacts_directory"`

	DefaultTimeout     time.Duration `json:"default_timeout"`
	SleepTimeout       time.Duration `json:"sleep_timeout"`
	DetectTimeout      time.Duration `json:"detect_timeout"`
	CfPushTimeout      time.Duration `json:"cf_push_timeout"`
	LongCurlTimeout    time.Duration `json:"long_curl_timeout"`
	BrokerStartTimeout time.Duration `json:"broker_start_timeout"`

	TimeoutScale float64 `json:"timeout_scale"`

	SecureAddress string `json:"secure_address"`

	DockerExecutable      string   `json:"docker_executable"`
	DockerParameters      []string `json:"docker_parameters"`
	DockerRegistryAddress string   `json:"docker_registry_address"`
	DockerPrivateImage    string   `json:"docker_private_image"`
	DockerUser            string   `json:"docker_user"`
	DockerPassword        string   `json:"docker_password"`
	DockerEmail           string   `json:"docker_email"`

	StaticFileBuildpackName string `json:"staticfile_buildpack_name"`
	JavaBuildpackName       string `json:"java_buildpack_name"`
	RubyBuildpackName       string `json:"ruby_buildpack_name"`
	NodejsBuildpackName     string `json:"nodejs_buildpack_name"`
	GoBuildpackName         string `json:"go_buildpack_name"`
	PythonBuildpackName     string `json:"python_buildpack_name"`
	PhpBuildpackName        string `json:"php_buildpack_name"`
	BinaryBuildpackName     string `json:"binary_buildpack_name"`

	IncludeBackendCompatiblity bool `json:"include_backend_compatibility"`
	IncludeDetect              bool `json:"include_detect"`
	IncludeDocker              bool `json:"include_docker"`
	IncludeInternetDependent   bool `json:"include_internet_dependent"`
	IncludeRouteServices       bool `json:"include_route_services"`
	IncludeRouting             bool `json:"include_routing"`
	IncludeSecurityGroups      bool `json:"include_security_groups"`
	IncludeServices            bool `json:"include_services"`
	IncludeSsh                 bool `json:"include_ssh"`
	IncludeV3                  bool `json:"include_v3"`

	NamePrefix string `json:"name_prefix"`
}

var defaults = Config{
	PersistentAppHost:      "CATS-persistent-app",
	PersistentAppSpace:     "CATS-persistent-space",
	PersistentAppOrg:       "CATS-persistent-org",
	PersistentAppQuotaName: "CATS-persistent-quota",

	StaticFileBuildpackName: "staticfile_buildpack",
	JavaBuildpackName:       "java_buildpack",
	RubyBuildpackName:       "ruby_buildpack",
	NodejsBuildpackName:     "nodejs_buildpack",
	GoBuildpackName:         "go_buildpack",
	PythonBuildpackName:     "python_buildpack",
	PhpBuildpackName:        "php_buildpack",
	BinaryBuildpackName:     "binary_buildpack",

	IncludeBackendCompatiblity: true,
	IncludeDetect:              true,
	IncludeDocker:              true,
	IncludeInternetDependent:   true,
	IncludeRouteServices:       true,
	IncludeRouting:             true,
	IncludeSecurityGroups:      true,
	IncludeServices:            true,
	IncludeSsh:                 true,
	IncludeV3:                  true,

	ArtifactsDirectory: filepath.Join("..", "results"),

	NamePrefix: "CATS",
}

func (c Config) ScaledTimeout(timeout time.Duration) time.Duration {
	return time.Duration(float64(timeout) * c.TimeoutScale)
}

var loadedConfig *Config

func Load(path string, config interface{}) error {
	c, ok := config.(*Config)
	if !ok {
		val := reflect.ValueOf(config).Elem().FieldByName("Config").Addr()
		c = val.Interface().(*Config)
	}

	*c = defaults
	err := loadConfigFromPath(path, config)
	if err != nil {
		return err
	}

	if c.ApiEndpoint == "" {
		return fmt.Errorf("missing configuration 'api'")
	}

	if c.AdminUser == "" {
		return fmt.Errorf("missing configuration 'admin_user'")
	}

	if c.AdminPassword == "" {
		return fmt.Errorf("missing configuration 'admin_password'")
	}

	if c.TimeoutScale <= 0 {
		c.TimeoutScale = 1.0
	}

	runner.SkipSSLValidation = c.SkipSSLValidation
	return nil
}

func LoadConfig() Config {
	if loadedConfig != nil {
		return *loadedConfig
	}

	var config Config

	err := Load(ConfigPath(), &config)
	if err != nil {
		panic(err)
	}

	loadedConfig = &config
	return config
}

func (c Config) Protocol() string {
	if c.UseHttp {
		return "http://"
	} else {
		return "https://"
	}
}

func loadConfigFromPath(path string, config interface{}) error {
	configFile, err := os.Open(path)
	if err != nil {
		return err
	}
	defer configFile.Close()

	decoder := json.NewDecoder(configFile)
	return decoder.Decode(config)
}

func ConfigPath() string {
	path := os.Getenv("CONFIG")
	if path == "" {
		panic("Must set $CONFIG to point to an integration config .json file.")
	}

	return path
}
<|endoftext|>


<commit_before>package mss

import (
	"fmt"

	"github.com/omniscale/magnacarto/color"
)

type codeType uint8

const (
	typeUnknown codeType = iota
	typeVar
	typeNum
	typePercent
	typeColor
	typeBool
	typeFunction
	typeFunctionEnd
	typeURL
	typeKeyword
	typeField
	typeFieldExpr
	typeString
	typeList
	typeStop

	typeNegation
	typeAdd
	typeSubtract
	typeMultiply
	typeDivide
)

func (t codeType) String() string {
	switch t {
	case typeNegation:
		return "!"
	case typeAdd:
		return "+"
	case typeSubtract:
		return "-"
	case typeMultiply:
		return "*"
	case typeDivide:
		return "/"
	case typeVar:
		return "v"
	case typeNum:
		return "n"
	case typePercent:
		return "%"
	case typeColor:
		return "c"
	case typeBool:
		return "b"
	case typeFunction:
		return "{"
	case typeFunctionEnd:
		return "}"
	case typeURL:
		return "@"
	case typeKeyword:
		return "#"
	case typeField:
		return "["
	case typeList:
		return "L"
	case typeString:
		return "\""
	case typeStop:
		return "S"
	case typeUnknown:
		return "?"
	default:
		return fmt.Sprintf("%#v", t)
	}
}

type expression struct {
	code []code
	pos  position
}

func (e *expression) addOperator(t codeType) {
	e.code = append(e.code, code{T: t})
}

func (e *expression) addValue(val interface{}, t codeType) {
	e.code = append(e.code, code{T: t, Value: val})
}

func (e *expression) clear() {
	e.code = e.code[:0]
}

func (e *expression) evaluate() (Value, error) {
	codes, _, err := evaluate(e.code)
	if err != nil {
		return nil, err
	}
	if len(codes) > 1 {
		// create copy since c points to internal code slice
		l := make([]Value, 0, len(codes))
		for _, c := range codes {
			l = append(l, c.Value)
		}
		return l, nil
	}
	return codes[0].Value, nil
}

type Field string

func evaluate(codes []code) ([]code, int, error) {
	top := 0
	for i := 0; i < len(codes); i++ {
		c := codes[i]
		switch c.T {
		case typeNum, typeColor, typePercent, typeString, typeKeyword, typeURL, typeBool, typeField, typeList:
			codes[top] = c
			top++
			continue
		case typeNegation:
			a := codes[top-1]
			a.Value = -a.Value.(float64)
			codes[top-1] = a
			continue
		case typeFunction:
			v, parsed, err := evaluate(codes[i+1:])
			i += parsed + 1
			if err != nil {
				return v, 0, err
			}
			if colorF, ok := colorFuncs[c.Value.(string)]; ok {
				if len(v) != 2 {
					return nil, 0, fmt.Errorf("function %s takes exactly two arguments, got %d", c.Value.(string), len(v))
				}
				if v[0].T != typeColor {
					return nil, 0, fmt.Errorf("function %s requires color as first argument, got %v", c.Value.(string), v[0])
				}
				if v[1].T != typeNum && v[1].T != typePercent {
					return nil, 0, fmt.Errorf("function %s requires number/percent as second argument, got %v", c.Value.(string), v[1])
				}
				v = []code{{Value: colorF(v[0].Value.(color.RGBA), v[1].Value.(float64)/100), T: typeColor}}
			} else if c.Value.(string) == "-mc-set-hue" {
				if len(v) != 2 {
					return nil, 0, fmt.Errorf("function %s takes exactly two arguments, got %d", c.Value.(string), len(v))
				}
				if v[0].T != typeColor {
					return nil, 0, fmt.Errorf("function %s requires color as first argument, got %v", c.Value.(string), v[0])
				}
				if v[1].T != typeColor {
					return nil, 0, fmt.Errorf("function %s requires color as second argument, got %v", c.Value.(string), v[1])
				}
				v = []code{{Value: color.SetHue(v[0].Value.(color.RGBA), v[1].Value.(color.RGBA)), T: typeColor}}
			} else if c.Value.(string) == "rgb" || c.Value.(string) == "rgba" {
				if c.Value.(string) == "rgb" && len(v) != 3 {
					return nil, 0, fmt.Errorf("rgb takes exactly three arguments, got %d", len(v))
				}
				if c.Value.(string) == "rgba" && len(v) != 4 {
					return nil, 0, fmt.Errorf("rgba takes exactly four arguments, got %d", len(v))
				}
				c := [4]float64{1, 1, 1, 1}
				for i := range v {
					if v[i].T == typeNum {
						if i < 3 {
							c[i] = v[i].Value.(float64) / 255
						} else {
							c[i] = v[i].Value.(float64) // alpha value is from 0.0-1.0
							if c[i] > 1.0 {
								c[i] /= 255 // TODO or clamp? compat with Carto?
							}
						}
					} else if v[i].T == typePercent {
						c[i] = v[i].Value.(float64) / 100
					} else {
						return nil, 0, fmt.Errorf("rgb/rgba takes float or percent arguments only, got %v", v[i])
					}
					if c[i] < 0 {
						c[i] = 0
					} else if c[i] > 255 {
						c[i] = 255
					}
				}
				v = []code{{
					Value: color.RGBA{
						c[0],
						c[1],
						c[2],
						c[3],
					},
					T: typeColor}}
			} else if c.Value.(string) == "stop" {
				if len(v) != 2 {
					return nil, 0, fmt.Errorf("stop takes exactly two arguments, got %d", len(v))
				}
				if v[0].T != typeNum {
					return nil, 0, fmt.Errorf("stop takes int as first argument only, got %v", v[i])
				}
				if v[1].T != typeColor {
					return nil, 0, fmt.Errorf("stop takes color as second argument only, got %v", v[i])
				}
				val := int(v[0].Value.(float64))
				c := v[1].Value.(color.RGBA)
				v = []code{{
					Value: Stop{Value: val, Color: c},
					T:     typeStop},
				}
			} else if c.Value.(string) == "__echo__" {
				// pass
			} else {
				return nil, 0, fmt.Errorf("unknown function %s", c.Value.(string))
			}
			for i, v := range v {
				codes[top+i] = v
			}
			top += len(v)
		case typeFunctionEnd:
			return codes[0:top], i, nil
		case typeAdd, typeSubtract, typeMultiply, typeDivide:
			a, b := codes[top-2], codes[top-1]
			top -= 2
			if a.T == typeNum && b.T == typeNum {
				switch c.T {
				case typeAdd:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) + b.Value.(float64)}
				case typeSubtract:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) - b.Value.(float64)}
				case typeMultiply:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) * b.Value.(float64)}
				case typeDivide:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) / b.Value.(float64)}
				}
			} else if c.T == typeAdd && a.T == typeString && b.T == typeString {
				// string concatenation
				codes[top] = code{T: typeString, Value: a.Value.(string) + b.Value.(string)}
			} else if c.T == typeAdd && a.T == typeString && b.T == typeField {
				codes[top] = code{T: typeFieldExpr, Value: []Value{a.Value.(string), Field(b.Value.(string))}}
			} else if c.T == typeAdd && a.T == typeField && b.T == typeString {
				codes[top] = code{T: typeFieldExpr, Value: []Value{Field(a.Value.(string)), b.Value.(string)}}
			} else if c.T == typeAdd && a.T == typeField && b.T == typeField {
				codes[top] = code{T: typeFieldExpr, Value: []Value{Field(a.Value.(string)), Field(b.Value.(string))}}
			} else if c.T == typeAdd && a.T == typeFieldExpr && b.T == typeField {
				codes[top] = code{T: typeFieldExpr, Value: append(a.Value.([]Value), Field(b.Value.(string)))}
			} else if c.T == typeAdd && a.T == typeFieldExpr && b.T == typeString {
				codes[top] = code{T: typeFieldExpr, Value: append(a.Value.([]Value), b.Value.(string))}
			} else if c.T == typeMultiply && a.T == typeColor && b.T == typeNum {
				c := a.Value.(color.RGBA)
				f := b.Value.(float64)
				c = color.Multiply(c, f)
				codes[top] = code{T: typeColor, Value: c}
			} else {
				return nil, 0, fmt.Errorf("unsupported operation %v for %v and %v", c, a, b)
			}
			top++
		}
	}
	return codes[:top], 0, nil
}

type Stop struct {
	Value int
	Color color.RGBA
}

type functype func(args []code) ([]code, error)

var colorFuncs map[string]colorFunc

type colorFunc func(color.RGBA, float64) color.RGBA

func init() {
	colorFuncs = map[string]colorFunc{
		"lighten":    color.Lighten,
		"darken":     color.Darken,
		"saturate":   color.Saturate,
		"desaturate": color.Desaturate,
		"fadein":     color.FadeIn,
		"fadeout":    color.FadeOut,
		"spin":       color.Spin,
	}
}

type code struct {
	T     codeType
	Value interface{}
}
<commit_msg>prevent runtime error when evaluate returns no code<commit_after>package mss

import (
	"fmt"

	"github.com/omniscale/magnacarto/color"
)

type codeType uint8

const (
	typeUnknown codeType = iota
	typeVar
	typeNum
	typePercent
	typeColor
	typeBool
	typeFunction
	typeFunctionEnd
	typeURL
	typeKeyword
	typeField
	typeFieldExpr
	typeString
	typeList
	typeStop

	typeNegation
	typeAdd
	typeSubtract
	typeMultiply
	typeDivide
)

func (t codeType) String() string {
	switch t {
	case typeNegation:
		return "!"
	case typeAdd:
		return "+"
	case typeSubtract:
		return "-"
	case typeMultiply:
		return "*"
	case typeDivide:
		return "/"
	case typeVar:
		return "v"
	case typeNum:
		return "n"
	case typePercent:
		return "%"
	case typeColor:
		return "c"
	case typeBool:
		return "b"
	case typeFunction:
		return "{"
	case typeFunctionEnd:
		return "}"
	case typeURL:
		return "@"
	case typeKeyword:
		return "#"
	case typeField:
		return "["
	case typeList:
		return "L"
	case typeString:
		return "\""
	case typeStop:
		return "S"
	case typeUnknown:
		return "?"
	default:
		return fmt.Sprintf("%#v", t)
	}
}

type expression struct {
	code []code
	pos  position
}

func (e *expression) addOperator(t codeType) {
	e.code = append(e.code, code{T: t})
}

func (e *expression) addValue(val interface{}, t codeType) {
	e.code = append(e.code, code{T: t, Value: val})
}

func (e *expression) clear() {
	e.code = e.code[:0]
}

func (e *expression) evaluate() (Value, error) {
	codes, _, err := evaluate(e.code)
	if err != nil {
		return nil, err
	}
	if len(codes) == 0 {
		return nil, fmt.Errorf("unable to evaluate expression")
	}
	if len(codes) > 1 {
		// create copy since c points to internal code slice
		l := make([]Value, 0, len(codes))
		for _, c := range codes {
			l = append(l, c.Value)
		}
		return l, nil
	}
	return codes[0].Value, nil
}

type Field string

func evaluate(codes []code) ([]code, int, error) {
	top := 0
	for i := 0; i < len(codes); i++ {
		c := codes[i]
		switch c.T {
		case typeNum, typeColor, typePercent, typeString, typeKeyword, typeURL, typeBool, typeField, typeList:
			codes[top] = c
			top++
			continue
		case typeNegation:
			a := codes[top-1]
			a.Value = -a.Value.(float64)
			codes[top-1] = a
			continue
		case typeFunction:
			v, parsed, err := evaluate(codes[i+1:])
			i += parsed + 1
			if err != nil {
				return v, 0, err
			}
			if colorF, ok := colorFuncs[c.Value.(string)]; ok {
				if len(v) != 2 {
					return nil, 0, fmt.Errorf("function %s takes exactly two arguments, got %d", c.Value.(string), len(v))
				}
				if v[0].T != typeColor {
					return nil, 0, fmt.Errorf("function %s requires color as first argument, got %v", c.Value.(string), v[0])
				}
				if v[1].T != typeNum && v[1].T != typePercent {
					return nil, 0, fmt.Errorf("function %s requires number/percent as second argument, got %v", c.Value.(string), v[1])
				}
				v = []code{{Value: colorF(v[0].Value.(color.RGBA), v[1].Value.(float64)/100), T: typeColor}}
			} else if c.Value.(string) == "-mc-set-hue" {
				if len(v) != 2 {
					return nil, 0, fmt.Errorf("function %s takes exactly two arguments, got %d", c.Value.(string), len(v))
				}
				if v[0].T != typeColor {
					return nil, 0, fmt.Errorf("function %s requires color as first argument, got %v", c.Value.(string), v[0])
				}
				if v[1].T != typeColor {
					return nil, 0, fmt.Errorf("function %s requires color as second argument, got %v", c.Value.(string), v[1])
				}
				v = []code{{Value: color.SetHue(v[0].Value.(color.RGBA), v[1].Value.(color.RGBA)), T: typeColor}}
			} else if c.Value.(string) == "rgb" || c.Value.(string) == "rgba" {
				if c.Value.(string) == "rgb" && len(v) != 3 {
					return nil, 0, fmt.Errorf("rgb takes exactly three arguments, got %d", len(v))
				}
				if c.Value.(string) == "rgba" && len(v) != 4 {
					return nil, 0, fmt.Errorf("rgba takes exactly four arguments, got %d", len(v))
				}
				c := [4]float64{1, 1, 1, 1}
				for i := range v {
					if v[i].T == typeNum {
						if i < 3 {
							c[i] = v[i].Value.(float64) / 255
						} else {
							c[i] = v[i].Value.(float64) // alpha value is from 0.0-1.0
							if c[i] > 1.0 {
								c[i] /= 255 // TODO or clamp? compat with Carto?
							}
						}
					} else if v[i].T == typePercent {
						c[i] = v[i].Value.(float64) / 100
					} else {
						return nil, 0, fmt.Errorf("rgb/rgba takes float or percent arguments only, got %v", v[i])
					}
					if c[i] < 0 {
						c[i] = 0
					} else if c[i] > 255 {
						c[i] = 255
					}
				}
				v = []code{{
					Value: color.RGBA{
						c[0],
						c[1],
						c[2],
						c[3],
					},
					T: typeColor}}
			} else if c.Value.(string) == "stop" {
				if len(v) != 2 {
					return nil, 0, fmt.Errorf("stop takes exactly two arguments, got %d", len(v))
				}
				if v[0].T != typeNum {
					return nil, 0, fmt.Errorf("stop takes int as first argument only, got %v", v[i])
				}
				if v[1].T != typeColor {
					return nil, 0, fmt.Errorf("stop takes color as second argument only, got %v", v[i])
				}
				val := int(v[0].Value.(float64))
				c := v[1].Value.(color.RGBA)
				v = []code{{
					Value: Stop{Value: val, Color: c},
					T:     typeStop},
				}
			} else if c.Value.(string) == "__echo__" {
				// pass
			} else {
				return nil, 0, fmt.Errorf("unknown function %s", c.Value.(string))
			}
			for i, v := range v {
				codes[top+i] = v
			}
			top += len(v)
		case typeFunctionEnd:
			return codes[0:top], i, nil
		case typeAdd, typeSubtract, typeMultiply, typeDivide:
			a, b := codes[top-2], codes[top-1]
			top -= 2
			if a.T == typeNum && b.T == typeNum {
				switch c.T {
				case typeAdd:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) + b.Value.(float64)}
				case typeSubtract:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) - b.Value.(float64)}
				case typeMultiply:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) * b.Value.(float64)}
				case typeDivide:
					codes[top] = code{T: typeNum, Value: a.Value.(float64) / b.Value.(float64)}
				}
			} else if c.T == typeAdd && a.T == typeString && b.T == typeString {
				// string concatenation
				codes[top] = code{T: typeString, Value: a.Value.(string) + b.Value.(string)}
			} else if c.T == typeAdd && a.T == typeString && b.T == typeField {
				codes[top] = code{T: typeFieldExpr, Value: []Value{a.Value.(string), Field(b.Value.(string))}}
			} else if c.T == typeAdd && a.T == typeField && b.T == typeString {
				codes[top] = code{T: typeFieldExpr, Value: []Value{Field(a.Value.(string)), b.Value.(string)}}
			} else if c.T == typeAdd && a.T == typeField && b.T == typeField {
				codes[top] = code{T: typeFieldExpr, Value: []Value{Field(a.Value.(string)), Field(b.Value.(string))}}
			} else if c.T == typeAdd && a.T == typeFieldExpr && b.T == typeField {
				codes[top] = code{T: typeFieldExpr, Value: append(a.Value.([]Value), Field(b.Value.(string)))}
			} else if c.T == typeAdd && a.T == typeFieldExpr && b.T == typeString {
				codes[top] = code{T: typeFieldExpr, Value: append(a.Value.([]Value), b.Value.(string))}
			} else if c.T == typeMultiply && a.T == typeColor && b.T == typeNum {
				c := a.Value.(color.RGBA)
				f := b.Value.(float64)
				c = color.Multiply(c, f)
				codes[top] = code{T: typeColor, Value: c}
			} else {
				return nil, 0, fmt.Errorf("unsupported operation %v for %v and %v", c, a, b)
			}
			top++
		}
	}
	return codes[:top], 0, nil
}

type Stop struct {
	Value int
	Color color.RGBA
}

type functype func(args []code) ([]code, error)

var colorFuncs map[string]colorFunc

type colorFunc func(color.RGBA, float64) color.RGBA

func init() {
	colorFuncs = map[string]colorFunc{
		"lighten":    color.Lighten,
		"darken":     color.Darken,
		"saturate":   color.Saturate,
		"desaturate": color.Desaturate,
		"fadein":     color.FadeIn,
		"fadeout":    color.FadeOut,
		"spin":       color.Spin,
	}
}

type code struct {
	T     codeType
	Value interface{}
}
<|endoftext|>


<commit_before>// Tool for compiling dub sources.
package main

import (
	"evergreen/compiler"
	"evergreen/dub/flow"
	"evergreen/dub/transform"
	"evergreen/dub/transform/golang"
	"evergreen/dub/tree"
	gocore "evergreen/go/core"
	goflow "evergreen/go/flow"
	gotransform "evergreen/go/transform"
	gotree "evergreen/go/tree"
	"evergreen/graph"
	"evergreen/io"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"
)

func dumpProgram(status compiler.PassStatus, runner *compiler.TaskRunner, program *flow.DubProgram, outputDir []string) {
	status.Begin()
	defer status.End()

	for _, dubPkg := range program.Packages {
		for _, f := range dubPkg.Funcs {
			styler := &flow.DotStyler{Decl: f, Core: program.Core}
			dot := graph.GraphToDot(f.CFG, styler)
			parts := append(outputDir, "dub_frontend")
			parts = append(parts, dubPkg.Path...)
			parts = append(parts, fmt.Sprintf("%s.svg", f.Name))
			outfile := filepath.Join(parts...)

			runner.Run(func() {
				io.WriteDot(dot, outfile)
			})
		}
	}
}

func analyizeProgram(program *flow.DubProgram) {
	for _, dubPkg := range program.Packages {
		for _, s := range dubPkg.Structs {
			if s.Implements != nil {
				s.Implements.IsParent = true
			}
		}
	}
}

func goFlowFuncName(cf *gocore.Function, f *goflow.FlowFunc) string {
	base := "func"
	if f.Recv != goflow.NoRegister {
		ref := f.Register_Scope.Get(f.Recv)
		at := ref.T
		pt, ok := at.(*gocore.PointerType)
		if !ok {
			panic(at)
		}
		st, ok := pt.Element.(*gocore.StructType)
		if !ok {
			panic(pt.Element)
		}
		base = "meth_" + st.Name
	}
	return base + "_" + cf.Name
}

func dumpFlowFuncs(status compiler.PassStatus, runner *compiler.TaskRunner, goFlowProgram *goflow.FlowProgram, goCoreProg *gocore.CoreProgram, outputDir []string) {
	status.Begin()
	defer status.End()

	iter := goFlowProgram.FlowFunc_Scope.Iter()
	for iter.Next() {
		fIndex, f := iter.Value()
		cf := goCoreProg.Function_Scope.Get(gocore.Function_Ref(fIndex))
		dot := graph.GraphToDot(f.CFG, &goflow.DotStyler{Func: f, Core: goCoreProg})
		parts := append(outputDir, "dub_to_go")
		p := goCoreProg.Package_Scope.Get(cf.Package)
		parts = append(parts, p.Path...)
		parts = append(parts, fmt.Sprintf("%s.svg", goFlowFuncName(cf, f)))
		outfile := filepath.Join(parts...)

		runner.Run(func() {
			io.WriteDot(dot, outfile)
		})
	}
}

func processProgram(status compiler.PassStatus, p compiler.LocationProvider, runner *compiler.TaskRunner, config *EGCConfig) {
	program, coreProg := tree.DubProgramFrontend(status.Pass("dub_frontend"), p, config.InputDir)
	if status.ShouldHalt() {
		return
	}
	flowProgram := transform.LowerProgram(status.Pass("lower"), program, coreProg)

	flow.TrimFlow(status.Pass("trim_flow"), flowProgram)

	if config.Dump {
		dumpProgram(status.Pass("dump"), runner, flowProgram, config.DumpDir)
	}

	analyizeProgram(flowProgram)

	goFlowProg, goCoreProg, bypass := golang.GenerateGo(status.Pass("dub_to_go"), flowProgram, coreProg, config.RootPackage, config.GenerateTests)
	if config.Dump {
		dumpFlowFuncs(status.Pass("dump_go"), runner, goFlowProg, goCoreProg, config.DumpDir)
	}
	goTreeProg := gotransform.FlowToTree(status.Pass("flow_to_tree"), goFlowProg, goCoreProg, bypass)

	gotree.GoProgramBackend(status.Pass("go_backend"), goTreeProg, goCoreProg, config.OutputDir, runner)
}

func entryPoint(p compiler.LocationProvider, status compiler.PassStatus, config *EGCConfig) {
	status.Begin()
	defer status.End()

	runner := compiler.CreateTaskRunner(config.Jobs)

	processProgram(status, p, runner, config)
	runner.Kill()
}

func mainLoop(config *EGCConfig, profiling bool) {
	p := compiler.MakeProvider()
	status := compiler.MakeStatus(p)

	start := time.Now()
	for i := 0; ; i++ {
		entryPoint(p, status.Pass("egc"), config)
		if profiling && time.Since(start) < time.Second*10 {
			fmt.Println("Re-running to improve profiling data", i)
		} else {
			break
		}
	}

	if status.ShouldHalt() {
		fmt.Printf("%d errors\n", status.ErrorCount())
		os.Exit(1)
	}
}

type EGCConfig struct {
	Dump          bool
	InputDir      string
	OutputDir     string
	RootPackage   []string
	DumpDir       []string
	GenerateTests bool
	Jobs          int
}

func flagError(message string) {
	fmt.Println(message)
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	config := &EGCConfig{
		DumpDir: []string{"output"},
	}

	var rootPackage string
	var verbosity int
	var cpuprofile string
	var memprofile string

	flag.StringVar(&config.InputDir, "indir", "", "Directory containing input files.")
	flag.StringVar(&config.OutputDir, "outdir", "", "Directory to output generated files.")
	flag.StringVar(&rootPackage, "gopackage", "", "Root package for generate sources.")

	flag.BoolVar(&config.Dump, "dump", false, "Dump flowgraphs to output/... (requires graphviz).")
	flag.BoolVar(&config.GenerateTests, "gentests", false, "Generate dub tests.")

	flag.StringVar(&cpuprofile, "cpuprofile", "", "Write cpu profile to file.")
	flag.StringVar(&memprofile, "memprofile", "", "Write memory profile to this file.")
	flag.IntVar(&config.Jobs, "j", runtime.NumCPU(), "Number of threads.")
	flag.IntVar(&verbosity, "v", 0, "Verbosity level.")

	flag.Parse()

	if config.InputDir == "" {
		flagError("-indir is required")
	}
	if config.OutputDir == "" {
		flagError("-outdir is required")
	}
	if rootPackage == "" {
		flagError("-gopackage is required")
	}
	config.RootPackage = strings.Split(rootPackage, "/")

	runtime.GOMAXPROCS(config.Jobs)
	compiler.Verbosity = verbosity

	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			fmt.Println(err.Error())
			return
		} else {
			pprof.StartCPUProfile(f)
			defer pprof.StopCPUProfile()
		}
	}

	mainLoop(config, cpuprofile != "")

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			fmt.Println(err.Error())
			return
		} else {
			pprof.WriteHeapProfile(f)
			f.Close()
		}
	}
}
<commit_msg>Stop dumping graphs of empty functions.<commit_after>// Tool for compiling dub sources.
package main

import (
	"evergreen/compiler"
	"evergreen/dub/flow"
	"evergreen/dub/transform"
	"evergreen/dub/transform/golang"
	"evergreen/dub/tree"
	gocore "evergreen/go/core"
	goflow "evergreen/go/flow"
	gotransform "evergreen/go/transform"
	gotree "evergreen/go/tree"
	"evergreen/graph"
	"evergreen/io"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"strings"
	"time"
)

func dumpProgram(status compiler.PassStatus, runner *compiler.TaskRunner, program *flow.DubProgram, outputDir []string) {
	status.Begin()
	defer status.End()

	for _, dubPkg := range program.Packages {
		for _, f := range dubPkg.Funcs {
			// Don't output empty functions.
			if f.CFG.NumNodes() <= 2 {
				continue
			}
			styler := &flow.DotStyler{Decl: f, Core: program.Core}
			dot := graph.GraphToDot(f.CFG, styler)
			parts := append(outputDir, "dub_frontend")
			parts = append(parts, dubPkg.Path...)
			parts = append(parts, fmt.Sprintf("%s.svg", f.Name))
			outfile := filepath.Join(parts...)

			runner.Run(func() {
				io.WriteDot(dot, outfile)
			})
		}
	}
}

func analyizeProgram(program *flow.DubProgram) {
	for _, dubPkg := range program.Packages {
		for _, s := range dubPkg.Structs {
			if s.Implements != nil {
				s.Implements.IsParent = true
			}
		}
	}
}

func goFlowFuncName(cf *gocore.Function, f *goflow.FlowFunc) string {
	base := "func"
	if f.Recv != goflow.NoRegister {
		ref := f.Register_Scope.Get(f.Recv)
		at := ref.T
		pt, ok := at.(*gocore.PointerType)
		if !ok {
			panic(at)
		}
		st, ok := pt.Element.(*gocore.StructType)
		if !ok {
			panic(pt.Element)
		}
		base = "meth_" + st.Name
	}
	return base + "_" + cf.Name
}

func dumpFlowFuncs(status compiler.PassStatus, runner *compiler.TaskRunner, goFlowProgram *goflow.FlowProgram, goCoreProg *gocore.CoreProgram, outputDir []string) {
	status.Begin()
	defer status.End()

	iter := goFlowProgram.FlowFunc_Scope.Iter()
	for iter.Next() {
		fIndex, f := iter.Value()
		// Don't output empty functions.
		if f.CFG.NumNodes() <= 2 {
			continue
		}
		cf := goCoreProg.Function_Scope.Get(gocore.Function_Ref(fIndex))
		dot := graph.GraphToDot(f.CFG, &goflow.DotStyler{Func: f, Core: goCoreProg})
		parts := append(outputDir, "dub_to_go")
		p := goCoreProg.Package_Scope.Get(cf.Package)
		parts = append(parts, p.Path...)
		parts = append(parts, fmt.Sprintf("%s.svg", goFlowFuncName(cf, f)))
		outfile := filepath.Join(parts...)

		runner.Run(func() {
			io.WriteDot(dot, outfile)
		})
	}
}

func processProgram(status compiler.PassStatus, p compiler.LocationProvider, runner *compiler.TaskRunner, config *EGCConfig) {
	program, coreProg := tree.DubProgramFrontend(status.Pass("dub_frontend"), p, config.InputDir)
	if status.ShouldHalt() {
		return
	}
	flowProgram := transform.LowerProgram(status.Pass("lower"), program, coreProg)

	flow.TrimFlow(status.Pass("trim_flow"), flowProgram)

	if config.Dump {
		dumpProgram(status.Pass("dump"), runner, flowProgram, config.DumpDir)
	}

	analyizeProgram(flowProgram)

	goFlowProg, goCoreProg, bypass := golang.GenerateGo(status.Pass("dub_to_go"), flowProgram, coreProg, config.RootPackage, config.GenerateTests)
	if config.Dump {
		dumpFlowFuncs(status.Pass("dump_go"), runner, goFlowProg, goCoreProg, config.DumpDir)
	}
	goTreeProg := gotransform.FlowToTree(status.Pass("flow_to_tree"), goFlowProg, goCoreProg, bypass)

	gotree.GoProgramBackend(status.Pass("go_backend"), goTreeProg, goCoreProg, config.OutputDir, runner)
}

func entryPoint(p compiler.LocationProvider, status compiler.PassStatus, config *EGCConfig) {
	status.Begin()
	defer status.End()

	runner := compiler.CreateTaskRunner(config.Jobs)

	processProgram(status, p, runner, config)
	runner.Kill()
}

func mainLoop(config *EGCConfig, profiling bool) {
	p := compiler.MakeProvider()
	status := compiler.MakeStatus(p)

	start := time.Now()
	for i := 0; ; i++ {
		entryPoint(p, status.Pass("egc"), config)
		if profiling && time.Since(start) < time.Second*10 {
			fmt.Println("Re-running to improve profiling data", i)
		} else {
			break
		}
	}

	if status.ShouldHalt() {
		fmt.Printf("%d errors\n", status.ErrorCount())
		os.Exit(1)
	}
}

type EGCConfig struct {
	Dump          bool
	InputDir      string
	OutputDir     string
	RootPackage   []string
	DumpDir       []string
	GenerateTests bool
	Jobs          int
}

func flagError(message string) {
	fmt.Println(message)
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	config := &EGCConfig{
		DumpDir: []string{"output"},
	}

	var rootPackage string
	var verbosity int
	var cpuprofile string
	var memprofile string

	flag.StringVar(&config.InputDir, "indir", "", "Directory containing input files.")
	flag.StringVar(&config.OutputDir, "outdir", "", "Directory to output generated files.")
	flag.StringVar(&rootPackage, "gopackage", "", "Root package for generate sources.")

	flag.BoolVar(&config.Dump, "dump", false, "Dump flowgraphs to output/... (requires graphviz).")
	flag.BoolVar(&config.GenerateTests, "gentests", false, "Generate dub tests.")

	flag.StringVar(&cpuprofile, "cpuprofile", "", "Write cpu profile to file.")
	flag.StringVar(&memprofile, "memprofile", "", "Write memory profile to this file.")
	flag.IntVar(&config.Jobs, "j", runtime.NumCPU(), "Number of threads.")
	flag.IntVar(&verbosity, "v", 0, "Verbosity level.")

	flag.Parse()

	if config.InputDir == "" {
		flagError("-indir is required")
	}
	if config.OutputDir == "" {
		flagError("-outdir is required")
	}
	if rootPackage == "" {
		flagError("-gopackage is required")
	}
	config.RootPackage = strings.Split(rootPackage, "/")

	runtime.GOMAXPROCS(config.Jobs)
	compiler.Verbosity = verbosity

	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			fmt.Println(err.Error())
			return
		} else {
			pprof.StartCPUProfile(f)
			defer pprof.StopCPUProfile()
		}
	}

	mainLoop(config, cpuprofile != "")

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			fmt.Println(err.Error())
			return
		} else {
			pprof.WriteHeapProfile(f)
			f.Close()
		}
	}
}
<|endoftext|>


<commit_before>package selenoid

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"sort"

	"errors"
	"github.com/aerokube/selenoid/config"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
	"github.com/heroku/docker-registry-client/registry"
	"strings"
	"time"
	. "vbom.ml/util/sortorder"
	"strconv"
)

const (
	Latest                = "latest"
	firefox               = "firefox"
	opera                 = "opera"
	tag_1216              = "12.16"
	selenoidImage         = "aerokube/selenoid"
	selenoidContainerName = "selenoid"
	selenoidContainerPort = 4444
)

type SelenoidConfig map[string]config.Versions

type DockerConfigurator struct {
	Logger
	ConfigDirAware
	VersionAware
	DownloadAware
	RequestedBrowsersAware
	LastVersions int
	Pull         bool
	RegistryUrl  string
	Tmpfs        int
	docker       *client.Client
	reg          *registry.Registry
}

func NewDockerConfigurator(config *LifecycleConfig) (*DockerConfigurator, error) {
	c := &DockerConfigurator{
		Logger:                 Logger{Quiet: config.Quiet},
		ConfigDirAware:         ConfigDirAware{ConfigDir: config.ConfigDir},
		VersionAware:           VersionAware{Version: config.Version},
		DownloadAware:          DownloadAware{DownloadNeeded: config.Download},
		RequestedBrowsersAware: RequestedBrowsersAware{Browsers: config.Browsers},
		RegistryUrl:            config.RegistryUrl,
		LastVersions:           config.LastVersions,
		Tmpfs:                  config.Tmpfs,
	}
	if c.Quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}
	err := c.initDockerClient()
	if err != nil {
		return nil, fmt.Errorf("new configurator: %v", err)
	}
	err = c.initRegistryClient()
	if err != nil {
		return nil, fmt.Errorf("new configurator: %v", err)
	}
	return c, nil
}

func (c *DockerConfigurator) initDockerClient() error {
	docker, err := client.NewEnvClient()
	if err != nil {
		return fmt.Errorf("failed to init Docker client: %v", err)
	}
	c.docker = docker
	return nil
}

func (c *DockerConfigurator) initRegistryClient() error {
	reg, err := registry.New(c.RegistryUrl, "", "")
	if err != nil {
		return fmt.Errorf("Docker Registry is not available: %v", err)
	}
	c.reg = reg
	return nil
}

func (c *DockerConfigurator) Close() error {
	if c.docker != nil {
		return c.docker.Close()
	}
	return nil
}

func (c *DockerConfigurator) IsDownloaded() bool {
	return c.getSelenoidImage() != nil
}

func (c *DockerConfigurator) getSelenoidImage() *types.ImageSummary {
	images, err := c.docker.ImageList(context.Background(), types.ImageListOptions{})
	if err != nil {
		c.Printf("Failed to list images: %v\n", err)
		return nil
	}
	for _, img := range images {
		if len(img.RepoTags) > 0 && strings.Contains(img.RepoTags[0], selenoidImage) {
			return &img
		}
	}
	return nil
}

func (c *DockerConfigurator) Download() (string, error) {
	version := c.Version
	if version == Latest {
		latestVersion := c.getLatestSelenoidVersion()
		if latestVersion != nil {
			version = *latestVersion
		}
	}
	ref := selenoidImage
	if version != Latest {
		ref = fmt.Sprintf("%s:%s", ref, version)
	}
	if !c.pullImage(context.Background(), ref) {
		return "", errors.New("failed to pull Selenoid image")
	}
	return ref, nil
}

func (c *DockerConfigurator) getLatestSelenoidVersion() *string {
	tags := c.fetchImageTags(selenoidImage)
	if len(tags) > 0 {
		return &tags[0]
	}
	return nil
}

func (c *DockerConfigurator) IsConfigured() bool {
	return fileExists(getSelenoidConfigPath(c.ConfigDir))
}

func (c *DockerConfigurator) Configure() (*SelenoidConfig, error) {
	cfg := c.createConfig()
	data, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return nil, fmt.Errorf("failed to marshal json: %v\n", err)
	}
	err = c.createConfigDir()
	if err != nil {
		return nil, fmt.Errorf("failed to create output directory: %v\n", err)
	}
	return &cfg, ioutil.WriteFile(getSelenoidConfigPath(c.ConfigDir), data, 0644)
}

func (c *DockerConfigurator) createConfig() SelenoidConfig {
	supportedBrowsers := c.getSupportedBrowsers()
	browsers := make(map[string]config.Versions)
	browsersToIterate := supportedBrowsers
	if c.Browsers != "" {
		requestedBrowsers := strings.Split(c.Browsers, comma)
		if len(requestedBrowsers) > 0 {
			browsersToIterate = make(map[string]string)
			for _, rb := range requestedBrowsers {
				if image, ok := supportedBrowsers[rb]; ok {
					browsersToIterate[rb] = image
					continue
				}
				c.Printf("Unsupported browser: %s\n", rb)
			}
		}
	}
	for browserName, image := range browsersToIterate {
		c.Printf("Processing browser \"%s\"...\n", browserName)
		tags := c.fetchImageTags(image)
		pulledTags := tags
		if c.DownloadNeeded {
			pulledTags = c.pullImages(image, tags)
		} else if c.LastVersions > 0 && c.LastVersions <= len(tags) {
			pulledTags = tags[:c.LastVersions]
		}

		if len(pulledTags) > 0 {
			browsers[browserName] = c.createVersions(browserName, image, pulledTags)
		}
	}
	return browsers
}

func (c *DockerConfigurator) getSupportedBrowsers() map[string]string {
	return map[string]string{
		"firefox": "selenoid/firefox",
		"chrome":  "selenoid/chrome",
		"opera":   "selenoid/opera",
	}
}

func (c *DockerConfigurator) fetchImageTags(image string) []string {
	c.Printf("Fetching tags for image \"%s\"...\n", image)
	tags, err := c.reg.Tags(image)
	if err != nil {
		c.Printf("Failed to fetch tags for image \"%s\"\n", image)
		return nil
	}
	tagsWithoutLatest := filterOutLatest(tags)
	strSlice := Natural(tagsWithoutLatest)
	sort.Sort(sort.Reverse(strSlice))
	return tagsWithoutLatest
}

func filterOutLatest(tags []string) []string {
	ret := []string{}
	for _, tag := range tags {
		if tag != Latest {
			ret = append(ret, tag)
		}
	}
	return ret
}

func (c *DockerConfigurator) createVersions(browserName string, image string, tags []string) config.Versions {
	versions := config.Versions{
		Default:  tags[0],
		Versions: make(map[string]*config.Browser),
	}
	for _, tag := range tags {
		browser := &config.Browser{
			Image: imageWithTag(image, tag),
			Port:  "4444",
			Path:  "/",
		}
		if browserName == firefox || (browserName == opera && tag == tag_1216) {
			browser.Path = "/wd/hub"
		}
		if c.Tmpfs > 0 {
			tmpfs := make(map[string]string)
			tmpfs["/tmp"] = fmt.Sprintf("size=%dm", c.Tmpfs)
			browser.Tmpfs = tmpfs
		}
		versions.Versions[tag] = browser
	}
	return versions
}

func imageWithTag(image string, tag string) string {
	return fmt.Sprintf("%s:%s", image, tag)
}

func (c *DockerConfigurator) pullImages(image string, tags []string) []string {
	pulledTags := []string{}
	ctx := context.Background()
loop:
	for _, tag := range tags {
		ref := imageWithTag(image, tag)
		c.Printf("Pulling image \"%s\"...\n", ref)
		if !c.pullImage(ctx, ref) {
			continue
		}
		pulledTags = append(pulledTags, tag)
		if c.LastVersions > 0 && len(pulledTags) == c.LastVersions {
			break loop
		}
	}
	return pulledTags
}

func (c *DockerConfigurator) pullImage(ctx context.Context, ref string) bool {
	resp, err := c.docker.ImagePull(ctx, ref, types.ImagePullOptions{})
	if err != nil {
		c.Printf("Failed to pull image \"%s\": %v", ref, err)
		return false
	}
	defer resp.Close()
	var row struct {
		Id     string `json:"id"`
		Status string `json:"status"`
	}
	scanner := bufio.NewScanner(resp)
	for prev := ""; scanner.Scan(); {
		err := json.Unmarshal(scanner.Bytes(), &row)
		if err != nil {
			return false
		}
		select {
		case <-ctx.Done():
			{
				c.Printf("Pulling \"%s\" interrupted: %v", ref, ctx.Err())
				return false
			}
		default:
			{
				if prev != row.Status {
					prev = row.Status
					c.Printf("%s: %s\n", row.Status, row.Id)
				}
			}
		}
	}
	if err := scanner.Err(); err != nil {
		c.Printf("Failed to pull image \"%s\": %v", ref, err)
	}
	return true
}

func (c *DockerConfigurator) IsRunning() bool {
	return c.getSelenoidContainer() != nil
}

func (c *DockerConfigurator) getSelenoidContainer() *types.Container {
	f := filters.NewArgs()
	f.Add("name", selenoidContainerName)
	containers, err := c.docker.ContainerList(context.Background(), types.ContainerListOptions{Filters: f})
	if err != nil {
		return nil
	}
	for _, c := range containers {
		for _, p := range c.Ports {
			if p.PublicPort == selenoidContainerPort {
				return &c
			}
		}
	}
	return nil
}

func (c *DockerConfigurator) Start() error {
	image := c.getSelenoidImage()
	if image == nil {
		return errors.New("Selenoid image is not downloaded: this is probably a bug")
	}
	env := []string{
		fmt.Sprintf("TZ=%s", time.Local),
	}
	portString := strconv.Itoa(selenoidContainerPort)
	port, err := nat.NewPort("tcp", portString)
	if err != nil {
		return fmt.Errorf("failed to init Selenoid port: %v", err)
	}
	exposedPorts := map[nat.Port]struct{}{port: {}}
	portBindings := nat.PortMap{}
	portBindings[port] = []nat.PortBinding{{HostIP: "0.0.0.0", HostPort: portString}}
	volumes := []string{fmt.Sprintf("%s:/etc/selenoid:ro", c.ConfigDir)}
	ctx := context.Background()
	ctr, err := c.docker.ContainerCreate(ctx,
		&container.Config{
			Hostname:     "localhost",
			Image:        image.RepoTags[0],
			Env:          env,
			ExposedPorts: exposedPorts,
		},
		&container.HostConfig{
			Binds: volumes,
			PortBindings: portBindings,
		},
		&network.NetworkingConfig{}, selenoidContainerName)
	if err != nil {
		return fmt.Errorf("failed to create container: %v", err)
	}
	err = c.docker.ContainerStart(ctx, ctr.ID, types.ContainerStartOptions{})
	if err != nil {
		c.removeContainer(ctr.ID)
		return fmt.Errorf("failed to start container: %v", err)
	}
	return nil
}

func (c *DockerConfigurator) removeContainer(id string) error {
	return c.docker.ContainerRemove(context.Background(), id, types.ContainerRemoveOptions{RemoveVolumes: true, Force: true})
}

func (c *DockerConfigurator) Stop() error {
	sc := c.getSelenoidContainer()
	if sc != nil {
		err := c.removeContainer(sc.ID)
		if err != nil {
			return fmt.Errorf("failed to stop container: %v", err)
		}
	}
	return nil
}
<commit_msg>Adding docker socket when necessary (fixes #13)<commit_after>package selenoid

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"sort"

	"errors"
	"github.com/aerokube/selenoid/config"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/nat"
	"github.com/heroku/docker-registry-client/registry"
	"strings"
	"time"
	. "vbom.ml/util/sortorder"
	"strconv"
)

const (
	Latest                = "latest"
	firefox               = "firefox"
	opera                 = "opera"
	tag_1216              = "12.16"
	selenoidImage         = "aerokube/selenoid"
	selenoidContainerName = "selenoid"
	selenoidContainerPort = 4444
)

type SelenoidConfig map[string]config.Versions

type DockerConfigurator struct {
	Logger
	ConfigDirAware
	VersionAware
	DownloadAware
	RequestedBrowsersAware
	LastVersions int
	Pull         bool
	RegistryUrl  string
	Tmpfs        int
	docker       *client.Client
	reg          *registry.Registry
}

func NewDockerConfigurator(config *LifecycleConfig) (*DockerConfigurator, error) {
	c := &DockerConfigurator{
		Logger:                 Logger{Quiet: config.Quiet},
		ConfigDirAware:         ConfigDirAware{ConfigDir: config.ConfigDir},
		VersionAware:           VersionAware{Version: config.Version},
		DownloadAware:          DownloadAware{DownloadNeeded: config.Download},
		RequestedBrowsersAware: RequestedBrowsersAware{Browsers: config.Browsers},
		RegistryUrl:            config.RegistryUrl,
		LastVersions:           config.LastVersions,
		Tmpfs:                  config.Tmpfs,
	}
	if c.Quiet {
		log.SetFlags(0)
		log.SetOutput(ioutil.Discard)
	}
	err := c.initDockerClient()
	if err != nil {
		return nil, fmt.Errorf("new configurator: %v", err)
	}
	err = c.initRegistryClient()
	if err != nil {
		return nil, fmt.Errorf("new configurator: %v", err)
	}
	return c, nil
}

func (c *DockerConfigurator) initDockerClient() error {
	docker, err := client.NewEnvClient()
	if err != nil {
		return fmt.Errorf("failed to init Docker client: %v", err)
	}
	c.docker = docker
	return nil
}

func (c *DockerConfigurator) initRegistryClient() error {
	reg, err := registry.New(c.RegistryUrl, "", "")
	if err != nil {
		return fmt.Errorf("Docker Registry is not available: %v", err)
	}
	c.reg = reg
	return nil
}

func (c *DockerConfigurator) Close() error {
	if c.docker != nil {
		return c.docker.Close()
	}
	return nil
}

func (c *DockerConfigurator) IsDownloaded() bool {
	return c.getSelenoidImage() != nil
}

func (c *DockerConfigurator) getSelenoidImage() *types.ImageSummary {
	images, err := c.docker.ImageList(context.Background(), types.ImageListOptions{})
	if err != nil {
		c.Printf("Failed to list images: %v\n", err)
		return nil
	}
	for _, img := range images {
		if len(img.RepoTags) > 0 && strings.Contains(img.RepoTags[0], selenoidImage) {
			return &img
		}
	}
	return nil
}

func (c *DockerConfigurator) Download() (string, error) {
	version := c.Version
	if version == Latest {
		latestVersion := c.getLatestSelenoidVersion()
		if latestVersion != nil {
			version = *latestVersion
		}
	}
	ref := selenoidImage
	if version != Latest {
		ref = fmt.Sprintf("%s:%s", ref, version)
	}
	if !c.pullImage(context.Background(), ref) {
		return "", errors.New("failed to pull Selenoid image")
	}
	return ref, nil
}

func (c *DockerConfigurator) getLatestSelenoidVersion() *string {
	tags := c.fetchImageTags(selenoidImage)
	if len(tags) > 0 {
		return &tags[0]
	}
	return nil
}

func (c *DockerConfigurator) IsConfigured() bool {
	return fileExists(getSelenoidConfigPath(c.ConfigDir))
}

func (c *DockerConfigurator) Configure() (*SelenoidConfig, error) {
	cfg := c.createConfig()
	data, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return nil, fmt.Errorf("failed to marshal json: %v\n", err)
	}
	err = c.createConfigDir()
	if err != nil {
		return nil, fmt.Errorf("failed to create output directory: %v\n", err)
	}
	return &cfg, ioutil.WriteFile(getSelenoidConfigPath(c.ConfigDir), data, 0644)
}

func (c *DockerConfigurator) createConfig() SelenoidConfig {
	supportedBrowsers := c.getSupportedBrowsers()
	browsers := make(map[string]config.Versions)
	browsersToIterate := supportedBrowsers
	if c.Browsers != "" {
		requestedBrowsers := strings.Split(c.Browsers, comma)
		if len(requestedBrowsers) > 0 {
			browsersToIterate = make(map[string]string)
			for _, rb := range requestedBrowsers {
				if image, ok := supportedBrowsers[rb]; ok {
					browsersToIterate[rb] = image
					continue
				}
				c.Printf("Unsupported browser: %s\n", rb)
			}
		}
	}
	for browserName, image := range browsersToIterate {
		c.Printf("Processing browser \"%s\"...\n", browserName)
		tags := c.fetchImageTags(image)
		pulledTags := tags
		if c.DownloadNeeded {
			pulledTags = c.pullImages(image, tags)
		} else if c.LastVersions > 0 && c.LastVersions <= len(tags) {
			pulledTags = tags[:c.LastVersions]
		}

		if len(pulledTags) > 0 {
			browsers[browserName] = c.createVersions(browserName, image, pulledTags)
		}
	}
	return browsers
}

func (c *DockerConfigurator) getSupportedBrowsers() map[string]string {
	return map[string]string{
		"firefox": "selenoid/firefox",
		"chrome":  "selenoid/chrome",
		"opera":   "selenoid/opera",
	}
}

func (c *DockerConfigurator) fetchImageTags(image string) []string {
	c.Printf("Fetching tags for image \"%s\"...\n", image)
	tags, err := c.reg.Tags(image)
	if err != nil {
		c.Printf("Failed to fetch tags for image \"%s\"\n", image)
		return nil
	}
	tagsWithoutLatest := filterOutLatest(tags)
	strSlice := Natural(tagsWithoutLatest)
	sort.Sort(sort.Reverse(strSlice))
	return tagsWithoutLatest
}

func filterOutLatest(tags []string) []string {
	ret := []string{}
	for _, tag := range tags {
		if tag != Latest {
			ret = append(ret, tag)
		}
	}
	return ret
}

func (c *DockerConfigurator) createVersions(browserName string, image string, tags []string) config.Versions {
	versions := config.Versions{
		Default:  tags[0],
		Versions: make(map[string]*config.Browser),
	}
	for _, tag := range tags {
		browser := &config.Browser{
			Image: imageWithTag(image, tag),
			Port:  "4444",
			Path:  "/",
		}
		if browserName == firefox || (browserName == opera && tag == tag_1216) {
			browser.Path = "/wd/hub"
		}
		if c.Tmpfs > 0 {
			tmpfs := make(map[string]string)
			tmpfs["/tmp"] = fmt.Sprintf("size=%dm", c.Tmpfs)
			browser.Tmpfs = tmpfs
		}
		versions.Versions[tag] = browser
	}
	return versions
}

func imageWithTag(image string, tag string) string {
	return fmt.Sprintf("%s:%s", image, tag)
}

func (c *DockerConfigurator) pullImages(image string, tags []string) []string {
	pulledTags := []string{}
	ctx := context.Background()
loop:
	for _, tag := range tags {
		ref := imageWithTag(image, tag)
		c.Printf("Pulling image \"%s\"...\n", ref)
		if !c.pullImage(ctx, ref) {
			continue
		}
		pulledTags = append(pulledTags, tag)
		if c.LastVersions > 0 && len(pulledTags) == c.LastVersions {
			break loop
		}
	}
	return pulledTags
}

func (c *DockerConfigurator) pullImage(ctx context.Context, ref string) bool {
	resp, err := c.docker.ImagePull(ctx, ref, types.ImagePullOptions{})
	if err != nil {
		c.Printf("Failed to pull image \"%s\": %v", ref, err)
		return false
	}
	defer resp.Close()
	var row struct {
		Id     string `json:"id"`
		Status string `json:"status"`
	}
	scanner := bufio.NewScanner(resp)
	for prev := ""; scanner.Scan(); {
		err := json.Unmarshal(scanner.Bytes(), &row)
		if err != nil {
			return false
		}
		select {
		case <-ctx.Done():
			{
				c.Printf("Pulling \"%s\" interrupted: %v", ref, ctx.Err())
				return false
			}
		default:
			{
				if prev != row.Status {
					prev = row.Status
					c.Printf("%s: %s\n", row.Status, row.Id)
				}
			}
		}
	}
	if err := scanner.Err(); err != nil {
		c.Printf("Failed to pull image \"%s\": %v", ref, err)
	}
	return true
}

func (c *DockerConfigurator) IsRunning() bool {
	return c.getSelenoidContainer() != nil
}

func (c *DockerConfigurator) getSelenoidContainer() *types.Container {
	f := filters.NewArgs()
	f.Add("name", selenoidContainerName)
	containers, err := c.docker.ContainerList(context.Background(), types.ContainerListOptions{Filters: f})
	if err != nil {
		return nil
	}
	for _, c := range containers {
		for _, p := range c.Ports {
			if p.PublicPort == selenoidContainerPort {
				return &c
			}
		}
	}
	return nil
}

func (c *DockerConfigurator) Start() error {
	image := c.getSelenoidImage()
	if image == nil {
		return errors.New("Selenoid image is not downloaded: this is probably a bug")
	}
	env := []string{
		fmt.Sprintf("TZ=%s", time.Local),
	}
	portString := strconv.Itoa(selenoidContainerPort)
	port, err := nat.NewPort("tcp", portString)
	if err != nil {
		return fmt.Errorf("failed to init Selenoid port: %v", err)
	}
	exposedPorts := map[nat.Port]struct{}{port: {}}
	portBindings := nat.PortMap{}
	portBindings[port] = []nat.PortBinding{{HostIP: "0.0.0.0", HostPort: portString}}
	volumes := []string{fmt.Sprintf("%s:/etc/selenoid:ro", c.ConfigDir)}
	const dockerSocket = "/var/run/docker.sock"
	if fileExists(dockerSocket) {
		volumes = append(volumes, fmt.Sprintf("%s:%s", dockerSocket, dockerSocket))
	}
	ctx := context.Background()
	ctr, err := c.docker.ContainerCreate(ctx,
		&container.Config{
			Hostname:     "localhost",
			Image:        image.RepoTags[0],
			Env:          env,
			ExposedPorts: exposedPorts,
		},
		&container.HostConfig{
			Binds: volumes,
			PortBindings: portBindings,
		},
		&network.NetworkingConfig{}, selenoidContainerName)
	if err != nil {
		return fmt.Errorf("failed to create container: %v", err)
	}
	err = c.docker.ContainerStart(ctx, ctr.ID, types.ContainerStartOptions{})
	if err != nil {
		c.removeContainer(ctr.ID)
		return fmt.Errorf("failed to start container: %v", err)
	}
	return nil
}

func (c *DockerConfigurator) removeContainer(id string) error {
	return c.docker.ContainerRemove(context.Background(), id, types.ContainerRemoveOptions{RemoveVolumes: true, Force: true})
}

func (c *DockerConfigurator) Stop() error {
	sc := c.getSelenoidContainer()
	if sc != nil {
		err := c.removeContainer(sc.ID)
		if err != nil {
			return fmt.Errorf("failed to stop container: %v", err)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	proto         = "tcp"
	retryInterval = 1
)

var (
	listeningAddress    = flag.String("listeningAddress", ":8080", "Address on which to expose Prometheus metrics.")
	listeningPath       = flag.String("listeningPath", "/metrics", "Path on which to expose Prometheus metrics.")
	muninAddress        = flag.String("muninAddress", "localhost:4949", "munin-node address.")
	muninScrapeInterval = flag.Int("muninScrapeInterval", 60, "Interval in seconds between scrapes.")
	globalConn          net.Conn
	hostname            string
	graphs              []string
	gaugePerMetric      map[string]*prometheus.GaugeVec
	counterPerMetric    map[string]*prometheus.CounterVec
	muninBanner         *regexp.Regexp
)

func init() {
	flag.Parse()
	var err error
	gaugePerMetric = map[string]*prometheus.GaugeVec{}
	counterPerMetric = map[string]*prometheus.CounterVec{}
	muninBanner = regexp.MustCompile(`# munin node at (.*)`)

	err = connect()
	if err != nil {
		log.Fatalf("Could not connect to %s: %s", *muninAddress, err)
	}
}

func serveStatus() {
	http.Handle(*listeningPath, prometheus.Handler())
	http.ListenAndServe(*listeningAddress, nil)
}

func connect() (err error) {
	log.Printf("Connecting...")
	globalConn, err = net.Dial(proto, *muninAddress)
	if err != nil {
		return
	}
	log.Printf("connected!")

	reader := bufio.NewReader(globalConn)
	head, err := reader.ReadString('\n')
	if err != nil {
		return
	}

	matches := muninBanner.FindStringSubmatch(head)
	if len(matches) != 2 { // expect: # munin node at <hostname>
		return fmt.Errorf("Unexpected line: %s", head)
	}
	hostname = matches[1]
	log.Printf("Found hostname: %s", hostname)
	return
}

func muninCommand(cmd string) (reader *bufio.Reader, err error) {
	reader = bufio.NewReader(globalConn)

	fmt.Fprintf(globalConn, cmd+"\n")

	_, err = reader.Peek(1)
	switch err {
	case io.EOF:
		log.Printf("not connected anymore, closing connection")
		globalConn.Close()
		for {
			err = connect()
			if err == nil {
				break
			}
			log.Printf("Couldn't reconnect: %s", err)
			time.Sleep(retryInterval * time.Second)
		}

		return muninCommand(cmd)
	case nil: //no error
		break
	default:
		log.Fatalf("Unexpected error: %s", err)
	}

	return
}

func muninList() (items []string, err error) {
	munin, err := muninCommand("list")
	if err != nil {
		log.Printf("couldn't get list")
		return
	}

	response, err := munin.ReadString('\n') // we are only interested in the first line
	if err != nil {
		log.Printf("couldn't read response")
		return
	}

	if response[0] == '#' { // # not expected here
		err = fmt.Errorf("Error getting items: %s", response)
		return
	}
	items = strings.Fields(strings.TrimRight(response, "\n"))
	return
}

func muninConfig(name string) (config map[string]map[string]string, graphConfig map[string]string, err error) {
	graphConfig = make(map[string]string)
	config = make(map[string]map[string]string)

	resp, err := muninCommand("config " + name)
	if err != nil {
		log.Printf("couldn't get config for %s", name)
		return
	}

	for {
		line, err := resp.ReadString('\n')
		if err == io.EOF {
			log.Fatalf("unexpected EOF, retrying")
			return muninConfig(name)
		}
		if err != nil {
			return nil, nil, err
		}
		if line == ".\n" { // munin end marker
			break
		}
		if line[0] == '#' { // here it's just a comment, so ignore it
			continue
		}
		parts := strings.Fields(line)
		if len(parts) < 2 {
			return nil, nil, fmt.Errorf("Line unexpected: %s", line)
		}
		key, value := parts[0], strings.TrimRight(strings.Join(parts[1:], " "), "\n")

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 { // it's a metric config (metric.label etc)
			if _, ok := config[keyParts[0]]; !ok { //FIXME: is there no better way?
				config[keyParts[0]] = make(map[string]string)
			}
			config[keyParts[0]][keyParts[1]] = value
		} else {
			graphConfig[keyParts[0]] = value
		}
	}
	return
}

func registerMetrics() (err error) {
	items, err := muninList()
	if err != nil {
		return
	}

	for _, name := range items {
		graphs = append(graphs, name)
		configs, graphConfig, err := muninConfig(name)
		if err != nil {
			return err
		}

		for metric, config := range configs {
			metricName := strings.Replace(name+"_"+metric, "-", "_", -1)
			desc := graphConfig["graph_title"] + ": " + config["label"]
			if config["info"] != "" {
				desc = desc + ", " + config["info"]
			}
			muninType := strings.ToLower(config["type"])
			// muninType can be empty and defaults to gauge
			if muninType == "counter" || muninType == "derive" {
				gv := prometheus.NewCounterVec(
					prometheus.CounterOpts{
						Name:        metricName,
						Help:        desc,
						ConstLabels: prometheus.Labels{"type": muninType},
					},
					[]string{"hostname", "graphname", "muninlabel"},
				)
				log.Printf("Registered counter %s: %s", metricName, desc)
				counterPerMetric[metricName] = gv
				prometheus.Register(gv)

			} else {
				gv := prometheus.NewGaugeVec(
					prometheus.GaugeOpts{
						Name:        metricName,
						Help:        desc,
						ConstLabels: prometheus.Labels{"type": "gauge"},
					},
					[]string{"hostname", "graphname", "muninlabel"},
				)
				log.Printf("Registered gauge %s: %s", metricName, desc)
				gaugePerMetric[metricName] = gv
				prometheus.Register(gv)
			}
		}
	}
	return nil
}

func fetchMetrics() (err error) {
	for _, graph := range graphs {
		munin, err := muninCommand("fetch " + graph)
		if err != nil {
			return err
		}

		for {
			line, err := munin.ReadString('\n')
			line = strings.TrimRight(line, "\n")
			if err == io.EOF {
				log.Fatalf("unexpected EOF, retrying")
				return fetchMetrics()
			}
			if err != nil {
				return err
			}
			if len(line) == 1 && line[0] == '.' {
				log.Printf("End of list")
				break
			}

			parts := strings.Fields(line)
			if len(parts) != 2 {
				log.Printf("unexpected line: %s", line)
				continue
			}
			key, valueString := strings.Split(parts[0], ".")[0], parts[1]
			value, err := strconv.ParseFloat(valueString, 64)
			if err != nil {
				log.Printf("Couldn't parse value in line %s, malformed?", line)
				continue
			}
			name := strings.Replace(graph+"_"+key, "-", "_", -1)
			log.Printf("%s: %f\n", name, value)
			_, isGauge := gaugePerMetric[name]
			if isGauge {
				gaugePerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
			} else {
				counterPerMetric[name].WithLabelValues(hostname, graph, key).Add(value)
			}
		}
	}
	return
}

func main() {
	flag.Parse()
	err := registerMetrics()
	if err != nil {
		log.Fatalf("Could not register metrics: %s", err)
	}

	go serveStatus()

	func() {
		for {
			log.Printf("Scraping")
			err := fetchMetrics()
			if err != nil {
				log.Printf("Error occured when trying to fetch metrics: %s", err)
			}
			time.Sleep(time.Duration(*muninScrapeInterval) * time.Second)
		}
	}()
}
<commit_msg>Adding check if Counter is actually registered before trying to add value to it<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const (
	proto         = "tcp"
	retryInterval = 1
)

var (
	listeningAddress    = flag.String("listeningAddress", ":8080", "Address on which to expose Prometheus metrics.")
	listeningPath       = flag.String("listeningPath", "/metrics", "Path on which to expose Prometheus metrics.")
	muninAddress        = flag.String("muninAddress", "localhost:4949", "munin-node address.")
	muninScrapeInterval = flag.Int("muninScrapeInterval", 60, "Interval in seconds between scrapes.")
	globalConn          net.Conn
	hostname            string
	graphs              []string
	gaugePerMetric      map[string]*prometheus.GaugeVec
	counterPerMetric    map[string]*prometheus.CounterVec
	muninBanner         *regexp.Regexp
)

func init() {
	flag.Parse()
	var err error
	gaugePerMetric = map[string]*prometheus.GaugeVec{}
	counterPerMetric = map[string]*prometheus.CounterVec{}
	muninBanner = regexp.MustCompile(`# munin node at (.*)`)

	err = connect()
	if err != nil {
		log.Fatalf("Could not connect to %s: %s", *muninAddress, err)
	}
}

func serveStatus() {
	http.Handle(*listeningPath, prometheus.Handler())
	http.ListenAndServe(*listeningAddress, nil)
}

func connect() (err error) {
	log.Printf("Connecting...")
	globalConn, err = net.Dial(proto, *muninAddress)
	if err != nil {
		return
	}
	log.Printf("connected!")

	reader := bufio.NewReader(globalConn)
	head, err := reader.ReadString('\n')
	if err != nil {
		return
	}

	matches := muninBanner.FindStringSubmatch(head)
	if len(matches) != 2 { // expect: # munin node at <hostname>
		return fmt.Errorf("Unexpected line: %s", head)
	}
	hostname = matches[1]
	log.Printf("Found hostname: %s", hostname)
	return
}

func muninCommand(cmd string) (reader *bufio.Reader, err error) {
	reader = bufio.NewReader(globalConn)

	fmt.Fprintf(globalConn, cmd+"\n")

	_, err = reader.Peek(1)
	switch err {
	case io.EOF:
		log.Printf("not connected anymore, closing connection")
		globalConn.Close()
		for {
			err = connect()
			if err == nil {
				break
			}
			log.Printf("Couldn't reconnect: %s", err)
			time.Sleep(retryInterval * time.Second)
		}

		return muninCommand(cmd)
	case nil: //no error
		break
	default:
		log.Fatalf("Unexpected error: %s", err)
	}

	return
}

func muninList() (items []string, err error) {
	munin, err := muninCommand("list")
	if err != nil {
		log.Printf("couldn't get list")
		return
	}

	response, err := munin.ReadString('\n') // we are only interested in the first line
	if err != nil {
		log.Printf("couldn't read response")
		return
	}

	if response[0] == '#' { // # not expected here
		err = fmt.Errorf("Error getting items: %s", response)
		return
	}
	items = strings.Fields(strings.TrimRight(response, "\n"))
	return
}

func muninConfig(name string) (config map[string]map[string]string, graphConfig map[string]string, err error) {
	graphConfig = make(map[string]string)
	config = make(map[string]map[string]string)

	resp, err := muninCommand("config " + name)
	if err != nil {
		log.Printf("couldn't get config for %s", name)
		return
	}

	for {
		line, err := resp.ReadString('\n')
		if err == io.EOF {
			log.Fatalf("unexpected EOF, retrying")
			return muninConfig(name)
		}
		if err != nil {
			return nil, nil, err
		}
		if line == ".\n" { // munin end marker
			break
		}
		if line[0] == '#' { // here it's just a comment, so ignore it
			continue
		}
		parts := strings.Fields(line)
		if len(parts) < 2 {
			return nil, nil, fmt.Errorf("Line unexpected: %s", line)
		}
		key, value := parts[0], strings.TrimRight(strings.Join(parts[1:], " "), "\n")

		keyParts := strings.Split(key, ".")
		if len(keyParts) > 1 { // it's a metric config (metric.label etc)
			if _, ok := config[keyParts[0]]; !ok { //FIXME: is there no better way?
				config[keyParts[0]] = make(map[string]string)
			}
			config[keyParts[0]][keyParts[1]] = value
		} else {
			graphConfig[keyParts[0]] = value
		}
	}
	return
}

func registerMetrics() (err error) {
	items, err := muninList()
	if err != nil {
		return
	}

	for _, name := range items {
		graphs = append(graphs, name)
		configs, graphConfig, err := muninConfig(name)
		if err != nil {
			return err
		}

		for metric, config := range configs {
			metricName := strings.Replace(name+"_"+metric, "-", "_", -1)
			desc := graphConfig["graph_title"] + ": " + config["label"]
			if config["info"] != "" {
				desc = desc + ", " + config["info"]
			}
			muninType := strings.ToLower(config["type"])
			// muninType can be empty and defaults to gauge
			if muninType == "counter" || muninType == "derive" {
				gv := prometheus.NewCounterVec(
					prometheus.CounterOpts{
						Name:        metricName,
						Help:        desc,
						ConstLabels: prometheus.Labels{"type": muninType},
					},
					[]string{"hostname", "graphname", "muninlabel"},
				)
				log.Printf("Registered counter %s: %s", metricName, desc)
				counterPerMetric[metricName] = gv
				prometheus.Register(gv)

			} else {
				gv := prometheus.NewGaugeVec(
					prometheus.GaugeOpts{
						Name:        metricName,
						Help:        desc,
						ConstLabels: prometheus.Labels{"type": "gauge"},
					},
					[]string{"hostname", "graphname", "muninlabel"},
				)
				log.Printf("Registered gauge %s: %s", metricName, desc)
				gaugePerMetric[metricName] = gv
				prometheus.Register(gv)
			}
		}
	}
	return nil
}

func fetchMetrics() (err error) {
	for _, graph := range graphs {
		munin, err := muninCommand("fetch " + graph)
		if err != nil {
			return err
		}

		for {
			line, err := munin.ReadString('\n')
			line = strings.TrimRight(line, "\n")
			if err == io.EOF {
				log.Fatalf("unexpected EOF, retrying")
				return fetchMetrics()
			}
			if err != nil {
				return err
			}
			if len(line) == 1 && line[0] == '.' {
				log.Printf("End of list")
				break
			}

			parts := strings.Fields(line)
			if len(parts) != 2 {
				log.Printf("unexpected line: %s", line)
				continue
			}
			key, valueString := strings.Split(parts[0], ".")[0], parts[1]
			value, err := strconv.ParseFloat(valueString, 64)
			if err != nil {
				log.Printf("Couldn't parse value in line %s, malformed?", line)
				continue
			}
			name := strings.Replace(graph+"_"+key, "-", "_", -1)
			log.Printf("%s: %f\n", name, value)
			_, isGauge := gaugePerMetric[name]
			if isGauge {
				gaugePerMetric[name].WithLabelValues(hostname, graph, key).Set(value)
				continue
			}
			_, isCounter := counterPerMetric[name]
			if isCounter {
				counterPerMetric[name].WithLabelValues(hostname, graph, key).Add(value)
				continue
			}
		}
	}
	return
}

func main() {
	flag.Parse()
	err := registerMetrics()
	if err != nil {
		log.Fatalf("Could not register metrics: %s", err)
	}

	go serveStatus()

	func() {
		for {
			log.Printf("Scraping")
			err := fetchMetrics()
			if err != nil {
				log.Printf("Error occured when trying to fetch metrics: %s", err)
			}
			time.Sleep(time.Duration(*muninScrapeInterval) * time.Second)
		}
	}()
}
<|endoftext|>


<commit_before>// repmbox handles repbin mailboxes.
//
// Release 20150606.
// You can reach me with
// repclient --recipientPubKey 7VW3oPLzQc7VS2anLyDtrdARDdSwa7QTF7h3N2t6J2VN_DTshmJFEDa7XM2w9nHBq4CgtvK4kYdBp8G3wFPkYcGd1
// Don't forget to put your own key into your message! (create with repmbox -add)
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path"
	"strconv"
	"strings"
	"time"
)

type sender struct {
	Sender     string
	PublicKey  string
	PrivateKey string
}

type config struct {
	PrivateKey string
	Server     string
	Start      int
	Count      int
	Sender     []sender
}

type senderKey struct {
	user    string
	privKey string
}

func genConfig() (*config, error) {
	// generate private key with repclient
	cmd := exec.Command("repclient", "--genkey", "--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	var cfg config
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			cfg.PrivateKey = parts[1]
			break
		}
	}
	if cfg.PrivateKey == "" {
		return nil, fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	cfg.Start = -1
	cfg.Count = 100
	return &cfg, nil
}

func loadConfig(configFile, configDir string) (*config, error) {
	if _, err := os.Stat(configFile); err != nil {
		return nil, fmt.Errorf("config file '%s' doesn't exist, generate with:\n"+
			"mkdir -p %s && repmbox --genconfig > %s", configFile, configDir, configFile)
	}
	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, err
	}
	var cfg config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}
	// check server entry
	if cfg.Server == "" {
		return nil, fmt.Errorf("specify Server entry in config file '%s'", configFile)
	}
	return &cfg, nil
}

func (cfg *config) save(configFile string) error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(configFile, jsn, 0600); err != nil {
		return err
	}
	return nil
}

func (cfg *config) show() error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	fmt.Println(string(jsn))
	return nil
}

func (cfg *config) addSender(addSender, configFile string) error {
	var snd sender
	snd.Sender = addSender
	// generate key pair for sender
	cmd := exec.Command("repclient", "--gentemp",
		"--privkey", cfg.PrivateKey,
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return err
	}
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			snd.PrivateKey = parts[1]
		}
		if len(parts) >= 2 && parts[0] == "STATUS (PublicKey):" {
			snd.PublicKey = parts[1]
			break
		}
	}
	if snd.PrivateKey == "" {
		return fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	if snd.PublicKey == "" {
		return fmt.Errorf("could not parse public key from output:\n%s", out.String())
	}
	cfg.Sender = append(cfg.Sender, snd)
	if err := cfg.save(configFile); err != nil {
		return err
	}
	fmt.Printf("mailbox for %s added:\n", addSender)
	fmt.Println(snd.PublicKey)
	return nil
}

func (cfg *config) getList() ([]string, bool, error) {
	cmd := exec.Command("repclient",
		"--index",
		"--server", cfg.Server,
		"--privkey", cfg.PrivateKey,
		"--start", strconv.Itoa(cfg.Start+1),
		"--count", strconv.Itoa(cfg.Count),
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		// check if error indicates no new messages
		lines := strings.Split(out.String(), "\n")
		for _, line := range lines {
			if line == "FATAL: List error: listparse: No entries in list" {
				// return empty list
				return nil, false, nil
			}
		}
		// show and return error
		fmt.Print(out.String())
		return nil, false, err
	}
	var list []string
	var more bool
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 {
			if parts[0] == "STATUS (MessageList):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 2 {
					list = append(list, particles[1])
				} else {
					return nil, false, fmt.Errorf("could not parse line: %s", line)
				}
			} else if parts[0] == "STATUS (ListResult):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 3 {
					var err error
					more, err = strconv.ParseBool(particles[2])
					if err != nil {
						return nil, false, err
					}
					break
				}
			}
		}
	}
	return list, more, nil
}

func (cfg *config) getSenderKeys() map[string]senderKey {
	keys := make(map[string]senderKey)
	for _, sdr := range cfg.Sender {
		pubKeys := strings.Split(sdr.PublicKey, "_")
		privKeys := strings.Split(sdr.PrivateKey, "_")
		keys[pubKeys[0]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[0],
		}
		keys[pubKeys[1]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[1],
		}
	}
	return keys
}

func (cfg *config) getMessages(list []string, outdir, stmdir string, verbose bool) error {
	// make sure STM directory exists
	if err := os.MkdirAll(stmdir, 0700); err != nil {
		return err
	}
	for _, msg := range list {
		fmt.Printf("retrieving message %s\n", msg)
		cmd := exec.Command("repclient",
			"--decrypt",
			"--keymgt", "3",
			"--server", cfg.Server,
			"--stmdir", stmdir,
			"--appdata",
			msg)
		var out bytes.Buffer
		cmd.Stdout = &out
		stderr, err := cmd.StderrPipe()
		if err != nil {
			return err
		}
		r, w, err := os.Pipe()
		if err != nil {
			return err
		}
		cmd.ExtraFiles = append(cmd.ExtraFiles, r)
		if err := cmd.Start(); err != nil {
			return err
		}
		var user string
		var stmfile string
		var keyNotAvailable bool
		keys := cfg.getSenderKeys()
		scanner := bufio.NewScanner(stderr)
		for scanner.Scan() {
			line := scanner.Text()
			if verbose {
				fmt.Println(line)
			}
			parts := strings.Split(line, "\t")
			if len(parts) >= 2 {
				if parts[0] == "STATUS (KeyMGTRequest):" {
					key, ok := keys[parts[1]]
					if !ok {
						if _, err := w.WriteString("not available\n"); err != nil {
							return err
						}
						keyNotAvailable = true
						break
					} else {
						if _, err := w.WriteString(key.privKey + "\n"); err != nil {
							return err
						}
						user = key.user
					}
				} else if parts[0] == "STATUS (STMFile):" {
					stmfile = parts[1]
				}
			} else {
				return fmt.Errorf("could not parse repclient output: %s", line)
			}
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		err = cmd.Wait()
		if !keyNotAvailable && err != nil {
			return fmt.Errorf("repclient call failed with: %s", err)
		}
		if stmfile != "" {
			fmt.Printf("new repost message from %s written to:\n%s\n", user, stmfile)
		} else if keyNotAvailable {
			fmt.Println("key not available, skipping message forever!")
		} else {
			// make sure output directory exists
			resdir := path.Join(outdir, user)
			if err := os.MkdirAll(resdir, 0700); err != nil {
				return err
			}
			// write message
			filename := path.Join(resdir, msg)
			if err := ioutil.WriteFile(path.Join(resdir, msg), out.Bytes(), 0600); err != nil {
				return err
			}
			fmt.Printf("new message from %s written to:\n%s\n", user, filename)
		}
	}
	return nil
}

func (cfg *config) downloadNewMessages(outdir, stmdir, configFile string, verbose bool) error {
	more := true
	for more {
		var list []string
		var err error
		// get list
		list, more, err = cfg.getList()
		if err != nil {
			return err
		}
		if len(list) == 0 {
			if verbose {
				fmt.Println("no new messages")
			}
			break
		}
		// get new messages
		if err := cfg.getMessages(list, outdir, stmdir, verbose); err != nil {
			return err
		}
		// increase start
		cfg.Start += len(list)
		if err := cfg.save(configFile); err != nil {
			return err
		}
		if verbose {
			fmt.Printf("more: %s\n", strconv.FormatBool(more))
		}
	}
	return nil
}

func appDataDir(appName string) (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return path.Join(user.HomeDir, ".config", appName), nil
}

func fatal(err error) {
	fmt.Fprintf(os.Stderr, "%s: error: %s\n", os.Args[0], err)
	os.Exit(1)
}

func mainFunc(configDir string) error {
	// define options
	addSender := flag.String("add", "", "Add new mailbox for given sender name")
	configFile := flag.String("configfile", path.Join(configDir, "repmbox.config"), "Path to configuration file")
	genconfig := flag.Bool("genconfig", false, "Generate config file")
	listSender := flag.Bool("list", false, "List mailboxes")
	loop := flag.Bool("loop", false, "Retrieve new messages every 5 minutes")
	outdir := flag.String("outdir", path.Join(configDir, "messages"), "Directory for downloaded messages")
	stmdir := flag.String("stmdir", path.Join(configDir, "stmdir"), "Directory for STM messages")
	verbose := flag.Bool("v", false, "Be verbose")
	// parse options
	flag.Parse()
	if *genconfig {
		cfg, err := genConfig()
		if err != nil {
			return err
		}
		if err := cfg.show(); err != nil {
			return err
		}
		return nil
	}
	// load config file
	cfg, err := loadConfig(*configFile, configDir)
	if err != nil {
		return err
	}
	if *addSender != "" {
		if err := cfg.addSender(*addSender, *configFile); err != nil {
			return err
		}
	} else if *listSender {
		for _, snd := range cfg.Sender {
			fmt.Println(snd.Sender, snd.PublicKey)
		}
	} else {
		for true {
			if err := cfg.downloadNewMessages(*outdir, *stmdir, *configFile, *verbose); err != nil {
				return err
			}
			if *loop {
				time.Sleep(5 * time.Minute)
			} else {
				break
			}
		}
	}
	return nil
}

func main() {
	configDir, err := appDataDir("repmbox")
	if err != nil {
		fatal(err)
	}
	if err := mainFunc(configDir); err != nil {
		fatal(err)
	}
}
<commit_msg>write PID file<commit_after>// repmbox handles repbin mailboxes.
//
// Release 20150606.
// You can reach me with
// repclient --recipientPubKey 7VW3oPLzQc7VS2anLyDtrdARDdSwa7QTF7h3N2t6J2VN_DTshmJFEDa7XM2w9nHBq4CgtvK4kYdBp8G3wFPkYcGd1
// Don't forget to put your own key into your message! (create with repmbox -add)
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"os/user"
	"path"
	"strconv"
	"strings"
	"time"
)

type sender struct {
	Sender     string
	PublicKey  string
	PrivateKey string
}

type config struct {
	PrivateKey string
	Server     string
	Start      int
	Count      int
	Sender     []sender
}

type senderKey struct {
	user    string
	privKey string
}

func genConfig() (*config, error) {
	// generate private key with repclient
	cmd := exec.Command("repclient", "--genkey", "--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return nil, err
	}
	var cfg config
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			cfg.PrivateKey = parts[1]
			break
		}
	}
	if cfg.PrivateKey == "" {
		return nil, fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	cfg.Start = -1
	cfg.Count = 100
	return &cfg, nil
}

func loadConfig(configFile, configDir string) (*config, error) {
	if _, err := os.Stat(configFile); err != nil {
		return nil, fmt.Errorf("config file '%s' doesn't exist, generate with:\n"+
			"mkdir -p %s && repmbox --genconfig > %s", configFile, configDir, configFile)
	}
	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		return nil, err
	}
	var cfg config
	if err := json.Unmarshal(data, &cfg); err != nil {
		return nil, err
	}
	// check server entry
	if cfg.Server == "" {
		return nil, fmt.Errorf("specify Server entry in config file '%s'", configFile)
	}
	return &cfg, nil
}

func (cfg *config) save(configFile string) error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	if err := ioutil.WriteFile(configFile, jsn, 0600); err != nil {
		return err
	}
	return nil
}

func (cfg *config) show() error {
	jsn, err := json.MarshalIndent(cfg, "", "    ")
	if err != nil {
		return err
	}
	fmt.Println(string(jsn))
	return nil
}

func (cfg *config) addSender(addSender, configFile string) error {
	var snd sender
	snd.Sender = addSender
	// generate key pair for sender
	cmd := exec.Command("repclient", "--gentemp",
		"--privkey", cfg.PrivateKey,
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		return err
	}
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 && parts[0] == "STATUS (PrivateKey):" {
			snd.PrivateKey = parts[1]
		}
		if len(parts) >= 2 && parts[0] == "STATUS (PublicKey):" {
			snd.PublicKey = parts[1]
			break
		}
	}
	if snd.PrivateKey == "" {
		return fmt.Errorf("could not parse private key from output:\n%s", out.String())
	}
	if snd.PublicKey == "" {
		return fmt.Errorf("could not parse public key from output:\n%s", out.String())
	}
	cfg.Sender = append(cfg.Sender, snd)
	if err := cfg.save(configFile); err != nil {
		return err
	}
	fmt.Printf("mailbox for %s added:\n", addSender)
	fmt.Println(snd.PublicKey)
	return nil
}

func (cfg *config) getList() ([]string, bool, error) {
	cmd := exec.Command("repclient",
		"--index",
		"--server", cfg.Server,
		"--privkey", cfg.PrivateKey,
		"--start", strconv.Itoa(cfg.Start+1),
		"--count", strconv.Itoa(cfg.Count),
		"--appdata")
	var out bytes.Buffer
	cmd.Stderr = &out
	if err := cmd.Run(); err != nil {
		// check if error indicates no new messages
		lines := strings.Split(out.String(), "\n")
		for _, line := range lines {
			if line == "FATAL: List error: listparse: No entries in list" {
				// return empty list
				return nil, false, nil
			}
		}
		// show and return error
		fmt.Print(out.String())
		return nil, false, err
	}
	var list []string
	var more bool
	lines := strings.Split(out.String(), "\n")
	for _, line := range lines {
		parts := strings.Split(line, "\t")
		if len(parts) >= 2 {
			if parts[0] == "STATUS (MessageList):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 2 {
					list = append(list, particles[1])
				} else {
					return nil, false, fmt.Errorf("could not parse line: %s", line)
				}
			} else if parts[0] == "STATUS (ListResult):" {
				particles := strings.Split(parts[1], " ")
				if len(particles) >= 3 {
					var err error
					more, err = strconv.ParseBool(particles[2])
					if err != nil {
						return nil, false, err
					}
					break
				}
			}
		}
	}
	return list, more, nil
}

func (cfg *config) getSenderKeys() map[string]senderKey {
	keys := make(map[string]senderKey)
	for _, sdr := range cfg.Sender {
		pubKeys := strings.Split(sdr.PublicKey, "_")
		privKeys := strings.Split(sdr.PrivateKey, "_")
		keys[pubKeys[0]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[0],
		}
		keys[pubKeys[1]] = senderKey{
			user:    sdr.Sender,
			privKey: privKeys[1],
		}
	}
	return keys
}

func (cfg *config) getMessages(list []string, outdir, stmdir string, verbose bool) error {
	// make sure STM directory exists
	if err := os.MkdirAll(stmdir, 0700); err != nil {
		return err
	}
	for _, msg := range list {
		fmt.Printf("retrieving message %s\n", msg)
		cmd := exec.Command("repclient",
			"--decrypt",
			"--keymgt", "3",
			"--server", cfg.Server,
			"--stmdir", stmdir,
			"--appdata",
			msg)
		var out bytes.Buffer
		cmd.Stdout = &out
		stderr, err := cmd.StderrPipe()
		if err != nil {
			return err
		}
		r, w, err := os.Pipe()
		if err != nil {
			return err
		}
		cmd.ExtraFiles = append(cmd.ExtraFiles, r)
		if err := cmd.Start(); err != nil {
			return err
		}
		var user string
		var stmfile string
		var keyNotAvailable bool
		keys := cfg.getSenderKeys()
		scanner := bufio.NewScanner(stderr)
		for scanner.Scan() {
			line := scanner.Text()
			if verbose {
				fmt.Println(line)
			}
			parts := strings.Split(line, "\t")
			if len(parts) >= 2 {
				if parts[0] == "STATUS (KeyMGTRequest):" {
					key, ok := keys[parts[1]]
					if !ok {
						if _, err := w.WriteString("not available\n"); err != nil {
							return err
						}
						keyNotAvailable = true
						break
					} else {
						if _, err := w.WriteString(key.privKey + "\n"); err != nil {
							return err
						}
						user = key.user
					}
				} else if parts[0] == "STATUS (STMFile):" {
					stmfile = parts[1]
				}
			} else {
				return fmt.Errorf("could not parse repclient output: %s", line)
			}
		}
		if err := scanner.Err(); err != nil {
			return err
		}
		err = cmd.Wait()
		if !keyNotAvailable && err != nil {
			return fmt.Errorf("repclient call failed with: %s", err)
		}
		if stmfile != "" {
			fmt.Printf("new repost message from %s written to:\n%s\n", user, stmfile)
		} else if keyNotAvailable {
			fmt.Println("key not available, skipping message forever!")
		} else {
			// make sure output directory exists
			resdir := path.Join(outdir, user)
			if err := os.MkdirAll(resdir, 0700); err != nil {
				return err
			}
			// write message
			filename := path.Join(resdir, msg)
			if err := ioutil.WriteFile(path.Join(resdir, msg), out.Bytes(), 0600); err != nil {
				return err
			}
			fmt.Printf("new message from %s written to:\n%s\n", user, filename)
		}
	}
	return nil
}

func (cfg *config) downloadNewMessages(outdir, stmdir, configFile string, verbose bool) error {
	more := true
	for more {
		var list []string
		var err error
		// get list
		list, more, err = cfg.getList()
		if err != nil {
			return err
		}
		if len(list) == 0 {
			if verbose {
				fmt.Println("no new messages")
			}
			break
		}
		// get new messages
		if err := cfg.getMessages(list, outdir, stmdir, verbose); err != nil {
			return err
		}
		// increase start
		cfg.Start += len(list)
		if err := cfg.save(configFile); err != nil {
			return err
		}
		if verbose {
			fmt.Printf("more: %s\n", strconv.FormatBool(more))
		}
	}
	return nil
}

func appDataDir(appName string) (string, error) {
	user, err := user.Current()
	if err != nil {
		return "", err
	}
	return path.Join(user.HomeDir, ".config", appName), nil
}

func writePIDFile(configDir string) (string, error) {
	// make sure config directory exists
	if err := os.MkdirAll(configDir, 0700); err != nil {
		return "", err
	}
	pidFile := path.Join(configDir, "repmbox.pid")
	if _, err := os.Stat(pidFile); err == nil {
		return "", fmt.Errorf("PID file '%s' already exists, another instance running?", pidFile)
	}
	fp, err := os.Create(pidFile)
	if err != nil {
		return "", err
	}
	if _, err := fp.WriteString(strconv.Itoa(os.Getpid()) + "\n"); err != nil {
		return "", err
	}
	if err := fp.Close(); err != nil {
		return "", err
	}
	return pidFile, nil
}

func mainFunc(configDir string) error {
	// define options
	addSender := flag.String("add", "", "Add new mailbox for given sender name")
	configFile := flag.String("configfile", path.Join(configDir, "repmbox.config"), "Path to configuration file")
	genconfig := flag.Bool("genconfig", false, "Generate config file")
	listSender := flag.Bool("list", false, "List mailboxes")
	loop := flag.Bool("loop", false, "Retrieve new messages every 5 minutes")
	outdir := flag.String("outdir", path.Join(configDir, "messages"), "Directory for downloaded messages")
	stmdir := flag.String("stmdir", path.Join(configDir, "stmdir"), "Directory for STM messages")
	verbose := flag.Bool("v", false, "Be verbose")
	// parse options
	flag.Parse()
	if *genconfig {
		cfg, err := genConfig()
		if err != nil {
			return err
		}
		if err := cfg.show(); err != nil {
			return err
		}
		return nil
	}
	// write PID file
	pidFile, err := writePIDFile(configDir)
	if err != nil {
		return err
	}
	defer os.Remove(pidFile)
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c // blocks until signal is received
		os.Remove(pidFile)
		os.Exit(2)
	}()
	// load config file
	cfg, err := loadConfig(*configFile, configDir)
	if err != nil {
		return err
	}
	if *addSender != "" {
		if err := cfg.addSender(*addSender, *configFile); err != nil {
			return err
		}
	} else if *listSender {
		for _, snd := range cfg.Sender {
			fmt.Println(snd.Sender, snd.PublicKey)
		}
	} else {
		for true {
			if err := cfg.downloadNewMessages(*outdir, *stmdir, *configFile, *verbose); err != nil {
				return err
			}
			if *loop {
				time.Sleep(5 * time.Minute)
			} else {
				break
			}
		}
	}
	return nil
}

func fatal(err error) {
	fmt.Fprintf(os.Stderr, "%s: error: %s\n", os.Args[0], err)
	os.Exit(1)
}

func main() {
	configDir, err := appDataDir("repmbox")
	if err != nil {
		fatal(err)
	}
	if err := mainFunc(configDir); err != nil {
		fatal(err)
	}
}
<|endoftext|>