2017-12-09 21:42:07 +00:00
|
|
|
package worker
|
|
|
|
|
|
|
|
|
|
import (
|
2020-01-28 04:18:20 +00:00
|
|
|
"context"
|
|
|
|
|
"fmt"
|
|
|
|
|
"testing"
|
|
|
|
|
"time"
|
|
|
|
|
|
2017-12-09 21:42:07 +00:00
|
|
|
"github.com/gophish/gophish/config"
|
2020-01-28 04:18:20 +00:00
|
|
|
"github.com/gophish/gophish/mailer"
|
2017-12-09 21:42:07 +00:00
|
|
|
"github.com/gophish/gophish/models"
|
|
|
|
|
"github.com/stretchr/testify/suite"
|
|
|
|
|
)
|
|
|
|
|
|
2020-01-28 04:18:20 +00:00
|
|
|
type logMailer struct {
|
|
|
|
|
queue chan []mailer.Mail
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (m *logMailer) Start(ctx context.Context) {
|
|
|
|
|
return
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (m *logMailer) Queue(ms []mailer.Mail) {
|
|
|
|
|
m.queue <- ms
|
|
|
|
|
}
|
|
|
|
|
|
2017-12-09 21:42:07 +00:00
|
|
|
// WorkerSuite is a suite of tests to cover API related functions
|
|
|
|
|
type WorkerSuite struct {
|
|
|
|
|
suite.Suite
|
2018-12-15 21:42:32 +00:00
|
|
|
config *config.Config
|
2017-12-09 21:42:07 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (s *WorkerSuite) SetupSuite() {
|
2018-12-15 21:42:32 +00:00
|
|
|
conf := &config.Config{
|
|
|
|
|
DBName: "sqlite3",
|
|
|
|
|
DBPath: ":memory:",
|
|
|
|
|
MigrationsPath: "../db/db_sqlite3/migrations/",
|
|
|
|
|
}
|
|
|
|
|
err := models.Setup(conf)
|
2017-12-09 21:42:07 +00:00
|
|
|
if err != nil {
|
|
|
|
|
s.T().Fatalf("Failed creating database: %v", err)
|
|
|
|
|
}
|
2018-12-15 21:42:32 +00:00
|
|
|
s.config = conf
|
2017-12-09 21:42:07 +00:00
|
|
|
s.Nil(err)
|
2020-01-28 04:18:20 +00:00
|
|
|
s.setupCampaignDependencies()
|
2017-12-09 21:42:07 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (s *WorkerSuite) TearDownTest() {
|
|
|
|
|
campaigns, _ := models.GetCampaigns(1)
|
|
|
|
|
for _, campaign := range campaigns {
|
|
|
|
|
models.DeleteCampaign(campaign.Id)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
2020-01-28 04:18:20 +00:00
|
|
|
func (s *WorkerSuite) setupCampaignDependencies() {
|
2018-12-15 21:42:32 +00:00
|
|
|
s.config.TestFlag = true
|
2017-12-09 21:42:07 +00:00
|
|
|
// Add a group
|
|
|
|
|
group := models.Group{Name: "Test Group"}
|
2020-01-28 04:18:20 +00:00
|
|
|
for i := 0; i < 10; i++ {
|
|
|
|
|
group.Targets = append(group.Targets, models.Target{
|
|
|
|
|
BaseRecipient: models.BaseRecipient{
|
|
|
|
|
Email: fmt.Sprintf("test%d@example.com", i),
|
|
|
|
|
FirstName: "First",
|
|
|
|
|
LastName: "Example"}})
|
2017-12-09 21:42:07 +00:00
|
|
|
}
|
|
|
|
|
group.UserId = 1
|
|
|
|
|
models.PostGroup(&group)
|
|
|
|
|
|
|
|
|
|
// Add a template
|
|
|
|
|
t := models.Template{Name: "Test Template"}
|
|
|
|
|
t.Subject = "Test subject"
|
|
|
|
|
t.Text = "Text text"
|
|
|
|
|
t.HTML = "<html>Test</html>"
|
|
|
|
|
t.UserId = 1
|
|
|
|
|
models.PostTemplate(&t)
|
|
|
|
|
|
|
|
|
|
// Add a landing page
|
|
|
|
|
p := models.Page{Name: "Test Page"}
|
|
|
|
|
p.HTML = "<html>Test</html>"
|
|
|
|
|
p.UserId = 1
|
|
|
|
|
models.PostPage(&p)
|
|
|
|
|
|
|
|
|
|
// Add a sending profile
|
|
|
|
|
smtp := models.SMTP{Name: "Test Page"}
|
|
|
|
|
smtp.UserId = 1
|
|
|
|
|
smtp.Host = "example.com"
|
|
|
|
|
smtp.FromAddress = "test@test.com"
|
|
|
|
|
models.PostSMTP(&smtp)
|
2020-01-28 04:18:20 +00:00
|
|
|
}
|
2017-12-09 21:42:07 +00:00
|
|
|
|
2020-01-28 04:18:20 +00:00
|
|
|
func (s *WorkerSuite) setupCampaign(id int) (*models.Campaign, error) {
|
2017-12-09 21:42:07 +00:00
|
|
|
// Setup and "launch" our campaign
|
|
|
|
|
// Set the status such that no emails are attempted
|
2020-01-28 04:18:20 +00:00
|
|
|
c := models.Campaign{Name: fmt.Sprintf("Test campaign - %d", id)}
|
2017-12-09 21:42:07 +00:00
|
|
|
c.UserId = 1
|
2020-01-28 04:18:20 +00:00
|
|
|
template, err := models.GetTemplate(1, 1)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
c.Template = template
|
|
|
|
|
|
|
|
|
|
page, err := models.GetPage(1, 1)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
c.Page = page
|
|
|
|
|
|
|
|
|
|
smtp, err := models.GetSMTP(1, 1)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
2017-12-09 21:42:07 +00:00
|
|
|
c.SMTP = smtp
|
2020-01-28 04:18:20 +00:00
|
|
|
|
|
|
|
|
group, err := models.GetGroup(1, 1)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
2017-12-09 21:42:07 +00:00
|
|
|
c.Groups = []models.Group{group}
|
2020-01-28 04:18:20 +00:00
|
|
|
err = models.PostCampaign(&c, c.UserId)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
err = c.UpdateStatus(models.CampaignEmailsSent)
|
|
|
|
|
return &c, err
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func (s *WorkerSuite) TestMailLogGrouping() {
|
|
|
|
|
// Create the campaigns and unlock the maillogs so that they're picked up
|
|
|
|
|
// by the worker
|
|
|
|
|
for i := 0; i < 10; i++ {
|
|
|
|
|
campaign, err := s.setupCampaign(i)
|
|
|
|
|
s.Nil(err)
|
|
|
|
|
ms, err := models.GetMailLogsByCampaign(campaign.Id)
|
|
|
|
|
s.Nil(err)
|
|
|
|
|
for _, m := range ms {
|
|
|
|
|
m.Unlock()
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
lm := &logMailer{queue: make(chan []mailer.Mail)}
|
|
|
|
|
worker := &DefaultWorker{}
|
|
|
|
|
worker.mailer = lm
|
|
|
|
|
|
|
|
|
|
// Trigger the worker, generating the maillogs and sending them to the
|
|
|
|
|
// mailer
|
|
|
|
|
worker.processCampaigns(time.Now())
|
|
|
|
|
|
|
|
|
|
// Verify that each slice of maillogs received belong to the same campaign
|
|
|
|
|
for i := 0; i < 10; i++ {
|
|
|
|
|
ms := <-lm.queue
|
|
|
|
|
maillog, ok := ms[0].(*models.MailLog)
|
|
|
|
|
if !ok {
|
|
|
|
|
s.T().Fatalf("unable to cast mail to models.MailLog")
|
|
|
|
|
}
|
|
|
|
|
expected := maillog.CampaignId
|
|
|
|
|
for _, m := range ms {
|
|
|
|
|
maillog, ok = m.(*models.MailLog)
|
|
|
|
|
if !ok {
|
|
|
|
|
s.T().Fatalf("unable to cast mail to models.MailLog")
|
|
|
|
|
}
|
|
|
|
|
got := maillog.CampaignId
|
|
|
|
|
s.Equal(expected, got)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func TestMailerSuite(t *testing.T) {
|
|
|
|
|
suite.Run(t, new(WorkerSuite))
|
2017-12-09 21:42:07 +00:00
|
|
|
}
|
