From af7a8f4c4e19b4ff92db0995adbc0cdeb3bd7d27 Mon Sep 17 00:00:00 2001 From: Jordan Date: Wed, 5 Feb 2014 10:57:53 -0600 Subject: [PATCH] Added easier support for Flashes Moving DB access (as much as possible) into `db` package. --- controllers/route.go | 24 ++++++++++-------- db/db.go | 54 ++++++++++++++++++++++++++++++++++++++++ middleware/middleware.go | 4 +-- templates/users.html | 2 +- 4 files changed, 70 insertions(+), 14 deletions(-) diff --git a/controllers/route.go b/controllers/route.go index 1a141ab4..d1c5cb60 100644 --- a/controllers/route.go +++ b/controllers/route.go @@ -9,6 +9,7 @@ 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" @@ -87,7 +88,7 @@ func Register(w http.ResponseWriter, r *http.Request) { } else { // Check the error m := "" - if err == auth.ErrUsernameTaken { + if err == db.ErrUsernameTaken { m = "Username already taken" } else { m = "Unknown error - please try again" @@ -110,11 +111,7 @@ func Logout(w http.ResponseWriter, r *http.Request) { // Now that we are all registered, we can log the user in session := ctx.Get(r, "session").(*sessions.Session) delete(session.Values, "id") - session.AddFlash(models.Flash{ - Type: "success", - Message: "You have successfully logged out.", - }) - session.Save(r, w) + Flash(w, r, "success", "You have successfully logged out") http.Redirect(w, r, "login", 302) } @@ -184,11 +181,7 @@ func Login(w http.ResponseWriter, r *http.Request) { session.Save(r, w) http.Redirect(w, r, "/", 302) } else { - session.AddFlash(models.Flash{ - Type: "danger", - Message: "Invalid Username/Password", - }) - session.Save(r, w) + Flash(w, r, "danger", "Invalid Username/Password") http.Redirect(w, r, "/login", 302) } } @@ -212,3 +205,12 @@ func checkError(e error, w http.ResponseWriter, m string) bool { } return false } + +func Flash(w http.ResponseWriter, r *http.Request, t string, m string) { + session := ctx.Get(r, "session").(*sessions.Session) + session.AddFlash(models.Flash{ + Type: t, + Message: m, + }) + session.Save(r, w) +} diff --git a/db/db.go b/db/db.go index 947d992d..8dc39e6d 100644 --- a/db/db.go +++ b/db/db.go @@ -2,6 +2,7 @@ package db import ( "database/sql" + "errors" "fmt" "os" "time" @@ -15,6 +16,7 @@ import ( var Conn *gorp.DbMap var DB *sql.DB var err error +var ErrUsernameTaken = errors.New("Username already taken") // Setup initializes the Conn object // It also populates the Gophish Config object @@ -67,3 +69,55 @@ func Setup() error { } 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 +} + +// GetUserByAPIKey returns the user that the given API Key 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 +} + +func PutUser(u *models.User) error { + _, err := Conn.Update(u) + return err +} + +func GetCampaigns(key interface{}) ([]models.Campaign, error) { + cs := []models.Campaign{} + _, err := Conn.Select(&cs, "SELECT c.id, name, created_date, completed_date, status, template FROM campaigns c, users u WHERE c.uid=u.id AND u.api_key=?", key) + return cs, err +} + +func GetCampaign(id int64) { + +} diff --git a/middleware/middleware.go b/middleware/middleware.go index 2f74417c..040012bc 100644 --- a/middleware/middleware.go +++ b/middleware/middleware.go @@ -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 := auth.GetUserById(id.(int64)) + u, err := db.GetUser(id.(int64)) if err != nil { ctx.Set(r, "user", nil) } @@ -47,7 +47,7 @@ func RequireAPIKey(handler http.Handler) http.HandlerFunc { } else { id, err := db.Conn.SelectInt("SELECT id FROM users WHERE api_key=?", ak) if id == 0 || err != nil { - http.Error(w, "Error: Invalid API Key", http.StatusInternalServerError) + JSONError(w, 500, "Invalid API Key") return } ctx.Set(r, "user_id", id) diff --git a/templates/users.html b/templates/users.html index 1c23e57b..f1c1a6c2 100644 --- a/templates/users.html +++ b/templates/users.html @@ -36,7 +36,7 @@ {{group.name}} - {{group.status}} + {{group.targets}} {{group.modified_date | date:'medium'}}