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

View File

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

View File

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