circolog/formatter/format.go

package formatter

import (
	"encoding/json"
	"fmt"
	"hash/fnv"
	"io"
	"text/template"
	"time"

	"git.lattuga.net/boyska/circolog/data"
	"github.com/mgutz/ansi"
	"gopkg.in/mgo.v2/bson"
)

// Formatter is an interface, so that multiple implementations can exist
type Formatter func(data.Message) string

var tmplFuncs template.FuncMap
var syslogTmpl *template.Template

func init() {
	tmplFuncs := template.FuncMap{
		"dateFormat": func(dt time.Time, fmt string) string {
			return dt.Format(fmt)
		},
		"rfc822": func(dt time.Time) string {
			return dt.Format(time.RFC822)
		},
		"sevName": func(s int) string {
			names := []string{"emerg ", "alert ", "crit  ", "err   ", "warn  ", "notice", "info  ", "dbg   "}
			switch {
			case s < 2: // emerg..alert
				return ansi.Color(names[s], "red+b")
			case s < 4: // emerg..err
				return ansi.Color(names[s], "red")
			case s < 6: // warn..notice
				return ansi.Color(names[s], "white+b")
			case s >= len(names):
				return "???"
			default:
				return names[s]
			}
		},
		"color": func(color, text string) string {
			return ansi.Color(text, color) // slow; should use colorfunc
		},
		"red": ansi.ColorFunc("red+b"),
		"autoColor": func(s string) string {
			// from https://weechat.org/blog/post/2011/08/28/Beautify-your-WeeChat
			palette := []string{"31", "35", "38", "40", "49", "63", "70", "80", "92", "99", "112", "126", "130", "138", "142", "148", "167", "169", "174", "176", "178", "184", "186", "210", "212", "215", "247"}
			hash := fnv.New32()
			hash.Write([]byte(s))
			picked := palette[int(hash.Sum32())%len(palette)]
			return ansi.Color(s, picked)
		},
	}
	syslogTmpl = template.Must(template.New("syslog").Funcs(tmplFuncs).Parse(
		"{{color \"yellow\" (rfc822 (index . \"time\")) }} {{index . \"host\"}} " +
			"{{index . \"prog\" | autoColor}}" +
			"{{ if (ne (index . \"proc_id\") \"-\")}}[{{index . \"proc_id\"}}]{{end}}: " +
			"{{ sevName (index . \"sev\") }} " +
			"{{index . \"msg\"}}",
	))
}

type Format int

func (rf *Format) Set(v string) error {
	newvalue, err := parseFormat(v)
	if err != nil {
		return err
	}
	*rf = newvalue
	return nil
}

func (rf Format) String() string {
	switch rf {
	case FormatJSON:
		return "json"
	case FormatSyslog:
		return "syslog"
	case FormatBSON:
		return "bson"
	}
	return ""
}

func (rf Format) WriteFormatted(w io.Writer, msg data.Message) error {
	return WriteFormatted(w, rf, msg)
}

func parseFormat(s string) (Format, error) {
	switch s {
	case "json":
		return FormatJSON, nil
	case "syslog":
		return FormatSyslog, nil
	case "bson":
		return FormatBSON, nil
	default:
		return 0, fmt.Errorf("Undefined format `%s`", s)
	}
}

const (
	FormatSyslog = iota // 0
	FormatJSON   = iota
	FormatBSON   = iota
)

func WriteFormatted(w io.Writer, f Format, msg data.Message) error {
	switch f {
	case FormatSyslog:
		return syslogTmpl.Execute(w, msg)
	case FormatJSON:
		enc := json.NewEncoder(w)
		enc.SetIndent("", "  ")
		return enc.Encode(msg)
	case FormatBSON:
		enc, err := bson.Marshal(msg)
		if err != nil {
			return err
		}
		_, err = w.Write(enc)
		return err
	}
	return nil
}