diff --git a/config/config.go b/config/config.go index 3884c9f7..a92e3a41 100644 --- a/config/config.go +++ b/config/config.go @@ -4,7 +4,6 @@ import ( "encoding/json" "fmt" "io/ioutil" - "os" ) type SMTPServer struct { @@ -27,7 +26,6 @@ func init() { config_file, err := ioutil.ReadFile("./config.json") if err != nil { fmt.Printf("File error: %v\n", err) - os.Exit(1) } json.Unmarshal(config_file, &Conf) } diff --git a/gophish.go b/gophish.go index 9433e5ef..acc45e62 100644 --- a/gophish.go +++ b/gophish.go @@ -38,7 +38,7 @@ import ( func main() { //Setup the global variables and settings err := models.Setup() - defer models.DB.Close() + //defer models.db.DB().Close() if err != nil { fmt.Println(err) } diff --git a/gophish_test.go b/gophish_test.go deleted file mode 100644 index 97510b9a..00000000 --- a/gophish_test.go +++ /dev/null @@ -1,14 +0,0 @@ -package main - -import ( - "testing" - - "github.com/jordan-wright/gophish/models" -) - -func TestDBSetup(t *testing.T) { - err := models.Setup() - if err != nil { - t.Fatalf("Failed creating database: %v", err) - } -} diff --git a/models/campaign.go b/models/campaign.go index e3846b90..46d91887 100644 --- a/models/campaign.go +++ b/models/campaign.go @@ -33,7 +33,7 @@ func GetCampaigns(uid int64) ([]Campaign, error) { if err != nil { fmt.Println(err) } - */for i, _ := range cs { + for i, _ := range cs { err := db.Model(&cs[i]).Related(&cs[i].Results).Error if err != nil { fmt.Println(err) @@ -47,11 +47,11 @@ func GetCampaigns(uid int64) ([]Campaign, error) { func GetCampaign(id int64, uid int64) (Campaign, error) { c := Campaign{} err := db.Where("id = ?", id).Where("user_id = ?", uid).Find(&c).Error - */if err != nil { + if err != nil { return c, err } err = db.Model(&c).Related(&c.Results).Error - */return c, err + return c, err } // PostCampaign inserts a campaign and all associated records into the database. diff --git a/models/models.go b/models/models.go index 24f98668..ecf572a3 100644 --- a/models/models.go +++ b/models/models.go @@ -1,7 +1,6 @@ package models import ( - "database/sql" "errors" "log" "os" @@ -14,7 +13,6 @@ import ( var Conn *gorp.DbMap var db gorm.DB -var DB *sql.DB var err error var ErrUsernameTaken = errors.New("username already taken") var Logger = log.New(os.Stdout, " ", log.Ldate|log.Ltime|log.Lshortfile) @@ -29,6 +27,7 @@ type Flash struct { // It also populates the Gophish Config object func Setup() error { db, err = gorm.Open("sqlite3", config.Conf.DBPath) + db.LogMode(true) db.SetLogger(Logger) if err != nil { Logger.Println(err) diff --git a/models/models_test.go b/models/models_test.go new file mode 100644 index 00000000..6c2d6733 --- /dev/null +++ b/models/models_test.go @@ -0,0 +1,38 @@ +package models + +import ( + "os" + "testing" + + "github.com/jordan-wright/gophish/config" + "launchpad.net/gocheck" +) + +// Hook up gocheck into the "go test" runner. +func Test(t *testing.T) { gocheck.TestingT(t) } + +type ModelsSuite struct{} + +var _ = gocheck.Suite(&ModelsSuite{}) + +func (s *ModelsSuite) SetUpSuite(c *gocheck.C) { + config.Conf.DBPath = "../gophish_test.db" + err := Setup() + if err != nil { + c.Fatalf("Failed creating database: %v", err) + } +} + +func (s *ModelsSuite) TestGetUser(c *gocheck.C) { + u, err := GetUser(1) + c.Assert(err, gocheck.Equals, nil) + c.Assert(u.Username, gocheck.Equals, "admin") +} + +func (s *ModelsSuite) TearDownSuite(c *gocheck.C) { + db.DB().Close() + err := os.Remove(config.Conf.DBPath) + if err != nil { + c.Fatalf("Failed deleting test database: %v", err) + } +}