From 1ff6247199ef7d5c1525f0b0317906cd1e599b07 Mon Sep 17 00:00:00 2001 From: Jordan Wright Date: Fri, 5 Oct 2018 18:03:06 -0500 Subject: [PATCH] WIP - Refactoring migrations to support custom logic. --- .../2018092312000000_0.8.0_result_urls.sql | 6 ++ gophish.go | 12 ++- .../2018092312000000_0.8.0_result_urls.go | 85 ++++++++++++++++ migrations/migrate.go | 99 +++++++++++++++++++ models/models.go | 37 ------- models/result.go | 1 + 6 files changed, 200 insertions(+), 40 deletions(-) create mode 100644 db/db_sqlite3/migrations/2018092312000000_0.8.0_result_urls.sql create mode 100644 migrations/2018092312000000_0.8.0_result_urls.go create mode 100644 migrations/migrate.go diff --git a/db/db_sqlite3/migrations/2018092312000000_0.8.0_result_urls.sql b/db/db_sqlite3/migrations/2018092312000000_0.8.0_result_urls.sql new file mode 100644 index 00000000..652fbf58 --- /dev/null +++ b/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. diff --git a/gophish.go b/gophish.go index 2df7933a..091fb48c 100644 --- a/gophish.go +++ b/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 - if !*disableMailer { - go mailer.Mailer.Start(ctx) + // Run the database migrations to ensure the schema is up-to-date + err = migrations.Migrate() + 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() diff --git a/migrations/2018092312000000_0.8.0_result_urls.go b/migrations/2018092312000000_0.8.0_result_urls.go new file mode 100644 index 00000000..a85fed9d --- /dev/null +++ b/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 +} diff --git a/migrations/migrate.go b/migrations/migrate.go new file mode 100644 index 00000000..ca17d66f --- /dev/null +++ b/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 +} diff --git a/models/models.go b/models/models.go index eb7bfe2a..08efda1a 100644 --- a/models/models.go +++ b/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) diff --git a/models/result.go b/models/result.go index 988ddffd..dc5db1b0 100644 --- a/models/result.go +++ b/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 }