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

pull/2162/head
Pavel Tsakalidis 2021-04-08 13:51:10 +01:00
parent db63ee978d
commit e198ba98af
9 changed files with 143 additions and 13 deletions

View File

@ -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

View File

@ -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

View File

@ -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, "")
}

View File

@ -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
}

View File

@ -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) {
const alphaNum = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
k := make([]byte, 7)
func generateResultId(alphaNum string, r_id_length int64) (string, error) {
if len(alphaNum) == 0 || r_id_length <= 0 {
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
}

View File

@ -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) {

File diff suppressed because one or more lines are too long

View File

@ -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')
})
})

View File

@ -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">