jilo-server/main.go

172 lines
4.7 KiB
Go
Raw Normal View History

2024-10-12 22:25:52 +00:00
package main
import (
2024-10-13 17:06:36 +00:00
"database/sql"
2024-10-14 15:23:09 +00:00
"flag"
"fmt"
2024-10-12 22:25:52 +00:00
"io/ioutil"
"log"
"net/http"
2024-10-14 15:23:09 +00:00
"os"
2024-10-12 22:25:52 +00:00
"time"
_ "github.com/mattn/go-sqlite3"
"gopkg.in/yaml.v2"
)
// Structures
type Agent struct {
Endpoint string `yaml:"endpoint"`
CheckPeriod int `yaml:"check_period"`
2024-10-12 22:25:52 +00:00
}
type Server struct {
Agents map[string]Agent `yaml:"agents"`
}
type Config struct {
Servers map[string]Server `yaml:"servers"`
DatabasePath string `yaml:"database_path"`
2024-10-12 22:25:52 +00:00
}
// Loading the config file
2024-10-12 22:25:52 +00:00
func readConfig(filePath string) Config {
var config Config
2024-10-12 22:25:52 +00:00
file, err := ioutil.ReadFile(filePath)
if err != nil {
log.Println("Can't read config file, using defaults.")
return config
}
err = yaml.Unmarshal(file, &config)
if err != nil {
log.Println("Can't parse the config file, using defaults.")
return config
}
return config
}
2024-10-14 15:23:09 +00:00
func setupDatabase(dbPath string, initDB bool) (*sql.DB, error) {
2024-10-13 17:06:36 +00:00
// Open the database
db, err := sql.Open("sqlite3", dbPath)
if err != nil {
return nil, err
}
2024-10-14 15:23:09 +00:00
// Check if the table exists
tableExists := checkTableExists(db)
if !tableExists && !initDB {
// Ask if we should create the table
fmt.Print("Table not found. Do you want to create it? (y/n): ")
var response string
fmt.Scanln(&response)
if response != "y" && response != "Y" {
log.Println("Exiting because the table is missing, but mandatory.")
os.Exit(1)
}
}
2024-10-13 17:06:36 +00:00
// If the table is not there, initialize it
createTable := `
CREATE TABLE IF NOT EXISTS endpoint_data (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
status_code INTEGER,
response_time_ms INTEGER
);`
_, err = db.Exec(createTable)
if err != nil {
return nil, err
}
return db, nil
}
2024-10-14 15:23:09 +00:00
// Check for the table
func checkTableExists(db *sql.DB) bool {
sql := `
SELECT name
FROM sqlite_master
WHERE type='table'
AND name='endpoint_data';`
row := db.QueryRow(sql)
var name string
err := row.Scan(&name)
return err == nil && name == "endpoint_data"
}
2024-10-12 22:25:52 +00:00
func checkEndpoint(endpoint string) (int, int64) {
log.Println("Sending HTTP get request to Jilo agent:", endpoint)
2024-10-12 22:25:52 +00:00
start := time.Now()
resp, err := http.Get(endpoint)
if err != nil {
2024-10-13 17:06:36 +00:00
log.Println("Failed to check the endpoint:", err)
2024-10-12 22:25:52 +00:00
return 0, 0
}
defer resp.Body.Close()
elapsed := time.Since(start).Milliseconds()
log.Printf("Received response: %d, Time taken: %d ms", resp.StatusCode, elapsed)
2024-10-12 22:25:52 +00:00
return resp.StatusCode, elapsed
}
2024-10-13 17:06:36 +00:00
func saveData(db *sql.DB, statusCode int, responseTime int64) {
_, err := db.Exec("INSERT INTO endpoint_data (status_code, response_time_ms) VALUES (?, ?)", statusCode, responseTime)
if err != nil {
log.Println("Failed to insert data into the database:", err)
}
}
// Main routine
2024-10-12 22:25:52 +00:00
func main() {
// First flush all the logs
log.SetFlags(log.LstdFlags | log.Lshortfile)
2024-10-14 15:23:09 +00:00
// Init the DB, "--init-db" creates the table
initDB := flag.Bool("init-db", false, "Create database table if not present without prompting")
flag.Parse()
// Config file
log.Println("Reading the config file...")
2024-10-12 22:25:52 +00:00
config := readConfig("jilo-server.conf")
2024-10-13 17:06:36 +00:00
// Connect to or setup the database
log.Println("Initializing the database...")
2024-10-14 15:23:09 +00:00
db, err := setupDatabase(config.DatabasePath, *initDB)
2024-10-13 17:06:36 +00:00
if err != nil {
log.Fatal("Failed to initialize the database:", err)
}
defer db.Close()
2024-10-12 22:25:52 +00:00
log.Println("Starting endpoint checker...")
// Iterate over the servers and agents
for serverName, server := range config.Servers {
for agentName, agent := range server.Agents {
go func(serverName, agentName string, agent Agent) {
// Ticker for the periodic checks
ticker := time.NewTicker(time.Duration(agent.CheckPeriod) * time.Minute)
defer ticker.Stop()
for {
log.Printf("Checking agent [%s - %s]: %s", serverName, agentName, agent.Endpoint)
statusCode, responseTime := checkEndpoint(agent.Endpoint)
log.Printf("Agent [%s - %s]: Status code: %d, Response time: %d ms", serverName, agentName, statusCode, responseTime)
saveData(db, statusCode, responseTime)
// Sleep until the next tick
<-ticker.C
}
}(serverName, agentName, agent)
}
2024-10-12 22:25:52 +00:00
}
// Prevent the main from exiting
select {}
2024-10-12 22:25:52 +00:00
}