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!
pull/24/head
Jordan 2014-03-24 22:31:33 -05:00
parent 4dd9a5cc53
commit 584d7dbc23
12 changed files with 499 additions and 498 deletions

View File

@ -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

View File

@ -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

View File

@ -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
}

View File

@ -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"

404
db/db.go
View File

@ -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
}

View File

@ -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)

View File

@ -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

101
models/campaign.go Normal file
View File

@ -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
}

193
models/group.go Normal file
View File

@ -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
}

View File

@ -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"`
}

48
models/template.go Normal file
View File

@ -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
}

52
models/user.go Normal file
View File

@ -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
}