Added easier support for Flashes

Moving DB access (as much as possible) into `db` package.

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)
+}

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) {
+
+}

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)

templates/users.html

@@ -36,7 +36,7 @@
                 <tbody>
                     <tr ng-repeat="group in groups">
                         <td>{{group.name}}</td>
-                        <td>{{group.status}}</td>
+                        <td>{{group.targets}}</td>
                         <td>{{group.modified_date | date:'medium'}}</td>
                     </tr>
                 </tbody>