<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
	cloudresourcemanager "github.com/google/google-api-go-client/cloudresourcemanager/v1beta1"
	compute "github.com/google/google-api-go-client/compute/v1"
	"github.com/mitchellh/go-homedir"
	"github.com/olekukonko/tablewriter"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

var Commands = []cli.Command{
	commandCache,
	commandProject,
	commandSsh,
	commandCurrentProject,
}

var commandCache = cli.Command{
	Name:  "cache",
	Usage: "",
	Description: `
`,
	Action: doCache,
}

var commandProject = cli.Command{
	Name:  "project",
	Usage: "",
	Description: `
`,
	Action: doProject,
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name: "zsh-widget, z",
		},
	},
}

var commandSsh = cli.Command{
	Name:  "ssh",
	Usage: "",
	Description: `
`,
	Action: doSsh,
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name: "zsh-widget, z",
		},
	},
}

var commandCurrentProject = cli.Command{
	Name:  "current",
	Usage: "",
	Description: `
`,
	Action: doCurrentProject,
}

// JSON struct
// config
type configCore struct {
	Project string `json:"project"`
	Account string `json:"account"`
}
type configRoot struct {
	Core configCore `json:"core"`
}

// cache
type t_cache struct {
	CacheDir  string
	Instances []*compute.Instance
	Projects  []*cloudresourcemanager.Project
}

func LoadCache() (*t_cache, error) {
	cachedir, err := homedir.Expand("~/.cache/geco/")
	if err != nil {
		return nil, err
	}
	cache := t_cache{CacheDir: cachedir}
	_, err = os.Stat(cache.CacheDir)
	if os.IsNotExist(err) {
		os.MkdirAll(cache.CacheDir, 0700)
	}
	fileinfos, _ := ioutil.ReadDir(cache.CacheDir)
	for _, fileinfo := range fileinfos {
		if fileinfo.IsDir() {
			continue
		}
		if fileinfo.Name() == "projects.json" {
			projects_json, err := ioutil.ReadFile(filepath.Join(cache.CacheDir, fileinfo.Name()))
			if err != nil {
				return nil, err
			}
			err = json.Unmarshal(projects_json, &cache.Projects)
			if err != nil {
				return nil, err
			}
		}
		if fileinfo.Name() == "instances.json" {
			instances_json, err := ioutil.ReadFile(filepath.Join(cache.CacheDir, fileinfo.Name()))
			if err != nil {
				return nil, err
			}
			err = json.Unmarshal(instances_json, &cache.Instances)
			if err != nil {
				return nil, err
			}
		}
	}

	return &cache, nil
}

func SaveCache(cache *t_cache) error {
	projects_json, _ := json.Marshal(cache.Projects)
	err := ioutil.WriteFile(filepath.Join(cache.CacheDir, "projects.json"), projects_json, 0600)
	if err != nil {
		return err
	}
	instances_json, _ := json.Marshal(cache.Instances)
	err = ioutil.WriteFile(filepath.Join(cache.CacheDir, "instances.json"), instances_json, 0600)
	if err != nil {
		return err
	}
	return nil
}

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

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

func doCache(c *cli.Context) {
	cache, err := LoadCache()
	if err != nil {
		panic(err)
	}

	ctx := oauth2.NoContext
	scopes := []string{compute.ComputeReadonlyScope}
	client, err := google.DefaultClient(ctx, scopes...)
	if err != nil {
		panic(err)
	}

	// gcloud projects list
	log.Println("loading projects...")
	service, err := cloudresourcemanager.New(client)
	if err != nil {
		panic(err)
	}

	res, err := service.Projects.List().Do()
	if err != nil {
		panic(err)
	}
	cache.Projects = res.Projects
	log.Printf("loaded projects, %d projects found.\n", len(cache.Projects))

	// gcloud instances list
	for _, project := range res.Projects {
		log.Printf("loading instances in %s (%s)...\n", project.Name, project.ProjectId)
		service, err := compute.New(client)
		if err != nil {
			log.Print(err)
			continue
		}

		res, err := service.Instances.AggregatedList(project.ProjectId).Do()
		if err != nil {
			log.Print(err)
			continue
		}

		num_instances := 0
		for _, instances_scoped_list := range res.Items {
			num_instances = num_instances + len(instances_scoped_list.Instances)
			cache.Instances = append(cache.Instances, instances_scoped_list.Instances...)
		}
		log.Printf("loaded instances in %s (%s), %d instances found.\n", project.Name, project.ProjectId, num_instances)
	}

	SaveCache(cache)
	log.Println("saved cache.")
}

func doProject(c *cli.Context) {
	cache, err := LoadCache()
	if err != nil {
		panic(err)
	}

	buf := renderProjectTable(cache.Projects)
	out := PecoCommand(buf)
	projectLine := strings.Fields(out)
	project_id := projectLine[1]

	cmd := []string{"gcloud", "config", "set", "project", project_id}

	if c.Bool("zsh-widget") {
		fmt.Println(strings.Join(cmd, " "))
	} else {
		log.Println(strings.Join(cmd, " "))
		exec.Command(cmd[0], cmd[1:]...).Output()
	}
}

func doSsh(c *cli.Context) {
	cache, err := LoadCache()
	if err != nil {
		panic(err)
	}

	config := LoadConfig()
	project := config.Core.Project

	buf := renderInstanceTable(project, cache.Instances)
	out := PecoCommand(buf)
	instanceLine := strings.Fields(out)

	if len(instanceLine) == 0 {
		os.Exit(1)
	}

	var instance, zone string
	if project == "" {
		project = instanceLine[1]
		instance = instanceLine[3]
		zone = instanceLine[5]
	} else {
		instance = instanceLine[1]
		zone = instanceLine[3]
	}

	cmd := []string{"gcloud", "compute", "ssh", instance, "--project=" + project, "--zone=" + zone}

	if c.Bool("zsh-widget") {
		fmt.Println(strings.Join(cmd, " "))
	} else {
		log.Println(strings.Join(cmd, " "))
		sshCom := exec.Command(cmd[0], cmd[1:]...)
		sshCom.Stdout = os.Stdout
		sshCom.Stderr = os.Stderr
		sshCom.Stdin = os.Stdin
		err = sshCom.Run()
		if err != nil {
			log.Fatal(err)
		}
	}
}

func doCurrentProject(c *cli.Context) {
	config := LoadConfig()
	project := config.Core.Project
	fmt.Println("project: " + project)
}

func LoadConfig() configRoot {
	out, _ := exec.Command("gcloud", "config", "list", "--format=json").Output()
	buf := string(out)
	decoder := json.NewDecoder(strings.NewReader(buf))
	var d configRoot
	decoder.Decode(&d)
	return d
}

func renderProjectTable(projects []*cloudresourcemanager.Project) []byte {
	var buf bytes.Buffer
	table := tablewriter.NewWriter(&buf)
	table.SetHeader([]string{"PROJECT_ID", "NAME", "PROJECT_NUMBER"})
	for _, p := range projects {
		row := []string{p.ProjectId, p.Name, fmt.Sprintf("%v", p.ProjectNumber)}
		table.Append(row)
	}
	table.Render()

	return buf.Bytes()
}

func renderInstanceTable(project_id string, instances []*compute.Instance) []byte {
	var buf bytes.Buffer
	table := tablewriter.NewWriter(&buf)
	if project_id == "" {
		table.SetHeader([]string{"PROJECT", "NAME", "ZONE", "MACHINE_TYPE", "INTERNAL_IP", "EXTERNAL_IP", "STATUS"})
	} else {
		table.SetHeader([]string{"NAME", "ZONE", "MACHINE_TYPE", "INTERNAL_IP", "EXTERNAL_IP", "STATUS"})
	}
	for _, ins := range instances {
		p := (func(selflink string) string {
			return strings.Split(strings.Split(selflink, "https://www.googleapis.com/compute/v1/projects/")[1], "/")[0]
		})(ins.SelfLink)
		zone := (func(a []string) string { return a[len(a)-1] })(strings.Split(ins.Zone, "/"))
		machine_type := (func(a []string) string { return a[len(a)-1] })(strings.Split(ins.MachineType, "/"))
		internal_ip := ins.NetworkInterfaces[0].NetworkIP
		external_ip := ins.NetworkInterfaces[0].AccessConfigs[0].NatIP
		var row []string
		if project_id == "" {
			row = []string{p, ins.Name, zone, machine_type, internal_ip, external_ip, ins.Status}
			table.Append(row)
		} else {
			if project_id == p {
				row = []string{ins.Name, zone, machine_type, internal_ip, external_ip, ins.Status}
				table.Append(row)
			}
		}
	}
	table.Render()

	return buf.Bytes()
}

func PecoCommand(into []byte) string {
	var buff bytes.Buffer
	pecoCom := exec.Command("peco")
	pecoCom.Stdin = bytes.NewReader(into) // , _ = c.StdoutPipe()
	pecoCom.Stdout = &buff

	_ = pecoCom.Start()
	_ = pecoCom.Wait()

	out := buff.String()
	return out
}
<commit_msg>Clear projects, instances on `geco cache`<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
	cloudresourcemanager "github.com/google/google-api-go-client/cloudresourcemanager/v1beta1"
	compute "github.com/google/google-api-go-client/compute/v1"
	"github.com/mitchellh/go-homedir"
	"github.com/olekukonko/tablewriter"
	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
)

var Commands = []cli.Command{
	commandCache,
	commandProject,
	commandSsh,
	commandCurrentProject,
}

var commandCache = cli.Command{
	Name:  "cache",
	Usage: "",
	Description: `
`,
	Action: doCache,
}

var commandProject = cli.Command{
	Name:  "project",
	Usage: "",
	Description: `
`,
	Action: doProject,
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name: "zsh-widget, z",
		},
	},
}

var commandSsh = cli.Command{
	Name:  "ssh",
	Usage: "",
	Description: `
`,
	Action: doSsh,
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name: "zsh-widget, z",
		},
	},
}

var commandCurrentProject = cli.Command{
	Name:  "current",
	Usage: "",
	Description: `
`,
	Action: doCurrentProject,
}

// JSON struct
// config
type configCore struct {
	Project string `json:"project"`
	Account string `json:"account"`
}
type configRoot struct {
	Core configCore `json:"core"`
}

// cache
type t_cache struct {
	CacheDir  string
	Instances []*compute.Instance
	Projects  []*cloudresourcemanager.Project
}

func LoadCache() (*t_cache, error) {
	cachedir, err := homedir.Expand("~/.cache/geco/")
	if err != nil {
		return nil, err
	}
	cache := t_cache{CacheDir: cachedir}
	_, err = os.Stat(cache.CacheDir)
	if os.IsNotExist(err) {
		os.MkdirAll(cache.CacheDir, 0700)
	}
	fileinfos, _ := ioutil.ReadDir(cache.CacheDir)
	for _, fileinfo := range fileinfos {
		if fileinfo.IsDir() {
			continue
		}
		if fileinfo.Name() == "projects.json" {
			projects_json, err := ioutil.ReadFile(filepath.Join(cache.CacheDir, fileinfo.Name()))
			if err != nil {
				return nil, err
			}
			err = json.Unmarshal(projects_json, &cache.Projects)
			if err != nil {
				return nil, err
			}
		}
		if fileinfo.Name() == "instances.json" {
			instances_json, err := ioutil.ReadFile(filepath.Join(cache.CacheDir, fileinfo.Name()))
			if err != nil {
				return nil, err
			}
			err = json.Unmarshal(instances_json, &cache.Instances)
			if err != nil {
				return nil, err
			}
		}
	}

	return &cache, nil
}

func SaveCache(cache *t_cache) error {
	projects_json, _ := json.Marshal(cache.Projects)
	err := ioutil.WriteFile(filepath.Join(cache.CacheDir, "projects.json"), projects_json, 0600)
	if err != nil {
		return err
	}
	instances_json, _ := json.Marshal(cache.Instances)
	err = ioutil.WriteFile(filepath.Join(cache.CacheDir, "instances.json"), instances_json, 0600)
	if err != nil {
		return err
	}
	return nil
}

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

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

func doCache(c *cli.Context) {
	cache, err := LoadCache()
	if err != nil {
		panic(err)
	}
	cache.Projects = []*cloudresourcemanager.Project{}
	cache.Instances = []*compute.Instance{}

	ctx := oauth2.NoContext
	scopes := []string{compute.ComputeReadonlyScope}
	client, err := google.DefaultClient(ctx, scopes...)
	if err != nil {
		panic(err)
	}

	// gcloud projects list
	log.Println("loading projects...")
	service, err := cloudresourcemanager.New(client)
	if err != nil {
		panic(err)
	}

	res, err := service.Projects.List().Do()
	if err != nil {
		panic(err)
	}
	cache.Projects = res.Projects
	log.Printf("loaded projects, %d projects found.\n", len(cache.Projects))

	// gcloud instances list
	for _, project := range res.Projects {
		log.Printf("loading instances in %s (%s)...\n", project.Name, project.ProjectId)
		service, err := compute.New(client)
		if err != nil {
			log.Print(err)
			continue
		}

		res, err := service.Instances.AggregatedList(project.ProjectId).Do()
		if err != nil {
			log.Print(err)
			continue
		}

		num_instances := 0
		for _, instances_scoped_list := range res.Items {
			num_instances = num_instances + len(instances_scoped_list.Instances)
			cache.Instances = append(cache.Instances, instances_scoped_list.Instances...)
		}
		log.Printf("loaded instances in %s (%s), %d instances found.\n", project.Name, project.ProjectId, num_instances)
	}

	SaveCache(cache)
	log.Println("saved cache.")
}

func doProject(c *cli.Context) {
	cache, err := LoadCache()
	if err != nil {
		panic(err)
	}

	buf := renderProjectTable(cache.Projects)
	out := PecoCommand(buf)
	projectLine := strings.Fields(out)
	project_id := projectLine[1]

	cmd := []string{"gcloud", "config", "set", "project", project_id}

	if c.Bool("zsh-widget") {
		fmt.Println(strings.Join(cmd, " "))
	} else {
		log.Println(strings.Join(cmd, " "))
		exec.Command(cmd[0], cmd[1:]...).Output()
	}
}

func doSsh(c *cli.Context) {
	cache, err := LoadCache()
	if err != nil {
		panic(err)
	}

	config := LoadConfig()
	project := config.Core.Project

	buf := renderInstanceTable(project, cache.Instances)
	out := PecoCommand(buf)
	instanceLine := strings.Fields(out)

	if len(instanceLine) == 0 {
		os.Exit(1)
	}

	var instance, zone string
	if project == "" {
		project = instanceLine[1]
		instance = instanceLine[3]
		zone = instanceLine[5]
	} else {
		instance = instanceLine[1]
		zone = instanceLine[3]
	}

	cmd := []string{"gcloud", "compute", "ssh", instance, "--project=" + project, "--zone=" + zone}

	if c.Bool("zsh-widget") {
		fmt.Println(strings.Join(cmd, " "))
	} else {
		log.Println(strings.Join(cmd, " "))
		sshCom := exec.Command(cmd[0], cmd[1:]...)
		sshCom.Stdout = os.Stdout
		sshCom.Stderr = os.Stderr
		sshCom.Stdin = os.Stdin
		err = sshCom.Run()
		if err != nil {
			log.Fatal(err)
		}
	}
}

func doCurrentProject(c *cli.Context) {
	config := LoadConfig()
	project := config.Core.Project
	fmt.Println("project: " + project)
}

func LoadConfig() configRoot {
	out, _ := exec.Command("gcloud", "config", "list", "--format=json").Output()
	buf := string(out)
	decoder := json.NewDecoder(strings.NewReader(buf))
	var d configRoot
	decoder.Decode(&d)
	return d
}

func renderProjectTable(projects []*cloudresourcemanager.Project) []byte {
	var buf bytes.Buffer
	table := tablewriter.NewWriter(&buf)
	table.SetHeader([]string{"PROJECT_ID", "NAME", "PROJECT_NUMBER"})
	for _, p := range projects {
		row := []string{p.ProjectId, p.Name, fmt.Sprintf("%v", p.ProjectNumber)}
		table.Append(row)
	}
	table.Render()

	return buf.Bytes()
}

func renderInstanceTable(project_id string, instances []*compute.Instance) []byte {
	var buf bytes.Buffer
	table := tablewriter.NewWriter(&buf)
	if project_id == "" {
		table.SetHeader([]string{"PROJECT", "NAME", "ZONE", "MACHINE_TYPE", "INTERNAL_IP", "EXTERNAL_IP", "STATUS"})
	} else {
		table.SetHeader([]string{"NAME", "ZONE", "MACHINE_TYPE", "INTERNAL_IP", "EXTERNAL_IP", "STATUS"})
	}
	for _, ins := range instances {
		p := (func(selflink string) string {
			return strings.Split(strings.Split(selflink, "https://www.googleapis.com/compute/v1/projects/")[1], "/")[0]
		})(ins.SelfLink)
		zone := (func(a []string) string { return a[len(a)-1] })(strings.Split(ins.Zone, "/"))
		machine_type := (func(a []string) string { return a[len(a)-1] })(strings.Split(ins.MachineType, "/"))
		internal_ip := ins.NetworkInterfaces[0].NetworkIP
		external_ip := ins.NetworkInterfaces[0].AccessConfigs[0].NatIP
		var row []string
		if project_id == "" {
			row = []string{p, ins.Name, zone, machine_type, internal_ip, external_ip, ins.Status}
			table.Append(row)
		} else {
			if project_id == p {
				row = []string{ins.Name, zone, machine_type, internal_ip, external_ip, ins.Status}
				table.Append(row)
			}
		}
	}
	table.Render()

	return buf.Bytes()
}

func PecoCommand(into []byte) string {
	var buff bytes.Buffer
	pecoCom := exec.Command("peco")
	pecoCom.Stdin = bytes.NewReader(into) // , _ = c.StdoutPipe()
	pecoCom.Stdout = &buff

	_ = pecoCom.Start()
	_ = pecoCom.Wait()

	out := buff.String()
	return out
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/bwmarrin/discordgo"
	"io"
	"os"
	"os/exec"
)

// generic command struct which contains name, description, and a function
type command struct {
	name             string                                                                                // human-readable name of the command
	description      string                                                                                // description of command's function
	usage            string                                                                                // example of how to correctly use command - [] for optional arguments, <> for required arguments
	verbs            []string                                                                                // all verbs which are mapped to the same command
	requiresDatabase bool                                                                                // does this command require database access?
	function         func([]string, *discordgo.Channel, *discordgo.MessageCreate, *discordgo.Session) *commandOutput // function which receives a slice of arguments and returns a string to display to the user
}

// output returned by all command functions, can contain a file to be uploaded
type commandOutput struct {
	response string
	file     io.Reader
	embed    *discordgo.MessageEmbed
}

func initCommands() map[string]*command {
	commandList := []*command{}

	commandList = append(commandList,

		// Define all commands here in the order they will be displayed by the help command
		// The 'usage' field should use the default verb
		// Do not include the command prefix

		&command{
			name:             "Display help",
			description:      "Lists all commands and their purposes.\nCan also display detailed info about a given command.",
			usage:            "help [verb]",
			verbs:            []string{"help", "commands"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				DebugPrint("Running help command.")

				if len(args) <= 0 {

					DebugPrint("No arguments; listing commands.")

					embed := NewEmbed().
						SetTitle("Source").
						SetAuthor("Sunbot " + version).
						//SetDescription("Database enabled: " + strconv.FormatBool(redisEnabled)).
						SetURL("https://github.com/techniponi/sunbot").
						SetImage(discordSession.State.User.AvatarURL("128"))

					for _, cmd := range commandList {
						/*
						if cmd.requiresDatabase && !redisEnabled {
							// Database is not enabled, this command needs it
						} else {
						*/
							embed.AddField(cmd.name, "`"+cfg.DefaultPrefix+cmd.usage+"`")
						//}
					}

					return &commandOutput{embed: embed.MessageEmbed}
				}

				DebugPrint("Verb was given...")

				// check if command exists
				if cmd, ok := commands[args[0]]; ok {

					embed := NewEmbed().
						SetTitle(cmd.name).
						SetDescription(cmd.description).
						AddField("Usage", "`"+cfg.DefaultPrefix+cmd.usage+"`")

					DebugPrint("Providing help for given verb.")

					// compile verbs
					verbOutput := ""
					for index, verb := range cmd.verbs {
						// don't add a comma if it's the last one
						if index == (len(cmd.verbs) - 1) {
							verbOutput += "`" + cfg.DefaultPrefix + verb + "`"
						} else {
							verbOutput += "`" + cfg.DefaultPrefix + verb + "`, "
						}
					}
					embed.AddField("Verbs", verbOutput)
					return &commandOutput{embed: embed.MessageEmbed}
				}
				DebugPrint("Given verb was not found.")
				return &commandOutput{response: "That isn't a valid command."}
			},
		},

		&command{
			name:             "Derpibooru search",
			description:      "Searches Derpibooru with the given tags as the query, chooses a random result to display.\nUse commas to separate tags like you would on the website.",
			usage:            "derpi <tags>",
			verbs:            []string{"derpi", "db", "derpibooru"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {
				if len(args) < 1 {
					DebugPrint("User ran derpibooru command with no tags given.")
					return &commandOutput{response: "Error: no tags specified"}
				}
				DebugPrint("User is running derpibooru command...")

				searchQuery := ""

				for _, arg := range args {
					searchQuery += arg + " "
				}

				// enforce 'safe' tag if channel is not nsfw
				if !channel.NSFW {
					DebugPrint("Channel #" + channel.Name + " is SFW, adding safe tag...")
					searchQuery += ",safe"
				}

				DebugPrint("Searching with tags:\n" + searchQuery)

				// use derpibooru.go to perform search
				results, err := DerpiSearchWithTags(searchQuery, cfg.DerpiApiKey)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error: " + err.Error()}
				}

				// check for results
				if len(results.Search) <= 0 {
					DebugPrint("Derpibooru returned no results.")
					return &commandOutput{response: "Error: no results."}
				}
				DebugPrint("Derpibooru returned results; parsed successfully.")
				// pick one randomly
				output := "http:" + results.Search[RandomRange(0, len(results.Search))].Image

				return &commandOutput{response: output}
			},
		},

		&command{
			name: "Exec",
			description: "Execute a shell command on my server.",
			usage: "exec <command>",
			verbs: []string{"exec"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				hasPermission := false

				// get user object
				DebugPrint("Getting user object")
				user, err := discordSession.State.Member(channel.GuildID, msgEvent.Author.ID)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error getting user from exec command"}
				}

				// get roles from that user
				DebugPrint("Getting user's roles")
				for _, roleID := range user.Roles {
					role, err := discordSession.State.Role(channel.GuildID, roleID)
					if err != nil {
						fmt.Println(err)
						return &commandOutput{response: "Error getting roles from user"}
					}

					DebugPrint("Checking for admin permission")
					if role.Permissions&discordgo.PermissionAdministrator == 0 {
						hasPermission = true
					}
				}

				if hasPermission {

					// convert slice to single string
					fullCommand := ""
					for _, arg := range args {
						fullCommand += arg + " "
					}

					cmd := exec.Command("/bin/bash", "-c", fullCommand)
					stdout, err := cmd.Output()

					if err != nil {
						fmt.Println(err)
						return &commandOutput{response: "Error running command"}
					}

					return &commandOutput{
						response: "```sh\n" + string(stdout) + "\n```",
					}
				}else{
					return &commandOutput{response: "Sorry, but only administrators can use that command."}
				}
			},
		},

		&command{
			name:             "Join",
			description:      "I will join the voice channel of the sender.",
			usage:            "join",
			verbs:            []string{"join"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				voiceConnection, err := JoinUserVoiceChannel(discordSession, msgEvent.Author.ID)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error joining voice channel - are you in one in this server?"}
				}
				fmt.Println("Joined channel:")
				fmt.Println(voiceConnection)

				// TODO: fix audio
/*
				DebugPrint("Setting options")
				options := dca.StdEncodeOptions
				options.RawOutput = true
				options.Bitrate = 96
				options.Application = "lowdelay"

				DebugPrint("Getting info")
				videoInfo, err := ytdl.GetVideoInfo("https://www.youtube.com/watch?v=gWBZJkfzhNY")
				if err != nil {
					DebugPrint("Error getting info")
					fmt.Println(err)
				}

				DebugPrint("Setting format")
				format := videoInfo.Formats.Extremes(ytdl.FormatAudioBitrateKey, true)[0]
				downloadURL, err := videoInfo.GetDownloadURL(format)
				if err != nil {
					DebugPrint("Error setting format")
					fmt.Println(err)
				}

				DebugPrint("Encoding file")
				encodingSession, err := dca.EncodeFile(downloadURL.String(), options)
				if err != nil {
					DebugPrint("Error encoding file")
					fmt.Println(err)
				}
				defer encodingSession.Cleanup()

				DebugPrint("Initiating stream")
				done := make(chan error)
				dca.NewStream(encodingSession, voiceConnection, done)
				playErr := <- done
				if playErr != nil && playErr != io.EOF {
					fmt.Println(playErr)
				}
*/
				return &commandOutput{response: "Joining your channel!"}
			},
		},

		&command{
			name:             "Leave",
			description:      "I will leave the voice channel.",
			usage:            "leave",
			verbs:            []string{"leave", "disconnect", "quit", "exit"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				for _, voiceChannel := range discordSession.VoiceConnections {
					DebugPrint("Looking for voice channel in this guild...")
					if voiceChannel.GuildID == channel.GuildID {
						DebugPrint("Found a channel in this guild!")
						voiceChannel.Disconnect()
						return &commandOutput{response: "Bye!"}
					}
				}
				return &commandOutput{response: ""}
			},
		},

		&command{
			name:             "Gay",
			description:      "Posts a very gay image.",
			usage:            "gay",
			verbs:            []string{"gay"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {
				file, err := os.Open("img/gaybats.png") // TODO: move this to database; allow users to add images (permission system?)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error opening file"}

				}
				return &commandOutput{file: file}
			},
		},

	/*
		&command{
			name:             "User stats",
			description:      "Displays the statistics of the user.",
			usage:            "stats [user]", // TODO: implement pinging users
			verbs:            []string{"stats"},
			requiresDatabase: true,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				if len(args) > 0 {
					if len(msgEvent.Mentions) > 0 {
						// User tagged someone else
						taggedUser := msgEvent.Mentions[0] // only the first one

						userDb, err := GetUser(taggedUser, false)
						if err != nil {
							return &commandOutput{response: "That user doesn't exist in the database yet. They need to chat some!"}
						}

						posts := userDb.Val()["posts"]
						return &commandOutput{response: taggedUser.Username + " has made " + posts + " posts!"} // TODO: format as embed, show more values
					}
					// user didn't tag anyone
					// TODO: accept aliases as well as mentions
					return &commandOutput{response: "To see someone's stats, tag the person directly!"}
				}
				// User's own stats
				userDb, err := GetUser(msgEvent.Author, false)
				if err != nil {
					return &commandOutput{response: "You don't exist in the database yet. You need to chat some!"}
				}
				posts := userDb.Val()["posts"]
				return &commandOutput{response: "You have made " + posts + " posts!"} // TODO: format as embed, show more values
			},
		},
	*/
	)

	// Map for matching verbs to commands
	commandMap := make(map[string]*command)

	// Loop through commandList to get each verb
	for _, cmd := range commandList {
		for _, verb := range cmd.verbs {
			commandMap[verb] = cmd
			DebugPrint("Mapped '" + verb + "' to '" + cmd.name + "'")
		}
	}

	return commandMap
}
<commit_msg>fix admin permission check<commit_after>package main

import (
	"fmt"
	"github.com/bwmarrin/discordgo"
	"io"
	"os"
	"os/exec"
)

// generic command struct which contains name, description, and a function
type command struct {
	name             string                                                                                // human-readable name of the command
	description      string                                                                                // description of command's function
	usage            string                                                                                // example of how to correctly use command - [] for optional arguments, <> for required arguments
	verbs            []string                                                                                // all verbs which are mapped to the same command
	requiresDatabase bool                                                                                // does this command require database access?
	function         func([]string, *discordgo.Channel, *discordgo.MessageCreate, *discordgo.Session) *commandOutput // function which receives a slice of arguments and returns a string to display to the user
}

// output returned by all command functions, can contain a file to be uploaded
type commandOutput struct {
	response string
	file     io.Reader
	embed    *discordgo.MessageEmbed
}

func initCommands() map[string]*command {
	commandList := []*command{}

	commandList = append(commandList,

		// Define all commands here in the order they will be displayed by the help command
		// The 'usage' field should use the default verb
		// Do not include the command prefix

		&command{
			name:             "Display help",
			description:      "Lists all commands and their purposes.\nCan also display detailed info about a given command.",
			usage:            "help [verb]",
			verbs:            []string{"help", "commands"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				DebugPrint("Running help command.")

				if len(args) <= 0 {

					DebugPrint("No arguments; listing commands.")

					embed := NewEmbed().
						SetTitle("Source").
						SetAuthor("Sunbot " + version).
						//SetDescription("Database enabled: " + strconv.FormatBool(redisEnabled)).
						SetURL("https://github.com/techniponi/sunbot").
						SetImage(discordSession.State.User.AvatarURL("128"))

					for _, cmd := range commandList {
						/*
						if cmd.requiresDatabase && !redisEnabled {
							// Database is not enabled, this command needs it
						} else {
						*/
							embed.AddField(cmd.name, "`"+cfg.DefaultPrefix+cmd.usage+"`")
						//}
					}

					return &commandOutput{embed: embed.MessageEmbed}
				}

				DebugPrint("Verb was given...")

				// check if command exists
				if cmd, ok := commands[args[0]]; ok {

					embed := NewEmbed().
						SetTitle(cmd.name).
						SetDescription(cmd.description).
						AddField("Usage", "`"+cfg.DefaultPrefix+cmd.usage+"`")

					DebugPrint("Providing help for given verb.")

					// compile verbs
					verbOutput := ""
					for index, verb := range cmd.verbs {
						// don't add a comma if it's the last one
						if index == (len(cmd.verbs) - 1) {
							verbOutput += "`" + cfg.DefaultPrefix + verb + "`"
						} else {
							verbOutput += "`" + cfg.DefaultPrefix + verb + "`, "
						}
					}
					embed.AddField("Verbs", verbOutput)
					return &commandOutput{embed: embed.MessageEmbed}
				}
				DebugPrint("Given verb was not found.")
				return &commandOutput{response: "That isn't a valid command."}
			},
		},

		&command{
			name:             "Derpibooru search",
			description:      "Searches Derpibooru with the given tags as the query, chooses a random result to display.\nUse commas to separate tags like you would on the website.",
			usage:            "derpi <tags>",
			verbs:            []string{"derpi", "db", "derpibooru"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {
				if len(args) < 1 {
					DebugPrint("User ran derpibooru command with no tags given.")
					return &commandOutput{response: "Error: no tags specified"}
				}
				DebugPrint("User is running derpibooru command...")

				searchQuery := ""

				for _, arg := range args {
					searchQuery += arg + " "
				}

				// enforce 'safe' tag if channel is not nsfw
				if !channel.NSFW {
					DebugPrint("Channel #" + channel.Name + " is SFW, adding safe tag...")
					searchQuery += ",safe"
				}

				DebugPrint("Searching with tags:\n" + searchQuery)

				// use derpibooru.go to perform search
				results, err := DerpiSearchWithTags(searchQuery, cfg.DerpiApiKey)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error: " + err.Error()}
				}

				// check for results
				if len(results.Search) <= 0 {
					DebugPrint("Derpibooru returned no results.")
					return &commandOutput{response: "Error: no results."}
				}
				DebugPrint("Derpibooru returned results; parsed successfully.")
				// pick one randomly
				output := "http:" + results.Search[RandomRange(0, len(results.Search))].Image

				return &commandOutput{response: output}
			},
		},

		&command{
			name: "Exec",
			description: "Execute a shell command on my server.",
			usage: "exec <command>",
			verbs: []string{"exec"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				hasPermission := false

				// get user object
				DebugPrint("Getting user object")
				user, err := discordSession.State.Member(channel.GuildID, msgEvent.Author.ID)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error getting user from exec command"}
				}

				// get roles from that user
				DebugPrint("Getting user's roles")
				for _, roleID := range user.Roles {
					role, err := discordSession.State.Role(channel.GuildID, roleID)
					if err != nil {
						fmt.Println(err)
						return &commandOutput{response: "Error getting roles from user"}
					}

					DebugPrint("Checking for admin permission")
					if role.Permissions&discordgo.PermissionAdministrator != 0 {
						hasPermission = true
					}
				}

				if hasPermission {

					// convert slice to single string
					fullCommand := ""
					for _, arg := range args {
						fullCommand += arg + " "
					}

					cmd := exec.Command("/bin/bash", "-c", fullCommand)
					stdout, err := cmd.Output()

					if err != nil {
						fmt.Println(err)
						return &commandOutput{response: "Error running command"}
					}

					return &commandOutput{
						response: "```sh\n" + string(stdout) + "\n```",
					}
				}else{
					return &commandOutput{response: "Sorry, but only administrators can use that command."}
				}
			},
		},

		&command{
			name:             "Join",
			description:      "I will join the voice channel of the sender.",
			usage:            "join",
			verbs:            []string{"join"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				voiceConnection, err := JoinUserVoiceChannel(discordSession, msgEvent.Author.ID)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error joining voice channel - are you in one in this server?"}
				}
				fmt.Println("Joined channel:")
				fmt.Println(voiceConnection)

				// TODO: fix audio
/*
				DebugPrint("Setting options")
				options := dca.StdEncodeOptions
				options.RawOutput = true
				options.Bitrate = 96
				options.Application = "lowdelay"

				DebugPrint("Getting info")
				videoInfo, err := ytdl.GetVideoInfo("https://www.youtube.com/watch?v=gWBZJkfzhNY")
				if err != nil {
					DebugPrint("Error getting info")
					fmt.Println(err)
				}

				DebugPrint("Setting format")
				format := videoInfo.Formats.Extremes(ytdl.FormatAudioBitrateKey, true)[0]
				downloadURL, err := videoInfo.GetDownloadURL(format)
				if err != nil {
					DebugPrint("Error setting format")
					fmt.Println(err)
				}

				DebugPrint("Encoding file")
				encodingSession, err := dca.EncodeFile(downloadURL.String(), options)
				if err != nil {
					DebugPrint("Error encoding file")
					fmt.Println(err)
				}
				defer encodingSession.Cleanup()

				DebugPrint("Initiating stream")
				done := make(chan error)
				dca.NewStream(encodingSession, voiceConnection, done)
				playErr := <- done
				if playErr != nil && playErr != io.EOF {
					fmt.Println(playErr)
				}
*/
				return &commandOutput{response: "Joining your channel!"}
			},
		},

		&command{
			name:             "Leave",
			description:      "I will leave the voice channel.",
			usage:            "leave",
			verbs:            []string{"leave", "disconnect", "quit", "exit"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				for _, voiceChannel := range discordSession.VoiceConnections {
					DebugPrint("Looking for voice channel in this guild...")
					if voiceChannel.GuildID == channel.GuildID {
						DebugPrint("Found a channel in this guild!")
						voiceChannel.Disconnect()
						return &commandOutput{response: "Bye!"}
					}
				}
				return &commandOutput{response: ""}
			},
		},

		&command{
			name:             "Gay",
			description:      "Posts a very gay image.",
			usage:            "gay",
			verbs:            []string{"gay"},
			requiresDatabase: false,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {
				file, err := os.Open("img/gaybats.png") // TODO: move this to database; allow users to add images (permission system?)
				if err != nil {
					fmt.Println(err)
					return &commandOutput{response: "Error opening file"}

				}
				return &commandOutput{file: file}
			},
		},

	/*
		&command{
			name:             "User stats",
			description:      "Displays the statistics of the user.",
			usage:            "stats [user]", // TODO: implement pinging users
			verbs:            []string{"stats"},
			requiresDatabase: true,
			function: func(args []string, channel *discordgo.Channel, msgEvent *discordgo.MessageCreate, discordSession *discordgo.Session) *commandOutput {

				if len(args) > 0 {
					if len(msgEvent.Mentions) > 0 {
						// User tagged someone else
						taggedUser := msgEvent.Mentions[0] // only the first one

						userDb, err := GetUser(taggedUser, false)
						if err != nil {
							return &commandOutput{response: "That user doesn't exist in the database yet. They need to chat some!"}
						}

						posts := userDb.Val()["posts"]
						return &commandOutput{response: taggedUser.Username + " has made " + posts + " posts!"} // TODO: format as embed, show more values
					}
					// user didn't tag anyone
					// TODO: accept aliases as well as mentions
					return &commandOutput{response: "To see someone's stats, tag the person directly!"}
				}
				// User's own stats
				userDb, err := GetUser(msgEvent.Author, false)
				if err != nil {
					return &commandOutput{response: "You don't exist in the database yet. You need to chat some!"}
				}
				posts := userDb.Val()["posts"]
				return &commandOutput{response: "You have made " + posts + " posts!"} // TODO: format as embed, show more values
			},
		},
	*/
	)

	// Map for matching verbs to commands
	commandMap := make(map[string]*command)

	// Loop through commandList to get each verb
	for _, cmd := range commandList {
		for _, verb := range cmd.verbs {
			commandMap[verb] = cmd
			DebugPrint("Mapped '" + verb + "' to '" + cmd.name + "'")
		}
	}

	return commandMap
}
<|endoftext|>


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

package RuntimeTest

import (
	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
)

// ExpectPolicyEnforcementUpdated sets PolicyEnforcement mode on the provided vm
// running Cilium to policyEnforcementType, and waits for endpoints on vm to
// be in 'ready' state. It asserts whether the configuration update was successful,
// and if the endpoints are ready.
func ExpectPolicyEnforcementUpdated(vm *helpers.SSHMeta, policyEnforcementType string) {
	By("Setting PolicyEnforcement=%s", policyEnforcementType)
	areEndpointsReady := vm.SetPolicyEnforcementAndWait(policyEnforcementType)
	ExpectWithOffset(1, areEndpointsReady).Should(BeTrue(), "unable to set PolicyEnforcement=%s", policyEnforcementType)
}

// ExpectEndpointSummary asserts whether the amount of endpoints managed by
// Cilium running on vm with PolicyEnforcement equal to policyEnforcementType
// is equal to numWithPolicyEnforcementType.
func ExpectEndpointSummary(vm *helpers.SSHMeta, policyEnforcementType string, numWithPolicyEnforcementType int) {
	endpoints, err := vm.PolicyEndpointsSummary()
	ExpectWithOffset(1, err).Should(BeNil(), "error getting endpoint summary")
	ExpectWithOffset(1, endpoints[policyEnforcementType]).To(Equal(numWithPolicyEnforcementType), "number of endpoints with %s=%s does not match", helpers.PolicyEnforcement, policyEnforcementType)
}

// ExpectCiliumReady asserts that cilium status is ready
func ExpectCiliumReady(vm *helpers.SSHMeta) {
	err := vm.WaitUntilReady(helpers.CiliumStartTimeout)
	ExpectWithOffset(1, err).To(BeNil(), "Cilium-agent cannot be started")

	vm.NetworkCreate(helpers.CiliumDockerNetwork, "")
	res := vm.NetworkGet(helpers.CiliumDockerNetwork)
	ExpectWithOffset(1, res).To(helpers.CMDSuccess(), "Cilium docker network is not created")

	res = vm.SetPolicyEnforcement(helpers.PolicyEnforcementDefault)
	ExpectWithOffset(1, res).To(helpers.CMDSuccess(), "Cannot set policy enforcement default")
}

// ExpectCiliumNotRunning checks that cilium has stopped with a grace period.
func ExpectCiliumNotRunning(vm *helpers.SSHMeta) {
	body := func() bool {
		res := vm.Exec("systemctl is-active cilium")
		return !res.WasSuccessful()
	}
	err := helpers.WithTimeout(body, "Cilium is still running after timeout", &helpers.TimeoutConfig{Timeout: helpers.CiliumStartTimeout})
	ExpectWithOffset(1, err).To(BeNil(), "Cilium is still running after timeout")
}

// ExpectDockerContainersMatchCiliumEndpoints asserts that docker containers in
// Cilium network match with the endpoint list
func ExpectDockerContainersMatchCiliumEndpoints(vm *helpers.SSHMeta) {
	ExpectWithOffset(1, vm.ValidateEndpointsAreCorrect(helpers.CiliumDockerNetwork)).To(BeNil(),
		"Docker containers mistmach with Cilium Endpoints")
}
<commit_msg>test: Fix flake in ValidateEndpointsAreCorrect<commit_after>// Copyright 2018-2019 Authors of Cilium
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package RuntimeTest

import (
	. "github.com/cilium/cilium/test/ginkgo-ext"
	"github.com/cilium/cilium/test/helpers"

	. "github.com/onsi/gomega"
)

// ExpectPolicyEnforcementUpdated sets PolicyEnforcement mode on the provided vm
// running Cilium to policyEnforcementType, and waits for endpoints on vm to
// be in 'ready' state. It asserts whether the configuration update was successful,
// and if the endpoints are ready.
func ExpectPolicyEnforcementUpdated(vm *helpers.SSHMeta, policyEnforcementType string) {
	By("Setting PolicyEnforcement=%s", policyEnforcementType)
	areEndpointsReady := vm.SetPolicyEnforcementAndWait(policyEnforcementType)
	ExpectWithOffset(1, areEndpointsReady).Should(BeTrue(), "unable to set PolicyEnforcement=%s", policyEnforcementType)
}

// ExpectEndpointSummary asserts whether the amount of endpoints managed by
// Cilium running on vm with PolicyEnforcement equal to policyEnforcementType
// is equal to numWithPolicyEnforcementType.
func ExpectEndpointSummary(vm *helpers.SSHMeta, policyEnforcementType string, numWithPolicyEnforcementType int) {
	endpoints, err := vm.PolicyEndpointsSummary()
	ExpectWithOffset(1, err).Should(BeNil(), "error getting endpoint summary")
	ExpectWithOffset(1, endpoints[policyEnforcementType]).To(Equal(numWithPolicyEnforcementType), "number of endpoints with %s=%s does not match", helpers.PolicyEnforcement, policyEnforcementType)
}

// ExpectCiliumReady asserts that cilium status is ready
func ExpectCiliumReady(vm *helpers.SSHMeta) {
	err := vm.WaitUntilReady(helpers.CiliumStartTimeout)
	ExpectWithOffset(1, err).To(BeNil(), "Cilium-agent cannot be started")

	vm.NetworkCreate(helpers.CiliumDockerNetwork, "")
	res := vm.NetworkGet(helpers.CiliumDockerNetwork)
	ExpectWithOffset(1, res).To(helpers.CMDSuccess(), "Cilium docker network is not created")

	res = vm.SetPolicyEnforcement(helpers.PolicyEnforcementDefault)
	ExpectWithOffset(1, res).To(helpers.CMDSuccess(), "Cannot set policy enforcement default")
}

// ExpectCiliumNotRunning checks that cilium has stopped with a grace period.
func ExpectCiliumNotRunning(vm *helpers.SSHMeta) {
	body := func() bool {
		res := vm.Exec("systemctl is-active cilium")
		return !res.WasSuccessful()
	}
	err := helpers.WithTimeout(body, "Cilium is still running after timeout", &helpers.TimeoutConfig{Timeout: helpers.CiliumStartTimeout})
	ExpectWithOffset(1, err).To(BeNil(), "Cilium is still running after timeout")
}

// ExpectDockerContainersMatchCiliumEndpoints asserts that docker containers in
// Cilium network match with the endpoint list
func ExpectDockerContainersMatchCiliumEndpoints(vm *helpers.SSHMeta) {
	ExpectWithOffset(1, vm.WaitEndpointsReady()).Should(BeTrue(), "Endpoints are not ready after timeout")
	ExpectWithOffset(1, vm.ValidateEndpointsAreCorrect(helpers.CiliumDockerNetwork)).To(BeNil(),
		"Docker containers mistmach with Cilium Endpoints")
}
<|endoftext|>


<commit_before>package bitops

// msb32 returns the number of bits required to store the value x
func msb32(x uint32) int {
	pos := 32
	temp := x >> 16
	if temp != 0 {
		pos -= 16
		x = temp
	}
	temp = x >> 8
	if temp != 0 {
		pos -= 8
		x = temp
	}
	temp = x >> 4
	if temp != 0 {
		pos -= 4
		x = temp
	}
	temp = x >> 2
	if temp != 0 {
		pos -= 2
		x = temp
	}
	temp = x >> 1
	if temp != 0 {
		return pos - 2
	}

	return int(uint32(pos) - x)
}

// msb64 returns the number of bits required to store the value x
func msb64(n uint64) int {
	if n <= 0 {
		return -1
	}
	var r, v uint
	if n >= 1<<32 {
		r += 32
		v = uint(n >> 32)
	} else {
		v = uint(n)
	}
	if v >= 1<<16 {
		r += 16
		v >>= 16
	}
	if v >= 1<<8 {
		r += 8
		v >>= 8
	}
	if v >= 1<<4 {
		r += 4
		v >>= 4
	}
	if v >= 1<<2 {
		r += 2
		v >>= 2
	}
	r += v >> 1
	return int(r)
}
<commit_msg>Add 64bit zig zag encoding helpers<commit_after>package bitops

// msb32 returns the number of bits required to store the value x
func msb32(x uint32) int {
	pos := 32
	temp := x >> 16
	if temp != 0 {
		pos -= 16
		x = temp
	}
	temp = x >> 8
	if temp != 0 {
		pos -= 8
		x = temp
	}
	temp = x >> 4
	if temp != 0 {
		pos -= 4
		x = temp
	}
	temp = x >> 2
	if temp != 0 {
		pos -= 2
		x = temp
	}
	temp = x >> 1
	if temp != 0 {
		return pos - 2
	}

	return int(uint32(pos) - x)
}

// msb64 returns the number of bits required to store the value x
func msb64(n uint64) int {
	if n <= 0 {
		return -1
	}
	var r, v uint
	if n >= 1<<32 {
		r += 32
		v = uint(n >> 32)
	} else {
		v = uint(n)
	}
	if v >= 1<<16 {
		r += 16
		v >>= 16
	}
	if v >= 1<<8 {
		r += 8
		v >>= 8
	}
	if v >= 1<<4 {
		r += 4
		v >>= 4
	}
	if v >= 1<<2 {
		r += 2
		v >>= 2
	}
	r += v >> 1
	return int(r)
}

func ZigZagEncode64(x int64) uint64 {
	return uint64(uint64(x<<1) ^ uint64((int64(x) >> 63)))
}

func ZigZagDecode64(v uint64) int64 {
	return int64((v >> 1) ^ uint64((int64(v&1)<<63)>>63))
}
<|endoftext|>


<commit_before>package scenarios

import (
	"net"
	"testing"

	v2net "github.com/v2ray/v2ray-core/common/net"
	v2testing "github.com/v2ray/v2ray-core/testing"
	"github.com/v2ray/v2ray-core/testing/assert"
	"github.com/v2ray/v2ray-core/testing/servers/tcp"
)

func TestRouter(t *testing.T) {
	v2testing.Current(t)

	tcpServer := &tcp.Server{
		Port: v2net.Port(50024),
		MsgProcessor: func(data []byte) []byte {
			buffer := make([]byte, 0, 2048)
			buffer = append(buffer, []byte("Processed: ")...)
			buffer = append(buffer, data...)
			return buffer
		},
	}
	_, err := tcpServer.Start()
	assert.Error(err).IsNil()
	defer tcpServer.Close()

	tcpServer2Accessed := false
	tcpServer2 := &tcp.Server{
		Port: v2net.Port(50025),
		MsgProcessor: func(data []byte) []byte {
			tcpServer2Accessed = true
			return data
		},
	}
	_, err = tcpServer2.Start()
	assert.Error(err).IsNil()
	defer tcpServer2.Close()

	assert.Error(InitializeServerSetOnce("test_3")).IsNil()

	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(50020),
	})

	payload := "direct dokodemo request."
	nBytes, err := conn.Write([]byte(payload))
	assert.Error(err).IsNil()
	assert.Int(nBytes).Equals(len(payload))

	conn.CloseWrite()

	response := make([]byte, 1024)
	nBytes, err = conn.Read(response)
	assert.Error(err).IsNil()
	assert.StringLiteral("Processed: " + payload).Equals(string(response[:nBytes]))
	conn.Close()

	conn, err = net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(50022),
	})

	payload = "blocked dokodemo request."
	nBytes, err = conn.Write([]byte(payload))
	assert.Error(err).IsNil()
	assert.Int(nBytes).Equals(len(payload))

	conn.CloseWrite()

	response = make([]byte, 1024)
	nBytes, err = conn.Read(response)
	assert.Error(err).IsNotNil()
	assert.Int(nBytes).Equals(0)
	assert.Bool(tcpServer2Accessed).IsFalse()
	conn.Close()

	CloseAllServers()
}
<commit_msg>assert error is nil<commit_after>package scenarios

import (
	"net"
	"testing"

	v2net "github.com/v2ray/v2ray-core/common/net"
	v2testing "github.com/v2ray/v2ray-core/testing"
	"github.com/v2ray/v2ray-core/testing/assert"
	"github.com/v2ray/v2ray-core/testing/servers/tcp"
)

func TestRouter(t *testing.T) {
	v2testing.Current(t)

	tcpServer := &tcp.Server{
		Port: v2net.Port(50024),
		MsgProcessor: func(data []byte) []byte {
			buffer := make([]byte, 0, 2048)
			buffer = append(buffer, []byte("Processed: ")...)
			buffer = append(buffer, data...)
			return buffer
		},
	}
	_, err := tcpServer.Start()
	assert.Error(err).IsNil()
	defer tcpServer.Close()

	tcpServer2Accessed := false
	tcpServer2 := &tcp.Server{
		Port: v2net.Port(50025),
		MsgProcessor: func(data []byte) []byte {
			tcpServer2Accessed = true
			return data
		},
	}
	_, err = tcpServer2.Start()
	assert.Error(err).IsNil()
	defer tcpServer2.Close()

	assert.Error(InitializeServerSetOnce("test_3")).IsNil()

	conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(50020),
	})
	assert.Error(err).IsNil()

	payload := "direct dokodemo request."
	nBytes, err := conn.Write([]byte(payload))
	assert.Error(err).IsNil()
	assert.Int(nBytes).Equals(len(payload))

	conn.CloseWrite()

	response := make([]byte, 1024)
	nBytes, err = conn.Read(response)
	assert.Error(err).IsNil()
	assert.StringLiteral("Processed: " + payload).Equals(string(response[:nBytes]))
	conn.Close()

	conn, err = net.DialTCP("tcp", nil, &net.TCPAddr{
		IP:   []byte{127, 0, 0, 1},
		Port: int(50022),
	})

	payload = "blocked dokodemo request."
	nBytes, err = conn.Write([]byte(payload))
	assert.Error(err).IsNil()
	assert.Int(nBytes).Equals(len(payload))

	conn.CloseWrite()

	response = make([]byte, 1024)
	nBytes, err = conn.Read(response)
	assert.Error(err).IsNotNil()
	assert.Int(nBytes).Equals(0)
	assert.Bool(tcpServer2Accessed).IsFalse()
	conn.Close()

	CloseAllServers()
}
<|endoftext|>


<commit_before>package mcat_test

import (
	"github.com/cloudfoundry/hm9000/config"
	"github.com/cloudfoundry/hm9000/models"
	"github.com/cloudfoundry/hm9000/store"
	"github.com/cloudfoundry/hm9000/testhelpers/desiredstateserver"
	"github.com/cloudfoundry/sonde-go/events"
	"github.com/cloudfoundry/storeadapter/storerunner/etcdstorerunner"
	"github.com/cloudfoundry/yagnats"
	. "github.com/onsi/gomega"
)

type Simulator struct {
	conf                   *config.Config
	storeRunner            *etcdstorerunner.ETCDClusterRunner
	store                  store.Store
	desiredStateServer     *desiredstateserver.DesiredStateServer
	currentHeartbeats      []*models.Heartbeat
	currentTimestamp       int
	cliRunner              *CLIRunner
	TicksToAttainFreshness int
	TicksToExpireHeartbeat int
	GracePeriod            int
	messageBus             yagnats.NATSConn
}

func NewSimulator(conf *config.Config, storeRunner *etcdstorerunner.ETCDClusterRunner, store store.Store, desiredStateServer *desiredstateserver.DesiredStateServer, cliRunner *CLIRunner, messageBus yagnats.NATSConn) *Simulator {
	desiredStateServer.Reset()

	return &Simulator{
		currentTimestamp:       100,
		conf:                   conf,
		storeRunner:            storeRunner,
		store:                  store,
		desiredStateServer:     desiredStateServer,
		cliRunner:              cliRunner,
		TicksToAttainFreshness: int(conf.ActualFreshnessTTLInHeartbeats) + 1,
		TicksToExpireHeartbeat: int(conf.HeartbeatTTLInHeartbeats),
		GracePeriod:            int(conf.GracePeriodInHeartbeats),
		messageBus:             messageBus,
	}
}

func (s *Simulator) Tick(numTicks int) {
	timeBetweenTicks := int(s.conf.HeartbeatPeriod)

	for i := 0; i < numTicks; i++ {
		s.currentTimestamp += timeBetweenTicks
		s.storeRunner.FastForwardTime(timeBetweenTicks)
		s.sendHeartbeats()
		//		s.sendHeartbeatsHttp()
		s.cliRunner.Run("analyze", s.currentTimestamp)
	}
}

func (s *Simulator) sendHeartbeats() {
	s.cliRunner.StartListener(s.currentTimestamp)
	metronAgent.Reset()
	for _, heartbeat := range s.currentHeartbeats {
		s.messageBus.Publish("dea.heartbeat", heartbeat.ToJSON())
	}

	nHeartbeats := len(s.currentHeartbeats)
	Eventually(func() bool {
		return metronAgent.MatchEvent("listener", events.Envelope_ValueMetric, "SavedHeartbeats", float64(nHeartbeats))
	}, IterationTimeout, 1.05).Should(BeTrue())

	s.cliRunner.StopListener()
}

//func (s *Simulator) sendHeartbeatsHttp() {
//	s.cliRunner.StartListener(s.currentTimestamp)
//	metronAgent.Reset()
//	for _, heartbeat := range s.currentHeartbeats {
//		_, err := http.Post("http://localhost:5335/dea/hearbeat", "application/json", bytes.NewBuffer(heartbeat.ToJSON()))
//		Expect(err).NotTo(HaveOccurred())
//		// Do HTTP-y stuff
//	}
//
//	nHeartbeats := len(s.currentHeartbeats)
//	Eventually(func() bool {
//		return metronAgent.MatchEvent("listener", events.Envelope_ValueMetric, "SavedHeartbeats", float64(nHeartbeats))
//	}, IterationTimeout, 0.05).Should(BeTrue())
//
//	s.cliRunner.StopListener()
//}
//
func (s *Simulator) SetDesiredState(desiredStates ...models.DesiredAppState) {
	s.desiredStateServer.SetDesiredState(desiredStates)
}

func (s *Simulator) SetCurrentHeartbeats(heartbeats ...*models.Heartbeat) {
	s.currentHeartbeats = heartbeats
}
<commit_msg>Remove http_heartbeat stuff<commit_after>package mcat_test

import (
	"github.com/cloudfoundry/hm9000/config"
	"github.com/cloudfoundry/hm9000/models"
	"github.com/cloudfoundry/hm9000/store"
	"github.com/cloudfoundry/hm9000/testhelpers/desiredstateserver"
	"github.com/cloudfoundry/sonde-go/events"
	"github.com/cloudfoundry/storeadapter/storerunner/etcdstorerunner"
	"github.com/cloudfoundry/yagnats"
	. "github.com/onsi/gomega"
)

type Simulator struct {
	conf                   *config.Config
	storeRunner            *etcdstorerunner.ETCDClusterRunner
	store                  store.Store
	desiredStateServer     *desiredstateserver.DesiredStateServer
	currentHeartbeats      []*models.Heartbeat
	currentTimestamp       int
	cliRunner              *CLIRunner
	TicksToAttainFreshness int
	TicksToExpireHeartbeat int
	GracePeriod            int
	messageBus             yagnats.NATSConn
}

func NewSimulator(conf *config.Config, storeRunner *etcdstorerunner.ETCDClusterRunner, store store.Store, desiredStateServer *desiredstateserver.DesiredStateServer, cliRunner *CLIRunner, messageBus yagnats.NATSConn) *Simulator {
	desiredStateServer.Reset()

	return &Simulator{
		currentTimestamp:       100,
		conf:                   conf,
		storeRunner:            storeRunner,
		store:                  store,
		desiredStateServer:     desiredStateServer,
		cliRunner:              cliRunner,
		TicksToAttainFreshness: int(conf.ActualFreshnessTTLInHeartbeats) + 1,
		TicksToExpireHeartbeat: int(conf.HeartbeatTTLInHeartbeats),
		GracePeriod:            int(conf.GracePeriodInHeartbeats),
		messageBus:             messageBus,
	}
}

func (s *Simulator) Tick(numTicks int) {
	timeBetweenTicks := int(s.conf.HeartbeatPeriod)

	for i := 0; i < numTicks; i++ {
		s.currentTimestamp += timeBetweenTicks
		s.storeRunner.FastForwardTime(timeBetweenTicks)
		s.sendHeartbeats()
		s.cliRunner.Run("analyze", s.currentTimestamp)
	}
}

func (s *Simulator) sendHeartbeats() {
	s.cliRunner.StartListener(s.currentTimestamp)
	metronAgent.Reset()
	for _, heartbeat := range s.currentHeartbeats {
		s.messageBus.Publish("dea.heartbeat", heartbeat.ToJSON())
	}

	nHeartbeats := len(s.currentHeartbeats)
	Eventually(func() bool {
		return metronAgent.MatchEvent("listener", events.Envelope_ValueMetric, "SavedHeartbeats", float64(nHeartbeats))
	}, IterationTimeout, 1.05).Should(BeTrue())

	s.cliRunner.StopListener()
}

func (s *Simulator) SetDesiredState(desiredStates ...models.DesiredAppState) {
	s.desiredStateServer.SetDesiredState(desiredStates)
}

func (s *Simulator) SetCurrentHeartbeats(heartbeats ...*models.Heartbeat) {
	s.currentHeartbeats = heartbeats
}
<|endoftext|>


<commit_before>package SchemaTool

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/xitongsys/parquet-go/SchemaHandler"
	"github.com/xitongsys/parquet-go/parquet"
)

func ParquetTypeToParquetTypeStr(pT *parquet.Type, cT *parquet.ConvertedType) (string, string) {
	var pTStr, cTStr string
	if pT != nil {
		switch *pT {
		case parquet.Type_BOOLEAN:
			pTStr = "BOOLEAN"
		case parquet.Type_INT32:
			pTStr = "INT32"
		case parquet.Type_INT64:
			pTStr = "INT64"
		case parquet.Type_INT96:
			pTStr = "INT96"
		case parquet.Type_FLOAT:
			pTStr = "FLOAT"
		case parquet.Type_DOUBLE:
			pTStr = "DOUBLE"
		case parquet.Type_BYTE_ARRAY:
			pTStr = "BYTE_ARRAY"
		case parquet.Type_FIXED_LEN_BYTE_ARRAY:
			pTStr = "FIXED_LEN_BYTE_ARRAY"
		}
	}
	if cT != nil {
		switch *cT {
		case parquet.ConvertedType_UTF8:
			cTStr = "UTF8"
		case parquet.ConvertedType_INT_8:
			cTStr = "INT_8"
		case parquet.ConvertedType_INT_16:
			cTStr = "INT_16"
		case parquet.ConvertedType_INT_32:
			cTStr = "INT_32"
		case parquet.ConvertedType_INT_64:
			cTStr = "INT_64"
		case parquet.ConvertedType_UINT_8:
			cTStr = "UINT_8"
		case parquet.ConvertedType_UINT_16:
			cTStr = "UINT_16"
		case parquet.ConvertedType_UINT_32:
			cTStr = "UINT_32"
		case parquet.ConvertedType_UINT_64:
			cTStr = "UINT_64"
		case parquet.ConvertedType_DATE:
			cTStr = "DATE"
		case parquet.ConvertedType_TIME_MILLIS:
			cTStr = "TIME_MILLIS"
		case parquet.ConvertedType_TIME_MICROS:
			cTStr = "TIME_MICROS"
		case parquet.ConvertedType_TIMESTAMP_MILLIS:
			cTStr = "TIMESTAMP_MILLIS"
		case parquet.ConvertedType_TIMESTAMP_MICROS:
			cTStr = "TIMESTAMP_MICROS"
		case parquet.ConvertedType_INTERVAL:
			cTStr = "INTERVAL"
		case parquet.ConvertedType_DECIMAL:
			cTStr = "DECIMAL"
		case parquet.ConvertedType_MAP:
			cTStr = "MAP"
		case parquet.ConvertedType_LIST:
			cTStr = "LIST"
		}
	}
	return pTStr, cTStr
}

func ParquetTypeToGoTypeStr(pT *parquet.Type, cT *parquet.ConvertedType) string {
	res := ""
	if pT != nil {
		switch *pT {
		case parquet.Type_BOOLEAN:
			res = "bool"
		case parquet.Type_INT32:
			res = "int32"
		case parquet.Type_INT64:
			res = "int64"
		case parquet.Type_INT96:
			res = "string"
		case parquet.Type_FLOAT:
			res = "float32"
		case parquet.Type_DOUBLE:
			res = "float64"
		case parquet.Type_BYTE_ARRAY:
			res = "string"
		case parquet.Type_FIXED_LEN_BYTE_ARRAY:
			res = "string"
		}
	}
	if cT != nil {
		switch *cT {
		case parquet.ConvertedType_UINT_8:
			res = "uint32"
		case parquet.ConvertedType_UINT_16:
			res = "uint32"
		case parquet.ConvertedType_UINT_32:
			res = "uint32"
		case parquet.ConvertedType_UINT_64:
			res = "uint64"
		}
	}
	return res
}

type Node struct {
	Indent   string
	SE       *parquet.SchemaElement
	Children []*Node
}

func NewNode(schema *parquet.SchemaElement) *Node {
	node := &(Node{
		Indent:   "",
		SE:       schema,
		Children: []*Node{},
	})
	return node
}

func (self *Node) OutputJsonSchema() string {
	res := "{\"Tag\":"
	pT, cT := self.SE.Type, self.SE.ConvertedType
	rTStr := "REQUIRED"
	if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_OPTIONAL {
		rTStr = "OPTIONAL"
	} else if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_REPEATED {
		rTStr = "REPEATED"
	}

	pTStr, cTStr := ParquetTypeToParquetTypeStr(pT, cT)
	tagStr := "\"name=%s, type=%s, repetitiontype=%s\""
	name := self.SE.GetName()

	if len(self.Children) == 0 {
		if *pT == parquet.Type_FIXED_LEN_BYTE_ARRAY && cT == nil {
			length := self.SE.GetTypeLength()
			tagStr = "\"name=%s, type=%s, length=%d, repetitiontype=%s\""
			res += fmt.Sprintf(tagStr, name, pTStr, length, rTStr) + "}"

		} else if cT != nil && *cT == parquet.ConvertedType_DECIMAL {
			scale, precision := self.SE.GetScale(), self.SE.GetPrecision()
			if *pT == parquet.Type_FIXED_LEN_BYTE_ARRAY {
				length := self.SE.GetTypeLength()
				tagStr = "\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, length=%d, repetitiontype=%s\""
				res += fmt.Sprintf(tagStr, name, cTStr, pTStr, scale, precision, length, rTStr) + "}"
			} else {
				tagStr = "\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, repetitiontype\""
				res += fmt.Sprintf(tagStr, name, cTStr, pTStr, scale, precision, rTStr) + "}"
			}

		} else {
			typeStr := pTStr
			if cT != nil {
				typeStr = cTStr
			}
			res += fmt.Sprintf(tagStr, name, typeStr, rTStr) + "}"

		}
	} else {
		if cT != nil {
			tagStr = "\"name=%s, type=%s, repetitiontype=%s\""
			res += fmt.Sprintf(tagStr, name, cTStr, rTStr)
		} else {
			tagStr = "\"name=%s, repetitiontype=%s\""
			res += fmt.Sprintf(tagStr, name, rTStr)
		}
		res += ",\n"
		res += "\"Fields\":[\n"

		nodes := self.Children
		if cT != nil {
			nodes = self.Children[0].Children
		}

		for i := 0; i < len(nodes); i++ {
			cNode := nodes[i]
			if i == len(nodes)-1 {
				res += cNode.OutputJsonSchema() + "\n"
			} else {
				res += cNode.OutputJsonSchema() + ",\n"
			}
		}

		res += "]\n"
		res += "}"

	}
	return res
}

func (cNode *Node) getStructTags() string {
	rTStr := "REQUIRED"
	if cNode.SE.GetRepetitionType() == parquet.FieldRepetitionType_OPTIONAL {
		rTStr = "OPTIONAL"
	} else if cNode.SE.GetRepetitionType() == parquet.FieldRepetitionType_REPEATED {
		rTStr = "REPEATED"
	}

	tags := fmt.Sprintf("`parquet:\"name=%s, type=%s, repetitiontype=%s\"`", cNode.SE.Name, *cNode.SE.Type, rTStr)

	pTStr, cTStr := ParquetTypeToParquetTypeStr(cNode.SE.Type, cNode.SE.ConvertedType)
	if *cNode.SE.Type == parquet.Type_FIXED_LEN_BYTE_ARRAY && cNode.SE.ConvertedType == nil {
		length := cNode.SE.GetTypeLength()
		tagStr := "`parquet:\"name=%s, type=%s, length=%d, repetitiontype=%s\"`"
		tags = fmt.Sprintf(tagStr, cNode.SE.Name, pTStr, length, rTStr)
	} else if cNode.SE.ConvertedType != nil && *cNode.SE.ConvertedType == parquet.ConvertedType_DECIMAL {
		scale, precision := cNode.SE.GetScale(), cNode.SE.GetPrecision()
		if *cNode.SE.Type == parquet.Type_FIXED_LEN_BYTE_ARRAY {
			length := cNode.SE.GetTypeLength()
			tagStr := "`parquet:\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, length=%d, repetitiontype=%s\"`"
			tags = fmt.Sprintf(tagStr, cNode.SE.Name, cTStr, pTStr, scale, precision, length, rTStr)
		} else {
			tagStr := "`parquet:\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, repetitiontype\"`"
			tags = fmt.Sprintf(tagStr, cNode.SE.Name, cNode.SE.Type, pTStr, scale, precision, rTStr)
		}
	}

	return tags
}

func (self *Node) OutputStruct(withName bool, withTags bool) string {
	name := self.SE.GetName()

	res := ""
	if withName {
		res += strings.Title(name)
	}

	pT, cT := self.SE.Type, self.SE.ConvertedType
	rTStr := " "
	if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_OPTIONAL {
		rTStr = " *"
	} else if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_REPEATED {
		rTStr = " []"
	}
	if !withName {
		rTStr = rTStr[1:]
	}

	if pT == nil && cT == nil {
		res += rTStr + "struct{\n"
		for _, cNode := range self.Children {
			if withTags {
				res += cNode.OutputStruct(true, withTags) + " " + cNode.getStructTags() + "\n"
			} else {
				res += cNode.OutputStruct(true, withTags) + "\n"
			}
		}
		res += "}"

	} else if cT != nil && *cT == parquet.ConvertedType_MAP {
		keyNode := self.Children[0].Children[0]
		keyPT, keyCT := keyNode.SE.Type, keyNode.SE.ConvertedType
		keyGoTypeStr := ParquetTypeToGoTypeStr(keyPT, keyCT)
		valNode := self.Children[0].Children[1]
		res += rTStr + "map[" + keyGoTypeStr + "]" + valNode.OutputStruct(false, withTags)

	} else if cT != nil && *cT == parquet.ConvertedType_LIST {
		cNode := self.Children[0].Children[0]
		res += rTStr + "[]" + cNode.OutputStruct(false, withTags)

	} else {
		goTypeStr := ParquetTypeToGoTypeStr(pT, cT)
		res += rTStr + goTypeStr
	}

	ress := strings.Split(res, "\n")
	for i := 0; i < len(ress); i++ {
		if i > 0 || withName {
			ress[i] = self.Indent + ress[i]
		}
	}

	res = strings.Join(ress, "\n")
	return res
}

type SchemaTree struct {
	Root *Node
}

func CreateSchemaTree(schemas []*parquet.SchemaElement) *SchemaTree {
	pos := 0
	stack := make([]*Node, 0)
	root := NewNode(schemas[0])
	stack = append(stack, root)
	pos++

	for len(stack) > 0 {
		node := stack[len(stack)-1]
		numChildren := int(node.SE.GetNumChildren())
		lc := len(node.Children)
		if lc < numChildren {
			newNode := NewNode(schemas[pos])
			newNode.Indent += "  "
			node.Children = append(node.Children, newNode)
			stack = append(stack, newNode)
			pos++
		} else {
			stack = stack[:len(stack)-1]
		}
	}

	st := new(SchemaTree)
	st.Root = root
	return st
}

func (self *SchemaTree) OutputJsonSchema() string {
	jsonStr := self.Root.OutputJsonSchema()
	var obj SchemaHandler.JSONSchemaItemType
	json.Unmarshal([]byte(jsonStr), &obj)
	res, _ := json.MarshalIndent(&obj, "", "  ")
	return string(res)

}

func (self *SchemaTree) OutputStruct(withTags bool) string {
	return self.Root.OutputStruct(true, withTags)
}
<commit_msg>Fixes #89<commit_after>package SchemaTool

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/xitongsys/parquet-go/SchemaHandler"
	"github.com/xitongsys/parquet-go/parquet"
)

func ParquetTypeToParquetTypeStr(pT *parquet.Type, cT *parquet.ConvertedType) (string, string) {
	var pTStr, cTStr string
	if pT != nil {
		switch *pT {
		case parquet.Type_BOOLEAN:
			pTStr = "BOOLEAN"
		case parquet.Type_INT32:
			pTStr = "INT32"
		case parquet.Type_INT64:
			pTStr = "INT64"
		case parquet.Type_INT96:
			pTStr = "INT96"
		case parquet.Type_FLOAT:
			pTStr = "FLOAT"
		case parquet.Type_DOUBLE:
			pTStr = "DOUBLE"
		case parquet.Type_BYTE_ARRAY:
			pTStr = "BYTE_ARRAY"
		case parquet.Type_FIXED_LEN_BYTE_ARRAY:
			pTStr = "FIXED_LEN_BYTE_ARRAY"
		}
	}
	if cT != nil {
		switch *cT {
		case parquet.ConvertedType_UTF8:
			cTStr = "UTF8"
		case parquet.ConvertedType_INT_8:
			cTStr = "INT_8"
		case parquet.ConvertedType_INT_16:
			cTStr = "INT_16"
		case parquet.ConvertedType_INT_32:
			cTStr = "INT_32"
		case parquet.ConvertedType_INT_64:
			cTStr = "INT_64"
		case parquet.ConvertedType_UINT_8:
			cTStr = "UINT_8"
		case parquet.ConvertedType_UINT_16:
			cTStr = "UINT_16"
		case parquet.ConvertedType_UINT_32:
			cTStr = "UINT_32"
		case parquet.ConvertedType_UINT_64:
			cTStr = "UINT_64"
		case parquet.ConvertedType_DATE:
			cTStr = "DATE"
		case parquet.ConvertedType_TIME_MILLIS:
			cTStr = "TIME_MILLIS"
		case parquet.ConvertedType_TIME_MICROS:
			cTStr = "TIME_MICROS"
		case parquet.ConvertedType_TIMESTAMP_MILLIS:
			cTStr = "TIMESTAMP_MILLIS"
		case parquet.ConvertedType_TIMESTAMP_MICROS:
			cTStr = "TIMESTAMP_MICROS"
		case parquet.ConvertedType_INTERVAL:
			cTStr = "INTERVAL"
		case parquet.ConvertedType_DECIMAL:
			cTStr = "DECIMAL"
		case parquet.ConvertedType_MAP:
			cTStr = "MAP"
		case parquet.ConvertedType_LIST:
			cTStr = "LIST"
		}
	}
	return pTStr, cTStr
}

func ParquetTypeToGoTypeStr(pT *parquet.Type, cT *parquet.ConvertedType) string {
	res := ""
	if pT != nil {
		switch *pT {
		case parquet.Type_BOOLEAN:
			res = "bool"
		case parquet.Type_INT32:
			res = "int32"
		case parquet.Type_INT64:
			res = "int64"
		case parquet.Type_INT96:
			res = "string"
		case parquet.Type_FLOAT:
			res = "float32"
		case parquet.Type_DOUBLE:
			res = "float64"
		case parquet.Type_BYTE_ARRAY:
			res = "string"
		case parquet.Type_FIXED_LEN_BYTE_ARRAY:
			res = "string"
		}
	}
	if cT != nil {
		switch *cT {
		case parquet.ConvertedType_UINT_8:
			res = "uint32"
		case parquet.ConvertedType_UINT_16:
			res = "uint32"
		case parquet.ConvertedType_UINT_32:
			res = "uint32"
		case parquet.ConvertedType_UINT_64:
			res = "uint64"
		}
	}
	return res
}

type Node struct {
	Indent   string
	SE       *parquet.SchemaElement
	Children []*Node
}

func NewNode(schema *parquet.SchemaElement) *Node {
	node := &(Node{
		Indent:   "",
		SE:       schema,
		Children: []*Node{},
	})
	return node
}

func (self *Node) OutputJsonSchema() string {
	res := "{\"Tag\":"
	pT, cT := self.SE.Type, self.SE.ConvertedType
	rTStr := "REQUIRED"
	if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_OPTIONAL {
		rTStr = "OPTIONAL"
	} else if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_REPEATED {
		rTStr = "REPEATED"
	}

	pTStr, cTStr := ParquetTypeToParquetTypeStr(pT, cT)
	tagStr := "\"name=%s, type=%s, repetitiontype=%s\""
	name := self.SE.GetName()

	if len(self.Children) == 0 {
		if *pT == parquet.Type_FIXED_LEN_BYTE_ARRAY && cT == nil {
			length := self.SE.GetTypeLength()
			tagStr = "\"name=%s, type=%s, length=%d, repetitiontype=%s\""
			res += fmt.Sprintf(tagStr, name, pTStr, length, rTStr) + "}"

		} else if cT != nil && *cT == parquet.ConvertedType_DECIMAL {
			scale, precision := self.SE.GetScale(), self.SE.GetPrecision()
			if *pT == parquet.Type_FIXED_LEN_BYTE_ARRAY {
				length := self.SE.GetTypeLength()
				tagStr = "\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, length=%d, repetitiontype=%s\""
				res += fmt.Sprintf(tagStr, name, cTStr, pTStr, scale, precision, length, rTStr) + "}"
			} else {
				tagStr = "\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, repetitiontype\""
				res += fmt.Sprintf(tagStr, name, cTStr, pTStr, scale, precision, rTStr) + "}"
			}

		} else {
			typeStr := pTStr
			if cT != nil {
				typeStr = cTStr
			}
			res += fmt.Sprintf(tagStr, name, typeStr, rTStr) + "}"

		}
	} else {
		if cT != nil {
			tagStr = "\"name=%s, type=%s, repetitiontype=%s\""
			res += fmt.Sprintf(tagStr, name, cTStr, rTStr)
		} else {
			tagStr = "\"name=%s, repetitiontype=%s\""
			res += fmt.Sprintf(tagStr, name, rTStr)
		}
		res += ",\n"
		res += "\"Fields\":[\n"

		nodes := self.Children
		if cT != nil {
			nodes = self.Children[0].Children
		}

		for i := 0; i < len(nodes); i++ {
			cNode := nodes[i]
			if i == len(nodes)-1 {
				res += cNode.OutputJsonSchema() + "\n"
			} else {
				res += cNode.OutputJsonSchema() + ",\n"
			}
		}

		res += "]\n"
		res += "}"

	}
	return res
}

func (cNode *Node) getStructTags() string {
	rTStr := "REQUIRED"
	if cNode.SE.GetRepetitionType() == parquet.FieldRepetitionType_OPTIONAL {
		rTStr = "OPTIONAL"
	} else if cNode.SE.GetRepetitionType() == parquet.FieldRepetitionType_REPEATED {
		rTStr = "REPEATED"
	}

	seType := cNode.SE.GetType()
	tags := fmt.Sprintf("`parquet:\"name=%s, type=%s, repetitiontype=%s\"`", cNode.SE.Name, seType, rTStr)

	pTStr, cTStr := ParquetTypeToParquetTypeStr(cNode.SE.Type, cNode.SE.ConvertedType)
	if seType == parquet.Type_FIXED_LEN_BYTE_ARRAY && cNode.SE.ConvertedType == nil {
		length := cNode.SE.GetTypeLength()
		tagStr := "`parquet:\"name=%s, type=%s, length=%d, repetitiontype=%s\"`"
		tags = fmt.Sprintf(tagStr, cNode.SE.Name, pTStr, length, rTStr)
	} else if cNode.SE.ConvertedType != nil && *cNode.SE.ConvertedType == parquet.ConvertedType_DECIMAL {
		scale, precision := cNode.SE.GetScale(), cNode.SE.GetPrecision()
		if seType == parquet.Type_FIXED_LEN_BYTE_ARRAY {
			length := cNode.SE.GetTypeLength()
			tagStr := "`parquet:\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, length=%d, repetitiontype=%s\"`"
			tags = fmt.Sprintf(tagStr, cNode.SE.Name, cTStr, pTStr, scale, precision, length, rTStr)
		} else {
			tagStr := "`parquet:\"name=%s, type=%s, basetype=%s, scale=%d, precision=%d, repetitiontype\"`"
			tags = fmt.Sprintf(tagStr, cNode.SE.Name, cNode.SE.Type, pTStr, scale, precision, rTStr)
		}
	}

	return tags
}

func (self *Node) OutputStruct(withName bool, withTags bool) string {
	name := self.SE.GetName()

	res := ""
	if withName {
		res += strings.Title(name)
	}

	pT, cT := self.SE.Type, self.SE.ConvertedType
	rTStr := " "
	if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_OPTIONAL {
		rTStr = " *"
	} else if self.SE.GetRepetitionType() == parquet.FieldRepetitionType_REPEATED {
		rTStr = " []"
	}
	if !withName {
		rTStr = rTStr[1:]
	}

	if pT == nil && cT == nil {
		res += rTStr + "struct {\n"
		for _, cNode := range self.Children {
			if withTags {
				res += cNode.OutputStruct(true, withTags) + " " + cNode.getStructTags() + "\n"
			} else {
				res += cNode.OutputStruct(true, withTags) + "\n"
			}
		}
		res += "}"

	} else if cT != nil && *cT == parquet.ConvertedType_MAP && self.Children != nil {
		keyNode := self.Children[0].Children[0]
		keyPT, keyCT := keyNode.SE.Type, keyNode.SE.ConvertedType
		keyGoTypeStr := ParquetTypeToGoTypeStr(keyPT, keyCT)
		valNode := self.Children[0].Children[1]
		res += rTStr + "map[" + keyGoTypeStr + "]" + valNode.OutputStruct(false, withTags)

	} else if cT != nil && *cT == parquet.ConvertedType_LIST && self.Children != nil {
		cNode := self.Children[0].Children[0]
		res += rTStr + "[]" + cNode.OutputStruct(false, withTags)

	} else {
		goTypeStr := ParquetTypeToGoTypeStr(pT, cT)
		res += rTStr + goTypeStr
	}

	ress := strings.Split(res, "\n")
	for i := 0; i < len(ress); i++ {
		if i > 0 || withName {
			ress[i] = self.Indent + ress[i]
		}
	}

	res = strings.Join(ress, "\n")
	return res
}

type SchemaTree struct {
	Root *Node
}

func CreateSchemaTree(schemas []*parquet.SchemaElement) *SchemaTree {
	pos := 0
	stack := make([]*Node, 0)
	root := NewNode(schemas[0])
	stack = append(stack, root)
	pos++

	for len(stack) > 0 {
		node := stack[len(stack)-1]
		numChildren := int(node.SE.GetNumChildren())
		lc := len(node.Children)
		if lc < numChildren {
			newNode := NewNode(schemas[pos])
			newNode.Indent += "  "
			node.Children = append(node.Children, newNode)
			stack = append(stack, newNode)
			pos++
		} else {
			stack = stack[:len(stack)-1]
		}
	}

	st := new(SchemaTree)
	st.Root = root
	return st
}

func (self *SchemaTree) OutputJsonSchema() string {
	jsonStr := self.Root.OutputJsonSchema()
	var obj SchemaHandler.JSONSchemaItemType
	json.Unmarshal([]byte(jsonStr), &obj)
	res, _ := json.MarshalIndent(&obj, "", "  ")
	return string(res)

}

func (self *SchemaTree) OutputStruct(withTags bool) string {
	return self.Root.OutputStruct(true, withTags)
}
<|endoftext|>


<commit_before>package goutil

import (
	"errors"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

//Gostrap creates a temporary environment to run the go(1) tool in.
//
//It embeds an exec.Cmd configured to run the go tool in this environment.
type Gostrap struct {
	root, old string
	*exec.Cmd
}

//NewGostrap creates a new Gostrap.
//
//The embedded exec.Cmd still needs the tool (ie, run, install, etc) and its
//arguments appended to its Args list. It defaults to using os.Stdout
//and os.Stderr for the Stdout and Stderr of the Cmd, respectively.
//
//You should investigate Run and Install first, as they are likely what you
//need, and if not their source, and WithGostrap's, show how Gostrap is used.
func NewGostrap() (*Gostrap, error) {
	t, err := ioutil.TempDir("", "lib-goutil")
	if err != nil {
		return nil, err
	}
	if err = os.MkdirAll(filepath.Join(t, "src"), 0777); err != nil {
		os.RemoveAll(t)
		return nil, err
	}
	gop, err := exec.LookPath("go")
	if err != nil {
		return nil, err
	}
	env := os.Environ()
	for i, e := range env {
		if strings.HasPrefix(e, "GOPATH=") {
			env[i] = e + ":" + t
			break
		}
	}
	tmp := &Gostrap{
		root: t,
		Cmd: &exec.Cmd{
			Path:   gop,
			Env:    env,
			Args:   []string{"go"},
			Stdout: os.Stdout,
			Stderr: os.Stderr,
		},
	}
	return tmp, nil
}

//Destroy the Gostrap and everything in it and everything it loves.
func (t *Gostrap) Destroy() {
	//ignore error, we're just being nice by calling this
	t.Popdir()

	//don't want to do anything vicious if this the zero value
	if t.root != "" {
		//There is nothing to do if this fails and no semantic consequence,
		//so it may do so in silence.
		os.RemoveAll(t.root)
	}
}

//Chdir makes the temp location the current working directory.
//It saves the current directory so this can be undone with Popdir.
func (t *Gostrap) Chdir() error {
	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if err = os.Chdir(t.JoinPath()); err != nil {
		return err
	}
	t.old = cwd
	return nil
}

//Popdir returns the working directory to what it was before Chdir was called.
func (t *Gostrap) Popdir() error {
	if t.old == "" {
		return errors.New("Never changed directory")
	}
	if err := os.Chdir(t.old); err != nil {
		return err
	}
	t.old = ""
	return nil
}

//JoinPath joins path to the temp directory.
func (t *Gostrap) JoinPath(path ...string) string {
	return filepath.Join(t.root, "src", filepath.Join(path...))
}

//AddFile creates a file in the temp location.
func (t *Gostrap) AddFile(name string, contents []byte) error {
	return ioutil.WriteFile(t.JoinPath(name), contents, 0666)
}

//WithGostrap creates a new Gostrap, calls Chdir, and defers
//calls to PopDir and Destroy, in the appropriate order, before invoking f.
func WithGostrap(f func(*Gostrap) error) error {
	t, err := NewGostrap()
	if err != nil {
		return err
	}
	defer t.Destroy()

	err = t.Chdir()
	if err != nil {
		return err
	}
	defer t.Popdir()

	return f(t)
}

func tagargs(tags []string) (out []string) {
	if len(tags) > 0 {
		//XXX what the hell is the format for multiple tags?
		out = append(out, "-tags", strings.Join(tags, " "))
	}
	return
}

//Run calls go run on a single input file, with optional build tags,
//in a Gostrap that's created on the fly and destroyed after.
func Run(file []byte, tags ...string) error {
	return WithGostrap(func(t *Gostrap) error {
		err := t.AddFile("main.go", file)
		if err != nil {
			return err
		}

		t.Args = append(t.Args, "run")
		t.Args = append(t.Args, tagargs(tags)...)
		t.Args = append(t.Args, "main.go")

		return t.Run()
	})
}

//Install calls go install on a single input file, with optional build tags,
//in a Gostrap that's created on the fly and destroyed after.
func Install(name, location string, file []byte, tags ...string) error {
	if location == "" || location == "." {
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}
		location = pwd
	}
	return WithGostrap(func(t *Gostrap) error {
		err := t.AddFile("main.go", file)
		if err != nil {
			return err
		}

		t.Args = append(t.Args, "build")
		t.Args = append(t.Args, tagargs(tags)...)
		t.Args = append(t.Args, "-o", name, "main.go")

		err = t.Run()
		if err != nil {
			return err
		}
		//need to copy the file
		return os.Rename(name, filepath.Join(location, name))
	})
}
<commit_msg>force build of all dependencies since that doesn't work across gopaths<commit_after>package goutil

import (
	"errors"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

//Gostrap creates a temporary environment to run the go(1) tool in.
//
//It embeds an exec.Cmd configured to run the go tool in this environment.
type Gostrap struct {
	root, old string
	*exec.Cmd
}

//NewGostrap creates a new Gostrap.
//
//The embedded exec.Cmd still needs the tool (ie, run, install, etc) and its
//arguments appended to its Args list. It defaults to using os.Stdout
//and os.Stderr for the Stdout and Stderr of the Cmd, respectively.
//
//Note that the go command will not rebuild out of date packages from
//other workspaces, and that Gostrap executes in its own workspace, so
//you may have to include the -a flag to force rebuilding of all dependent
//packages.
//
//You should investigate Run and Install first, as they are likely what you
//need, and if not their source, and WithGostrap's, show how Gostrap is used.
func NewGostrap() (*Gostrap, error) {
	t, err := ioutil.TempDir("", "lib-goutil")
	if err != nil {
		return nil, err
	}
	if err = os.MkdirAll(filepath.Join(t, "src"), 0777); err != nil {
		os.RemoveAll(t)
		return nil, err
	}
	gop, err := exec.LookPath("go")
	if err != nil {
		return nil, err
	}
	env := os.Environ()
	for i, e := range env {
		if strings.HasPrefix(e, "GOPATH=") {
			env[i] = e + ":" + t
			break
		}
	}
	tmp := &Gostrap{
		root: t,
		Cmd: &exec.Cmd{
			Path:   gop,
			Env:    env,
			Args:   []string{"go"},
			Stdout: os.Stdout,
			Stderr: os.Stderr,
		},
	}
	return tmp, nil
}

//Destroy the Gostrap and everything in it and everything it loves.
func (t *Gostrap) Destroy() {
	//ignore error, we're just being nice by calling this
	t.Popdir()

	//don't want to do anything vicious if this the zero value
	if t.root != "" {
		//There is nothing to do if this fails and no semantic consequence,
		//so it may do so in silence.
		os.RemoveAll(t.root)
	}
}

//Chdir makes the temp location the current working directory.
//It saves the current directory so this can be undone with Popdir.
func (t *Gostrap) Chdir() error {
	cwd, err := os.Getwd()
	if err != nil {
		return err
	}
	if err = os.Chdir(t.JoinPath()); err != nil {
		return err
	}
	t.old = cwd
	return nil
}

//Popdir returns the working directory to what it was before Chdir was called.
func (t *Gostrap) Popdir() error {
	if t.old == "" {
		return errors.New("Never changed directory")
	}
	if err := os.Chdir(t.old); err != nil {
		return err
	}
	t.old = ""
	return nil
}

//JoinPath joins path to the temp directory.
func (t *Gostrap) JoinPath(path ...string) string {
	return filepath.Join(t.root, "src", filepath.Join(path...))
}

//AddFile creates a file in the temp location.
func (t *Gostrap) AddFile(name string, contents []byte) error {
	return ioutil.WriteFile(t.JoinPath(name), contents, 0666)
}

//WithGostrap creates a new Gostrap, calls Chdir, and defers
//calls to PopDir and Destroy, in the appropriate order, before invoking f.
func WithGostrap(f func(*Gostrap) error) error {
	t, err := NewGostrap()
	if err != nil {
		return err
	}
	defer t.Destroy()

	err = t.Chdir()
	if err != nil {
		return err
	}
	defer t.Popdir()

	return f(t)
}

func tagargs(tags []string) (out []string) {
	if len(tags) > 0 {
		//XXX what the hell is the format for multiple tags?
		out = append(out, "-tags", strings.Join(tags, " "))
	}
	return
}

//Run calls go run on a single input file, with optional build tags,
//in a Gostrap that's created on the fly and destroyed after.
func Run(file []byte, tags ...string) error {
	return WithGostrap(func(t *Gostrap) error {
		err := t.AddFile("main.go", file)
		if err != nil {
			return err
		}

		t.Args = append(t.Args, "run", "-a")
		t.Args = append(t.Args, tagargs(tags)...)
		t.Args = append(t.Args, "main.go")

		return t.Run()
	})
}

//Install calls go install on a single input file, with optional build tags,
//in a Gostrap that's created on the fly and destroyed after.
func Install(name, location string, file []byte, tags ...string) error {
	if location == "" || location == "." {
		pwd, err := os.Getwd()
		if err != nil {
			return err
		}
		location = pwd
	}
	return WithGostrap(func(t *Gostrap) error {
		err := t.AddFile("main.go", file)
		if err != nil {
			return err
		}

		t.Args = append(t.Args, "build", "-a")
		t.Args = append(t.Args, tagargs(tags)...)
		t.Args = append(t.Args, "-o", name, "main.go")

		err = t.Run()
		if err != nil {
			return err
		}
		//need to copy the file
		return os.Rename(name, filepath.Join(location, name))
	})
}
<|endoftext|>


<commit_before>// Copyright 2010  The "goconfig" Authors
//
// Use of this source code is governed by the Simplified BSD License
// that can be found in the LICENSE file.
//
// This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
// OR CONDITIONS OF ANY KIND, either express or implied. See the License
// for more details.

package config

import (
	"os"
	"strconv"
	"strings"
)

// Bool has the same behaviour as String but converts the response to bool.
// See "boolString" for string values converted to bool.
func (self *Config) Bool(section string, option string) (value bool, err os.Error) {
	sv, err := self.String(section, option)
	if err != nil {
		return false, err
	}

	value, ok := boolString[strings.ToLower(sv)]
	if !ok {
		return false, os.NewError("could not parse bool value: " + sv)
	}

	return value, nil
}

// Float has the same behaviour as String but converts the response to float.
func (self *Config) Float(section string, option string) (value float, err os.Error) {
	sv, err := self.String(section, option)
	if err == nil {
		value, err = strconv.Atof(sv)
	}

	return value, err
}

// Int has the same behaviour as String but converts the response to int.
func (self *Config) Int(section string, option string) (value int, err os.Error) {
	sv, err := self.String(section, option)
	if err == nil {
		value, err = strconv.Atoi(sv)
	}

	return value, err
}

// RawString gets the (raw) string value for the given option in the section.
// The raw string value is not subjected to unfolding, which was illustrated in
// the beginning of this documentation.
//
// It returns an error if either the section or the option do not exist.
func (self *Config) RawString(section string, option string) (value string, err os.Error) {
	if _, ok := self.data[section]; ok {
		if tValue, ok := self.data[section][option]; ok {
			return tValue.v, nil
		}
		return "", os.NewError(optionError(option).String())
	}
	return "", os.NewError(sectionError(section).String())
}

// String gets the string value for the given option in the section.
// If the value needs to be unfolded (see e.g. %(host)s example in the beginning
// of this documentation), then String does this unfolding automatically, up to
// _DEPTH_VALUES number of iterations.
//
// It returns an error if either the section or the option do not exist, or the
// unfolding cycled.
func (self *Config) String(section string, option string) (value string, err os.Error) {
	value, err = self.RawString(section, option)
	if err != nil {
		return "", err
	}

	var i int

	for i = 0; i < _DEPTH_VALUES; i++ { // keep a sane depth
		vr := varRegExp.FindString(value)
		if len(vr) == 0 {
			break
		}

		// Take off leading '%(' and trailing ')s'
		noption := strings.TrimLeft(vr, "%(")
		noption = strings.TrimRight(noption, ")s")

		// Search variable in default section
		nvalue, _ := self.data[_DEFAULT_SECTION][noption]
		if _, ok := self.data[section][noption]; ok {
			nvalue = self.data[section][noption]
		}
		if nvalue.v == "" {
			return "", os.NewError(optionError(noption).String())
		}

		// substitute by new value and take off leading '%(' and trailing ')s'
		value = strings.Replace(value, vr, nvalue.v, -1)
	}

	if i == _DEPTH_VALUES {
		return "", os.NewError("possible cycle while unfolding variables: " +
			"max depth of " + strconv.Itoa(_DEPTH_VALUES) + " reached")
	}

	return value, nil
}

<commit_msg>float is deprecated in favor of float64 by raylu<commit_after>// Copyright 2010  The "goconfig" Authors
//
// Use of this source code is governed by the Simplified BSD License
// that can be found in the LICENSE file.
//
// This software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
// OR CONDITIONS OF ANY KIND, either express or implied. See the License
// for more details.

package config

import (
	"os"
	"strconv"
	"strings"
)

// Bool has the same behaviour as String but converts the response to bool.
// See "boolString" for string values converted to bool.
func (self *Config) Bool(section string, option string) (value bool, err os.Error) {
	sv, err := self.String(section, option)
	if err != nil {
		return false, err
	}

	value, ok := boolString[strings.ToLower(sv)]
	if !ok {
		return false, os.NewError("could not parse bool value: " + sv)
	}

	return value, nil
}

// Float has the same behaviour as String but converts the response to float.
func (self *Config) Float(section string, option string) (value float64, err os.Error) {
	sv, err := self.String(section, option)
	if err == nil {
		value, err = strconv.Atof64(sv)
	}

	return value, err
}

// Int has the same behaviour as String but converts the response to int.
func (self *Config) Int(section string, option string) (value int, err os.Error) {
	sv, err := self.String(section, option)
	if err == nil {
		value, err = strconv.Atoi(sv)
	}

	return value, err
}

// RawString gets the (raw) string value for the given option in the section.
// The raw string value is not subjected to unfolding, which was illustrated in
// the beginning of this documentation.
//
// It returns an error if either the section or the option do not exist.
func (self *Config) RawString(section string, option string) (value string, err os.Error) {
	if _, ok := self.data[section]; ok {
		if tValue, ok := self.data[section][option]; ok {
			return tValue.v, nil
		}
		return "", os.NewError(optionError(option).String())
	}
	return "", os.NewError(sectionError(section).String())
}

// String gets the string value for the given option in the section.
// If the value needs to be unfolded (see e.g. %(host)s example in the beginning
// of this documentation), then String does this unfolding automatically, up to
// _DEPTH_VALUES number of iterations.
//
// It returns an error if either the section or the option do not exist, or the
// unfolding cycled.
func (self *Config) String(section string, option string) (value string, err os.Error) {
	value, err = self.RawString(section, option)
	if err != nil {
		return "", err
	}

	var i int

	for i = 0; i < _DEPTH_VALUES; i++ { // keep a sane depth
		vr := varRegExp.FindString(value)
		if len(vr) == 0 {
			break
		}

		// Take off leading '%(' and trailing ')s'
		noption := strings.TrimLeft(vr, "%(")
		noption = strings.TrimRight(noption, ")s")

		// Search variable in default section
		nvalue, _ := self.data[_DEFAULT_SECTION][noption]
		if _, ok := self.data[section][noption]; ok {
			nvalue = self.data[section][noption]
		}
		if nvalue.v == "" {
			return "", os.NewError(optionError(noption).String())
		}

		// substitute by new value and take off leading '%(' and trailing ')s'
		value = strings.Replace(value, vr, nvalue.v, -1)
	}

	if i == _DEPTH_VALUES {
		return "", os.NewError("possible cycle while unfolding variables: " +
			"max depth of " + strconv.Itoa(_DEPTH_VALUES) + " reached")
	}

	return value, nil
}

<|endoftext|>


<commit_before>package peco

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/nsf/termbox-go"
)

func TestReadRC(t *testing.T) {
	txt := `
{
	"Keymap": {
		"C-j": "peco.Finish",
		"C-x,C-c": "peco.Finish"
	},
	"Style": {
		"Basic": ["on_default", "default"],
		"Selected": ["underline", "on_cyan", "black"],
		"Query": ["yellow", "bold"],
		"Matched": ["cyan", "bold", "on_red"]
	},
	"Prompt": "[peco]"
}
`
	cfg := NewConfig()
	if err := json.Unmarshal([]byte(txt), cfg); err != nil {
		t.Fatalf("Error unmarshaling json: %s", err)
	}
	t.Logf("%#q", cfg)
}

type stringsToStyleTest struct {
	strings []string
	style   *Style
}

func TestStringsToStyle(t *testing.T) {
	tests := []stringsToStyleTest{
		stringsToStyleTest{
			strings: []string{"on_default", "default"},
			style:   &Style{fg: termbox.ColorDefault, bg: termbox.ColorDefault},
		},
		stringsToStyleTest{
			strings: []string{"bold", "on_blue", "yellow"},
			style:   &Style{fg: termbox.ColorYellow | termbox.AttrBold, bg: termbox.ColorBlue},
		},
		stringsToStyleTest{
			strings: []string{"underline", "on_cyan", "black"},
			style:   &Style{fg: termbox.ColorBlack | termbox.AttrUnderline, bg: termbox.ColorCyan},
		},
		stringsToStyleTest{
			strings: []string{"blink", "on_red", "white"},
			style:   &Style{fg: termbox.ColorWhite | termbox.AttrReverse, bg: termbox.ColorRed},
		},
	}

	t.Logf("Checking strings -> color mapping...")
	for _, test := range tests {
		t.Logf("    checking %s...", test.strings)
		if a := stringsToStyle(test.strings); *a != *test.style {
			t.Errorf("Expected '%s' to be '%#v', but got '%#v'", test.strings, test.style, a)
		}
	}
}

func TestLocateRcfile(t *testing.T) {
	dir, err := ioutil.TempDir("", "peco-")
	if err != nil {
		t.Fatalf("Failed to create temporary directory: %s", err)
	}

	homedirFunc = func() (string, error) {
		return dir, nil
	}

	expected := []string{
		filepath.Join(dir, "peco"),
		filepath.Join(dir, "1", "peco"),
		filepath.Join(dir, "2", "peco"),
		filepath.Join(dir, "3", "peco"),
		filepath.Join(dir, ".peco"),
	}

	i := 0
	_locateRcfileIn = func(dir string) (string, error) {
		t.Logf("looking for file in %s", dir)
		if i > len(expected)-1 {
			t.Fatalf("Got %d directories, only have %d", i+1, len(expected))
		}

		if expected[i] != dir {
			t.Errorf("Expected %s, got %s", expected[i], dir)
		}
		i++
		return "", fmt.Errorf("error: Not found")
	}

	os.Setenv("XDG_CONFIG_HOME", dir)
	os.Setenv("XDG_CONFIG_DIRS", strings.Join(
		[]string{
			filepath.Join(dir, "1"),
			filepath.Join(dir, "2"),
			filepath.Join(dir, "3"),
		},
		fmt.Sprintf("%c", filepath.ListSeparator),
	))

	LocateRcfile()
	expected[0] = filepath.Join(dir, ".config", "peco")
	os.Setenv("XDG_CONFIG_HOME", "")
	i = 0
	LocateRcfile()

}
<commit_msg>add test case ("on_blink" and "reverse" attribute)<commit_after>package peco

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/nsf/termbox-go"
)

func TestReadRC(t *testing.T) {
	txt := `
{
	"Keymap": {
		"C-j": "peco.Finish",
		"C-x,C-c": "peco.Finish"
	},
	"Style": {
		"Basic": ["on_default", "default"],
		"Selected": ["underline", "on_cyan", "black"],
		"Query": ["yellow", "bold"],
		"Matched": ["cyan", "bold", "on_red"]
	},
	"Prompt": "[peco]"
}
`
	cfg := NewConfig()
	if err := json.Unmarshal([]byte(txt), cfg); err != nil {
		t.Fatalf("Error unmarshaling json: %s", err)
	}
	t.Logf("%#q", cfg)
}

type stringsToStyleTest struct {
	strings []string
	style   *Style
}

func TestStringsToStyle(t *testing.T) {
	tests := []stringsToStyleTest{
		stringsToStyleTest{
			strings: []string{"on_default", "default"},
			style:   &Style{fg: termbox.ColorDefault, bg: termbox.ColorDefault},
		},
		stringsToStyleTest{
			strings: []string{"bold", "on_blue", "yellow"},
			style:   &Style{fg: termbox.ColorYellow | termbox.AttrBold, bg: termbox.ColorBlue},
		},
		stringsToStyleTest{
			strings: []string{"underline", "on_cyan", "black"},
			style:   &Style{fg: termbox.ColorBlack | termbox.AttrUnderline, bg: termbox.ColorCyan},
		},
		stringsToStyleTest{
			strings: []string{"reverse", "on_red", "white"},
			style:   &Style{fg: termbox.ColorWhite | termbox.AttrReverse, bg: termbox.ColorRed},
		},
		stringsToStyleTest{
			strings: []string{"on_bold", "on_magenta", "green"},
			style:   &Style{fg: termbox.ColorGreen, bg: termbox.ColorMagenta | termbox.AttrBold},
		},
	}

	t.Logf("Checking strings -> color mapping...")
	for _, test := range tests {
		t.Logf("    checking %s...", test.strings)
		if a := stringsToStyle(test.strings); *a != *test.style {
			t.Errorf("Expected '%s' to be '%#v', but got '%#v'", test.strings, test.style, a)
		}
	}
}

func TestLocateRcfile(t *testing.T) {
	dir, err := ioutil.TempDir("", "peco-")
	if err != nil {
		t.Fatalf("Failed to create temporary directory: %s", err)
	}

	homedirFunc = func() (string, error) {
		return dir, nil
	}

	expected := []string{
		filepath.Join(dir, "peco"),
		filepath.Join(dir, "1", "peco"),
		filepath.Join(dir, "2", "peco"),
		filepath.Join(dir, "3", "peco"),
		filepath.Join(dir, ".peco"),
	}

	i := 0
	_locateRcfileIn = func(dir string) (string, error) {
		t.Logf("looking for file in %s", dir)
		if i > len(expected)-1 {
			t.Fatalf("Got %d directories, only have %d", i+1, len(expected))
		}

		if expected[i] != dir {
			t.Errorf("Expected %s, got %s", expected[i], dir)
		}
		i++
		return "", fmt.Errorf("error: Not found")
	}

	os.Setenv("XDG_CONFIG_HOME", dir)
	os.Setenv("XDG_CONFIG_DIRS", strings.Join(
		[]string{
			filepath.Join(dir, "1"),
			filepath.Join(dir, "2"),
			filepath.Join(dir, "3"),
		},
		fmt.Sprintf("%c", filepath.ListSeparator),
	))

	LocateRcfile()
	expected[0] = filepath.Join(dir, ".config", "peco")
	os.Setenv("XDG_CONFIG_HOME", "")
	i = 0
	LocateRcfile()

}
<|endoftext|>