From e137126a90f8fcda941d790ca2f26eda6cd120c9 Mon Sep 17 00:00:00 2001 From: Jordan Date: Tue, 25 Mar 2014 23:53:51 -0500 Subject: [PATCH 1/8] Working on gorm integration TODO: [ ] Finish up groups (many-to-many with group_targets) [ ] Convert Template models --- auth/auth.go | 4 +- controllers/api.go | 2 +- middleware/middleware.go | 7 ++-- models/campaign.go | 65 ++++++++++++++++++++++----------- models/group.go | 74 +++++++++++++++++++++++++++++-------- models/models.go | 79 ++++++++++++++++++++++++++++++++++++++-- models/template.go | 8 ++-- models/user.go | 18 +++++---- templates/base.html | 2 +- 9 files changed, 199 insertions(+), 60 deletions(-) diff --git a/auth/auth.go b/auth/auth.go index ab268d57..d6409c72 100644 --- a/auth/auth.go +++ b/auth/auth.go @@ -33,7 +33,7 @@ func Login(r *http.Request) (bool, error) { username, password := r.FormValue("username"), r.FormValue("password") session, _ := Store.Get(r, "gophish") u, err := models.GetUserByUsername(username) - if err != models.ErrUsernameTaken { + if err != nil && err != models.ErrUsernameTaken { return false, err } //If we've made it here, we should have a valid user stored in u @@ -61,7 +61,7 @@ func Register(r *http.Request) (bool, error) { h, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost) u.Username = username u.Hash = string(h) - u.APIKey = GenerateSecureKey() + u.ApiKey = GenerateSecureKey() if err != nil { return false, err } diff --git a/controllers/api.go b/controllers/api.go index 17a880d0..e33ba787 100644 --- a/controllers/api.go +++ b/controllers/api.go @@ -39,7 +39,7 @@ func API_Reset(w http.ResponseWriter, r *http.Request) { switch { case r.Method == "POST": u := ctx.Get(r, "user").(models.User) - u.APIKey = auth.GenerateSecureKey() + u.ApiKey = auth.GenerateSecureKey() err := models.PutUser(&u) if err != nil { Flash(w, r, "danger", "Error resetting API Key") diff --git a/middleware/middleware.go b/middleware/middleware.go index 0d69d905..797bc3bf 100644 --- a/middleware/middleware.go +++ b/middleware/middleware.go @@ -52,12 +52,13 @@ func RequireAPIKey(handler http.Handler) http.HandlerFunc { if ak == "" { JSONError(w, 400, "API Key not set") } else { - id, err := models.Conn.SelectInt("SELECT id FROM users WHERE api_key=?", ak) - if id == 0 || err != nil { + u, err := models.GetUserByAPIKey(ak) + /* id, err := models.Conn.SelectInt("SELECT id FROM users WHERE api_key=?", ak) + */if err != nil { JSONError(w, 400, "Invalid API Key") return } - ctx.Set(r, "user_id", id) + ctx.Set(r, "user_id", u.Id) ctx.Set(r, "api_key", ak) handler.ServeHTTP(w, r) } diff --git a/models/campaign.go b/models/campaign.go index 9db54ee0..7c2e064b 100644 --- a/models/campaign.go +++ b/models/campaign.go @@ -9,39 +9,54 @@ import ( //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"` + UserId int64 `json:"-"` + 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 Status string `json:"status"` - Results []Result `json:"results,omitempty" db:"-"` - Groups []Group `json:"groups,omitempty" db:"-"` + Results []Result `json:"results,omitempty"` + Groups []Group `json:"groups,omitempty"` } type Result struct { - Target - Status string `json:"status"` + Id int64 `json:"-"` + CampaignId int64 `json:"-"` + Email string `json:"email"` + Status string `json:"status" sql:"not null"` } // 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) + err := db.Debug().Model(&User{Id: uid}).Related(&cs).Error + if err != nil { + fmt.Println(err) } + /* _, 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 := db.Debug().Model(&cs[i]).Related(&cs[i].Results).Error + if err != nil { + fmt.Println(err) + } + /* _, err = Conn.Select(&cs[i].Results, "SELECT r.email, r.status FROM campaign_results r WHERE r.cid=?", cs[i].Id) + */ + } + fmt.Printf("%v", cs) 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 { + err := db.Debug().Where("id = ?", id).Where("user_id = ?", uid).Find(&c).Error + /* 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 + err = db.Debug().Model(&c).Related(&c.Results).Error + /* _, 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. @@ -58,7 +73,8 @@ func PostCampaign(c *Campaign, uid int64) error { } } // Insert into the DB - err = Conn.Insert(c) + /*err = Conn.Insert(c)*/ + err = db.Debug().Save(&c).Error if err != nil { Logger.Println(err) return err @@ -67,35 +83,40 @@ func PostCampaign(c *Campaign, uid int64) error { 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"} + r := Result{Email: t.Email, Status: "Unknown", CampaignId: c.Id} 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) + err := db.Debug().Save(&r).Error + /*_, 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) + /*_, 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) + err := db.Debug().Delete(&Result{CampaignId: id}).Error + /*_, err := Conn.Exec("DELETE FROM campaign_results WHERE cid=?", id)*/ if err != nil { + Logger.Println(err) return err } // Delete the reference to the campaign in the user_campaigns table - _, err = Conn.Exec("DELETE FROM user_campaigns WHERE cid=?", id) + err = db.Debug().Delete(&Campaign{Id: id}).Error + /*_, err = Conn.Exec("DELETE FROM user_campaigns WHERE cid=?", id)*/ if err != nil { + Logger.Panicln(err) return err } // Delete the campaign itself - _, err = Conn.Exec("DELETE FROM campaigns WHERE id=?", id) + /*_, err = Conn.Exec("DELETE FROM campaigns WHERE id=?", id)*/ return err } diff --git a/models/group.go b/models/group.go index 1305dda4..bb2295bc 100644 --- a/models/group.go +++ b/models/group.go @@ -8,8 +8,18 @@ import ( 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:"-"` + ModifiedDate time.Time `json:"modified_date"` + Targets []Target `json:"targets" sql:"-"` +} + +type UserGroup struct { + UserId int64 `json:"-"` + GroupId int64 `json:"-"` +} + +type GroupTarget struct { + GroupId int64 `json:"-"` + TargetId int64 `json:"-"` } type Target struct { @@ -20,13 +30,14 @@ type Target struct { // 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) + err := db.Debug().Table("groups g").Select("g.id, g.name, g.modified_date").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=?", uid).Scan(&gs).Error + /*_, 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) + gs[i].Targets, err = GetTargets(gs[i].Id) if err != nil { Logger.Println(err) } @@ -37,12 +48,14 @@ func GetGroups(uid int64) ([]Group, error) { // 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) + err := db.Debug().Table("groups g").Select("g.id, g.name, g.modified_date").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.id=?", uid, id).Scan(&g).Error + /*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) + g.Targets, err = GetTargets(g.Id) + /*_, 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) } @@ -52,12 +65,14 @@ func GetGroup(id int64, uid int64) (Group, error) { // 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) + err := db.Debug().Table("groups g").Select("g.id, g.name, g.modified_date").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.name=?", uid, n).Scan(&g).Error + /*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) + g.Targets, err = GetTargets(g.Id) + /*_, 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) } @@ -67,16 +82,22 @@ func GetGroupByName(n string, uid int64) (Group, error) { // PostGroup creates a new group in the database. func PostGroup(g *Group, uid int64) error { // Insert into the DB - err = Conn.Insert(g) + /*err = Conn.Insert(g)*/ + err = db.Save(&g).Error 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) + err = db.Save(&UserGroup{GroupId: g.Id, UserId: uid}).Error + if err != nil { + Logger.Println(err) + return err + } + /*_, 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) } @@ -90,7 +111,8 @@ func PutGroup(g *Group, uid int64) error { // 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) + ts, err = GetTargets(g.Id) + /*_, 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 @@ -133,7 +155,8 @@ func PutGroup(g *Group, uid int64) error { } // Update the group g.ModifiedDate = time.Now() - _, err = Conn.Update(g) + err = db.Debug().Update(&g).Error + /*_, err = Conn.Update(g)*/ if err != nil { Logger.Println(err) return err @@ -178,16 +201,35 @@ func insertTargetIntoGroup(t Target, gid int64) error { // 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) + err := db.Debug().Where("group_id=?", id).Delete(&GroupTarget{}).Error + /*_, err := Conn.Exec("DELETE FROM group_targets WHERE gid=?", id)*/ if err != nil { + Logger.Println(err) return err } // Delete the reference to the group in the user_group table - _, err = Conn.Exec("DELETE FROM user_groups WHERE gid=?", id) + err = db.Debug().Where("group_id=?", id).Delete(&UserGroup{}).Error + /*_, err = Conn.Exec("DELETE FROM user_groups WHERE gid=?", id)*/ if err != nil { + Logger.Println(err) return err } // Delete the group itself - _, err = Conn.Exec("DELETE FROM groups WHERE id=?", id) + err = db.Debug().Delete(&Group{Id: id}).Error + /*_, err = Conn.Exec("DELETE FROM groups WHERE id=?", id)*/ + if err != nil { + Logger.Println(err) + return err + } return err } + +func GetTargets(gid int64) ([]Target, error) { + ts := []Target{} + err := db.Debug().Table("targets t").Select("t.id, t.email").Joins("left join group_targets gt ON t.id = gt.target_id").Where("gt.group_id=?", gid).Scan(&ts).Error + /*_, 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 { + return nil, err + } + return ts, nil +} diff --git a/models/models.go b/models/models.go index daf8343d..ce34aea9 100644 --- a/models/models.go +++ b/models/models.go @@ -3,24 +3,41 @@ package models import ( "database/sql" "errors" + "fmt" "log" "os" + "time" "github.com/coopernurse/gorp" + "github.com/jinzhu/gorm" "github.com/jordan-wright/gophish/config" _ "github.com/mattn/go-sqlite3" ) var Conn *gorp.DbMap +var db gorm.DB 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) +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) + if err != nil { + fmt.Printf("%v", err) + } + db, err = gorm.Open("sqlite3", "gophish_gorm.db") + db.SetLogger(Logger) + db.CreateTable(User{}) + db.CreateTable(Target{}) + db.CreateTable(Result{}) + db.CreateTable(Group{}) + db.CreateTable(GroupTarget{}) + db.CreateTable(UserGroup{}) + db.CreateTable(Template{}) + db.CreateTable(Campaign{}) Conn = &gorp.DbMap{Db: DB, Dialect: gorp.SqliteDialect{}} //If the file already exists, delete it and recreate it _, err = os.Stat(config.Conf.DBPath) @@ -57,9 +74,63 @@ func Setup() error { init_user := User{ Username: "admin", Hash: "$2a$10$IYkPp0.QsM81lYYPrQx6W.U6oQGw7wMpozrKhKAHUBVL4mkm/EvAS", //gophish - APIKey: "12345678901234567890123456789012", + ApiKey: "12345678901234567890123456789012", } - Conn.Insert(&init_user) + err = db.Save(&init_user).Error + if err != nil { + Logger.Println(err) + } + init_campaign := Campaign{ + UserId: 1, + Name: "First Campaign", + CreatedDate: time.Now(), //gophish + CompletedDate: time.Now(), + Template: "", + Status: "In Progress", + Results: []Result{}, + Groups: []Group{}, + } + err = db.Debug().Save(&init_campaign).Error + if err != nil { + Logger.Println(err) + } + init_result := Result{ + CampaignId: 1, + Email: "test@example.com", + Status: "Unsuccessful", + } + err = db.Debug().Save(&init_result).Error + if err != nil { + Logger.Println(err) + } + init_group := Group{ + Name: "New Group", + ModifiedDate: time.Now(), + } + err = db.Debug().Save(&init_group).Error + if err != nil { + Logger.Println(err) + } + init_ug := UserGroup{ + UserId: 1, + GroupId: 1, + } + err = db.Debug().Save(&init_ug).Error + if err != nil { + Logger.Println(err) + } + init_target := Target{ + Email: "test@example.com", + } + err = db.Debug().Save(&init_target).Error + if err != nil { + Logger.Println(err) + } + init_gt := GroupTarget{ + GroupId: 1, + TargetId: 1, + } + err = db.Debug().Save(&init_gt).Error if err != nil { Logger.Println(err) } diff --git a/models/template.go b/models/template.go index 39f01dfd..eb97476f 100644 --- a/models/template.go +++ b/models/template.go @@ -4,10 +4,10 @@ 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"` + Name string `json:"name"` + Text string `json:"text"` + Html string `json:"html"` + ModifiedDate time.Time `json:"modified_date"` } // GetTemplates returns the templates owned by the given user. diff --git a/models/user.go b/models/user.go index 2f329905..c3e52ed2 100644 --- a/models/user.go +++ b/models/user.go @@ -5,16 +5,17 @@ import "database/sql" // User represents the user model for gophish. type User struct { Id int64 `json:"id"` - Username string `json:"username"` + Username string `json:"username"` /* sql:"not null;unique"`*/ Hash string `json:"-"` - APIKey string `json:"api_key" db:"api_key"` + ApiKey string `json:"api_key" db:"api_key"` /* sql:"not null;unique"`*/ } // 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) + /* err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE id=?", id)*/ + err := db.Where("id=?", id).First(&u).Error if err != nil { return u, err } @@ -23,9 +24,10 @@ func GetUser(id int64) (User, error) { // 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) { +func GetUserByAPIKey(key string) (User, error) { u := User{} - err := Conn.SelectOne(&u, "SELECT id, username, api_key FROM Users WHERE apikey=?", key) + /* err := Conn.SelectOne(&u, "SELECT id, username, api_key FROM Users WHERE apikey=?", key)*/ + err := db.Debug().Where("api_key = ?", key).First(&u).Error if err != nil { return u, err } @@ -36,7 +38,8 @@ func GetUserByAPIKey(key []byte) (User, error) { // error is thrown. func GetUserByUsername(username string) (User, error) { u := User{} - err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE username=?", username) + /* err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE username=?", username)*/ + err := db.Where("username = ?", username).First(&u).Error if err != sql.ErrNoRows { return u, ErrUsernameTaken } else if err != nil { @@ -47,6 +50,7 @@ func GetUserByUsername(username string) (User, error) { // PutUser updates the given user func PutUser(u *User) error { - _, err := Conn.Update(u) + err := db.Debug().Update(&u).Error + /*_, err := Conn.Update(u)*/ return err } diff --git a/templates/base.html b/templates/base.html index fa1404e6..9f7fc38a 100644 --- a/templates/base.html +++ b/templates/base.html @@ -20,7 +20,7 @@ {{%if .User%}} - + {{%end%}} From 6e5361b210f513f0ceed17cf23f240e438d733da Mon Sep 17 00:00:00 2001 From: Jordan Date: Wed, 26 Mar 2014 14:32:48 -0500 Subject: [PATCH 2/8] Continuing gorm integration Todo: Templates --- models/group.go | 46 ++++++++++++++++++++++++++-------------------- models/models.go | 27 ++++++--------------------- 2 files changed, 32 insertions(+), 41 deletions(-) diff --git a/models/group.go b/models/group.go index bb2295bc..17e22dd3 100644 --- a/models/group.go +++ b/models/group.go @@ -3,6 +3,8 @@ package models import ( "net/mail" "time" + + "github.com/jinzhu/gorm" ) type Group struct { @@ -30,7 +32,7 @@ type Target struct { // GetGroups returns the groups owned by the given user. func GetGroups(uid int64) ([]Group, error) { gs := []Group{} - err := db.Debug().Table("groups g").Select("g.id, g.name, g.modified_date").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=?", uid).Scan(&gs).Error + err := db.Debug().Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=?", uid).Scan(&gs).Error /*_, 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) @@ -48,7 +50,7 @@ func GetGroups(uid int64) ([]Group, error) { // 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 := db.Debug().Table("groups g").Select("g.id, g.name, g.modified_date").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.id=?", uid, id).Scan(&g).Error + err := db.Debug().Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.id=?", uid, id).Scan(&g).Error /*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) @@ -65,7 +67,7 @@ func GetGroup(id int64, uid int64) (Group, error) { // 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 := db.Debug().Table("groups g").Select("g.id, g.name, g.modified_date").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.name=?", uid, n).Scan(&g).Error + err := db.Debug().Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.name=?", uid, n).Scan(&g).Error /*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) @@ -83,7 +85,7 @@ func GetGroupByName(n string, uid int64) (Group, error) { func PostGroup(g *Group, uid int64) error { // Insert into the DB /*err = Conn.Insert(g)*/ - err = db.Save(&g).Error + err = db.Save(g).Error if err != nil { Logger.Println(err) return err @@ -131,7 +133,8 @@ func PutGroup(g *Group, uid int64) error { } // 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) + err = db.Debug().Where("group_id=? and target_id=?", g.Id, t.Id).Delete(&GroupTarget{}).Error + /*_, 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) } @@ -155,7 +158,7 @@ func PutGroup(g *Group, uid int64) error { } // Update the group g.ModifiedDate = time.Now() - err = db.Debug().Update(&g).Error + err = db.Debug().Save(g).Error /*_, err = Conn.Update(g)*/ if err != nil { Logger.Println(err) @@ -169,28 +172,34 @@ func insertTargetIntoGroup(t Target, gid int64) error { Logger.Printf("Invalid email %s\n", t.Email) return err } - trans, err := Conn.Begin() + trans := db.Begin() + trans.Debug().Where(t).FirstOrCreate(&t) + Logger.Printf("ID of Target after FirstOrCreate: %d", t.Id) + /*_, err = trans.Exec("INSERT OR IGNORE INTO targets VALUES (null, ?)", t.Email)*/ if err != nil { - Logger.Println(err) + Logger.Printf("Error adding target: %s\n", t.Email) 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.) + /*// 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.Debug().Where("group_id=? and target_id=?", gid, t.Id).Find(&GroupTarget{}).Error + if err == gorm.RecordNotFound { + err = trans.Debug().Save(&GroupTarget{GroupId: gid, TargetId: t.Id}).Error + if err != nil { + Logger.Println(err) + return err + } } - _, err = trans.Exec("INSERT OR IGNORE INTO group_targets VALUES (?,?)", gid, t.Id) + /*_, 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() + err = trans.Commit().Error if err != nil { Logger.Printf("Error committing db changes\n") return err @@ -228,8 +237,5 @@ func GetTargets(gid int64) ([]Target, error) { ts := []Target{} err := db.Debug().Table("targets t").Select("t.id, t.email").Joins("left join group_targets gt ON t.id = gt.target_id").Where("gt.group_id=?", gid).Scan(&ts).Error /*_, 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 { - return nil, err - } - return ts, nil + return ts, err } diff --git a/models/models.go b/models/models.go index ce34aea9..75ef0056 100644 --- a/models/models.go +++ b/models/models.go @@ -106,31 +106,16 @@ func Setup() error { init_group := Group{ Name: "New Group", ModifiedDate: time.Now(), + Targets: []Target{Target{ + Email: "test@example.com"}, + }, } - err = db.Debug().Save(&init_group).Error + err = PostGroup(&init_group, 1) if err != nil { Logger.Println(err) } - init_ug := UserGroup{ - UserId: 1, - GroupId: 1, - } - err = db.Debug().Save(&init_ug).Error - if err != nil { - Logger.Println(err) - } - init_target := Target{ - Email: "test@example.com", - } - err = db.Debug().Save(&init_target).Error - if err != nil { - Logger.Println(err) - } - init_gt := GroupTarget{ - GroupId: 1, - TargetId: 1, - } - err = db.Debug().Save(&init_gt).Error + init_group.Targets = []Target{Target{Email: "test2@example.com"}} + err = PutGroup(&init_group, 1) if err != nil { Logger.Println(err) } From 5740ee32738076cc23f8bf82199cd6e52b7ead4b Mon Sep 17 00:00:00 2001 From: Jordan Date: Wed, 26 Mar 2014 14:50:16 -0500 Subject: [PATCH 3/8] Fixed bug in POSTing campaign Removed gorp initialization code from models.go --- controllers/api.go | 1 + models/campaign.go | 2 +- models/models.go | 109 +++++++++------------------------------------ 3 files changed, 22 insertions(+), 90 deletions(-) diff --git a/controllers/api.go b/controllers/api.go index e33ba787..a2c5e701 100644 --- a/controllers/api.go +++ b/controllers/api.go @@ -80,6 +80,7 @@ func API_Campaigns(w http.ResponseWriter, r *http.Request) { c.CreatedDate = time.Now() c.CompletedDate = time.Time{} c.Status = IN_PROGRESS + 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 diff --git a/models/campaign.go b/models/campaign.go index 7c2e064b..bbef7d8d 100644 --- a/models/campaign.go +++ b/models/campaign.go @@ -74,7 +74,7 @@ func PostCampaign(c *Campaign, uid int64) error { } // Insert into the DB /*err = Conn.Insert(c)*/ - err = db.Debug().Save(&c).Error + err = db.Debug().Save(c).Error if err != nil { Logger.Println(err) return err diff --git a/models/models.go b/models/models.go index 75ef0056..24f98668 100644 --- a/models/models.go +++ b/models/models.go @@ -3,10 +3,8 @@ package models import ( "database/sql" "errors" - "fmt" "log" "os" - "time" "github.com/coopernurse/gorp" "github.com/jinzhu/gorm" @@ -21,55 +19,33 @@ var err error var ErrUsernameTaken = errors.New("username already taken") var Logger = log.New(os.Stdout, " ", log.Ldate|log.Ltime|log.Lshortfile) +// Flash is used to hold flash information for use in templates. +type Flash struct { + Type string + Message string +} + // Setup initializes the Conn object // It also populates the Gophish Config object func Setup() error { - DB, err := sql.Open("sqlite3", config.Conf.DBPath) - if err != nil { - fmt.Printf("%v", err) - } - db, err = gorm.Open("sqlite3", "gophish_gorm.db") + db, err = gorm.Open("sqlite3", config.Conf.DBPath) db.SetLogger(Logger) - db.CreateTable(User{}) - db.CreateTable(Target{}) - db.CreateTable(Result{}) - db.CreateTable(Group{}) - db.CreateTable(GroupTarget{}) - db.CreateTable(UserGroup{}) - db.CreateTable(Template{}) - db.CreateTable(Campaign{}) - Conn = &gorp.DbMap{Db: DB, Dialect: gorp.SqliteDialect{}} + if err != nil { + Logger.Println(err) + return err + } //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 err - } - } + Logger.Printf("Database not found... creating db at %s\n", config.Conf.DBPath) + db.CreateTable(User{}) + db.CreateTable(Target{}) + db.CreateTable(Result{}) + db.CreateTable(Group{}) + db.CreateTable(GroupTarget{}) + db.CreateTable(UserGroup{}) + db.CreateTable(Template{}) + db.CreateTable(Campaign{}) //Create the default user init_user := User{ Username: "admin", @@ -80,51 +56,6 @@ func Setup() error { if err != nil { Logger.Println(err) } - init_campaign := Campaign{ - UserId: 1, - Name: "First Campaign", - CreatedDate: time.Now(), //gophish - CompletedDate: time.Now(), - Template: "", - Status: "In Progress", - Results: []Result{}, - Groups: []Group{}, - } - err = db.Debug().Save(&init_campaign).Error - if err != nil { - Logger.Println(err) - } - init_result := Result{ - CampaignId: 1, - Email: "test@example.com", - Status: "Unsuccessful", - } - err = db.Debug().Save(&init_result).Error - if err != nil { - Logger.Println(err) - } - init_group := Group{ - Name: "New Group", - ModifiedDate: time.Now(), - Targets: []Target{Target{ - Email: "test@example.com"}, - }, - } - err = PostGroup(&init_group, 1) - if err != nil { - Logger.Println(err) - } - init_group.Targets = []Target{Target{Email: "test2@example.com"}} - err = PutGroup(&init_group, 1) - if err != nil { - Logger.Println(err) - } } return nil } - -// Flash is used to hold flash information for use in templates. -type Flash struct { - Type string - Message string -} From 776f591f68b93c70b525252581ee3e4d7a36c005 Mon Sep 17 00:00:00 2001 From: Jordan Date: Wed, 26 Mar 2014 15:01:49 -0500 Subject: [PATCH 4/8] Cleaned up old gorp code Bug fix in settings template --- models/campaign.go | 36 ++++++++++--------------------- models/group.go | 48 +++++++++++------------------------------ models/user.go | 12 ++++------- templates/settings.html | 2 +- 4 files changed, 28 insertions(+), 70 deletions(-) diff --git a/models/campaign.go b/models/campaign.go index bbef7d8d..e3846b90 100644 --- a/models/campaign.go +++ b/models/campaign.go @@ -29,18 +29,15 @@ type Result struct { // GetCampaigns returns the campaigns owned by the given user. func GetCampaigns(uid int64) ([]Campaign, error) { cs := []Campaign{} - err := db.Debug().Model(&User{Id: uid}).Related(&cs).Error + err := db.Model(&User{Id: uid}).Related(&cs).Error if err != nil { fmt.Println(err) } - /* _, 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 := db.Debug().Model(&cs[i]).Related(&cs[i].Results).Error + err := db.Model(&cs[i]).Related(&cs[i].Results).Error if err != nil { fmt.Println(err) } - /* _, err = Conn.Select(&cs[i].Results, "SELECT r.email, r.status FROM campaign_results r WHERE r.cid=?", cs[i].Id) - */ } fmt.Printf("%v", cs) return cs, err @@ -49,13 +46,11 @@ func GetCampaigns(uid int64) ([]Campaign, error) { // 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 := db.Debug().Where("id = ?", id).Where("user_id = ?", uid).Find(&c).Error - /* 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) + err := db.Where("id = ?", id).Where("user_id = ?", uid).Find(&c).Error */if err != nil { return c, err } - err = db.Debug().Model(&c).Related(&c.Results).Error - /* _, err = Conn.Select(&c.Results, "SELECT r.email, r.status FROM campaign_results r WHERE r.cid=?", c.Id) + err = db.Model(&c).Related(&c.Results).Error */return c, err } @@ -73,8 +68,7 @@ func PostCampaign(c *Campaign, uid int64) error { } } // Insert into the DB - /*err = Conn.Insert(c)*/ - err = db.Debug().Save(c).Error + err = db.Save(c).Error if err != nil { Logger.Println(err) return err @@ -86,37 +80,29 @@ func PostCampaign(c *Campaign, uid int64) error { r := Result{Email: t.Email, Status: "Unknown", CampaignId: c.Id} c.Results = append(c.Results, r) fmt.Printf("%v", c.Results) - err := db.Debug().Save(&r).Error - /*_, err = Conn.Exec("INSERT INTO campaign_results VALUES (?,?,?)", c.Id, r.Email, r.Status)*/ + err := db.Save(&r).Error 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 } +//DeleteCampaign deletes the specified campaign func DeleteCampaign(id int64) error { - // Delete all the campaign_results entries for this group - err := db.Debug().Delete(&Result{CampaignId: id}).Error - /*_, err := Conn.Exec("DELETE FROM campaign_results WHERE cid=?", id)*/ + // Delete all the campaign results + err := db.Delete(&Result{CampaignId: id}).Error if err != nil { Logger.Println(err) return err } - // Delete the reference to the campaign in the user_campaigns table - err = db.Debug().Delete(&Campaign{Id: id}).Error - /*_, err = Conn.Exec("DELETE FROM user_campaigns WHERE cid=?", id)*/ + // Delete the campaign + err = db.Delete(&Campaign{Id: id}).Error if err != nil { Logger.Panicln(err) 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 index 17e22dd3..8341dc92 100644 --- a/models/group.go +++ b/models/group.go @@ -32,8 +32,7 @@ type Target struct { // GetGroups returns the groups owned by the given user. func GetGroups(uid int64) ([]Group, error) { gs := []Group{} - err := db.Debug().Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=?", uid).Scan(&gs).Error - /*_, 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)*/ + err := db.Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=?", uid).Scan(&gs).Error if err != nil { Logger.Println(err) return gs, err @@ -50,14 +49,12 @@ func GetGroups(uid int64) ([]Group, error) { // 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 := db.Debug().Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.id=?", uid, id).Scan(&g).Error - /*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)*/ + err := db.Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.id=?", uid, id).Scan(&g).Error if err != nil { Logger.Println(err) return g, err } g.Targets, err = GetTargets(g.Id) - /*_, 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) } @@ -67,14 +64,12 @@ func GetGroup(id int64, uid int64) (Group, error) { // 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 := db.Debug().Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.name=?", uid, n).Scan(&g).Error - /*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)*/ + err := db.Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.name=?", uid, n).Scan(&g).Error if err != nil { Logger.Println(err) return g, err } g.Targets, err = GetTargets(g.Id) - /*_, 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) } @@ -84,7 +79,6 @@ func GetGroupByName(n string, uid int64) (Group, error) { // PostGroup creates a new group in the database. func PostGroup(g *Group, uid int64) error { // Insert into the DB - /*err = Conn.Insert(g)*/ err = db.Save(g).Error if err != nil { Logger.Println(err) @@ -96,10 +90,6 @@ func PostGroup(g *Group, uid int64) error { Logger.Println(err) return err } - /*_, 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) } @@ -114,7 +104,6 @@ func PutGroup(g *Group, uid int64) error { // Get all the targets currently in the database for the group ts := []Target{} ts, err = GetTargets(g.Id) - /*_, 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 @@ -133,8 +122,7 @@ func PutGroup(g *Group, uid int64) error { } // If the target does not exist in the group any longer, we delete it if !tExists { - err = db.Debug().Where("group_id=? and target_id=?", g.Id, t.Id).Delete(&GroupTarget{}).Error - /*_, err = Conn.Exec("DELETE FROM group_targets WHERE gid=? AND tid=?", g.Id, t.Id)*/ + err = db.Where("group_id=? and target_id=?", g.Id, t.Id).Delete(&GroupTarget{}).Error if err != nil { Logger.Printf("Error deleting email %s\n", t.Email) } @@ -158,7 +146,7 @@ func PutGroup(g *Group, uid int64) error { } // Update the group g.ModifiedDate = time.Now() - err = db.Debug().Save(g).Error + err = db.Save(g).Error /*_, err = Conn.Update(g)*/ if err != nil { Logger.Println(err) @@ -173,28 +161,20 @@ func insertTargetIntoGroup(t Target, gid int64) error { return err } trans := db.Begin() - trans.Debug().Where(t).FirstOrCreate(&t) + trans.Where(t).FirstOrCreate(&t) Logger.Printf("ID of Target after FirstOrCreate: %d", t.Id) - /*_, err = trans.Exec("INSERT OR IGNORE INTO targets VALUES (null, ?)", t.Email)*/ if err != nil { Logger.Printf("Error adding target: %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.Debug().Where("group_id=? and target_id=?", gid, t.Id).Find(&GroupTarget{}).Error + err = trans.Where("group_id=? and target_id=?", gid, t.Id).Find(&GroupTarget{}).Error if err == gorm.RecordNotFound { - err = trans.Debug().Save(&GroupTarget{GroupId: gid, TargetId: t.Id}).Error + err = trans.Save(&GroupTarget{GroupId: gid, TargetId: t.Id}).Error if err != nil { Logger.Println(err) 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 @@ -210,22 +190,19 @@ func insertTargetIntoGroup(t Target, gid int64) error { // 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 := db.Debug().Where("group_id=?", id).Delete(&GroupTarget{}).Error - /*_, err := Conn.Exec("DELETE FROM group_targets WHERE gid=?", id)*/ + err := db.Where("group_id=?", id).Delete(&GroupTarget{}).Error if err != nil { Logger.Println(err) return err } // Delete the reference to the group in the user_group table - err = db.Debug().Where("group_id=?", id).Delete(&UserGroup{}).Error - /*_, err = Conn.Exec("DELETE FROM user_groups WHERE gid=?", id)*/ + err = db.Where("group_id=?", id).Delete(&UserGroup{}).Error if err != nil { Logger.Println(err) return err } // Delete the group itself - err = db.Debug().Delete(&Group{Id: id}).Error - /*_, err = Conn.Exec("DELETE FROM groups WHERE id=?", id)*/ + err = db.Delete(&Group{Id: id}).Error if err != nil { Logger.Println(err) return err @@ -235,7 +212,6 @@ func DeleteGroup(id int64) error { func GetTargets(gid int64) ([]Target, error) { ts := []Target{} - err := db.Debug().Table("targets t").Select("t.id, t.email").Joins("left join group_targets gt ON t.id = gt.target_id").Where("gt.group_id=?", gid).Scan(&ts).Error - /*_, 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)*/ + err := db.Table("targets t").Select("t.id, t.email").Joins("left join group_targets gt ON t.id = gt.target_id").Where("gt.group_id=?", gid).Scan(&ts).Error return ts, err } diff --git a/models/user.go b/models/user.go index c3e52ed2..9f16a1a6 100644 --- a/models/user.go +++ b/models/user.go @@ -5,16 +5,15 @@ import "database/sql" // User represents the user model for gophish. type User struct { Id int64 `json:"id"` - Username string `json:"username"` /* sql:"not null;unique"`*/ + Username string `json:"username" sql:"not null;unique"` Hash string `json:"-"` - ApiKey string `json:"api_key" db:"api_key"` /* sql:"not null;unique"`*/ + ApiKey string `json:"api_key" sql:"not null;unique"` } // 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)*/ err := db.Where("id=?", id).First(&u).Error if err != nil { return u, err @@ -26,8 +25,7 @@ func GetUser(id int64) (User, error) { // error is thrown. func GetUserByAPIKey(key string) (User, error) { u := User{} - /* err := Conn.SelectOne(&u, "SELECT id, username, api_key FROM Users WHERE apikey=?", key)*/ - err := db.Debug().Where("api_key = ?", key).First(&u).Error + err := db.Where("api_key = ?", key).First(&u).Error if err != nil { return u, err } @@ -38,7 +36,6 @@ func GetUserByAPIKey(key string) (User, error) { // error is thrown. func GetUserByUsername(username string) (User, error) { u := User{} - /* err := Conn.SelectOne(&u, "SELECT * FROM Users WHERE username=?", username)*/ err := db.Where("username = ?", username).First(&u).Error if err != sql.ErrNoRows { return u, ErrUsernameTaken @@ -50,7 +47,6 @@ func GetUserByUsername(username string) (User, error) { // PutUser updates the given user func PutUser(u *User) error { - err := db.Debug().Update(&u).Error - /*_, err := Conn.Update(u)*/ + err := db.Update(&u).Error return err } diff --git a/templates/settings.html b/templates/settings.html index 8cc15477..1c02c380 100644 --- a/templates/settings.html +++ b/templates/settings.html @@ -27,7 +27,7 @@
- +
From cb9c405f46118aef4cee1fa352bcc8551bc4cf1a Mon Sep 17 00:00:00 2001 From: Jordan Date: Wed, 26 Mar 2014 21:42:07 -0500 Subject: [PATCH 5/8] Added better testing with gocheck Fixed some typos leading to syntax errors TODO: Finish up gorm integration into templates --- config/config.go | 2 -- gophish.go | 2 +- gophish_test.go | 14 -------------- models/campaign.go | 6 +++--- models/models.go | 3 +-- models/models_test.go | 38 ++++++++++++++++++++++++++++++++++++++ 6 files changed, 43 insertions(+), 22 deletions(-) delete mode 100644 gophish_test.go create mode 100644 models/models_test.go diff --git a/config/config.go b/config/config.go index 3884c9f7..a92e3a41 100644 --- a/config/config.go +++ b/config/config.go @@ -4,7 +4,6 @@ import ( "encoding/json" "fmt" "io/ioutil" - "os" ) type SMTPServer struct { @@ -27,7 +26,6 @@ func init() { config_file, err := ioutil.ReadFile("./config.json") if err != nil { fmt.Printf("File error: %v\n", err) - os.Exit(1) } json.Unmarshal(config_file, &Conf) } diff --git a/gophish.go b/gophish.go index 9433e5ef..acc45e62 100644 --- a/gophish.go +++ b/gophish.go @@ -38,7 +38,7 @@ import ( func main() { //Setup the global variables and settings err := models.Setup() - defer models.DB.Close() + //defer models.db.DB().Close() if err != nil { fmt.Println(err) } diff --git a/gophish_test.go b/gophish_test.go deleted file mode 100644 index 97510b9a..00000000 --- a/gophish_test.go +++ /dev/null @@ -1,14 +0,0 @@ -package main - -import ( - "testing" - - "github.com/jordan-wright/gophish/models" -) - -func TestDBSetup(t *testing.T) { - err := models.Setup() - if err != nil { - t.Fatalf("Failed creating database: %v", err) - } -} diff --git a/models/campaign.go b/models/campaign.go index e3846b90..46d91887 100644 --- a/models/campaign.go +++ b/models/campaign.go @@ -33,7 +33,7 @@ func GetCampaigns(uid int64) ([]Campaign, error) { if err != nil { fmt.Println(err) } - */for i, _ := range cs { + for i, _ := range cs { err := db.Model(&cs[i]).Related(&cs[i].Results).Error if err != nil { fmt.Println(err) @@ -47,11 +47,11 @@ func GetCampaigns(uid int64) ([]Campaign, error) { func GetCampaign(id int64, uid int64) (Campaign, error) { c := Campaign{} err := db.Where("id = ?", id).Where("user_id = ?", uid).Find(&c).Error - */if err != nil { + if err != nil { return c, err } err = db.Model(&c).Related(&c.Results).Error - */return c, err + return c, err } // PostCampaign inserts a campaign and all associated records into the database. diff --git a/models/models.go b/models/models.go index 24f98668..ecf572a3 100644 --- a/models/models.go +++ b/models/models.go @@ -1,7 +1,6 @@ package models import ( - "database/sql" "errors" "log" "os" @@ -14,7 +13,6 @@ import ( var Conn *gorp.DbMap var db gorm.DB -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) @@ -29,6 +27,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.SetLogger(Logger) if err != nil { Logger.Println(err) diff --git a/models/models_test.go b/models/models_test.go new file mode 100644 index 00000000..6c2d6733 --- /dev/null +++ b/models/models_test.go @@ -0,0 +1,38 @@ +package models + +import ( + "os" + "testing" + + "github.com/jordan-wright/gophish/config" + "launchpad.net/gocheck" +) + +// Hook up gocheck into the "go test" runner. +func Test(t *testing.T) { gocheck.TestingT(t) } + +type ModelsSuite struct{} + +var _ = gocheck.Suite(&ModelsSuite{}) + +func (s *ModelsSuite) SetUpSuite(c *gocheck.C) { + config.Conf.DBPath = "../gophish_test.db" + err := Setup() + if err != nil { + c.Fatalf("Failed creating database: %v", err) + } +} + +func (s *ModelsSuite) TestGetUser(c *gocheck.C) { + u, err := GetUser(1) + c.Assert(err, gocheck.Equals, nil) + c.Assert(u.Username, gocheck.Equals, "admin") +} + +func (s *ModelsSuite) TearDownSuite(c *gocheck.C) { + db.DB().Close() + err := os.Remove(config.Conf.DBPath) + if err != nil { + c.Fatalf("Failed deleting test database: %v", err) + } +} From 1b8b24756216c73675c5555c8c7599c2d5252100 Mon Sep 17 00:00:00 2001 From: Jordan Date: Wed, 26 Mar 2014 21:49:39 -0500 Subject: [PATCH 6/8] Added gocheck to be installed in Travis-CI --- .travis.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.travis.yml b/.travis.yml index 9a18d00e..2f9151ea 100644 --- a/.travis.yml +++ b/.travis.yml @@ -2,4 +2,7 @@ language: go go: - 1.1 - - tip \ No newline at end of file + - tip + +install: + - go get launchpad.net/gocheck \ No newline at end of file From a0d8dab6332e200646fa81d39fcd8af2550edd3c Mon Sep 17 00:00:00 2001 From: Jordan Date: Wed, 26 Mar 2014 21:53:23 -0500 Subject: [PATCH 7/8] Forgot to specify original build strings as well --- .travis.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.travis.yml b/.travis.yml index 2f9151ea..cb2cc2e7 100644 --- a/.travis.yml +++ b/.travis.yml @@ -5,4 +5,5 @@ go: - tip install: + - go get -d -v ./... && go build -v ./... - go get launchpad.net/gocheck \ No newline at end of file From f0e7ac46c8b06b902d31296ac8ba7958cdf33173 Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 27 Mar 2014 13:19:57 -0500 Subject: [PATCH 8/8] Integrated gorm with templates - Gorm integration complete Updated groups to belong to one user. I may make a Team and TeamGroups, and TeamUsers m2m relationships later. Added another test - more on the way soon. --- controllers/api.go | 11 ++++--- models/group.go | 69 +++++++++++++++---------------------------- models/models.go | 2 +- models/models_test.go | 9 ++++++ models/template.go | 17 ++++++----- models/user.go | 2 +- 6 files changed, 50 insertions(+), 60 deletions(-) diff --git a/controllers/api.go b/controllers/api.go index a2c5e701..b535ff17 100644 --- a/controllers/api.go +++ b/controllers/api.go @@ -176,7 +176,8 @@ func API_Groups(w http.ResponseWriter, r *http.Request) { return } g.ModifiedDate = time.Now() - err = models.PostGroup(&g, ctx.Get(r, "user_id").(int64)) + g.UserId = ctx.Get(r, "user_id").(int64) + err = models.PostGroup(&g) if checkError(err, w, "Error inserting group", http.StatusInternalServerError) { return } @@ -205,11 +206,11 @@ func API_Groups_Id(w http.ResponseWriter, r *http.Request) { } writeJSON(w, gj) case r.Method == "DELETE": - _, err := models.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 } - err = models.DeleteGroup(id) + err = models.DeleteGroup(&g) if checkError(err, w, "Error deleting group", http.StatusInternalServerError) { return } @@ -230,7 +231,9 @@ func API_Groups_Id(w http.ResponseWriter, r *http.Request) { http.Error(w, "Error: No targets specified", http.StatusBadRequest) return } - err = models.PutGroup(&g, ctx.Get(r, "user_id").(int64)) + g.ModifiedDate = time.Now() + g.UserId = ctx.Get(r, "user_id").(int64) + err = models.PutGroup(&g) if checkError(err, w, "Error updating group", http.StatusInternalServerError) { return } diff --git a/models/group.go b/models/group.go index 8341dc92..dc80c501 100644 --- a/models/group.go +++ b/models/group.go @@ -9,16 +9,12 @@ import ( type Group struct { Id int64 `json:"id"` + UserId int64 `json:"-"` Name string `json:"name"` ModifiedDate time.Time `json:"modified_date"` Targets []Target `json:"targets" sql:"-"` } -type UserGroup struct { - UserId int64 `json:"-"` - GroupId int64 `json:"-"` -} - type GroupTarget struct { GroupId int64 `json:"-"` TargetId int64 `json:"-"` @@ -32,7 +28,7 @@ type Target struct { // GetGroups returns the groups owned by the given user. func GetGroups(uid int64) ([]Group, error) { gs := []Group{} - err := db.Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=?", uid).Scan(&gs).Error + err := db.Where("user_id=?", uid).Find(&gs).Error if err != nil { Logger.Println(err) return gs, err @@ -49,7 +45,7 @@ func GetGroups(uid int64) ([]Group, error) { // 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 := db.Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.id=?", uid, id).Scan(&g).Error + err := db.Where("user_id=? and id=?", uid, id).Find(&g).Error if err != nil { Logger.Println(err) return g, err @@ -64,7 +60,7 @@ func GetGroup(id int64, uid int64) (Group, error) { // 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 := db.Table("groups g").Select("g.*").Joins("left join user_groups ug ON g.id = ug.group_id").Where("ug.user_id=? and g.name=?", uid, n).Scan(&g).Error + err := db.Where("user_id=? and name=?", uid, n).Find(&g).Error if err != nil { Logger.Println(err) return g, err @@ -77,19 +73,13 @@ func GetGroupByName(n string, uid int64) (Group, error) { } // PostGroup creates a new group in the database. -func PostGroup(g *Group, uid int64) error { +func PostGroup(g *Group) error { // Insert into the DB err = db.Save(g).Error if err != nil { Logger.Println(err) return err } - // Now, let's add the user->user_groups->group mapping - err = db.Save(&UserGroup{GroupId: g.Id, UserId: uid}).Error - if err != nil { - Logger.Println(err) - return err - } for _, t := range g.Targets { insertTargetIntoGroup(t, g.Id) } @@ -97,11 +87,7 @@ func PostGroup(g *Group, uid int64) error { } // 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 +func PutGroup(g *Group) error { ts := []Target{} ts, err = GetTargets(g.Id) if err != nil { @@ -144,8 +130,6 @@ func PutGroup(g *Group, uid int64) error { insertTargetIntoGroup(nt, g.Id) } } - // Update the group - g.ModifiedDate = time.Now() err = db.Save(g).Error /*_, err = Conn.Update(g)*/ if err != nil { @@ -155,6 +139,23 @@ func PutGroup(g *Group, uid int64) error { return nil } +// DeleteGroup deletes a given group by group ID and user ID +func DeleteGroup(g *Group) error { + // Delete all the group_targets entries for this group + err := db.Where("group_id=?", g.Id).Delete(&GroupTarget{}).Error + if err != nil { + Logger.Println(err) + return err + } + // Delete the group itself + err = db.Delete(g).Error + if err != nil { + Logger.Println(err) + return err + } + return err +} + func insertTargetIntoGroup(t Target, gid int64) error { if _, err = mail.ParseAddress(t.Email); err != nil { Logger.Printf("Invalid email %s\n", t.Email) @@ -162,7 +163,6 @@ func insertTargetIntoGroup(t Target, gid int64) error { } trans := db.Begin() trans.Where(t).FirstOrCreate(&t) - Logger.Printf("ID of Target after FirstOrCreate: %d", t.Id) if err != nil { Logger.Printf("Error adding target: %s\n", t.Email) return err @@ -187,29 +187,6 @@ func insertTargetIntoGroup(t Target, gid int64) error { 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 := db.Where("group_id=?", id).Delete(&GroupTarget{}).Error - if err != nil { - Logger.Println(err) - return err - } - // Delete the reference to the group in the user_group table - err = db.Where("group_id=?", id).Delete(&UserGroup{}).Error - if err != nil { - Logger.Println(err) - return err - } - // Delete the group itself - err = db.Delete(&Group{Id: id}).Error - if err != nil { - Logger.Println(err) - return err - } - return err -} - func GetTargets(gid int64) ([]Target, error) { ts := []Target{} err := db.Table("targets t").Select("t.id, t.email").Joins("left join group_targets gt ON t.id = gt.target_id").Where("gt.group_id=?", gid).Scan(&ts).Error diff --git a/models/models.go b/models/models.go index ecf572a3..87cbc372 100644 --- a/models/models.go +++ b/models/models.go @@ -42,8 +42,8 @@ func Setup() error { db.CreateTable(Result{}) db.CreateTable(Group{}) db.CreateTable(GroupTarget{}) - db.CreateTable(UserGroup{}) db.CreateTable(Template{}) + db.CreateTable(UserTemplate{}) db.CreateTable(Campaign{}) //Create the default user init_user := User{ diff --git a/models/models_test.go b/models/models_test.go index 6c2d6733..1bd8bd6d 100644 --- a/models/models_test.go +++ b/models/models_test.go @@ -29,6 +29,15 @@ func (s *ModelsSuite) TestGetUser(c *gocheck.C) { c.Assert(u.Username, gocheck.Equals, "admin") } +func (s *ModelsSuite) TestPutUser(c *gocheck.C) { + u, err := GetUser(1) + u.Username = "admin_changed" + err = PutUser(&u) + c.Assert(err, gocheck.Equals, nil) + u, err = GetUser(1) + c.Assert(u.Username, gocheck.Equals, "admin_changed") +} + func (s *ModelsSuite) TearDownSuite(c *gocheck.C) { db.DB().Close() err := os.Remove(config.Conf.DBPath) diff --git a/models/template.go b/models/template.go index eb97476f..a26c1d8e 100644 --- a/models/template.go +++ b/models/template.go @@ -10,33 +10,34 @@ type Template struct { ModifiedDate time.Time `json:"modified_date"` } +type UserTemplate struct { + UserId int64 `json:"-"` + TemplateId int64 `json:"-"` +} + // 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) + err := db.Table("templates t").Select("t.*").Joins("left join user_templates ut ON t.id = ut.template_id").Where("ut.user_id=?", uid).Scan(&ts).Error 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 - } + err := db.Table("templates t").Select("t.*").Joins("left join user_templates ut ON t.id = ut.template_id").Where("ut.user_id=? and t.id=?", uid, id).Scan(&t).Error 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) + err := db.Save(t).Error 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) + err = db.Exec("INSERT OR IGNORE INTO user_templates VALUES (?,?)", uid, t.Id).Error if err != nil { Logger.Printf("Error adding many-many mapping for template %s\n", t.Name) } diff --git a/models/user.go b/models/user.go index 9f16a1a6..a3915c2c 100644 --- a/models/user.go +++ b/models/user.go @@ -47,6 +47,6 @@ func GetUserByUsername(username string) (User, error) { // PutUser updates the given user func PutUser(u *User) error { - err := db.Update(&u).Error + err := db.Save(u).Error return err }