From f0cd6bcdfdc513246dc33ccabcca9824d461c829 Mon Sep 17 00:00:00 2001 From: Jordan Wright Date: Fri, 29 Apr 2016 21:09:17 -0500 Subject: [PATCH] Building out Task functionality and starting to introduce it into campaigns --- controllers/route.go | 10 +- models/campaign.go | 100 ++++------------ models/group.go | 3 + models/models_test.go | 72 ++++++++++++ models/page.go | 3 + models/smtp.go | 3 + models/task.go | 259 ++++++++++++++++++++++++++++++++++++++++++ models/template.go | 3 + worker/worker.go | 9 +- 9 files changed, 375 insertions(+), 87 deletions(-) create mode 100644 models/task.go diff --git a/controllers/route.go b/controllers/route.go index 257663fc..8bf73c01 100644 --- a/controllers/route.go +++ b/controllers/route.go @@ -161,10 +161,12 @@ func PhishHandler(w http.ResponseWriter, r *http.Request) { if err != nil { Logger.Println(err) } - p, err := models.GetPage(c.PageId, c.UserId) - if err != nil { - Logger.Println(err) - } + /* + p, err := models.GetPage(c.PageId, c.UserId) + if err != nil { + Logger.Println(err) + }*/ + p := models.Page{} switch { case r.Method == "GET": err = c.AddEvent(models.Event{Email: rs.Email, Message: models.EVENT_CLICKED}) diff --git a/models/campaign.go b/models/campaign.go index fe212f6b..a3b2f92b 100644 --- a/models/campaign.go +++ b/models/campaign.go @@ -15,17 +15,12 @@ type Campaign struct { Name string `json:"name" sql:"not null"` CreatedDate time.Time `json:"created_date"` CompletedDate time.Time `json:"completed_date"` - TemplateId int64 `json:"-"` - Template Template `json:"template"` - PageId int64 `json:"-"` - Page Page `json:"page"` + TaskId int64 `json:"-"` + Tasks []Task `json:"tasks"` Status string `json:"status"` Results []Result `json:"results,omitempty"` Groups []Group `json:"groups,omitempty"` Events []Event `json:"timeline,omitemtpy"` - SMTPId int64 `json:"-"` - SMTP SMTP `json:"smtp"` - URL string `json:"url"` } // ErrCampaignNameNotSpecified indicates there was no template given by the user @@ -43,17 +38,11 @@ var ErrPageNotSpecified = errors.New("No landing page specified") // ErrSMTPNotSpecified indicates a sending profile was not provided for the campaign var ErrSMTPNotSpecified = errors.New("No sending profile specified") -// ErrTemplateNotFound indicates the template specified does not exist in the database -var ErrTemplateNotFound = errors.New("Template not found") +// ErrTasksNotSpecified indicates there were no tasks given by the user +var ErrTasksNotSpecified = errors.New("No tasks specified") -// ErrGroupnNotFound indicates a group specified by the user does not exist in the database -var ErrGroupNotFound = errors.New("Group not found") - -// ErrPageNotFound indicates a page specified by the user does not exist in the database -var ErrPageNotFound = errors.New("Page not found") - -// ErrSMTPNotFound indicates a sending profile specified by the user does not exist in the database -var ErrSMTPNotFound = errors.New("Sending profile not found") +// ErrInvalidStartTask indicates the starting task was not sending an email +var ErrInvalidStartTask = errors.New("All campaigns must start by sending an email") // Validate checks to make sure there are no invalid fields in a submitted campaign func (c *Campaign) Validate() error { @@ -62,12 +51,10 @@ func (c *Campaign) Validate() error { return ErrCampaignNameNotSpecified case len(c.Groups) == 0: return ErrGroupNotSpecified - case c.Template.Name == "": - return ErrTemplateNotSpecified - case c.Page.Name == "": - return ErrPageNotSpecified - case c.SMTP.Name == "": - return ErrSMTPNotSpecified + case len(c.Tasks) == 0: + return ErrTasksNotSpecified + case c.Tasks[0].Type != "SEND_EMAIL": + return ErrInvalidStartTask } return nil } @@ -105,7 +92,7 @@ func (c *Campaign) UpdateStatus(s string) error { func (c *Campaign) AddEvent(e Event) error { e.CampaignId = c.Id e.Time = time.Now() - return db.Debug().Save(&e).Error + return db.Save(&e).Error } // getDetails retrieves the related attributes of the campaign @@ -123,30 +110,10 @@ func (c *Campaign) getDetails() error { Logger.Printf("%s: events not found for campaign\n", err) return err } - err = db.Table("templates").Where("id=?", c.TemplateId).Find(&c.Template).Error + c.Tasks, err = GetTasks(c.UserId, c.Id) if err != nil { - if err != gorm.ErrRecordNotFound { - return err - } - c.Template = Template{Name: "[Deleted]"} - Logger.Printf("%s: template not found for campaign\n", err) - } - err = db.Table("pages").Where("id=?", c.PageId).Find(&c.Page).Error - if err != nil { - if err != gorm.ErrRecordNotFound { - return err - } - c.Page = Page{Name: "[Deleted]"} - Logger.Printf("%s: page not found for campaign\n", err) - } - err = db.Table("SMTP").Where("id=?", c.SMTPId).Find(&c.SMTP).Error - if err != nil { - // Check if the SMTP was deleted - if err != gorm.ErrRecordNotFound { - return err - } - c.SMTP = SMTP{Name: "[Deleted]"} - Logger.Printf("%s: sending profile not found for campaign\n", err) + Logger.Println(err) + return err } return nil } @@ -211,39 +178,13 @@ func PostCampaign(c *Campaign, uid int64) error { return err } } - // Check to make sure the template exists - t, err := GetTemplateByName(c.Template.Name, uid) - if err == gorm.ErrRecordNotFound { - Logger.Printf("Error - Template %s does not exist", t.Name) - return ErrTemplateNotFound - } else if err != nil { - Logger.Println(err) - return err + for _, t := range c.Tasks { + err = PostTask(&t) + if err != nil { + Logger.Println(err) + return err + } } - c.Template = t - c.TemplateId = t.Id - // Check to make sure the page exists - p, err := GetPageByName(c.Page.Name, uid) - if err == gorm.ErrRecordNotFound { - Logger.Printf("Error - Page %s does not exist", p.Name) - return ErrPageNotFound - } else if err != nil { - Logger.Println(err) - return err - } - c.Page = p - c.PageId = p.Id - // Check to make sure the sending profile exists - s, err := GetSMTPByName(c.SMTP.Name, uid) - if err == gorm.ErrRecordNotFound { - Logger.Printf("Error - Sending profile %s does not exist", s.Name) - return ErrPageNotFound - } else if err != nil { - Logger.Println(err) - return err - } - c.SMTP = s - c.SMTPId = s.Id // Insert into the DB err = db.Save(c).Error if err != nil { @@ -285,6 +226,7 @@ func DeleteCampaign(id int64) error { Logger.Println(err) return err } + // TODO Delete all the flows associated with the campaign // Delete the campaign err = db.Delete(&Campaign{Id: id}).Error if err != nil { diff --git a/models/group.go b/models/group.go index f7669a47..f453daf4 100644 --- a/models/group.go +++ b/models/group.go @@ -43,6 +43,9 @@ var ErrGroupNameNotSpecified = errors.New("Group name not specified") // ErrNoTargetsSpecified is thrown when no targets are specified by the user var ErrNoTargetsSpecified = errors.New("No targets specified") +// ErrGroupnNotFound indicates a group specified by the user does not exist in the database +var ErrGroupNotFound = errors.New("Group not found") + // Validate performs validation on a group given by the user func (g *Group) Validate() error { switch { diff --git a/models/models_test.go b/models/models_test.go index 9ab58d92..3526996e 100644 --- a/models/models_test.go +++ b/models/models_test.go @@ -174,6 +174,78 @@ func (s *ModelsSuite) TestPostPage(c *check.C) { }) } +func (s *ModelsSuite) TestPostTaskMissingValues(c *check.C) { + // Missing template_id + t := Task{ + UserId: 1, + CampaignId: 1, + Type: "SEND_EMAIL", + Metadata: `{ + "smtp_id" : 1 + }`, + } + err = PostTask(&t) + c.Assert(err, check.Equals, ErrTemplateIdNotSpecified) + // Missing smtp_id + t.Metadata = `{ + "template_id" : 1 + }` + err = PostTask(&t) + c.Assert(err, check.Equals, ErrSMTPIdNotSpecified) +} + +func (s *ModelsSuite) TestPostTasks(c *check.C) { + temp := Template{ + Name: "Test Template", + Text: "Testing", + HTML: "Testing", + UserId: 1, + } + err := PostTemplate(&temp) + c.Assert(err, check.Equals, nil) + c.Assert(temp.Id, check.Equals, int64(1)) + smtp := SMTP{ + Name: "Test SMTP", + Host: "1.1.1.1:25", + FromAddress: "Foo Bar ", + UserId: 1, + } + err = PostSMTP(&smtp) + c.Assert(err, check.Equals, nil) + c.Assert(smtp.Id, check.Equals, int64(2)) + t := Task{ + UserId: 1, + CampaignId: 1, + Type: "SEND_EMAIL", + Metadata: `{ + "smtp_id" : 2, + "template_id" : 1 + }`, + } + st := Task{ + UserId: 1, + CampaignId: 1, + Type: "SEND_EMAIL", + Metadata: `{ + "smtp_id" : 2, + "template_id" : 1 + }`, + } + err = PostTasks([]*Task{&t, &st}) + c.Assert(err, check.Equals, nil) + c.Assert(t.Id, check.Equals, int64(1)) + c.Assert(t.NextId, check.Equals, int64(2)) + c.Assert(t.PreviousId, check.Equals, int64(0)) + c.Assert(st.NextId, check.Equals, int64(0)) + c.Assert(st.PreviousId, check.Equals, int64(1)) + c.Assert(st.Id, check.Equals, int64(2)) + // Check retrieving a value from the database + t, err = GetTask(t.Id, t.UserId) + c.Assert(err, check.Equals, nil) + c.Assert(t.NextId, check.Equals, int64(2)) + c.Assert(t.PreviousId, check.Equals, int64(0)) +} + func (s *ModelsSuite) TestPutUser(c *check.C) { u, err := GetUser(1) u.Username = "admin_changed" diff --git a/models/page.go b/models/page.go index 29ee4c3b..caf25e14 100644 --- a/models/page.go +++ b/models/page.go @@ -23,6 +23,9 @@ type Page struct { // ErrPageNameNotSpecified is thrown if the name of the landing page is blank. var ErrPageNameNotSpecified = errors.New("Page Name not specified") +// ErrPageNotFound indicates a page specified by the user does not exist in the database +var ErrPageNotFound = errors.New("Page not found") + // parseHTML parses the page HTML on save to handle the // capturing (or lack thereof!) of credentials and passwords func (p *Page) parseHTML() error { diff --git a/models/smtp.go b/models/smtp.go index 65bb7e82..740290dc 100644 --- a/models/smtp.go +++ b/models/smtp.go @@ -28,6 +28,9 @@ var ErrFromAddressNotSpecified = errors.New("No From Address specified") // in the SMTP configuration var ErrHostNotSpecified = errors.New("No SMTP Host specified") +// ErrSMTPNotFound indicates a sending profile specified by the user does not exist in the database +var ErrSMTPNotFound = errors.New("Sending profile not found") + // TableName specifies the database tablename for Gorm to use func (s SMTP) TableName() string { return "smtp" diff --git a/models/task.go b/models/task.go new file mode 100644 index 00000000..3a8d6c72 --- /dev/null +++ b/models/task.go @@ -0,0 +1,259 @@ +package models + +import ( + "encoding/json" + "errors" + + "github.com/jinzhu/gorm" +) + +// Task contains the fields used for a Task model +// Currently, the following tasks are supported: +// - LANDING_PAGE - Point users to a landing page +// - SEND_EMAIL - Send an email to users +// +// Tasks are stored in a list format in the database. +// Each task points to both its previous task and its next task. +type Task struct { + Id int64 `json:"id" gorm:"column:id; primary_key:yes"` + UserId int64 `json:"-" gorm:"column:user_id"` + CampaignId int64 `json:"campaign_id" gorm:"column:campaign_id"` + Type string `json:"type"` + PreviousId int64 `json:"previous_id" gorm:"column:previous_id"` + NextId int64 `json:"next_id" gorm:"column:next_id"` + Metadata string `json:"metadata" gorm:column:metadata"` +} + +// ErrTaskTypeNotSpecified occurs when a type is not provided in a task +var ErrTaskTypeNotSpecfied = errors.New("No type specified for task") + +// ErrInvalidTaskType occurs when an invalid task type is specified +var ErrInvalidTaskType = errors.New("Invalid task type") + +// PageMetadata contains the attributes for the metadata on a LANDING_PAGE +// task +type PageMetadata struct { + URL string `json:"url"` + PageId int64 `json:"page_id"` + UserId int64 `json:"-"` +} + +// ErrUrlNotSpecified occurs when a URL is not provided in a LANDING_PAGE +// task +var ErrUrlNotSpecified = errors.New("No URL specfied") + +// ErrPageIdNotSpecified occurs when a page id is not provided in a LANDING_PAGE +// task +var ErrPageIdNotSpecified = errors.New("Page Id not specified") + +// Validate validates that there exists a URL and a +// PageId in the metadata +// We also validate that the PageId is valid for the +// given UserId +func (p *PageMetadata) Validate() error { + switch { + case p.URL == "": + return ErrUrlNotSpecified + case p.PageId == 0: + return ErrPageIdNotSpecified + } + _, err := GetPage(p.PageId, p.UserId) + if err == gorm.ErrRecordNotFound { + return ErrPageNotFound + } + return err +} + +// SMTPMetadata contains the attributes for the metadata of a SEND_EMAIL +// task +type SMTPMetadata struct { + SMTPId int64 `json:"smtp_id"` + TemplateId int64 `json:"template_id"` + UserId int64 `json:"-"` +} + +// ErrSMTPIdNotSpecified occurs when an SMTP Id is not specified in +// a SEND_EMAIL task +var ErrSMTPIdNotSpecified = errors.New("SMTP Id not specified") + +// ErrTemplateIdNotSpecified occurs when a template id is not specified in +// a SEND_EMAIL task +var ErrTemplateIdNotSpecified = errors.New("Template Id not specified") + +// Validate validates that there exists an SMTPId and a +// TemplateId in the task metadata +// We also validate that the SMTPId and TemplateId are +// valid for the given UserId +func (s *SMTPMetadata) Validate() error { + // Check that the values are provided + switch { + case s.SMTPId == 0: + return ErrSMTPIdNotSpecified + case s.TemplateId == 0: + return ErrTemplateIdNotSpecified + } + // Check that the template and smtp are valid + _, err := GetTemplate(s.TemplateId, s.UserId) + if err == gorm.ErrRecordNotFound { + return ErrTemplateNotFound + } + if err != nil { + return err + } + _, err = GetSMTP(s.SMTPId, s.UserId) + if err == gorm.ErrRecordNotFound { + return ErrSMTPNotFound + } + return err +} + +// Validate validates that the required metadata and core information +// is present in a Task +func (t *Task) Validate() error { + switch { + case t.Type == "LANDING_PAGE": + p := PageMetadata{UserId: t.UserId} + err := json.Unmarshal([]byte(t.Metadata), &p) + if err != nil { + return err + } + return p.Validate() + case t.Type == "SEND_EMAIL": + s := SMTPMetadata{UserId: t.UserId} + err := json.Unmarshal([]byte(t.Metadata), &s) + if err != nil { + return err + } + return s.Validate() + } + return ErrInvalidTaskType +} + +// Next returns the next task in the flow +func (t *Task) Next() (Task, error) { + n := Task{} + err := db.Debug().Where("id=?", t.NextId).Find(&n).Error + if err != nil { + Logger.Println(err) + } + return n, err +} + +// Previous returns the previous task in the flow +func (t *Task) Previous() (Task, error) { + p := Task{} + err := db.Debug().Where("id=?", t.PreviousId).Find(&p).Error + if err != nil { + Logger.Println(err) + } + return p, err +} + +// GetTasks returns all the tasks in the campaign flow +func GetTasks(uid int64, cid int64) ([]Task, error) { + ts := []Task{} + t := Task{} + // Get the campaign to find the starting task ID + c := Campaign{} + err := db.Where("id=? and user_id=?", cid, uid).Find(&c).Error + if err != nil { + Logger.Println(err) + return ts, err + } + // Get the first task + err = db.Debug().Where("id=? and user_id=?", c.TaskId, uid).Find(&t).Error + if err != nil { + Logger.Println(err) + return ts, err + } + ts = append(ts, t) + // Enumerate through all the rest of the tasks, appending them to our list + for t.NextId != 0 && err != nil { + t, err = t.Next() + ts = append(ts, t) + } + // Return the results + return ts, err +} + +// GetTask returns the task, if it exists, specified by the given id and user_id. +func GetTask(id int64, uid int64) (Task, error) { + t := Task{} + err := db.Where("user_id=? and id=?", uid, id).Find(&t).Error + if err != nil { + Logger.Println(err) + } + return t, err +} + +// PostTask creates a new task and saves it to the database +// Additionally, if there is a previous id the task points to, +// it will update the previous task's "NextId" to point to itself. +func PostTask(t *Task) error { + err := t.Validate() + if err != nil { + Logger.Println(err) + return err + } + err = db.Save(t).Error + if err != nil { + Logger.Println(err) + return err + } + if t.PreviousId == 0 { + return nil + } + p := Task{} + err = db.Where("user_id=? and id=?", t.UserId, t.PreviousId).Find(&p).Error + if err != nil { + return err + } + p.NextId = t.Id + err = db.Save(&p).Error + return err +} + +// PostTasks is a helper to automatically handle the setting +// of task PreviousId and NextId. It validates tasks before +// saving them to the database. +func PostTasks(ts []*Task) error { + // Validate all the tasks + for _, t := range ts { + if err := t.Validate(); err != nil { + Logger.Println(err) + return err + } + } + // Now, we can insert all the tasks + for i, t := range ts { + // The first element does not have a PreviousId + if i > 0 { + ts[i].PreviousId = ts[i-1].Id + } + // Insert the task + err := PostTask(t) + if err != nil { + return err + } + // Finally, we have to update the previous task with the + // NextId that was just automatically set + + if t.PreviousId != 0 { + err = db.Where("user_id=? and id=?", t.UserId, t.PreviousId).Find(&ts[i-1]).Error + if err != nil { + return err + } + } + } + return nil +} + +// DeleteTask deletes an existing task in the database. +// An error is returned if a page with the given user id and task id is not found. +func DeleteTask(id int64, uid int64) error { + err = db.Where("user_id=?", uid).Delete(Task{Id: id}).Error + if err != nil { + Logger.Println(err) + } + return err +} diff --git a/models/template.go b/models/template.go index aa28867b..f9a3f323 100644 --- a/models/template.go +++ b/models/template.go @@ -25,6 +25,9 @@ var ErrTemplateNameNotSpecified = errors.New("Template name not specified") // ErrTemplateMissingParameter is thrown when a needed parameter is not provided var ErrTemplateMissingParameter = errors.New("Need to specify at least plaintext or HTML content") +// ErrTemplateNotFound indicates the template specified does not exist in the database +var ErrTemplateNotFound = errors.New("Template not found") + // Validate checks the given template to make sure values are appropriate and complete func (t *Template) Validate() error { switch { diff --git a/worker/worker.go b/worker/worker.go index 214828e6..cd75a43f 100644 --- a/worker/worker.go +++ b/worker/worker.go @@ -4,7 +4,6 @@ import ( "bytes" "crypto/tls" "encoding/base64" - "encoding/json" "errors" "log" "net/mail" @@ -37,10 +36,12 @@ func New() *Worker { func (w *Worker) Start() { Logger.Println("Background Worker Started Successfully - Waiting for Campaigns") for { - processCampaign(<-w.Queue) + <-w.Queue + //processCampaign(<-w.Queue) } } +/* func processCampaign(c *models.Campaign) { Logger.Printf("Worker received: %s", c.Name) err := c.UpdateStatus(models.CAMPAIGN_IN_PROGRESS) @@ -124,7 +125,7 @@ func processCampaign(c *models.Campaign) { } } Logger.Printf("Sending Email to %s\n", t.Email) - err = e.SendWithTLS(c.SMTP.Host, auth, tc) + //err = e.SendWithTLS(c.SMTP.Host, auth, tc) if err != nil { Logger.Println(err) es := struct { @@ -160,7 +161,7 @@ func processCampaign(c *models.Campaign) { Logger.Println(err) } } - +*/ func SendTestEmail(s *models.SendTestEmailRequest) error { e := email.Email{ Subject: s.Template.Subject,