github.com/edermi/gophish_mods@v0.7.0/models/models.go (about)

     1  package models
     2  
     3  import (
     4  	"crypto/rand"
     5  	"fmt"
     6  	"io"
     7  
     8  	"bitbucket.org/liamstask/goose/lib/goose"
     9  
    10  	_ "github.com/go-sql-driver/mysql" // Blank import needed to import mysql
    11  	"github.com/gophish/gophish/config"
    12  	log "github.com/gophish/gophish/logger"
    13  	"github.com/jinzhu/gorm"
    14  	_ "github.com/mattn/go-sqlite3" // Blank import needed to import sqlite3
    15  )
    16  
    17  var db *gorm.DB
    18  var err error
    19  
    20  const (
    21  	CAMPAIGN_IN_PROGRESS string = "In progress"
    22  	CAMPAIGN_QUEUED      string = "Queued"
    23  	CAMPAIGN_CREATED     string = "Created"
    24  	CAMPAIGN_EMAILS_SENT string = "Emails Sent"
    25  	CAMPAIGN_COMPLETE    string = "Completed"
    26  	EVENT_SENT           string = "Email Sent"
    27  	EVENT_SENDING_ERROR  string = "Error Sending Email"
    28  	EVENT_OPENED         string = "Email Opened"
    29  	EVENT_CLICKED        string = "Clicked Link"
    30  	EVENT_DATA_SUBMIT    string = "Submitted Data"
    31  	EVENT_REPORTED       string = "Email Reported"
    32  	EVENT_PROXY_REQUEST  string = "Proxied request"
    33  	STATUS_SUCCESS       string = "Success"
    34  	STATUS_QUEUED        string = "Queued"
    35  	STATUS_SENDING       string = "Sending"
    36  	STATUS_UNKNOWN       string = "Unknown"
    37  	STATUS_SCHEDULED     string = "Scheduled"
    38  	STATUS_RETRY         string = "Retrying"
    39  	ERROR                string = "Error"
    40  )
    41  
    42  // Flash is used to hold flash information for use in templates.
    43  type Flash struct {
    44  	Type    string
    45  	Message string
    46  }
    47  
    48  // Response contains the attributes found in an API response
    49  type Response struct {
    50  	Message string      `json:"message"`
    51  	Success bool        `json:"success"`
    52  	Data    interface{} `json:"data"`
    53  }
    54  
    55  // Copy of auth.GenerateSecureKey to prevent cyclic import with auth library
    56  func generateSecureKey() string {
    57  	k := make([]byte, 32)
    58  	io.ReadFull(rand.Reader, k)
    59  	return fmt.Sprintf("%x", k)
    60  }
    61  
    62  func chooseDBDriver(name, openStr string) goose.DBDriver {
    63  	d := goose.DBDriver{Name: name, OpenStr: openStr}
    64  
    65  	switch name {
    66  	case "mysql":
    67  		d.Import = "github.com/go-sql-driver/mysql"
    68  		d.Dialect = &goose.MySqlDialect{}
    69  
    70  	// Default database is sqlite3
    71  	default:
    72  		d.Import = "github.com/mattn/go-sqlite3"
    73  		d.Dialect = &goose.Sqlite3Dialect{}
    74  	}
    75  
    76  	return d
    77  }
    78  
    79  // Setup initializes the Conn object
    80  // It also populates the Gophish Config object
    81  func Setup() error {
    82  	// Setup the goose configuration
    83  	migrateConf := &goose.DBConf{
    84  		MigrationsDir: config.Conf.MigrationsPath,
    85  		Env:           "production",
    86  		Driver:        chooseDBDriver(config.Conf.DBName, config.Conf.DBPath),
    87  	}
    88  	// Get the latest possible migration
    89  	latest, err := goose.GetMostRecentDBVersion(migrateConf.MigrationsDir)
    90  	if err != nil {
    91  		log.Error(err)
    92  		return err
    93  	}
    94  	// Open our database connection
    95  	db, err = gorm.Open(config.Conf.DBName, config.Conf.DBPath)
    96  	db.LogMode(false)
    97  	db.SetLogger(log.Logger)
    98  	db.DB().SetMaxOpenConns(1)
    99  	if err != nil {
   100  		log.Error(err)
   101  		return err
   102  	}
   103  	// Migrate up to the latest version
   104  	err = goose.RunMigrationsOnDb(migrateConf, migrateConf.MigrationsDir, latest, db.DB())
   105  	if err != nil {
   106  		log.Error(err)
   107  		return err
   108  	}
   109  	// Create the admin user if it doesn't exist
   110  	var userCount int64
   111  	db.Model(&User{}).Count(&userCount)
   112  	if userCount == 0 {
   113  		initUser := User{
   114  			Username: "admin",
   115  			Hash:     "$2a$10$IYkPp0.QsM81lYYPrQx6W.U6oQGw7wMpozrKhKAHUBVL4mkm/EvAS", //gophish
   116  		}
   117  		initUser.ApiKey = generateSecureKey()
   118  		err = db.Save(&initUser).Error
   119  		if err != nil {
   120  			log.Error(err)
   121  			return err
   122  		}
   123  	}
   124  	return nil
   125  }