From 584d7dbc236fad0e0d79b9b96d9e515677f3d26c Mon Sep 17 00:00:00 2001 From: Jordan Date: Mon, 24 Mar 2014 22:31:33 -0500 Subject: [PATCH] Major refactoring - modularized models into separate files. Removed db package (moved to models) I will be looking to migrate to gorm (instead of gorp) soon! --- auth/auth.go | 13 +- config/config.go | 17 +- controllers/api.go | 31 ++- controllers/route.go | 3 +- db/db.go | 404 --------------------------------------- gophish.go | 9 +- middleware/middleware.go | 6 +- models/campaign.go | 101 ++++++++++ models/group.go | 193 +++++++++++++++++++ models/models.go | 120 ++++++------ models/template.go | 48 +++++ models/user.go | 52 +++++ 12 files changed, 499 insertions(+), 498 deletions(-) delete mode 100644 db/db.go create mode 100644 models/campaign.go create mode 100644 models/group.go create mode 100644 models/template.go create mode 100644 models/user.go diff --git a/auth/auth.go b/auth/auth.go index 1fba1f78..ab268d57 100644 --- a/auth/auth.go +++ b/auth/auth.go @@ -8,12 +8,11 @@ import ( "io" "net/http" - "code.google.com/p/go.crypto/bcrypt" "crypto/rand" + "code.google.com/p/go.crypto/bcrypt" ctx "github.com/gorilla/context" "github.com/gorilla/securecookie" "github.com/gorilla/sessions" - "github.com/jordan-wright/gophish/db" "github.com/jordan-wright/gophish/models" ) @@ -33,8 +32,8 @@ var ErrInvalidPassword = errors.New("Invalid Password") func Login(r *http.Request) (bool, error) { username, password := r.FormValue("username"), r.FormValue("password") session, _ := Store.Get(r, "gophish") - u, err := db.GetUserByUsername(username) - if err != db.ErrUsernameTaken { + u, err := models.GetUserByUsername(username) + if err != models.ErrUsernameTaken { return false, err } //If we've made it here, we should have a valid user stored in u @@ -52,7 +51,7 @@ func Login(r *http.Request) (bool, error) { // Register attempts to register the user given a request. func Register(r *http.Request) (bool, error) { username, password := r.FormValue("username"), r.FormValue("password") - u, err := db.GetUserByUsername(username) + u, err := models.GetUserByUsername(username) // If we have an error which is not simply indicating that no user was found, report it if err != sql.ErrNoRows { return false, err @@ -66,7 +65,7 @@ func Register(r *http.Request) (bool, error) { if err != nil { return false, err } - err = db.Conn.Insert(&u) + err = models.Conn.Insert(&u) if err != nil { return false, err } @@ -94,7 +93,7 @@ func ChangePassword(r *http.Request) error { return err } u.Hash = string(h) - if err = db.PutUser(&u); err != nil { + if err = models.PutUser(&u); err != nil { return err } return nil diff --git a/config/config.go b/config/config.go index 21e6f701..3884c9f7 100644 --- a/config/config.go +++ b/config/config.go @@ -5,11 +5,22 @@ import ( "fmt" "io/ioutil" "os" - - "github.com/jordan-wright/gophish/models" ) -var Conf models.Config +type SMTPServer struct { + Host string `json:"host"` + User string `json:"user"` + Password string `json:"password"` +} + +// Config represents the configuration information. +type Config struct { + URL string `json:"url"` + SMTP SMTPServer `json:"smtp"` + DBPath string `json:"dbpath"` +} + +var Conf Config func init() { // Get the config file diff --git a/controllers/api.go b/controllers/api.go index f5cd2842..17a880d0 100644 --- a/controllers/api.go +++ b/controllers/api.go @@ -11,7 +11,6 @@ import ( ctx "github.com/gorilla/context" "github.com/gorilla/mux" "github.com/jordan-wright/gophish/auth" - "github.com/jordan-wright/gophish/db" "github.com/jordan-wright/gophish/models" ) @@ -41,7 +40,7 @@ func API_Reset(w http.ResponseWriter, r *http.Request) { case r.Method == "POST": u := ctx.Get(r, "user").(models.User) u.APIKey = auth.GenerateSecureKey() - err := db.PutUser(&u) + err := models.PutUser(&u) if err != nil { Flash(w, r, "danger", "Error resetting API Key") } else { @@ -56,7 +55,7 @@ func API_Reset(w http.ResponseWriter, r *http.Request) { func API_Campaigns(w http.ResponseWriter, r *http.Request) { switch { case r.Method == "GET": - cs, err := db.GetCampaigns(ctx.Get(r, "user_id").(int64)) + cs, err := models.GetCampaigns(ctx.Get(r, "user_id").(int64)) if err != nil { fmt.Println(err) } @@ -81,7 +80,7 @@ func API_Campaigns(w http.ResponseWriter, r *http.Request) { c.CreatedDate = time.Now() c.CompletedDate = time.Time{} c.Status = IN_PROGRESS - err = db.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) { return } @@ -101,7 +100,7 @@ func API_Campaigns_Id(w http.ResponseWriter, r *http.Request) { switch { case r.Method == "GET": c := models.Campaign{} - c, err := db.GetCampaign(id, ctx.Get(r, "user_id").(int64)) + c, err := models.GetCampaign(id, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "No campaign found", http.StatusNotFound) { return } @@ -111,11 +110,11 @@ func API_Campaigns_Id(w http.ResponseWriter, r *http.Request) { } writeJSON(w, cj) case r.Method == "DELETE": - _, err := db.GetCampaign(id, ctx.Get(r, "user_id").(int64)) + _, err := models.GetCampaign(id, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "No campaign found", http.StatusNotFound) { return } - err = db.DeleteCampaign(id) + err = models.DeleteCampaign(id) if checkError(err, w, "Error deleting campaign", http.StatusInternalServerError) { return } @@ -153,7 +152,7 @@ RESULT { "name" : "Test Group", func API_Groups(w http.ResponseWriter, r *http.Request) { switch { case r.Method == "GET": - gs, err := db.GetGroups(ctx.Get(r, "user_id").(int64)) + gs, err := models.GetGroups(ctx.Get(r, "user_id").(int64)) if checkError(err, w, "Groups not found", http.StatusNotFound) { return } @@ -176,7 +175,7 @@ func API_Groups(w http.ResponseWriter, r *http.Request) { return } g.ModifiedDate = time.Now() - err = db.PostGroup(&g, ctx.Get(r, "user_id").(int64)) + err = models.PostGroup(&g, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "Error inserting group", http.StatusInternalServerError) { return } @@ -195,7 +194,7 @@ func API_Groups_Id(w http.ResponseWriter, r *http.Request) { id, _ := strconv.ParseInt(vars["id"], 0, 64) switch { case r.Method == "GET": - g, err := db.GetGroup(id, ctx.Get(r, "user_id").(int64)) + g, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "No group found", http.StatusNotFound) { return } @@ -205,17 +204,17 @@ func API_Groups_Id(w http.ResponseWriter, r *http.Request) { } writeJSON(w, gj) case r.Method == "DELETE": - _, err := db.GetGroup(id, ctx.Get(r, "user_id").(int64)) + _, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "No group found", http.StatusNotFound) { return } - err = db.DeleteGroup(id) + err = models.DeleteGroup(id) if checkError(err, w, "Error deleting group", http.StatusInternalServerError) { return } writeJSON(w, []byte("{\"success\" : \"true\"}")) case r.Method == "PUT": - _, err := db.GetGroup(id, ctx.Get(r, "user_id").(int64)) + _, err := models.GetGroup(id, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "No group found", http.StatusNotFound) { return } @@ -230,7 +229,7 @@ func API_Groups_Id(w http.ResponseWriter, r *http.Request) { http.Error(w, "Error: No targets specified", http.StatusBadRequest) return } - err = db.PutGroup(&g, ctx.Get(r, "user_id").(int64)) + err = models.PutGroup(&g, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "Error updating group", http.StatusInternalServerError) { return } @@ -245,7 +244,7 @@ func API_Groups_Id(w http.ResponseWriter, r *http.Request) { func API_Templates(w http.ResponseWriter, r *http.Request) { switch { case r.Method == "GET": - ts, err := db.GetTemplates(ctx.Get(r, "user_id").(int64)) + ts, err := models.GetTemplates(ctx.Get(r, "user_id").(int64)) if checkError(err, w, "Templates not found", http.StatusNotFound) { return } @@ -263,7 +262,7 @@ func API_Templates(w http.ResponseWriter, r *http.Request) { return } t.ModifiedDate = time.Now() - err = db.PostTemplate(&t, ctx.Get(r, "user_id").(int64)) + err = models.PostTemplate(&t, ctx.Get(r, "user_id").(int64)) if checkError(err, w, "Error inserting template", http.StatusInternalServerError) { return } diff --git a/controllers/route.go b/controllers/route.go index 96974fb7..16070b3b 100644 --- a/controllers/route.go +++ b/controllers/route.go @@ -9,7 +9,6 @@ import ( "github.com/gorilla/mux" "github.com/gorilla/sessions" "github.com/jordan-wright/gophish/auth" - "github.com/jordan-wright/gophish/db" mid "github.com/jordan-wright/gophish/middleware" "github.com/jordan-wright/gophish/models" "github.com/justinas/nosurf" @@ -92,7 +91,7 @@ func Register(w http.ResponseWriter, r *http.Request) { } else { // Check the error m := "" - if err == db.ErrUsernameTaken { + if err == models.ErrUsernameTaken { m = "Username already taken" } else { m = "Unknown error - please try again" diff --git a/db/db.go b/db/db.go deleted file mode 100644 index 6dc81720..00000000 --- a/db/db.go +++ /dev/null @@ -1,404 +0,0 @@ -package db - -import ( - "database/sql" - "errors" - "fmt" - "log" - "net/mail" - "os" - "time" - - "github.com/coopernurse/gorp" - "github.com/jordan-wright/gophish/config" - "github.com/jordan-wright/gophish/models" - _ "github.com/mattn/go-sqlite3" -) - -var Conn *gorp.DbMap -var DB *sql.DB -var err error -var ErrUsernameTaken = errors.New("Username already taken") -var Logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile) - -// Setup initializes the Conn object -// It also populates the Gophish Config object -func Setup() error { - DB, err := sql.Open("sqlite3", config.Conf.DBPath) - Conn = &gorp.DbMap{Db: DB, Dialect: gorp.SqliteDialect{}} - //If the file already exists, delete it and recreate it - _, err = os.Stat(config.Conf.DBPath) - Conn.AddTableWithName(models.User{}, "users").SetKeys(true, "Id") - Conn.AddTableWithName(models.Campaign{}, "campaigns").SetKeys(true, "Id") - Conn.AddTableWithName(models.Group{}, "groups").SetKeys(true, "Id") - Conn.AddTableWithName(models.Template{}, "templates").SetKeys(true, "Id") - if err != nil { - Logger.Println("Database not found, recreating...") - createTablesSQL := []string{ - //Create tables - `CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT NOT NULL, hash VARCHAR(60) NOT NULL, api_key VARCHAR(32), UNIQUE(username), UNIQUE(api_key));`, - `CREATE TABLE campaigns (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, created_date TIMESTAMP NOT NULL, completed_date TIMESTAMP, template TEXT, status TEXT NOT NULL);`, - `CREATE TABLE targets (id INTEGER PRIMARY KEY AUTOINCREMENT, email TEXT NOT NULL, UNIQUE(email));`, - `CREATE TABLE groups (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, modified_date TIMESTAMP NOT NULL);`, - `CREATE TABLE campaign_results (cid INTEGER NOT NULL, email TEXT NOT NULL, status TEXT NOT NULL, FOREIGN KEY (cid) REFERENCES campaigns(id), UNIQUE(cid, email, status))`, - `CREATE TABLE templates (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, modified_date TIMESTAMP NOT NULL, html TEXT NOT NULL, text TEXT NOT NULL);`, - `CREATE TABLE files (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, path TEXT NOT NULL);`, - `CREATE TABLE user_campaigns (uid INTEGER NOT NULL, cid INTEGER NOT NULL, FOREIGN KEY (uid) REFERENCES users(id), FOREIGN KEY (cid) REFERENCES campaigns(id), UNIQUE(uid, cid))`, - `CREATE TABLE user_groups (uid INTEGER NOT NULL, gid INTEGER NOT NULL, FOREIGN KEY (uid) REFERENCES users(id), FOREIGN KEY (gid) REFERENCES groups(id), UNIQUE(uid, gid))`, - `CREATE TABLE group_targets (gid INTEGER NOT NULL, tid INTEGER NOT NULL, FOREIGN KEY (gid) REFERENCES groups(id), FOREIGN KEY (tid) REFERENCES targets(id), UNIQUE(gid, tid));`, - `CREATE TABLE user_templates (uid INTEGER NOT NULL, tid INTEGER NOT NULL, FOREIGN KEY (uid) REFERENCES users(id), FOREIGN KEY (tid) REFERENCES templates(id), UNIQUE(uid, tid));`, - `CREATE TABLE template_files (tid INTEGER NOT NULL, fid INTEGER NOT NULL, FOREIGN KEY (tid) REFERENCES templates(id), FOREIGN KEY(fid) REFERENCES files(id), UNIQUE(tid, fid));`, - } - Logger.Printf("Creating db at %s\n", config.Conf.DBPath) - //Create the tables needed - for _, stmt := range createTablesSQL { - _, err = DB.Exec(stmt) - if err != nil { - return err - } - } - //Create the default user - init_user := models.User{ - Username: "admin", - Hash: "$2a$10$IYkPp0.QsM81lYYPrQx6W.U6oQGw7wMpozrKhKAHUBVL4mkm/EvAS", //gophish - APIKey: "12345678901234567890123456789012", - } - Conn.Insert(&init_user) - if err != nil { - Logger.Println(err) - } - } - return nil -} - -// API Functions (GET, POST, PUT, DELETE) - -// GetUser returns the user that the given id corresponds to. If no user is found, an -// error is thrown. -func GetUser(id int64) (models.User, error) { - u := models.User{} - err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE id=?", id) - if err != nil { - return u, err - } - return u, nil -} - -// GetUserByAPIKey returns the user that the given API Key corresponds to. If no user is found, an -// error is thrown. -func GetUserByAPIKey(key []byte) (models.User, error) { - u := models.User{} - err := Conn.SelectOne(&u, "SELECT id, username, api_key FROM Users WHERE apikey=?", key) - if err != nil { - return u, err - } - return u, nil -} - -// GetUserByUsername returns the user that the given username corresponds to. If no user is found, an -// error is thrown. -func GetUserByUsername(username string) (models.User, error) { - u := models.User{} - err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE username=?", username) - if err != sql.ErrNoRows { - return u, ErrUsernameTaken - } else if err != nil { - return u, err - } - return u, nil -} - -// PutUser updates the given user -func PutUser(u *models.User) error { - _, err := Conn.Update(u) - return err -} - -// GetCampaigns returns the campaigns owned by the given user. -func GetCampaigns(uid int64) ([]models.Campaign, error) { - cs := []models.Campaign{} - _, err := Conn.Select(&cs, "SELECT c.id, name, created_date, completed_date, status, template FROM campaigns c, user_campaigns uc, users u WHERE uc.uid=u.id AND uc.cid=c.id AND u.id=?", uid) - for i, _ := range cs { - _, err = Conn.Select(&cs[i].Results, "SELECT r.email, r.status FROM campaign_results r WHERE r.cid=?", cs[i].Id) - } - return cs, err -} - -// GetCampaign returns the campaign, if it exists, specified by the given id and user_id. -func GetCampaign(id int64, uid int64) (models.Campaign, error) { - c := models.Campaign{} - err := Conn.SelectOne(&c, "SELECT c.id, name, created_date, completed_date, status, template FROM campaigns c, user_campaigns uc, users u WHERE uc.uid=u.id AND uc.cid=c.id AND c.id=? AND u.id=?", id, uid) - if err != nil { - return c, err - } - _, err = Conn.Select(&c.Results, "SELECT r.email, r.status FROM campaign_results r WHERE r.cid=?", c.Id) - return c, err -} - -// PostCampaign inserts a campaign and all associated records into the database. -func PostCampaign(c *models.Campaign, uid int64) error { - // Check to make sure all the groups already exist - for i, g := range c.Groups { - c.Groups[i], err = GetGroupByName(g.Name, uid) - if err == sql.ErrNoRows { - Logger.Printf("Error - Group %s does not exist", g.Name) - return err - } else if err != nil { - Logger.Println(err) - return err - } - } - // Insert into the DB - err = Conn.Insert(c) - if err != nil { - Logger.Println(err) - return err - } - // Insert all the results - for _, g := range c.Groups { - // Insert a result for each target in the group - for _, t := range g.Targets { - r := models.Result{Target: t, Status: "Unknown"} - c.Results = append(c.Results, r) - fmt.Printf("%v", c.Results) - _, err = Conn.Exec("INSERT INTO campaign_results VALUES (?,?,?)", c.Id, r.Email, r.Status) - if err != nil { - Logger.Printf("Error adding result record for target %s\n", t.Email) - Logger.Println(err) - } - } - } - _, err = Conn.Exec("INSERT OR IGNORE INTO user_campaigns VALUES (?,?)", uid, c.Id) - if err != nil { - Logger.Printf("Error adding many-many mapping for campaign %s\n", c.Name) - } - return nil -} - -func DeleteCampaign(id int64) error { - // Delete all the campaign_results entries for this group - _, err := Conn.Exec("DELETE FROM campaign_results WHERE cid=?", id) - if err != nil { - return err - } - // Delete the reference to the campaign in the user_campaigns table - _, err = Conn.Exec("DELETE FROM user_campaigns WHERE cid=?", id) - if err != nil { - return err - } - // Delete the campaign itself - _, err = Conn.Exec("DELETE FROM campaigns WHERE id=?", id) - return err -} - -// GetGroups returns the groups owned by the given user. -func GetGroups(uid int64) ([]models.Group, error) { - gs := []models.Group{} - _, err := Conn.Select(&gs, "SELECT g.id, g.name, g.modified_date FROM groups g, user_groups ug, users u WHERE ug.uid=u.id AND ug.gid=g.id AND u.id=?", uid) - if err != nil { - Logger.Println(err) - return gs, err - } - for i, _ := range gs { - _, err := Conn.Select(&gs[i].Targets, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", gs[i].Id) - if err != nil { - Logger.Println(err) - } - } - return gs, nil -} - -// GetGroup returns the group, if it exists, specified by the given id and user_id. -func GetGroup(id int64, uid int64) (models.Group, error) { - g := models.Group{} - err := Conn.SelectOne(&g, "SELECT g.id, g.name, g.modified_date FROM groups g, user_groups ug, users u WHERE ug.uid=u.id AND ug.gid=g.id AND g.id=? AND u.id=?", id, uid) - if err != nil { - Logger.Println(err) - return g, err - } - _, err = Conn.Select(&g.Targets, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", g.Id) - if err != nil { - Logger.Println(err) - } - return g, nil -} - -// GetGroupByName returns the group, if it exists, specified by the given name and user_id. -func GetGroupByName(n string, uid int64) (models.Group, error) { - g := models.Group{} - err := Conn.SelectOne(&g, "SELECT g.id, g.name, g.modified_date FROM groups g, user_groups ug, users u WHERE ug.uid=u.id AND ug.gid=g.id AND g.name=? AND u.id=?", n, uid) - if err != nil { - Logger.Println(err) - return g, err - } - _, err = Conn.Select(&g.Targets, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", g.Id) - if err != nil { - Logger.Println(err) - } - return g, nil -} - -// PostGroup creates a new group in the database. -func PostGroup(g *models.Group, uid int64) error { - // Insert into the DB - err = Conn.Insert(g) - if err != nil { - Logger.Println(err) - return err - } - // Now, let's add the user->user_groups->group mapping - _, err = Conn.Exec("INSERT OR IGNORE INTO user_groups VALUES (?,?)", uid, g.Id) - if err != nil { - Logger.Printf("Error adding many-many mapping for group %s\n", g.Name) - } - for _, t := range g.Targets { - insertTargetIntoGroup(t, g.Id) - } - return nil -} - -// PutGroup updates the given group if found in the database. -func PutGroup(g *models.Group, uid int64) error { - // Update all the foreign keys, and many to many relationships - // We will only delete the group->targets entries. We keep the actual targets - // since they are needed by the Results table - // Get all the targets currently in the database for the group - ts := []models.Target{} - _, err = Conn.Select(&ts, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", g.Id) - if err != nil { - Logger.Printf("Error getting targets from group ID: %d", g.Id) - return err - } - // Enumerate through, removing any entries that are no longer in the group - // For every target in the database - tExists := false - for _, t := range ts { - tExists = false - // Is the target still in the group? - for _, nt := range g.Targets { - if t.Email == nt.Email { - tExists = true - break - } - } - // If the target does not exist in the group any longer, we delete it - if !tExists { - _, err = Conn.Exec("DELETE FROM group_targets WHERE gid=? AND tid=?", g.Id, t.Id) - if err != nil { - Logger.Printf("Error deleting email %s\n", t.Email) - } - } - } - // Insert any entries that are not in the database - // For every target in the new group - for _, nt := range g.Targets { - // Check and see if the target already exists in the db - tExists = false - for _, t := range ts { - if t.Email == nt.Email { - tExists = true - break - } - } - // If the target is not in the db, we add it - if !tExists { - insertTargetIntoGroup(nt, g.Id) - } - } - // Update the group - g.ModifiedDate = time.Now() - _, err = Conn.Update(g) - if err != nil { - Logger.Println(err) - return err - } - return nil -} - -// GetTemplates returns the templates owned by the given user. -func GetTemplates(uid int64) ([]models.Template, error) { - ts := []models.Template{} - _, err := Conn.Select(&ts, "SELECT t.id, t.name, t.modified_date, t.text, t.html FROM templates t, user_templates ut, users u WHERE ut.uid=u.id AND ut.tid=t.id AND u.id=?", uid) - return ts, err -} - -// GetTemplate returns the template, if it exists, specified by the given id and user_id. -func GetTemplate(id int64, uid int64) (models.Template, error) { - t := models.Template{} - err := Conn.SelectOne(&t, "SELECT t.id, t.name, t.modified_date, t.text, t.html FROM templates t, user_templates ut, users u WHERE ut.uid=u.id AND ut.tid=t.id AND t.id=? AND u.id=?", id, uid) - if err != nil { - return t, err - } - return t, err -} - -// PostTemplate creates a new template in the database. -func PostTemplate(t *models.Template, uid int64) error { - // Insert into the DB - err = Conn.Insert(t) - if err != nil { - Logger.Println(err) - return err - } - // Now, let's add the user->user_templates->template mapping - _, err = Conn.Exec("INSERT OR IGNORE INTO user_templates VALUES (?,?)", uid, t.Id) - if err != nil { - Logger.Printf("Error adding many-many mapping for template %s\n", t.Name) - } - return nil -} - -func PutTemplate(t *models.Template, uid int64) error { - return nil -} - -func insertTargetIntoGroup(t models.Target, gid int64) error { - if _, err = mail.ParseAddress(t.Email); err != nil { - Logger.Printf("Invalid email %s\n", t.Email) - return err - } - trans, err := Conn.Begin() - if err != nil { - Logger.Println(err) - return err - } - _, err = trans.Exec("INSERT OR IGNORE INTO targets VALUES (null, ?)", t.Email) - if err != nil { - Logger.Printf("Error adding email: %s\n", t.Email) - return err - } - // Bug: res.LastInsertId() does not work for this, so we need to select it manually (how frustrating.) - t.Id, err = trans.SelectInt("SELECT id FROM targets WHERE email=?", t.Email) - if err != nil { - Logger.Printf("Error getting id for email: %s\n", t.Email) - return err - } - _, err = trans.Exec("INSERT OR IGNORE INTO group_targets VALUES (?,?)", gid, t.Id) - if err != nil { - Logger.Printf("Error adding many-many mapping for %s\n", t.Email) - return err - } - err = trans.Commit() - if err != nil { - Logger.Printf("Error committing db changes\n") - return err - } - return nil -} - -// DeleteGroup deletes a given group by group ID and user ID -func DeleteGroup(id int64) error { - // Delete all the group_targets entries for this group - _, err := Conn.Exec("DELETE FROM group_targets WHERE gid=?", id) - if err != nil { - return err - } - // Delete the reference to the group in the user_group table - _, err = Conn.Exec("DELETE FROM user_groups WHERE gid=?", id) - if err != nil { - return err - } - // Delete the group itself - _, err = Conn.Exec("DELETE FROM groups WHERE id=?", id) - return err -} diff --git a/gophish.go b/gophish.go index 7b003dda..0a69e67c 100644 --- a/gophish.go +++ b/gophish.go @@ -31,17 +31,16 @@ import ( "github.com/jordan-wright/gophish/config" "github.com/jordan-wright/gophish/controllers" - "github.com/jordan-wright/gophish/db" "github.com/jordan-wright/gophish/middleware" + "github.com/jordan-wright/gophish/models" ) func main() { //Setup the global variables and settings - err := db.Setup() - defer db.DB.Close() - if err != nil { + defer models.DB.Close() + /* if err != nil { fmt.Println(err) - } + }*/ fmt.Printf("Gophish server started at http://%s\n", config.Conf.URL) http.Handle("/", controllers.Use(controllers.CreateRouter().ServeHTTP, middleware.GetContext)) http.ListenAndServe(config.Conf.URL, nil) diff --git a/middleware/middleware.go b/middleware/middleware.go index 13110a89..0d69d905 100644 --- a/middleware/middleware.go +++ b/middleware/middleware.go @@ -5,7 +5,7 @@ import ( ctx "github.com/gorilla/context" "github.com/jordan-wright/gophish/auth" - "github.com/jordan-wright/gophish/db" + "github.com/jordan-wright/gophish/models" ) // GetContext wraps each request in a function which fills in the context for a given request. @@ -24,7 +24,7 @@ func GetContext(handler http.Handler) http.HandlerFunc { // Put the session in the context so that ctx.Set(r, "session", session) if id, ok := session.Values["id"]; ok { - u, err := db.GetUser(id.(int64)) + u, err := models.GetUser(id.(int64)) if err != nil { ctx.Set(r, "user", nil) } @@ -52,7 +52,7 @@ func RequireAPIKey(handler http.Handler) http.HandlerFunc { if ak == "" { JSONError(w, 400, "API Key not set") } else { - id, err := db.Conn.SelectInt("SELECT id FROM users WHERE api_key=?", ak) + id, err := models.Conn.SelectInt("SELECT id FROM users WHERE api_key=?", ak) if id == 0 || err != nil { JSONError(w, 400, "Invalid API Key") return diff --git a/models/campaign.go b/models/campaign.go new file mode 100644 index 00000000..9db54ee0 --- /dev/null +++ b/models/campaign.go @@ -0,0 +1,101 @@ +package models + +import ( + "database/sql" + "fmt" + "time" +) + +//Campaign is a struct representing a created campaign +type Campaign struct { + Id int64 `json:"id"` + Name string `json:"name"` + CreatedDate time.Time `json:"created_date" db:"created_date"` + CompletedDate time.Time `json:"completed_date" db:"completed_date"` + Template string `json:"template"` //This may change + Status string `json:"status"` + Results []Result `json:"results,omitempty" db:"-"` + Groups []Group `json:"groups,omitempty" db:"-"` +} + +type Result struct { + Target + Status string `json:"status"` +} + +// GetCampaigns returns the campaigns owned by the given user. +func GetCampaigns(uid int64) ([]Campaign, error) { + cs := []Campaign{} + _, err := Conn.Select(&cs, "SELECT c.id, name, created_date, completed_date, status, template FROM campaigns c, user_campaigns uc, users u WHERE uc.uid=u.id AND uc.cid=c.id AND u.id=?", uid) + for i, _ := range cs { + _, err = Conn.Select(&cs[i].Results, "SELECT r.email, r.status FROM campaign_results r WHERE r.cid=?", cs[i].Id) + } + return cs, err +} + +// GetCampaign returns the campaign, if it exists, specified by the given id and user_id. +func GetCampaign(id int64, uid int64) (Campaign, error) { + c := Campaign{} + err := Conn.SelectOne(&c, "SELECT c.id, name, created_date, completed_date, status, template FROM campaigns c, user_campaigns uc, users u WHERE uc.uid=u.id AND uc.cid=c.id AND c.id=? AND u.id=?", id, uid) + if err != nil { + return c, err + } + _, err = Conn.Select(&c.Results, "SELECT r.email, r.status FROM campaign_results r WHERE r.cid=?", c.Id) + return c, err +} + +// PostCampaign inserts a campaign and all associated records into the database. +func PostCampaign(c *Campaign, uid int64) error { + // Check to make sure all the groups already exist + for i, g := range c.Groups { + c.Groups[i], err = GetGroupByName(g.Name, uid) + if err == sql.ErrNoRows { + Logger.Printf("Error - Group %s does not exist", g.Name) + return err + } else if err != nil { + Logger.Println(err) + return err + } + } + // Insert into the DB + err = Conn.Insert(c) + if err != nil { + Logger.Println(err) + return err + } + // Insert all the results + for _, g := range c.Groups { + // Insert a result for each target in the group + for _, t := range g.Targets { + r := Result{Target: t, Status: "Unknown"} + c.Results = append(c.Results, r) + fmt.Printf("%v", c.Results) + _, err = Conn.Exec("INSERT INTO campaign_results VALUES (?,?,?)", c.Id, r.Email, r.Status) + if err != nil { + Logger.Printf("Error adding result record for target %s\n", t.Email) + Logger.Println(err) + } + } + } + _, err = Conn.Exec("INSERT OR IGNORE INTO user_campaigns VALUES (?,?)", uid, c.Id) + if err != nil { + Logger.Printf("Error adding many-many mapping for campaign %s\n", c.Name) + } + return nil +} + +func DeleteCampaign(id int64) error { + // Delete all the campaign_results entries for this group + _, err := Conn.Exec("DELETE FROM campaign_results WHERE cid=?", id) + if err != nil { + return err + } + // Delete the reference to the campaign in the user_campaigns table + _, err = Conn.Exec("DELETE FROM user_campaigns WHERE cid=?", id) + if err != nil { + return err + } + // Delete the campaign itself + _, err = Conn.Exec("DELETE FROM campaigns WHERE id=?", id) + return err +} diff --git a/models/group.go b/models/group.go new file mode 100644 index 00000000..1305dda4 --- /dev/null +++ b/models/group.go @@ -0,0 +1,193 @@ +package models + +import ( + "net/mail" + "time" +) + +type Group struct { + Id int64 `json:"id"` + Name string `json:"name"` + ModifiedDate time.Time `json:"modified_date" db:"modified_date"` + Targets []Target `json:"targets" db:"-"` +} + +type Target struct { + Id int64 `json:"-"` + Email string `json:"email"` +} + +// GetGroups returns the groups owned by the given user. +func GetGroups(uid int64) ([]Group, error) { + gs := []Group{} + _, err := Conn.Select(&gs, "SELECT g.id, g.name, g.modified_date FROM groups g, user_groups ug, users u WHERE ug.uid=u.id AND ug.gid=g.id AND u.id=?", uid) + if err != nil { + Logger.Println(err) + return gs, err + } + for i, _ := range gs { + _, err := Conn.Select(&gs[i].Targets, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", gs[i].Id) + if err != nil { + Logger.Println(err) + } + } + return gs, nil +} + +// GetGroup returns the group, if it exists, specified by the given id and user_id. +func GetGroup(id int64, uid int64) (Group, error) { + g := Group{} + err := Conn.SelectOne(&g, "SELECT g.id, g.name, g.modified_date FROM groups g, user_groups ug, users u WHERE ug.uid=u.id AND ug.gid=g.id AND g.id=? AND u.id=?", id, uid) + if err != nil { + Logger.Println(err) + return g, err + } + _, err = Conn.Select(&g.Targets, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", g.Id) + if err != nil { + Logger.Println(err) + } + return g, nil +} + +// GetGroupByName returns the group, if it exists, specified by the given name and user_id. +func GetGroupByName(n string, uid int64) (Group, error) { + g := Group{} + err := Conn.SelectOne(&g, "SELECT g.id, g.name, g.modified_date FROM groups g, user_groups ug, users u WHERE ug.uid=u.id AND ug.gid=g.id AND g.name=? AND u.id=?", n, uid) + if err != nil { + Logger.Println(err) + return g, err + } + _, err = Conn.Select(&g.Targets, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", g.Id) + if err != nil { + Logger.Println(err) + } + return g, nil +} + +// PostGroup creates a new group in the database. +func PostGroup(g *Group, uid int64) error { + // Insert into the DB + err = Conn.Insert(g) + if err != nil { + Logger.Println(err) + return err + } + // Now, let's add the user->user_groups->group mapping + _, err = Conn.Exec("INSERT OR IGNORE INTO user_groups VALUES (?,?)", uid, g.Id) + if err != nil { + Logger.Printf("Error adding many-many mapping for group %s\n", g.Name) + } + for _, t := range g.Targets { + insertTargetIntoGroup(t, g.Id) + } + return nil +} + +// PutGroup updates the given group if found in the database. +func PutGroup(g *Group, uid int64) error { + // Update all the foreign keys, and many to many relationships + // We will only delete the group->targets entries. We keep the actual targets + // since they are needed by the Results table + // Get all the targets currently in the database for the group + ts := []Target{} + _, err = Conn.Select(&ts, "SELECT t.id, t.email FROM targets t, group_targets gt WHERE gt.gid=? AND gt.tid=t.id", g.Id) + if err != nil { + Logger.Printf("Error getting targets from group ID: %d", g.Id) + return err + } + // Enumerate through, removing any entries that are no longer in the group + // For every target in the database + tExists := false + for _, t := range ts { + tExists = false + // Is the target still in the group? + for _, nt := range g.Targets { + if t.Email == nt.Email { + tExists = true + break + } + } + // If the target does not exist in the group any longer, we delete it + if !tExists { + _, err = Conn.Exec("DELETE FROM group_targets WHERE gid=? AND tid=?", g.Id, t.Id) + if err != nil { + Logger.Printf("Error deleting email %s\n", t.Email) + } + } + } + // Insert any entries that are not in the database + // For every target in the new group + for _, nt := range g.Targets { + // Check and see if the target already exists in the db + tExists = false + for _, t := range ts { + if t.Email == nt.Email { + tExists = true + break + } + } + // If the target is not in the db, we add it + if !tExists { + insertTargetIntoGroup(nt, g.Id) + } + } + // Update the group + g.ModifiedDate = time.Now() + _, err = Conn.Update(g) + if err != nil { + Logger.Println(err) + return err + } + return nil +} + +func insertTargetIntoGroup(t Target, gid int64) error { + if _, err = mail.ParseAddress(t.Email); err != nil { + Logger.Printf("Invalid email %s\n", t.Email) + return err + } + trans, err := Conn.Begin() + if err != nil { + Logger.Println(err) + return err + } + _, err = trans.Exec("INSERT OR IGNORE INTO targets VALUES (null, ?)", t.Email) + if err != nil { + Logger.Printf("Error adding email: %s\n", t.Email) + return err + } + // Bug: res.LastInsertId() does not work for this, so we need to select it manually (how frustrating.) + t.Id, err = trans.SelectInt("SELECT id FROM targets WHERE email=?", t.Email) + if err != nil { + Logger.Printf("Error getting id for email: %s\n", t.Email) + return err + } + _, err = trans.Exec("INSERT OR IGNORE INTO group_targets VALUES (?,?)", gid, t.Id) + if err != nil { + Logger.Printf("Error adding many-many mapping for %s\n", t.Email) + return err + } + err = trans.Commit() + if err != nil { + Logger.Printf("Error committing db changes\n") + return err + } + return nil +} + +// DeleteGroup deletes a given group by group ID and user ID +func DeleteGroup(id int64) error { + // Delete all the group_targets entries for this group + _, err := Conn.Exec("DELETE FROM group_targets WHERE gid=?", id) + if err != nil { + return err + } + // Delete the reference to the group in the user_group table + _, err = Conn.Exec("DELETE FROM user_groups WHERE gid=?", id) + if err != nil { + return err + } + // Delete the group itself + _, err = Conn.Exec("DELETE FROM groups WHERE id=?", id) + return err +} diff --git a/models/models.go b/models/models.go index fb4f354a..f1cd123e 100644 --- a/models/models.go +++ b/models/models.go @@ -1,29 +1,70 @@ package models -import +import ( + "database/sql" + "errors" + "log" + "os" -// SMTPServer is used to provide a default SMTP server preference. -"time" + "github.com/coopernurse/gorp" + "github.com/jordan-wright/gophish/config" + _ "github.com/mattn/go-sqlite3" +) -type SMTPServer struct { - Host string `json:"host"` - User string `json:"user"` - Password string `json:"password"` -} +var Conn *gorp.DbMap +var DB *sql.DB +var err error +var ErrUsernameTaken = errors.New("Username already taken") +var Logger = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lshortfile) -// Config represents the configuration information. -type Config struct { - URL string `json:"url"` - SMTP SMTPServer `json:"smtp"` - DBPath string `json:"dbpath"` -} - -// User represents the user model for gophish. -type User struct { - Id int64 `json:"id"` - Username string `json:"username"` - Hash string `json:"-"` - APIKey string `json:"api_key" db:"api_key"` +// Setup initializes the Conn object +// It also populates the Gophish Config object +func init() { + DB, err := sql.Open("sqlite3", config.Conf.DBPath) + Conn = &gorp.DbMap{Db: DB, Dialect: gorp.SqliteDialect{}} + //If the file already exists, delete it and recreate it + _, err = os.Stat(config.Conf.DBPath) + Conn.AddTableWithName(User{}, "users").SetKeys(true, "Id") + Conn.AddTableWithName(Campaign{}, "campaigns").SetKeys(true, "Id") + Conn.AddTableWithName(Group{}, "groups").SetKeys(true, "Id") + Conn.AddTableWithName(Template{}, "templates").SetKeys(true, "Id") + if err != nil { + Logger.Println("Database not found, recreating...") + createTablesSQL := []string{ + //Create tables + `CREATE TABLE users (id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT NOT NULL, hash VARCHAR(60) NOT NULL, api_key VARCHAR(32), UNIQUE(username), UNIQUE(api_key));`, + `CREATE TABLE campaigns (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, created_date TIMESTAMP NOT NULL, completed_date TIMESTAMP, template TEXT, status TEXT NOT NULL);`, + `CREATE TABLE targets (id INTEGER PRIMARY KEY AUTOINCREMENT, email TEXT NOT NULL, UNIQUE(email));`, + `CREATE TABLE groups (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, modified_date TIMESTAMP NOT NULL);`, + `CREATE TABLE campaign_results (cid INTEGER NOT NULL, email TEXT NOT NULL, status TEXT NOT NULL, FOREIGN KEY (cid) REFERENCES campaigns(id), UNIQUE(cid, email, status))`, + `CREATE TABLE templates (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, modified_date TIMESTAMP NOT NULL, html TEXT NOT NULL, text TEXT NOT NULL);`, + `CREATE TABLE files (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, path TEXT NOT NULL);`, + `CREATE TABLE user_campaigns (uid INTEGER NOT NULL, cid INTEGER NOT NULL, FOREIGN KEY (uid) REFERENCES users(id), FOREIGN KEY (cid) REFERENCES campaigns(id), UNIQUE(uid, cid))`, + `CREATE TABLE user_groups (uid INTEGER NOT NULL, gid INTEGER NOT NULL, FOREIGN KEY (uid) REFERENCES users(id), FOREIGN KEY (gid) REFERENCES groups(id), UNIQUE(uid, gid))`, + `CREATE TABLE group_targets (gid INTEGER NOT NULL, tid INTEGER NOT NULL, FOREIGN KEY (gid) REFERENCES groups(id), FOREIGN KEY (tid) REFERENCES targets(id), UNIQUE(gid, tid));`, + `CREATE TABLE user_templates (uid INTEGER NOT NULL, tid INTEGER NOT NULL, FOREIGN KEY (uid) REFERENCES users(id), FOREIGN KEY (tid) REFERENCES templates(id), UNIQUE(uid, tid));`, + `CREATE TABLE template_files (tid INTEGER NOT NULL, fid INTEGER NOT NULL, FOREIGN KEY (tid) REFERENCES templates(id), FOREIGN KEY(fid) REFERENCES files(id), UNIQUE(tid, fid));`, + } + Logger.Printf("Creating db at %s\n", config.Conf.DBPath) + //Create the tables needed + for _, stmt := range createTablesSQL { + _, err = DB.Exec(stmt) + if err != nil { + /* return nil, err*/ + } + } + //Create the default user + init_user := User{ + Username: "admin", + Hash: "$2a$10$IYkPp0.QsM81lYYPrQx6W.U6oQGw7wMpozrKhKAHUBVL4mkm/EvAS", //gophish + APIKey: "12345678901234567890123456789012", + } + Conn.Insert(&init_user) + if err != nil { + Logger.Println(err) + } + } + /* return Conn, nil*/ } // Flash is used to hold flash information for use in templates. @@ -31,40 +72,3 @@ type Flash struct { Type string Message string } - -//Campaign is a struct representing a created campaign -type Campaign struct { - Id int64 `json:"id"` - Name string `json:"name"` - CreatedDate time.Time `json:"created_date" db:"created_date"` - CompletedDate time.Time `json:"completed_date" db:"completed_date"` - Template string `json:"template"` //This may change - Status string `json:"status"` - Results []Result `json:"results,omitempty" db:"-"` - Groups []Group `json:"groups,omitempty" db:"-"` -} - -type Result struct { - Target - Status string `json:"status"` -} - -type Group struct { - Id int64 `json:"id"` - Name string `json:"name"` - ModifiedDate time.Time `json:"modified_date" db:"modified_date"` - Targets []Target `json:"targets" db:"-"` -} - -type Target struct { - Id int64 `json:"-"` - Email string `json:"email"` -} - -type Template struct { - Id int64 `json:"id"` - Name string `json:"name" db:"name"` - Text string `json:"text" db:"text"` - Html string `json:"html" db:"html"` - ModifiedDate time.Time `json:"modified_date" db:"modified_date"` -} diff --git a/models/template.go b/models/template.go new file mode 100644 index 00000000..39f01dfd --- /dev/null +++ b/models/template.go @@ -0,0 +1,48 @@ +package models + +import "time" + +type Template struct { + Id int64 `json:"id"` + Name string `json:"name" db:"name"` + Text string `json:"text" db:"text"` + Html string `json:"html" db:"html"` + ModifiedDate time.Time `json:"modified_date" db:"modified_date"` +} + +// GetTemplates returns the templates owned by the given user. +func GetTemplates(uid int64) ([]Template, error) { + ts := []Template{} + _, err := Conn.Select(&ts, "SELECT t.id, t.name, t.modified_date, t.text, t.html FROM templates t, user_templates ut, users u WHERE ut.uid=u.id AND ut.tid=t.id AND u.id=?", uid) + return ts, err +} + +// GetTemplate returns the template, if it exists, specified by the given id and user_id. +func GetTemplate(id int64, uid int64) (Template, error) { + t := Template{} + err := Conn.SelectOne(&t, "SELECT t.id, t.name, t.modified_date, t.text, t.html FROM templates t, user_templates ut, users u WHERE ut.uid=u.id AND ut.tid=t.id AND t.id=? AND u.id=?", id, uid) + if err != nil { + return t, err + } + return t, err +} + +// PostTemplate creates a new template in the database. +func PostTemplate(t *Template, uid int64) error { + // Insert into the DB + err = Conn.Insert(t) + if err != nil { + Logger.Println(err) + return err + } + // Now, let's add the user->user_templates->template mapping + _, err = Conn.Exec("INSERT OR IGNORE INTO user_templates VALUES (?,?)", uid, t.Id) + if err != nil { + Logger.Printf("Error adding many-many mapping for template %s\n", t.Name) + } + return nil +} + +func PutTemplate(t *Template, uid int64) error { + return nil +} diff --git a/models/user.go b/models/user.go new file mode 100644 index 00000000..2f329905 --- /dev/null +++ b/models/user.go @@ -0,0 +1,52 @@ +package models + +import "database/sql" + +// User represents the user model for gophish. +type User struct { + Id int64 `json:"id"` + Username string `json:"username"` + Hash string `json:"-"` + APIKey string `json:"api_key" db:"api_key"` +} + +// GetUser returns the user that the given id corresponds to. If no user is found, an +// error is thrown. +func GetUser(id int64) (User, error) { + u := User{} + err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE id=?", id) + if err != nil { + return u, err + } + return u, nil +} + +// GetUserByAPIKey returns the user that the given API Key corresponds to. If no user is found, an +// error is thrown. +func GetUserByAPIKey(key []byte) (User, error) { + u := User{} + err := Conn.SelectOne(&u, "SELECT id, username, api_key FROM Users WHERE apikey=?", key) + if err != nil { + return u, err + } + return u, nil +} + +// GetUserByUsername returns the user that the given username corresponds to. If no user is found, an +// error is thrown. +func GetUserByUsername(username string) (User, error) { + u := User{} + err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE username=?", username) + if err != sql.ErrNoRows { + return u, ErrUsernameTaken + } else if err != nil { + return u, err + } + return u, nil +} + +// PutUser updates the given user +func PutUser(u *User) error { + _, err := Conn.Update(u) + return err +}