cristiancmoises
/
gophish
mirror of https://github.com/gophish/gophish
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
gophish/worker/worker_test.go

166 lines
3.7 KiB
Go
Raw Normal View History

Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
package worker
import (
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
"context"
"fmt"
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
"testing"
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
"time"
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
"github.com/gophish/gophish/config"
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
"github.com/gophish/gophish/mailer"
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
"github.com/gophish/gophish/models"
)
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +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
}
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
// testContext is context to cover API related functions
type testContext struct {
Refactor servers (#1321) * Refactoring servers to support custom workers and graceful shutdown. * Refactoring workers to support custom mailers. * Refactoring mailer to be an interface, with proper instances instead of a single global instance * Cleaning up a few things. Locking maillogs for campaigns set to launch immediately to prevent a race condition. * Cleaning up API middleware to be simpler * Moving template parameters to separate struct * Changed LoadConfig to return config object * Cleaned up some error handling, removing uninitialized global error in models package * Changed static file serving to use the unindexed package
2018-12-15 21:42:32 +00:00
config *config.Config
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
}
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
func setupTest(t *testing.T) *testContext {
Refactor servers (#1321) * Refactoring servers to support custom workers and graceful shutdown. * Refactoring workers to support custom mailers. * Refactoring mailer to be an interface, with proper instances instead of a single global instance * Cleaning up a few things. Locking maillogs for campaigns set to launch immediately to prevent a race condition. * Cleaning up API middleware to be simpler * Moving template parameters to separate struct * Changed LoadConfig to return config object * Cleaned up some error handling, removing uninitialized global error in models package * Changed static file serving to use the unindexed package
2018-12-15 21:42:32 +00:00
conf := &config.Config{
DBName: "sqlite3",
DBPath: ":memory:",
MigrationsPath: "../db/db_sqlite3/migrations/",
}
err := models.Setup(conf)
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
if err != nil {
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
t.Fatalf("Failed creating database: %v", err)
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
}
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
ctx := &testContext{}
ctx.config = conf
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
createTestData(t, ctx)
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
return ctx
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
}
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
func createTestData(t *testing.T, ctx *testContext) {
ctx.config.TestFlag = true
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
// Add a group
group := models.Group{Name: "Test Group"}
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +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"}})
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
}
group.UserId = 1
models.PostGroup(&group)
// Add a template
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
template := models.Template{Name: "Test Template"}
template.Subject = "Test subject"
template.Text = "Text text"
template.HTML = "<html>Test</html>"
template.UserId = 1
models.PostTemplate(&template)
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
// 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)
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
}
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
func setupCampaign(id int) (*models.Campaign, error) {
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
// Setup and "launch" our campaign
// Set the status such that no emails are attempted
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
c := models.Campaign{Name: fmt.Sprintf("Test campaign - %d", id)}
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
c.UserId = 1
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
template, err := models.GetTemplate(1, 1)
if err != nil {
return nil, err
}
Refactoring tests to remove stretchr/testify dependency
2020-02-02 03:44:50 +00:00
c.Template = template
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
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
}
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
c.SMTP = smtp
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
group, err := models.GetGroup(1, 1)
if err != nil {
return nil, err
}
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
c.Groups = []models.Group{group}
Improve the Performance of Sending Emails (#1753) Adding the ability to cache campaigns on maillogs to greatly improve generation time (and, by extension, sending speed and memory usage).
2020-03-01 02:19:54 +00:00
err = models.PostCampaign(&c, c.UserId)
if err != nil {
return nil, err
}
err = c.UpdateStatus(models.CampaignEmailsSent)
return &c, err
}
func TestMailLogGrouping(t *testing.T) {
setupTest(t)
// Create the campaigns and unlock the maillogs so that they're picked up
// by the worker
for i := 0; i < 10; i++ {
campaign, err := setupCampaign(i)
if err != nil {
t.Fatalf("error creating campaign: %v", err)
}
ms, err := models.GetMailLogsByCampaign(campaign.Id)
if err != nil {
t.Fatalf("error getting maillogs for campaign: %v", 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 {
t.Fatalf("unable to cast mail to models.MailLog")
}
expected := maillog.CampaignId
for _, m := range ms {
maillog, ok = m.(*models.MailLog)
if !ok {
t.Fatalf("unable to cast mail to models.MailLog")
}
got := maillog.CampaignId
if got != expected {
t.Fatalf("unexpected campaign ID received for maillog: got %d expected %d", got, expected)
}
}
}
Email refactoring (#878) The initial pass at refactoring the way we send emails.
2017-12-09 21:42:07 +00:00
}