diff --git a/config.toml b/config.toml index 8490224..4d604d5 100644 --- a/config.toml +++ b/config.toml @@ -6,6 +6,9 @@ NUMBER_OF_WORKERS=20 SUPRESS_CUDA=1 +[ServiceUser] +USER="service" + [Worker] PULLING_TIME="500ms" NUMBER_OF_WORKERS=1 diff --git a/logic/utils/config.go b/logic/utils/config.go index c764186..ea1f62b 100644 --- a/logic/utils/config.go +++ b/logic/utils/config.go @@ -1,7 +1,11 @@ package utils import ( + "database/sql" "os" + "strings" + + . "git.andr3h3nriqu3s.com/andr3/fyp/logic/db_types" toml "github.com/BurntSushi/toml" "github.com/charmbracelet/log" @@ -12,6 +16,11 @@ type WorkerConfig struct { Pulling string `toml:"pulling_time"` } +type ServiceUser struct { + User string `toml:"user"` + UserId string `toml:"__user__id__"` +} + type Config struct { Hostname string Port int @@ -19,6 +28,8 @@ type Config struct { SupressCuda int `toml:"supress_cuda"` GpuWorker WorkerConfig `toml:"Worker"` + + ServiceUser ServiceUser `toml:"ServiceUser"` } func LoadConfig() Config { @@ -37,17 +48,62 @@ func LoadConfig() Config { NumberOfWorkers: 1, Pulling: "500ms", }, + ServiceUser: ServiceUser{ + User: "Service", + UserId: "", + }, } } var conf Config _, err = toml.Decode(string(dat), &conf) - if conf.SupressCuda == 1 { - log.Warn("Supressing Cuda Messages!") - os.Setenv("TF_CPP_MIN_VLOG_LEVEL", "3") - os.Setenv("TF_CPP_MIN_LOG_LEVEL", "3") - } + if conf.SupressCuda == 1 { + log.Warn("Supressing Cuda Messages!") + os.Setenv("TF_CPP_MIN_VLOG_LEVEL", "3") + os.Setenv("TF_CPP_MIN_LOG_LEVEL", "3") + } return conf } + +func (c *Config) GenerateToken(db *sql.DB) { + if c.ServiceUser.User == "" { + log.Fatal("A user needs to be set in a configuration file") + } + + var user struct { + Password string `db:"password"` + UserId string `db:"id"` + } + + err := GetDBOnce(db, &user, "users where username=$1;", c.ServiceUser.User) + if err == NotFoundError { + var newUser = struct { + Username string + Email string + Salt string + Password string + UserType UserType `db:"user_type"` + }{ + c.ServiceUser.User, + c.ServiceUser.User, + "", + "", + User_Admin, + } + id, err := InsertReturnId(db, &newUser, "users", "id") + if err != nil { + log.Fatal("Failed to create user", "err", err) + } + c.ServiceUser.UserId = id + } else if err != nil { + log.Fatal("To get user name", "err", err) + return + } else { + if len(strings.ReplaceAll(user.Password, " ", "")) != 0 { + log.Fatal("User already exists and is not the service user", "user", user) + } + c.ServiceUser.UserId = user.UserId + } +} diff --git a/main.go b/main.go index 494aa8a..fe44f34 100644 --- a/main.go +++ b/main.go @@ -37,6 +37,7 @@ func main() { config := LoadConfig() log.Info("Config loaded!", "config", config) + config.GenerateToken(db) StartRunners(db, config)