/
models.go
131 lines (116 loc) · 3.46 KB
/
models.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package models
import (
"crypto/rand"
"errors"
"fmt"
"io"
"log"
"os"
"bitbucket.org/liamstask/goose/lib/goose"
_ "github.com/go-sql-driver/mysql"
"github.com/gophish/gophish/config"
"github.com/jinzhu/gorm"
_ "github.com/mattn/go-sqlite3" // Blank import needed to import sqlite3
)
var db *gorm.DB
var err error
// ErrUsernameTaken is thrown when a user attempts to register a username that is taken.
var ErrUsernameTaken = errors.New("username already taken")
// Logger is a global logger used to show informational, warning, and error messages
var Logger = log.New(os.Stdout, " ", log.Ldate|log.Ltime|log.Lshortfile)
const (
CAMPAIGN_IN_PROGRESS string = "In progress"
CAMPAIGN_QUEUED string = "Queued"
CAMPAIGN_CREATED string = "Created"
CAMPAIGN_EMAILS_SENT string = "Emails Sent"
CAMPAIGN_COMPLETE string = "Completed"
EVENT_SENT string = "Email Sent"
EVENT_SENDING_ERROR string = "Error Sending Email"
EVENT_OPENED string = "Email Opened"
EVENT_CLICKED string = "Clicked Link"
EVENT_DATA_SUBMIT string = "Submitted Data"
STATUS_SUCCESS string = "Success"
STATUS_SENDING string = "Sending"
STATUS_UNKNOWN string = "Unknown"
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() error {
create_db := false
if _, err = os.Stat(config.Conf.DBPath); err != nil || config.Conf.DBPath == ":memory:" {
create_db = true
}
// Setup the goose configuration
migrateConf := &goose.DBConf{
MigrationsDir: config.Conf.MigrationsPath,
Env: "production",
Driver: chooseDBDriver(config.Conf.DBName, config.Conf.DBPath),
}
// Get the latest possible migration
latest, err := goose.GetMostRecentDBVersion(migrateConf.MigrationsDir)
if err != nil {
Logger.Println(err)
return err
}
// Open our database connection
db, err = gorm.Open(config.Conf.DBName, config.Conf.DBPath)
db.LogMode(false)
db.SetLogger(Logger)
db.DB().SetMaxOpenConns(1)
if err != nil {
Logger.Println(err)
return err
}
// Migrate up to the latest version
err = goose.RunMigrationsOnDb(migrateConf, migrateConf.MigrationsDir, latest, db.DB())
if err != nil {
Logger.Println(err)
return err
}
//If the database didn't exist, we need to create the admin user
if create_db {
//Create the default user
initUser := User{
Username: "admin",
Hash: "$2a$10$IYkPp0.QsM81lYYPrQx6W.U6oQGw7wMpozrKhKAHUBVL4mkm/EvAS", //gophish
}
initUser.ApiKey = generateSecureKey()
err = db.Save(&initUser).Error
if err != nil {
Logger.Println(err)
return err
}
}
return nil
}