diff --git a/models/template_context.go b/models/template_context.go index e7b0252e..37c59f64 100644 --- a/models/template_context.go +++ b/models/template_context.go @@ -23,6 +23,7 @@ type PhishingTemplateContext struct { Tracker string TrackingURL string RId string + BaseURL string BaseRecipient } @@ -37,22 +38,29 @@ func NewPhishingTemplateContext(ctx TemplateContext, r BaseRecipient, rid string if fn == "" { fn = f.Address } - baseURL, err := ExecuteTemplate(ctx.getBaseURL(), r) + templateURL, err := ExecuteTemplate(ctx.getBaseURL(), r) if err != nil { return PhishingTemplateContext{}, err } - phishURL, _ := url.Parse(baseURL) + // For the base URL, we'll reset the the path and the query + // This will create a URL in the form of http://example.com + baseURL, _ := url.Parse(templateURL) + baseURL.Path = "" + baseURL.RawQuery = "" + + phishURL, _ := url.Parse(templateURL) q := phishURL.Query() q.Set(RecipientParameter, rid) phishURL.RawQuery = q.Encode() - trackingURL, _ := url.Parse(baseURL) + trackingURL, _ := url.Parse(templateURL) trackingURL.Path = path.Join(trackingURL.Path, "/track") trackingURL.RawQuery = q.Encode() return PhishingTemplateContext{ BaseRecipient: r, + BaseURL: baseURL.String(), URL: phishURL.String(), TrackingURL: trackingURL.String(), Tracker: "", diff --git a/models/template_context_test.go b/models/template_context_test.go new file mode 100644 index 00000000..0b6ac4e1 --- /dev/null +++ b/models/template_context_test.go @@ -0,0 +1,47 @@ +package models + +import ( + "fmt" + + check "gopkg.in/check.v1" +) + +type mockTemplateContext struct { + URL string + FromAddress string +} + +func (m mockTemplateContext) getFromAddress() string { + return m.FromAddress +} + +func (m mockTemplateContext) getBaseURL() string { + return m.URL +} + +func (s *ModelsSuite) TestNewTemplateContext(c *check.C) { + r := Result{ + BaseRecipient: BaseRecipient{ + FirstName: "Foo", + LastName: "Bar", + Email: "foo@bar.com", + }, + RId: "1234567", + } + ctx := mockTemplateContext{ + URL: "http://example.com", + FromAddress: "From Address ", + } + expected := PhishingTemplateContext{ + URL: fmt.Sprintf("%s?rid=%s", ctx.URL, r.RId), + BaseURL: ctx.URL, + BaseRecipient: r.BaseRecipient, + TrackingURL: fmt.Sprintf("%s/track?rid=%s", ctx.URL, r.RId), + From: "From Address", + RId: r.RId, + } + expected.Tracker = "" + got, err := NewPhishingTemplateContext(ctx, r.BaseRecipient, r.RId) + c.Assert(err, check.Equals, nil) + c.Assert(got, check.DeepEquals, expected) +}