From 939fb2f8471ab5f265c98210f0e64ad2286b63d1 Mon Sep 17 00:00:00 2001
From: Simon Let <simon.a.let@gmail.com>
Date: Mon, 9 Sep 2019 23:19:36 +0200
Subject: [PATCH] add sanitized flag to record, add Enrich() to record

---
 common/resh-common.go                     | 34 +++++++++++++++++++----
 sanitize-history/resh-sanitize-history.go |  4 ++-
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/common/resh-common.go b/common/resh-common.go
index ccfc21c..1bd3f4a 100644
--- a/common/resh-common.go
+++ b/common/resh-common.go
@@ -1,6 +1,11 @@
 package common
 
-import "strconv"
+import (
+	"log"
+	"strconv"
+
+	"github.com/mattn/go-shellwords"
+)
 
 // Record representing single executed command with its metadata
 type Record struct {
@@ -65,7 +70,11 @@ type Record struct {
 	ReshRevision string `json:"reshRevision"`
 
 	// added by sanitizatizer
-	CmdLength int `json:"cmdLength"`
+	Sanitized bool `json:"sanitized"`
+	CmdLength int  `json:"cmdLength"`
+
+	// enriching fields - added "later"
+	FirstWord string `json:"firstWord"`
 }
 
 // FallbackRecord when record is too old and can't be parsed into regular Record
@@ -131,9 +140,6 @@ type FallbackRecord struct {
 	ReshUuid     string `json:"reshUuid"`
 	ReshVersion  string `json:"reshVersion"`
 	ReshRevision string `json:"reshRevision"`
-
-	// added by sanitizatizer
-	CmdLength int `json:"cmdLength"`
 }
 
 // ConvertRecord from FallbackRecord to Record
@@ -202,6 +208,24 @@ func ConvertRecord(r *FallbackRecord) Record {
 	}
 }
 
+// Enrich - adds additional fields to the record
+func (r *Record) Enrich() {
+	// Get command/first word from commandline
+	r.FirstWord = GetCommandFromCommandLine(r.CmdLine)
+}
+
+// GetCommandFromCommandLine func
+func GetCommandFromCommandLine(cmdLine string) string {
+	args, err := shellwords.Parse(cmdLine)
+	if err != nil {
+		log.Fatal("shellwords Error:", err)
+	}
+	if len(args) > 0 {
+		return args[0]
+	}
+	return ""
+}
+
 // Config struct
 type Config struct {
 	Port int
diff --git a/sanitize-history/resh-sanitize-history.go b/sanitize-history/resh-sanitize-history.go
index 03926b1..c9fc057 100644
--- a/sanitize-history/resh-sanitize-history.go
+++ b/sanitize-history/resh-sanitize-history.go
@@ -176,6 +176,8 @@ func (s *sanitizer) sanitizeRecord(record *common.Record) error {
 	if err != nil {
 		log.Fatal("Cmd:", record.CmdLine, "; sanitization error:", err)
 	}
+
+	record.Sanitized = true
 	return nil
 }
 
@@ -418,4 +420,4 @@ func (s *sanitizer) trimHash(hash string) string {
 
 func (s *sanitizer) isInWhitelist(token string) bool {
 	return s.whitelist[strings.ToLower(token)] == true
-}
\ No newline at end of file
+}