github.com/gophish/gophish@v0.12.2-0.20230915144530-8e7929441393/models/attachment_test.go (about)

     1  package models
     2  
     3  import (
     4  	"bufio"
     5  	"encoding/base64"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"log"
     9  	"os"
    10  	"path/filepath"
    11  	"strings"
    12  
    13  	"gopkg.in/check.v1"
    14  )
    15  
    16  func (s *ModelsSuite) TestAttachment(c *check.C) {
    17  	ptx := PhishingTemplateContext{
    18  		BaseRecipient: BaseRecipient{
    19  			FirstName: "Foo",
    20  			LastName:  "Bar",
    21  			Email:     "foo@bar.com",
    22  			Position:  "Space Janitor",
    23  		},
    24  		BaseURL:     "http://testurl.com",
    25  		URL:         "http://testurl.com/?rid=1234567",
    26  		TrackingURL: "http://testurl.local/track?rid=1234567",
    27  		Tracker:     "<img alt='' style='display: none' src='http://testurl.local/track?rid=1234567'/>",
    28  		From:        "From Address",
    29  		RId:         "1234567",
    30  	}
    31  
    32  	files, err := ioutil.ReadDir("testdata")
    33  	if err != nil {
    34  		log.Fatalf("Failed to open attachment folder 'testdata': %v\n", err)
    35  	}
    36  	for _, ff := range files {
    37  		if !ff.IsDir() && !strings.Contains(ff.Name(), "templated") {
    38  			fname := ff.Name()
    39  			fmt.Printf("Checking attachment file -> %s\n", fname)
    40  			data := readFile("testdata/" + fname)
    41  			if filepath.Ext(fname) == ".b64" {
    42  				fname = fname[:len(fname)-4]
    43  			}
    44  			a := Attachment{
    45  				Content: data,
    46  				Name:    fname,
    47  			}
    48  			t, err := a.ApplyTemplate(ptx)
    49  			c.Assert(err, check.Equals, nil)
    50  			c.Assert(a.vanillaFile, check.Equals, strings.Contains(fname, "without-vars"))
    51  			c.Assert(a.vanillaFile, check.Not(check.Equals), strings.Contains(fname, "with-vars"))
    52  
    53  			// Verfify template was applied as expected
    54  			tt, err := ioutil.ReadAll(t)
    55  			if err != nil {
    56  				log.Fatalf("Failed to parse templated file '%s': %v\n", fname, err)
    57  			}
    58  			templatedFile := base64.StdEncoding.EncodeToString(tt)
    59  			expectedOutput := readFile("testdata/" + strings.TrimSuffix(ff.Name(), filepath.Ext(ff.Name())) + ".templated" + filepath.Ext(ff.Name())) // e.g text-file-with-vars.templated.txt
    60  			c.Assert(templatedFile, check.Equals, expectedOutput)
    61  		}
    62  	}
    63  }
    64  
    65  func readFile(fname string) string {
    66  	f, err := os.Open(fname)
    67  	if err != nil {
    68  		log.Fatalf("Failed to open file '%s': %v\n", fname, err)
    69  	}
    70  	reader := bufio.NewReader(f)
    71  	content, err := ioutil.ReadAll(reader)
    72  	if err != nil {
    73  		log.Fatalf("Failed to read file '%s': %v\n", fname, err)
    74  	}
    75  	data := ""
    76  	if filepath.Ext(fname) == ".b64" {
    77  		data = string(content)
    78  	} else {
    79  		data = base64.StdEncoding.EncodeToString(content)
    80  	}
    81  	return data
    82  }