From 3d7f9d96f197d6770e534a39c16951599c3b362c Mon Sep 17 00:00:00 2001
From: Hon <8292703+hxnyk@users.noreply.github.com>
Date: Tue, 4 Jun 2024 14:54:39 -0700
Subject: [PATCH] Add elasticsearch to tui (#2915)

Co-authored-by: mcastorina <m.castorina93@gmail.com>
---
 pkg/tui/pages/source_select/source_select.go  |   1 +
 .../sources/elasticsearch/elasticsearch.go    | 116 ++++++++++++++++++
 pkg/tui/sources/sources.go                    |   5 +
 3 files changed, 122 insertions(+)
 create mode 100644 pkg/tui/sources/elasticsearch/elasticsearch.go

diff --git a/pkg/tui/pages/source_select/source_select.go b/pkg/tui/pages/source_select/source_select.go
index 1cbafc752..8668421a4 100644
--- a/pkg/tui/pages/source_select/source_select.go
+++ b/pkg/tui/pages/source_select/source_select.go
@@ -65,6 +65,7 @@ func New(c common.Common) *SourceSelect {
 		OssItem("Git", "Scan git repositories."),
 		OssItem("GitHub", "Scan GitHub repositories and/or organizations."),
 		OssItem("Filesystem", "Scan your filesystem by selecting what directories to scan."),
+		OssItem("Elasticsearch", "Scan your Elasticsearch cluster or Elastic Cloud instance."),
 		OssItem("Postman", "Scan a collection, workspace, or environment from Postman, the API platform."),
 		OssItem("GitLab", "Scan GitLab repositories."),
 		OssItem("AWS S3", "Scan Amazon S3 buckets."),
diff --git a/pkg/tui/sources/elasticsearch/elasticsearch.go b/pkg/tui/sources/elasticsearch/elasticsearch.go
new file mode 100644
index 000000000..105174108
--- /dev/null
+++ b/pkg/tui/sources/elasticsearch/elasticsearch.go
@@ -0,0 +1,116 @@
+package elasticsearch
+
+import (
+	"strings"
+
+	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/common"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/components/textinputs"
+)
+
+type elasticSearchCmdModel struct {
+	textinputs.Model
+}
+
+func GetNote() string {
+	return "To connect to a local cluster, please provide the node IPs and either (username AND password) OR service token. ⭐\n⭐ To connect to a cloud cluster, please provide cloud ID AND API key."
+}
+
+func GetFields() elasticSearchCmdModel {
+	return elasticSearchCmdModel{textinputs.New([]textinputs.InputConfig{
+		{
+			Label:    "Elastic node(s)",
+			Key:      "nodes",
+			Required: false,
+			Help:     "Elastic node IPs - for scanning local clusters. Separate by space if multiple.",
+		},
+		{
+			Label:    "Username",
+			Key:      "username",
+			Required: false,
+			Help:     "Elasticsearch username. Pairs with password. For scanning local clusters.",
+		},
+		{
+			Label:    "Password",
+			Key:      "password",
+			Required: false,
+			Help:     "Elasticsearch password. Pairs with username. For scanning local clusters.",
+		},
+		{
+			Label:    "Service Token",
+			Key:      "serviceToken",
+			Required: false,
+			Help:     "Elastic service token. For scanning local clusters.",
+		},
+		{
+			Label:    "Cloud ID",
+			Key:      "cloudId",
+			Required: false,
+			Help:     "Elastic cloud ID. Pairs with API key. For scanning cloud clusters.",
+		},
+		{
+			Label:    "API Key",
+			Key:      "apiKey",
+			Required: false,
+			Help:     "Elastic API key. Pairs with cloud ID. For scanning cloud clusters.",
+		}})}
+}
+
+func findFirstNonEmptyKey(inputs map[string]textinputs.Input, keys []string) string {
+	for _, key := range keys {
+		if val, ok := inputs[key]; ok && val.Value != "" {
+			return key
+		}
+	}
+	return ""
+}
+
+func getConnectionKeys(inputs map[string]textinputs.Input) []string {
+	keys := []string{"username", "password", "serviceToken", "cloudId", "apiKey"}
+	key := findFirstNonEmptyKey(inputs, keys)
+
+	keyMap := map[string][]string{
+		"username":     {"username", "password", "nodes"},
+		"password":     {"username", "password", "nodes"},
+		"serviceToken": {"serviceToken", "nodes"},
+		"cloudId":      {"cloudId", "apiKey"},
+		"apiKey":       {"cloudId", "apiKey"},
+	}
+
+	if val, ok := keyMap[key]; ok {
+		return val
+	}
+
+	return nil
+}
+
+func (m elasticSearchCmdModel) Cmd() string {
+	var command []string
+	command = append(command, "trufflehog", "elasticsearch")
+	inputs := m.GetInputs()
+
+	for _, key := range getConnectionKeys(inputs) {
+		val, ok := inputs[key]
+		if !ok || val.Value == "" {
+			continue
+		}
+
+		if key == "nodes" {
+			nodes := strings.Fields(val.Value)
+			for _, node := range nodes {
+				command = append(command, "--nodes="+node)
+			}
+		} else {
+			command = append(command, "--"+key+"="+val.Value)
+		}
+	}
+
+	return strings.Join(command, " ")
+}
+
+func (m elasticSearchCmdModel) Summary() string {
+	inputs := m.GetInputs()
+	labels := m.GetLabels()
+
+	summaryKeys := getConnectionKeys(inputs)
+	return common.SummarizeSource(summaryKeys, inputs, labels)
+}
diff --git a/pkg/tui/sources/sources.go b/pkg/tui/sources/sources.go
index b65bf811f..0364afe73 100644
--- a/pkg/tui/sources/sources.go
+++ b/pkg/tui/sources/sources.go
@@ -6,6 +6,7 @@ import (
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/sources/circleci"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/sources/docker"
+	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/sources/elasticsearch"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/sources/filesystem"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/sources/gcs"
 	"github.com/trufflesecurity/trufflehog/v3/pkg/tui/sources/git"
@@ -23,6 +24,8 @@ func GetSourceNotes(sourceName string) string {
 		return github.GetNote()
 	case "postman":
 		return postman.GetNote()
+	case "elasticsearch":
+		return elasticsearch.GetNote()
 
 	default:
 		return ""
@@ -45,6 +48,8 @@ func GetSourceFields(sourceName string) CmdModel {
 		return circleci.GetFields()
 	case "docker":
 		return docker.GetFields()
+	case "elasticsearch":
+		return elasticsearch.GetFields()
 	case "filesystem":
 		return filesystem.GetFields()
 	case "gcs (google cloud storage)":