From c59415a133e15d43d86bfd6fb4e83808f34db514 Mon Sep 17 00:00:00 2001 From: Jordan Date: Thu, 30 Jan 2014 15:08:14 -0600 Subject: [PATCH] Adding some models - Incorporated use of `gorp` package to allow ORM'ish functionality --- auth/auth.go | 18 +++--------------- controllers/route.go | 3 ++- db/db.go | 38 ++++++++++++++++++++------------------ gophish.go | 8 ++------ models/models.go | 37 ++++++++++++++++++++++++++++--------- static/css/main.css | 3 +++ templates/nav.html | 17 ++++++++++++++--- 7 files changed, 72 insertions(+), 52 deletions(-) diff --git a/auth/auth.go b/auth/auth.go index 33033051..8bec4e3f 100644 --- a/auth/auth.go +++ b/auth/auth.go @@ -28,12 +28,8 @@ var Store = sessions.NewCookieStore( func Login(r *http.Request) (bool, error) { username, password := r.FormValue("username"), r.FormValue("password") session, _ := Store.Get(r, "gophish") - stmt, err := db.Conn.Prepare("SELECT * FROM Users WHERE username=?") - if err != nil { - return false, err - } u := models.User{} - err = stmt.QueryRow(username).Scan(&u.Id, &u.Username, &u.Hash, &u.APIKey) + err := db.Conn.SelectOne(&u, "SELECT * FROM Users WHERE username=?", username) if err == sql.ErrNoRows { //Return false, but don't return an error return false, nil @@ -57,11 +53,7 @@ func Login(r *http.Request) (bool, error) { // error is thrown. func GetUserById(id int) (models.User, error) { u := models.User{} - stmt, err := db.Conn.Prepare("SELECT id, username, apikey FROM Users WHERE id=?") - if err != nil { - return u, err - } - err = stmt.QueryRow(id).Scan(&u.Id, &u.Username, &u.APIKey) + err := db.Conn.SelectOne(&u, "SELECT id, username, apikey FROM Users WHERE id=?", id) if err != nil { return u, err } @@ -72,11 +64,7 @@ func GetUserById(id int) (models.User, error) { // error is thrown. func GetUserByAPIKey(key []byte) (models.User, error) { u := models.User{} - stmt, err := db.Conn.Prepare("SELECT id, username, apikey FROM Users WHERE apikey=?") - if err != nil { - return u, err - } - err = stmt.QueryRow(key).Scan(&u.Id, &u.Username, &u.APIKey) + err := db.Conn.SelectOne(&u, "SELECT id, username, apikey FROM Users WHERE apikey=?", key) if err != nil { return u, err } diff --git a/controllers/route.go b/controllers/route.go index a4770fc5..1937da87 100644 --- a/controllers/route.go +++ b/controllers/route.go @@ -85,7 +85,7 @@ func Base(w http.ResponseWriter, r *http.Request) { Flashes []interface{} }{Title: "Dashboard", User: ctx.Get(r, "user").(models.User)} fmt.Println(params.User.Username) - getTemplate(w, "dashboard").ExecuteTemplate(w, "base", nil) + getTemplate(w, "dashboard").ExecuteTemplate(w, "base", params) } func Users(w http.ResponseWriter, r *http.Request) { @@ -123,6 +123,7 @@ func Login(w http.ResponseWriter, r *http.Request) { } succ, err := auth.Login(r) if err != nil { + fmt.Println(err) http.Error(w, "Error logging in", http.StatusInternalServerError) } //If we've logged in, save the session and redirect to the dashboard diff --git a/db/db.go b/db/db.go index 88742cfa..4c5f7331 100644 --- a/db/db.go +++ b/db/db.go @@ -5,47 +5,49 @@ import ( "fmt" "os" + "github.com/coopernurse/gorp" "github.com/jordan-wright/gophish/config" + "github.com/jordan-wright/gophish/models" _ "github.com/mattn/go-sqlite3" ) -var Conn *sql.DB +var Conn *gorp.DbMap +var DB *sql.DB +var err error // Setup initializes the Conn object // It also populates the Gophish Config object -func Setup(reset bool) error { +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) - if err == nil { - os.Remove(config.Conf.DBPath) - } - Conn, err = sql.Open("sqlite3", config.Conf.DBPath) + _, err = os.Stat(config.Conf.DBPath) + Conn.AddTableWithName(models.User{}, "users").SetKeys(true, "Id") + Conn.AddTableWithName(models.Campaign{}, "campaigns").SetKeys(true, "Id") if err != nil { - return err - } - if reset { + fmt.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, apikey VARCHAR(32));`, - `CREATE TABLE Campaigns (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, created_date TEXT NOT NULL, completed_date TEXT, status TEXT NOT NULL);`, + `CREATE TABLE Campaigns (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, created_date TEXT NOT NULL, completed_date TEXT, template TEXT, status TEXT NOT NULL, uid INTEGER, FOREIGN KEY (uid) REFERENCES Users(id));`, } fmt.Println("Creating db at " + config.Conf.DBPath) //Create the tables needed for _, stmt := range createTablesSQL { - _, err = Conn.Exec(stmt) + _, err = DB.Exec(stmt) if err != nil { return err } } //Create the default user - stmt, err := Conn.Prepare(`INSERT INTO Users (username, hash, apikey) VALUES (?, ?, ?);`) - defer stmt.Close() - if err != nil { - return err + init_user := models.User{ + Username: "jordan", + Hash: "$2a$10$d4OtT.RkEOQn.iruVWIQ5u8CeV/85ZYF41y8wKeUwsAPqPNFvTccW", + APIKey: "12345678901234567890123456789012", } - _, err = stmt.Exec("jordan", "$2a$10$d4OtT.RkEOQn.iruVWIQ5u8CeV/85ZYF41y8wKeUwsAPqPNFvTccW", "12345678901234567890123456789012") + Conn.Insert(&init_user) if err != nil { - return err + fmt.Println(err) } } return nil diff --git a/gophish.go b/gophish.go index 6ca935aa..7b003dda 100644 --- a/gophish.go +++ b/gophish.go @@ -26,7 +26,6 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */ import ( - "flag" "fmt" "net/http" @@ -36,13 +35,10 @@ import ( "github.com/jordan-wright/gophish/middleware" ) -var setupFlag = flag.Bool("setup", false, "Starts the initial setup process for Gophish") - func main() { //Setup the global variables and settings - flag.Parse() - err := db.Setup(*setupFlag) - defer db.Conn.Close() + err := db.Setup() + defer db.DB.Close() if err != nil { fmt.Println(err) } diff --git a/models/models.go b/models/models.go index dc97b33a..b6761f47 100644 --- a/models/models.go +++ b/models/models.go @@ -5,8 +5,6 @@ import ( // SMTPServer is used to provide a default SMTP server preference. "time" - - "github.com/jordan-wright/email" ) type SMTPServer struct { @@ -38,11 +36,32 @@ type Flash struct { //Campaign is a struct representing a created campaign type Campaign struct { - Id int `json:"id"` - Name string `json:"name"` - CreatedDate time.Time `json:"created_date"` - CompletedDate time.Time `json:"completed_date"` - Targets []mail.Address `json:"targets"` - Template *email.Email `json:"template"` //This may change - Status string `json:"status"` + Id int `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"` + Uid int +} + +type UserCampaigns struct { + CampaignId int + UserId int +} + +type Result struct { + Id int + TargetId int + Status string `json:"status"` +} + +type CampaignResults struct { + CampaignId int + TargetId int +} + +type Target struct { + Id int `json:"-"` + Email mail.Address `json:"email"` } diff --git a/static/css/main.css b/static/css/main.css index 48f3e29a..d10a3566 100644 --- a/static/css/main.css +++ b/static/css/main.css @@ -7,6 +7,9 @@ padding-top:8px; padding-bottom:0px; } + #navbar-dropdown { + margin-top:8px; + } .sans { font-family:'Open Sans', sans-serif !important; } diff --git a/templates/nav.html b/templates/nav.html index 97b5baa2..2883b1d8 100644 --- a/templates/nav.html +++ b/templates/nav.html @@ -14,9 +14,20 @@