1package prefixed
2
3import (
4	"bytes"
5	"fmt"
6	"io"
7	"os"
8	"regexp"
9	"runtime"
10	"sort"
11	"strings"
12	"sync"
13	"time"
14
15	"github.com/mgutz/ansi"
16	"github.com/sirupsen/logrus"
17	"golang.org/x/crypto/ssh/terminal"
18)
19
20const defaultTimestampFormat = time.RFC3339
21
22var (
23	baseTimestamp      time.Time    = time.Now()
24	defaultColorScheme *ColorScheme = &ColorScheme{
25		InfoLevelStyle:  "green",
26		WarnLevelStyle:  "yellow",
27		ErrorLevelStyle: "red",
28		FatalLevelStyle: "red",
29		PanicLevelStyle: "red",
30		DebugLevelStyle: "blue",
31		PrefixStyle:     "cyan",
32		TimestampStyle:  "black+h",
33	}
34	noColorsColorScheme *compiledColorScheme = &compiledColorScheme{
35		InfoLevelColor:  ansi.ColorFunc(""),
36		WarnLevelColor:  ansi.ColorFunc(""),
37		ErrorLevelColor: ansi.ColorFunc(""),
38		FatalLevelColor: ansi.ColorFunc(""),
39		PanicLevelColor: ansi.ColorFunc(""),
40		DebugLevelColor: ansi.ColorFunc(""),
41		PrefixColor:     ansi.ColorFunc(""),
42		TimestampColor:  ansi.ColorFunc(""),
43	}
44	defaultCompiledColorScheme *compiledColorScheme = compileColorScheme(defaultColorScheme)
45)
46
47func miniTS() int {
48	return int(time.Since(baseTimestamp) / time.Second)
49}
50
51type ColorScheme struct {
52	InfoLevelStyle  string
53	WarnLevelStyle  string
54	ErrorLevelStyle string
55	FatalLevelStyle string
56	PanicLevelStyle string
57	DebugLevelStyle string
58	PrefixStyle     string
59	TimestampStyle  string
60}
61
62type compiledColorScheme struct {
63	InfoLevelColor  func(string) string
64	WarnLevelColor  func(string) string
65	ErrorLevelColor func(string) string
66	FatalLevelColor func(string) string
67	PanicLevelColor func(string) string
68	DebugLevelColor func(string) string
69	PrefixColor     func(string) string
70	TimestampColor  func(string) string
71}
72
73type TextFormatter struct {
74	// Set to true to bypass checking for a TTY before outputting colors.
75	ForceColors bool
76
77	// Force disabling colors. For a TTY colors are enabled by default.
78	DisableColors bool
79
80	// Force formatted layout, even for non-TTY output.
81	ForceFormatting bool
82
83	// Disable timestamp logging. useful when output is redirected to logging
84	// system that already adds timestamps.
85	DisableTimestamp bool
86
87	// Disable the conversion of the log levels to uppercase
88	DisableUppercase bool
89
90	// Enable logging the full timestamp when a TTY is attached instead of just
91	// the time passed since beginning of execution.
92	FullTimestamp bool
93
94	// Timestamp format to use for display when a full timestamp is printed.
95	TimestampFormat string
96
97	// The fields are sorted by default for a consistent output. For applications
98	// that log extremely frequently and don't use the JSON formatter this may not
99	// be desired.
100	DisableSorting bool
101
102	// Wrap empty fields in quotes if true.
103	QuoteEmptyFields bool
104
105	// Can be set to the override the default quoting character "
106	// with something else. For example: ', or `.
107	QuoteCharacter string
108
109	// Pad msg field with spaces on the right for display.
110	// The value for this parameter will be the size of padding.
111	// Its default value is zero, which means no padding will be applied for msg.
112	SpacePadding int
113
114	// Pad prefix field with spaces on the right for display.
115	// The value for this parameter will be the size of padding.
116	// Its default value is zero, which means no padding will be applied for prefix.
117	PrefixPadding int
118
119	// Color scheme to use.
120	colorScheme *compiledColorScheme
121
122	// Whether the logger's out is to a terminal.
123	isTerminal bool
124
125	// CallerPrettyfier can be set by the user to modify the content
126	// of the function and file keys in the data when ReportCaller is
127	// activated. If any of the returned value is the empty string the
128	// corresponding key will be removed from fields.
129	CallerPrettyfier func(*runtime.Frame) (function string, file string)
130
131	CallerFormatter func(function, file string) string
132
133	sync.Once
134}
135
136func getCompiledColor(main string, fallback string) func(string) string {
137	var style string
138	if main != "" {
139		style = main
140	} else {
141		style = fallback
142	}
143	return ansi.ColorFunc(style)
144}
145
146func compileColorScheme(s *ColorScheme) *compiledColorScheme {
147	return &compiledColorScheme{
148		InfoLevelColor:  getCompiledColor(s.InfoLevelStyle, defaultColorScheme.InfoLevelStyle),
149		WarnLevelColor:  getCompiledColor(s.WarnLevelStyle, defaultColorScheme.WarnLevelStyle),
150		ErrorLevelColor: getCompiledColor(s.ErrorLevelStyle, defaultColorScheme.ErrorLevelStyle),
151		FatalLevelColor: getCompiledColor(s.FatalLevelStyle, defaultColorScheme.FatalLevelStyle),
152		PanicLevelColor: getCompiledColor(s.PanicLevelStyle, defaultColorScheme.PanicLevelStyle),
153		DebugLevelColor: getCompiledColor(s.DebugLevelStyle, defaultColorScheme.DebugLevelStyle),
154		PrefixColor:     getCompiledColor(s.PrefixStyle, defaultColorScheme.PrefixStyle),
155		TimestampColor:  getCompiledColor(s.TimestampStyle, defaultColorScheme.TimestampStyle),
156	}
157}
158
159func (f *TextFormatter) init(entry *logrus.Entry) {
160	if len(f.QuoteCharacter) == 0 {
161		f.QuoteCharacter = "\""
162	}
163	if entry.Logger != nil {
164		f.isTerminal = f.checkIfTerminal(entry.Logger.Out)
165	}
166}
167
168func (f *TextFormatter) checkIfTerminal(w io.Writer) bool {
169	switch v := w.(type) {
170	case *os.File:
171		return terminal.IsTerminal(int(v.Fd()))
172	default:
173		return false
174	}
175}
176
177func (f *TextFormatter) SetColorScheme(colorScheme *ColorScheme) {
178	f.colorScheme = compileColorScheme(colorScheme)
179}
180
181func (f *TextFormatter) Format(entry *logrus.Entry) ([]byte, error) {
182	var b *bytes.Buffer
183	var keys []string = make([]string, 0, len(entry.Data))
184	for k := range entry.Data {
185		keys = append(keys, k)
186	}
187	lastKeyIdx := len(keys) - 1
188
189	if !f.DisableSorting {
190		sort.Strings(keys)
191	}
192	if entry.Buffer != nil {
193		b = entry.Buffer
194	} else {
195		b = &bytes.Buffer{}
196	}
197
198	prefixFieldClashes(entry.Data)
199
200	f.Do(func() { f.init(entry) })
201
202	isFormatted := f.ForceFormatting || f.isTerminal
203
204	timestampFormat := f.TimestampFormat
205	if timestampFormat == "" {
206		timestampFormat = defaultTimestampFormat
207	}
208	if isFormatted {
209		isColored := (f.ForceColors || f.isTerminal) && !f.DisableColors
210		var colorScheme *compiledColorScheme
211		if isColored {
212			if f.colorScheme == nil {
213				colorScheme = defaultCompiledColorScheme
214			} else {
215				colorScheme = f.colorScheme
216			}
217		} else {
218			colorScheme = noColorsColorScheme
219		}
220		f.printColored(b, entry, keys, timestampFormat, colorScheme)
221	} else {
222		if !f.DisableTimestamp {
223			f.appendKeyValue(b, "time", entry.Time.Format(timestampFormat), true)
224		}
225		f.appendKeyValue(b, "level", entry.Level.String(), true)
226		if entry.Message != "" {
227			f.appendKeyValue(b, "msg", entry.Message, lastKeyIdx >= 0)
228		}
229
230		if entry.HasCaller() {
231			var funcVal, fileVal string
232			if f.CallerPrettyfier != nil {
233				funcVal, fileVal = f.CallerPrettyfier(entry.Caller)
234			} else {
235				funcVal, fileVal = extractCallerInfo(entry.Caller)
236			}
237
238			if funcVal != "" {
239				f.appendKeyValue(b, "func", funcVal, true)
240			}
241			if fileVal != "" {
242				f.appendKeyValue(b, "file", fileVal, true)
243			}
244		}
245
246		for i, key := range keys {
247			f.appendKeyValue(b, key, entry.Data[key], lastKeyIdx != i)
248		}
249	}
250
251	b.WriteByte('\n')
252	return b.Bytes(), nil
253}
254
255func (f *TextFormatter) printColored(b *bytes.Buffer, entry *logrus.Entry, keys []string, timestampFormat string, colorScheme *compiledColorScheme) {
256	var levelColor func(string) string
257	var levelText string
258	switch entry.Level {
259	case logrus.InfoLevel:
260		levelColor = colorScheme.InfoLevelColor
261	case logrus.WarnLevel:
262		levelColor = colorScheme.WarnLevelColor
263	case logrus.ErrorLevel:
264		levelColor = colorScheme.ErrorLevelColor
265	case logrus.FatalLevel:
266		levelColor = colorScheme.FatalLevelColor
267	case logrus.PanicLevel:
268		levelColor = colorScheme.PanicLevelColor
269	default:
270		levelColor = colorScheme.DebugLevelColor
271	}
272
273	if entry.Level != logrus.WarnLevel {
274		levelText = entry.Level.String()
275	} else {
276		levelText = "warn"
277	}
278
279	if !f.DisableUppercase {
280		levelText = strings.ToUpper(levelText)
281	}
282
283	level := levelColor(fmt.Sprintf("%5s", levelText))
284	prefix := ""
285	message := entry.Message
286
287	adjustedPrefixPadding := f.PrefixPadding //compensate for ANSI color sequences
288
289	if prefixValue, ok := entry.Data["prefix"]; ok {
290		rawPrefixLength := len(prefixValue.(string))
291		prefix = colorScheme.PrefixColor(" " + prefixValue.(string) + ":")
292		adjustedPrefixPadding = f.PrefixPadding + (len(prefix) - rawPrefixLength - 1)
293	} else {
294		prefixValue, trimmedMsg := extractPrefix(entry.Message)
295		rawPrefixLength := len(prefixValue)
296		if len(prefixValue) > 0 {
297			prefix = colorScheme.PrefixColor(" " + prefixValue + ":")
298			message = trimmedMsg
299		}
300		adjustedPrefixPadding = f.PrefixPadding + (len(prefix) - rawPrefixLength - 1)
301	}
302
303	prefixFormat := "%s"
304	if f.PrefixPadding != 0 {
305		prefixFormat = fmt.Sprintf("%%-%ds%%s", adjustedPrefixPadding)
306	}
307
308	messageFormat := "%s"
309	if f.SpacePadding != 0 {
310		messageFormat = fmt.Sprintf("%%-%ds", f.SpacePadding)
311	}
312
313	caller := ""
314	if entry.HasCaller() {
315		var funcVal, fileVal string
316		if f.CallerPrettyfier != nil {
317			funcVal, fileVal = f.CallerPrettyfier(entry.Caller)
318		} else {
319			funcVal, fileVal = extractCallerInfo(entry.Caller)
320		}
321
322		if f.CallerFormatter != nil {
323			caller = f.CallerFormatter(funcVal, fileVal)
324		} else {
325			caller = fmt.Sprintf(" (%s: %s)", fileVal, funcVal)
326		}
327	}
328
329	if f.DisableTimestamp {
330		fmt.Fprintf(b, "%s"+prefixFormat+" "+messageFormat, level, prefix, caller, message)
331	} else {
332		var timestamp string
333		if !f.FullTimestamp {
334			timestamp = fmt.Sprintf("[%04d]", miniTS())
335		} else {
336			timestamp = fmt.Sprintf("[%s]", entry.Time.Format(timestampFormat))
337		}
338		fmt.Fprintf(b, "%s %s"+prefixFormat+" "+messageFormat, colorScheme.TimestampColor(timestamp), level, prefix, caller, message)
339	}
340	for _, k := range keys {
341		if k != "prefix" {
342			v := entry.Data[k]
343			fmt.Fprintf(b, " %s=%+v", levelColor(k), v)
344		}
345	}
346}
347
348func (f *TextFormatter) needsQuoting(text string) bool {
349	if f.QuoteEmptyFields && len(text) == 0 {
350		return true
351	}
352	for _, ch := range text {
353		if !((ch >= 'a' && ch <= 'z') ||
354			(ch >= 'A' && ch <= 'Z') ||
355			(ch >= '0' && ch <= '9') ||
356			ch == '-' || ch == '.') {
357			return true
358		}
359	}
360	return false
361}
362
363func extractCallerInfo(caller *runtime.Frame) (string, string) {
364	funcVal := caller.Function
365	fileVal := fmt.Sprintf("%s:%d", caller.File, caller.Line)
366	return funcVal, fileVal
367}
368
369func extractPrefix(msg string) (string, string) {
370	prefix := ""
371	regex := regexp.MustCompile("^\\[(.*?)\\]")
372	if regex.MatchString(msg) {
373		match := regex.FindString(msg)
374		prefix, msg = match[1:len(match)-1], strings.TrimSpace(msg[len(match):])
375	}
376	return prefix, msg
377}
378
379func (f *TextFormatter) appendKeyValue(b *bytes.Buffer, key string, value interface{}, appendSpace bool) {
380	b.WriteString(key)
381	b.WriteByte('=')
382	f.appendValue(b, value)
383
384	if appendSpace {
385		b.WriteByte(' ')
386	}
387}
388
389func (f *TextFormatter) appendValue(b *bytes.Buffer, value interface{}) {
390	switch value := value.(type) {
391	case string:
392		if !f.needsQuoting(value) {
393			b.WriteString(value)
394		} else {
395			fmt.Fprintf(b, "%s%v%s", f.QuoteCharacter, value, f.QuoteCharacter)
396		}
397	case error:
398		errmsg := value.Error()
399		if !f.needsQuoting(errmsg) {
400			b.WriteString(errmsg)
401		} else {
402			fmt.Fprintf(b, "%s%v%s", f.QuoteCharacter, errmsg, f.QuoteCharacter)
403		}
404	default:
405		fmt.Fprint(b, value)
406	}
407}
408
409// This is to not silently overwrite `time`, `msg` and `level` fields when
410// dumping it. If this code wasn't there doing:
411//
412//  logrus.WithField("level", 1).Info("hello")
413//
414// would just silently drop the user provided level. Instead with this code
415// it'll be logged as:
416//
417//  {"level": "info", "fields.level": 1, "msg": "hello", "time": "..."}
418func prefixFieldClashes(data logrus.Fields) {
419	if t, ok := data["time"]; ok {
420		data["fields.time"] = t
421	}
422
423	if m, ok := data["msg"]; ok {
424		data["fields.msg"] = m
425	}
426
427	if l, ok := data["level"]; ok {
428		data["fields.level"] = l
429	}
430}
431