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 ", 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 ", 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 := ` ` 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 ", 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") }