#2398 Added processing of .zip attachments including files contained

models/attachment.go

@@ -46,31 +46,15 @@ func (a Attachment) Validate() error {
 	return err
 }
 
-// ApplyTemplate parses different attachment files and applies the supplied phishing template.
+// ApplyTemplateMSOffice parses MS Office attachment files and applies the supplied phishing template.
-func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, error) {
+func (a *Attachment) ApplyTemplateMSOffice(ptx PhishingTemplateContext, decodedAttachment []byte) (io.Reader, error) {
-
-	decodedAttachment := base64.NewDecoder(base64.StdEncoding, strings.NewReader(a.Content))
-
-	// If we've already determined there are no template variables in this attachment return it immediately
-	if a.vanillaFile == true {
-		return decodedAttachment, nil
-	}
-
-	// Decided to use the file extension rather than the content type, as there seems to be quite
-	//  a bit of variability with types. e.g sometimes a Word docx file would have:
-	//   "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
-	fileExtension := filepath.Ext(a.Name)
-
-	switch fileExtension {
-
-	case ".docx", ".docm", ".pptx", ".xlsx", ".xlsm":
 	// Most modern office formats are xml based and can be unarchived.
 	// .docm and .xlsm files are comprised of xml, and a binary blob for the macro code
 
 	// Zip archives require random access for reading, so it's hard to stream bytes. Solution seems to be to use a buffer.
 	// See https://stackoverflow.com/questions/16946978/how-to-unzip-io-readcloser
 	b := new(bytes.Buffer)
-		b.ReadFrom(decodedAttachment)
+	b.Write(decodedAttachment)
 	zipReader, err := zip.NewReader(bytes.NewReader(b.Bytes()), int64(b.Len())) // Create a new zip reader from the file
 
 	if err != nil {
@@ -83,7 +67,6 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
 	// i. Read each file from the Word document archive
 	// ii. Apply the template to it
 	// iii. Add the templated content to a new zip Word archive
-		a.vanillaFile = true
 	for _, zipFile := range zipReader.File {
 		ff, err := zipFile.Open()
 		if err != nil {
@@ -135,20 +118,104 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
 	}
 	zipWriter.Close()
 	return bytes.NewReader(newZipArchive.Bytes()), err
+}
+
+// ApplyTemplateTextFiles applies the supplied phishing template to text file attachments (txt, html, ics)
+func (a *Attachment) ApplyTemplateTextFiles(ptx PhishingTemplateContext, decodedAttachment []byte) (io.Reader, error) {
+
+	processedAttachment, err := ExecuteTemplate(string(decodedAttachment), ptx)
+	if err != nil {
+		return nil, err
+	}
+	if processedAttachment != string(decodedAttachment) {
+		a.vanillaFile = false
+	}
+	return strings.NewReader(processedAttachment), nil
+}
+
+// ApplyTemplate parses different attachment files and applies the supplied phishing template.
+func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, error) {
+
+	decodedAttachment := base64.NewDecoder(base64.StdEncoding, strings.NewReader(a.Content))
+
+	// If we've already determined there are no template variables in this attachment return it immediately
+	if a.vanillaFile == true {
+		return decodedAttachment, nil
+	}
+
+	// Decided to use the file extension rather than the content type, as there seems to be quite
+	//  a bit of variability with types. e.g sometimes a Word docx file would have:
+	//   "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+	fileExtension := filepath.Ext(a.Name)
 
-	case ".txt", ".html", ".ics":
 	b, err := ioutil.ReadAll(decodedAttachment)
 	if err != nil {
 		return nil, err
 	}
-		processedAttachment, err := ExecuteTemplate(string(b), ptx)
+
+	// Initially assume that the attachment is vanilla and alter the state if changes are made to files
+	a.vanillaFile = true
+
+	switch fileExtension {
+
+	case ".zip":
+		zipReader, err := zip.NewReader(bytes.NewReader(b), int64(len(b)))
 		if err != nil {
 			return nil, err
 		}
-		if processedAttachment == string(b) {
+
-			a.vanillaFile = true
+		newZipArchive := new(bytes.Buffer)
+		zipWriter := zip.NewWriter(newZipArchive) // For writing the new archive
+
+		// Iterate over every file in the zip and apply the template depending on the filetype
+		for _, zipFile := range zipReader.File {
+			ff, err := zipFile.Open()
+			if err != nil {
+				return nil, err
 			}
-		return strings.NewReader(processedAttachment), nil
+			contents, err := io.ReadAll(ff)
+			if err != nil {
+				return nil, err
+			}
+			defer ff.Close()
+
+			subfileExtension := filepath.Ext(zipFile.Name)
+			var tFile io.Reader
+			switch subfileExtension {
+
+			case ".docx", ".docm", ".pptx", ".xlsx", ".xlsm":
+				tFile, err = a.ApplyTemplateMSOffice(ptx, contents)
+				if err != nil {
+					zipWriter.Close()
+					return nil, err
+				}
+			case ".txt", ".html", ".ics":
+				tFile, err = a.ApplyTemplateTextFiles(ptx, contents)
+				if err != nil {
+					zipWriter.Close()
+					return nil, err
+				}
+			default:
+				tFile = bytes.NewReader(contents)
+			}
+
+			// Write the possibly changed file to the new zip file
+			tmp, err := zipWriter.Create(zipFile.Name)
+			if err != nil {
+				zipWriter.Close()
+				return nil, err
+			}
+			if _, err := io.Copy(tmp, tFile); err != nil {
+				zipWriter.Close()
+				return nil, err
+			}
+		}
+		zipWriter.Close()
+		return bytes.NewReader(newZipArchive.Bytes()), err
+	case ".docx", ".docm", ".pptx", ".xlsx", ".xlsm":
+		return a.ApplyTemplateMSOffice(ptx, b)
+	case ".txt", ".html", ".ics":
+		return a.ApplyTemplateTextFiles(ptx, b)
 	default:
 		return decodedAttachment, nil // Default is to simply return the file
 	}