<commit_before>package envetcd

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"net"
	"net/url"
	"os"
	"regexp"
	"strconv"
	"strings"
	"text/template"

	"github.com/coreos/etcd/pkg/transport"
	"github.com/coreos/go-etcd/etcd"
)

// order of precedence:
// global < system < service < host

// Regexp for invalid characters in keys
var invalidRegexp = regexp.MustCompile(`[^a-zA-Z0-9_]`)

// Config contains all of the parameters needed to run GetKeyPairs
type Config struct {
	Peers             []string
	Sync              bool
	Sanitize          bool
	Upcase            bool
	UseDefaultGateway bool
	Prefix            string
	System            string
	Service           string
	Hostname          string
	TLS               *transport.TLSInfo
}

var (
	etcdKeyTemplates = [...]string{
		"{{.Prefix}}/global",
		"{{if .System}}{{.Prefix}}/system/{{.System}}{{end}}",
		"{{if .Service}}{{.Prefix}}/service/{{.Service}}{{end}}",
		"{{if .Hostname}}{{.Prefix}}/host/{{.Hostname}}{{end}}",
	}

	gatewayIP *net.IP
)

// KeyPairs is a slice of KeyPair pointers
type KeyPairs map[string]string

func init() {
	if ip, err := getDefaultRouteGateway(); err != nil {
		log.Printf("[INFO] envetcd error getting default gateway: %v\n", err)
	} else {
		gatewayIP = &ip
	}
}

func massagePeers(config *Config) error {
	for i, ep := range config.Peers {

		u, err := url.Parse(ep)
		if err != nil {
			return err
		}

		if u.Scheme == "" {
			u.Scheme = "http"
		}

		config.Peers[i] = u.String()
	}

	return nil
}

func getClient(config *Config) (*etcd.Client, error) {
	if config == nil {
		return nil, errors.New("config is nil")
	}

	if err := massagePeers(config); err != nil {
		return nil, err
	}

	client := etcd.NewClient(config.Peers)

	if config.TLS != nil {
		tr, err := transport.NewTransport(*config.TLS)
		if err != nil {
			return nil, err
		}

		client.SetTransport(tr)
	}

	// Sync cluster.
	if config.Sync {
		if ok := client.SyncCluster(); !ok {
			return nil, errors.New("cannot sync with the cluster using endpoints \"" + strings.Join(config.Peers, "\", \"") + "\"")
		}
	}

	return client, nil
}

// Set modifies the current environment with variables retrieved from etcd. Set
// will not overwrite existing variables.
// On linux systems, the default gateway will be automatically used as the etcd
// endpoint.
// If $ETCD_ENDPOINT is set, it will override the default gateway.
// $ETCD_ENDPOINT should look like "http://127.0.0.1:4001".
// service should be set by the application calling Set and not derived from
// an environment variable.
// Set will also use some other environment variables if they exist.
// $ETCD_PREFIX defaults to "/config"
// $HOSTNAME will be honored if it is set.
// An error is returned only if there was an actual error. Inability to
// determine the etcd endpoint as tolerated and not considered an error. In this
// case Set will simply not do anyting and it is up to the application to ensure
// that it has been configured properly through other means such as environment
// variables or command line flags.
func Set(service string) error {
	etcdEndpoint := os.Getenv("ETCD_ENDPOINT")

	useDefaultGateway := true
	if len(os.Getenv("ENVETCD_USE_DEFAULT_GATEWAY")) > 0 {
		if val, err := strconv.ParseBool(os.Getenv("ENVETCD_USE_DEFAULT_GATEWAY")); err != nil {
			log.Printf("[INFO] envetcd.Set could not parse $ENVETCD_USE_DEFAULT_GATEWAY, defaulting to true: %v\n", err)
		} else {
			useDefaultGateway = val
		}
	}

	if gatewayIP != nil && useDefaultGateway && len(etcdEndpoint) == 0 {
		etcdEndpoint = fmt.Sprintf("http://%s:4001", gatewayIP.String())
	}

	if len(etcdEndpoint) == 0 {
		log.Println("[INFO] envetcd.Set returned after it could not determine the etcd endpoint")
		return nil
	}

	config := &Config{
		Peers:             []string{etcdEndpoint},
		Sync:              true,
		Sanitize:          true,
		Upcase:            true,
		UseDefaultGateway: useDefaultGateway,
		Prefix:            os.Getenv("ETCD_PREFIX"),
		Service:           service,
		Hostname:          os.Getenv("HOSTNAME"),
	}

	if len(config.Peers[0]) == 0 {
		config.Peers[0] = "http://127.0.0.1:4001"
	}

	if len(config.Prefix) == 0 {
		config.Prefix = "/config"
	}

	keyPairs, err := GetKeyPairs(config)
	if err != nil {
		return err
	}

	for key, value := range keyPairs {
		if len(os.Getenv(key)) == 0 {
			os.Setenv(key, value)
		}
	}

	return nil
}

// GetKeyPairs takes a given config and client, and returns all key pairs
func GetKeyPairs(config *Config) (KeyPairs, error) {
	const noSort = false
	const recursive = true

	if len(config.Service) > 0 && len(config.System) == 0 {
		if index := strings.Index(config.Service, "-"); index != -1 {
			config.System = config.Service[:index]
		}
	}

	client, err := getClient(config)
	if err != nil {
		return nil, err
	}

	keyPairs := make(KeyPairs)

	// For each etcd directory key template
	// fetch from etcd and add to keypairs
	for i, tmpl := range etcdKeyTemplates {
		t := template.Must(template.New("etcdKey" + strconv.Itoa(i)).Parse(tmpl))
		var buf bytes.Buffer
		err := t.Execute(&buf, config)
		if err != nil {
			continue
		}

		if buf.Len() == 0 {
			continue
		}

		dir := buf.String()

		resp, err := client.Get(dir, noSort, recursive)
		if err != nil {
			continue
		}

		addKeyPair(config, keyPairs, dir, resp.Node)
	}

	// Expose --service --system and --hostname command variables to the child process' emvironment,
	// ignoring if these are not set in etcd (or set in etcd. Command line arguments take precedence.)
	if "" != config.Service {
		keyPairs["ENVETCD_SERVICE"] = config.Service
	}

	if "" != config.System {
		keyPairs["ENVETCD_SYSTEM"] = config.System
	}

	if "" != config.Hostname {
		keyPairs["ENVETCD_HOSTNAME"] = config.Hostname
	}

	if config.UseDefaultGateway && gatewayIP != nil {
		keyPairs["ENVETCD_DEFAULT_GATEWAY"] = gatewayIP.String()
	}

	return keyPairs, nil
}

func addKeyPair(config *Config, keyPairs KeyPairs, dir string, node *etcd.Node) {
	if node.Dir {
		for _, child := range node.Nodes {
			addKeyPair(config, keyPairs, dir, child)
		}
		return
	}

	key := strings.TrimPrefix(node.Key, dir) // strip the prefix directory from the key
	if key == "" {
		return
	}

	key = strings.TrimLeft(key, "/")         // strip any leading slashes
	key = strings.Replace(key, "/", "_", -1) // convert any remaining slashes to underscores (for any nested keys)

	if config.Sanitize {
		key = invalidRegexp.ReplaceAllString(key, "_")
	}

	if config.Upcase {
		key = strings.ToUpper(key)
	}

	keyPairs[key] = node.Value
}
<commit_msg>add some debugging<commit_after>package envetcd

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"net"
	"net/url"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"text/template"

	"github.com/coreos/etcd/pkg/transport"
	"github.com/coreos/go-etcd/etcd"
)

// order of precedence:
// global < system < service < host

// Regexp for invalid characters in keys
var invalidRegexp = regexp.MustCompile(`[^a-zA-Z0-9_]`)

// Config contains all of the parameters needed to run GetKeyPairs
type Config struct {
	Peers             []string
	Sync              bool
	Sanitize          bool
	Upcase            bool
	UseDefaultGateway bool
	Prefix            string
	System            string
	Service           string
	Hostname          string
	TLS               *transport.TLSInfo
}

var (
	etcdKeyTemplates = [...]string{
		"{{.Prefix}}/global",
		"{{if .System}}{{.Prefix}}/system/{{.System}}{{end}}",
		"{{if .Service}}{{.Prefix}}/service/{{.Service}}{{end}}",
		"{{if .Hostname}}{{.Prefix}}/host/{{.Hostname}}{{end}}",
	}

	gatewayIP *net.IP
)

// KeyPairs is a slice of KeyPair pointers
type KeyPairs map[string]string

func init() {
	if ip, err := getDefaultRouteGateway(); err != nil {
		log.Printf("[INFO] envetcd error getting default gateway: %v\n", err)
	} else {
		gatewayIP = &ip
	}
}

func massagePeers(config *Config) error {
	for i, ep := range config.Peers {

		u, err := url.Parse(ep)
		if err != nil {
			return err
		}

		if u.Scheme == "" {
			u.Scheme = "http"
		}

		config.Peers[i] = u.String()
	}

	return nil
}

func getClient(config *Config) (*etcd.Client, error) {
	if config == nil {
		return nil, errors.New("config is nil")
	}

	if err := massagePeers(config); err != nil {
		return nil, err
	}

	client := etcd.NewClient(config.Peers)

	if config.TLS != nil {
		tr, err := transport.NewTransport(*config.TLS)
		if err != nil {
			return nil, err
		}

		client.SetTransport(tr)
	}

	// Sync cluster.
	if config.Sync {
		if ok := client.SyncCluster(); !ok {
			return nil, errors.New("cannot sync with the cluster using endpoints \"" + strings.Join(config.Peers, "\", \"") + "\"")
		}
	}

	return client, nil
}

// Set modifies the current environment with variables retrieved from etcd. Set
// will not overwrite existing variables.
// On linux systems, the default gateway will be automatically used as the etcd
// endpoint.
// If $ETCD_ENDPOINT is set, it will override the default gateway.
// $ETCD_ENDPOINT should look like "http://127.0.0.1:4001".
// service should be set by the application calling Set and not derived from
// an environment variable.
// Set will also use some other environment variables if they exist.
// $ETCD_PREFIX defaults to "/config"
// $HOSTNAME will be honored if it is set.
// An error is returned only if there was an actual error. Inability to
// determine the etcd endpoint as tolerated and not considered an error. In this
// case Set will simply not do anyting and it is up to the application to ensure
// that it has been configured properly through other means such as environment
// variables or command line flags.
func Set(service string) error {
	etcdEndpoint := os.Getenv("ETCD_ENDPOINT")

	useDefaultGateway := true
	if len(os.Getenv("ENVETCD_USE_DEFAULT_GATEWAY")) > 0 {
		if val, err := strconv.ParseBool(os.Getenv("ENVETCD_USE_DEFAULT_GATEWAY")); err != nil {
			log.Printf("[INFO] envetcd.Set could not parse $ENVETCD_USE_DEFAULT_GATEWAY, defaulting to true: %v\n", err)
		} else {
			useDefaultGateway = val
		}
	}

	if gatewayIP != nil && useDefaultGateway && len(etcdEndpoint) == 0 {
		etcdEndpoint = fmt.Sprintf("http://%s:4001", gatewayIP.String())
	}

	if len(etcdEndpoint) == 0 {
		log.Println("[INFO] envetcd.Set returned after it could not determine the etcd endpoint")
		return nil
	}

	config := &Config{
		Peers:             []string{etcdEndpoint},
		Sync:              true,
		Sanitize:          true,
		Upcase:            true,
		UseDefaultGateway: useDefaultGateway,
		Prefix:            os.Getenv("ETCD_PREFIX"),
		Service:           service,
		Hostname:          os.Getenv("HOSTNAME"),
	}

	if len(config.Peers[0]) == 0 {
		config.Peers[0] = "http://127.0.0.1:4001"
	}

	if len(config.Prefix) == 0 {
		config.Prefix = "/config"
	}

	keyPairs, err := GetKeyPairs(config)
	if err != nil {
		return err
	}

	for key, value := range keyPairs {
		if len(os.Getenv(key)) == 0 {
			os.Setenv(key, value)
		}
	}

	return nil
}

// GetKeyPairs takes a given config and client, and returns all key pairs
func GetKeyPairs(config *Config) (KeyPairs, error) {
	const noSort = false
	const recursive = true

	if len(config.Service) > 0 && len(config.System) == 0 {
		if index := strings.Index(config.Service, "-"); index != -1 {
			config.System = config.Service[:index]
		}
	}

	client, err := getClient(config)
	if err != nil {
		return nil, err
	}

	keyPairs := make(KeyPairs)

	// For each etcd directory key template
	// fetch from etcd and add to keypairs
	for i, tmpl := range etcdKeyTemplates {
		t := template.Must(template.New("etcdKey" + strconv.Itoa(i)).Parse(tmpl))
		var buf bytes.Buffer
		err := t.Execute(&buf, config)
		if err != nil {
			continue
		}

		if buf.Len() == 0 {
			continue
		}

		dir := buf.String()

		resp, err := client.Get(dir, noSort, recursive)
		if err != nil {
			continue
		}

		addKeyPair(config, keyPairs, dir, resp.Node)
	}

	// Expose --service --system and --hostname command variables to the child process' emvironment,
	// ignoring if these are not set in etcd (or set in etcd. Command line arguments take precedence.)
	if "" != config.Service {
		keyPairs["ENVETCD_SERVICE"] = config.Service
	}

	if "" != config.System {
		keyPairs["ENVETCD_SYSTEM"] = config.System
	}

	if "" != config.Hostname {
		keyPairs["ENVETCD_HOSTNAME"] = config.Hostname
	}

	if config.UseDefaultGateway && gatewayIP != nil {
		keyPairs["ENVETCD_DEFAULT_GATEWAY"] = gatewayIP.String()
	}

	var keys []string
	for key := range keyPairs {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, key := range keys {
		log.Printf("[DEBUG] envetcd: %v => %v\n", key, keyPairs[key])
	}

	return keyPairs, nil
}

func addKeyPair(config *Config, keyPairs KeyPairs, dir string, node *etcd.Node) {
	if node.Dir {
		for _, child := range node.Nodes {
			addKeyPair(config, keyPairs, dir, child)
		}
		return
	}

	key := strings.TrimPrefix(node.Key, dir) // strip the prefix directory from the key
	if key == "" {
		return
	}

	key = strings.TrimLeft(key, "/")         // strip any leading slashes
	key = strings.Replace(key, "/", "_", -1) // convert any remaining slashes to underscores (for any nested keys)

	if config.Sanitize {
		key = invalidRegexp.ReplaceAllString(key, "_")
	}

	if config.Upcase {
		key = strings.ToUpper(key)
	}

	keyPairs[key] = node.Value
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"os"
	"sort"
)

var (
	client       *etcd.Client
	printVersion bool
	outputFormat string
	debug        bool
	cluster      = ClusterValue{"http://localhost:4001"}
	curlChan     chan string
)

func output(resp *etcd.Response) {
	switch outputFormat {
	case "json":
		b, err := json.Marshal(resp)
		if err != nil {
			panic(err)
		}
		fmt.Printf("%s\n", b)
	case "full":
		fmt.Printf("Index: %v\nValue: %v\n", resp.Index, resp.Value)
	case "value-only":
		fmt.Println(resp.Value)
	}
}

func main() {
	flag.BoolVar(&printVersion, "version", false, "print the version and exit")
	flag.Var(&cluster, "C", "a comma seperated list of machine addresses in the cluster e.g. 127.0.0.1:4001,127.0.0.1:4002")
	flag.StringVar(&outputFormat, "format", "json", "Output server response in the given format, either `json`, `full`, or `value-only`")
	flag.BoolVar(&debug, "debug", false, "Output cURL commands which can be used to re-produce the request")
	flag.Parse()

	if printVersion {
		fmt.Println(releaseVersion)
		os.Exit(0)
	}

	if debug {
		// Making the channel buffered to avoid potential deadlocks
		curlChan = make(chan string, 10)
		etcd.SetCurlChan(curlChan)
	}

	client = etcd.NewClient(cluster.GetMachines())

	args := flag.Args()

	if len(args) == 0 {
		fmt.Println("Usage: etcdctl [flags] [command] [flags for command]\n")
		fmt.Println("Available flags include:\n")
		flag.PrintDefaults()
		fmt.Println()
		fmt.Println("Available commands:")
		fmt.Println()
		slice := make([]string, len(commands))
		i := 0
		for k, _ := range commands {
			slice[i] = k
			i++
		}
		sort.Strings(slice)
		for _, c := range slice {
			fmt.Println(c)
		}
		fmt.Println()
		fmt.Println(`To see the full usage for a specific command, run "etcdctl [command]"`)
		fmt.Println()
		os.Exit(1)
	}

	commandName := args[0]

	command, ok := commands[commandName]

	if !ok {
		fmt.Println("wrong command provided")
		os.Exit(MalformedEtcdctlArguments)
	}

	commandArgs := args[1:]

	if len(commandArgs) > command.maxArgs || len(commandArgs) < command.minArgs {
		fmt.Println(command.usage)
		os.Exit(MalformedEtcdctlArguments)
	}

	if !client.SyncCluster() {
		fmt.Println("cannot sync with the given cluster")
		os.Exit(FailedToConnectToHost)
	}

	err := command.f(commandArgs)

	if err != nil {
		fmt.Print("Error: ")
		fmt.Println(err)
		os.Exit(ErrorFromEtcd)
	}
}
<commit_msg>Accommodate changes in go-etcd<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"os"
	"sort"
)

var (
	client       *etcd.Client
	printVersion bool
	outputFormat string
	debug        bool
	cluster      = ClusterValue{"http://localhost:4001"}
	curlChan     chan string
)

func output(resp *etcd.Response) {
	switch outputFormat {
	case "json":
		b, err := json.Marshal(resp)
		if err != nil {
			panic(err)
		}
		fmt.Printf("%s\n", b)
	case "full":
		fmt.Printf("Index: %v\nValue: %v\n", resp.ModifiedIndex, resp.Value)
	case "value-only":
		fmt.Println(resp.Value)
	}
}

func main() {
	flag.BoolVar(&printVersion, "version", false, "print the version and exit")
	flag.Var(&cluster, "C", "a comma seperated list of machine addresses in the cluster e.g. 127.0.0.1:4001,127.0.0.1:4002")
	flag.StringVar(&outputFormat, "format", "json", "Output server response in the given format, either `json`, `full`, or `value-only`")
	flag.BoolVar(&debug, "debug", false, "Output cURL commands which can be used to re-produce the request")
	flag.Parse()

	if printVersion {
		fmt.Println(releaseVersion)
		os.Exit(0)
	}

	if debug {
		// Making the channel buffered to avoid potential deadlocks
		curlChan = make(chan string, 10)
		etcd.SetCurlChan(curlChan)
	}

	client = etcd.NewClient(cluster.GetMachines())

	args := flag.Args()

	if len(args) == 0 {
		fmt.Println("Usage: etcdctl [flags] [command] [flags for command]\n")
		fmt.Println("Available flags include:\n")
		flag.PrintDefaults()
		fmt.Println()
		fmt.Println("Available commands:")
		fmt.Println()
		slice := make([]string, len(commands))
		i := 0
		for k, _ := range commands {
			slice[i] = k
			i++
		}
		sort.Strings(slice)
		for _, c := range slice {
			fmt.Println(c)
		}
		fmt.Println()
		fmt.Println(`To see the full usage for a specific command, run "etcdctl [command]"`)
		fmt.Println()
		os.Exit(1)
	}

	commandName := args[0]

	command, ok := commands[commandName]

	if !ok {
		fmt.Println("wrong command provided")
		os.Exit(MalformedEtcdctlArguments)
	}

	commandArgs := args[1:]

	if len(commandArgs) > command.maxArgs || len(commandArgs) < command.minArgs {
		fmt.Println(command.usage)
		os.Exit(MalformedEtcdctlArguments)
	}

	if !client.SyncCluster() {
		fmt.Println("cannot sync with the given cluster")
		os.Exit(FailedToConnectToHost)
	}

	err := command.f(commandArgs)

	if err != nil {
		fmt.Print("Error: ")
		fmt.Println(err)
		os.Exit(ErrorFromEtcd)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

var sep = flag.Bool("s", false, "separate arguments with spaces")
var key = flag.String("key", "", "etcd key")
var host = flag.String("host", "", "etcd host")
var hosts []string

func main() {
	if len(os.Args) > 1 && strings.HasPrefix(os.Args[1], "-s ") {
		args := []string{}
		for _, arg := range os.Args {
			args = append(args, strings.Split(arg, " ")...)
		}
		os.Args = args
	}
	flag.Parse()

	if *key == "" {
		*key = os.Getenv("ETCDENV_KEY")
	}
	if *key == "" {
		fmt.Fprintln(os.Stderr, "etcdenv [-key=key] [...]")
		flag.PrintDefaults()
		os.Exit(1)
	}

	if *host == "" {
		*host = os.Getenv("ETCDENV_HOST")
	}
	if *host != "" {
		hosts = []string{*host}
	}
	client := etcd.NewClient(hosts)
	res, err := client.Get(*key, true, true)
	if err != nil {
		fmt.Fprintf(os.Stderr, "etcdenv: %s\n", err)
		os.Exit(1)
	}

	envs := os.Environ()
	for _, n := range res.Node.Nodes {
		key := strings.Split(n.Key, "/")
		k, v := strings.ToUpper(key[len(key)-1]), n.Value
		envs = append(envs, k+"="+v)
	}

	if flag.NArg() == 0 {
		for _, env := range envs {
			line := fmt.Sprintf("%q", env)
			fmt.Println(line[1 : len(line)-1])
		}
		os.Exit(0)
	}

	cmd := exec.Command(flag.Args()[0], flag.Args()[1:]...)
	cmd.Env = envs
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	err = cmd.Run()
	if cmd.Process == nil {
		fmt.Fprintf(os.Stderr, "etcdenv: %s\n", err)
		os.Exit(1)
	}
	os.Exit(cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus())
}
<commit_msg>Make -key flag optional<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/coreos/go-etcd/etcd"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

var sep = flag.Bool("s", false, "separate arguments with spaces")
var key = flag.String("key", "", "etcd key")
var host = flag.String("host", "", "etcd host")
var hosts []string

func main() {
	if len(os.Args) > 1 && strings.HasPrefix(os.Args[1], "-s ") {
		args := []string{}
		for _, arg := range os.Args {
			args = append(args, strings.Split(arg, " ")...)
		}
		os.Args = args
	}
	flag.Parse()

	if *key == "" {
		*key = os.Getenv("ETCDENV_KEY")
	}
	if *host == "" {
		*host = os.Getenv("ETCDENV_HOST")
	}
	if *host != "" {
		hosts = []string{*host}
	}
	client := etcd.NewClient(hosts)
	res, err := client.Get(*key, true, true)
	if err != nil {
		fmt.Fprintf(os.Stderr, "etcdenv: %s\n", err)
		os.Exit(1)
	}

	envs := os.Environ()
	for _, n := range res.Node.Nodes {
		if !n.Dir {
			key := strings.Split(n.Key, "/")
			k, v := strings.ToUpper(key[len(key)-1]), n.Value
			envs = append(envs, k+"="+v)
		}
	}

	if flag.NArg() == 0 {
		for _, env := range envs {
			line := fmt.Sprintf("%q", env)
			fmt.Println(line[1 : len(line)-1])
		}
		os.Exit(0)
	}

	cmd := exec.Command(flag.Args()[0], flag.Args()[1:]...)
	cmd.Env = envs
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	cmd.Stdin = os.Stdin
	err = cmd.Run()
	if cmd.Process == nil {
		fmt.Fprintf(os.Stderr, "etcdenv: %s\n", err)
		os.Exit(1)
	}
	os.Exit(cmd.ProcessState.Sys().(syscall.WaitStatus).ExitStatus())
}
<|endoftext|>


<commit_before>package event_lib

import (
	"bytes"
	"errors"
	"log"
	"text/template"
	"time"
	//"net/url"

	"gopkg.in/mgo.v2/bson"

	"scal"
	"scal/settings"
	"scal/storage"
	. "scal/user_lib"
)

/*
// PublicJoinPolicy
const (
    FreeToJoin = "FreeToJoin"
    NoJoin = "NoJoin"
    //JoinNeedsVerify = "JoinNeedsVerify" // makes implementation more complex
)
*/

type EventMetaModel struct {
	EventId      bson.ObjectId    `bson:"_id"`
	EventType    string           `bson:"eventType"`
	CreationTime time.Time        `bson:"creationTime"`
	OwnerEmail   string           `bson:"ownerEmail"`
	IsPublic     bool             `bson:"isPublic"`
	AccessEmails []string         `bson:"accessEmails"`
	GroupId      *bson.ObjectId   `bson:"groupId"`
	GroupModel   *EventGroupModel `bson:"-"`

	//PublicJoinPolicy string         `bson:"publicJoinPolicy"` // not indexed
	PublicJoinOpen bool `bson:"publicJoinOpen"`
	MaxAttendees   int  `bson:"maxAttendees"`
}

type InviteEmailTemplateParams struct {
	EventModel  *BaseEventModel
	SenderEmail string
	SenderName  string
	Email       string
	Name        string
	JoinURL     string
}

func (self EventMetaModel) UniqueM() scal.M {
	return scal.M{
		"_id": self.EventId,
	}
}
func (self EventMetaModel) Collection() string {
	return storage.C_eventMeta
}
func (self EventMetaModel) GroupIdHex() string {
	if self.GroupId != nil {
		return self.GroupId.Hex()
	}
	return ""
}
func (self EventMetaModel) JsonM() scal.M {
	return scal.M{
		"creationTime":   self.CreationTime,
		"ownerEmail":     self.OwnerEmail,
		"isPublic":       self.IsPublic,
		"accessEmails":   self.AccessEmails,
		"groupId":        self.GroupIdHex(),
		"publicJoinOpen": self.PublicJoinOpen,
		"maxAttendees":   self.MaxAttendees,
	}
}
func (self *EventMetaModel) CanReadFull(email string) bool {
	if email == self.OwnerEmail {
		return true
	}
	for _, aEmail := range self.AccessEmails {
		if email == aEmail {
			return true
		}
	}
	if self.GroupModel != nil {
		for _, aEmail := range self.GroupModel.ReadAccessEmails {
			if email == aEmail {
				return true
			}
		}
	}
	return false
}
func (self *EventMetaModel) CanRead(email string) bool {
	if self.IsPublic {
		return true
	}
	return self.CanReadFull(email)
}
func (self *EventMetaModel) GetAttending(
	db storage.Database,
	email string,
) string {
	// returns YES, NO, or MAYBE
	attendingModel, _ := LoadEventAttendingModel(db, self.EventId, email)
	return attendingModel.Attending
}
func (self *EventMetaModel) SetAttending(
	db storage.Database,
	email string,
	attending string,
) error {
	// attending: YES, NO, or MAYBE
	attendingModel, err := LoadEventAttendingModel(db, self.EventId, email)
	if err != nil {
		return err
	}
	attendingModel.Attending = attending
	attendingModel.ModifiedTime = time.Now()
	err = attendingModel.Save(db)
	return err
}
func (self *EventMetaModel) AttendingStatusCount(
	db storage.Database,
	attending string,
) (int, error) {
	return db.FindCount(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
	)
}
func (self *EventMetaModel) Join(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == YES {
		return errors.New("you have already joined this event")
	}
	if !self.CanReadFull(email) {
		if self.IsPublic {
			if !self.PublicJoinOpen {
				return errors.New("this public event is not open for joining")
			}
		} else {
			return errors.New("no access, no join")
		}
	}
	if self.MaxAttendees > 0 {
		attendingCount, err := self.AttendingStatusCount(db, YES)
		if err != nil {
			return err
		}
		if attendingCount >= self.MaxAttendees {
			return errors.New("maximum attendees exceeded, can not join event")
		}
	}
	self.SetAttending(db, email, YES)
	return nil
}
func (self *EventMetaModel) Leave(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == NO {
		if self.CanReadFull(email) {
			return errors.New("you are not attending for this event")
		}
	}
	self.SetAttending(db, email, NO)
	return nil
}
func (self *EventMetaModel) PublicCanJoin() bool {
	return self.IsPublic && self.PublicJoinOpen
}
func (self *EventMetaModel) Invite(
	db storage.Database,
	email string,
	inviteEmails *[]string,
	remoteIp string,
	host string,
) (error, int) {
	/*
		returns (err, errCode)
		errCode values:
			scal.BadRequest
			scal.Forbidden
			scal.InternalServerError
	*/
	var err error
	if inviteEmails == nil {
		return errors.New("missing 'inviteEmails'"), scal.BadRequest
	}
	if len(*inviteEmails) == 0 {
		return errors.New("empty 'inviteEmails'"), scal.BadRequest
	}

	now := time.Now()

	fullAc := self.CanReadFull(email)
	public := self.PublicCanJoin()
	if !(fullAc || public) {
		return errors.New("not allowed to invite to this event"), scal.Forbidden
	}
	if fullAc {
		accessEmailsMap := make(map[string]bool)
		newAccessEmails := make(
			[]string,
			0,
			len(self.AccessEmails)+len(*inviteEmails),
		)
		for _, aEmail := range self.AccessEmails {
			accessEmailsMap[aEmail] = true
			newAccessEmails = append(newAccessEmails, aEmail)
		}
		for _, inviteEmail := range *inviteEmails {
			_, found := accessEmailsMap[inviteEmail]
			if !found {
				newAccessEmails = append(newAccessEmails, inviteEmail)
			}
		}
		metaChangeLog := EventMetaChangeLogModel{
			Time:     now,
			Email:    email,
			RemoteIp: remoteIp,
			EventId:  self.EventId,
			FuncName: "SetEventAccess",
			AccessEmails: &[2][]string{
				self.AccessEmails,
				newAccessEmails,
			},
		}
		self.AccessEmails = newAccessEmails
		err = db.Insert(metaChangeLog)
		if err != nil {
			return err, scal.InternalServerError
		}
		err = db.Update(self)
		if err != nil {
			return err, scal.InternalServerError
		}
	}

	eventRev, err := LoadLastRevisionModel(db, &self.EventId)
	if err != nil {
		if db.IsNotFound(err) {
			return errors.New("event not found"), scal.BadRequest
		} else {
			return err, scal.InternalServerError
		}
	}
	eventModel, err := LoadBaseEventModel(db, eventRev.Sha1)
	if err != nil {
		return err, scal.InternalServerError
	}

	tplText := settings.EVENT_INVITE_EMAIL_TEMPLATE
	user := UserModelByEmail(email, db)
	if user == nil {
		return errors.New("user not found"), scal.InternalServerError
	}
	eventUrl := host +
		"/event/" +
		self.EventType + "/" +
		self.EventId.Hex() + "/"
	joinURL := eventUrl + "join"
	for _, inviteEmail := range *inviteEmails {
		subject := "Invitation to event: " + eventModel.Summary
		tpl, err := template.New(subject).Parse(tplText)
		if err != nil {
			return err, scal.InternalServerError
		}
		var inviteName string
		inviteUser := UserModelByEmail(inviteEmail, db)
		if inviteUser == nil {
			//return errors.New("invited email not found"), scal.BadRequest
			// FIXME
			inviteName = ""
		} else {
			inviteName = inviteUser.FullName
		}

		tplParams := InviteEmailTemplateParams{
			EventModel:  eventModel,
			SenderEmail: email,
			SenderName:  user.FullName,
			Email:       inviteEmail,
			Name:        inviteName,
			JoinURL:     joinURL,
		}
		buf := bytes.NewBufferString("")
		err = tpl.Execute(buf, tplParams)
		if err != nil {
			return err, scal.InternalServerError
		}
		emailBody := buf.String()
		db.Insert(EventInvitationModel{
			Time:         now,
			SenderEmail:  email,
			InvitedEmail: inviteEmail,
			EventId:      self.EventId,
		})

		scal.SendEmail(
			inviteEmail,
			subject,
			false, // isHtml
			emailBody,
		)
	}
	return nil, 200
}
func (self *EventMetaModel) GetEmailsByAttendingStatus(
	db storage.Database,
	attending string,
) []string {
	emailStructs := []struct {
		Email string `bson:"email"`
	}{}
	err := db.FindAll(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
		&emailStructs,
	)
	if err != nil {
		log.Printf(
			"Internal Error: GetAttendingEmails: eventId=%v: %s\n",
			self.EventId,
			err.Error(),
		)
	}
	emails := make([]string, len(emailStructs))
	for i, m := range emailStructs {
		emails[i] = m.Email
	}
	return emails
}
func (self *EventMetaModel) GetAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, YES)
}
func (self *EventMetaModel) GetNotAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, NO)
}
func (self *EventMetaModel) GetMaybeAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, MAYBE)
}

func LoadEventMetaModel(
	db storage.Database,
	eventId *bson.ObjectId,
	loadGroup bool,
) (*EventMetaModel, error) {
	var err error
	eventMeta := EventMetaModel{
		EventId: *eventId,
	}
	err = db.Get(&eventMeta)
	if err != nil {
		return nil, err
	}
	if loadGroup && eventMeta.GroupId != nil {
		// groupModel, err, internalErr
		groupModel, err, _ := LoadGroupModelById(
			"groupId",
			db,
			eventMeta.GroupId,
		)
		if err != nil {
			return nil, err
		}
		eventMeta.GroupModel = groupModel
	}
	return &eventMeta, nil
}

type EventMetaChangeLogModel struct {
	Time     time.Time     `bson:"time"`
	Email    string        `bson:"email"`
	RemoteIp string        `bson:"remoteIp"`
	EventId  bson.ObjectId `bson:"eventId"`
	FuncName string        `bson:"funcName"`

	GroupId        *[2]*bson.ObjectId `bson:"groupId,omitempty"`
	OwnerEmail     *[2]*string        `bson:"ownerEmail,omitempty"`
	IsPublic       *[2]bool           `bson:"isPublic,omitempty"`
	AccessEmails   *[2][]string       `bson:"accessEmails,omitempty"`
	PublicJoinOpen *[2]bool           `bson:"publicJoinOpen,omitempty"`
	MaxAttendees   *[2]int            `bson:"maxAttendees,omitempty"`
}

func (model EventMetaChangeLogModel) Collection() string {
	return storage.C_eventMetaChangeLog
}

type EventInvitationModel struct {
	Time         time.Time     `bson:"time"`
	SenderEmail  string        `bson:"senderEmail"`
	InvitedEmail string        `bson:"invitedEmail"`
	EventId      bson.ObjectId `bson:"eventId"`
}

func (model EventInvitationModel) Collection() string {
	return storage.C_invitation
}

func GetEventMetaPipeResults(
	db storage.Database,
	pipeline *[]scal.M,
) (*[]scal.M, error) {
	results := []scal.M{}
	for res := range db.PipeIter(storage.C_eventMeta, pipeline) {
		if err := res.Err; err != nil {
			return nil, err
		}
		if eventId, ok := res.M["_id"]; ok {
			res.M["eventId"] = eventId
			delete(res.M, "_id")
		}
		if dataI, ok := res.M["data"]; ok {
			data := dataI.(scal.M)
			delete(data, "_id")
			res.M["data"] = data
		}
		results = append(results, res.M)
	}
	return &results, nil
}
<commit_msg>add EventMetaModel.FieldsMtime, with fmt and style fixes<commit_after>package event_lib

import (
	"bytes"
	"errors"
	"log"
	"text/template"
	"time"
	//"net/url"

	"gopkg.in/mgo.v2/bson"

	"scal"
	"scal/settings"
	"scal/storage"
	. "scal/user_lib"
)

/*
// PublicJoinPolicy
const (
    FreeToJoin = "FreeToJoin"
    NoJoin = "NoJoin"
    //JoinNeedsVerify = "JoinNeedsVerify" // makes implementation more complex
)
*/

type EventMetaModel struct {
	EventId   bson.ObjectId `bson:"_id"`
	EventType string        `bson:"eventType"`

	CreationTime time.Time            `bson:"creationTime"`
	FieldsMtime  map[string]time.Time `bson:"fieldsMtime"`

	OwnerEmail   string   `bson:"ownerEmail"`
	IsPublic     bool     `bson:"isPublic"`
	AccessEmails []string `bson:"accessEmails"`

	GroupId    *bson.ObjectId   `bson:"groupId"`
	GroupModel *EventGroupModel `bson:"-"`

	//PublicJoinPolicy string         `bson:"publicJoinPolicy"` // not indexed
	PublicJoinOpen bool `bson:"publicJoinOpen"`
	MaxAttendees   int  `bson:"maxAttendees"`
}

type InviteEmailTemplateParams struct {
	EventModel  *BaseEventModel
	SenderEmail string
	SenderName  string
	Email       string
	Name        string
	JoinURL     string
}

func (self EventMetaModel) UniqueM() scal.M {
	return scal.M{
		"_id": self.EventId,
	}
}
func (self EventMetaModel) Collection() string {
	return storage.C_eventMeta
}
func (self EventMetaModel) GroupIdHex() string {
	if self.GroupId != nil {
		return self.GroupId.Hex()
	}
	return ""
}
func (self EventMetaModel) JsonM() scal.M {
	return scal.M{
		"creationTime":   self.CreationTime,
		"ownerEmail":     self.OwnerEmail,
		"isPublic":       self.IsPublic,
		"accessEmails":   self.AccessEmails,
		"groupId":        self.GroupIdHex(),
		"publicJoinOpen": self.PublicJoinOpen,
		"maxAttendees":   self.MaxAttendees,
	}
}
func (self *EventMetaModel) CanReadFull(email string) bool {
	if email == self.OwnerEmail {
		return true
	}
	for _, aEmail := range self.AccessEmails {
		if email == aEmail {
			return true
		}
	}
	if self.GroupModel != nil {
		for _, aEmail := range self.GroupModel.ReadAccessEmails {
			if email == aEmail {
				return true
			}
		}
	}
	return false
}
func (self *EventMetaModel) CanRead(email string) bool {
	if self.IsPublic {
		return true
	}
	return self.CanReadFull(email)
}
func (self *EventMetaModel) GetAttending(
	db storage.Database,
	email string,
) string {
	// returns YES, NO, or MAYBE
	attendingModel, _ := LoadEventAttendingModel(db, self.EventId, email)
	return attendingModel.Attending
}
func (self *EventMetaModel) SetAttending(
	db storage.Database,
	email string,
	attending string,
) error {
	// attending: YES, NO, or MAYBE
	attendingModel, err := LoadEventAttendingModel(db, self.EventId, email)
	if err != nil {
		return err
	}
	attendingModel.Attending = attending
	attendingModel.ModifiedTime = time.Now()
	err = attendingModel.Save(db)
	return err
}
func (self *EventMetaModel) AttendingStatusCount(
	db storage.Database,
	attending string,
) (int, error) {
	return db.FindCount(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
	)
}
func (self *EventMetaModel) Join(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == YES {
		return errors.New("you have already joined this event")
	}
	if !self.CanReadFull(email) {
		if self.IsPublic {
			if !self.PublicJoinOpen {
				return errors.New("this public event is not open for joining")
			}
		} else {
			return errors.New("no access, no join")
		}
	}
	if self.MaxAttendees > 0 {
		attendingCount, err := self.AttendingStatusCount(db, YES)
		if err != nil {
			return err
		}
		if attendingCount >= self.MaxAttendees {
			return errors.New("maximum attendees exceeded, can not join event")
		}
	}
	self.SetAttending(db, email, YES)
	return nil
}
func (self *EventMetaModel) Leave(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == NO {
		if self.CanReadFull(email) {
			return errors.New("you are not attending for this event")
		}
	}
	self.SetAttending(db, email, NO)
	return nil
}
func (self *EventMetaModel) PublicCanJoin() bool {
	return self.IsPublic && self.PublicJoinOpen
}
func (self *EventMetaModel) Invite(
	db storage.Database,
	email string,
	inviteEmails *[]string,
	remoteIp string,
	host string,
) (error, int) {
	/*
		returns (err, errCode)
		errCode values:
			scal.BadRequest
			scal.Forbidden
			scal.InternalServerError
	*/
	var err error
	if inviteEmails == nil {
		return errors.New("missing 'inviteEmails'"), scal.BadRequest
	}
	if len(*inviteEmails) == 0 {
		return errors.New("empty 'inviteEmails'"), scal.BadRequest
	}

	now := time.Now()

	fullAc := self.CanReadFull(email)
	public := self.PublicCanJoin()
	if !(fullAc || public) {
		return errors.New("not allowed to invite to this event"), scal.Forbidden
	}
	if fullAc {
		accessEmailsMap := make(map[string]bool)
		newAccessEmails := make(
			[]string,
			0,
			len(self.AccessEmails)+len(*inviteEmails),
		)
		for _, aEmail := range self.AccessEmails {
			accessEmailsMap[aEmail] = true
			newAccessEmails = append(newAccessEmails, aEmail)
		}
		for _, inviteEmail := range *inviteEmails {
			_, found := accessEmailsMap[inviteEmail]
			if !found {
				newAccessEmails = append(newAccessEmails, inviteEmail)
			}
		}
		metaChangeLog := EventMetaChangeLogModel{
			Time:     now,
			Email:    email,
			RemoteIp: remoteIp,
			EventId:  self.EventId,
			FuncName: "SetEventAccess",
			AccessEmails: &[2][]string{
				self.AccessEmails,
				newAccessEmails,
			},
		}
		self.AccessEmails = newAccessEmails
		err = db.Insert(metaChangeLog)
		if err != nil {
			return err, scal.InternalServerError
		}
		err = db.Update(self)
		if err != nil {
			return err, scal.InternalServerError
		}
	}

	eventRev, err := LoadLastRevisionModel(db, &self.EventId)
	if err != nil {
		if db.IsNotFound(err) {
			return errors.New("event not found"), scal.BadRequest
		} else {
			return err, scal.InternalServerError
		}
	}
	eventModel, err := LoadBaseEventModel(db, eventRev.Sha1)
	if err != nil {
		return err, scal.InternalServerError
	}

	tplText := settings.EVENT_INVITE_EMAIL_TEMPLATE
	user := UserModelByEmail(email, db)
	if user == nil {
		return errors.New("user not found"), scal.InternalServerError
	}
	eventUrl := host +
		"/event/" +
		self.EventType + "/" +
		self.EventId.Hex() + "/"
	joinURL := eventUrl + "join"
	for _, inviteEmail := range *inviteEmails {
		subject := "Invitation to event: " + eventModel.Summary
		tpl, err := template.New(subject).Parse(tplText)
		if err != nil {
			return err, scal.InternalServerError
		}
		var inviteName string
		inviteUser := UserModelByEmail(inviteEmail, db)
		if inviteUser == nil {
			//return errors.New("invited email not found"), scal.BadRequest
			// FIXME
			inviteName = ""
		} else {
			inviteName = inviteUser.FullName
		}

		tplParams := InviteEmailTemplateParams{
			EventModel:  eventModel,
			SenderEmail: email,
			SenderName:  user.FullName,
			Email:       inviteEmail,
			Name:        inviteName,
			JoinURL:     joinURL,
		}
		buf := bytes.NewBufferString("")
		err = tpl.Execute(buf, tplParams)
		if err != nil {
			return err, scal.InternalServerError
		}
		emailBody := buf.String()
		db.Insert(EventInvitationModel{
			Time:         now,
			SenderEmail:  email,
			InvitedEmail: inviteEmail,
			EventId:      self.EventId,
		})

		scal.SendEmail(
			inviteEmail,
			subject,
			false, // isHtml
			emailBody,
		)
	}
	return nil, 200
}
func (self *EventMetaModel) GetEmailsByAttendingStatus(
	db storage.Database,
	attending string,
) []string {
	emailStructs := []struct {
		Email string `bson:"email"`
	}{}
	err := db.FindAll(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
		&emailStructs,
	)
	if err != nil {
		log.Printf(
			"Internal Error: GetAttendingEmails: eventId=%v: %s\n",
			self.EventId,
			err.Error(),
		)
	}
	emails := make([]string, len(emailStructs))
	for i, m := range emailStructs {
		emails[i] = m.Email
	}
	return emails
}
func (self *EventMetaModel) GetAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, YES)
}
func (self *EventMetaModel) GetNotAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, NO)
}
func (self *EventMetaModel) GetMaybeAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, MAYBE)
}

func LoadEventMetaModel(
	db storage.Database,
	eventId *bson.ObjectId,
	loadGroup bool,
) (*EventMetaModel, error) {
	var err error
	eventMeta := EventMetaModel{
		EventId: *eventId,
	}
	err = db.Get(&eventMeta)
	if err != nil {
		return nil, err
	}
	if loadGroup && eventMeta.GroupId != nil {
		// groupModel, err, internalErr
		groupModel, err, _ := LoadGroupModelById(
			"groupId",
			db,
			eventMeta.GroupId,
		)
		if err != nil {
			return nil, err
		}
		eventMeta.GroupModel = groupModel
	}
	return &eventMeta, nil
}

type EventMetaChangeLogModel struct {
	Time     time.Time     `bson:"time"`
	Email    string        `bson:"email"`
	RemoteIp string        `bson:"remoteIp"`
	EventId  bson.ObjectId `bson:"eventId"`
	FuncName string        `bson:"funcName"`

	GroupId        *[2]*bson.ObjectId `bson:"groupId,omitempty"`
	OwnerEmail     *[2]*string        `bson:"ownerEmail,omitempty"`
	IsPublic       *[2]bool           `bson:"isPublic,omitempty"`
	AccessEmails   *[2][]string       `bson:"accessEmails,omitempty"`
	PublicJoinOpen *[2]bool           `bson:"publicJoinOpen,omitempty"`
	MaxAttendees   *[2]int            `bson:"maxAttendees,omitempty"`
}

func (model EventMetaChangeLogModel) Collection() string {
	return storage.C_eventMetaChangeLog
}

type EventInvitationModel struct {
	Time         time.Time     `bson:"time"`
	SenderEmail  string        `bson:"senderEmail"`
	InvitedEmail string        `bson:"invitedEmail"`
	EventId      bson.ObjectId `bson:"eventId"`
}

func (model EventInvitationModel) Collection() string {
	return storage.C_invitation
}

func GetEventMetaPipeResults(
	db storage.Database,
	pipeline *[]scal.M,
) (*[]scal.M, error) {
	results := []scal.M{}
	for res := range db.PipeIter(storage.C_eventMeta, pipeline) {
		if err := res.Err; err != nil {
			return nil, err
		}
		if eventId, ok := res.M["_id"]; ok {
			res.M["eventId"] = eventId
			delete(res.M, "_id")
		}
		if dataI, ok := res.M["data"]; ok {
			data := dataI.(scal.M)
			delete(data, "_id")
			res.M["data"] = data
		}
		results = append(results, res.M)
	}
	return &results, nil
}
<|endoftext|>


<commit_before>package event_lib

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"net/url"
	"text/template"
	"time"

	"github.com/ilius/restpc"
	//"net/url"

	"gopkg.in/mgo.v2/bson"

	"scal"
	"scal/settings"
	"scal/storage"
	. "scal/user_lib"
)

/*
// PublicJoinPolicy
const (
    FreeToJoin = "FreeToJoin"
    NoJoin = "NoJoin"
    //JoinNeedsVerify = "JoinNeedsVerify" // makes implementation more complex
)
*/

type EventMetaModel struct {
	EventId   bson.ObjectId `bson:"_id"`
	EventType string        `bson:"eventType"`

	CreationTime time.Time            `bson:"creationTime"`
	FieldsMtime  map[string]time.Time `bson:"fieldsMtime"`

	OwnerEmail   string   `bson:"ownerEmail"`
	IsPublic     bool     `bson:"isPublic"`
	AccessEmails []string `bson:"accessEmails"`

	GroupId    *bson.ObjectId   `bson:"groupId"`
	GroupModel *EventGroupModel `bson:"-"`

	//PublicJoinPolicy string         `bson:"publicJoinPolicy"` // not indexed
	PublicJoinOpen bool `bson:"publicJoinOpen"`
	MaxAttendees   int  `bson:"maxAttendees"`
}

type InviteEmailTemplateParams struct {
	EventModel   *BaseEventModel
	SenderEmail  string
	SenderName   string
	Email        string
	Name         string
	EventType    string
	EventId      string
	TokenEscaped string
	Host         string
}

func (self EventMetaModel) UniqueM() scal.M {
	return scal.M{
		"_id": self.EventId,
	}
}
func (self EventMetaModel) Collection() string {
	return storage.C_eventMeta
}
func (self EventMetaModel) GroupIdHex() string {
	if self.GroupId != nil {
		return self.GroupId.Hex()
	}
	return ""
}
func (self EventMetaModel) JsonM() scal.M {
	return scal.M{
		"creationTime":   self.CreationTime,
		"ownerEmail":     self.OwnerEmail,
		"isPublic":       self.IsPublic,
		"accessEmails":   self.AccessEmails,
		"groupId":        self.GroupIdHex(),
		"publicJoinOpen": self.PublicJoinOpen,
		"maxAttendees":   self.MaxAttendees,
	}
}
func (self *EventMetaModel) CanReadFull(email string) bool {
	if email == self.OwnerEmail {
		return true
	}
	for _, aEmail := range self.AccessEmails {
		if email == aEmail {
			return true
		}
	}
	if self.GroupModel != nil {
		for _, aEmail := range self.GroupModel.ReadAccessEmails {
			if email == aEmail {
				return true
			}
		}
	}
	return false
}
func (self *EventMetaModel) CanRead(email string) bool {
	if self.IsPublic {
		return true
	}
	return self.CanReadFull(email)
}
func (self *EventMetaModel) GetAttending(
	db storage.Database,
	email string,
) string {
	// returns YES, NO, or MAYBE
	attendingModel, _ := LoadEventAttendingModel(db, self.EventId, email)
	return attendingModel.Attending
}
func (self *EventMetaModel) SetAttending(
	db storage.Database,
	email string,
	attending string,
) error {
	// attending: YES, NO, or MAYBE
	attendingModel, err := LoadEventAttendingModel(db, self.EventId, email)
	if err != nil {
		return err
	}
	attendingModel.Attending = attending
	attendingModel.ModifiedTime = time.Now()
	err = attendingModel.Save(db)
	return err
}
func (self *EventMetaModel) AttendingStatusCount(
	db storage.Database,
	attending string,
) (int, error) {
	return db.FindCount(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
	)
}
func (self *EventMetaModel) Join(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == YES {
		return errors.New("you have already joined this event")
	}
	if !self.CanReadFull(email) {
		if self.IsPublic {
			if !self.PublicJoinOpen {
				return errors.New("this public event is not open for joining")
			}
		} else {
			return errors.New("no access, no join")
		}
	}
	if self.MaxAttendees > 0 {
		attendingCount, err := self.AttendingStatusCount(db, YES)
		if err != nil {
			return err
		}
		if attendingCount >= self.MaxAttendees {
			return errors.New("maximum attendees exceeded, can not join event")
		}
	}
	self.SetAttending(db, email, YES)
	return nil
}
func (self *EventMetaModel) Leave(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == NO {
		if self.CanReadFull(email) {
			return errors.New("you are not attending for this event")
		}
	}
	self.SetAttending(db, email, NO)
	return nil
}
func (self *EventMetaModel) PublicCanJoin() bool {
	return self.IsPublic && self.PublicJoinOpen
}
func (self *EventMetaModel) Invite(
	db storage.Database,
	email string,
	inviteEmails []string,
	remoteIp string,
	host string,
) error {
	var err error
	if len(inviteEmails) == 0 {
		return restpc.NewError(restpc.InvalidArgument, "missing 'inviteEmails'", nil)
	}

	now := time.Now()

	fullAc := self.CanReadFull(email)
	public := self.PublicCanJoin()
	if !(fullAc || public) {
		return restpc.NewError(restpc.PermissionDenied, "not allowed to invite to this event", nil)
	}
	if fullAc {
		accessEmailsMap := make(map[string]bool)
		newAccessEmails := make(
			[]string,
			0,
			len(self.AccessEmails)+len(inviteEmails),
		)
		for _, aEmail := range self.AccessEmails {
			accessEmailsMap[aEmail] = true
			newAccessEmails = append(newAccessEmails, aEmail)
		}
		for _, inviteEmail := range inviteEmails {
			_, found := accessEmailsMap[inviteEmail]
			if !found {
				newAccessEmails = append(newAccessEmails, inviteEmail)
			}
		}
		metaChangeLog := EventMetaChangeLogModel{
			Time:     now,
			Email:    email,
			RemoteIp: remoteIp,
			EventId:  self.EventId,
			FuncName: "SetEventAccess",
			AccessEmails: &[2][]string{
				self.AccessEmails,
				newAccessEmails,
			},
		}
		self.AccessEmails = newAccessEmails
		err = db.Insert(metaChangeLog)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
		err = db.Update(self)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
	}

	eventRev, err := LoadLastRevisionModel(db, &self.EventId)
	if err != nil {
		if db.IsNotFound(err) {
			return restpc.NewError(restpc.NotFound, "event not found", err)
		}
		return restpc.NewError(restpc.Internal, "", err)
	}
	eventModel, err := LoadBaseEventModel(db, eventRev.Sha1)
	if err != nil {
		return restpc.NewError(restpc.Internal, "", err)
	}

	tplText := settings.EVENT_INVITE_EMAIL_TEMPLATE
	user := UserModelByEmail(email, db)
	if user == nil {
		return restpc.NewError(restpc.NotFound, "user not found", nil)
		// FIXME: or Internal?
	}
	for _, inviteEmail := range inviteEmails {
		subject := "Invitation to event: " + eventModel.Summary
		tpl, err := template.New(subject).Parse(tplText)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
		var inviteName string
		inviteUser := UserModelByEmail(inviteEmail, db)
		if inviteUser == nil {
			//return errors.New("invited email not found"), scal.BadRequest
			// FIXME
			inviteName = ""
		} else {
			inviteName = inviteUser.FullName
		}

		tplParams := InviteEmailTemplateParams{
			EventModel:  eventModel,
			SenderEmail: email,
			SenderName:  user.FullName,
			Email:       inviteEmail,
			Name:        inviteName,
			EventType:   self.EventType,
			EventId:     self.EventId.Hex(),
			Host:        host,
		}

		{
			token, _ := newEventInvitationToken(eventModel.Id, inviteEmail)
			tokenEscaped := url.QueryEscape(token) // Go < 1.8
			// tokenEscaped := url.PathEscape(token) // Go 1.8
			tplParams.TokenEscaped = tokenEscaped
		}

		buf := bytes.NewBufferString("")
		err = tpl.Execute(buf, tplParams)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
		emailBody := buf.String()
		db.Insert(EventInvitationModel{
			Time:         now,
			SenderEmail:  email,
			InvitedEmail: inviteEmail,
			EventId:      self.EventId,
		})

		err = scal.SendEmail(
			inviteEmail,
			subject,
			false, // isHtml
			emailBody,
		)
		if err != nil {
			fmt.Println("Failed to send email:", err)
			fmt.Println(emailBody)
		}
	}
	return nil
}
func (self *EventMetaModel) GetEmailsByAttendingStatus(
	db storage.Database,
	attending string,
) []string {
	emailStructs := []struct {
		Email string `bson:"email"`
	}{}
	err := db.FindAll(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
		&emailStructs,
	)
	if err != nil {
		log.Printf(
			"Internal Error: GetAttendingEmails: eventId=%v: %s\n",
			self.EventId,
			err.Error(),
		)
	}
	emails := make([]string, len(emailStructs))
	for i, m := range emailStructs {
		emails[i] = m.Email
	}
	return emails
}
func (self *EventMetaModel) GetAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, YES)
}
func (self *EventMetaModel) GetNotAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, NO)
}
func (self *EventMetaModel) GetMaybeAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, MAYBE)
}

func LoadEventMetaModel(
	db storage.Database,
	eventId *bson.ObjectId,
	loadGroup bool,
) (*EventMetaModel, error) {
	var err error
	eventMeta := EventMetaModel{
		EventId: *eventId,
	}
	err = db.Get(&eventMeta)
	if err != nil {
		return nil, err
	}
	if loadGroup && eventMeta.GroupId != nil {
		// groupModel, err, internalErr
		groupModel, err := LoadGroupModelById(
			"groupId",
			db,
			eventMeta.GroupId,
		)
		if err != nil {
			return nil, err
		}
		eventMeta.GroupModel = groupModel
	}
	return &eventMeta, nil
}

type EventMetaChangeLogModel struct {
	Time     time.Time     `bson:"time"`
	Email    string        `bson:"email"`
	RemoteIp string        `bson:"remoteIp"`
	EventId  bson.ObjectId `bson:"eventId"`
	FuncName string        `bson:"funcName"`

	GroupId        *[2]*bson.ObjectId `bson:"groupId,omitempty"`
	OwnerEmail     *[2]*string        `bson:"ownerEmail,omitempty"`
	IsPublic       *[2]bool           `bson:"isPublic,omitempty"`
	AccessEmails   *[2][]string       `bson:"accessEmails,omitempty"`
	PublicJoinOpen *[2]bool           `bson:"publicJoinOpen,omitempty"`
	MaxAttendees   *[2]int            `bson:"maxAttendees,omitempty"`
}

func (model EventMetaChangeLogModel) Collection() string {
	return storage.C_eventMetaChangeLog
}

type EventInvitationModel struct {
	Time         time.Time     `bson:"time"`
	SenderEmail  string        `bson:"senderEmail"`
	InvitedEmail string        `bson:"invitedEmail"`
	EventId      bson.ObjectId `bson:"eventId"`
}

func (model EventInvitationModel) Collection() string {
	return storage.C_invitation
}

func GetEventMetaPipeResults(
	db storage.Database,
	pipeline *[]scal.M,
) (*[]scal.M, error) {
	results := []scal.M{}
	for res := range db.PipeIter(storage.C_eventMeta, pipeline) {
		if err := res.Err; err != nil {
			return nil, err
		}
		if eventId, ok := res.M["_id"]; ok {
			res.M["eventId"] = eventId
			delete(res.M, "_id")
		}
		if dataI, ok := res.M["data"]; ok {
			data := dataI.(scal.M)
			delete(data, "_id")
			res.M["data"] = data
		}
		results = append(results, res.M)
	}
	return &results, nil
}
<commit_msg>fix non-working event invitation token<commit_after>package event_lib

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"net/url"
	"text/template"
	"time"

	"github.com/ilius/restpc"
	//"net/url"

	"gopkg.in/mgo.v2/bson"

	"scal"
	"scal/settings"
	"scal/storage"
	. "scal/user_lib"
)

/*
// PublicJoinPolicy
const (
    FreeToJoin = "FreeToJoin"
    NoJoin = "NoJoin"
    //JoinNeedsVerify = "JoinNeedsVerify" // makes implementation more complex
)
*/

type EventMetaModel struct {
	EventId   bson.ObjectId `bson:"_id"`
	EventType string        `bson:"eventType"`

	CreationTime time.Time            `bson:"creationTime"`
	FieldsMtime  map[string]time.Time `bson:"fieldsMtime"`

	OwnerEmail   string   `bson:"ownerEmail"`
	IsPublic     bool     `bson:"isPublic"`
	AccessEmails []string `bson:"accessEmails"`

	GroupId    *bson.ObjectId   `bson:"groupId"`
	GroupModel *EventGroupModel `bson:"-"`

	//PublicJoinPolicy string         `bson:"publicJoinPolicy"` // not indexed
	PublicJoinOpen bool `bson:"publicJoinOpen"`
	MaxAttendees   int  `bson:"maxAttendees"`
}

type InviteEmailTemplateParams struct {
	EventModel   *BaseEventModel
	SenderEmail  string
	SenderName   string
	Email        string
	Name         string
	EventType    string
	EventId      string
	TokenEscaped string
	Host         string
}

func (self EventMetaModel) UniqueM() scal.M {
	return scal.M{
		"_id": self.EventId,
	}
}
func (self EventMetaModel) Collection() string {
	return storage.C_eventMeta
}
func (self EventMetaModel) GroupIdHex() string {
	if self.GroupId != nil {
		return self.GroupId.Hex()
	}
	return ""
}
func (self EventMetaModel) JsonM() scal.M {
	return scal.M{
		"creationTime":   self.CreationTime,
		"ownerEmail":     self.OwnerEmail,
		"isPublic":       self.IsPublic,
		"accessEmails":   self.AccessEmails,
		"groupId":        self.GroupIdHex(),
		"publicJoinOpen": self.PublicJoinOpen,
		"maxAttendees":   self.MaxAttendees,
	}
}
func (self *EventMetaModel) CanReadFull(email string) bool {
	if email == self.OwnerEmail {
		return true
	}
	for _, aEmail := range self.AccessEmails {
		if email == aEmail {
			return true
		}
	}
	if self.GroupModel != nil {
		for _, aEmail := range self.GroupModel.ReadAccessEmails {
			if email == aEmail {
				return true
			}
		}
	}
	return false
}
func (self *EventMetaModel) CanRead(email string) bool {
	if self.IsPublic {
		return true
	}
	return self.CanReadFull(email)
}
func (self *EventMetaModel) GetAttending(
	db storage.Database,
	email string,
) string {
	// returns YES, NO, or MAYBE
	attendingModel, _ := LoadEventAttendingModel(db, self.EventId, email)
	return attendingModel.Attending
}
func (self *EventMetaModel) SetAttending(
	db storage.Database,
	email string,
	attending string,
) error {
	// attending: YES, NO, or MAYBE
	attendingModel, err := LoadEventAttendingModel(db, self.EventId, email)
	if err != nil {
		return err
	}
	attendingModel.Attending = attending
	attendingModel.ModifiedTime = time.Now()
	err = attendingModel.Save(db)
	return err
}
func (self *EventMetaModel) AttendingStatusCount(
	db storage.Database,
	attending string,
) (int, error) {
	return db.FindCount(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
	)
}
func (self *EventMetaModel) Join(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == YES {
		return errors.New("you have already joined this event")
	}
	if !self.CanReadFull(email) {
		if self.IsPublic {
			if !self.PublicJoinOpen {
				return errors.New("this public event is not open for joining")
			}
		} else {
			return errors.New("no access, no join")
		}
	}
	if self.MaxAttendees > 0 {
		attendingCount, err := self.AttendingStatusCount(db, YES)
		if err != nil {
			return err
		}
		if attendingCount >= self.MaxAttendees {
			return errors.New("maximum attendees exceeded, can not join event")
		}
	}
	self.SetAttending(db, email, YES)
	return nil
}
func (self *EventMetaModel) Leave(db storage.Database, email string) error {
	// does not make any changes on self
	if self.GetAttending(db, email) == NO {
		if self.CanReadFull(email) {
			return errors.New("you are not attending for this event")
		}
	}
	self.SetAttending(db, email, NO)
	return nil
}
func (self *EventMetaModel) PublicCanJoin() bool {
	return self.IsPublic && self.PublicJoinOpen
}
func (self *EventMetaModel) Invite(
	db storage.Database,
	email string,
	inviteEmails []string,
	remoteIp string,
	host string,
) error {
	var err error
	if len(inviteEmails) == 0 {
		return restpc.NewError(restpc.InvalidArgument, "missing 'inviteEmails'", nil)
	}

	now := time.Now()

	fullAc := self.CanReadFull(email)
	public := self.PublicCanJoin()
	if !(fullAc || public) {
		return restpc.NewError(restpc.PermissionDenied, "not allowed to invite to this event", nil)
	}
	if fullAc {
		accessEmailsMap := make(map[string]bool)
		newAccessEmails := make(
			[]string,
			0,
			len(self.AccessEmails)+len(inviteEmails),
		)
		for _, aEmail := range self.AccessEmails {
			accessEmailsMap[aEmail] = true
			newAccessEmails = append(newAccessEmails, aEmail)
		}
		for _, inviteEmail := range inviteEmails {
			_, found := accessEmailsMap[inviteEmail]
			if !found {
				newAccessEmails = append(newAccessEmails, inviteEmail)
			}
		}
		metaChangeLog := EventMetaChangeLogModel{
			Time:     now,
			Email:    email,
			RemoteIp: remoteIp,
			EventId:  self.EventId,
			FuncName: "SetEventAccess",
			AccessEmails: &[2][]string{
				self.AccessEmails,
				newAccessEmails,
			},
		}
		self.AccessEmails = newAccessEmails
		err = db.Insert(metaChangeLog)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
		err = db.Update(self)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
	}

	eventRev, err := LoadLastRevisionModel(db, &self.EventId)
	if err != nil {
		if db.IsNotFound(err) {
			return restpc.NewError(restpc.NotFound, "event not found", err)
		}
		return restpc.NewError(restpc.Internal, "", err)
	}
	eventModel, err := LoadBaseEventModel(db, eventRev.Sha1)
	if err != nil {
		return restpc.NewError(restpc.Internal, "", err)
	}

	tplText := settings.EVENT_INVITE_EMAIL_TEMPLATE
	user := UserModelByEmail(email, db)
	if user == nil {
		return restpc.NewError(restpc.NotFound, "user not found", nil)
		// FIXME: or Internal?
	}
	for _, inviteEmail := range inviteEmails {
		subject := "Invitation to event: " + eventModel.Summary
		tpl, err := template.New(subject).Parse(tplText)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
		var inviteName string
		inviteUser := UserModelByEmail(inviteEmail, db)
		if inviteUser == nil {
			//return errors.New("invited email not found"), scal.BadRequest
			// FIXME
			inviteName = ""
		} else {
			inviteName = inviteUser.FullName
		}

		tplParams := InviteEmailTemplateParams{
			EventModel:  eventModel,
			SenderEmail: email,
			SenderName:  user.FullName,
			Email:       inviteEmail,
			Name:        inviteName,
			EventType:   self.EventType,
			EventId:     self.EventId.Hex(),
			Host:        host,
		}

		{
			token, _ := newEventInvitationToken(self.EventId, inviteEmail)
			tokenEscaped := url.QueryEscape(token) // Go < 1.8
			// tokenEscaped := url.PathEscape(token) // Go 1.8
			tplParams.TokenEscaped = tokenEscaped
		}

		buf := bytes.NewBufferString("")
		err = tpl.Execute(buf, tplParams)
		if err != nil {
			return restpc.NewError(restpc.Internal, "", err)
		}
		emailBody := buf.String()
		db.Insert(EventInvitationModel{
			Time:         now,
			SenderEmail:  email,
			InvitedEmail: inviteEmail,
			EventId:      self.EventId,
		})

		err = scal.SendEmail(
			inviteEmail,
			subject,
			false, // isHtml
			emailBody,
		)
		if err != nil {
			fmt.Println("Failed to send email:", err)
			fmt.Println(emailBody)
		}
	}
	return nil
}
func (self *EventMetaModel) GetEmailsByAttendingStatus(
	db storage.Database,
	attending string,
) []string {
	emailStructs := []struct {
		Email string `bson:"email"`
	}{}
	err := db.FindAll(
		storage.C_attending,
		scal.M{
			"eventId":   self.EventId,
			"attending": attending,
		},
		&emailStructs,
	)
	if err != nil {
		log.Printf(
			"Internal Error: GetAttendingEmails: eventId=%v: %s\n",
			self.EventId,
			err.Error(),
		)
	}
	emails := make([]string, len(emailStructs))
	for i, m := range emailStructs {
		emails[i] = m.Email
	}
	return emails
}
func (self *EventMetaModel) GetAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, YES)
}
func (self *EventMetaModel) GetNotAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, NO)
}
func (self *EventMetaModel) GetMaybeAttendingEmails(db storage.Database) []string {
	return self.GetEmailsByAttendingStatus(db, MAYBE)
}

func LoadEventMetaModel(
	db storage.Database,
	eventId *bson.ObjectId,
	loadGroup bool,
) (*EventMetaModel, error) {
	var err error
	eventMeta := EventMetaModel{
		EventId: *eventId,
	}
	err = db.Get(&eventMeta)
	if err != nil {
		return nil, err
	}
	if loadGroup && eventMeta.GroupId != nil {
		// groupModel, err, internalErr
		groupModel, err := LoadGroupModelById(
			"groupId",
			db,
			eventMeta.GroupId,
		)
		if err != nil {
			return nil, err
		}
		eventMeta.GroupModel = groupModel
	}
	return &eventMeta, nil
}

type EventMetaChangeLogModel struct {
	Time     time.Time     `bson:"time"`
	Email    string        `bson:"email"`
	RemoteIp string        `bson:"remoteIp"`
	EventId  bson.ObjectId `bson:"eventId"`
	FuncName string        `bson:"funcName"`

	GroupId        *[2]*bson.ObjectId `bson:"groupId,omitempty"`
	OwnerEmail     *[2]*string        `bson:"ownerEmail,omitempty"`
	IsPublic       *[2]bool           `bson:"isPublic,omitempty"`
	AccessEmails   *[2][]string       `bson:"accessEmails,omitempty"`
	PublicJoinOpen *[2]bool           `bson:"publicJoinOpen,omitempty"`
	MaxAttendees   *[2]int            `bson:"maxAttendees,omitempty"`
}

func (model EventMetaChangeLogModel) Collection() string {
	return storage.C_eventMetaChangeLog
}

type EventInvitationModel struct {
	Time         time.Time     `bson:"time"`
	SenderEmail  string        `bson:"senderEmail"`
	InvitedEmail string        `bson:"invitedEmail"`
	EventId      bson.ObjectId `bson:"eventId"`
}

func (model EventInvitationModel) Collection() string {
	return storage.C_invitation
}

func GetEventMetaPipeResults(
	db storage.Database,
	pipeline *[]scal.M,
) (*[]scal.M, error) {
	results := []scal.M{}
	for res := range db.PipeIter(storage.C_eventMeta, pipeline) {
		if err := res.Err; err != nil {
			return nil, err
		}
		if eventId, ok := res.M["_id"]; ok {
			res.M["eventId"] = eventId
			delete(res.M, "_id")
		}
		if dataI, ok := res.M["data"]; ok {
			data := dataI.(scal.M)
			delete(data, "_id")
			res.M["data"] = data
		}
		results = append(results, res.M)
	}
	return &results, nil
}
<|endoftext|>


<commit_before>// Juniper cPTX for KNE
// Copyright (c) Juniper Networks, Inc., 2021. All rights reserved.

package cptx

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

	tpb "github.com/openconfig/kne/proto/topo"
	"github.com/openconfig/kne/topo/node"
	scraplinetwork "github.com/scrapli/scrapligo/driver/network"
	scrapliopts "github.com/scrapli/scrapligo/driver/options"
	scrapliutil "github.com/scrapli/scrapligo/util"
	scraplicfg "github.com/scrapli/scrapligocfg"
	log "github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/pointer"
)

// ErrIncompatibleCliConn raised when an invalid scrapligo cli transport type is found.
var ErrIncompatibleCliConn = errors.New("incompatible cli connection in use")

var (
	// For committing a very large config
	scrapliOperationTimeout = 300 * time.Second
)

const (
	defaultInitContainerImage = "networkop/init-wait:latest"

	scrapliPlatformName = "juniper_junos"
)

func New(nodeImpl *node.Impl) (node.Node, error) {
	if nodeImpl == nil {
		return nil, fmt.Errorf("nodeImpl cannot be nil")
	}
	if nodeImpl.Proto == nil {
		return nil, fmt.Errorf("nodeImpl.Proto cannot be nil")
	}
	cfg := defaults(nodeImpl.Proto)
	nodeImpl.Proto = cfg
	n := &Node{
		Impl: nodeImpl,
	}
	return n, nil
}

type Node struct {
	*node.Impl
	cliConn *scraplinetwork.Driver

	// scrapli options used in testing
	testOpts []scrapliutil.Option
}

// Add validations for interfaces the node provides
var (
	_ node.ConfigPusher = (*Node)(nil)
)

// SpawnCLIConn spawns a CLI connection towards a Network OS using `kubectl exec` terminal and ensures CLI is ready
// to accept inputs.
// scrapligo options can be provided to this function for a caller to modify scrapligo platform.
// For example, mock transport can be set via options
func (n *Node) SpawnCLIConn() error {
	opts := []scrapliutil.Option{
		scrapliopts.WithAuthBypass(),
		scrapliopts.WithTimeoutOps(scrapliOperationTimeout),
	}

	// add options defined in test package
	opts = append(opts, n.testOpts...)

	opts = n.PatchCLIConnOpen("kubectl", []string{"Cli", "-c"}, opts)

	var err error
	n.cliConn, err = n.GetCLIConn(scrapliPlatformName, opts)

	return err
}

func (n *Node) ConfigPush(ctx context.Context, r io.Reader) error {
	log.Infof("%s - pushing config", n.Name())

	cfg, err := io.ReadAll(r)
	cfgs := string(cfg)

	log.Debug(cfgs)

	if err != nil {
		return err
	}

	err = n.SpawnCLIConn()
	if err != nil {
		return err
	}

	defer n.cliConn.Close()

	// use a static candidate file name for test transport
	var candidateConfigFile string
	if len(n.testOpts) != 0 {
		candidateConfigFile = "scrapli_cfg_testing"
	}

	c, err := scraplicfg.NewCfg(
		n.cliConn,
		"juniper_junos",
		scraplicfg.WithCandidateName(candidateConfigFile),
	)
	if err != nil {
		return err
	}

	err = c.Prepare()
	if err != nil {
		return err
	}

	resp, err := c.LoadConfig(
		cfgs,
		true, // load replace
	)
	if err != nil {
		return err
	}
	if resp.Failed != nil {
		return resp.Failed
	}

	resp, err = c.CommitConfig()
	if err != nil {
		return err
	}
	if resp.Failed != nil {
		return resp.Failed
	}

	log.Infof("%s - finished config push", n.Name())

	return nil
}

func (n *Node) ResetCfg(ctx context.Context) error {
	log.Infof("%s - resetting config", n.Name())

	err := n.SpawnCLIConn()
	if err != nil {
		return err
	}

	defer n.cliConn.Close()

	cfgs := []string{
		"load factory-default",
		"delete system commit factory-settings",
		// Plaintext password is 'Google123'
		// Setting a plaintext password would require an interactive prompt to enter the password.
		"set system root-authentication encrypted-password $6$7uA5z8vs$cmHIvL0aLU4ioWAHPR0PLeU/mJj.JO/5pQVQoqRlInK3fJNTLYLhwiDi.Q6gHhltSB3S1P/.raEsuDSH7akcJ/",
		// Without SSH enabled the cli binary returns with the message 'ssh is disabled'
		"set system services ssh root-login allow",
		"commit",
	}
	resp, err := n.cliConn.SendConfigs(cfgs)
	if err != nil {
		return err
	}

	if resp.Failed == nil {
		log.Infof("%s - finshed resetting config", n.Name())
	}

	return resp.Failed
}

func (n *Node) Create(ctx context.Context) error {
	log.Infof("Creating cPTX node resource %s", n.Name())

	if err := n.CreateConfig(ctx); err != nil {
		return fmt.Errorf("node %s failed to create config-map %w", n.Name(), err)
	}
	log.Infof("Created cPTX node %s configmap", n.Name())

	pb := n.Proto
	initContainerImage := pb.Config.InitImage
	if initContainerImage == "" {
		initContainerImage = defaultInitContainerImage
	}

	// downward api - pass some useful values to container
	if n.isChannelized() {
		pb.Config.Env["CPTX_CHANNELIZED"] = "1"
	}
	pb.Config.Env["CPTX_CPU_LIMIT"] = pb.Constraints["cpu"]
	pb.Config.Env["CPTX_MEMORY_LIMIT"] = pb.Constraints["memory"]
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name: n.Name(),
			Labels: map[string]string{
				"app":  n.Name(),
				"topo": n.Namespace,
			},
		},
		Spec: corev1.PodSpec{
			InitContainers: []corev1.Container{{
				Name:  fmt.Sprintf("init-%s", n.Name()),
				Image: initContainerImage,
				Args: []string{
					fmt.Sprintf("%d", len(pb.GetInterfaces())+1),
					fmt.Sprintf("%d", pb.GetConfig().Sleep),
				},
				ImagePullPolicy: "IfNotPresent",
			}},
			Containers: []corev1.Container{{
				Name:            n.Name(),
				Image:           pb.Config.Image,
				Command:         pb.Config.Command,
				Args:            pb.Config.Args,
				Env:             node.ToEnvVar(pb.Config.Env),
				Resources:       node.ToResourceRequirements(pb.Constraints),
				ImagePullPolicy: "IfNotPresent",
				SecurityContext: &corev1.SecurityContext{
					Privileged: pointer.Bool(true),
					RunAsUser:  pointer.Int64(0),
					Capabilities: &corev1.Capabilities{
						Add: []corev1.Capability{"SYS_ADMIN"},
					},
				},
				VolumeMounts: []corev1.VolumeMount{{
					Name:      fmt.Sprintf("%s-run-mount", pb.Name),
					ReadOnly:  false,
					MountPath: "/run",
				}},
			}},
			Volumes: []corev1.Volume{{
				Name: fmt.Sprintf("%s-run-mount", pb.Name),
				VolumeSource: corev1.VolumeSource{
					EmptyDir: &corev1.EmptyDirVolumeSource{
						Medium: "Memory",
					},
				},
			}},
			TerminationGracePeriodSeconds: pointer.Int64(0),
			NodeSelector:                  map[string]string{},
			Affinity: &corev1.Affinity{
				PodAntiAffinity: &corev1.PodAntiAffinity{
					PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{{
						Weight: 100,
						PodAffinityTerm: corev1.PodAffinityTerm{
							LabelSelector: &metav1.LabelSelector{
								MatchExpressions: []metav1.LabelSelectorRequirement{{
									Key:      "topo",
									Operator: "In",
									Values:   []string{pb.Name},
								}},
							},
							TopologyKey: "kubernetes.io/hostname",
						},
					}},
				},
			},
		},
	}
	if pb.Config.ConfigData != nil {
		pod.Spec.Volumes = append(pod.Spec.Volumes, corev1.Volume{
			Name: "startup-config-volume",
			VolumeSource: corev1.VolumeSource{
				ConfigMap: &corev1.ConfigMapVolumeSource{
					LocalObjectReference: corev1.LocalObjectReference{
						Name: fmt.Sprintf("%s-config", pb.Name),
					},
				},
			},
		})
		for i, c := range pod.Spec.Containers {
			pod.Spec.Containers[i].VolumeMounts = append(c.VolumeMounts, corev1.VolumeMount{
				Name:      "startup-config-volume",
				MountPath: pb.Config.ConfigPath + "/" + pb.Config.ConfigFile,
				SubPath:   pb.Config.ConfigFile,
				ReadOnly:  true,
			})
		}
	}
	sPod, err := n.KubeClient.CoreV1().Pods(n.Namespace).Create(ctx, pod, metav1.CreateOptions{})
	if err != nil {
		return fmt.Errorf("failed to create pod for %q: %w", pb.Name, err)
	}
	log.Debugf("Pod created:\n%+v\n", sPod)
	log.Infof("Created cPTX node resource %s pod", n.Name())
	if err := n.CreateService(ctx); err != nil {
		return err
	}
	log.Infof("Created cPTX node resource %s services", n.Name())
	return nil
}

func defaults(pb *tpb.Node) *tpb.Node {
	if pb == nil {
		pb = &tpb.Node{
			Name: "default_cptx_node",
		}
	}
	if pb.Constraints == nil {
		pb.Constraints = map[string]string{
			"cpu":    "8",
			"memory": "8Gi",
		}
	}
	if pb.Services == nil {
		pb.Services = map[uint32]*tpb.Service{
			443: {
				Name:   "ssl",
				Inside: 443,
			},
			22: {
				Name:   "ssh",
				Inside: 22,
			},
			50051: {
				Name:   "gnmi",
				Inside: 50051,
			},
		}
	}
	if pb.Labels == nil {
		pb.Labels = map[string]string{
			"type":   tpb.Node_JUNIPER_CEVO.String(),
			"vendor": tpb.Vendor_JUNIPER.String(),
		}
	} else {
		if pb.Labels["type"] == "" {
			pb.Labels["type"] = tpb.Node_JUNIPER_CEVO.String()
		}
		if pb.Labels["vendor"] == "" {
			pb.Labels["vendor"] = tpb.Vendor_JUNIPER.String()
		}
	}
	if pb.Config == nil {
		pb.Config = &tpb.Config{}
	}
	if len(pb.Config.GetCommand()) == 0 {
		pb.Config.Command = []string{
			"/entrypoint.sh",
		}
	}
	if pb.Config.Image == "" {
		pb.Config.Image = "cptx:latest"
	}
	if pb.Config.Env == nil {
		pb.Config.Env = map[string]string{
			"CPTX": "1",
		}
	}
	if pb.Config.EntryCommand == "" {
		pb.Config.EntryCommand = fmt.Sprintf("kubectl exec -it %s -- cli -c", pb.Name)
	}
	if pb.Config.ConfigPath == "" {
		pb.Config.ConfigPath = "/home/evo/configdisk"
	}
	if pb.Config.ConfigFile == "" {
		pb.Config.ConfigFile = "juniper.conf"
	}
	return pb
}

// isChannelized is a helper function that returns 1 if cptx is channelized
func (n *Node) isChannelized() bool {
	interfaces := n.GetProto().GetInterfaces()
	for key, value := range interfaces {
		if strings.Contains(key, "eth") && strings.Contains(value.Name, ":") {
			return true
		}
	}
	return false
}

func init() {
	node.Register(tpb.Node_JUNIPER_CEVO, New)
	node.Vendor(tpb.Vendor_JUNIPER, New)
}
<commit_msg>add validation for resetter interface for cptx (#191)<commit_after>// Juniper cPTX for KNE
// Copyright (c) Juniper Networks, Inc., 2021. All rights reserved.

package cptx

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

	tpb "github.com/openconfig/kne/proto/topo"
	"github.com/openconfig/kne/topo/node"
	scraplinetwork "github.com/scrapli/scrapligo/driver/network"
	scrapliopts "github.com/scrapli/scrapligo/driver/options"
	scrapliutil "github.com/scrapli/scrapligo/util"
	scraplicfg "github.com/scrapli/scrapligocfg"
	log "github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/pointer"
)

// ErrIncompatibleCliConn raised when an invalid scrapligo cli transport type is found.
var ErrIncompatibleCliConn = errors.New("incompatible cli connection in use")

var (
	// For committing a very large config
	scrapliOperationTimeout = 300 * time.Second
)

const (
	defaultInitContainerImage = "networkop/init-wait:latest"

	scrapliPlatformName = "juniper_junos"
)

func New(nodeImpl *node.Impl) (node.Node, error) {
	if nodeImpl == nil {
		return nil, fmt.Errorf("nodeImpl cannot be nil")
	}
	if nodeImpl.Proto == nil {
		return nil, fmt.Errorf("nodeImpl.Proto cannot be nil")
	}
	cfg := defaults(nodeImpl.Proto)
	nodeImpl.Proto = cfg
	n := &Node{
		Impl: nodeImpl,
	}
	return n, nil
}

type Node struct {
	*node.Impl
	cliConn *scraplinetwork.Driver

	// scrapli options used in testing
	testOpts []scrapliutil.Option
}

// Add validations for interfaces the node provides
var (
	_ node.ConfigPusher = (*Node)(nil)
	_ node.Resetter     = (*Node)(nil)
)

// SpawnCLIConn spawns a CLI connection towards a Network OS using `kubectl exec` terminal and ensures CLI is ready
// to accept inputs.
// scrapligo options can be provided to this function for a caller to modify scrapligo platform.
// For example, mock transport can be set via options
func (n *Node) SpawnCLIConn() error {
	opts := []scrapliutil.Option{
		scrapliopts.WithAuthBypass(),
		scrapliopts.WithTimeoutOps(scrapliOperationTimeout),
	}

	// add options defined in test package
	opts = append(opts, n.testOpts...)

	opts = n.PatchCLIConnOpen("kubectl", []string{"Cli", "-c"}, opts)

	var err error
	n.cliConn, err = n.GetCLIConn(scrapliPlatformName, opts)

	return err
}

func (n *Node) ConfigPush(ctx context.Context, r io.Reader) error {
	log.Infof("%s - pushing config", n.Name())

	cfg, err := io.ReadAll(r)
	cfgs := string(cfg)

	log.Debug(cfgs)

	if err != nil {
		return err
	}

	err = n.SpawnCLIConn()
	if err != nil {
		return err
	}

	defer n.cliConn.Close()

	// use a static candidate file name for test transport
	var candidateConfigFile string
	if len(n.testOpts) != 0 {
		candidateConfigFile = "scrapli_cfg_testing"
	}

	c, err := scraplicfg.NewCfg(
		n.cliConn,
		"juniper_junos",
		scraplicfg.WithCandidateName(candidateConfigFile),
	)
	if err != nil {
		return err
	}

	err = c.Prepare()
	if err != nil {
		return err
	}

	resp, err := c.LoadConfig(
		cfgs,
		true, // load replace
	)
	if err != nil {
		return err
	}
	if resp.Failed != nil {
		return resp.Failed
	}

	resp, err = c.CommitConfig()
	if err != nil {
		return err
	}
	if resp.Failed != nil {
		return resp.Failed
	}

	log.Infof("%s - finished config push", n.Name())

	return nil
}

func (n *Node) ResetCfg(ctx context.Context) error {
	log.Infof("%s - resetting config", n.Name())

	err := n.SpawnCLIConn()
	if err != nil {
		return err
	}

	defer n.cliConn.Close()

	cfgs := []string{
		"load factory-default",
		"delete system commit factory-settings",
		// Plaintext password is 'Google123'
		// Setting a plaintext password would require an interactive prompt to enter the password.
		"set system root-authentication encrypted-password $6$7uA5z8vs$cmHIvL0aLU4ioWAHPR0PLeU/mJj.JO/5pQVQoqRlInK3fJNTLYLhwiDi.Q6gHhltSB3S1P/.raEsuDSH7akcJ/",
		// Without SSH enabled the cli binary returns with the message 'ssh is disabled'
		"set system services ssh root-login allow",
		"commit",
	}
	resp, err := n.cliConn.SendConfigs(cfgs)
	if err != nil {
		return err
	}

	if resp.Failed == nil {
		log.Infof("%s - finshed resetting config", n.Name())
	}

	return resp.Failed
}

func (n *Node) Create(ctx context.Context) error {
	log.Infof("Creating cPTX node resource %s", n.Name())

	if err := n.CreateConfig(ctx); err != nil {
		return fmt.Errorf("node %s failed to create config-map %w", n.Name(), err)
	}
	log.Infof("Created cPTX node %s configmap", n.Name())

	pb := n.Proto
	initContainerImage := pb.Config.InitImage
	if initContainerImage == "" {
		initContainerImage = defaultInitContainerImage
	}

	// downward api - pass some useful values to container
	if n.isChannelized() {
		pb.Config.Env["CPTX_CHANNELIZED"] = "1"
	}
	pb.Config.Env["CPTX_CPU_LIMIT"] = pb.Constraints["cpu"]
	pb.Config.Env["CPTX_MEMORY_LIMIT"] = pb.Constraints["memory"]
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name: n.Name(),
			Labels: map[string]string{
				"app":  n.Name(),
				"topo": n.Namespace,
			},
		},
		Spec: corev1.PodSpec{
			InitContainers: []corev1.Container{{
				Name:  fmt.Sprintf("init-%s", n.Name()),
				Image: initContainerImage,
				Args: []string{
					fmt.Sprintf("%d", len(pb.GetInterfaces())+1),
					fmt.Sprintf("%d", pb.GetConfig().Sleep),
				},
				ImagePullPolicy: "IfNotPresent",
			}},
			Containers: []corev1.Container{{
				Name:            n.Name(),
				Image:           pb.Config.Image,
				Command:         pb.Config.Command,
				Args:            pb.Config.Args,
				Env:             node.ToEnvVar(pb.Config.Env),
				Resources:       node.ToResourceRequirements(pb.Constraints),
				ImagePullPolicy: "IfNotPresent",
				SecurityContext: &corev1.SecurityContext{
					Privileged: pointer.Bool(true),
					RunAsUser:  pointer.Int64(0),
					Capabilities: &corev1.Capabilities{
						Add: []corev1.Capability{"SYS_ADMIN"},
					},
				},
				VolumeMounts: []corev1.VolumeMount{{
					Name:      fmt.Sprintf("%s-run-mount", pb.Name),
					ReadOnly:  false,
					MountPath: "/run",
				}},
			}},
			Volumes: []corev1.Volume{{
				Name: fmt.Sprintf("%s-run-mount", pb.Name),
				VolumeSource: corev1.VolumeSource{
					EmptyDir: &corev1.EmptyDirVolumeSource{
						Medium: "Memory",
					},
				},
			}},
			TerminationGracePeriodSeconds: pointer.Int64(0),
			NodeSelector:                  map[string]string{},
			Affinity: &corev1.Affinity{
				PodAntiAffinity: &corev1.PodAntiAffinity{
					PreferredDuringSchedulingIgnoredDuringExecution: []corev1.WeightedPodAffinityTerm{{
						Weight: 100,
						PodAffinityTerm: corev1.PodAffinityTerm{
							LabelSelector: &metav1.LabelSelector{
								MatchExpressions: []metav1.LabelSelectorRequirement{{
									Key:      "topo",
									Operator: "In",
									Values:   []string{pb.Name},
								}},
							},
							TopologyKey: "kubernetes.io/hostname",
						},
					}},
				},
			},
		},
	}
	if pb.Config.ConfigData != nil {
		pod.Spec.Volumes = append(pod.Spec.Volumes, corev1.Volume{
			Name: "startup-config-volume",
			VolumeSource: corev1.VolumeSource{
				ConfigMap: &corev1.ConfigMapVolumeSource{
					LocalObjectReference: corev1.LocalObjectReference{
						Name: fmt.Sprintf("%s-config", pb.Name),
					},
				},
			},
		})
		for i, c := range pod.Spec.Containers {
			pod.Spec.Containers[i].VolumeMounts = append(c.VolumeMounts, corev1.VolumeMount{
				Name:      "startup-config-volume",
				MountPath: pb.Config.ConfigPath + "/" + pb.Config.ConfigFile,
				SubPath:   pb.Config.ConfigFile,
				ReadOnly:  true,
			})
		}
	}
	sPod, err := n.KubeClient.CoreV1().Pods(n.Namespace).Create(ctx, pod, metav1.CreateOptions{})
	if err != nil {
		return fmt.Errorf("failed to create pod for %q: %w", pb.Name, err)
	}
	log.Debugf("Pod created:\n%+v\n", sPod)
	log.Infof("Created cPTX node resource %s pod", n.Name())
	if err := n.CreateService(ctx); err != nil {
		return err
	}
	log.Infof("Created cPTX node resource %s services", n.Name())
	return nil
}

func defaults(pb *tpb.Node) *tpb.Node {
	if pb == nil {
		pb = &tpb.Node{
			Name: "default_cptx_node",
		}
	}
	if pb.Constraints == nil {
		pb.Constraints = map[string]string{
			"cpu":    "8",
			"memory": "8Gi",
		}
	}
	if pb.Services == nil {
		pb.Services = map[uint32]*tpb.Service{
			443: {
				Name:   "ssl",
				Inside: 443,
			},
			22: {
				Name:   "ssh",
				Inside: 22,
			},
			50051: {
				Name:   "gnmi",
				Inside: 50051,
			},
		}
	}
	if pb.Labels == nil {
		pb.Labels = map[string]string{
			"type":   tpb.Node_JUNIPER_CEVO.String(),
			"vendor": tpb.Vendor_JUNIPER.String(),
		}
	} else {
		if pb.Labels["type"] == "" {
			pb.Labels["type"] = tpb.Node_JUNIPER_CEVO.String()
		}
		if pb.Labels["vendor"] == "" {
			pb.Labels["vendor"] = tpb.Vendor_JUNIPER.String()
		}
	}
	if pb.Config == nil {
		pb.Config = &tpb.Config{}
	}
	if len(pb.Config.GetCommand()) == 0 {
		pb.Config.Command = []string{
			"/entrypoint.sh",
		}
	}
	if pb.Config.Image == "" {
		pb.Config.Image = "cptx:latest"
	}
	if pb.Config.Env == nil {
		pb.Config.Env = map[string]string{
			"CPTX": "1",
		}
	}
	if pb.Config.EntryCommand == "" {
		pb.Config.EntryCommand = fmt.Sprintf("kubectl exec -it %s -- cli -c", pb.Name)
	}
	if pb.Config.ConfigPath == "" {
		pb.Config.ConfigPath = "/home/evo/configdisk"
	}
	if pb.Config.ConfigFile == "" {
		pb.Config.ConfigFile = "juniper.conf"
	}
	return pb
}

// isChannelized is a helper function that returns 1 if cptx is channelized
func (n *Node) isChannelized() bool {
	interfaces := n.GetProto().GetInterfaces()
	for key, value := range interfaces {
		if strings.Contains(key, "eth") && strings.Contains(value.Name, ":") {
			return true
		}
	}
	return false
}

func init() {
	node.Register(tpb.Node_JUNIPER_CEVO, New)
	node.Vendor(tpb.Vendor_JUNIPER, New)
}
<|endoftext|>


<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package main

import (
	"bytes"
	"flag"
	"os"
	"strconv"
	"strings"
	"testing"

	"github.com/google/mtail/metrics"
	"github.com/google/mtail/mtail"
	"github.com/google/mtail/vm"
	"github.com/google/mtail/watcher"
	"github.com/kylelemons/godebug/pretty"
)

var (
	recordBenchmark = flag.Bool("record_benchmark", false, "Record the benchmark results to 'benchmark_results.csv'.")
)

var exampleProgramTests = []struct {
	programfile string // Example program file.
	logfile     string // Sample log input.
	jsonfile    string // Expected metrics after processing.
}{
	{
		"examples/rsyncd.mtail",
		"testdata/rsyncd.log",
		"testdata/rsyncd.json",
	},
	{
		"examples/sftp.mtail",
		"testdata/sftp_chroot.log",
		"testdata/sftp_chroot.json",
	},
	{
		"examples/dhcpd.mtail",
		"testdata/anonymised_dhcpd_log",
		"testdata/anonymised_dhcpd_log.json",
	},
}

func TestExamplePrograms(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode")
	}
	for _, tc := range exampleProgramTests {
		w := watcher.NewFakeWatcher()
		store := &metrics.Store{}
		o := mtail.Options{Progs: tc.programfile, W: w, Store: store}
		mtail, err := mtail.New(o)
		if err != nil {
			t.Fatalf("create mtail failed: %s", err)
		}

		if err := mtail.OneShot(tc.logfile); err != nil {
			t.Errorf("Oneshot failed for %s: %s", tc.logfile, err)
			continue
		}

		// Dirty hack to create json files :)
		if false {
			j, err := os.Create(tc.jsonfile)
			if err != nil {
				t.Errorf("%s: could not open json file: %s", tc.jsonfile, err)
				continue
			}
			defer j.Close()
			if err := mtail.WriteMetrics(j); err != nil {
				t.Errorf("couldn't marshall metrics: %q", err)
				continue
			}
		}

		j, err := os.Open(tc.jsonfile)
		if err != nil {
			t.Fatalf("%s: could not open json file: %s", tc.jsonfile, err)
		}
		defer j.Close()
		var ex bytes.Buffer

		if _, err := ex.ReadFrom(j); err != nil {
			t.Fatalf("Couldn't read from json: %s", err)
		}

		mtail.Close()
		var m bytes.Buffer
		mtail.WriteMetrics(&m)
		m.WriteString("\n") // Golden data has trailing newline.

		diff := pretty.Compare(
			strings.Split(ex.String(), "\n"), // want
			strings.Split(m.String(), "\n"))  // got
		if len(diff) > 0 {
			t.Errorf("%s: metrics don't match:\n%s", tc.programfile, diff)

			t.Errorf("Store metrics: %#v", store.Metrics)
		}
	}
}

func benchmarkProgram(b *testing.B, programfile string, logfile string) {
	w := watcher.NewFakeWatcher()
	o := mtail.Options{Progs: programfile, W: w}
	mtail, err := mtail.New(o)
	if err != nil {
		b.Fatalf("Failed to create mtail: %s", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := mtail.OneShot(logfile); err != nil {
			b.Errorf("OneShot log parse failed: %s", err)
		}
	}
	b.StopTimer()
	mtail.Close()
	l, err := strconv.ParseInt(vm.LineCount.String(), 10, 64)
	if err != nil {
		b.Fatalf("strconv.ParseInt failed: %s", err)
		return
	}
	b.SetBytes(l)
}

func BenchmarkRsyncdProgram(b *testing.B) {
	benchmarkProgram(b, "examples/rsync.mtail", "testdata/rsyncd.log")
}
func BenchmarkSftpProgram(b *testing.B) {
	benchmarkProgram(b, "examples/sftp.mtail", "testdata/sftp_chroot.log")
}
func BenchmarkDhcpdProgram(b *testing.B) {
	benchmarkProgram(b, "examples/dhcpd.mtail", "testdata/anonymised_dhcpd_log")
}
<commit_msg>Fix broken use of defer in ex-test.<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package main

import (
	"bytes"
	"flag"
	"os"
	"strconv"
	"strings"
	"testing"

	"github.com/google/mtail/metrics"
	"github.com/google/mtail/mtail"
	"github.com/google/mtail/vm"
	"github.com/google/mtail/watcher"
	"github.com/kylelemons/godebug/pretty"
)

var (
	recordBenchmark = flag.Bool("record_benchmark", false, "Record the benchmark results to 'benchmark_results.csv'.")
)

var exampleProgramTests = []struct {
	programfile string // Example program file.
	logfile     string // Sample log input.
	jsonfile    string // Expected metrics after processing.
}{
	{
		"examples/rsyncd.mtail",
		"testdata/rsyncd.log",
		"testdata/rsyncd.json",
	},
	{
		"examples/sftp.mtail",
		"testdata/sftp_chroot.log",
		"testdata/sftp_chroot.json",
	},
	{
		"examples/dhcpd.mtail",
		"testdata/anonymised_dhcpd_log",
		"testdata/anonymised_dhcpd_log.json",
	},
}

func TestExamplePrograms(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode")
	}
	for _, tc := range exampleProgramTests {
		w := watcher.NewFakeWatcher()
		store := &metrics.Store{}
		o := mtail.Options{Progs: tc.programfile, W: w, Store: store}
		mtail, err := mtail.New(o)
		if err != nil {
			t.Fatalf("create mtail failed: %s", err)
		}

		if err := mtail.OneShot(tc.logfile); err != nil {
			t.Errorf("Oneshot failed for %s: %s", tc.logfile, err)
			continue
		}

		// Dirty hack to create json files :)
		if false {
			j, err := os.Create(tc.jsonfile)
			if err != nil {
				t.Errorf("%s: could not open json file: %s", tc.jsonfile, err)
				continue
			}
			if err := mtail.WriteMetrics(j); err != nil {
				t.Errorf("couldn't marshall metrics: %q", err)
				continue
			}
			j.Close()
		}

		j, err := os.Open(tc.jsonfile)
		if err != nil {
			t.Fatalf("%s: could not open json file: %s", tc.jsonfile, err)
		}
		defer j.Close()
		var ex bytes.Buffer

		if _, err := ex.ReadFrom(j); err != nil {
			t.Fatalf("Couldn't read from json: %s", err)
		}

		mtail.Close()
		var m bytes.Buffer
		mtail.WriteMetrics(&m)
		m.WriteString("\n") // Golden data has trailing newline.

		diff := pretty.Compare(
			strings.Split(ex.String(), "\n"), // want
			strings.Split(m.String(), "\n"))  // got
		if len(diff) > 0 {
			t.Errorf("%s: metrics don't match:\n%s", tc.programfile, diff)

			t.Errorf("Store metrics: %#v", store.Metrics)
		}
	}
}

func benchmarkProgram(b *testing.B, programfile string, logfile string) {
	w := watcher.NewFakeWatcher()
	o := mtail.Options{Progs: programfile, W: w}
	mtail, err := mtail.New(o)
	if err != nil {
		b.Fatalf("Failed to create mtail: %s", err)
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		if err := mtail.OneShot(logfile); err != nil {
			b.Errorf("OneShot log parse failed: %s", err)
		}
	}
	b.StopTimer()
	mtail.Close()
	l, err := strconv.ParseInt(vm.LineCount.String(), 10, 64)
	if err != nil {
		b.Fatalf("strconv.ParseInt failed: %s", err)
		return
	}
	b.SetBytes(l)
}

func BenchmarkRsyncdProgram(b *testing.B) {
	benchmarkProgram(b, "examples/rsync.mtail", "testdata/rsyncd.log")
}
func BenchmarkSftpProgram(b *testing.B) {
	benchmarkProgram(b, "examples/sftp.mtail", "testdata/sftp_chroot.log")
}
func BenchmarkDhcpdProgram(b *testing.B) {
	benchmarkProgram(b, "examples/dhcpd.mtail", "testdata/anonymised_dhcpd_log")
}
<|endoftext|>