From ca7da7fc623a57024918faa247f73d816f0432fd Mon Sep 17 00:00:00 2001 From: Stefan Etringer Date: Mon, 26 May 2025 14:19:49 +0000 Subject: [PATCH] moved from mysql to sqlite for the agents database --- main.go | 9 +++- src/server/api/agentApi.go | 3 +- src/server/database/database.go | 74 +++++++++++++++++++++++++++------ 3 files changed, 71 insertions(+), 15 deletions(-) diff --git a/main.go b/main.go index 2b7a74a..3d2b2f3 100644 --- a/main.go +++ b/main.go @@ -70,6 +70,9 @@ func init() { if err != nil { log.Fatal(err) } + + // Agents database + db = database.InitSQLiteDB("/tmp/gontrol_agents.db") } func renderTemplate(w http.ResponseWriter, tmpl string, data interface{}) { @@ -289,6 +292,8 @@ func main() { // sqlite3 has been initialized in init() defer logger.CloseDB() + defer db.Close() + var cfg Config readEnv(&cfg) @@ -310,8 +315,8 @@ func main() { webMux.Handle("/static/", http.FileServer(http.FS(staticFiles))) webMux.HandleFunc("/proxyAgent", proxyAgentHandler) - db = database.InitDB (cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Name) - defer db.Close() + // db := database.InitDB (cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Name) + // defer db.Close() name := randomname.GenerateRandomName() fmt.Sprintf("Server instance: %s", name) diff --git a/src/server/api/agentApi.go b/src/server/api/agentApi.go index 752534c..a3541f8 100644 --- a/src/server/api/agentApi.go +++ b/src/server/api/agentApi.go @@ -55,7 +55,8 @@ func CreateAgent(db *sql.DB, w http.ResponseWriter, r * http.Request) (http.Resp // initalContact := r.FormValue("initialContact") // lastContact := r.FormValue("lastContact") - query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, addPort, initialContact, lastContact) VALUES (?, ?, ?, ?, ?, NOW(), NOW())" + // query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, addPort, initialContact, lastContact) VALUES (?, ?, ?, ?, ?, NOW(), NOW())" + query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, addPort, initialContact, lastContact) VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" _, err = db.Exec(query, agentId, agentName, agentType, IPv4Address, addPort) if err != nil { log.Printf("Database err is: %s", err) diff --git a/src/server/database/database.go b/src/server/database/database.go index c4d5757..aff92a0 100644 --- a/src/server/database/database.go +++ b/src/server/database/database.go @@ -1,26 +1,76 @@ package database import ( - "fmt" + // "fmt" "log" "database/sql" _ "github.com/go-sql-driver/mysql" + _ "github.com/mattn/go-sqlite3" +) + +var ( + Agent_db *sql.DB + ) -func InitDB (dbUser string, dbPassword string, dbHost string, dbPort int16, dbName string ) *sql.DB { - var db *sql.DB - var err error - dbOpen := fmt.Sprintf("%s:%s@(%s:%d)/%s?parseTime=true", dbUser,dbPassword, dbHost, dbPort, dbName) - db, err = sql.Open("mysql", dbOpen) +/* This is the connection initializer for mysql */ +// func InitDB (dbUser string, dbPassword string, dbHost string, dbPort int16, dbName string ) *sql.DB { +// var db *sql.DB +// var err error +// dbOpen := fmt.Sprintf("%s:%s@(%s:%d)/%s?parseTime=true", dbUser,dbPassword, dbHost, dbPort, dbName) +// db, err = sql.Open("mysql", dbOpen) +// if err != nil { +// log.Fatal(err) +// } + +// if err = db.Ping(); err != nil { +// log.Fatal(err) +// } +// log.Println("Connected to database") + +// return db +// } + +func InitSQLiteDB(dbPath string) *sql.DB { + var err error + Agent_db, err = sql.Open("sqlite3", dbPath) if err != nil { - log.Fatal(err) + log.Println("Error opening DB: %w", err) } - if err = db.Ping(); err != nil { - log.Fatal(err) - } - log.Println("Connected to database") + // CreateTableQuery := `CREATE TABLE IF NOT EXISTS logs ( + // id INTEGER PRIMARY KEY AUTOINCREMENT, + // message TEXT, + // timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, + // level TEXT + // );` - return db + // id TEXT PRIMARY KEY DEFAULT (lower(hex(randomblob(16)))), + CreateTableQuery :=`CREATE TABLE IF NOT EXISTS agents ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + agentId INTEGER UNIQUE, + agentType TEXT, + agentName TEXT, + IPv4Address TEXT, + initialContact datetime, + lastContact datetime, + addPort TEXT + );` + + _, err = Agent_db.Exec(CreateTableQuery) + if err != nil { + log.Println("Error creating table: %w", err) + } + + return Agent_db +} + +func CloseDB() { + if Agent_db != nil { + err := Agent_db.Close() + if err != nil { + log.Println("Error closing database: %v", err) + } + } }