Small fixes

attachment-template-support
Glenn Wilkinson 2020-11-20 15:27:25 +01:00
parent 304aa3b6b7
commit aabf8925ad
2 changed files with 67 additions and 86 deletions

View File

@ -21,9 +21,8 @@ type Attachment struct {
vanillaFile bool // Vanilla file has no template variables 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 { func (a Attachment) Validate() error {
vc := ValidationContext{ vc := ValidationContext{
FromAddress: "foo@bar.com", FromAddress: "foo@bar.com",
BaseURL: "http://example.com", BaseURL: "http://example.com",
@ -53,7 +52,7 @@ 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 we've already determined there are no template variables in this attachment return it immediately
if a.vanillaFile == true { if a.vanillaFile == true {
return decodedAttachment, nil 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:
@ -70,8 +69,8 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
// 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 { if err != nil {
return nil, err return nil, err
} }
@ -82,7 +81,6 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
// i. Read each file from the Word document archive // i. Read each file from the Word document archive
// ii. Apply the template to it // ii. Apply the template to it
// iii. Add the templated content to a new zip Word archive // iii. Add the templated content to a new zip Word archive
fileContainedTemplatesVars := false
for _, zipFile := range zipReader.File { for _, zipFile := range zipReader.File {
ff, err := zipFile.Open() ff, err := zipFile.Open()
if err != nil { if err != nil {
@ -99,11 +97,12 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
// For each file apply the template. // For each file apply the template.
tFile, err = ExecuteTemplate(string(contents), ptx) tFile, err = ExecuteTemplate(string(contents), ptx)
if err != nil { 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 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 // 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) { if tFile != string(contents) {
fileContainedTemplatesVars = true a.vanillaFile = true
} }
} else { } else {
tFile = string(contents) // Could move this to the declaration of tFile, but might be confusing to read tFile = string(contents) // Could move this to the declaration of tFile, but might be confusing to read
@ -120,32 +119,24 @@ func (a *Attachment) ApplyTemplate(ptx PhishingTemplateContext) (io.Reader, erro
return nil, err return nil, err
} }
} }
// If no files in the archive had template variables, we set the 'parent' file to not be checked in the future
if fileContainedTemplatesVars == false {
a.vanillaFile = true
}
zipWriter.Close() zipWriter.Close()
return bytes.NewReader(newZipArchive.Bytes()), err return bytes.NewReader(newZipArchive.Bytes()), err
//processedAttachment = newZipArchive.String()
case ".txt", ".html": case ".txt", ".html":
// Feels like a lot of Reader --> String --> Reader going on here b, err := ioutil.ReadAll(decodedAttachment)
buf := new(strings.Builder)
_, err := io.Copy(buf, decodedAttachment)
if err != nil { if err != nil {
return nil, err return nil, err
} }
processedAttachment, err := ExecuteTemplate(buf.String(), ptx) processedAttachment, err := ExecuteTemplate(string(b), ptx)
if err != nil { if err != nil {
return nil, err return nil, err
} }
if processedAttachment == string(buf.String()) { if processedAttachment == string(string(b)) {
a.vanillaFile = true a.vanillaFile = true
} }
return strings.NewReader(processedAttachment), nil return strings.NewReader(processedAttachment), nil
default: default:
return decodedAttachment, nil // Default is to simply return the file return decodedAttachment, nil // Default is to simply return the file
} }
}
} }

View File

@ -1,11 +1 @@
There are no variables in this file. 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.