From e11d470c08818fc30688fe66ce301b8390518d3b Mon Sep 17 00:00:00 2001
From: mathew <meta@pobox.com>
Date: Mon, 12 Oct 2020 08:16:04 -0500
Subject: [PATCH] Add syslog CEE support (#262)

---
 syslog.go      | 42 ++++++++++++++++++++++++++++++++----------
 syslog_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/syslog.go b/syslog.go
index ef3b2c8..c408283 100644
--- a/syslog.go
+++ b/syslog.go
@@ -7,6 +7,10 @@ import (
 	"io"
 )
 
+// See http://cee.mitre.org/language/1.0-beta1/clt.html#syslog
+// or https://www.rsyslog.com/json-elasticsearch/
+const ceePrefix = "@cee:"
+
 // SyslogWriter is an interface matching a syslog.Writer struct.
 type SyslogWriter interface {
 	io.Writer
@@ -19,17 +23,34 @@ type SyslogWriter interface {
 }
 
 type syslogWriter struct {
-	w SyslogWriter
+	w      SyslogWriter
+	prefix string
 }
 
 // SyslogLevelWriter wraps a SyslogWriter and call the right syslog level
 // method matching the zerolog level.
 func SyslogLevelWriter(w SyslogWriter) LevelWriter {
-	return syslogWriter{w}
+	return syslogWriter{w, ""}
+}
+
+// SyslogCEEWriter wraps a SyslogWriter with a SyslogLevelWriter that adds a
+// MITRE CEE prefix for JSON syslog entries, compatible with rsyslog 
+// and syslog-ng JSON logging support. 
+// See https://www.rsyslog.com/json-elasticsearch/
+func SyslogCEEWriter(w SyslogWriter) LevelWriter {
+	return syslogWriter{w, ceePrefix}
 }
 
 func (sw syslogWriter) Write(p []byte) (n int, err error) {
-	return sw.w.Write(p)
+	var pn int
+	if sw.prefix != "" {
+		pn, err = sw.w.Write([]byte(sw.prefix))
+		if err != nil {
+			return pn, err
+		}
+	}
+	n, err = sw.w.Write(p)
+	return pn + n, err
 }
 
 // WriteLevel implements LevelWriter interface.
@@ -37,22 +58,23 @@ func (sw syslogWriter) WriteLevel(level Level, p []byte) (n int, err error) {
 	switch level {
 	case TraceLevel:
 	case DebugLevel:
-		err = sw.w.Debug(string(p))
+		err = sw.w.Debug(sw.prefix + string(p))
 	case InfoLevel:
-		err = sw.w.Info(string(p))
+		err = sw.w.Info(sw.prefix + string(p))
 	case WarnLevel:
-		err = sw.w.Warning(string(p))
+		err = sw.w.Warning(sw.prefix + string(p))
 	case ErrorLevel:
-		err = sw.w.Err(string(p))
+		err = sw.w.Err(sw.prefix + string(p))
 	case FatalLevel:
-		err = sw.w.Emerg(string(p))
+		err = sw.w.Emerg(sw.prefix + string(p))
 	case PanicLevel:
-		err = sw.w.Crit(string(p))
+		err = sw.w.Crit(sw.prefix + string(p))
 	case NoLevel:
-		err = sw.w.Info(string(p))
+		err = sw.w.Info(sw.prefix + string(p))
 	default:
 		panic("invalid level")
 	}
+	// Any CEE prefix is not part of the message, so we don't include its length
 	n = len(p)
 	return
 }
diff --git a/syslog_test.go b/syslog_test.go
index 927a9c3..c168ba6 100644
--- a/syslog_test.go
+++ b/syslog_test.go
@@ -4,7 +4,9 @@
 package zerolog
 
 import (
+	"bytes"
 	"reflect"
+	"strings"
 	"testing"
 )
 
@@ -68,3 +70,39 @@ func TestSyslogWriter(t *testing.T) {
 		t.Errorf("Invalid syslog message routing: want %v, got %v", want, got)
 	}
 }
+
+type testCEEwriter struct {
+	buf *bytes.Buffer
+}
+
+// Only implement one method as we're just testing the prefixing
+func (c testCEEwriter) Debug(m string) error { return nil }
+
+func (c testCEEwriter) Info(m string) error {
+	_, err := c.buf.Write([]byte(m))
+	return err
+}
+
+func (c testCEEwriter) Warning(m string) error { return nil }
+
+func (c testCEEwriter) Err(m string) error { return nil }
+
+func (c testCEEwriter) Emerg(m string) error { return nil }
+
+func (c testCEEwriter) Crit(m string) error { return nil }
+
+func (c testCEEwriter) Write(b []byte) (int, error) {
+	return c.buf.Write(b)
+}
+
+func TestSyslogWriter_WithCEE(t *testing.T) {
+	var buf bytes.Buffer
+	sw := testCEEwriter{&buf}
+	log := New(SyslogCEEWriter(sw))
+	log.Info().Str("key", "value").Msg("message string")
+	got := string(buf.Bytes())
+	want := "@cee:{"
+	if !strings.HasPrefix(got, want) {
+		t.Errorf("Bad CEE message start: want %v, got %v", want, got)
+	}
+}