mirror of https://github.com/gophish/gophish
257 lines
6.8 KiB
Go
257 lines
6.8 KiB
Go
package models
|
|
|
|
import (
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/PuerkitoBio/goquery"
|
|
"github.com/gophish/gophish/config"
|
|
"gopkg.in/check.v1"
|
|
)
|
|
|
|
// Hook up gocheck into the "go test" runner.
|
|
func Test(t *testing.T) { check.TestingT(t) }
|
|
|
|
type ModelsSuite struct{}
|
|
|
|
var _ = check.Suite(&ModelsSuite{})
|
|
|
|
func (s *ModelsSuite) SetUpSuite(c *check.C) {
|
|
config.Conf.DBPath = ":memory:"
|
|
config.Conf.MigrationsPath = "../db/migrations/"
|
|
err := Setup()
|
|
if err != nil {
|
|
c.Fatalf("Failed creating database: %v", err)
|
|
}
|
|
}
|
|
|
|
func (s *ModelsSuite) TestGetUser(c *check.C) {
|
|
u, err := GetUser(1)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(u.Username, check.Equals, "admin")
|
|
}
|
|
|
|
func (s *ModelsSuite) TestGeneratedAPIKey(c *check.C) {
|
|
u, err := GetUser(1)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(u.ApiKey, check.Not(check.Equals), "12345678901234567890123456789012")
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostGroup(c *check.C) {
|
|
g := Group{Name: "Test Group"}
|
|
g.Targets = []Target{Target{Email: "test@example.com"}}
|
|
g.UserId = 1
|
|
err := PostGroup(&g)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(g.Name, check.Equals, "Test Group")
|
|
c.Assert(g.Targets[0].Email, check.Equals, "test@example.com")
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostGroupNoName(c *check.C) {
|
|
g := Group{Name: ""}
|
|
g.Targets = []Target{Target{Email: "test@example.com"}}
|
|
g.UserId = 1
|
|
err := PostGroup(&g)
|
|
c.Assert(err, check.Equals, ErrGroupNameNotSpecified)
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostGroupNoTargets(c *check.C) {
|
|
g := Group{Name: "No Target Group"}
|
|
g.Targets = []Target{}
|
|
g.UserId = 1
|
|
err := PostGroup(&g)
|
|
c.Assert(err, check.Equals, ErrNoTargetsSpecified)
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostSMTP(c *check.C) {
|
|
smtp := SMTP{
|
|
Name: "Test SMTP",
|
|
Host: "1.1.1.1:25",
|
|
FromAddress: "Foo Bar <foo@example.com>",
|
|
UserId: 1,
|
|
}
|
|
err = PostSMTP(&smtp)
|
|
c.Assert(err, check.Equals, nil)
|
|
ss, err := GetSMTPs(1)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(len(ss), check.Equals, 1)
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostSMTPNoHost(c *check.C) {
|
|
smtp := SMTP{
|
|
Name: "Test SMTP",
|
|
FromAddress: "Foo Bar <foo@example.com>",
|
|
UserId: 1,
|
|
}
|
|
err = PostSMTP(&smtp)
|
|
c.Assert(err, check.Equals, ErrHostNotSpecified)
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostSMTPNoFrom(c *check.C) {
|
|
smtp := SMTP{
|
|
Name: "Test SMTP",
|
|
UserId: 1,
|
|
Host: "1.1.1.1:25",
|
|
}
|
|
err = PostSMTP(&smtp)
|
|
c.Assert(err, check.Equals, ErrFromAddressNotSpecified)
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostPage(c *check.C) {
|
|
html := `<html>
|
|
<head></head>
|
|
<body><form action="example.com">
|
|
<input name="username"/>
|
|
<input name="password" type="password"/>
|
|
</form></body>
|
|
</html>`
|
|
p := Page{
|
|
Name: "Test Page",
|
|
HTML: html,
|
|
RedirectURL: "http://example.com",
|
|
}
|
|
// Check the capturing credentials and passwords
|
|
p.CaptureCredentials = true
|
|
p.CapturePasswords = true
|
|
err := PostPage(&p)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(p.RedirectURL, check.Equals, "http://example.com")
|
|
d, err := goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
|
|
c.Assert(err, check.Equals, nil)
|
|
forms := d.Find("form")
|
|
forms.Each(func(i int, f *goquery.Selection) {
|
|
// Check the action has been set
|
|
a, _ := f.Attr("action")
|
|
c.Assert(a, check.Equals, "")
|
|
// Check the password still has a name
|
|
_, ok := f.Find("input[type=\"password\"]").Attr("name")
|
|
c.Assert(ok, check.Equals, true)
|
|
// Check the username is still correct
|
|
u, ok := f.Find("input").Attr("name")
|
|
c.Assert(ok, check.Equals, true)
|
|
c.Assert(u, check.Equals, "username")
|
|
})
|
|
// Check what happens when we don't capture passwords
|
|
p.CapturePasswords = false
|
|
p.HTML = html
|
|
p.RedirectURL = ""
|
|
err = PutPage(&p)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(p.RedirectURL, check.Equals, "")
|
|
d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
|
|
c.Assert(err, check.Equals, nil)
|
|
forms = d.Find("form")
|
|
forms.Each(func(i int, f *goquery.Selection) {
|
|
// Check the action has been set
|
|
a, _ := f.Attr("action")
|
|
c.Assert(a, check.Equals, "")
|
|
// Check the password still has a name
|
|
_, ok := f.Find("input[type=\"password\"]").Attr("name")
|
|
c.Assert(ok, check.Equals, false)
|
|
// Check the username is still correct
|
|
u, ok := f.Find("input").Attr("name")
|
|
c.Assert(ok, check.Equals, true)
|
|
c.Assert(u, check.Equals, "username")
|
|
})
|
|
// Finally, check when we don't capture credentials
|
|
p.CaptureCredentials = false
|
|
p.HTML = html
|
|
err = PutPage(&p)
|
|
c.Assert(err, check.Equals, nil)
|
|
d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML))
|
|
c.Assert(err, check.Equals, nil)
|
|
forms = d.Find("form")
|
|
forms.Each(func(i int, f *goquery.Selection) {
|
|
// Check the action has been set
|
|
a, _ := f.Attr("action")
|
|
c.Assert(a, check.Equals, "")
|
|
// Check the password still has a name
|
|
_, ok := f.Find("input[type=\"password\"]").Attr("name")
|
|
c.Assert(ok, check.Equals, false)
|
|
// Check the username is still correct
|
|
_, ok = f.Find("input").Attr("name")
|
|
c.Assert(ok, check.Equals, false)
|
|
})
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostTaskMissingValues(c *check.C) {
|
|
// Missing template_id
|
|
t := Task{
|
|
UserId: 1,
|
|
CampaignId: 1,
|
|
Type: "SEND_EMAIL",
|
|
Metadata: `{
|
|
"smtp_id" : 1
|
|
}`,
|
|
}
|
|
err = PostTask(&t)
|
|
c.Assert(err, check.Equals, ErrTemplateIdNotSpecified)
|
|
// Missing smtp_id
|
|
t.Metadata = `{
|
|
"template_id" : 1
|
|
}`
|
|
err = PostTask(&t)
|
|
c.Assert(err, check.Equals, ErrSMTPIdNotSpecified)
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPostTasks(c *check.C) {
|
|
temp := Template{
|
|
Name: "Test Template",
|
|
Text: "Testing",
|
|
HTML: "Testing",
|
|
UserId: 1,
|
|
}
|
|
err := PostTemplate(&temp)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(temp.Id, check.Equals, int64(1))
|
|
smtp := SMTP{
|
|
Name: "Test SMTP",
|
|
Host: "1.1.1.1:25",
|
|
FromAddress: "Foo Bar <foo@example.com>",
|
|
UserId: 1,
|
|
}
|
|
err = PostSMTP(&smtp)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(smtp.Id, check.Equals, int64(2))
|
|
t := Task{
|
|
UserId: 1,
|
|
CampaignId: 1,
|
|
Type: "SEND_EMAIL",
|
|
Metadata: `{
|
|
"smtp_id" : 2,
|
|
"template_id" : 1
|
|
}`,
|
|
}
|
|
st := Task{
|
|
UserId: 1,
|
|
CampaignId: 1,
|
|
Type: "SEND_EMAIL",
|
|
Metadata: `{
|
|
"smtp_id" : 2,
|
|
"template_id" : 1
|
|
}`,
|
|
}
|
|
err = PostTasks([]*Task{&t, &st})
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(t.Id, check.Equals, int64(1))
|
|
c.Assert(t.NextId, check.Equals, int64(2))
|
|
c.Assert(t.PreviousId, check.Equals, int64(0))
|
|
c.Assert(st.NextId, check.Equals, int64(0))
|
|
c.Assert(st.PreviousId, check.Equals, int64(1))
|
|
c.Assert(st.Id, check.Equals, int64(2))
|
|
// Check retrieving a value from the database
|
|
t, err = GetTask(t.Id, t.UserId)
|
|
c.Assert(err, check.Equals, nil)
|
|
c.Assert(t.NextId, check.Equals, int64(2))
|
|
c.Assert(t.PreviousId, check.Equals, int64(0))
|
|
}
|
|
|
|
func (s *ModelsSuite) TestPutUser(c *check.C) {
|
|
u, err := GetUser(1)
|
|
u.Username = "admin_changed"
|
|
err = PutUser(&u)
|
|
c.Assert(err, check.Equals, nil)
|
|
u, err = GetUser(1)
|
|
c.Assert(u.Username, check.Equals, "admin_changed")
|
|
}
|