mirror of https://github.com/gophish/gophish
86 lines
2.1 KiB
Go
86 lines
2.1 KiB
Go
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
|
|
}
|