Skip to content

Commit

Permalink
Add color support for messages to console. (Velocidex#551)
Browse files Browse the repository at this point in the history
  • Loading branch information
scudette authored Aug 12, 2020
1 parent 38897a1 commit c2302f3
Show file tree
Hide file tree
Showing 9 changed files with 182 additions and 24 deletions.
3 changes: 1 addition & 2 deletions bin/artifacts.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,7 @@ var (
"show", "Show an artifact")

artifact_command_show_name = artifact_command_show.Arg(
"name", "Name to show.").
HintAction(listArtifactsHint).String()
"name", "Name to show.").Required().String()

artifact_command_list_name = artifact_command_list.Arg(
"regex", "Regex of names to match.").
Expand Down
6 changes: 6 additions & 0 deletions bin/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ var (
artifact_definitions_dir = app.Flag(
"definitions", "A directory containing artifact definitions").String()

no_color_flag = app.Flag("nocolor", "Disable color output").Bool()

verbose_flag = app.Flag(
"verbose", "Enabled verbose logging for client.").Short('v').
Default("false").Bool()
Expand Down Expand Up @@ -127,6 +129,10 @@ func main() {

command := kingpin.MustParse(app.Parse(args))

if *no_color_flag {
logging.NoColor = true
}

// Most commands load a config in the folloing order
DefaultConfigLoader = new(config.Loader).WithVerbose(*verbose_flag).
WithFileLoader(*config_path).
Expand Down
2 changes: 1 addition & 1 deletion bin/query.go
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ func doQuery() {
builder := artifacts.ScopeBuilder{
Config: config_obj,
ACLManager: vql_subsystem.NullACLManager{},
Logger: log.New(&LogWriter{config_obj}, "Velociraptor: ", 0),
Logger: log.New(&LogWriter{config_obj}, "", 0),
Env: ordereddict.NewDict(),
}

Expand Down
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ require (
github.com/golang/snappy v0.0.1
github.com/google/rpmpack v0.0.0-20200615183209-0c831d19bd44
github.com/google/uuid v1.1.1 // indirect
github.com/gookit/color v1.2.7
github.com/gorilla/csrf v1.6.2
github.com/gorilla/schema v1.1.0
github.com/grpc-ecosystem/grpc-gateway v1.11.3
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,8 @@ github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+
github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
github.com/googleapis/gax-go/v2 v2.0.5 h1:sjZBwGj9Jlw33ImPtvFviGYvseOtDM7hkSKB7+Tv3SM=
github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
github.com/gookit/color v1.2.7 h1:4qePMNWZhrmbfYJDix+J4V2l0iVW+6jQGjicELlN14E=
github.com/gookit/color v1.2.7/go.mod h1:AhIE+pS6D4Ql0SQWbBeXPHw7gY0/sjHoA4s/n1KB7xg=
github.com/gorilla/csrf v1.6.2 h1:QqQ/OWwuFp4jMKgBFAzJVW3FMULdyUW7JoM4pEWuqKg=
github.com/gorilla/csrf v1.6.2/go.mod h1:7tSf8kmjNYr7IWDCYhd3U8Ck34iQ/Yw5CJu7bAkHEGI=
github.com/gorilla/schema v1.1.0 h1:CamqUDOFUBqzrvxuz2vEwo8+SUdwsluFh7IlzJh30LY=
Expand Down
32 changes: 11 additions & 21 deletions logging/logging.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,12 @@
package logging

import (
"bytes"
"fmt"
"io/ioutil"
"log"
"os"
"path/filepath"
"strings"
"regexp"
"sync"
"time"

Expand All @@ -33,11 +32,11 @@ import (
"github.com/rifflock/lfshook"
"github.com/sirupsen/logrus"
config_proto "www.velocidex.com/golang/velociraptor/config/proto"
"www.velocidex.com/golang/velociraptor/json"
)

var (
SuppressLogging = false
NoColor = false

GenericComponent = "Velociraptor"
FrontendComponent = "VelociraptorFrontend"
Expand All @@ -53,6 +52,9 @@ var (

mu sync.Mutex
prelogs []string

tag_regex = regexp.MustCompile("<([^>]+)>")
closing_tag_regex = regexp.MustCompile("</>")
)

func InitLogging(config_obj *config_proto.Config) error {
Expand Down Expand Up @@ -230,28 +232,11 @@ func (self *LogManager) makeNewComponent(
stderr_map[logrus.ErrorLevel] = os.Stderr
}

Log.Hooks.Add(lfshook.NewHook(stderr_map, &Formatter{}))
Log.Hooks.Add(lfshook.NewHook(stderr_map, &Formatter{stderr_map}))

return &LogContext{Log}, nil
}

type Formatter struct{}

func (self *Formatter) Format(entry *logrus.Entry) ([]byte, error) {
b := &bytes.Buffer{}

levelText := strings.ToUpper(entry.Level.String())
fmt.Fprintf(b, "[%s] %v %s ", levelText, entry.Time.Format(time.RFC3339),
strings.TrimRight(entry.Message, "\r\n"))

if len(entry.Data) > 0 {
serialized, _ := json.Marshal(entry.Data)
fmt.Fprintf(b, "%s", serialized)
}

return append(b.Bytes(), '\n'), nil
}

type logWriter struct {
logger *LogContext
}
Expand Down Expand Up @@ -295,3 +280,8 @@ func GetStackTrace(err error) string {
}
return ""
}

func clearTag(message string) string {
message = tag_regex.ReplaceAllString(message, "")
return closing_tag_regex.ReplaceAllString(message, "")
}
67 changes: 67 additions & 0 deletions logging/logging_generic.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
// +build !windows

package logging

import (
"bytes"
"encoding/json"
"fmt"
"strings"
"time"

"github.com/rifflock/lfshook"
"github.com/sirupsen/logrus"
)

var (
color_map = map[string]string{
"reset": "\033[0m",
"red": "\033[31m",
"green": "\033[32m",
"yellow": "\033[33m",
"blue": "\033[34m",
"purple": "\033[35m",
"cyan": "\033[36m",
"gray": "\033[37m",
"white": "\033[97m",
}
)

type Formatter struct {
stderr_map lfshook.WriterMap
}

func (self *Formatter) Format(entry *logrus.Entry) ([]byte, error) {
b := &bytes.Buffer{}

levelText := strings.ToUpper(entry.Level.String())
fmt.Fprintf(b, "[%s] %v %s ", levelText, entry.Time.Format(time.RFC3339),
replaceTagWithCode(strings.TrimRight(entry.Message, "\r\n")))

if len(entry.Data) > 0 {
serialized, _ := json.Marshal(entry.Data)
fmt.Fprintf(b, "%s", serialized)
}

return append(b.Bytes(), '\n'), nil
}

func replaceTagWithCode(message string) string {
if NoColor {
return clearTag(message)
}

result := tag_regex.ReplaceAllStringFunc(message, func(hit string) string {
matches := tag_regex.FindStringSubmatch(hit)
if len(matches) > 1 {
code, pres := color_map[matches[1]]
if pres {
return code
}
}
return hit
})

reset := color_map["reset"]
return closing_tag_regex.ReplaceAllString(result, reset) + reset
}
60 changes: 60 additions & 0 deletions logging/logging_windows.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
// +build windows

package logging

import (
"bytes"
"encoding/json"
"fmt"
"strings"
"time"

"github.com/gookit/color"
"github.com/rifflock/lfshook"
"github.com/sirupsen/logrus"
)

type Formatter struct {
stderr_map lfshook.WriterMap
}

func (self *Formatter) Format(entry *logrus.Entry) ([]byte, error) {
b := &bytes.Buffer{}

levelText := strings.ToUpper(entry.Level.String())
fmt.Fprintf(b, "[%s] %v %s ", levelText, entry.Time.Format(time.RFC3339),
strings.TrimRight(entry.Message, "\r\n"))

if len(entry.Data) > 0 {
serialized, _ := json.Marshal(entry.Data)
fmt.Fprintf(b, "%s", serialized)
}

// Only print the result to the console, if there is an stderr
// map to it.
_, pres := self.stderr_map[entry.Level]
if pres {
if NoColor {
return []byte(clearTag(b.String())), nil
}
color.Println(normalize(b.String()))
}

return nil, nil
}

func normalize(line string) string {
// Get count of opening tags
opening_matches := tag_regex.FindAllString(line, -1)
closing_matches := closing_tag_regex.FindAllString(line, -1)

if len(opening_matches) > len(closing_matches) {
for i := 0; i < len(opening_matches)-len(closing_matches); i++ {
line += "</>"
}
} else if len(opening_matches) < len(closing_matches) {
line = closing_tag_regex.ReplaceAllString(line, "")
}

return line
}
33 changes: 33 additions & 0 deletions services/inventory/dummy.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
package inventory

import (
"context"

"github.com/pkg/errors"
artifacts_proto "www.velocidex.com/golang/velociraptor/artifacts/proto"
config_proto "www.velocidex.com/golang/velociraptor/config/proto"
"www.velocidex.com/golang/velociraptor/services"
)

type Dummy struct{}

func (self Dummy) Get() *artifacts_proto.ThirdParty {
return &artifacts_proto.ThirdParty{}
}

func (self Dummy) GetToolInfo(ctx context.Context, config_obj *config_proto.Config,
tool string) (*artifacts_proto.Tool, error) {
return nil, errors.New("Not found")
}

func (self Dummy) AddTool(config_obj *config_proto.Config, tool *artifacts_proto.Tool) error {
return errors.New("Inventory service not available")
}

func (self Dummy) RemoveTool(config_obj *config_proto.Config, tool_name string) error {
return errors.New("Inventory service not available")
}

func init() {
services.RegisterInventory(&Dummy{})
}

0 comments on commit c2302f3

Please sign in to comment.