Small fixes

models/attachment.go

@@ -21,9 +21,8 @@ type Attachment struct {
 	vanillaFile bool   // Vanilla file has no template variables
 }
 
-// ValidateAttachment ensures that the provided attachment uses the supported template variables correctly.
+// Validate ensures that the provided attachment uses the supported template variables correctly.
 func (a Attachment) Validate() error {
-
 	vc := ValidationContext{
 		FromAddress: "foo@bar.com",
 		BaseURL:     "http://example.com",
@@ -53,99 +52,91 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
 	// If we've already determined there are no template variables in this attachment return it immediately
 	if a.vanillaFile == true {
 		return decodedAttachment, nil
-	} else {
+	}
 
-		// Decided to use the file extension rather than the content type, as there seems to be quite
+	// 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:
+	//  a bit of variability with types. e.g sometimes a Word docx file would have:
-		//   "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+	//   "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
-		fileExtension := filepath.Ext(a.Name)
+	fileExtension := filepath.Ext(a.Name)
 
-		switch fileExtension {
+	switch fileExtension {
 
-		case ".docx", ".docm", ".pptx", ".xlsx", ".xlsm":
+	case ".docx", ".docm", ".pptx", ".xlsx", ".xlsm":
-			// Most modern office formats are xml based and can be unarchived.
+		// 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
+		// .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.
+		// 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
+		// See https://stackoverflow.com/questions/16946978/how-to-unzip-io-readcloser
-			b := new(bytes.Buffer)
+		b := new(bytes.Buffer)
-			b.ReadFrom(decodedAttachment)
+		b.ReadFrom(decodedAttachment)
-			buf := b.Bytes()
+		zipReader, err := zip.NewReader(bytes.NewReader(b.Bytes()), int64(b.Len())) // Create a new zip reader from the file
-			zipReader, err := zip.NewReader(bytes.NewReader(buf), int64(len(buf))) // Create a new zip reader from the file
+
+		if err != nil {
+			return nil, err
+		}
+
+		newZipArchive := new(bytes.Buffer)
+		zipWriter := zip.NewWriter(newZipArchive) // For writing the new archive
+
+		// 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
+		for _, zipFile := range zipReader.File {
+			ff, err := zipFile.Open()
 			if err != nil {
 				return nil, err
 			}
-
+			defer ff.Close()
-			newZipArchive := new(bytes.Buffer)
+			contents, err := ioutil.ReadAll(ff)
-			zipWriter := zip.NewWriter(newZipArchive) // For writing the new archive
+			if err != nil {
-
+				return nil, err
-			// i. Read each file from the Word document archive
+			}
-			// ii. Apply the template to it
+			subFileExtension := filepath.Ext(zipFile.Name)
-			// iii. Add the templated content to a new zip Word archive
+			var tFile string
-			fileContainedTemplatesVars := false
+			if subFileExtension == ".xml" || subFileExtension == ".rels" { // Ignore other files, e.g binary ones and images
-			for _, zipFile := range zipReader.File {
+				// For each file apply the template.
-				ff, err := zipFile.Open()
+				tFile, err = ExecuteTemplate(string(contents), ptx)
-				if err != nil {
-					return nil, err
-				}
-				defer ff.Close()
-				contents, err := ioutil.ReadAll(ff)
-				if err != nil {
-					return nil, err
-				}
-				subFileExtension := filepath.Ext(zipFile.Name)
-				var tFile string
-				if subFileExtension == ".xml" || subFileExtension == ".rels" { // Ignore other files, e.g binary ones and images
-					// For each file apply the template.
-					tFile, err = ExecuteTemplate(string(contents), ptx)
-					if err != nil {
-						return nil, err
-					}
-					// Check if the subfile changed. We only need this to be set once to know in the future to check the 'parent' file
-					if tFile != string(contents) {
-						fileContainedTemplatesVars = true
-					}
-				} else {
-					tFile = string(contents) // Could move this to the declaration of tFile, but might be confusing to read
-				}
-				// Write new Word archive
-				newZipFile, err := zipWriter.Create(zipFile.Name)
 				if err != nil {
 					zipWriter.Close() // Don't use defer when writing files https://www.joeshaw.org/dont-defer-close-on-writable-files/
 					return nil, err
 				}
-				_, err = newZipFile.Write([]byte(tFile))
+				// Check if the subfile changed. We only need this to be set once to know in the future to check the 'parent' file
-				if err != nil {
+				if tFile != string(contents) {
-					zipWriter.Close()
+					a.vanillaFile = true
-					return nil, err
 				}
+			} else {
+				tFile = string(contents) // Could move this to the declaration of tFile, but might be confusing to read
 			}
-			// If no files in the archive had template variables, we set the 'parent' file to not be checked in the future
+			// Write new Word archive
-			if fileContainedTemplatesVars == false {
+			newZipFile, err := zipWriter.Create(zipFile.Name)
-				a.vanillaFile = true
-			}
-			zipWriter.Close()
-			return bytes.NewReader(newZipArchive.Bytes()), err
-			//processedAttachment = newZipArchive.String()
-
-		case ".txt", ".html":
-			// Feels like a lot of Reader --> String --> Reader going on here
-			buf := new(strings.Builder)
-			_, err := io.Copy(buf, decodedAttachment)
 			if err != nil {
+				zipWriter.Close() // Don't use defer when writing files https://www.joeshaw.org/dont-defer-close-on-writable-files/
 				return nil, err
 			}
-			processedAttachment, err := ExecuteTemplate(buf.String(), ptx)
+			_, err = newZipFile.Write([]byte(tFile))
 			if err != nil {
+				zipWriter.Close()
 				return nil, err
 			}
-			if processedAttachment == string(buf.String()) {
-				a.vanillaFile = true
-			}
-			return strings.NewReader(processedAttachment), nil
-		default:
-			return decodedAttachment, nil // Default is to simply return the file
 		}
+		zipWriter.Close()
+		return bytes.NewReader(newZipArchive.Bytes()), err
+
+	case ".txt", ".html":
+		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(string(b)) {
+			a.vanillaFile = true
+		}
+		return strings.NewReader(processedAttachment), nil
+	default:
+		return decodedAttachment, nil // Default is to simply return the file
 	}
 
 }

models/testdata/text-file-without-vars.txt

@@ -1,11 +1 @@
 There are no variables in this file.
-
-Lorem ipsum dolor sit amet, consectetur adipiscing elit. Donec fermentum rhoncus mi, nec tristique nulla. Cras nec tempus ligula, non vulputate tellus. Maecenas lobortis quam quis diam tempus pellentesque. Nullam id elit lobortis, dictum purus quis, aliquam dui. Sed maximus tempus accumsan. Praesent a velit leo. Fusce malesuada dui lacus, sed auctor elit luctus sed. Pellentesque convallis commodo justo, eu lacinia leo aliquet sed. Vestibulum id lectus fermentum, porta odio vel, laoreet metus. Morbi sit amet mattis est. Suspendisse sem justo, viverra ullamcorper metus sit amet, porttitor efficitur tellus. Duis odio libero, scelerisque eget tempus a, vulputate in enim. Nunc ut erat quis dolor auctor ultricies.
-
-Suspendisse odio purus, luctus ut placerat nec, consectetur id nisl. Morbi sit amet interdum dolor, sed volutpat enim. Proin sit amet quam eros. Suspendisse placerat nisl ut lorem facilisis ullamcorper. Ut et leo consequat, pulvinar est vitae, tempus sem. Nullam hendrerit efficitur viverra. Nam pellentesque non tellus ac ullamcorper. Vivamus elementum felis eget ornare finibus. Donec sit amet purus id est congue vulputate ac quis enim. Duis ut enim libero.
-
-Duis dictum dolor ante. Nulla consequat varius dolor sed rutrum. Integer dictum consequat volutpat. Aenean fringilla ex id mauris consectetur maximus. Vivamus suscipit odio non leo congue, quis hendrerit lacus luctus. Nulla facilisi. Aliquam nec ipsum suscipit, pulvinar sapien et, bibendum ante. Nunc auctor velit ac varius eleifend. Fusce viverra cursus tristique. Morbi ornare ipsum odio, aliquet auctor neque vestibulum vitae. Praesent in turpis urna. Cras vel enim id risus iaculis dapibus bibendum mollis velit. Morbi vel quam ac odio ullamcorper aliquam eget id ex. Nulla facilisi. Morbi in ex quis felis ullamcorper laoreet vel eget ex. Aenean ultricies quam tortor, non molestie enim rutrum eu.
-
-Cras ipsum neque, pretium nec ornare eget, lacinia id magna. Praesent gravida neque orci, id pellentesque nisl ultricies ut. Etiam accumsan imperdiet accumsan. Suspendisse quis consequat arcu. Cras id nisi at augue ultrices faucibus et sit amet quam. Vivamus nec orci viverra, efficitur velit et, consequat urna. Etiam imperdiet vestibulum velit, in convallis mauris consequat ac. Nullam non egestas tortor, eget sagittis tortor. Etiam in lacinia urna. Nullam interdum lacus ut ullamcorper ornare. Vestibulum tincidunt nulla tincidunt, venenatis tellus vel, placerat tellus. Nam sit amet leo volutpat sapien varius fringilla iaculis sed sapien. Praesent accumsan, urna quis tincidunt pharetra, augue nisi iaculis odio, sed rhoncus purus mi vel est. Sed ac diam diam. Morbi ut pellentesque elit.
-
-Maecenas varius orci neque, eget feugiat mi dignissim sit amet. Donec ac euismod arcu, id ornare sem. Curabitur egestas lorem vitae mi molestie venenatis. Aenean volutpat sapien libero, sit amet varius ipsum dapibus in. Sed aliquam ante quis turpis varius, ornare vulputate nisi bibendum. Cras eget pharetra lorem. Nulla at odio lacinia, vestibulum lacus vitae, luctus ligula. Sed tincidunt dolor vitae ex feugiat fermentum.