Added ability to reset API token

Cleaned up session flash handling
pull/24/head
Jordan 2014-02-02 14:47:06 -06:00
parent 2cfe2b8f8f
commit e312e90570
4 changed files with 45 additions and 5 deletions

View File

@ -16,6 +16,7 @@ import (
//init registers the necessary models to be saved in the session later //init registers the necessary models to be saved in the session later
func init() { func init() {
gob.Register(&models.User{}) gob.Register(&models.User{})
gob.Register(&models.Flash{})
} }
var Store = sessions.NewCookieStore( var Store = sessions.NewCookieStore(

View File

@ -1,14 +1,17 @@
package controllers package controllers
import ( import (
"crypto/rand"
"encoding/json" "encoding/json"
"fmt" "fmt"
"io"
"net/http" "net/http"
"strconv" "strconv"
"time" "time"
ctx "github.com/gorilla/context" ctx "github.com/gorilla/context"
"github.com/gorilla/mux" "github.com/gorilla/mux"
"github.com/gorilla/sessions"
"github.com/jordan-wright/gophish/db" "github.com/jordan-wright/gophish/db"
"github.com/jordan-wright/gophish/models" "github.com/jordan-wright/gophish/models"
) )
@ -28,21 +31,46 @@ func API(w http.ResponseWriter, r *http.Request) {
} }
} }
// API (/api/reset) resets a user's API key
func API_Reset(w http.ResponseWriter, r *http.Request) {
switch {
case r.Method == "GET":
u := ctx.Get(r, "user").(models.User)
// Inspired from gorilla/securecookie
k := make([]byte, 32)
_, err := io.ReadFull(rand.Reader, k)
checkError(err, w, "Error setting new API key")
u.APIKey = fmt.Sprintf("%x", k)
db.Conn.Exec("UPDATE users SET api_key=? WHERE id=?", u.APIKey, u.Id)
session := ctx.Get(r, "session").(*sessions.Session)
session.AddFlash(models.Flash{
Type: "success",
Message: "API Key Successfully Reset",
})
session.Save(r, w)
http.Redirect(w, r, "/settings", 302)
}
}
// API_Campaigns returns a list of campaigns if requested via GET. // API_Campaigns returns a list of campaigns if requested via GET.
// If requested via POST, API_Campaigns creates a new campaign and returns a reference to it. // If requested via POST, API_Campaigns creates a new campaign and returns a reference to it.
func API_Campaigns(w http.ResponseWriter, r *http.Request) { func API_Campaigns(w http.ResponseWriter, r *http.Request) {
switch { switch {
case r.Method == "GET": case r.Method == "GET":
cs := []models.Campaign{} cs := []models.Campaign{}
_, err := db.Conn.Select(&cs, "SELECT campaigns.id, name, created_date, completed_date, status, template FROM campaigns, users WHERE campaigns.uid=users.id AND users.api_key=?", ctx.Get(r, "api_key")) _, err := db.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=?", ctx.Get(r, "api_key"))
if err != nil { if err != nil {
fmt.Println(err) fmt.Println(err)
} }
/*for c := range cs {
_, err := db.Conn.Select(&cs.Results, "SELECT r.id ")
}*/
cj, err := json.MarshalIndent(cs, "", " ") cj, err := json.MarshalIndent(cs, "", " ")
if checkError(err, w, "Error looking up campaigns") { if checkError(err, w, "Error looking up campaigns") {
return return
} }
writeJSON(w, cj) writeJSON(w, cj)
//POST: Create a new campaign and return it as JSON
case r.Method == "POST": case r.Method == "POST":
c := models.Campaign{} c := models.Campaign{}
// Put the request into a campaign // Put the request into a campaign

View File

@ -28,6 +28,7 @@ func CreateRouter() *mux.Router {
// Create the API routes // Create the API routes
api := router.PathPrefix("/api").Subrouter() api := router.PathPrefix("/api").Subrouter()
api.HandleFunc("/", Use(API, mid.RequireLogin)) api.HandleFunc("/", Use(API, mid.RequireLogin))
api.HandleFunc("/reset", Use(API_Reset, mid.RequireLogin))
api.HandleFunc("/campaigns", Use(API_Campaigns, mid.RequireAPIKey)) api.HandleFunc("/campaigns", Use(API_Campaigns, mid.RequireAPIKey))
api.HandleFunc("/campaigns/{id}", Use(API_Campaigns_Id, mid.RequireAPIKey)) api.HandleFunc("/campaigns/{id}", Use(API_Campaigns_Id, mid.RequireAPIKey))
@ -67,9 +68,13 @@ func Users(w http.ResponseWriter, r *http.Request) {
func Settings(w http.ResponseWriter, r *http.Request) { func Settings(w http.ResponseWriter, r *http.Request) {
params := struct { params := struct {
User models.User User models.User
Title string Title string
Flashes []interface{}
}{Title: "Settings", User: ctx.Get(r, "user").(models.User)} }{Title: "Settings", User: ctx.Get(r, "user").(models.User)}
session := ctx.Get(r, "session").(*sessions.Session)
params.Flashes = session.Flashes()
session.Save(r, w)
getTemplate(w, "settings").ExecuteTemplate(w, "base", params) getTemplate(w, "settings").ExecuteTemplate(w, "base", params)
} }
@ -91,6 +96,8 @@ func Login(w http.ResponseWriter, r *http.Request) {
session := ctx.Get(r, "session").(*sessions.Session) session := ctx.Get(r, "session").(*sessions.Session)
switch { switch {
case r.Method == "GET": case r.Method == "GET":
params.Flashes = session.Flashes()
session.Save(r, w)
getTemplate(w, "login").ExecuteTemplate(w, "base", params) getTemplate(w, "login").ExecuteTemplate(w, "base", params)
case r.Method == "POST": case r.Method == "POST":
//Attempt to login //Attempt to login
@ -111,8 +118,8 @@ func Login(w http.ResponseWriter, r *http.Request) {
Type: "danger", Type: "danger",
Message: "Invalid Username/Password", Message: "Invalid Username/Password",
}) })
params.Flashes = session.Flashes() session.Save(r, w)
getTemplate(w, "login").ExecuteTemplate(w, "base", params) http.Redirect(w, r, "/login", 302)
} }
} }
} }

View File

@ -20,6 +20,7 @@
</ul> </ul>
</div> </div>
<div class="col-md-9 sans"> <div class="col-md-9 sans">
{{%template "flashes" .Flashes%}}
<h1 style="margin:0px 0px 15px 0px;">User Settings</h1> <h1 style="margin:0px 0px 15px 0px;">User Settings</h1>
<div class="row"> <div class="row">
<div class="col-md-2"> <div class="col-md-2">
@ -39,6 +40,9 @@
<div class="col-md-6"> <div class="col-md-6">
<input type="text" value="{{%.User.APIKey%}}" class="form-control" readonly/> <input type="text" value="{{%.User.APIKey%}}" class="form-control" readonly/>
</div> </div>
<a href="/api/reset">
<button class="btn btn-primary"><i class="fa fa-refresh"></i> Reset</button>
</a>
</div> </div>
<br /> <br />
<button class="btn btn-primary">Save</button> <button class="btn btn-primary">Save</button>