odbtest/tools/bucket2rule/bucket2rule.go

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"
	"unicode/utf8"

	"git.wecise.com/wecise/odb-go/odb"
	"gitee.com/wecisecode/util/logger"
)

// https://stackoverflow.com/questions/22367337/last-item-in-a-template-range
var ErrBadPattern = path.ErrBadPattern

type RuleStru struct {
	Class   string
	Bucket  string
	VarName string

	Param []string

	NickName    string
	PrintID     string
	PrintAtTime string
	SnmpResult  []string // snmpResult0, ...
	TmplIndex   []string // {{index .value 0}} , ...
}

//-------------------------------------------------------------------------------------------------
// bucket2rule -class "class1, class2" -tmpl "myrule.tmpl" -type "etcd"/"file"  -out "/opt/syslog"
// ./bucket2rule -keyspace "matrix" -hosts "172.26.38.247" -class "/matrix/entity/linux" -except "*_baseline*"
// ./bucket2rule -keyspace "matrix" -hosts "172.26.38.247" -class "/matrix/entity/linux" -except "*_baseline*" -type "file" -out "/opt/matrix/agent/bin/output"
//-------------------------------------------------------------------------------------------------

func main() {

	class := flag.String("class", "", "class list, sep with ',' .")
	myhost := flag.String("hosts", "", "no defualt ")

	tmpl := flag.String("tmpl", "", "template file, defualt rule.tmpl in current dir")
	out_type := flag.String("type", "", "output file type, defualt json for etcd")
	out := flag.String("out", "", "output file dir if type=file or etcd path if type=json, default /$keyspace/rules ")
	varname := flag.String("varname", "", "default _DATA ")
	keyspace := flag.String("keyspace", "", "default matrix ")
	bucket := flag.String("bucket", "", "bucket field list, , sep with ',' default all fields")
	except := flag.String("except", "", "except bucket field list, sep with ',' default none .")

	flag.Parse()

	var clazz, buckets, excepts []string
	if *class == "" {
		logger.Errorf("class can't be empty .")
		os.Exit(1)
	} else {
		clazz = strings.Split(*class, ",")
	}

	if *bucket != "" {
		buckets = strings.Split(*bucket, ",")
	}

	if *except != "" {
		excepts = strings.Split(*except, ",")
	}

	var hosts []string
	if *myhost == "" {
		logger.Errorf("hosts can't be empty .")
		os.Exit(1)
	} else {
		hosts = strings.Split(*myhost, ",")
	}

	if *keyspace == "" {
		*keyspace = "matrix"
	}

	if *tmpl == "" {
		*tmpl = "rule.tmpl"
	}

	if *out_type == "" {
		*out_type = "json"
	}

	if *out == "" {
		*out = fmt.Sprintf("/%s/rules", *keyspace)
	}

	if *varname == "" {
		*varname = "_DATA"
	}

	handle, err := template.ParseFiles(*tmpl)

	if err != nil {
		logger.Error(err)
		os.Exit(1)
	}

	strus := ruleData(*keyspace, hosts, clazz, buckets, excepts)

	result := map[string]string{}
	for _, stru := range strus {
		stru.PrintID = `{{.id}}`
		stru.PrintAtTime = `{{.attime}}`
		stru.SnmpResult = make([]string, len(stru.Param))
		stru.TmplIndex = make([]string, len(stru.Param))
		stru.VarName = *varname

		for i := range stru.Param {
			stru.SnmpResult[i] = fmt.Sprintf("snmpResult%d", i)
			stru.TmplIndex[i] = fmt.Sprintf(`{{index .value %d}}`, i)
		}

		buf := &bytes.Buffer{}
		if err := handle.ExecuteTemplate(buf, *tmpl, stru); err != nil {
			logger.Error(err)
			os.Exit(1)
		} else {
			switch *out_type {
			case "file":
				key := fmt.Sprintf(`%s_%s.rule`, stru.Bucket, stru.NickName)
				result[key] = buf.String()
			case "json":
				key := fmt.Sprintf(`%s/%s_%s.rule`, *out, stru.Bucket, stru.NickName)
				result[key] = buf.String()
			}
		}
	}

	switch *out_type {
	case "file":
		for k, v := range result {
			file := filepath.Join(*out, k)
			if err = ioutil.WriteFile(file, []byte(v), 0666); err != nil {
				logger.Error(err)
				os.Exit(1)
			}
		}

	case "json":
		if bs, err := json.Marshal(result); err != nil {
			logger.Error(err)
		} else {
			if err = ioutil.WriteFile("output.json", bs, 0666); err != nil {
				logger.Error(err)
				os.Exit(1)
			}
		}
	}
}

func ruleData(keyspace string, hosts, clazz, buckets, excepts []string) []*RuleStru {

	rules := []*RuleStru{}
	client, err := odb.NewClient(&odb.Config{
		Keyspace: keyspace,
		Hosts:    hosts,
	})

	if err != nil {
		logger.Error(err)
		return rules
	}
	defer func() {
		if err := client.Close(); err != nil {
			logger.Error(err)
		}
	}()

	for _, class := range clazz {

		class = strings.TrimSpace(class)

		var nickname string
		if res, err := client.Query(fmt.Sprintf("select nickname from /system/class where name='%s'", class)).WithMeta(odb.QueryMeta{"meta": "false"}).Do(); err != nil {
			logger.Error(err)
		} else if len(res.Data) > 0 {
			nickname = res.Data[0]["nickname"].(string)
			if res, err := client.Query(fmt.Sprintf("select field, param from /system/tsdb where class='%s'", class)).WithMeta(odb.QueryMeta{"meta": "false"}).Do(); err != nil {
				logger.Error(err)
			} else {
				for _, row := range res.Data {
					field := row["field"].(string)
					if buckets != nil {
						find := false
						for _, bucket := range buckets {
							bucket = strings.TrimSpace(bucket)
							if bucket == field {
								find = true
								break
							}
						}
						if !find {
							continue
						}
					}

					if excepts != nil {
						find := false
						for _, except := range excepts {
							except = strings.TrimSpace(except)
							if strings.ContainsAny(except, "*?") {
								if match, _ := ExpMatch(except, field); match {
									find = true
									break
								}
							} else {
								if except == field {
									find = true
									break
								}
							}
						}
						if find {
							continue
						}
					}

					param := []string{}
					if v := row["param"]; v != nil {
						for _, val := range v.([]interface{}) {
							param = append(param, val.(string))
						}
					}
					rules = append(rules, &RuleStru{Class: class, NickName: nickname, Bucket: field, Param: param})
				}
			}
		}
	}

	return rules
}

func ExpMatch(pattern, name string) (bool, error) {
	// check some base cases
	patternLen, nameLen := len(pattern), len(name)
	if patternLen == 0 && nameLen == 0 {
		return true, nil
	}
	if patternLen == 0 {
		return false, nil
	}
	if nameLen == 0 && pattern != "*" {
		return false, nil
	}

	// check for matches one rune at a time
	patIdx, nameIdx := 0, 0
	for patIdx < patternLen && nameIdx < nameLen {
		patRune, patAdj := utf8.DecodeRuneInString(pattern[patIdx:])
		nameRune, nameAdj := utf8.DecodeRuneInString(name[nameIdx:])
		if patRune == '\\' {
			// handle escaped runes
			patIdx += patAdj
			patRune, patAdj = utf8.DecodeRuneInString(pattern[patIdx:])
			if patRune == utf8.RuneError {
				return false, ErrBadPattern
			} else if patRune == nameRune {
				patIdx += patAdj
				nameIdx += nameAdj
			} else {
				return false, nil
			}
		} else if patRune == '*' {
			// handle stars
			if patIdx += patAdj; patIdx >= patternLen {
				// a star at the end of a pattern will always
				// match the rest of the path
				return true, nil
			}

			// check if we can make any matches
			for ; nameIdx < nameLen; nameIdx += nameAdj {
				if m, _ := ExpMatch(pattern[patIdx:], name[nameIdx:]); m {
					return true, nil
				}
			}
			return false, nil
		} else if patRune == '[' {
			// handle character sets
			patIdx += patAdj
			endClass := indexRuneWithEscaping(pattern[patIdx:], ']')
			if endClass == -1 {
				return false, ErrBadPattern
			}
			endClass += patIdx
			classRunes := []rune(pattern[patIdx:endClass])
			classRunesLen := len(classRunes)
			if classRunesLen > 0 {
				classIdx := 0
				matchClass := false
				if classRunes[0] == '^' {
					classIdx++
				}
				for classIdx < classRunesLen {
					low := classRunes[classIdx]
					if low == '-' {
						return false, ErrBadPattern
					}
					classIdx++
					if low == '\\' {
						if classIdx < classRunesLen {
							low = classRunes[classIdx]
							classIdx++
						} else {
							return false, ErrBadPattern
						}
					}
					high := low
					if classIdx < classRunesLen && classRunes[classIdx] == '-' {
						// we have a range of runes
						if classIdx++; classIdx >= classRunesLen {
							return false, ErrBadPattern
						}
						high = classRunes[classIdx]
						if high == '-' {
							return false, ErrBadPattern
						}
						classIdx++
						if high == '\\' {
							if classIdx < classRunesLen {
								high = classRunes[classIdx]
								classIdx++
							} else {
								return false, ErrBadPattern
							}
						}
					}
					if low <= nameRune && nameRune <= high {
						matchClass = true
					}
				}
				if matchClass == (classRunes[0] == '^') {
					return false, nil
				}
			} else {
				return false, ErrBadPattern
			}
			patIdx = endClass + 1
			nameIdx += nameAdj
		} else if patRune == '?' || patRune == nameRune {
			// handle single-rune wildcard
			patIdx += patAdj
			nameIdx += nameAdj
		} else {
			return false, nil
		}
	}
	if patIdx >= patternLen && nameIdx >= nameLen {
		return true, nil
	}
	if nameIdx >= nameLen && pattern[patIdx:] == "*" || pattern[patIdx:] == "**" {
		return true, nil
	}
	return false, nil
}

// Find the first index of a rune in a string,
// ignoring any times the rune is escaped using "\".
func indexRuneWithEscaping(s string, r rune) int {
	end := strings.IndexRune(s, r)
	if end == -1 {
		return -1
	}
	if end > 0 && s[end-1] == '\\' {
		start := end + utf8.RuneLen(r)
		end = indexRuneWithEscaping(s[start:], r)
		if end != -1 {
			end += start
		}
	}
	return end
}