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

@@ -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

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)
+	}
+}

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)