Adding some models - Incorporated use of `gorp` package to allow ORM'ish functionality

auth/auth.go

@@ -28,12 +28,8 @@ var Store = sessions.NewCookieStore(
 func Login(r *http.Request) (bool, error) {
 	username, password := r.FormValue("username"), r.FormValue("password")
 	session, _ := Store.Get(r, "gophish")
-	stmt, err := db.Conn.Prepare("SELECT * FROM Users WHERE username=?")
-	if err != nil {
-		return false, err
-	}
 	u := models.User{}
-	err = stmt.QueryRow(username).Scan(&u.Id, &u.Username, &u.Hash, &u.APIKey)
+	err := db.Conn.SelectOne(&u, "SELECT * FROM Users WHERE username=?", username)
 	if err == sql.ErrNoRows {
 		//Return false, but don't return an error
 		return false, nil
@@ -57,11 +53,7 @@ func Login(r *http.Request) (bool, error) {
 // error is thrown.
 func GetUserById(id int) (models.User, error) {
 	u := models.User{}
-	stmt, err := db.Conn.Prepare("SELECT id, username, apikey FROM Users WHERE id=?")
+	err := db.Conn.SelectOne(&u, "SELECT id, username, apikey FROM Users WHERE id=?", id)
-	if err != nil {
-		return u, err
-	}
-	err = stmt.QueryRow(id).Scan(&u.Id, &u.Username, &u.APIKey)
 	if err != nil {
 		return u, err
 	}
@@ -72,11 +64,7 @@ func GetUserById(id int) (models.User, error) {
 // error is thrown.
 func GetUserByAPIKey(key []byte) (models.User, error) {
 	u := models.User{}
-	stmt, err := db.Conn.Prepare("SELECT id, username, apikey FROM Users WHERE apikey=?")
+	err := db.Conn.SelectOne(&u, "SELECT id, username, apikey FROM Users WHERE apikey=?", key)
-	if err != nil {
-		return u, err
-	}
-	err = stmt.QueryRow(key).Scan(&u.Id, &u.Username, &u.APIKey)
 	if err != nil {
 		return u, err
 	}

controllers/route.go

@@ -85,7 +85,7 @@ func Base(w http.ResponseWriter, r *http.Request) {
 		Flashes []interface{}
 	}{Title: "Dashboard", User: ctx.Get(r, "user").(models.User)}
 	fmt.Println(params.User.Username)
-	getTemplate(w, "dashboard").ExecuteTemplate(w, "base", nil)
+	getTemplate(w, "dashboard").ExecuteTemplate(w, "base", params)
 }
 
 func Users(w http.ResponseWriter, r *http.Request) {
@@ -123,6 +123,7 @@ func Login(w http.ResponseWriter, r *http.Request) {
 		}
 		succ, err := auth.Login(r)
 		if err != nil {
+			fmt.Println(err)
 			http.Error(w, "Error logging in", http.StatusInternalServerError)
 		}
 		//If we've logged in, save the session and redirect to the dashboard

db/db.go

@@ -5,47 +5,49 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/coopernurse/gorp"
 	"github.com/jordan-wright/gophish/config"
+	"github.com/jordan-wright/gophish/models"
 	_ "github.com/mattn/go-sqlite3"
 )
 
-var Conn *sql.DB
+var Conn *gorp.DbMap
+var DB *sql.DB
+var err error
 
 // Setup initializes the Conn object
 // It also populates the Gophish Config object
-func Setup(reset bool) error {
+func Setup() error {
+	DB, err := sql.Open("sqlite3", config.Conf.DBPath)
+	Conn = &gorp.DbMap{Db: DB, Dialect: gorp.SqliteDialect{}}
 	//If the file already exists, delete it and recreate it
-	_, err := os.Stat(config.Conf.DBPath)
+	_, err = os.Stat(config.Conf.DBPath)
-	if err == nil {
+	Conn.AddTableWithName(models.User{}, "users").SetKeys(true, "Id")
-		os.Remove(config.Conf.DBPath)
+	Conn.AddTableWithName(models.Campaign{}, "campaigns").SetKeys(true, "Id")
-	}
-	Conn, err = sql.Open("sqlite3", config.Conf.DBPath)
 	if err != nil {
-		return err
+		fmt.Println("Database not found, recreating...")
-	}
-	if reset {
 		createTablesSQL := []string{
 			//Create tables
 			`CREATE TABLE Users (id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT NOT NULL, hash VARCHAR(60) NOT NULL, apikey VARCHAR(32));`,
-			`CREATE TABLE Campaigns (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, created_date TEXT NOT NULL, completed_date TEXT, status TEXT NOT NULL);`,
+			`CREATE TABLE Campaigns (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL, created_date TEXT NOT NULL, completed_date TEXT, template TEXT, status TEXT NOT NULL, uid INTEGER, FOREIGN KEY (uid) REFERENCES Users(id));`,
 		}
 		fmt.Println("Creating db at " + config.Conf.DBPath)
 		//Create the tables needed
 		for _, stmt := range createTablesSQL {
-			_, err = Conn.Exec(stmt)
+			_, err = DB.Exec(stmt)
 			if err != nil {
 				return err
 			}
 		}
 		//Create the default user
-		stmt, err := Conn.Prepare(`INSERT INTO Users (username, hash, apikey) VALUES (?, ?, ?);`)
+		init_user := models.User{
-		defer stmt.Close()
+			Username: "jordan",
-		if err != nil {
+			Hash:     "$2a$10$d4OtT.RkEOQn.iruVWIQ5u8CeV/85ZYF41y8wKeUwsAPqPNFvTccW",
-			return err
+			APIKey:   "12345678901234567890123456789012",
 		}
-		_, err = stmt.Exec("jordan", "$2a$10$d4OtT.RkEOQn.iruVWIQ5u8CeV/85ZYF41y8wKeUwsAPqPNFvTccW", "12345678901234567890123456789012")
+		Conn.Insert(&init_user)
 		if err != nil {
-			return err
+			fmt.Println(err)
 		}
 	}
 	return nil

gophish.go

@@ -26,7 +26,6 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 */
 import (
-	"flag"
 	"fmt"
 	"net/http"
 
@@ -36,13 +35,10 @@ import (
 	"github.com/jordan-wright/gophish/middleware"
 )
 
-var setupFlag = flag.Bool("setup", false, "Starts the initial setup process for Gophish")
-
 func main() {
 	//Setup the global variables and settings
-	flag.Parse()
+	err := db.Setup()
-	err := db.Setup(*setupFlag)
+	defer db.DB.Close()
-	defer db.Conn.Close()
 	if err != nil {
 		fmt.Println(err)
 	}

models/models.go

@@ -5,8 +5,6 @@ import (
 
 	// SMTPServer is used to provide a default SMTP server preference.
 	"time"
-
-	"github.com/jordan-wright/email"
 )
 
 type SMTPServer struct {
@@ -38,11 +36,32 @@ type Flash struct {
 
 //Campaign is a struct representing a created campaign
 type Campaign struct {
-	Id            int            `json:"id"`
+	Id            int       `json:"id"`
-	Name          string         `json:"name"`
+	Name          string    `json:"name"`
-	CreatedDate   time.Time      `json:"created_date"`
+	CreatedDate   time.Time `json:"created_date" db:"created_date"`
-	CompletedDate time.Time      `json:"completed_date"`
+	CompletedDate time.Time `json:"completed_date" db:"completed_date"`
-	Targets       []mail.Address `json:"targets"`
+	Template      string    `json:"template"` //This may change
-	Template      *email.Email   `json:"template"` //This may change
+	Status        string    `json:"status"`
-	Status        string         `json:"status"`
+	Uid           int
+}
+
+type UserCampaigns struct {
+	CampaignId int
+	UserId     int
+}
+
+type Result struct {
+	Id       int
+	TargetId int
+	Status   string `json:"status"`
+}
+
+type CampaignResults struct {
+	CampaignId int
+	TargetId   int
+}
+
+type Target struct {
+	Id    int          `json:"-"`
+	Email mail.Address `json:"email"`
 }

static/css/main.css

@@ -7,6 +7,9 @@
       padding-top:8px;
       padding-bottom:0px;
   }
+  #navbar-dropdown {
+    margin-top:8px;
+  }
   .sans {
       font-family:'Open Sans', sans-serif !important;
   }

templates/nav.html

@@ -14,9 +14,20 @@
             <ul class="nav navbar-nav navbar-right">
                 <li>
                     {{if .}}
-                    <a href="/settings">
+                    <div class="btn-group" id="navbar-dropdown">
                         <button type="button" class="btn btn-primary"><i class="fa fa-user"></i> {{.Username}}</button>
-                    </a>
+                        <button type="button" class="btn btn-primary dropdown-toggle" data-toggle="dropdown">
+                            <span class="caret" style="border-top-color:#FFFFFF"></span>
+                            <span class="sr-only">Toggle Dropdown</span>
+                        </button>
+                        <ul class="dropdown-menu" role="menu">
+                            <li><a href="/settings">Settings</a>
+                            </li>
+                            <li class="divider"></li>
+                            <li><a href="/logout">Logout</a>
+                            </li>
+                        </ul>
+                    </div>
                     {{else}}
                     <a href="/login">
                         <button type="button" class="btn btn-primary">Login</button>