From 642d5fd7a4ff67768b5d3df1d3f9739e7f9675b1 Mon Sep 17 00:00:00 2001
From: MBharanya <mbharanya@users.noreply.github.com>
Date: Thu, 29 Sep 2022 13:25:35 +0200
Subject: [PATCH] Add configurable list of filetypes that allow substitutions
 in attachments

---
 config.json           | 14 +++++++++++++-
 config/config.go      |  6 ++++++
 config/config_test.go | 14 +++++++++++++-
 models/attachment.go  | 38 +++++++++++++++++++++++++-------------
 4 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/config.json b/config.json
index 8d6ff39d..cc8266bf 100644
--- a/config.json
+++ b/config.json
@@ -19,5 +19,17 @@
 	"logging": {
 		"filename": "",
 		"level": ""
+	},
+	"attachments": {
+		"plain_text_file_list": [
+			".txt",
+			".html",
+			".ics",
+			".ps1",
+			".bat",
+			".vbs",
+			".sh",
+			".py"
+		]
 	}
-}
+}
\ No newline at end of file
diff --git a/config/config.go b/config/config.go
index 62b7a850..bd25876e 100644
--- a/config/config.go
+++ b/config/config.go
@@ -26,6 +26,11 @@ type PhishServer struct {
 	KeyPath   string `json:"key_path"`
 }
 
+// Attachments represents the handling of attachments in emails
+type Attachments struct {
+	PlainTextFileList []string `json:"plain_text_file_list"`
+}
+
 // Config represents the configuration information.
 type Config struct {
 	AdminConf      AdminServer `json:"admin_server"`
@@ -37,6 +42,7 @@ type Config struct {
 	TestFlag       bool        `json:"test_flag"`
 	ContactAddress string      `json:"contact_address"`
 	Logging        *log.Config `json:"logging"`
+	Attachments    Attachments `json:"attachments"`
 }
 
 // Version contains the current gophish version
diff --git a/config/config_test.go b/config/config_test.go
index 5a30b885..3cb4360f 100644
--- a/config/config_test.go
+++ b/config/config_test.go
@@ -26,7 +26,19 @@ var validConfig = []byte(`{
 	"db_name": "sqlite3",
 	"db_path": "gophish.db",
 	"migrations_prefix": "db/db_",
-	"contact_address": ""
+	"contact_address": "",
+	"attachments": {
+		"plain_text_file_list": [
+			".txt",
+			".html",
+			".ics",
+			".ps1",
+			".bat",
+			".vbs",
+			".sh",
+			".py"
+		]
+	}
 }`)
 
 func createTemporaryConfig(t *testing.T) *os.File {
diff --git a/models/attachment.go b/models/attachment.go
index f008e74d..03b1f9e1 100644
--- a/models/attachment.go
+++ b/models/attachment.go
@@ -61,6 +61,21 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
 	//   "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
 	fileExtension := filepath.Ext(a.Name)
 
+	if array_contains(conf.Attachments.PlainTextFileList, fileExtension) {
+		b, err := ioutil.ReadAll(decodedAttachment)
+		if err != nil {
+			return nil, err
+		}
+		processedAttachment, err := ExecuteTemplate(string(b), ptx)
+		if err != nil {
+			return nil, err
+		}
+		if processedAttachment == string(b) {
+			a.vanillaFile = true
+		}
+		return strings.NewReader(processedAttachment), nil
+	}
+
 	switch fileExtension {
 
 	case ".docx", ".docm", ".pptx", ".xlsx", ".xlsm":
@@ -136,21 +151,18 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
 		zipWriter.Close()
 		return bytes.NewReader(newZipArchive.Bytes()), err
 
-	case ".txt", ".html", ".ics":
-		b, err := ioutil.ReadAll(decodedAttachment)
-		if err != nil {
-			return nil, err
-		}
-		processedAttachment, err := ExecuteTemplate(string(b), ptx)
-		if err != nil {
-			return nil, err
-		}
-		if processedAttachment == string(b) {
-			a.vanillaFile = true
-		}
-		return strings.NewReader(processedAttachment), nil
 	default:
 		return decodedAttachment, nil // Default is to simply return the file
 	}
 
 }
+
+func array_contains(s []string, str string) bool {
+	for _, v := range s {
+		if v == str {
+			return true
+		}
+	}
+
+	return false
+}