gophish/worker/worker.go

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