Changed Setup() to only return an error - I will need to find a way to handle the database connection either globally, or per package as needed. Not sure yet.

pull/24/head
Jordan 2014-01-09 01:08:32 -06:00
parent 6cf9ae4bb8
commit d4dbf30499
2 changed files with 8 additions and 8 deletions

View File

@ -38,7 +38,7 @@ var setupFlag = flag.Bool("setup", false, "Starts the initial setup process for
func main() { func main() {
//Setup the global variables and settings //Setup the global variables and settings
_, err = Setup() err = Setup()
if err != nil { if err != nil {
fmt.Println(err) fmt.Println(err)
} }

View File

@ -16,36 +16,36 @@ var Db sql.DB
//init registers the necessary models to be saved in the session later //init registers the necessary models to be saved in the session later
func init() { func init() {
gob.Register(&models.User{}) gob.Register(&models.User{})
Setup()
} }
// Setup creates and returns the database needed by Gophish. // Setup creates and returns the database needed by Gophish.
// It also populates the Gophish Config object // It also populates the Gophish Config object
func Setup() { func Setup() error {
//If the file already exists, delete it and recreate it //If the file already exists, delete it and recreate it
if _, err := os.Stat(config.Conf.DBPath); err == nil { if _, err := os.Stat(config.Conf.DBPath); err == nil {
os.Remove(Conf.DBPath) os.Remove(Conf.DBPath)
} }
fmt.Println("Creating db at " + config.Conf.DBPath) fmt.Println("Creating db at " + config.Conf.DBPath)
db, err := sql.Open("sqlite3", config.Conf.DBPath) db, err := sql.Open("sqlite3", config.Conf.DBPath)
defer db.Close()
if err != nil { if err != nil {
return nil, err return err
} }
//Create the tables needed //Create the tables needed
_, err = db.Exec( _, err = db.Exec(
`CREATE TABLE Users (id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT, hash VARCHAR(32), apikey VARCHAR(32));`) `CREATE TABLE Users (id INTEGER PRIMARY KEY AUTOINCREMENT, username TEXT, hash VARCHAR(32), apikey VARCHAR(32));`)
if err != nil { if err != nil {
return nil, err return err
} }
//Create the default user //Create the default user
stmt, err := db.Prepare(`INSERT INTO Users (username, hash, apikey) VALUES (?, ?, ?);`) stmt, err := db.Prepare(`INSERT INTO Users (username, hash, apikey) VALUES (?, ?, ?);`)
defer stmt.Close() defer stmt.Close()
if err != nil { if err != nil {
return nil, err return err
} }
_, err = stmt.Exec("jordan", "12345678901234567890123456789012", "12345678901234567890123456789012") _, err = stmt.Exec("jordan", "12345678901234567890123456789012", "12345678901234567890123456789012")
if err != nil { if err != nil {
return nil, err return err
} }
return db, nil return nil
} }