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

pull/24/head
Jordan 2014-01-30 15:08:14 -06:00
parent e993e8c898
commit c59415a133
7 changed files with 72 additions and 52 deletions

View File

@ -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=?")
if err != nil {
return u, err
}
err = stmt.QueryRow(id).Scan(&u.Id, &u.Username, &u.APIKey)
err := db.Conn.SelectOne(&u, "SELECT id, username, apikey FROM Users WHERE id=?", id)
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=?")
if err != nil {
return u, err
}
err = stmt.QueryRow(key).Scan(&u.Id, &u.Username, &u.APIKey)
err := db.Conn.SelectOne(&u, "SELECT id, username, apikey FROM Users WHERE apikey=?", key)
if err != nil {
return u, err
}

View File

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

View File

@ -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)
if err == nil {
os.Remove(config.Conf.DBPath)
}
Conn, err = sql.Open("sqlite3", config.Conf.DBPath)
_, err = os.Stat(config.Conf.DBPath)
Conn.AddTableWithName(models.User{}, "users").SetKeys(true, "Id")
Conn.AddTableWithName(models.Campaign{}, "campaigns").SetKeys(true, "Id")
if err != nil {
return err
}
if reset {
fmt.Println("Database not found, recreating...")
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 (?, ?, ?);`)
defer stmt.Close()
if err != nil {
return err
init_user := models.User{
Username: "jordan",
Hash: "$2a$10$d4OtT.RkEOQn.iruVWIQ5u8CeV/85ZYF41y8wKeUwsAPqPNFvTccW",
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

View File

@ -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(*setupFlag)
defer db.Conn.Close()
err := db.Setup()
defer db.DB.Close()
if err != nil {
fmt.Println(err)
}

View File

@ -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"`
Name string `json:"name"`
CreatedDate time.Time `json:"created_date"`
CompletedDate time.Time `json:"completed_date"`
Targets []mail.Address `json:"targets"`
Template *email.Email `json:"template"` //This may change
Status string `json:"status"`
Id int `json:"id"`
Name string `json:"name"`
CreatedDate time.Time `json:"created_date" db:"created_date"`
CompletedDate time.Time `json:"completed_date" db:"completed_date"`
Template string `json:"template"` //This may change
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"`
}

View File

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

View File

@ -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>
@ -27,4 +38,4 @@
</div>
</div>
</div>
{{end}}
{{end}}