WIP - Refactoring migrations to support custom logic.

db/db_sqlite3/migrations/2018092312000000_0.8.0_result_urls.sql

@@ -0,0 +1,6 @@
+-- +goose Up
+-- SQL in this section is executed when the migration is applied.
+ALTER TABLE results ADD COLUMN `url` VARCHAR(255);
+
+-- +goose Down
+-- SQL in this section is executed when the migration is rolled back.

gophish.go

@@ -41,6 +41,7 @@ import (
 	"github.com/gophish/gophish/controllers"
 	log "github.com/gophish/gophish/logger"
 	"github.com/gophish/gophish/mailer"
+	"github.com/gophish/gophish/migrations"
 	"github.com/gophish/gophish/models"
 	"github.com/gophish/gophish/util"
 	"github.com/gorilla/handlers"
@@ -71,15 +72,20 @@ func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	// Provide the option to disable the built-in mailer
+	// Run the database migrations to ensure the schema is up-to-date
-	if !*disableMailer {
+	err = migrations.Migrate()
-		go mailer.Mailer.Start(ctx)
+	if err != nil {
+		log.Fatal(err)
 	}
 	// Setup the global variables and settings
 	err = models.Setup()
 	if err != nil {
 		log.Fatal(err)
 	}
+	// Provide the option to disable the built-in mailer
+	if !*disableMailer {
+		go mailer.Mailer.Start(ctx)
+	}
 	// Unlock any maillogs that may have been locked for processing
 	// when Gophish was last shutdown.
 	err = models.UnlockAllMailLogs()

migrations/2018092312000000_0.8.0_result_urls.go

@@ -0,0 +1,85 @@
+package migrations
+
+import (
+	log "github.com/gophish/gophish/logger"
+	"github.com/gophish/gophish/models"
+	"github.com/jinzhu/gorm"
+)
+
+// Migration2018092312000000 backfills models.Result objects with the correctly
+// parsed URLs for use in USB drop campaign
+type Migration2018092312000000 struct{}
+
+func (m Migration2018092312000000) generateURL(campaign *models.Campaign, result *models.Result) error {
+	pctx, err := models.NewPhishingTemplateContext(campaign, result.BaseRecipient, result.RId)
+	if err != nil {
+		return err
+	}
+	result.URL = pctx.URL
+	return nil
+}
+
+func (m Migration2018092312000000) updateResults(db *gorm.DB, campaign models.Campaign) error {
+	tx := db.Begin()
+	defer func() {
+		if r := recover(); r != nil {
+			tx.Rollback()
+		}
+	}()
+	if tx.Error != nil {
+		return tx.Error
+	}
+	// Gather all of the results for this campaign
+	results, err := tx.Table("results").Where("campaign_id=?", campaign.Id).Rows()
+	if err != nil {
+		tx.Rollback()
+		return err
+	}
+	for results.Next() {
+		var result models.Result
+		if err = tx.ScanRows(results, &result); err != nil {
+			tx.Rollback()
+			return err
+		}
+		if err = m.generateURL(&campaign, &result); err != nil {
+			tx.Rollback()
+			return err
+		}
+		log.Infof("Campaign ID: %d Result: %s URL: %s\n", campaign.Id, result.RId, result.URL)
+		if err = tx.Save(&result).Error; err != nil {
+			tx.Rollback()
+			return err
+		}
+	}
+	err = results.Close()
+	if err != nil {
+		log.Error(err)
+	}
+	log.Info("committing")
+	return tx.Commit().Error
+}
+
+// Up backfills previous models.Result objects with the correct parsed URLs
+func (m Migration2018092312000000) Up(db *gorm.DB) error {
+	campaigns := []models.Campaign{}
+	err := db.Table("campaigns").Find(&campaigns).Error
+	if err != nil {
+		log.Error(err)
+		return err
+	}
+	// For each campaign, iterate over the results and parse the correct URL,
+	// storing it back in the database.
+	for _, campaign := range campaigns {
+		log.Infof("Getting results for %d\n", campaign.Id)
+		err = m.updateResults(db, campaign)
+		if err != nil {
+			log.Error(err)
+			return err
+		}
+	}
+	return nil
+}
+
+func (m Migration2018092312000000) Down(db *gorm.DB) error {
+	return nil
+}

migrations/migrate.go

@@ -0,0 +1,99 @@
+package migrations
+
+import (
+	"bitbucket.org/liamstask/goose/lib/goose"
+	"github.com/gophish/gophish/config"
+	log "github.com/gophish/gophish/logger"
+	"github.com/jinzhu/gorm"
+)
+
+// Migration is an interface that defines the needed operations for a custom
+// database migration.
+type Migration interface {
+	Up(db *gorm.DB) error
+	Down(db *gorm.DB) error
+}
+
+// CustomMigrations are the list of migrations we need to run that include
+// custom logic.
+// Any migrations in this list need a corresponding SQL migration in the
+// db/db_*/migrations/ directories. The corresponding SQL migration may include
+// any setup instructions that are then used in these custom migrations.
+var CustomMigrations = map[int64]Migration{
+	2018092312000000: Migration2018092312000000{},
+}
+
+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
+}
+
+// Migrate executes the database migrations, resulting in an up-to-date
+// instance of the database schema.
+func Migrate() error {
+	// Open a database connection for our migrations
+	db, err := gorm.Open(config.Conf.DBName, config.Conf.DBPath)
+	db.LogMode(false)
+	db.SetLogger(log.Logger)
+	db.DB().SetMaxOpenConns(1)
+	if err != nil {
+		log.Error(err)
+		return err
+	}
+	defer db.Close()
+	// 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 {
+		log.Error(err)
+		return err
+	}
+	currentVersion, err := goose.GetDBVersion(migrateConf)
+	if err != nil {
+		log.Error(err)
+		return err
+	}
+	// Collect all the outstanding migrations that need to be executed
+	ms, err := goose.CollectMigrations(migrateConf.MigrationsDir, currentVersion, latest)
+	if err != nil {
+		log.Errorf("Error collecting migrations: %s\n", err)
+		return err
+	}
+	for _, m := range ms {
+		if migration, ok := CustomMigrations[m.Version]; ok {
+			// Run all the migrations up to and including this point
+			log.Infof("Found custom migration %d. Running previous migrations\n", m.Version)
+			err = goose.RunMigrationsOnDb(migrateConf, migrateConf.MigrationsDir, m.Version, db.DB())
+			// After the setup migration runs, we can run our custom logic
+			err = migration.Up(db)
+			if err != nil {
+				log.Errorf("Error applying migration %d: %s\n", m.Version, err)
+				return err
+			}
+		}
+	}
+	// Finally, do one last pass to ensure that all the migrations up to the
+	// latest one are executed
+	err = goose.RunMigrationsOnDb(migrateConf, migrateConf.MigrationsDir, latest, db.DB())
+	if err != nil {
+		log.Error(err)
+		return err
+	}
+	return nil
+}

models/models.go

@@ -5,8 +5,6 @@ import (
 	"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"
@@ -59,38 +57,9 @@ func generateSecureKey() string {
 	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 {
-	// 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 {
-		log.Error(err)
-		return err
-	}
 	// Open our database connection
 	db, err = gorm.Open(config.Conf.DBName, config.Conf.DBPath)
 	db.LogMode(false)
@@ -100,12 +69,6 @@ func Setup() error {
 		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)

models/result.go

@@ -35,6 +35,7 @@ type Result struct {
 	SendDate     time.Time `json:"send_date"`
 	Reported     bool      `json:"reported" sql:"not null"`
 	ModifiedDate time.Time `json:"modified_date"`
+	URL          string    `json:"url"`
 	BaseRecipient
 }