forked from gophish/gophish
/
worker.go
103 lines (93 loc) · 2.76 KB
/
worker.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
package worker
import (
"log"
"os"
"time"
"github.com/gophish/gophish/mailer"
"github.com/gophish/gophish/models"
)
// Logger is the logger for the worker
var Logger = log.New(os.Stdout, " ", log.Ldate|log.Ltime|log.Lshortfile)
// Worker is the background worker that handles watching for new campaigns and sending emails appropriately.
type Worker struct{}
// New creates a new worker object to handle the creation of campaigns
func New() *Worker {
return &Worker{}
}
// Start launches the worker to poll the database every minute for any pending maillogs
// that need to be processed.
func (w *Worker) Start() {
Logger.Println("Background Worker Started Successfully - Waiting for Campaigns")
for t := range time.Tick(1 * time.Minute) {
ms, err := models.GetQueuedMailLogs(t.UTC())
if err != nil {
Logger.Println(err)
continue
}
// Lock the MailLogs (they will be unlocked after processing)
err = models.LockMailLogs(ms, true)
if err != nil {
Logger.Println(err)
continue
}
// We'll group the maillogs by campaign ID to (sort of) group
// them by sending profile. This lets the mailer re-use the Sender
// instead of having to re-connect to the SMTP server for every
// email.
msg := make(map[int64][]mailer.Mail)
for _, m := range ms {
msg[m.CampaignId] = append(msg[m.CampaignId], m)
}
// Next, we process each group of maillogs in parallel
for cid, msc := range msg {
go func(cid int64, msc []mailer.Mail) {
uid := msc[0].(*models.MailLog).UserId
c, err := models.GetCampaign(cid, uid)
if err != nil {
Logger.Println(err)
errorMail(err, msc)
return
}
if c.Status == models.CAMPAIGN_QUEUED {
err := c.UpdateStatus(models.CAMPAIGN_IN_PROGRESS)
if err != nil {
Logger.Println(err)
return
}
}
Logger.Printf("Sending %d maillogs to Mailer", len(msc))
mailer.Mailer.Queue <- msc
}(cid, msc)
}
}
}
// LaunchCampaign starts a campaign
func (w *Worker) LaunchCampaign(c models.Campaign) {
ms, err := models.GetMailLogsByCampaign(c.Id)
if err != nil {
Logger.Println(err)
return
}
models.LockMailLogs(ms, true)
// This is required since you cannot pass a slice of values
// that implements an interface as a slice of that interface.
mailEntries := []mailer.Mail{}
for _, m := range ms {
mailEntries = append(mailEntries, m)
}
mailer.Mailer.Queue <- mailEntries
}
// SendTestEmail sends a test email
func (w *Worker) SendTestEmail(s *models.SendTestEmailRequest) error {
go func() {
mailer.Mailer.Queue <- []mailer.Mail{s}
}()
return <-s.ErrorChan
}
// errorMail is a helper to handle erroring out a slice of Mail instances
// in the case that an unrecoverable error occurs.
func errorMail(err error, ms []mailer.Mail) {
for _, m := range ms {
m.Error(err)
}
}