Browse Source

Added easier support for Flashes

Moving DB access (as much as possible) into `db` package.
pull/24/head
Jordan 7 years ago
parent
commit
af7a8f4c4e
4 changed files with 70 additions and 14 deletions
  1. +13
    -11
      controllers/route.go
  2. +54
    -0
      db/db.go
  3. +2
    -2
      middleware/middleware.go
  4. +1
    -1
      templates/users.html

+ 13
- 11
controllers/route.go View File

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

+ 54
- 0
db/db.go View File

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

+ 2
- 2
middleware/middleware.go View File

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

+ 1
- 1
templates/users.html View File

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

Loading…
Cancel
Save