<commit_before>package main

import (
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

func importFile(filename string) {
	promptPassword()
	fileContents, err := ioutil.ReadFile(filename)
	if err != nil {
		logger.Error("%v", err)
		os.Exit(-1)
	}
	entries, _ := parseEntries(string(fileContents))
	for _, entry := range entries {
		writeEntry(entry, true)
	}
	logger.Info("Imported '%s' to %s.", filename, ConfigArgs.WorkingFile)
}

func exportFile(filename string) {
	promptPassword()
	fullText, _ := getFullEntry()
	err := ioutil.WriteFile(filename, []byte(fullText), 0644)
	if err != nil {
		logger.Error("%v", err)
		os.Exit(-1)
	}
	logger.Info("Exported '%s' to %s.", ConfigArgs.WorkingFile, filename)
}

func parseEntries(text string) ([]string, []int) {
	defer timeTrack(time.Now(), "Parsing entries")
	entry := ""
	entries := make(map[int]string)
	var gt int
	for _, line := range strings.Split(text, "\n") {
		entryWords := strings.Split(strings.TrimSpace(line), " ")
		if len(entryWords) > 1 {
			t1, e1 := time.Parse("2006-01-02 15:04:05", entryWords[0]+" "+entryWords[1])
			t2, e2 := time.Parse("2006-01-02 15:04", entryWords[0]+" "+entryWords[1])
			if e1 == nil || e2 == nil {
				if len(entry) > 0 {
					if _, ok := entries[gt]; ok {
						logger.Warn("Duplicate entry for %s", entryWords[0]+" "+entryWords[1])
					}
					entries[gt] = strings.TrimSpace(entry)
				}
				entry = ""
				if e1 == nil {
					gt = int(t1.Unix())
				} else {
					gt = int(t2.Unix())
				}
			}
		}
		entry += strings.TrimRight(line, " ") + "\n"
	}
	if len(entry) > 0 {
		entries[gt] = strings.TrimSpace(entry)
	}
	if len(entries) == 1 {
		return append([]string{}, entries[gt]), append([]int{}, gt)
	}

	entriesInOrder, gtsInOrder := sortEntries(entries)
	return entriesInOrder, gtsInOrder
}

func sortEntries(entries map[int]string) ([]string, []int) {
	// Sort the entries in order
	var keys []int
	for k := range entries {
		keys = append(keys, k)
	}
	sort.Ints(keys)

	entriesInOrder := []string{}
	gtsInOrder := []int{}
	for _, k := range keys {
		entriesInOrder = append(entriesInOrder, entries[k])
		gtsInOrder = append(gtsInOrder, k)
	}
	logger.Debug("Sorted %d entries.", len(entriesInOrder))
	return entriesInOrder, gtsInOrder
}

func cleanUp() error {
	dir := RuntimeArgs.TempPath
	d, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(dir, name))
		if err != nil {
			return err
		}
	}

	fileList := listFiles()
	for i, f := range fileList {
		files, _ := ioutil.ReadDir(path.Join(RuntimeArgs.WorkingPath, f))
		if len(files) < 2 {
			for _, file := range files {
				logger.Debug("Remove %s.", path.Join(RuntimeArgs.WorkingPath, f, file.Name()))
				err := os.Remove(path.Join(RuntimeArgs.WorkingPath, f, file.Name()))
				if err != nil {
					log.Fatal(err)
				}
			}
			logger.Debug("Remove %s.", path.Join(RuntimeArgs.WorkingPath, f))
			err := os.Remove(path.Join(RuntimeArgs.WorkingPath, f))
			if err != nil {
				log.Fatal(err)
			}
			if ConfigArgs.WorkingFile == f {
				if len(fileList) < 2 {
					ConfigArgs.WorkingFile = "notes.txt"
				} else {
					if i != 0 {
						ConfigArgs.WorkingFile = fileList[0]
					} else {
						ConfigArgs.WorkingFile = fileList[1]
					}
				}
				// Save current config parameters
				b, err := json.Marshal(ConfigArgs)
				if err != nil {
					log.Println(err)
				}
				ioutil.WriteFile(path.Join(RuntimeArgs.WorkingPath, "config.json"), b, 0644)

			}
		}
	}

	return nil
}

func editEntry() string {
	logger.Debug("Editing file")
	vimrc := `func! WordProcessorModeCLI()
		setlocal formatoptions=t1
		setlocal textwidth=80
		map j gj
		map k gk
		set formatprg=par
		setlocal wrap
		setlocal linebreak
		setlocal noexpandtab
		normal G$
endfu
com! WPCLI call WordProcessorModeCLI()`
	// Append to .vimrc file
	if exists(path.Join(RuntimeArgs.HomePath, ".vimrc")) {
		// Check if .vimrc file contains code
		logger.Debug("Found .vimrc.")
		fileContents, err := ioutil.ReadFile(path.Join(RuntimeArgs.HomePath, ".vimrc"))
		if err != nil {
			log.Fatal(err)
		}
		if !strings.Contains(string(fileContents), "com! WPCLI call WordProcessorModeCLI") {
			// Append to fileContents
			logger.Debug("WPCLI not found in .vimrc, adding it...")
			newvimrc := string(fileContents) + "\n" + vimrc
			err := ioutil.WriteFile(path.Join(RuntimeArgs.HomePath, ".vimrc"), []byte(newvimrc), 0644)
			if err != nil {
				log.Fatal(err)
			}
		} else {
			logger.Debug("WPCLI found in .vimrc.")
		}
	} else {
		logger.Debug("Can not find .vimrc, creating new .vimrc...")
		err := ioutil.WriteFile(path.Join(RuntimeArgs.HomePath, ".vimrc"), []byte(vimrc), 0644)
		if err != nil {
			log.Fatal(err)
		}
	}

	cmdArgs := []string{"-c", "WPCLI", "+startinsert", path.Join(RuntimeArgs.TempPath, "temp")}
	if len(RuntimeArgs.TextSearch) > 0 {
		searchTerms := strings.Split(RuntimeArgs.TextSearch, " ")
		cmdArgs = append([]string{"-c", "2match Keyword /\\c\\v(" + strings.Join(searchTerms, "|") + ")/"}, cmdArgs...)
	}
	cmd := exec.Command("vim", cmdArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
	fileContents, _ := ioutil.ReadFile(path.Join(RuntimeArgs.TempPath, "temp"))
	cleanUp()
	return string(fileContents)
}

func writeEntry(fileContents string, forceWrite bool) bool {
	// Hash contents to get filename
	h := sha1.New()
	h.Write([]byte(fileContents))
	sha1_hash := hex.EncodeToString(h.Sum(nil))
	fileName := string(sha1_hash) + ".gpg"
	if exists(path.Join(RuntimeArgs.FullPath, fileName)) {
		return false
	}

	// logger.Debug("Entry contains %d bytes.", len(fileContents))
	if len(fileContents) < 22 && !forceWrite {
		logger.Info("No data appended.")
		return false
	}

	encryptedText := encryptString(string(fileContents), RuntimeArgs.Passphrase)
	err := ioutil.WriteFile(path.Join(RuntimeArgs.FullPath, fileName), []byte(encryptedText), 0644)
	if err != nil {
		log.Fatal(err)
	}
	logger.Info("Wrote %s.", fileName)
	return true
}
<commit_msg>Fixed error that doesn't allow new file creation<commit_after>package main

import (
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

func importFile(filename string) {
	promptPassword()
	fileContents, err := ioutil.ReadFile(filename)
	if err != nil {
		logger.Error("%v", err)
		os.Exit(-1)
	}
	entries, _ := parseEntries(string(fileContents))
	for _, entry := range entries {
		writeEntry(entry, true)
	}
	logger.Info("Imported '%s' to %s.", filename, ConfigArgs.WorkingFile)
}

func exportFile(filename string) {
	promptPassword()
	fullText, _ := getFullEntry()
	err := ioutil.WriteFile(filename, []byte(fullText), 0644)
	if err != nil {
		logger.Error("%v", err)
		os.Exit(-1)
	}
	logger.Info("Exported '%s' to %s.", ConfigArgs.WorkingFile, filename)
}

func parseEntries(text string) ([]string, []int) {
	defer timeTrack(time.Now(), "Parsing entries")
	entry := ""
	entries := make(map[int]string)
	var gt int
	for _, line := range strings.Split(text, "\n") {
		entryWords := strings.Split(strings.TrimSpace(line), " ")
		if len(entryWords) > 1 {
			t1, e1 := time.Parse("2006-01-02 15:04:05", entryWords[0]+" "+entryWords[1])
			t2, e2 := time.Parse("2006-01-02 15:04", entryWords[0]+" "+entryWords[1])
			if e1 == nil || e2 == nil {
				if len(entry) > 0 {
					if _, ok := entries[gt]; ok {
						logger.Warn("Duplicate entry for %s", entryWords[0]+" "+entryWords[1])
					}
					entries[gt] = strings.TrimSpace(entry)
				}
				entry = ""
				if e1 == nil {
					gt = int(t1.Unix())
				} else {
					gt = int(t2.Unix())
				}
			}
		}
		entry += strings.TrimRight(line, " ") + "\n"
	}
	if len(entry) > 0 {
		entries[gt] = strings.TrimSpace(entry)
	}
	if len(entries) == 1 {
		return append([]string{}, entries[gt]), append([]int{}, gt)
	}

	entriesInOrder, gtsInOrder := sortEntries(entries)
	return entriesInOrder, gtsInOrder
}

func sortEntries(entries map[int]string) ([]string, []int) {
	// Sort the entries in order
	var keys []int
	for k := range entries {
		keys = append(keys, k)
	}
	sort.Ints(keys)

	entriesInOrder := []string{}
	gtsInOrder := []int{}
	for _, k := range keys {
		entriesInOrder = append(entriesInOrder, entries[k])
		gtsInOrder = append(gtsInOrder, k)
	}
	logger.Debug("Sorted %d entries.", len(entriesInOrder))
	return entriesInOrder, gtsInOrder
}

func cleanUp() error {
	logger.Debug("Cleaning..")
	dir := RuntimeArgs.TempPath
	d, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer d.Close()
	names, err := d.Readdirnames(-1)
	if err != nil {
		return err
	}
	for _, name := range names {
		err = os.RemoveAll(filepath.Join(dir, name))
		if err != nil {
			return err
		}
	}

	fileList := listFiles()
	for i, f := range fileList {
		files, _ := ioutil.ReadDir(path.Join(RuntimeArgs.WorkingPath, f))
		if len(files) < 2 {
			for _, file := range files {
				logger.Debug("Remove %s.", path.Join(RuntimeArgs.WorkingPath, f, file.Name()))
				err := os.Remove(path.Join(RuntimeArgs.WorkingPath, f, file.Name()))
				if err != nil {
					log.Fatal(err)
				}
			}
			logger.Debug("Remove %s.", path.Join(RuntimeArgs.WorkingPath, f))
			err := os.Remove(path.Join(RuntimeArgs.WorkingPath, f))
			if err != nil {
				log.Fatal(err)
			}
			if ConfigArgs.WorkingFile == f {
				if len(fileList) < 2 {
					ConfigArgs.WorkingFile = "notes.txt"
				} else {
					if i != 0 {
						ConfigArgs.WorkingFile = fileList[0]
					} else {
						ConfigArgs.WorkingFile = fileList[1]
					}
				}
				// Save current config parameters
				b, err := json.Marshal(ConfigArgs)
				if err != nil {
					log.Println(err)
				}
				ioutil.WriteFile(path.Join(RuntimeArgs.WorkingPath, "config.json"), b, 0644)

			}
		}
	}

	return nil
}

func editEntry() string {
	logger.Debug("Editing file")
	vimrc := `func! WordProcessorModeCLI()
		setlocal formatoptions=t1
		setlocal textwidth=80
		map j gj
		map k gk
		set formatprg=par
		setlocal wrap
		setlocal linebreak
		setlocal noexpandtab
		normal G$
endfu
com! WPCLI call WordProcessorModeCLI()`
	// Append to .vimrc file
	if exists(path.Join(RuntimeArgs.HomePath, ".vimrc")) {
		// Check if .vimrc file contains code
		logger.Debug("Found .vimrc.")
		fileContents, err := ioutil.ReadFile(path.Join(RuntimeArgs.HomePath, ".vimrc"))
		if err != nil {
			log.Fatal(err)
		}
		if !strings.Contains(string(fileContents), "com! WPCLI call WordProcessorModeCLI") {
			// Append to fileContents
			logger.Debug("WPCLI not found in .vimrc, adding it...")
			newvimrc := string(fileContents) + "\n" + vimrc
			err := ioutil.WriteFile(path.Join(RuntimeArgs.HomePath, ".vimrc"), []byte(newvimrc), 0644)
			if err != nil {
				log.Fatal(err)
			}
		} else {
			logger.Debug("WPCLI found in .vimrc.")
		}
	} else {
		logger.Debug("Can not find .vimrc, creating new .vimrc...")
		err := ioutil.WriteFile(path.Join(RuntimeArgs.HomePath, ".vimrc"), []byte(vimrc), 0644)
		if err != nil {
			log.Fatal(err)
		}
	}

	cmdArgs := []string{"-c", "WPCLI", "+startinsert", path.Join(RuntimeArgs.TempPath, "temp")}
	if len(RuntimeArgs.TextSearch) > 0 {
		searchTerms := strings.Split(RuntimeArgs.TextSearch, " ")
		cmdArgs = append([]string{"-c", "2match Keyword /\\c\\v(" + strings.Join(searchTerms, "|") + ")/"}, cmdArgs...)
	}
	cmd := exec.Command("vim", cmdArgs...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
	fileContents, _ := ioutil.ReadFile(path.Join(RuntimeArgs.TempPath, "temp"))
	return string(fileContents)
}

func writeEntry(fileContents string, forceWrite bool) bool {
	// Hash contents to get filename
	h := sha1.New()
	h.Write([]byte(fileContents))
	sha1_hash := hex.EncodeToString(h.Sum(nil))
	fileName := string(sha1_hash) + ".gpg"
	if exists(path.Join(RuntimeArgs.FullPath, fileName)) {
		return false
	}

	// logger.Debug("Entry contains %d bytes.", len(fileContents))
	if len(fileContents) < 22 && !forceWrite {
		logger.Info("No data appended.")
		return false
	}

	encryptedText := encryptString(string(fileContents), RuntimeArgs.Passphrase)
	err := ioutil.WriteFile(path.Join(RuntimeArgs.FullPath, fileName), []byte(encryptedText), 0644)
	if err != nil {
		log.Fatal(err)
	}
	logger.Info("Wrote %s.", fileName)
	return true
}
<|endoftext|>


<commit_before>package envh

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// ErrNotFound is triggered when environment variable cannot be found
var ErrNotFound = fmt.Errorf("Variable not found")

// ErrWrongType is triggered when we try to convert variable to a wrong type
var ErrWrongType = fmt.Errorf("Variable can't be converted")

// ErrDuplicated is triggered when a variable is already defined in a tree structure
var ErrDuplicated = fmt.Errorf("Variable was already defined before")

func parseVars() *map[string]string {
	results := map[string]string{}

	for _, v := range os.Environ() {
		e := strings.SplitN(v, "=", 2)

		results[e[0]] = e[1]
	}

	return &results
}

// Env manage environment variables
// by giving high level api to interact with them
type Env struct {
	envs *map[string]string
}

// NewEnv creates a new Env instance
func NewEnv() Env {
	return Env{parseVars()}
}

// GetAllValues retrieves a slice of all environment variables values
func (e Env) GetAllValues() []string {
	results := []string{}

	for _, v := range *e.envs {
		results = append(results, v)
	}

	return results
}

// GetAllKeys retrieves a slice of all environment variables keys
func (e Env) GetAllKeys() []string {
	results := []string{}

	for k := range *e.envs {
		results = append(results, k)
	}

	return results
}

// GetString return a string if variable exists
// or an error otherwise
func (e Env) GetString(key string) (string, error) {
	if v, ok := (*e.envs)[key]; ok {
		return v, nil
	}

	return "", ErrNotFound
}

// GetInt return an integer if variable exists
// or an error if value is not an integer or doesn't exist
func (e Env) GetInt(key string) (int, error) {
	v, ok := (*e.envs)[key]

	if !ok {
		return 0, ErrNotFound
	}

	i, err := strconv.Atoi(v)

	if err != nil {
		return 0, ErrWrongType
	}

	return i, nil
}

// GetFloat return a float if variable exists
// or an error if value is not a float or doesn't exist
func (e Env) GetFloat(key string) (float32, error) {
	v, ok := (*e.envs)[key]

	if !ok {
		return 0, ErrNotFound
	}

	f, err := strconv.ParseFloat(v, 32)

	if err != nil {
		return 0, ErrWrongType
	}

	return float32(f), nil
}

// GetBool return a boolean if variable exists
// or an error if value is not a boolean or doesn't exist
func (e Env) GetBool(key string) (bool, error) {
	v, ok := (*e.envs)[key]

	if !ok {
		return false, ErrNotFound
	}

	b, err := strconv.ParseBool(v)

	if err != nil {
		return false, ErrWrongType
	}

	return b, nil
}

// FindEntries retrieves all keys matching a given regexp and their
// corresponding values
func (e Env) FindEntries(reg string) (map[string]string, error) {
	results := map[string]string{}

	r, err := regexp.Compile(reg)

	if err != nil {
		return results, err
	}

	for k, v := range *e.envs {
		if r.MatchString(k) {
			results[k] = v
		}
	}

	return results, nil
}

type node struct {
	childs   []*node
	key      string
	value    string
	hasValue bool
	root     bool
}

func newNode() *node {
	return &node{childs: []*node{}}
}

func newRootNode() *node {
	return &node{childs: []*node{}, root: true}
}

func (n *node) findAllChildsByKey(key string, withValue bool) *[]*node {
	results := []*node{}
	nodes := n.childs

	for {
		tank := []*node{}

		for _, node := range nodes {
			if node.key == key {
				if withValue && node.hasValue || !withValue {
					results = append(results, node)
				}
			}

			tank = append(tank, node.childs...)
		}

		nodes = tank

		if len(tank) == 0 {
			return &results
		}
	}
}

func (n *node) findChildByKey(key string) (*node, bool) {
	for _, child := range n.childs {
		if child.key == key {
			return child, true
		}
	}

	return nil, false
}

func (n *node) appendChild(child *node) bool {
	if _, ok := n.findChildByKey(child.key); ok {
		return false
	}

	n.childs = append(n.childs, child)

	return true
}

func createTreeFromDelimiterFilteringByRegexp(reg *regexp.Regexp, delimiter string) (*node, error) {
	rootNode := newRootNode()

	for key, value := range *parseVars() {
		if reg.MatchString(key) {
			current := rootNode

			for _, component := range strings.Split(key, delimiter) {
				n, exists := current.findChildByKey(component)

				if exists {
					current = n
				} else {
					child := newNode()
					child.key = component
					current.appendChild(child)
					current = child
				}
			}

			if current.hasValue {
				return nil, ErrDuplicated
			}

			current.hasValue = true
			current.value = value
		}
	}

	return rootNode, nil
}

}
<commit_msg>Move logic to private functions<commit_after>package envh

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// ErrNotFound is triggered when environment variable cannot be found
var ErrNotFound = fmt.Errorf("Variable not found")

// ErrWrongType is triggered when we try to convert variable to a wrong type
var ErrWrongType = fmt.Errorf("Variable can't be converted")

// ErrDuplicated is triggered when a variable is already defined in a tree structure
var ErrDuplicated = fmt.Errorf("Variable was already defined before")

func parseVars() *map[string]string {
	results := map[string]string{}

	for _, v := range os.Environ() {
		e := strings.SplitN(v, "=", 2)

		results[e[0]] = e[1]
	}

	return &results
}

func getString(fun func() (string, bool)) (string, error) {
	if v, ok := fun(); ok {
		return v, nil
	}

	return "", ErrNotFound
}

func getInt(fun func() (string, bool)) (int, error) {
	v, ok := fun()

	if !ok {
		return 0, ErrNotFound
	}

	i, err := strconv.Atoi(v)

	if err != nil {
		return 0, ErrWrongType
	}

	return i, nil
}

func getFloat(fun func() (string, bool)) (float32, error) {
	v, ok := fun()

	if !ok {
		return 0, ErrNotFound
	}

	f, err := strconv.ParseFloat(v, 32)

	if err != nil {
		return 0, ErrWrongType
	}

	return float32(f), nil
}

func getBool(fun func() (string, bool)) (bool, error) {
	v, ok := fun()

	if !ok {
		return false, ErrNotFound
	}

	b, err := strconv.ParseBool(v)

	if err != nil {
		return false, ErrWrongType
	}

	return b, nil
}

// Env manage environment variables
// by giving high level api to interact with them
type Env struct {
	envs *map[string]string
}

// NewEnv creates a new Env instance
func NewEnv() Env {
	return Env{parseVars()}
}

// GetAllValues retrieves a slice of all environment variables values
func (e Env) GetAllValues() []string {
	results := []string{}

	for _, v := range *e.envs {
		results = append(results, v)
	}

	return results
}

// GetAllKeys retrieves a slice of all environment variables keys
func (e Env) GetAllKeys() []string {
	results := []string{}

	for k := range *e.envs {
		results = append(results, k)
	}

	return results
}

// GetString return a string if variable exists
// or an error otherwise
func (e Env) GetString(key string) (string, error) {
	return getString(func() (string, bool) {
		v, ok := (*e.envs)[key]

		return v, ok
	})
}

// GetInt return an integer if variable exists
// or an error if value is not an integer or doesn't exist
func (e Env) GetInt(key string) (int, error) {
	return getInt(func() (string, bool) {
		v, ok := (*e.envs)[key]

		return v, ok
	})
}

// GetFloat return a float if variable exists
// or an error if value is not a float or doesn't exist
func (e Env) GetFloat(key string) (float32, error) {
	return getFloat(func() (string, bool) {
		v, ok := (*e.envs)[key]

		return v, ok
	})
}

// GetBool return a boolean if variable exists
// or an error if value is not a boolean or doesn't exist
func (e Env) GetBool(key string) (bool, error) {
	return getBool(func() (string, bool) {
		v, ok := (*e.envs)[key]

		return v, ok
	})
}

// FindEntries retrieves all keys matching a given regexp and their
// corresponding values
func (e Env) FindEntries(reg string) (map[string]string, error) {
	results := map[string]string{}

	r, err := regexp.Compile(reg)

	if err != nil {
		return results, err
	}

	for k, v := range *e.envs {
		if r.MatchString(k) {
			results[k] = v
		}
	}

	return results, nil
}

type node struct {
	childs   []*node
	key      string
	value    string
	hasValue bool
	root     bool
}

func newNode() *node {
	return &node{childs: []*node{}}
}

func newRootNode() *node {
	return &node{childs: []*node{}, root: true}
}

func (n *node) findAllChildsByKey(key string, withValue bool) *[]*node {
	results := []*node{}
	nodes := n.childs

	for {
		tank := []*node{}

		for _, node := range nodes {
			if node.key == key {
				if withValue && node.hasValue || !withValue {
					results = append(results, node)
				}
			}

			tank = append(tank, node.childs...)
		}

		nodes = tank

		if len(tank) == 0 {
			return &results
		}
	}
}

func (n *node) findChildByKey(key string) (*node, bool) {
	for _, child := range n.childs {
		if child.key == key {
			return child, true
		}
	}

	return nil, false
}

func (n *node) appendChild(child *node) bool {
	if _, ok := n.findChildByKey(child.key); ok {
		return false
	}

	n.childs = append(n.childs, child)

	return true
}

func createTreeFromDelimiterFilteringByRegexp(reg *regexp.Regexp, delimiter string) (*node, error) {
	rootNode := newRootNode()

	for key, value := range *parseVars() {
		if reg.MatchString(key) {
			current := rootNode

			for _, component := range strings.Split(key, delimiter) {
				n, exists := current.findChildByKey(component)

				if exists {
					current = n
				} else {
					child := newNode()
					child.key = component
					current.appendChild(child)
					current = child
				}
			}

			if current.hasValue {
				return nil, ErrDuplicated
			}

			current.hasValue = true
			current.value = value
		}
	}

	return rootNode, nil
}

}
<|endoftext|>


<commit_before>package paza

import (
	"fmt"
	"regexp"
	"strconv"
	"sync/atomic"
	"unicode/utf8"
)

var (
	pt = fmt.Printf
)

type Set struct {
	parsers   map[string]Parser
	recursive map[string]bool
	serial    uint64
}

type Parser func(input *Input, start int) (ok bool, n int)

type stackEntry struct {
	parser string
	start  int
	ok     bool
	length int
}

type Input struct {
	text  []byte
	stack []stackEntry
}

func NewInput(text []byte) *Input {
	return &Input{
		text: text,
	}
}

func NewSet() *Set {
	return &Set{
		parsers:   make(map[string]Parser),
		recursive: make(map[string]bool),
	}
}

func (s *Set) Add(name string, parser Parser) {
	s.parsers[name] = parser
}

func (s *Set) AddRec(name string, parser Parser) {
	s.parsers[name] = parser
	s.recursive[name] = true
}

func (s *Set) Regex(re string) Parser {
	regex := regexp.MustCompile(re)
	return func(input *Input, start int) (bool, int) {
		if loc := regex.FindIndex(input.text[start:]); loc != nil && loc[0] == 0 {
			return true, loc[1]
		}
		return false, 0
	}
}

func (s *Set) Rune(r rune) Parser {
	return func(input *Input, start int) (bool, int) {
		ru, l := utf8.DecodeRune(input.text[start:])
		if ru == utf8.RuneError {
			return false, 0
		}
		if ru != r {
			return false, 0
		}
		return true, l
	}
}

func (s *Set) Call(name string, input *Input, start int) (bool, int) {
	fn, ok := s.parsers[name]
	if !ok {
		panic("parser not found " + name)
	}

	// non recursive parser
	if _, ok := s.recursive[name]; !ok {
		return fn(input, start)
	}

	// search stack
	for i := len(input.stack) - 1; i >= 0; i-- {
		mem := input.stack[i]
		if mem.parser == name && mem.start == start { // found
			return mem.ok, mem.length
		}
	}
	// not found, append a new entry
	entry := stackEntry{
		parser: name,
		start:  start,
		ok:     false,
		length: 0,
	}
	input.stack = append(input.stack, entry)
	// find the right bound
	for {
		stackSize := len(input.stack) // save stack size
		ok, l := fn(input, start)
		input.stack = input.stack[:stackSize] // unwind stack
		if !ok {
			return false, 0
		}
		if l < entry.length { // over bound
			return entry.ok, entry.length
		} else if l == entry.length { // not extending
			return ok, l
		}
		entry.ok = ok
		entry.length = l
		// update stack
		for i := len(input.stack) - 1; i >= 0; i-- {
			e := input.stack[i]
			if e.parser == name && e.start == start {
				input.stack[i].ok = ok
				input.stack[i].length = l
				break
			}
		}
	}

}

func (s *Set) getNames(parsers []interface{}) (ret []string) {
	for _, parser := range parsers {
		switch parser := parser.(type) {
		case string:
			ret = append(ret, parser)
		case Parser:
			name := "__parser__" + strconv.Itoa(int(atomic.AddUint64(&s.serial, 1)))
			s.Add(name, parser)
			ret = append(ret, name)
		default:
			panic(fmt.Sprintf("unknown parser type %T", parser))
		}
	}
	return
}

func (s *Set) Concat(parsers ...interface{}) Parser {
	names := s.getNames(parsers)
	return func(input *Input, start int) (bool, int) {
		index := start
		for _, name := range names {
			if ok, l := s.Call(name, input, index); !ok {
				return false, 0
			} else {
				index += l
			}
		}
		return true, index - start
	}
}

func (s *Set) OrdChoice(parsers ...interface{}) Parser {
	names := s.getNames(parsers)
	return func(input *Input, start int) (bool, int) {
		for _, name := range names {
			if ok, l := s.Call(name, input, start); ok {
				return ok, l
			}
		}
		return false, 0
	}
}
<commit_msg>optimize<commit_after>package paza

import (
	"fmt"
	"regexp"
	"strconv"
	"sync/atomic"
	"unicode/utf8"
)

var (
	pt = fmt.Printf
)

type parserInfo struct {
	parser    Parser
	recursive bool
}

type Set struct {
	parsers map[string]parserInfo
	serial  uint64
}

type Parser func(input *Input, start int) (ok bool, n int)

type stackEntry struct {
	parser string
	start  int
	ok     bool
	length int
}

type Input struct {
	text  []byte
	stack []stackEntry
}

func NewInput(text []byte) *Input {
	return &Input{
		text: text,
	}
}

func NewSet() *Set {
	return &Set{
		parsers: make(map[string]parserInfo),
	}
}

func (s *Set) Add(name string, parser Parser) {
	s.parsers[name] = parserInfo{parser, false}
}

func (s *Set) AddRec(name string, parser Parser) {
	s.parsers[name] = parserInfo{parser, true}
}

func (s *Set) Regex(re string) Parser {
	regex := regexp.MustCompile(re)
	return func(input *Input, start int) (bool, int) {
		if loc := regex.FindIndex(input.text[start:]); loc != nil && loc[0] == 0 {
			return true, loc[1]
		}
		return false, 0
	}
}

func (s *Set) Rune(r rune) Parser {
	return func(input *Input, start int) (bool, int) {
		ru, l := utf8.DecodeRune(input.text[start:])
		if ru == utf8.RuneError {
			return false, 0
		}
		if ru != r {
			return false, 0
		}
		return true, l
	}
}

func (s *Set) Call(name string, input *Input, start int) (bool, int) {
	info, ok := s.parsers[name]
	if !ok {
		panic("parser not found " + name)
	}

	// non recursive parser
	if !info.recursive {
		return info.parser(input, start)
	}

	// search stack
	for i := len(input.stack) - 1; i >= 0; i-- {
		mem := input.stack[i]
		if mem.parser == name && mem.start == start { // found
			return mem.ok, mem.length
		}
	}
	// not found, append a new entry
	input.stack = append(input.stack, stackEntry{
		parser: name,
		start:  start,
		ok:     false,
		length: 0,
	})
	// find the right bound
	lastOk := false
	lastLen := 0
	for {
		stackSize := len(input.stack) // save stack size
		ok, l := info.parser(input, start)
		input.stack = input.stack[:stackSize] // unwind stack
		if !ok {
			return false, 0
		}
		if l < lastLen { // over bound
			return lastOk, lastLen
		} else if l == lastLen { // not extending
			return ok, l
		}
		lastOk = ok
		lastLen = l
		// update stack
		for i := len(input.stack) - 1; i >= 0; i-- {
			e := input.stack[i]
			if e.parser == name && e.start == start {
				input.stack[i].ok = ok
				input.stack[i].length = l
				break
			}
		}
	}

}

func (s *Set) getNames(parsers []interface{}) (ret []string) {
	for _, parser := range parsers {
		switch parser := parser.(type) {
		case string:
			ret = append(ret, parser)
		case Parser:
			name := "__parser__" + strconv.Itoa(int(atomic.AddUint64(&s.serial, 1)))
			s.Add(name, parser)
			ret = append(ret, name)
		default:
			panic(fmt.Sprintf("unknown parser type %T", parser))
		}
	}
	return
}

func (s *Set) Concat(parsers ...interface{}) Parser {
	names := s.getNames(parsers)
	return func(input *Input, start int) (bool, int) {
		index := start
		for _, name := range names {
			if ok, l := s.Call(name, input, index); !ok {
				return false, 0
			} else {
				index += l
			}
		}
		return true, index - start
	}
}

func (s *Set) OrdChoice(parsers ...interface{}) Parser {
	names := s.getNames(parsers)
	return func(input *Input, start int) (bool, int) {
		for _, name := range names {
			if ok, l := s.Call(name, input, start); ok {
				return ok, l
			}
		}
		return false, 0
	}
}
<|endoftext|>


<commit_before>// Package pcap is a wrapper around the pcap library.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t * p, struct pcap_pkthdr **pkt_header,
		      u_char ** pkt_data)
{
	return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"

import (
	"errors"
	"fmt"
	"net"
	"syscall"
	"time"
	"unsafe"
)

type Pcap struct {
	cptr *C.pcap_t
}

type pcapError struct{ string }

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

func Version() string               { return C.GoString(C.pcap_lib_version()) }
func (p *Pcap) Datalink() int       { return int(C.pcap_datalink(p.cptr)) }
func (e *pcapError) Error() string  { return e.string }
func (p *Pcap) Geterror() error     { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

func Create(device string) (*Pcap, error) {
	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	cptr := C.pcap_create(dev, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr: cptr,
	}, nil
}

// Set buffer size (units in bytes) on activated handle.
func (p *Pcap) SetBufferSize(sz int32) error {
	if C.pcap_set_buffer_size(p.cptr, C.int(sz)) != 0 {
		return p.Geterror()
	}
	return nil
}

// If arg p is non-zero promiscuous mode will be set on capture handle when it
// is activated.
func (p *Pcap) SetPromisc(promisc bool) error {
	pro := int32(0)
	if promisc {
		pro = 1
	}

	if C.pcap_set_promisc(p.cptr, C.int(pro)) != 0 {
		return p.Geterror()
	}
	return nil
}

func (p *Pcap) SetSnapLen(s int32) error {
	if C.pcap_set_snaplen(p.cptr, C.int(s)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Set read timeout (milliseconds) that will be used on a capture handle when it
// is activated.
func (p *Pcap) SetReadTimeout(toMs int32) error {
	if C.pcap_set_timeout(p.cptr, C.int(toMs)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Activate a packet capture handle to look at packets on the network, with the
// options that were set on the handle being in effect.
func (p *Pcap) Activate() error {
	if C.pcap_activate(p.cptr) != 0 {
		return p.Geterror()
	}
	return nil
}

// OpenLive opens a device and returns a handler.
func OpenLive(device string, snaplen int32, promisc bool, timeout_ms int32) (*Pcap, error) {
	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	pro := int32(0)
	if promisc {
		pro = 1
	}

	cptr := C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr: cptr,
	}, nil
}

// OpenOffline provides a Pcap over a .pcap file.
func OpenOffline(file string) (*Pcap, error) {
	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	cptr := C.pcap_open_offline(cf, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr: cptr,
	}, nil
}

// Pcap closes a handler.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) NextEx() (*Packet, int32) {
	var pkthdr *C.struct_pcap_pkthdr
	var bufPtr *C.u_char
	result := int32(C.hack_pcap_next_ex(p.cptr, &pkthdr, &bufPtr))
	buf := unsafe.Pointer(bufPtr)
	if buf == nil {
		return nil, result
	}

	pkt := &Packet{
		Time:   time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)),
		Caplen: uint32(pkthdr.caplen),
		Len:    uint32(pkthdr.len),
		Data:   make([]byte, pkthdr.caplen),
	}
	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return pkt, result
}

func (p *Pcap) Getstats() (*Stat, error) {
	var cstats _Ctype_struct_pcap_stat
	if C.pcap_stats(p.cptr, &cstats) == -1 {
		return nil, p.Geterror()
	}

	return &Stat{
		PacketsReceived:  uint32(cstats.ps_recv),
		PacketsDropped:   uint32(cstats.ps_drop),
		PacketsIfDropped: uint32(cstats.ps_ifdrop),
	}, nil
}

func (p *Pcap) SetFilter(expr string) error {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) == -1 {
		return p.Geterror()
	}
	defer C.pcap_freecode(&bpf)

	if C.pcap_setfilter(p.cptr, &bpf) == -1 {
		return p.Geterror()
	}
	return nil
}

func (p *Pcap) SetDataLink(dlt int) error {
	if -1 == C.pcap_set_datalink(p.cptr, C.int(dlt)) {
		return p.Geterror()
	}
	return nil
}

func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

func FindAllDevs() ([]Interface, error) {
	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	var alldevsp *C.pcap_if_t
	if C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) == -1 {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))

	ifs := []Interface{}
	for dev := alldevsp; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		ifs = append(ifs, Interface{
			Name:        C.GoString(dev.name),
			Description: C.GoString(dev.description),
			Addresses:   findAllAddresses(dev.addresses),
			// TODO: add more elements
		})
	}
	return ifs, nil
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) []IFAddress {
	// TODO - make it support more than IPv4 and IPv6?
	a := []IFAddress{}
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}

		ip, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}
		netmask, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}

		a = append(a, IFAddress{
			IP:      ip,
			Netmask: netmask,
		})
	}
	return a
}

func sockaddrToIP(rsa *syscall.RawSockaddr) ([]byte, error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		ip := make([]byte, 4)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		ip := make([]byte, 16)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	default:
		return nil, fmt.Errorf("unsupported address type %d", rsa.Family)
	}
}

// Inject TODO
func (p *Pcap) Inject(data []byte) error {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))
	defer C.free(unsafe.Pointer(buf))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) == -1 {
		return p.Geterror()
	}
	return nil
}
<commit_msg>Missed a few fixes in pcap.go<commit_after>// Package pcap is a wrapper around the pcap library.
package pcap

/*
#cgo LDFLAGS: -lpcap
#include <stdlib.h>
#include <pcap.h>

// Workaround for not knowing how to cast to const u_char**
int hack_pcap_next_ex(pcap_t * p, struct pcap_pkthdr **pkt_header,
		      u_char ** pkt_data)
{
	return pcap_next_ex(p, pkt_header, (const u_char **)pkt_data);
}
*/
import "C"

import (
	"errors"
	"fmt"
	"net"
	"syscall"
	"time"
	"unsafe"
)

type Pcap struct {
	cptr *C.pcap_t
}

type pcapError struct{ string }

type Stat struct {
	PacketsReceived  uint32
	PacketsDropped   uint32
	PacketsIfDropped uint32
}

type Interface struct {
	Name        string
	Description string
	Addresses   []IFAddress
	// TODO: add more elements
}

type IFAddress struct {
	IP      net.IP
	Netmask net.IPMask
	// TODO: add broadcast + PtP dst ?
}

func Version() string               { return C.GoString(C.pcap_lib_version()) }
func (p *Pcap) Datalink() int       { return int(C.pcap_datalink(p.cptr)) }
func (e *pcapError) Error() string  { return e.string }
func (p *Pcap) Geterror() error     { return &pcapError{C.GoString(C.pcap_geterr(p.cptr))} }
func (p *Pcap) Next() (pkt *Packet) { rv, _ := p.NextEx(); return rv }

// Create TODO
func Create(device string) (*Pcap, error) {
	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	cptr := C.pcap_create(dev, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr: cptr,
	}, nil
}

// SetBufferSize sets buffer size (in bytes) on the activated handle.
func (p *Pcap) SetBufferSize(sz int32) error {
	if C.pcap_set_buffer_size(p.cptr, C.int(sz)) != 0 {
		return p.Geterror()
	}
	return nil
}

// SetPromisc sets promiscuous mode on the handle. It should be called before
// activation.
func (p *Pcap) SetPromisc(promisc bool) error {
	pro := int32(0)
	if promisc {
		pro = 1
	}

	if C.pcap_set_promisc(p.cptr, C.int(pro)) != 0 {
		return p.Geterror()
	}
	return nil
}

func (p *Pcap) SetSnapLen(s int32) error {
	if C.pcap_set_snaplen(p.cptr, C.int(s)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Set read timeout (milliseconds) that will be used on a capture handle when it
// is activated.
func (p *Pcap) SetReadTimeout(toMs int32) error {
	if C.pcap_set_timeout(p.cptr, C.int(toMs)) != 0 {
		return p.Geterror()
	}
	return nil
}

// Activate a packet capture handle to look at packets on the network, with the
// options that were set on the handle being in effect.
func (p *Pcap) Activate() error {
	if C.pcap_activate(p.cptr) != 0 {
		return p.Geterror()
	}
	return nil
}

// OpenLive opens a device and returns a handler.
func OpenLive(device string, snaplen int32, promisc bool, timeout_ms int32) (*Pcap, error) {
	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	dev := C.CString(device)
	defer C.free(unsafe.Pointer(dev))

	pro := int32(0)
	if promisc {
		pro = 1
	}

	cptr := C.pcap_open_live(dev, C.int(snaplen), C.int(pro), C.int(timeout_ms), buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr: cptr,
	}, nil
}

// OpenOffline provides a Pcap over a .pcap file.
func OpenOffline(file string) (*Pcap, error) {
	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	cf := C.CString(file)
	defer C.free(unsafe.Pointer(cf))

	cptr := C.pcap_open_offline(cf, buf)
	if cptr == nil {
		return nil, &pcapError{C.GoString(buf)}
	}

	return &Pcap{
		cptr: cptr,
	}, nil
}

// Pcap closes a handler.
func (p *Pcap) Close() {
	C.pcap_close(p.cptr)
}

func (p *Pcap) NextEx() (*Packet, int32) {
	var pkthdr *C.struct_pcap_pkthdr
	var bufPtr *C.u_char
	result := int32(C.hack_pcap_next_ex(p.cptr, &pkthdr, &bufPtr))
	buf := unsafe.Pointer(bufPtr)
	if buf == nil {
		return nil, result
	}

	pkt := &Packet{
		Time:   time.Unix(int64(pkthdr.ts.tv_sec), int64(pkthdr.ts.tv_usec)),
		Caplen: uint32(pkthdr.caplen),
		Len:    uint32(pkthdr.len),
		Data:   make([]byte, pkthdr.caplen),
	}
	for i := uint32(0); i < pkt.Caplen; i++ {
		pkt.Data[i] = *(*byte)(unsafe.Pointer(uintptr(buf) + uintptr(i)))
	}
	return pkt, result
}

func (p *Pcap) Getstats() (*Stat, error) {
	var cstats _Ctype_struct_pcap_stat
	if C.pcap_stats(p.cptr, &cstats) == -1 {
		return nil, p.Geterror()
	}

	return &Stat{
		PacketsReceived:  uint32(cstats.ps_recv),
		PacketsDropped:   uint32(cstats.ps_drop),
		PacketsIfDropped: uint32(cstats.ps_ifdrop),
	}, nil
}

func (p *Pcap) SetFilter(expr string) error {
	var bpf _Ctype_struct_bpf_program
	cexpr := C.CString(expr)
	defer C.free(unsafe.Pointer(cexpr))

	if C.pcap_compile(p.cptr, &bpf, cexpr, 1, 0) == -1 {
		return p.Geterror()
	}
	defer C.pcap_freecode(&bpf)

	if C.pcap_setfilter(p.cptr, &bpf) == -1 {
		return p.Geterror()
	}
	return nil
}

func (p *Pcap) SetDataLink(dlt int) error {
	if -1 == C.pcap_set_datalink(p.cptr, C.int(dlt)) {
		return p.Geterror()
	}
	return nil
}

func DatalinkValueToName(dlt int) string {
	if name := C.pcap_datalink_val_to_name(C.int(dlt)); name != nil {
		return C.GoString(name)
	}
	return ""
}

func DatalinkValueToDescription(dlt int) string {
	if desc := C.pcap_datalink_val_to_description(C.int(dlt)); desc != nil {
		return C.GoString(desc)
	}
	return ""
}

func FindAllDevs() ([]Interface, error) {
	buf := (*C.char)(C.calloc(ERRBUF_SIZE, 1))
	defer C.free(unsafe.Pointer(buf))

	var alldevsp *C.pcap_if_t
	if C.pcap_findalldevs((**C.pcap_if_t)(&alldevsp), buf) == -1 {
		return nil, errors.New(C.GoString(buf))
	}
	defer C.pcap_freealldevs((*C.pcap_if_t)(alldevsp))

	ifs := []Interface{}
	for dev := alldevsp; dev != nil; dev = (*C.pcap_if_t)(dev.next) {
		ifs = append(ifs, Interface{
			Name:        C.GoString(dev.name),
			Description: C.GoString(dev.description),
			Addresses:   findAllAddresses(dev.addresses),
			// TODO: add more elements
		})
	}
	return ifs, nil
}

func findAllAddresses(addresses *_Ctype_struct_pcap_addr) []IFAddress {
	// TODO - make it support more than IPv4 and IPv6?
	a := []IFAddress{}
	for curaddr := addresses; curaddr != nil; curaddr = (*_Ctype_struct_pcap_addr)(curaddr.next) {
		if curaddr.addr == nil {
			continue
		}

		ip, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}
		netmask, err := sockaddrToIP((*syscall.RawSockaddr)(unsafe.Pointer(curaddr.addr)))
		if err != nil {
			continue
		}

		a = append(a, IFAddress{
			IP:      ip,
			Netmask: netmask,
		})
	}
	return a
}

func sockaddrToIP(rsa *syscall.RawSockaddr) ([]byte, error) {
	switch rsa.Family {
	case syscall.AF_INET:
		pp := (*syscall.RawSockaddrInet4)(unsafe.Pointer(rsa))
		ip := make([]byte, 4)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	case syscall.AF_INET6:
		pp := (*syscall.RawSockaddrInet6)(unsafe.Pointer(rsa))
		ip := make([]byte, 16)
		for i := 0; i < len(ip); i++ {
			ip[i] = pp.Addr[i]
		}
		return ip, nil

	default:
		return nil, fmt.Errorf("unsupported address type %d", rsa.Family)
	}
}

// Inject TODO
func (p *Pcap) Inject(data []byte) error {
	buf := (*C.char)(C.malloc((C.size_t)(len(data))))
	defer C.free(unsafe.Pointer(buf))

	for i := 0; i < len(data); i++ {
		*(*byte)(unsafe.Pointer(uintptr(unsafe.Pointer(buf)) + uintptr(i))) = data[i]
	}

	if C.pcap_inject(p.cptr, unsafe.Pointer(buf), (C.size_t)(len(data))) == -1 {
		return p.Geterror()
	}
	return nil
}
<|endoftext|>


<commit_before>package xlattice_go

import "errors"

/**
 * A Peer is another Node, a neighbor.
 *
 * XXX As this has evolved, it begins to look like Node should
 * XXX simply extend Peer.
 *
 * @author Jim Dixon
 */

type Peer struct {
    nodeID      *NodeID
    pubkey      *PublicKey
    overlays    []*Overlay
    connectors  []*Connector
}

func NewNewPeer(id *NodeID) (*Peer, error){
    return NewPeer(id, nil, nil, nil)
}

func NewPeer(id *NodeID, k *PublicKey, o *[]*Overlay,
                                       c *[]*Connector) (*Peer, error) {

    // IDENTITY /////////////////////////////////////////////////////
    if id == nil {
        err := errors.New("IllegalArgument: nil NodeID")
        return nil, err
    }
    nodeID := (*id).Clone()
    pubkey := k
    var overlays []*Overlay       // an empty slice
    if o != nil {
        count := len(*o)
        for i := 0; i < count; i++ {
            overlays = append(overlays, (*o)[i])
        }
    }   // FOO
    var ctors []*Connector        // another empty slice
    if c != nil {
        count := len(*c)
        for i := 0; i < count; i++ {
            ctors = append(ctors, (*c)[i])
        }
    }   // FOO
    p := new(Peer)
    p.nodeID = nodeID           // the clone
    p.pubkey = pubkey
    p.overlays = overlays
    p.connectors = ctors
    return p, nil
}
func (p *Peer) GetNodeID() *NodeID {
    return p.nodeID
}
func (p *Peer) GetPublicKey() *PublicKey {
    return p.pubkey
}
func (p *Peer) SetPublicKey(k *PublicKey) error {
    err := error(nil)
    if k == nil {
        err = errors.New("IllegalArgument: nil PublicKey")
    } else {
        p.pubkey = k
    }
    return err
}
// OVERLAYS /////////////////////////////////////////////////////////
func (p *Peer) addOverlay (o *Overlay) error {
    if o == nil  {
        return errors.New("IllegalArgument: nil Overlay")
    }
    p.overlays = append(p.overlays, o)
    return nil
}
/**
 * @return a count of the number of overlays the peer can be
 *         accessed through
 */
func(p *Peer)sizeOverlays () int {
    return len(p.overlays)
}
/** @return how to access the peer (transport, protocol, address) */
func(p *Peer) GetOverlay(n int ) *Overlay {
    return p.overlays[n]
}

// CONNECTORS ///////////////////////////////////////////////////////
func (p *Peer) addConnector (c *Connector) error {
    if c == nil {
        return errors.New("IllegalArgument: nil Connector")
    }
    p.connectors = append(p.connectors, c)
    return nil
}
/** @return a count of known Connectors for this Peer */
func(p *Peer) SizeConnectors() int {
    return len(p.connectors)
}
/**
 * Return a Connector, an Address-Protocol pair identifying
 * an Acceptor for the Peer.  Connectors are arranged in order
 * of preference, with the zero-th Connector being the most
 * preferred.
 *
 * XXX Could as easily return an EndPoint.
 *
 * @return the Nth Connector
 */
func(p *Peer) GetConnector(n int) *Connector {
    return p.connectors[n]
}

// EQUAL ////////////////////////////////////////////////////////////
func (p *Peer) Equal(any interface{}) bool {
    if any == p   { return true }
    if any == nil { return false }
    switch v := any.(type) {
    case Peer:
        _ = v;
    default:
        return false
    }
    other := any.(Peer)         // type assertion
    // THINK ABOUT publicKey.equals(any.publicKey)
    if p.nodeID == other.nodeID                 { return true }
    if p.nodeID.Length() != other.nodeID.Length()     { return false }
    myVal := p.nodeID.Value()
    otherVal := other.nodeID.Value()
    for i := 0; i < p.nodeID.Length(); i++ {
        if myVal[i] != otherVal[i] { return false }
    }
    return false
}
<commit_msg>adding Node abstraction<commit_after>package xlattice_go

import "errors"

/**
 * A Peer is another Node, a neighbor.
 */

type Peer struct {
    nodeID      *NodeID
    pubkey      *PublicKey
    overlays    []*Overlay
    connectors  []*Connector
}

func NewNewPeer(id *NodeID) (*Peer, error){
    return NewPeer(id, nil, nil, nil)
}

func NewPeer(id *NodeID, k *PublicKey, o *[]*Overlay,
                                       c *[]*Connector) (*Peer, error) {

    // IDENTITY /////////////////////////////////////////////////////
    if id == nil {
        err := errors.New("IllegalArgument: nil NodeID")
        return nil, err
    }
    nodeID := (*id).Clone()
    pubkey := k
    var overlays []*Overlay       // an empty slice
    if o != nil {
        count := len(*o)
        for i := 0; i < count; i++ {
            overlays = append(overlays, (*o)[i])
        }
    }   // FOO
    var ctors []*Connector        // another empty slice
    if c != nil {
        count := len(*c)
        for i := 0; i < count; i++ {
            ctors = append(ctors, (*c)[i])
        }
    }   // FOO
    p := new(Peer)
    p.nodeID = nodeID           // the clone
    p.pubkey = pubkey
    p.overlays = overlays
    p.connectors = ctors
    return p, nil
}
func (p *Peer) GetNodeID() *NodeID {
    return p.nodeID
}
func (p *Peer) GetPublicKey() *PublicKey {
    return p.pubkey
}
func (p *Peer) SetPublicKey(k *PublicKey) error {
    err := error(nil)
    if k == nil {
        err = errors.New("IllegalArgument: nil PublicKey")
    } else {
        p.pubkey = k
    }
    return err
}
// OVERLAYS /////////////////////////////////////////////////////////
func (p *Peer) addOverlay (o *Overlay) error {
    if o == nil  {
        return errors.New("IllegalArgument: nil Overlay")
    }
    p.overlays = append(p.overlays, o)
    return nil
}
/**
 * @return a count of the number of overlays the peer can be
 *         accessed through
 */
func(p *Peer)sizeOverlays () int {
    return len(p.overlays)
}
/** @return how to access the peer (transport, protocol, address) */
func(p *Peer) GetOverlay(n int ) *Overlay {
    return p.overlays[n]
}

// CONNECTORS ///////////////////////////////////////////////////////
func (p *Peer) addConnector (c *Connector) error {
    if c == nil {
        return errors.New("IllegalArgument: nil Connector")
    }
    p.connectors = append(p.connectors, c)
    return nil
}
/** @return a count of known Connectors for this Peer */
func(p *Peer) SizeConnectors() int {
    return len(p.connectors)
}
/**
 * Return a Connector, an Address-Protocol pair identifying
 * an Acceptor for the Peer.  Connectors are arranged in order
 * of preference, with the zero-th Connector being the most
 * preferred.
 *
 * XXX Could as easily return an EndPoint.
 *
 * @return the Nth Connector
 */
func(p *Peer) GetConnector(n int) *Connector {
    return p.connectors[n]
}

// EQUAL ////////////////////////////////////////////////////////////
func (p *Peer) Equal(any interface{}) bool {
    if any == p   { return true }
    if any == nil { return false }
    switch v := any.(type) {
    case Peer:
        _ = v;
    default:
        return false
    }
    other := any.(Peer)         // type assertion
    // THINK ABOUT publicKey.equals(any.publicKey)
    if p.nodeID == other.nodeID                 { return true }
    if p.nodeID.Length() != other.nodeID.Length()     { return false }
    myVal := p.nodeID.Value()
    otherVal := other.nodeID.Value()
    for i := 0; i < p.nodeID.Length(); i++ {
        if myVal[i] != otherVal[i] { return false }
    }
    return false
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
)

// Instances of Msg store the transmited message and its source peer data
type Msg struct {
	Message string
	From    Peer
}

// Peer stores screen name and ip address data
type Peer struct {
	Name    string
	Address string
}

// Peers provides a map of known Peers accessed by an Address
type Peers map[string]Peer

// PeerData implements the p2p control data for a peer
type PeerData struct {
	Self            Peer
	Peers           Peers
	addPeers        chan (Peer)
	exitPeer        chan (Peer)
	currentPeers    chan (Peers)
	getCurrentPeers chan (bool)
	receivedMsg     chan (Msg)
	userMsg         chan (Msg)
}

// Startup client logic
func main() {
        //Parse comandline argumments
	port := flag.String("p", "7500", "Port number")
	name := flag.String("n", "Surematics", "ScreenName")
	peer := flag.String("i", "", "IP address of peer to connect to.")
	flag.Parse()

        //initialize client control handler
	peercontroller := NewPeerData(Peer{*name, getIp() + ":" + *port})
	peercontroller.start()

	if len(*peer) != 0 {
		peercontroller.addPeers <- Peer{"", *peer}
	}

        //scan stdin for input
	peercontroller.startStdinListener(peercontroller.Self)
}

// Get IP of this machine
func getIp() string {
    conn, err := net.Dial("udp", "8.8.8.8:80")
    if err != nil {
        return "localhost"
    }
    defer conn.Close()

    localAddr := conn.LocalAddr().String()
    idx := strings.LastIndex(localAddr, ":")

    return localAddr[0:idx]
}

// NewPeerData instantiate, initialize, and return PeerData control data
func NewPeerData(self Peer) *PeerData {
	peercontroller := new(PeerData)
	peercontroller.Self = self
	peercontroller.Peers = make(Peers)
	peercontroller.addPeers = make(chan (Peer))
	peercontroller.currentPeers = make(chan (Peers))
	peercontroller.getCurrentPeers = make(chan (bool))

	peercontroller.userMsg = make(chan (Msg))
	peercontroller.receivedMsg = make(chan (Msg))
	return peercontroller
}

// start api and command processing threads
func (peercontroller *PeerData) start() {
	go peercontroller.cmdLoop()
	go peercontroller.webListener()
	fmt.Printf("Peer node started on %s \n", peercontroller.Self.Address)
}

//Process detected events
func (peercontroller *PeerData) cmdLoop() {
	for {
		select {
		case peer := <-peercontroller.addPeers:
			if !peercontroller.existingPeer(peer) {
				peercontroller.Peers[peer.Address] = peer
				go peercontroller.sendAdd(peer)
				fmt.Printf("Connecting to: %s \n", peer.Address)
			}

		case <-peercontroller.getCurrentPeers:
			peercontroller.currentPeers <- peercontroller.Peers

		case peer := <-peercontroller.exitPeer:
			delete(peercontroller.Peers, peer.Address)

		case Msg := <-peercontroller.receivedMsg:
			fmt.Printf("%s: %s\n", Msg.From.Name, Msg.Message)

		case Msg := <-peercontroller.userMsg:
			for _, peer := range peercontroller.Peers {
				go peercontroller.sendTx(peer, Msg)
			}
		}
	}
}

func (peercontroller *PeerData) existingPeer(peer Peer) bool {
	if peer.Address == peercontroller.Self.Address {
		return true
	}
	_, existingPeer := peercontroller.Peers[peer.Address]
	return existingPeer
}

//connect to other peer nodes
func (peercontroller *PeerData) sendAdd(peer Peer) {
	URL := "http://" + peer.Address + "/add"

	j, _ := json.Marshal(peercontroller.Self)

	resp, err := http.Post(URL, "application/json", bytes.NewBuffer(j))
	if err != nil {
		peercontroller.exitPeer <- peer
		return
	}

	peercontroller.addPeers <- peer

	defer resp.Body.Close()
	otherPeers := Peers{}
	decode := json.NewDecoder(resp.Body)
	err = decode.Decode(&otherPeers)
	for _, peer := range otherPeers {
		peercontroller.addPeers <- peer
	}
}

//send message to peer
func (peercontroller *PeerData) sendTx(peer Peer, msg Msg) {
	URL := "http://" + peer.Address + "/tx"

	j, _ := json.Marshal(msg)

	_, err := http.Post(URL, "application/json", bytes.NewBuffer(j))

	if err != nil {
		peercontroller.exitPeer <- peer
		return
	}
}

//start api handlers to allow interaction with peers
func (peercontroller *PeerData) webListener() {
	http.HandleFunc("/tx", createTxHandler(peercontroller))
	http.HandleFunc("/add", createAddHandler(peercontroller))
	log.Fatal(http.ListenAndServe(peercontroller.Self.Address, nil))
}

//Implement add api method to update peercontroller peer data
func createAddHandler(peercontroller *PeerData) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		addingpeer := Peer{}
		decode := json.NewDecoder(r.Body)
		err := decode.Decode(&addingpeer)

		if err != nil {
			log.Printf("Add Peer Error: %v", err)
		}

		peercontroller.addPeers <- addingpeer
		peercontroller.getCurrentPeers <- true

		encode := json.NewEncoder(w)
		encode.Encode(<-peercontroller.currentPeers)
	}
}

//Implement Tx api method to update peercontroller msg data
func createTxHandler(peercontroller *PeerData) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		mg := Msg{}
		decode := json.NewDecoder(r.Body)
		err := decode.Decode(&mg)
		if err != nil {
			log.Printf("Message Error: %v", err)
		}

		peercontroller.receivedMsg <- mg
	}
}

//Parse stdin for text and update peercontroller msg data
func (peercontroller *PeerData) startStdinListener(sender Peer) {
	reader := bufio.NewReader(os.Stdin)

	for {
		line, _ := reader.ReadString('\n')
		message := line[:len(line)-1]
		peercontroller.userMsg <- Msg{message, sender}
	}
}
<commit_msg>add package<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
        "strings"
	"os"
)

// Instances of Msg store the transmited message and its source peer data
type Msg struct {
	Message string
	From    Peer
}

// Peer stores screen name and ip address data
type Peer struct {
	Name    string
	Address string
}

// Peers provides a map of known Peers accessed by an Address
type Peers map[string]Peer

// PeerData implements the p2p control data for a peer
type PeerData struct {
	Self            Peer
	Peers           Peers
	addPeers        chan (Peer)
	exitPeer        chan (Peer)
	currentPeers    chan (Peers)
	getCurrentPeers chan (bool)
	receivedMsg     chan (Msg)
	userMsg         chan (Msg)
}

// Startup client logic
func main() {
        //Parse comandline argumments
	port := flag.String("p", "7500", "Port number")
	name := flag.String("n", "Surematics", "ScreenName")
	peer := flag.String("i", "", "IP address of peer to connect to.")
	flag.Parse()

        //initialize client control handler
	peercontroller := NewPeerData(Peer{*name, getIp() + ":" + *port})
	peercontroller.start()

	if len(*peer) != 0 {
		peercontroller.addPeers <- Peer{"", *peer}
	}

        //scan stdin for input
	peercontroller.startStdinListener(peercontroller.Self)
}

// Get IP of this machine
func getIp() string {
    conn, err := net.Dial("udp", "8.8.8.8:80")
    if err != nil {
        return "localhost"
    }
    defer conn.Close()

    localAddr := conn.LocalAddr().String()
    idx := strings.LastIndex(localAddr, ":")

    return localAddr[0:idx]
}

// NewPeerData instantiate, initialize, and return PeerData control data
func NewPeerData(self Peer) *PeerData {
	peercontroller := new(PeerData)
	peercontroller.Self = self
	peercontroller.Peers = make(Peers)
	peercontroller.addPeers = make(chan (Peer))
	peercontroller.currentPeers = make(chan (Peers))
	peercontroller.getCurrentPeers = make(chan (bool))

	peercontroller.userMsg = make(chan (Msg))
	peercontroller.receivedMsg = make(chan (Msg))
	return peercontroller
}

// start api and command processing threads
func (peercontroller *PeerData) start() {
	go peercontroller.cmdLoop()
	go peercontroller.webListener()
	fmt.Printf("Peer node started on %s \n", peercontroller.Self.Address)
}

//Process detected events
func (peercontroller *PeerData) cmdLoop() {
	for {
		select {
		case peer := <-peercontroller.addPeers:
			if !peercontroller.existingPeer(peer) {
				peercontroller.Peers[peer.Address] = peer
				go peercontroller.sendAdd(peer)
				fmt.Printf("Connecting to: %s \n", peer.Address)
			}

		case <-peercontroller.getCurrentPeers:
			peercontroller.currentPeers <- peercontroller.Peers

		case peer := <-peercontroller.exitPeer:
			delete(peercontroller.Peers, peer.Address)

		case Msg := <-peercontroller.receivedMsg:
			fmt.Printf("%s: %s\n", Msg.From.Name, Msg.Message)

		case Msg := <-peercontroller.userMsg:
			for _, peer := range peercontroller.Peers {
				go peercontroller.sendTx(peer, Msg)
			}
		}
	}
}

func (peercontroller *PeerData) existingPeer(peer Peer) bool {
	if peer.Address == peercontroller.Self.Address {
		return true
	}
	_, existingPeer := peercontroller.Peers[peer.Address]
	return existingPeer
}

//connect to other peer nodes
func (peercontroller *PeerData) sendAdd(peer Peer) {
	URL := "http://" + peer.Address + "/add"

	j, _ := json.Marshal(peercontroller.Self)

	resp, err := http.Post(URL, "application/json", bytes.NewBuffer(j))
	if err != nil {
		peercontroller.exitPeer <- peer
		return
	}

	peercontroller.addPeers <- peer

	defer resp.Body.Close()
	otherPeers := Peers{}
	decode := json.NewDecoder(resp.Body)
	err = decode.Decode(&otherPeers)
	for _, peer := range otherPeers {
		peercontroller.addPeers <- peer
	}
}

//send message to peer
func (peercontroller *PeerData) sendTx(peer Peer, msg Msg) {
	URL := "http://" + peer.Address + "/tx"

	j, _ := json.Marshal(msg)

	_, err := http.Post(URL, "application/json", bytes.NewBuffer(j))

	if err != nil {
		peercontroller.exitPeer <- peer
		return
	}
}

//start api handlers to allow interaction with peers
func (peercontroller *PeerData) webListener() {
	http.HandleFunc("/tx", createTxHandler(peercontroller))
	http.HandleFunc("/add", createAddHandler(peercontroller))
	log.Fatal(http.ListenAndServe(peercontroller.Self.Address, nil))
}

//Implement add api method to update peercontroller peer data
func createAddHandler(peercontroller *PeerData) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		addingpeer := Peer{}
		decode := json.NewDecoder(r.Body)
		err := decode.Decode(&addingpeer)

		if err != nil {
			log.Printf("Add Peer Error: %v", err)
		}

		peercontroller.addPeers <- addingpeer
		peercontroller.getCurrentPeers <- true

		encode := json.NewEncoder(w)
		encode.Encode(<-peercontroller.currentPeers)
	}
}

//Implement Tx api method to update peercontroller msg data
func createTxHandler(peercontroller *PeerData) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		mg := Msg{}
		decode := json.NewDecoder(r.Body)
		err := decode.Decode(&mg)
		if err != nil {
			log.Printf("Message Error: %v", err)
		}

		peercontroller.receivedMsg <- mg
	}
}

//Parse stdin for text and update peercontroller msg data
func (peercontroller *PeerData) startStdinListener(sender Peer) {
	reader := bufio.NewReader(os.Stdin)

	for {
		line, _ := reader.ReadString('\n')
		message := line[:len(line)-1]
		peercontroller.userMsg <- Msg{message, sender}
	}
}
<|endoftext|>


<commit_before>package mpawsrds

import (
	mp "github.com/mackerelio/go-mackerel-plugin"
)

func (p RDSPlugin) rdsMetrics() (metrics []string) {
	for _, v := range p.GraphDefinition() {
		for _, vv := range v.Metrics {
			metrics = append(metrics, vv.Name)
		}
	}
	return
}

func mergeGraphDefs(a, b map[string]mp.Graphs) map[string]mp.Graphs {
	for k, v := range b {
		a[k] = v
	}
	return a
}

func (p RDSPlugin) mySQLGraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		p.Prefix + ".BinLogDiskUsage": {
			Label: p.LabelPrefix + " BinLogDiskUsage",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "BinLogDiskUsage", Label: "Usage"},
			},
		},
	}
}

func (p RDSPlugin) postgreSQLGraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		p.Prefix + ".MaximumUsedTransactionIDs": {
			Label: p.LabelPrefix + " Maximum Used Transaction IDs",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "MaximumUsedTransactionIDs", Label: "IDs"},
			},
		},
		p.Prefix + ".OldestReplicationSlotLag": {
			Label: p.LabelPrefix + " Oldest Replication Slot Lag",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "OldestReplicationSlotLag", Label: "SlotLag"},
			},
		},
		p.Prefix + ".ReplicationSlotDiskUsage": {
			Label: p.LabelPrefix + " Replication Slot Disk Usage",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "ReplicationSlotDiskUsage", Label: "SlotDiskUsage"},
			},
		},
		p.Prefix + ".TransactionLogsDiskUsage": {
			Label: p.LabelPrefix + " Transaction Logs Disk Usage",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "TransactionLogsDiskUsage", Label: "DiskUsage"},
			},
		},
		p.Prefix + ".TransactionLogsGeneration": {
			Label: p.LabelPrefix + " Transaction Logs Generation",
			Unit:  "bytes/sec",
			Metrics: []mp.Metrics{
				{Name: "TransactionLogsGeneration", Label: "Generation"},
			},
		},
	}
}

func (p RDSPlugin) auroraGraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		p.Prefix + ".Deadlocks": {
			Label: p.LabelPrefix + " Dead Locks",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "Deadlocks", Label: "Deadlocks"},
			},
		},
		p.Prefix + ".Transaction": {
			Label: p.LabelPrefix + " Transaction",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "ActiveTransactions", Label: "Active"},
				{Name: "BlockedTransactions", Label: "Blocked"},
			},
		},
		p.Prefix + ".EngineUptime": {
			Label: p.LabelPrefix + " Engine Uptime",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "EngineUptime", Label: "EngineUptime"},
			},
		},
		p.Prefix + ".Latency": {
			Label: p.LabelPrefix + " Latency in second",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "SelectThroughput", Label: "Select"},
				{Name: "InsertLatency", Label: "Insert"},
				{Name: "UpdateLatency", Label: "Update"},
				{Name: "DeleteLatency", Label: "Delete"},
				{Name: "CommitLatency", Label: "Commit"},
				{Name: "DDLLatency", Label: "DDL"},
				{Name: "DMLLatency", Label: "DML"},
			},
		},
		p.Prefix + ".Throughput": {
			Label: p.LabelPrefix + " Throughput",
			Unit:  "bytes/sec",
			Metrics: []mp.Metrics{
				{Name: "SelectThroughput", Label: "Select"},
				{Name: "InsertThroughput", Label: "Insert"},
				{Name: "UpdateThroughput", Label: "Update"},
				{Name: "DeleteThroughput", Label: "Delete"},
				{Name: "CommitThroughput", Label: "Commit"},
				{Name: "DDLThroughput", Label: "DDL"},
				{Name: "DMLThroughput", Label: "DML"},
			},
		},
		p.Prefix + ".Queries": {
			Label: p.LabelPrefix + " Queries",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "Queries", Label: "Queries"},
			},
		},
		p.Prefix + ".LoginFailures": {
			Label: p.LabelPrefix + " Login Failures",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "LoginFailures", Label: "LoginFailures"},
			},
		},
		p.Prefix + ".CacheHitRatio": {
			Label: p.LabelPrefix + " Cache Hit Ratio",
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				{Name: "ResultSetCacheHitRatio", Label: "ResultSet"},
				{Name: "BufferCacheHitRatio", Label: "Buffer"},
			},
		},
		p.Prefix + ".AuroraBinlogReplicaLag": {
			Label: p.LabelPrefix + " Aurora Binlog ReplicaLag",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "AuroraBinlogReplicaLag", Label: "BinlogReplicaLag"},
			},
		},
		p.Prefix + ".AuroraReplicaLag": {
			Label: p.LabelPrefix + " Aurora ReplicaLag",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "AuroraReplicaLagMaximum", Label: "Maximum"},
				{Name: "AuroraReplicaLagMinimum", Label: "Minimum"},
			},
		},
	}
}
<commit_msg>fix metric name<commit_after>package mpawsrds

import (
	mp "github.com/mackerelio/go-mackerel-plugin"
)

func (p RDSPlugin) rdsMetrics() (metrics []string) {
	for _, v := range p.GraphDefinition() {
		for _, vv := range v.Metrics {
			metrics = append(metrics, vv.Name)
		}
	}
	return
}

func mergeGraphDefs(a, b map[string]mp.Graphs) map[string]mp.Graphs {
	for k, v := range b {
		a[k] = v
	}
	return a
}

func (p RDSPlugin) mySQLGraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		p.Prefix + ".BinLogDiskUsage": {
			Label: p.LabelPrefix + " BinLogDiskUsage",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "BinLogDiskUsage", Label: "Usage"},
			},
		},
	}
}

func (p RDSPlugin) postgreSQLGraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		p.Prefix + ".MaximumUsedTransactionIDs": {
			Label: p.LabelPrefix + " Maximum Used Transaction IDs",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "MaximumUsedTransactionIDs", Label: "IDs"},
			},
		},
		p.Prefix + ".OldestReplicationSlotLag": {
			Label: p.LabelPrefix + " Oldest Replication Slot Lag",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "OldestReplicationSlotLag", Label: "SlotLag"},
			},
		},
		p.Prefix + ".ReplicationSlotDiskUsage": {
			Label: p.LabelPrefix + " Replication Slot Disk Usage",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "ReplicationSlotDiskUsage", Label: "SlotDiskUsage"},
			},
		},
		p.Prefix + ".TransactionLogsDiskUsage": {
			Label: p.LabelPrefix + " Transaction Logs Disk Usage",
			Unit:  "bytes",
			Metrics: []mp.Metrics{
				{Name: "TransactionLogsDiskUsage", Label: "DiskUsage"},
			},
		},
		p.Prefix + ".TransactionLogsGeneration": {
			Label: p.LabelPrefix + " Transaction Logs Generation",
			Unit:  "bytes/sec",
			Metrics: []mp.Metrics{
				{Name: "TransactionLogsGeneration", Label: "Generation"},
			},
		},
	}
}

func (p RDSPlugin) auroraGraphDefinition() map[string]mp.Graphs {
	return map[string]mp.Graphs{
		p.Prefix + ".Deadlocks": {
			Label: p.LabelPrefix + " Dead Locks",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "Deadlocks", Label: "Deadlocks"},
			},
		},
		p.Prefix + ".Transaction": {
			Label: p.LabelPrefix + " Transaction",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "ActiveTransactions", Label: "Active"},
				{Name: "BlockedTransactions", Label: "Blocked"},
			},
		},
		p.Prefix + ".EngineUptime": {
			Label: p.LabelPrefix + " Engine Uptime",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "EngineUptime", Label: "EngineUptime"},
			},
		},
		p.Prefix + ".Latency": {
			Label: p.LabelPrefix + " Latency in second",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "SelectLatency", Label: "Select"},
				{Name: "InsertLatency", Label: "Insert"},
				{Name: "UpdateLatency", Label: "Update"},
				{Name: "DeleteLatency", Label: "Delete"},
				{Name: "CommitLatency", Label: "Commit"},
				{Name: "DDLLatency", Label: "DDL"},
				{Name: "DMLLatency", Label: "DML"},
			},
		},
		p.Prefix + ".Throughput": {
			Label: p.LabelPrefix + " Throughput",
			Unit:  "bytes/sec",
			Metrics: []mp.Metrics{
				{Name: "SelectThroughput", Label: "Select"},
				{Name: "InsertThroughput", Label: "Insert"},
				{Name: "UpdateThroughput", Label: "Update"},
				{Name: "DeleteThroughput", Label: "Delete"},
				{Name: "CommitThroughput", Label: "Commit"},
				{Name: "DDLThroughput", Label: "DDL"},
				{Name: "DMLThroughput", Label: "DML"},
			},
		},
		p.Prefix + ".Queries": {
			Label: p.LabelPrefix + " Queries",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "Queries", Label: "Queries"},
			},
		},
		p.Prefix + ".LoginFailures": {
			Label: p.LabelPrefix + " Login Failures",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "LoginFailures", Label: "LoginFailures"},
			},
		},
		p.Prefix + ".CacheHitRatio": {
			Label: p.LabelPrefix + " Cache Hit Ratio",
			Unit:  "percentage",
			Metrics: []mp.Metrics{
				{Name: "ResultSetCacheHitRatio", Label: "ResultSet"},
				{Name: "BufferCacheHitRatio", Label: "Buffer"},
			},
		},
		p.Prefix + ".AuroraBinlogReplicaLag": {
			Label: p.LabelPrefix + " Aurora Binlog ReplicaLag",
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "AuroraBinlogReplicaLag", Label: "BinlogReplicaLag"},
			},
		},
		p.Prefix + ".AuroraReplicaLag": {
			Label: p.LabelPrefix + " Aurora ReplicaLag",
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "AuroraReplicaLagMaximum", Label: "Maximum"},
				{Name: "AuroraReplicaLagMinimum", Label: "Minimum"},
			},
		},
	}
}
<|endoftext|>


<commit_before>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

// FluentdMetrics plugin for fluentd
type FluentdMetrics struct {
	Target          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string

	plugins []FluentdPluginMetrics
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (fpm FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(fpm.EmitRecords)
	case "emit_count":
		return float64(fpm.EmitCount)
	case "write_count":
		return float64(fpm.WriteCount)
	case "rollback_count":
		return float64(fpm.RollbackCount)
	case "slow_flush_count":
		return float64(fpm.SlowFlushCount)
	case "flush_time_count":
		return float64(fpm.FlushTimeCount)
	case "buffer_stage_length":
		return float64(fpm.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(fpm.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(fpm.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return fpm.BufferAvailableBufferSpaceRatios
	default:
		logger.Warningf("extended-metrics %s not defined", name)
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdMetrics) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics["fluentd.retry_count."+pid] = float64(p.RetryCount)
		metrics["fluentd.buffer_queue_length."+pid] = float64(p.BufferQueueLength)
		metrics["fluentd.buffer_total_queued_size."+pid] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			key := strings.Join([]string{"fluentd", name, pid}, ".")
			metrics[key] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdMetrics) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

// FetchMetrics interface for mackerelplugin
func (f FluentdMetrics) FetchMetrics() (map[string]interface{}, error) {
	req, err := http.NewRequest(http.MethodGet, f.Target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

var defaultGraphs = map[string]mp.Graphs{
	"fluentd.retry_count": {
		Label: "Fluentd retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_length": {
		Label: "Fluentd queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_total_queued_size": {
		Label: "Fluentd buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"fluentd.emit_records": {
		Label: "Fluentd emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.emit_count": {
		Label: "Fluentd emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.write_count": {
		Label: "Fluentd write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.rollback_count": {
		Label: "Fluentd rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.slow_flush_count": {
		Label: "Fluentd slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.flush_time_count": {
		Label: "Fluentd buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.buffer_stage_length": {
		Label: "Fluentd length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_stage_byte_size": {
		Label: "Fluentd bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_byte_size": {
		Label: "Fluentd bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_available_buffer_space_ratios": {
		Label: "Fluentd available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdMetrics) GraphDefinition() map[string]mp.Graphs {
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		g := g
		graphs[key] = g
	}
	for _, name := range f.extendedMetrics {
		if g, ok := extendedGraphs["fluentd."+name]; ok {
			graphs["fluentd."+name] = g
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, strings.TrimPrefix(key, "fluentd."))
		}
	case "":
	default:
		extendedMetrics = strings.Split(*extendedMetricNames, ",")
	}
	f := FluentdMetrics{
		Target:          fmt.Sprintf("http://%s:%s/api/plugins.json", *host, *port),
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<commit_msg>exit 1 if -extended_metrics is not supported<commit_after>package mpfluentd

import (
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"regexp"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin-helper"
	"github.com/mackerelio/golib/logging"
)

var logger = logging.GetLogger("metrics.plugin.fluentd")

var prefix = "fluentd"

func metricName(names ...string) string {
	return prefix + "." + strings.Join(names, ".")
}

// FluentdMetrics plugin for fluentd
type FluentdMetrics struct {
	Target          string
	Tempfile        string
	pluginType      string
	pluginIDPattern *regexp.Regexp
	extendedMetrics []string

	plugins []FluentdPluginMetrics
}

// FluentdPluginMetrics metrics
type FluentdPluginMetrics struct {
	RetryCount            uint64 `json:"retry_count"`
	BufferQueueLength     uint64 `json:"buffer_queue_length"`
	BufferTotalQueuedSize uint64 `json:"buffer_total_queued_size"`
	OutputPlugin          bool   `json:"output_plugin"`
	Type                  string `json:"type"`
	PluginCategory        string `json:"plugin_category"`
	PluginID              string `json:"plugin_id"`
	normalizedPluginID    string

	// extended metrics fluentd >= 1.6
	// https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released
	EmitRecords                      uint64  `json:"emit_records"`
	EmitCount                        uint64  `json:"emit_count"`
	WriteCount                       uint64  `json:"write_count"`
	RollbackCount                    uint64  `json:"rollback_count"`
	SlowFlushCount                   uint64  `json:"slow_flush_count"`
	FlushTimeCount                   uint64  `json:"flush_time_count"`
	BufferStageLength                uint64  `json:"buffer_stage_length"`
	BufferStageByteSize              uint64  `json:"buffer_stage_byte_size"`
	BufferQueueByteSize              uint64  `json:"buffer_queue_byte_size"`
	BufferAvailableBufferSpaceRatios float64 `json:"buffer_available_buffer_space_ratios"`
}

func (fpm FluentdPluginMetrics) getExtended(name string) float64 {
	switch name {
	case "emit_records":
		return float64(fpm.EmitRecords)
	case "emit_count":
		return float64(fpm.EmitCount)
	case "write_count":
		return float64(fpm.WriteCount)
	case "rollback_count":
		return float64(fpm.RollbackCount)
	case "slow_flush_count":
		return float64(fpm.SlowFlushCount)
	case "flush_time_count":
		return float64(fpm.FlushTimeCount)
	case "buffer_stage_length":
		return float64(fpm.BufferStageLength)
	case "buffer_stage_byte_size":
		return float64(fpm.BufferStageByteSize)
	case "buffer_queue_byte_size":
		return float64(fpm.BufferQueueByteSize)
	case "buffer_available_buffer_space_ratios":
		return fpm.BufferAvailableBufferSpaceRatios
	}
	return 0
}

// FluentMonitorJSON monitor json
type FluentMonitorJSON struct {
	Plugins []FluentdPluginMetrics `json:"plugins"`
}

var normalizePluginIDRe = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizePluginID(in string) string {
	return normalizePluginIDRe.ReplaceAllString(in, "_")
}

func (fpm FluentdPluginMetrics) getNormalizedPluginID() string {
	if fpm.normalizedPluginID == "" {
		fpm.normalizedPluginID = normalizePluginID(fpm.PluginID)
	}
	return fpm.normalizedPluginID
}

func (f *FluentdMetrics) parseStats(body []byte) (map[string]interface{}, error) {
	var j FluentMonitorJSON
	err := json.Unmarshal(body, &j)
	f.plugins = j.Plugins

	metrics := make(map[string]interface{})
	for _, p := range f.plugins {
		if f.nonTargetPlugin(p) {
			continue
		}
		pid := p.getNormalizedPluginID()
		metrics[metricName("retry_count", pid)] = float64(p.RetryCount)
		metrics[metricName("buffer_queue_length", pid)] = float64(p.BufferQueueLength)
		metrics[metricName("buffer_total_queued_size", pid)] = float64(p.BufferTotalQueuedSize)
		for _, name := range f.extendedMetrics {
			key := strings.Join([]string{"fluentd", name, pid}, ".")
			metrics[key] = p.getExtended(name)
		}
	}
	return metrics, err
}

func (f *FluentdMetrics) nonTargetPlugin(plugin FluentdPluginMetrics) bool {
	if plugin.PluginCategory != "output" {
		return true
	}
	if f.pluginType != "" && f.pluginType != plugin.Type {
		return true
	}
	if f.pluginIDPattern != nil && !f.pluginIDPattern.MatchString(plugin.PluginID) {
		return true
	}
	return false
}

// FetchMetrics interface for mackerelplugin
func (f FluentdMetrics) FetchMetrics() (map[string]interface{}, error) {
	req, err := http.NewRequest(http.MethodGet, f.Target, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", "mackerel-plugin-fluentd")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return f.parseStats(body)
}

var defaultGraphs = map[string]mp.Graphs{
	"fluentd.retry_count": {
		Label: "Fluentd retry count",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_length": {
		Label: "Fluentd queue length",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_total_queued_size": {
		Label: "Fluentd buffer total queued size",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

var extendedGraphs = map[string]mp.Graphs{
	"fluentd.emit_records": {
		Label: "Fluentd emitted records",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.emit_count": {
		Label: "Fluentd emit calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.write_count": {
		Label: "Fluentd write/try_write calls",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.rollback_count": {
		Label: "Fluentd rollbacks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.slow_flush_count": {
		Label: "Fluentd slow flushes",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.flush_time_count": {
		Label: "Fluentd buffer flush time in msec",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: true},
		},
	},
	"fluentd.buffer_stage_length": {
		Label: "Fluentd length of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_stage_byte_size": {
		Label: "Fluentd bytesize of staged buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_queue_byte_size": {
		Label: "Fluentd bytesize of queued buffer chunks",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
	"fluentd.buffer_available_buffer_space_ratios": {
		Label: "Fluentd available space for buffer",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "*", Label: "%1", Diff: false},
		},
	},
}

// GraphDefinition interface for mackerelplugin
func (f FluentdMetrics) GraphDefinition() map[string]mp.Graphs {
	graphs := make(map[string]mp.Graphs, len(defaultGraphs))
	for key, g := range defaultGraphs {
		g := g
		graphs[key] = g
	}
	for _, name := range f.extendedMetrics {
		fullName := metricName(name)
		if g, ok := extendedGraphs[fullName]; ok {
			graphs[fullName] = g
		}
	}
	return graphs
}

// Do the plugin
func Do() {
	host := flag.String("host", "localhost", "fluentd monitor_agent host")
	port := flag.String("port", "24220", "fluentd monitor_agent port")
	pluginType := flag.String("plugin-type", "", "Gets the metric that matches this plugin type")
	pluginIDPatternString := flag.String("plugin-id-pattern", "", "Gets the metric that matches this plugin id pattern")
	tempFile := flag.String("tempfile", "", "Temp file name")
	extendedMetricNames := flag.String("extended_metrics", "", "extended metric names joind with ',' or 'all' (fluentd >= v1.6.0)")
	flag.Parse()

	var pluginIDPattern *regexp.Regexp
	var err error
	if *pluginIDPatternString != "" {
		pluginIDPattern, err = regexp.Compile(*pluginIDPatternString)
		if err != nil {
			fmt.Fprintf(os.Stderr, "failed to exec mackerel-plugin-fluentd: invalid plugin-id-pattern: %s\n", err)
			os.Exit(1)
		}
	}

	var extendedMetrics []string
	switch *extendedMetricNames {
	case "all":
		for key := range extendedGraphs {
			extendedMetrics = append(extendedMetrics, strings.TrimPrefix(key, prefix+"."))
		}
	case "":
	default:
		for _, name := range strings.Split(*extendedMetricNames, ",") {
			fullName := metricName(name)
			if _, exists := extendedGraphs[fullName]; !exists {
				fmt.Fprintf(os.Stderr, "extended_metrics %s is not supported. See also https://www.fluentd.org/blog/fluentd-v1.6.0-has-been-released\n", name)
				os.Exit(1)
			}
			extendedMetrics = append(extendedMetrics, name)
		}
	}
	f := FluentdMetrics{
		Target:          fmt.Sprintf("http://%s:%s/api/plugins.json", *host, *port),
		Tempfile:        *tempFile,
		pluginType:      *pluginType,
		pluginIDPattern: pluginIDPattern,
		extendedMetrics: extendedMetrics,
	}

	helper := mp.NewMackerelPlugin(f)

	helper.Tempfile = *tempFile
	if *tempFile == "" {
		tempFileSuffix := []string{*host, *port}
		if *pluginType != "" {
			tempFileSuffix = append(tempFileSuffix, *pluginType)
		}
		if *pluginIDPatternString != "" {
			tempFileSuffix = append(tempFileSuffix, fmt.Sprintf("%x", md5.Sum([]byte(*pluginIDPatternString))))
		}
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-fluentd-%s", strings.Join(tempFileSuffix, "-")))
	}

	helper.Run()
}
<|endoftext|>