package models import ( "errors" "fmt" "log" "os" "reflect" "ccsecrets/utils" "github.com/jmoiron/sqlx" "github.com/joho/godotenv" _ "modernc.org/sqlite" ) var db *sqlx.DB const ( sqlFile = "ccsecrets.db" ) const createRoles string = ` CREATE TABLE IF NOT EXISTS roles ( RoleId INTEGER PRIMARY KEY ASC, RoleName VARCHAR, ReadOnly BOOLEAN, Admin BOOLEAN ); ` const createUsers string = ` CREATE TABLE IF NOT EXISTS users ( UserId INTEGER PRIMARY KEY ASC, RoleId INTEGER, UserName VARCHAR, Password VARCHAR, FOREIGN KEY (RoleId) REFERENCES roles(RoleId) ); ` const createSecrets string = ` CREATE TABLE IF NOT EXISTS secrets ( SecretId INTEGER PRIMARY KEY ASC, RoleId INTEGER, DeviceName VARCHAR, DeviceCategory VARCHAR, UserName VARCHAR, Secret VARCHAR, FOREIGN KEY (RoleId) REFERENCES roles(RoleId) ); ` const createSchema string = ` CREATE TABLE IF NOT EXISTS schema ( Version INTEGER ); ` // Establish connection to sqlite database func ConnectDatabase() { var err error // Load data from environment file err = godotenv.Load(".env") if err != nil { log.Fatalf("Error loading .env file") } // Try using sqlite as our database sqlPath := utils.GetFilePath(sqlFile) db, err = sqlx.Open("sqlite", sqlPath) if err != nil { fmt.Printf("Error opening sqlite database connection to file '%s' : '%s'\n", sqlPath, err) os.Exit(1) } else { fmt.Printf("Connected to sqlite database file '%s'\n", sqlPath) } //sqlx.NameMapper = func(s string) string { return s } // Make sure our tables exist CreateTables() //defer db.Close() } func DisconnectDatabase() { fmt.Printf("DisconnectDatabase called") defer db.Close() } func CreateTables() { var err error var rowCount int // Create database tables if it doesn't exist // Roles table should go first since other tables refer to it if _, err = db.Exec(createRoles); err != nil { fmt.Printf("Error checking roles table : '%s'", err) os.Exit(1) } rowCount, _ = CheckCount("roles") if rowCount == 0 { if _, err = db.Exec("INSERT INTO roles VALUES(1, 'Admin', false, true);"); err != nil { fmt.Printf("Error adding initial admin role : '%s'", err) os.Exit(1) } if _, err = db.Exec("INSERT INTO roles VALUES(2, 'UserRole', false, false);"); err != nil { fmt.Printf("Error adding initial admin role : '%s'", err) os.Exit(1) } if _, err = db.Exec("INSERT INTO roles VALUES(3, 'GuestRole', true, false);"); err != nil { fmt.Printf("Error adding initial admin role : '%s'", err) os.Exit(1) } } // Users table if _, err = db.Exec(createUsers); err != nil { fmt.Printf("Error checking users table : '%s'", err) os.Exit(1) } rowCount, _ = CheckCount("users") if rowCount == 0 { if _, err = db.Exec("INSERT INTO users VALUES(1, 1, 'Administrator', '$2a$10$k1qldm.bWqZsQWrKPdahR.Pfz5LxkMUka2.8INEeSD7euzkiznIR.');"); err != nil { fmt.Printf("Error adding initial admin role : '%s'", err) os.Exit(1) } } // Secrets table if _, err = db.Exec(createSecrets); err != nil { fmt.Printf("Error checking secrets table : '%s'", err) os.Exit(1) } // Schema table should go last so we know if the database has a value in the schema table then everything was created properly if _, err = db.Exec(createSchema); err != nil { fmt.Printf("Error checking schema table : '%s'", err) os.Exit(1) } schemaCheck, _ := CheckColumnExists("schema", "version") if !schemaCheck { if _, err = db.Exec("INSERT INTO schema VALUES(1);"); err != nil { fmt.Printf("Error adding initial scehama version : '%s'", err) os.Exit(1) } } } // Count the number of records in the sqlite database // Borrowed from https://gist.github.com/trkrameshkumar/f4f1c00ef5d578561c96?permalink_comment_id=2687592#gistcomment-2687592 func CheckCount(tablename string) (int, error) { var count int stmt, err := db.Prepare("SELECT COUNT(*) as count FROM " + tablename) if err != nil { fmt.Printf("CheckCount error preparing sqlite statement : '%s'\n", err) return 0, err } err = stmt.QueryRow().Scan(&count) if err != nil { fmt.Printf("CheckCount error querying database record count : '%s'\n", err) return 0, err } stmt.Close() // or use defer rows.Close(), idc return count, nil } // From https://stackoverflow.com/a/60100045 func GenerateInsertMethod(q interface{}) (string, error) { if reflect.ValueOf(q).Kind() == reflect.Struct { query := fmt.Sprintf("INSERT INTO %s", reflect.TypeOf(q).Name()) fieldNames := "" fieldValues := "" v := reflect.ValueOf(q) for i := 0; i < v.NumField(); i++ { if i == 0 { fieldNames = fmt.Sprintf("%s%s", fieldNames, v.Type().Field(i).Name) } else { fieldNames = fmt.Sprintf("%s, %s", fieldNames, v.Type().Field(i).Name) } switch v.Field(i).Kind() { case reflect.Int: if i == 0 { fieldValues = fmt.Sprintf("%s%d", fieldValues, v.Field(i).Int()) } else { fieldValues = fmt.Sprintf("%s, %d", fieldValues, v.Field(i).Int()) } case reflect.String: if i == 0 { fieldValues = fmt.Sprintf("%s\"%s\"", fieldValues, v.Field(i).String()) } else { fieldValues = fmt.Sprintf("%s, \"%s\"", fieldValues, v.Field(i).String()) } case reflect.Bool: var boolSet int8 if v.Field(i).Bool() { boolSet = 1 } if i == 0 { fieldValues = fmt.Sprintf("%s%d", fieldValues, boolSet) } else { fieldValues = fmt.Sprintf("%s, %d", fieldValues, boolSet) } default: fmt.Printf("Unsupported type '%s'\n", v.Field(i).Kind()) } } query = fmt.Sprintf("%s(%s) VALUES (%s)", query, fieldNames, fieldValues) return query, nil } return "", errors.New("SqlGenerationError") } func CheckColumnExists(table string, column string) (bool, error) { var count int64 rows, err := db.Queryx("SELECT COUNT(*) AS CNTREC FROM pragma_table_info('" + table + "') WHERE name='" + column + "';") if err != nil { fmt.Printf("CheckColumnExists error querying database for existence of column '%s' : '%s'\n", column, err) return false, err } defer rows.Close() for rows.Next() { // cols is an []interface{} of all of the column results cols, _ := rows.SliceScan() fmt.Printf("CheckColumnExists Value is '%v'\n", cols[0].(int64)) count = cols[0].(int64) if count == 1 { return true, nil } else { return false, nil } } err = rows.Err() if err != nil { fmt.Printf("CheckColumnExists error getting results : '%s'\n", err) return false, err } return false, nil }