mirror of https://github.com/gophish/gophish
135 lines
3.4 KiB
Go
135 lines
3.4 KiB
Go
package models
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"fmt"
|
|
"io"
|
|
|
|
"bitbucket.org/liamstask/goose/lib/goose"
|
|
|
|
_ "github.com/go-sql-driver/mysql" // Blank import needed to import mysql
|
|
"github.com/gophish/gophish/config"
|
|
log "github.com/gophish/gophish/logger"
|
|
"github.com/jinzhu/gorm"
|
|
_ "github.com/mattn/go-sqlite3" // Blank import needed to import sqlite3
|
|
)
|
|
|
|
var db *gorm.DB
|
|
var conf *config.Config
|
|
|
|
const (
|
|
CampaignInProgress string = "In progress"
|
|
CampaignQueued string = "Queued"
|
|
CampaignCreated string = "Created"
|
|
CampaignEmailsSent string = "Emails Sent"
|
|
CampaignComplete string = "Completed"
|
|
EventSent string = "Email Sent"
|
|
EventSendingError string = "Error Sending Email"
|
|
EventOpened string = "Email Opened"
|
|
EventClicked string = "Clicked Link"
|
|
EventDataSubmit string = "Submitted Data"
|
|
EventReported string = "Email Reported"
|
|
EventProxyRequest string = "Proxied request"
|
|
StatusSuccess string = "Success"
|
|
StatusQueued string = "Queued"
|
|
StatusSending string = "Sending"
|
|
StatusUnknown string = "Unknown"
|
|
StatusScheduled string = "Scheduled"
|
|
StatusRetry string = "Retrying"
|
|
Error string = "Error"
|
|
)
|
|
|
|
// Flash is used to hold flash information for use in templates.
|
|
type Flash struct {
|
|
Type string
|
|
Message string
|
|
}
|
|
|
|
// Response contains the attributes found in an API response
|
|
type Response struct {
|
|
Message string `json:"message"`
|
|
Success bool `json:"success"`
|
|
Data interface{} `json:"data"`
|
|
}
|
|
|
|
// Copy of auth.GenerateSecureKey to prevent cyclic import with auth library
|
|
func generateSecureKey() string {
|
|
k := make([]byte, 32)
|
|
io.ReadFull(rand.Reader, k)
|
|
return fmt.Sprintf("%x", k)
|
|
}
|
|
|
|
func chooseDBDriver(name, openStr string) goose.DBDriver {
|
|
d := goose.DBDriver{Name: name, OpenStr: openStr}
|
|
|
|
switch name {
|
|
case "mysql":
|
|
d.Import = "github.com/go-sql-driver/mysql"
|
|
d.Dialect = &goose.MySqlDialect{}
|
|
|
|
// Default database is sqlite3
|
|
default:
|
|
d.Import = "github.com/mattn/go-sqlite3"
|
|
d.Dialect = &goose.Sqlite3Dialect{}
|
|
}
|
|
|
|
return d
|
|
}
|
|
|
|
// Setup initializes the Conn object
|
|
// It also populates the Gophish Config object
|
|
func Setup(c *config.Config) error {
|
|
// Setup the package-scoped config
|
|
conf = c
|
|
// Setup the goose configuration
|
|
migrateConf := &goose.DBConf{
|
|
MigrationsDir: conf.MigrationsPath,
|
|
Env: "production",
|
|
Driver: chooseDBDriver(conf.DBName, conf.DBPath),
|
|
}
|
|
// Get the latest possible migration
|
|
latest, err := goose.GetMostRecentDBVersion(migrateConf.MigrationsDir)
|
|
if err != nil {
|
|
log.Error(err)
|
|
return err
|
|
}
|
|
// Open our database connection
|
|
db, err = gorm.Open(conf.DBName, conf.DBPath)
|
|
db.LogMode(false)
|
|
db.SetLogger(log.Logger)
|
|
db.DB().SetMaxOpenConns(1)
|
|
if err != nil {
|
|
log.Error(err)
|
|
return err
|
|
}
|
|
// Migrate up to the latest version
|
|
err = goose.RunMigrationsOnDb(migrateConf, migrateConf.MigrationsDir, latest, db.DB())
|
|
if err != nil {
|
|
log.Error(err)
|
|
return err
|
|
}
|
|
// Create the admin user if it doesn't exist
|
|
var userCount int64
|
|
db.Model(&User{}).Count(&userCount)
|
|
adminRole, err := GetRoleBySlug(RoleAdmin)
|
|
if err != nil {
|
|
log.Error(err)
|
|
return err
|
|
}
|
|
if userCount == 0 {
|
|
initUser := User{
|
|
Username: "admin",
|
|
Hash: "$2a$10$IYkPp0.QsM81lYYPrQx6W.U6oQGw7wMpozrKhKAHUBVL4mkm/EvAS", //gophish
|
|
Role: adminRole,
|
|
RoleID: adminRole.ID,
|
|
}
|
|
initUser.ApiKey = generateSecureKey()
|
|
err = db.Save(&initUser).Error
|
|
if err != nil {
|
|
log.Error(err)
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|