github.com/amar224/phishing-tool@v0.9.0/models/models_test.go (about)

     1  package models
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/gophish/gophish/config"
     7  	"gopkg.in/check.v1"
     8  )
     9  
    10  // Hook up gocheck into the "go test" runner.
    11  func Test(t *testing.T) { check.TestingT(t) }
    12  
    13  type ModelsSuite struct {
    14  	config *config.Config
    15  }
    16  
    17  var _ = check.Suite(&ModelsSuite{})
    18  
    19  func (s *ModelsSuite) SetUpSuite(c *check.C) {
    20  	conf := &config.Config{
    21  		DBName:         "sqlite3",
    22  		DBPath:         ":memory:",
    23  		MigrationsPath: "../db/db_sqlite3/migrations/",
    24  	}
    25  	s.config = conf
    26  	err := Setup(conf)
    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  	db.Delete(Result{})
    41  	db.Delete(MailLog{})
    42  	db.Delete(Campaign{})
    43  
    44  	// Reset users table to default state.
    45  	db.Not("id", 1).Delete(User{})
    46  	db.Model(User{}).Update("username", "admin")
    47  }
    48  
    49  func (s *ModelsSuite) createCampaignDependencies(ch *check.C, optional ...string) Campaign {
    50  	// we use the optional parameter to pass an alternative subject
    51  	group := Group{Name: "Test Group"}
    52  	group.Targets = []Target{
    53  		Target{BaseRecipient: BaseRecipient{Email: "test1@example.com", FirstName: "First", LastName: "Example"}},
    54  		Target{BaseRecipient: BaseRecipient{Email: "test2@example.com", FirstName: "Second", LastName: "Example"}},
    55  		Target{BaseRecipient: BaseRecipient{Email: "test3@example.com", FirstName: "Second", LastName: "Example"}},
    56  		Target{BaseRecipient: BaseRecipient{Email: "test4@example.com", FirstName: "Second", LastName: "Example"}},
    57  	}
    58  	group.UserId = 1
    59  	ch.Assert(PostGroup(&group), check.Equals, nil)
    60  
    61  	// Add a template
    62  	t := Template{Name: "Test Template"}
    63  	if len(optional) > 0 {
    64  		t.Subject = optional[0]
    65  	} else {
    66  		t.Subject = "{{.RId}} - Subject"
    67  	}
    68  	t.Text = "{{.RId}} - Text"
    69  	t.HTML = "{{.RId}} - HTML"
    70  	t.UserId = 1
    71  	ch.Assert(PostTemplate(&t), check.Equals, nil)
    72  
    73  	// Add a landing page
    74  	p := Page{Name: "Test Page"}
    75  	p.HTML = "<html>Test</html>"
    76  	p.UserId = 1
    77  	ch.Assert(PostPage(&p), check.Equals, nil)
    78  
    79  	// Add a sending profile
    80  	smtp := SMTP{Name: "Test Page"}
    81  	smtp.UserId = 1
    82  	smtp.Host = "example.com"
    83  	smtp.FromAddress = "test@test.com"
    84  	ch.Assert(PostSMTP(&smtp), check.Equals, nil)
    85  
    86  	c := Campaign{Name: "Test campaign"}
    87  	c.UserId = 1
    88  	c.Template = t
    89  	c.Page = p
    90  	c.SMTP = smtp
    91  	c.Groups = []Group{group}
    92  	return c
    93  }
    94  
    95  func (s *ModelsSuite) createCampaign(ch *check.C) Campaign {
    96  	c := s.createCampaignDependencies(ch)
    97  	// Setup and "launch" our campaign
    98  	ch.Assert(PostCampaign(&c, c.UserId), check.Equals, nil)
    99  
   100  	// For comparing the dates, we need to fetch the campaign again. This is
   101  	// to solve an issue where the campaign object right now has time down to
   102  	// the microsecond, while in MySQL it's rounded down to the second.
   103  	c, _ = GetCampaign(c.Id, c.UserId)
   104  	return c
   105  }
   106  
   107  func setupBenchmark(b *testing.B) {
   108  	conf := &config.Config{
   109  		DBName:         "sqlite3",
   110  		DBPath:         ":memory:",
   111  		MigrationsPath: "../db/db_sqlite3/migrations/",
   112  	}
   113  	err := Setup(conf)
   114  	if err != nil {
   115  		b.Fatalf("Failed creating database: %v", err)
   116  	}
   117  }
   118  
   119  func tearDownBenchmark(b *testing.B) {
   120  	err := db.Close()
   121  	if err != nil {
   122  		b.Fatalf("error closing database: %v", err)
   123  	}
   124  }
   125  
   126  func resetBenchmark(b *testing.B) {
   127  	db.Delete(Group{})
   128  	db.Delete(Target{})
   129  	db.Delete(GroupTarget{})
   130  	db.Delete(SMTP{})
   131  	db.Delete(Page{})
   132  	db.Delete(Result{})
   133  	db.Delete(MailLog{})
   134  	db.Delete(Campaign{})
   135  
   136  	// Reset users table to default state.
   137  	db.Not("id", 1).Delete(User{})
   138  	db.Model(User{}).Update("username", "admin")
   139  }