diff --git a/auth/auth.go b/auth/auth.go index 7f87ff10..b9fc8701 100644 --- a/auth/auth.go +++ b/auth/auth.go @@ -16,6 +16,7 @@ import ( //init registers the necessary models to be saved in the session later func init() { gob.Register(&models.User{}) + gob.Register(&models.Flash{}) } var Store = sessions.NewCookieStore( diff --git a/controllers/api.go b/controllers/api.go index 78718bf9..236f321a 100644 --- a/controllers/api.go +++ b/controllers/api.go @@ -1,14 +1,17 @@ package controllers import ( + "crypto/rand" "encoding/json" "fmt" + "io" "net/http" "strconv" "time" ctx "github.com/gorilla/context" "github.com/gorilla/mux" + "github.com/gorilla/sessions" "github.com/jordan-wright/gophish/db" "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. // 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) { switch { case r.Method == "GET": 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 { fmt.Println(err) } + /*for c := range cs { + _, err := db.Conn.Select(&cs.Results, "SELECT r.id ") + }*/ cj, err := json.MarshalIndent(cs, "", " ") if checkError(err, w, "Error looking up campaigns") { return } writeJSON(w, cj) + //POST: Create a new campaign and return it as JSON case r.Method == "POST": c := models.Campaign{} // Put the request into a campaign diff --git a/controllers/route.go b/controllers/route.go index 3c33e91a..f3e9bbb1 100644 --- a/controllers/route.go +++ b/controllers/route.go @@ -28,6 +28,7 @@ func CreateRouter() *mux.Router { // Create the API routes api := router.PathPrefix("/api").Subrouter() 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/{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) { params := struct { - User models.User - Title string + User models.User + Title string + Flashes []interface{} }{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) } @@ -91,6 +96,8 @@ func Login(w http.ResponseWriter, r *http.Request) { session := ctx.Get(r, "session").(*sessions.Session) switch { case r.Method == "GET": + params.Flashes = session.Flashes() + session.Save(r, w) getTemplate(w, "login").ExecuteTemplate(w, "base", params) case r.Method == "POST": //Attempt to login @@ -111,8 +118,8 @@ func Login(w http.ResponseWriter, r *http.Request) { Type: "danger", Message: "Invalid Username/Password", }) - params.Flashes = session.Flashes() - getTemplate(w, "login").ExecuteTemplate(w, "base", params) + session.Save(r, w) + http.Redirect(w, r, "/login", 302) } } } diff --git a/templates/settings.html b/templates/settings.html index b9a61112..c372db28 100644 --- a/templates/settings.html +++ b/templates/settings.html @@ -20,6 +20,7 @@
+ {{%template "flashes" .Flashes%}}

User Settings

@@ -39,6 +40,9 @@
+ + +