Initial commit of worker.go (Still work to do)

Moved constants to models.go
Changed Campaign.Template to be an actual template (will need to adjust all the methods to handle it)
Added UpdateCampaignStatus function to update a campaign status
pull/24/head
Jordan 2014-03-27 23:31:51 -05:00
parent ac7fda0b04
commit 7c50f4dbb4
4 changed files with 67 additions and 10 deletions

View File

@ -12,14 +12,15 @@ import (
"github.com/gorilla/mux"
"github.com/jordan-wright/gophish/auth"
"github.com/jordan-wright/gophish/models"
"github.com/jordan-wright/gophish/worker"
)
const (
IN_PROGRESS string = "In progress"
WAITING string = "Waiting"
COMPLETE string = "Completed"
ERROR string = "Error"
)
var Worker *worker.Worker
func init() {
Worker = worker.New()
go Worker.Start()
}
// API (/api) provides access to api documentation
func API(w http.ResponseWriter, r *http.Request) {
@ -79,12 +80,13 @@ func API_Campaigns(w http.ResponseWriter, r *http.Request) {
// Fill in the details
c.CreatedDate = time.Now()
c.CompletedDate = time.Time{}
c.Status = IN_PROGRESS
c.Status = models.QUEUED
c.UserId = ctx.Get(r, "user_id").(int64)
err = models.PostCampaign(&c, ctx.Get(r, "user_id").(int64))
if checkError(err, w, "Cannot insert campaign into database", http.StatusInternalServerError) {
return
}
Worker.Queue <- &c
cj, err := json.MarshalIndent(c, "", " ")
if checkError(err, w, "Error creating JSON response", http.StatusInternalServerError) {
return

View File

@ -13,7 +13,7 @@ type Campaign struct {
Name string `json:"name" sql:"not null"`
CreatedDate time.Time `json:"created_date"`
CompletedDate time.Time `json:"completed_date"`
Template string `json:"template"` //This may change
Template Template `json:"template"` //This may change
Status string `json:"status"`
Results []Result `json:"results,omitempty"`
Groups []Group `json:"groups,omitempty"`
@ -39,7 +39,6 @@ func GetCampaigns(uid int64) ([]Campaign, error) {
fmt.Println(err)
}
}
fmt.Printf("%v", cs)
return cs, err
}
@ -90,6 +89,10 @@ func PostCampaign(c *Campaign, uid int64) error {
return nil
}
func UpdateCampaignStatus(c *Campaign, s string) error {
return db.Debug().Model(c).UpdateColumn("status", s).Error
}
//DeleteCampaign deletes the specified campaign
func DeleteCampaign(id int64) error {
// Delete all the campaign results

View File

@ -17,6 +17,13 @@ var err error
var ErrUsernameTaken = errors.New("username already taken")
var Logger = log.New(os.Stdout, " ", log.Ldate|log.Ltime|log.Lshortfile)
const (
IN_PROGRESS string = "In progress"
QUEUED string = "Queued"
COMPLETE string = "Completed"
ERROR string = "Error"
)
// Flash is used to hold flash information for use in templates.
type Flash struct {
Type string
@ -27,7 +34,7 @@ type Flash struct {
// It also populates the Gophish Config object
func Setup() error {
db, err = gorm.Open("sqlite3", config.Conf.DBPath)
db.LogMode(true)
db.LogMode(false)
db.SetLogger(Logger)
if err != nil {
Logger.Println(err)

45
worker/worker.go Normal file
View File

@ -0,0 +1,45 @@
package worker
import (
"log"
"os"
"github.com/jordan-wright/gophish/models"
)
var Logger = log.New(os.Stdout, " ", log.Ldate|log.Ltime|log.Lshortfile)
type Worker struct {
Queue chan *models.Campaign
}
func New() *Worker {
return &Worker{
Queue: make(chan *models.Campaign),
}
}
// Start launches the worker to monitor the database for any jobs.
// If a job is found, it launches the job
func (w *Worker) Start() {
Logger.Println("Background Worker Started Successfully - Waiting for Campaigns")
for {
processCampaign(<-w.Queue)
}
}
func processCampaign(c *models.Campaign) {
Logger.Printf("Worker received: %s", c.Name)
err := models.UpdateCampaignStatus(c, models.IN_PROGRESS)
if err != nil {
Logger.Println(err)
}
for _, t := range c.Results {
Logger.Println("Creating email using template")
/*e := email.Email{
Text: []byte(c.Template.Text),
HTML: []byte(c.Template.Html),
}*/
Logger.Printf("Sending Email to %s\n", t.Email)
}
}