Implement functionality to support custom RId - Ability to specify both the character set and length

db/db_mysql/migrations/20210408110421_0.11.1_custom_rid.sql

@@ -0,0 +1,9 @@
+
+-- +goose Up
+-- SQL in section 'Up' is executed when this migration is applied
+ALTER TABLE `campaigns` ADD COLUMN character_set VARCHAR(255);
+ALTER TABLE `campaigns` ADD COLUMN r_id_length INTEGER;
+
+-- +goose Down
+-- SQL section 'Down' is executed when this migration is rolled back
+

db/db_sqlite3/migrations/20210408110430_0.11.1_custom_rid.sql

@@ -0,0 +1,9 @@
+
+-- +goose Up
+-- SQL in section 'Up' is executed when this migration is applied
+ALTER TABLE campaigns ADD COLUMN character_set VARCHAR(255);
+ALTER TABLE campaigns ADD COLUMN r_id_length INTEGER;
+
+-- +goose Down
+-- SQL section 'Down' is executed when this migration is rolled back
+

models/campaign.go

@@ -2,7 +2,9 @@ package models
 
 import (
 	"errors"
+	"math"
 	"net/url"
+	"strings"
 	"time"
 
 	log "github.com/gophish/gophish/logger"
@@ -31,6 +33,9 @@ type Campaign struct {
 	SMTPId        int64     `json:"-"`
 	SMTP          SMTP      `json:"smtp"`
 	URL           string    `json:"url"`
+	CustomRId     bool      `json:"custom_rid" gorm:"-"`
+	CharacterSet  string	`json:"character_set" sql:"not null"`
+	RIdLength     int64		`json:"r_id_length"`
 }
 
 // CampaignResults is a struct representing the results from a campaign
@@ -58,6 +63,8 @@ type CampaignSummary struct {
 	Status        string        `json:"status"`
 	Name          string        `json:"name"`
 	Stats         CampaignStats `json:"stats"`
+	CharacterSet  string		`json:"character_set"`
+	RIdLength     int64		    `json:"r_id_length"`
 }
 
 // CampaignStats is a struct representing the statistics for a single campaign
@@ -126,6 +133,15 @@ var ErrSMTPNotFound = errors.New("Sending profile not found")
 // launch date
 var ErrInvalidSendByDate = errors.New("The launch date must be before the \"send emails by\" date")
 
+// ErrInvalidCharacterSet indicates that the user has entered an invalid character set, ie it's empty
+var ErrInvalidCharacterSet = errors.New("The defined character set is invalid")
+
+// ErrInvalidRIdLength indicates that an invalid RId length has been set (ie zero)
+var ErrInvalidRIdLength = errors.New("The RId length entered is invalid")
+
+// ErrInsufficientCharsetKeyspace indicates that there aren't enough RId combinations to match the total recipients, for the given CharacterSet and RIdLength.
+var ErrInsufficientCharsetKeyspace = errors.New("The specified CharacterSet/Length combination cannot cover the total recipients in this campaign. Consider increasing the Character Set, RId Length, or both");
+
 // RecipientParameter is the URL parameter that points to the result ID for a recipient.
 const RecipientParameter = "rid"
 
@@ -144,6 +160,12 @@ func (c *Campaign) Validate() error {
 		return ErrSMTPNotSpecified
 	case !c.SendByDate.IsZero() && !c.LaunchDate.IsZero() && c.SendByDate.Before(c.LaunchDate):
 		return ErrInvalidSendByDate
+	case c.CustomRId:
+		if c.CharacterSet == "" {
+			return ErrInvalidCharacterSet
+		} else if c.RIdLength <= 0 {
+			return ErrInvalidRIdLength
+		}
 	}
 	return nil
 }
@@ -324,7 +346,7 @@ func GetCampaignSummaries(uid int64) (CampaignSummaries, error) {
 	cs := []CampaignSummary{}
 	// Get the basic campaign information
 	query := db.Table("campaigns").Where("user_id = ?", uid)
-	query = query.Select("id, name, created_date, launch_date, send_by_date, completed_date, status")
+	query = query.Select("id, name, created_date, launch_date, send_by_date, completed_date, status, character_set, r_id_length")
 	err := query.Scan(&cs).Error
 	if err != nil {
 		log.Error(err)
@@ -347,7 +369,7 @@ func GetCampaignSummaries(uid int64) (CampaignSummaries, error) {
 func GetCampaignSummary(id int64, uid int64) (CampaignSummary, error) {
 	cs := CampaignSummary{}
 	query := db.Table("campaigns").Where("user_id = ? AND id = ?", uid, id)
-	query = query.Select("id, name, created_date, launch_date, send_by_date, completed_date, status")
+	query = query.Select("id, name, created_date, launch_date, send_by_date, completed_date, status, character_set, r_id_length")
 	err := query.Scan(&cs).Error
 	if err != nil {
 		log.Error(err)
@@ -487,6 +509,19 @@ func PostCampaign(c *Campaign, uid int64) error {
 		}
 		totalRecipients += len(c.Groups[i].Targets)
 	}
+
+	if c.CustomRId {
+		c.CharacterSet = CleanCharacterSet(c.CharacterSet)
+		// We need to check if the keyspace for the custom character set and rid is enough to match the totalRecipients.
+		// This is on purpose <= so that we don't come too close to the limit.
+		if int(math.Pow(float64(len(c.CharacterSet)), float64(c.RIdLength))) <= totalRecipients {
+			return ErrInsufficientCharsetKeyspace
+		}
+	} else {
+		c.CharacterSet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+		c.RIdLength = 7
+	}
+
 	// Check to make sure the template exists
 	t, err := GetTemplateByName(c.Template.Name, uid)
 	if err == gorm.ErrRecordNotFound {
@@ -564,7 +599,7 @@ func PostCampaign(c *Campaign, uid int64) error {
 				Reported:     false,
 				ModifiedDate: c.CreatedDate,
 			}
-			err = r.GenerateId(tx)
+			err = r.GenerateId(tx, c.CharacterSet, c.RIdLength)
 			if err != nil {
 				log.Error(err)
 				tx.Rollback()
@@ -667,3 +702,16 @@ func CompleteCampaign(id int64, uid int64) error {
 	}
 	return err
 }
+
+func CleanCharacterSet(character_set string) string {
+	keys := make(map[string]bool)
+	chars := strings.Split(character_set, "")
+	unique := []string{}
+	for _, char := range chars {
+		if _, value := keys[char]; !value {
+			keys[char] = true
+			unique = append(unique, char)
+		}
+	}
+	return strings.Join(unique, "")
+}

models/email_request.go

@@ -80,7 +80,7 @@ func (s *EmailRequest) Success() error {
 // PostEmailRequest stores a SendTestEmailRequest in the database.
 func PostEmailRequest(s *EmailRequest) error {
 	// Generate an ID to be used in the underlying Result object
-	rid, err := generateResultId()
+	rid, err := generateResultId("", 0)
 	if err != nil {
 		return err
 	}

models/result.go

@@ -3,6 +3,8 @@ package models
 import (
 	"crypto/rand"
 	"encoding/json"
+	"errors"
+	"math"
 	"math/big"
 	"net"
 	"time"
@@ -170,9 +172,12 @@ func (r *Result) UpdateGeo(addr string) error {
 	return db.Save(r).Error
 }
 
-func generateResultId() (string, error) {
+func generateResultId(alphaNum string, r_id_length int64) (string, error) {
-	const alphaNum = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+	if len(alphaNum) == 0 || r_id_length <= 0 {
-	k := make([]byte, 7)
+		alphaNum = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+		r_id_length = 7
+	}
+	k := make([]byte, r_id_length)
 	for i := range k {
 		idx, err := rand.Int(rand.Reader, big.NewInt(int64(len(alphaNum))))
 		if err != nil {
@@ -185,10 +190,12 @@ func generateResultId() (string, error) {
 
 // GenerateId generates a unique key to represent the result
 // in the database
-func (r *Result) GenerateId(tx *gorm.DB) error {
+func (r *Result) GenerateId(tx *gorm.DB, character_set string, r_id_length int64) error {
 	// Keep trying until we generate a unique key (shouldn't take more than one or two iterations)
+	max_iterations := int(math.Pow(float64(len(character_set)), float64(r_id_length)))
+	iteration_count := 0
 	for {
-		rid, err := generateResultId()
+		rid, err := generateResultId(character_set, r_id_length)
 		if err != nil {
 			return err
 		}
@@ -197,6 +204,11 @@ func (r *Result) GenerateId(tx *gorm.DB) error {
 		if err == gorm.ErrRecordNotFound {
 			break
 		}
+
+		iteration_count++
+		if iteration_count >= max_iterations {
+			return errors.New("Too many iterations - Consider increasing the character set and/or RId length")
+		}
 	}
 	return nil
 }

models/result_test.go

@@ -10,10 +10,20 @@ import (
 
 func (s *ModelsSuite) TestGenerateResultId(c *check.C) {
 	r := Result{}
-	r.GenerateId(db)
+	r.GenerateId(db, "", 0)
 	match, err := regexp.Match("[a-zA-Z0-9]{7}", []byte(r.RId))
 	c.Assert(err, check.Equals, nil)
 	c.Assert(match, check.Equals, true)
+
+	r.GenerateId(db, "0123456789", 4)
+	match, err = regexp.Match("[0-9]{4}", []byte(r.RId))
+	c.Assert(err, check.Equals, nil)
+	c.Assert(match, check.Equals, true)
+
+	r.GenerateId(db, "abcdefghijklmnopqrstuvwxyz", 3)
+	match, err = regexp.Match("[a-z]{3}", []byte(r.RId))
+	c.Assert(err, check.Equals, nil)
+	c.Assert(match, check.Equals, true)
 }
 
 func (s *ModelsSuite) TestFormatAddress(c *check.C) {

static/js/src/app/campaigns.js

@@ -37,6 +37,11 @@ function launch() {
                 if (send_by_date != "") {
                     send_by_date = moment(send_by_date, "MMMM Do YYYY, h:mm a").utc().format()
                 }
+
+                var custom_rid = $('#customize_rid').is(':checked');
+                var character_set = $('#character_set').val().trim();
+                var r_id_length = parseInt($('#r_id_length').val().trim()) || 7;
+
                 campaign = {
                     name: $("#name").val(),
                     template: {
@@ -52,6 +57,9 @@ function launch() {
                     launch_date: moment($("#launch_date").val(), "MMMM Do YYYY, h:mm a").utc().format(),
                     send_by_date: send_by_date || null,
                     groups: groups,
+                    custom_rid: custom_rid,
+                    character_set: character_set,
+                    r_id_length: r_id_length
                 }
                 // Submit the campaign
                 api.campaigns.post(campaign)
@@ -378,10 +386,18 @@ $(document).ready(function () {
                         var quickStats = launchDate + "<br><br>" + "Number of recipients: " + campaign.stats.total + "<br><br>" + "Emails opened: " + campaign.stats.opened + "<br><br>" + "Emails clicked: " + campaign.stats.clicked + "<br><br>" + "Submitted Credentials: " + campaign.stats.submitted_data + "<br><br>" + "Errors : " + campaign.stats.error + "<br><br>" + "Reported : " + campaign.stats.email_reported
                     }
 
+                    var charSetText = "";
+                    if (campaign.character_set == "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789") {
+                        charSetText = "default";
+                    } else {
+                        charSetText = escapeHtml(campaign.character_set) + " (Length: " + campaign.r_id_length + ")";
+                    }
+
                     var row = [
                         escapeHtml(campaign.name),
                         moment(campaign.created_date).format('MMMM Do YYYY, h:mm:ss a'),
                         "<span class=\"label " + label + "\" data-toggle=\"tooltip\" data-placement=\"right\" data-html=\"true\" title=\"" + quickStats + "\">" + campaign.status + "</span>",
+                        "<span>" + charSetText + "</span>",
                         "<div class='pull-right'><a class='btn btn-primary' href='/campaigns/" + campaign.id + "' data-toggle='tooltip' data-placement='left' title='View Results'>\
                     <i class='fa fa-bar-chart'></i>\
                     </a>\
@@ -424,4 +440,10 @@ $(document).ready(function () {
             return 0;
         });
     })
+
+    $('#customize_rid').change(function() {
+        $(this).is(':checked')
+            ? $('#fields-customize_rid').removeClass('hidden')
+            : $('#fields-customize_rid').addClass('hidden')
+    })
 })

templates/campaigns.html

@@ -37,8 +37,9 @@
                     <thead>
                         <tr>
                             <th class="col-md-3">Name</th>
-                            <th class="col-md-4">Created Date</th>
+                            <th class="col-md-2">Created Date</th>
                             <th class="col-md-2">Status</th>
+                            <th class="col-md-2">Character Set</th>
                             <th class="col-md-3 no-sort"></th>
                         </tr>
                     </thead>
@@ -60,8 +61,9 @@
                     <thead>
                         <tr>
                             <th class="col-md-3">Name</th>
-                            <th class="col-md-4">Created Date</th>
+                            <th class="col-md-2">Created Date</th>
                             <th class="col-md-2">Status</th>
+                            <th class="col-md-2">Character Set</th>
                             <th class="col-md-3 no-sort"></th>
                         </tr>
                     </thead>
@@ -124,6 +126,24 @@
                     </div>
                     <label class="control-label" for="users">Groups:</label>
                     <select class="form-control" id="users" multiple="multiple"></select>
+
+                    <div class="row">
+                        <div class="col-md-12">
+                            <input type="checkbox" id="customize_rid">
+                            <label for="customize_rid">Customize RId</label>
+                        </div>
+                    </div>
+
+                    <div class="row hidden" id="fields-customize_rid">
+                        <div class="col-md-8">
+                            <label for="character_set">RId Character Set</label>
+                            <input type="text" class="form-control" id="character_set" placeholder="Character set for RId variable" value="abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789">
+                        </div>
+                        <div class="col-md-4">
+                            <label for="r_id_length">RId Length</label>
+                            <input type="number" class="form-control" id="r_id_length" placeholder="Length for RId variable" value="7">
+                        </div>
+                    </div>
                 </div>
             </div>
             <div class="modal-footer">