github.com/jonathaningram/gophish@v0.3.1-0.20170829042651-ac3fe6aeae6c/models/models_test.go (about) 1 package models 2 3 import ( 4 "net/mail" 5 "regexp" 6 "strings" 7 "testing" 8 9 "github.com/PuerkitoBio/goquery" 10 "github.com/gophish/gophish/config" 11 "github.com/jinzhu/gorm" 12 "gopkg.in/check.v1" 13 ) 14 15 // Hook up gocheck into the "go test" runner. 16 func Test(t *testing.T) { check.TestingT(t) } 17 18 type ModelsSuite struct{} 19 20 var _ = check.Suite(&ModelsSuite{}) 21 22 func (s *ModelsSuite) SetUpSuite(c *check.C) { 23 config.Conf.DBName = "sqlite3" 24 config.Conf.DBPath = ":memory:" 25 config.Conf.MigrationsPath = "../db/db_sqlite3/migrations/" 26 err := Setup() 27 if err != nil { 28 c.Fatalf("Failed creating database: %v", err) 29 } 30 } 31 32 func (s *ModelsSuite) TearDownTest(c *check.C) { 33 // Clear database tables between each test. If new tables are 34 // used in this test suite they will need to be cleaned up here. 35 db.Delete(Group{}) 36 db.Delete(Target{}) 37 db.Delete(GroupTarget{}) 38 db.Delete(SMTP{}) 39 db.Delete(Page{}) 40 41 // Reset users table to default state. 42 db.Not("id", 1).Delete(User{}) 43 db.Model(User{}).Update("username", "admin") 44 } 45 46 func (s *ModelsSuite) TestGetUser(c *check.C) { 47 u, err := GetUser(1) 48 c.Assert(err, check.Equals, nil) 49 c.Assert(u.Username, check.Equals, "admin") 50 } 51 52 func (s *ModelsSuite) TestPutUser(c *check.C) { 53 u, err := GetUser(1) 54 u.Username = "admin_changed" 55 err = PutUser(&u) 56 c.Assert(err, check.Equals, nil) 57 u, err = GetUser(1) 58 c.Assert(u.Username, check.Equals, "admin_changed") 59 } 60 61 func (s *ModelsSuite) TestGeneratedAPIKey(c *check.C) { 62 u, err := GetUser(1) 63 c.Assert(err, check.Equals, nil) 64 c.Assert(u.ApiKey, check.Not(check.Equals), "12345678901234567890123456789012") 65 } 66 67 func (s *ModelsSuite) TestPostGroup(c *check.C) { 68 g := Group{Name: "Test Group"} 69 g.Targets = []Target{Target{Email: "test@example.com"}} 70 g.UserId = 1 71 err := PostGroup(&g) 72 c.Assert(err, check.Equals, nil) 73 c.Assert(g.Name, check.Equals, "Test Group") 74 c.Assert(g.Targets[0].Email, check.Equals, "test@example.com") 75 } 76 77 func (s *ModelsSuite) TestPostGroupNoName(c *check.C) { 78 g := Group{Name: ""} 79 g.Targets = []Target{Target{Email: "test@example.com"}} 80 g.UserId = 1 81 err := PostGroup(&g) 82 c.Assert(err, check.Equals, ErrGroupNameNotSpecified) 83 } 84 85 func (s *ModelsSuite) TestPostGroupNoTargets(c *check.C) { 86 g := Group{Name: "No Target Group"} 87 g.Targets = []Target{} 88 g.UserId = 1 89 err := PostGroup(&g) 90 c.Assert(err, check.Equals, ErrNoTargetsSpecified) 91 } 92 93 func (s *ModelsSuite) TestGetGroups(c *check.C) { 94 // Add groups. 95 PostGroup(&Group{ 96 Name: "Test Group 1", 97 Targets: []Target{Target{Email: "test1@example.com"}}, 98 UserId: 1, 99 }) 100 PostGroup(&Group{ 101 Name: "Test Group 2", 102 Targets: []Target{Target{Email: "test2@example.com"}}, 103 UserId: 1, 104 }) 105 106 // Get groups and test result. 107 groups, err := GetGroups(1) 108 c.Assert(err, check.Equals, nil) 109 c.Assert(len(groups), check.Equals, 2) 110 c.Assert(len(groups[0].Targets), check.Equals, 1) 111 c.Assert(len(groups[1].Targets), check.Equals, 1) 112 c.Assert(groups[0].Name, check.Equals, "Test Group 1") 113 c.Assert(groups[1].Name, check.Equals, "Test Group 2") 114 c.Assert(groups[0].Targets[0].Email, check.Equals, "test1@example.com") 115 c.Assert(groups[1].Targets[0].Email, check.Equals, "test2@example.com") 116 } 117 118 func (s *ModelsSuite) TestGetGroupsNoGroups(c *check.C) { 119 groups, err := GetGroups(1) 120 c.Assert(err, check.Equals, nil) 121 c.Assert(len(groups), check.Equals, 0) 122 } 123 124 func (s *ModelsSuite) TestGetGroup(c *check.C) { 125 // Add group. 126 PostGroup(&Group{ 127 Name: "Test Group", 128 Targets: []Target{Target{Email: "test@example.com"}}, 129 UserId: 1, 130 }) 131 132 // Get group and test result. 133 group, err := GetGroup(1, 1) 134 c.Assert(err, check.Equals, nil) 135 c.Assert(len(group.Targets), check.Equals, 1) 136 c.Assert(group.Name, check.Equals, "Test Group") 137 c.Assert(group.Targets[0].Email, check.Equals, "test@example.com") 138 } 139 140 func (s *ModelsSuite) TestGetGroupNoGroups(c *check.C) { 141 _, err := GetGroup(1, 1) 142 c.Assert(err, check.Equals, gorm.ErrRecordNotFound) 143 } 144 145 func (s *ModelsSuite) TestGetGroupByName(c *check.C) { 146 // Add group. 147 PostGroup(&Group{ 148 Name: "Test Group", 149 Targets: []Target{Target{Email: "test@example.com"}}, 150 UserId: 1, 151 }) 152 153 // Get group and test result. 154 group, err := GetGroupByName("Test Group", 1) 155 c.Assert(err, check.Equals, nil) 156 c.Assert(len(group.Targets), check.Equals, 1) 157 c.Assert(group.Name, check.Equals, "Test Group") 158 c.Assert(group.Targets[0].Email, check.Equals, "test@example.com") 159 } 160 161 func (s *ModelsSuite) TestGetGroupByNameNoGroups(c *check.C) { 162 _, err := GetGroupByName("Test Group", 1) 163 c.Assert(err, check.Equals, gorm.ErrRecordNotFound) 164 } 165 166 func (s *ModelsSuite) TestPutGroup(c *check.C) { 167 // Add test group. 168 group := Group{Name: "Test Group"} 169 group.Targets = []Target{ 170 Target{Email: "test1@example.com", FirstName: "First", LastName: "Example"}, 171 Target{Email: "test2@example.com", FirstName: "Second", LastName: "Example"}, 172 } 173 group.UserId = 1 174 PostGroup(&group) 175 176 // Update one of group's targets. 177 group.Targets[0].FirstName = "Updated" 178 err := PutGroup(&group) 179 c.Assert(err, check.Equals, nil) 180 181 // Verify updated target information. 182 targets, _ := GetTargets(group.Id) 183 c.Assert(targets[0].Email, check.Equals, "test1@example.com") 184 c.Assert(targets[0].FirstName, check.Equals, "Updated") 185 c.Assert(targets[0].LastName, check.Equals, "Example") 186 c.Assert(targets[1].Email, check.Equals, "test2@example.com") 187 c.Assert(targets[1].FirstName, check.Equals, "Second") 188 c.Assert(targets[1].LastName, check.Equals, "Example") 189 } 190 191 func (s *ModelsSuite) TestPutGroupEmptyAttribute(c *check.C) { 192 // Add test group. 193 group := Group{Name: "Test Group"} 194 group.Targets = []Target{ 195 Target{Email: "test1@example.com", FirstName: "First", LastName: "Example"}, 196 Target{Email: "test2@example.com", FirstName: "Second", LastName: "Example"}, 197 } 198 group.UserId = 1 199 PostGroup(&group) 200 201 // Update one of group's targets. 202 group.Targets[0].FirstName = "" 203 err := PutGroup(&group) 204 c.Assert(err, check.Equals, nil) 205 206 // Verify updated empty attribute was saved. 207 targets, _ := GetTargets(group.Id) 208 c.Assert(targets[0].Email, check.Equals, "test1@example.com") 209 c.Assert(targets[0].FirstName, check.Equals, "") 210 c.Assert(targets[0].LastName, check.Equals, "Example") 211 c.Assert(targets[1].Email, check.Equals, "test2@example.com") 212 c.Assert(targets[1].FirstName, check.Equals, "Second") 213 c.Assert(targets[1].LastName, check.Equals, "Example") 214 } 215 216 func (s *ModelsSuite) TestPostSMTP(c *check.C) { 217 smtp := SMTP{ 218 Name: "Test SMTP", 219 Host: "1.1.1.1:25", 220 FromAddress: "Foo Bar <foo@example.com>", 221 UserId: 1, 222 } 223 err = PostSMTP(&smtp) 224 c.Assert(err, check.Equals, nil) 225 ss, err := GetSMTPs(1) 226 c.Assert(err, check.Equals, nil) 227 c.Assert(len(ss), check.Equals, 1) 228 } 229 230 func (s *ModelsSuite) TestPostSMTPNoHost(c *check.C) { 231 smtp := SMTP{ 232 Name: "Test SMTP", 233 FromAddress: "Foo Bar <foo@example.com>", 234 UserId: 1, 235 } 236 err = PostSMTP(&smtp) 237 c.Assert(err, check.Equals, ErrHostNotSpecified) 238 } 239 240 func (s *ModelsSuite) TestPostSMTPNoFrom(c *check.C) { 241 smtp := SMTP{ 242 Name: "Test SMTP", 243 UserId: 1, 244 Host: "1.1.1.1:25", 245 } 246 err = PostSMTP(&smtp) 247 c.Assert(err, check.Equals, ErrFromAddressNotSpecified) 248 } 249 250 func (s *ModelsSuite) TestPostSMTPValidHeader(c *check.C) { 251 smtp := SMTP{ 252 Name: "Test SMTP", 253 Host: "1.1.1.1:25", 254 FromAddress: "Foo Bar <foo@example.com>", 255 UserId: 1, 256 Headers: []Header{ 257 Header{Key: "Reply-To", Value: "test@example.com"}, 258 Header{Key: "X-Mailer", Value: "gophish"}, 259 }, 260 } 261 err = PostSMTP(&smtp) 262 c.Assert(err, check.Equals, nil) 263 ss, err := GetSMTPs(1) 264 c.Assert(err, check.Equals, nil) 265 c.Assert(len(ss), check.Equals, 1) 266 } 267 268 func (s *ModelsSuite) TestPostPage(c *check.C) { 269 html := `<html> 270 <head></head> 271 <body><form action="example.com"> 272 <input name="username"/> 273 <input name="password" type="password"/> 274 </form></body> 275 </html>` 276 p := Page{ 277 Name: "Test Page", 278 HTML: html, 279 RedirectURL: "http://example.com", 280 } 281 // Check the capturing credentials and passwords 282 p.CaptureCredentials = true 283 p.CapturePasswords = true 284 err := PostPage(&p) 285 c.Assert(err, check.Equals, nil) 286 c.Assert(p.RedirectURL, check.Equals, "http://example.com") 287 d, err := goquery.NewDocumentFromReader(strings.NewReader(p.HTML)) 288 c.Assert(err, check.Equals, nil) 289 forms := d.Find("form") 290 forms.Each(func(i int, f *goquery.Selection) { 291 // Check the action has been set 292 a, _ := f.Attr("action") 293 c.Assert(a, check.Equals, "") 294 // Check the password still has a name 295 _, ok := f.Find("input[type=\"password\"]").Attr("name") 296 c.Assert(ok, check.Equals, true) 297 // Check the username is still correct 298 u, ok := f.Find("input").Attr("name") 299 c.Assert(ok, check.Equals, true) 300 c.Assert(u, check.Equals, "username") 301 }) 302 // Check what happens when we don't capture passwords 303 p.CapturePasswords = false 304 p.HTML = html 305 p.RedirectURL = "" 306 err = PutPage(&p) 307 c.Assert(err, check.Equals, nil) 308 c.Assert(p.RedirectURL, check.Equals, "") 309 d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML)) 310 c.Assert(err, check.Equals, nil) 311 forms = d.Find("form") 312 forms.Each(func(i int, f *goquery.Selection) { 313 // Check the action has been set 314 a, _ := f.Attr("action") 315 c.Assert(a, check.Equals, "") 316 // Check the password still has a name 317 _, ok := f.Find("input[type=\"password\"]").Attr("name") 318 c.Assert(ok, check.Equals, false) 319 // Check the username is still correct 320 u, ok := f.Find("input").Attr("name") 321 c.Assert(ok, check.Equals, true) 322 c.Assert(u, check.Equals, "username") 323 }) 324 // Finally, check when we don't capture credentials 325 p.CaptureCredentials = false 326 p.HTML = html 327 err = PutPage(&p) 328 c.Assert(err, check.Equals, nil) 329 d, err = goquery.NewDocumentFromReader(strings.NewReader(p.HTML)) 330 c.Assert(err, check.Equals, nil) 331 forms = d.Find("form") 332 forms.Each(func(i int, f *goquery.Selection) { 333 // Check the action has been set 334 a, _ := f.Attr("action") 335 c.Assert(a, check.Equals, "") 336 // Check the password still has a name 337 _, ok := f.Find("input[type=\"password\"]").Attr("name") 338 c.Assert(ok, check.Equals, false) 339 // Check the username is still correct 340 _, ok = f.Find("input").Attr("name") 341 c.Assert(ok, check.Equals, false) 342 }) 343 } 344 345 func (s *ModelsSuite) TestGenerateResultId(c *check.C) { 346 r := Result{} 347 r.GenerateId() 348 match, err := regexp.Match("[a-zA-Z0-9]{7}", []byte(r.RId)) 349 c.Assert(err, check.Equals, nil) 350 c.Assert(match, check.Equals, true) 351 } 352 353 func (s *ModelsSuite) TestFormatAddress(c *check.C) { 354 r := Result{ 355 FirstName: "John", 356 LastName: "Doe", 357 Email: "johndoe@example.com", 358 } 359 expected := &mail.Address{ 360 Name: "John Doe", 361 Address: "johndoe@example.com", 362 } 363 c.Assert(r.FormatAddress(), check.Equals, expected.String()) 364 365 r = Result{ 366 Email: "johndoe@example.com", 367 } 368 c.Assert(r.FormatAddress(), check.Equals, r.Email) 369 }