From b7c69662ced64155c4163d5ef8bdc91ae2ad092f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20J=C3=B3zsef=20J=C3=A1nv=C3=A1ri?= <4534880+dzsibi@users.noreply.github.com> Date: Wed, 1 Jun 2022 17:14:22 +0200 Subject: [PATCH] Embed or attach files based on their file extension (#1525) Embed or attach files based on their file extension: * Set 'Content-Disposition: inline' for images * Set 'Content-Disposition: attachment' for other files --- models/email_request.go | 13 ++--------- models/maillog.go | 52 ++++++++++++++++++++++++++++++----------- models/maillog_test.go | 44 ++++++++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 24 deletions(-) diff --git a/models/email_request.go b/models/email_request.go index 09440fcd..f441b61a 100644 --- a/models/email_request.go +++ b/models/email_request.go @@ -1,11 +1,8 @@ package models import ( - "encoding/base64" "fmt" - "io" "net/mail" - "strings" "github.com/gophish/gomail" "github.com/gophish/gophish/config" @@ -171,16 +168,10 @@ func (s *EmailRequest) Generate(msg *gomail.Message) error { msg.AddAlternative("text/html", html) } } + // Attach the files for _, a := range s.Template.Attachments { - msg.Attach(func(a Attachment) (string, gomail.FileSetting, gomail.FileSetting) { - h := map[string][]string{"Content-ID": {fmt.Sprintf("<%s>", a.Name)}} - return a.Name, gomail.SetCopyFunc(func(w io.Writer) error { - decoder := base64.NewDecoder(base64.StdEncoding, strings.NewReader(a.Content)) - _, err = io.Copy(w, decoder) - return err - }), gomail.SetHeader(h) - }(a)) + addAttachment(msg, a, ptx) } return nil diff --git a/models/maillog.go b/models/maillog.go index d87a5d0d..7ab3a2ef 100644 --- a/models/maillog.go +++ b/models/maillog.go @@ -9,6 +9,8 @@ import ( "math/big" "net/mail" "os" + "path/filepath" + "strings" "time" "github.com/gophish/gomail" @@ -25,6 +27,9 @@ var MaxSendAttempts = 8 // MailLog is exceeded. var ErrMaxSendAttempts = errors.New("max send attempts exceeded") +// Attachments with these file extensions have inline disposition +var embeddedFileExtensions = []string{".jpg", ".jpeg", ".png", ".gif"} + // MailLog is a struct that holds information about an email that is to be // sent out. type MailLog struct { @@ -251,19 +256,8 @@ func (m *MailLog) Generate(msg *gomail.Message) error { } } // Attach the files - for i, _ := range c.Template.Attachments { - a := &c.Template.Attachments[i] - msg.Attach(func(a *Attachment) (string, gomail.FileSetting, gomail.FileSetting) { - h := map[string][]string{"Content-ID": {fmt.Sprintf("<%s>", a.Name)}} - return a.Name, gomail.SetCopyFunc(func(w io.Writer) error { - content, err := a.ApplyTemplate(ptx) - if err != nil { - return err - } - _, err = io.Copy(w, content) - return err - }), gomail.SetHeader(h) - }(a)) + for _, a := range c.Template.Attachments { + addAttachment(msg, a, ptx) } return nil @@ -335,3 +329,35 @@ func (m *MailLog) generateMessageID() (string, error) { msgid := fmt.Sprintf("<%d.%d.%d@%s>", t, pid, rint, h) return msgid, nil } + +// Check if an attachment should have inline disposition based on +// its file extension. +func shouldEmbedAttachment(name string) bool { + ext := filepath.Ext(name) + for _, v := range embeddedFileExtensions { + if strings.EqualFold(ext, v) { + return true + } + } + return false +} + +// Add an attachment to a gomail message, with the Content-Disposition +// header set to inline or attachment depending on its file extension. +func addAttachment(msg *gomail.Message, a Attachment, ptx PhishingTemplateContext) { + copyFunc := gomail.SetCopyFunc(func(c Attachment) func(w io.Writer) error { + return func(w io.Writer) error { + reader, err := a.ApplyTemplate(ptx) + if err != nil { + return err + } + _, err = io.Copy(w, reader) + return err + } + }(a)) + if shouldEmbedAttachment(a.Name) { + msg.Embed(a.Name, copyFunc) + } else { + msg.Attach(a.Name, copyFunc) + } +} diff --git a/models/maillog_test.go b/models/maillog_test.go index 495e973c..88619f93 100644 --- a/models/maillog_test.go +++ b/models/maillog_test.go @@ -360,6 +360,50 @@ func (s *ModelsSuite) TestMailLogGenerateEmptySubject(ch *check.C) { ch.Assert(got.Subject, check.Equals, expected.Subject) } +func (s *ModelsSuite) TestShouldEmbedAttachment(ch *check.C) { + + // Supported file extensions + ch.Assert(shouldEmbedAttachment(".png"), check.Equals, true) + ch.Assert(shouldEmbedAttachment(".jpg"), check.Equals, true) + ch.Assert(shouldEmbedAttachment(".jpeg"), check.Equals, true) + ch.Assert(shouldEmbedAttachment(".gif"), check.Equals, true) + + // Some other file extensions + ch.Assert(shouldEmbedAttachment(".docx"), check.Equals, false) + ch.Assert(shouldEmbedAttachment(".txt"), check.Equals, false) + ch.Assert(shouldEmbedAttachment(".jar"), check.Equals, false) + ch.Assert(shouldEmbedAttachment(".exe"), check.Equals, false) + + // Invalid input + ch.Assert(shouldEmbedAttachment(""), check.Equals, false) + ch.Assert(shouldEmbedAttachment("png"), check.Equals, false) +} + +func (s *ModelsSuite) TestEmbedAttachment(ch *check.C) { + campaign := s.createCampaignDependencies(ch) + campaign.Template.Attachments = []Attachment{ + { + Name: "test.png", + Type: "image/png", + Content: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABAQMAAAAl21bKAAAAA1BMVEUAAACnej3aAAAAAXRSTlMAQObYZgAAAApJREFUCNdjYAAAAAIAAeIhvDMAAAAASUVORK5CYII=", + }, + { + Name: "test.txt", + Type: "text/plain", + Content: "VGVzdCB0ZXh0IGZpbGU=", + }, + } + PutTemplate(&campaign.Template) + ch.Assert(PostCampaign(&campaign, campaign.UserId), check.Equals, nil) + got := s.emailFromFirstMailLog(campaign, ch) + + // The email package simply ignores attachments where the Content-Disposition header is set + // to inline, so the best we can do without replacing the whole thing is to check that only + // the text file was added as an attachment. + ch.Assert(got.Attachments, check.HasLen, 1) + ch.Assert(got.Attachments[0].Filename, check.Equals, "test.txt") +} + func BenchmarkMailLogGenerate100(b *testing.B) { setupBenchmark(b) campaign := setupCampaign(b, 100)