jilo-server/main.go

332 lines
9.4 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-15 19:34:08 +00:00
"net/url"
2024-10-14 15:23:09 +00:00
"os"
2024-10-12 22:25:52 +00:00
"time"
"github.com/golang-jwt/jwt/v5"
2024-10-12 22:25:52 +00:00
_ "github.com/mattn/go-sqlite3"
"gopkg.in/yaml.v2"
)
// Structures
type Agent struct {
2024-10-16 06:45:07 +00:00
ID int
URL string
Secret string
CheckPeriod int
}
type Config struct {
2024-10-16 06:45:07 +00:00
DatabasePath string `yaml:"database_path"`
HealthCheckEnabled bool `yaml:"health_check_enabled"`
HealthCheckPort int `yaml:"health_check_port"`
HealthCheckEndpoint string `yaml:"health_check_endpoint"`
2024-10-15 17:48:28 +00:00
}
var defaultConfig = Config {
DatabasePath: "./jilo-server.db",
2024-10-16 06:45:07 +00:00
HealthCheckEnabled: false,
HealthCheckPort: 8080,
HealthCheckEndpoint: "/health",
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 {
2024-10-15 17:48:28 +00:00
config := defaultConfig
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-16 06:45:07 +00:00
// Start the health check
func startHealthCheckServer(port int, endpoint string) {
http.HandleFunc(endpoint, func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
w.Write([]byte("Jilo Server is running"))
})
address := fmt.Sprintf(":%d", port)
log.Printf("Starting health check server on %s%s", address, endpoint)
go http.ListenAndServe(address, nil)
}
// Database initialization
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 := `
2024-10-15 17:48:28 +00:00
CREATE TABLE IF NOT EXISTS jilo_agent_checks (
2024-10-13 17:06:36 +00:00
id INTEGER PRIMARY KEY AUTOINCREMENT,
2024-10-15 17:48:28 +00:00
agent_id INTEGER,
2024-10-13 17:06:36 +00:00
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
status_code INTEGER,
response_time_ms INTEGER,
2024-10-15 17:48:28 +00:00
response_content TEXT,
FOREIGN KEY(agent_id) REFERENCES jilo_agents(id)
);`
2024-10-13 17:06:36 +00:00
_, 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 {
2024-10-15 17:48:28 +00:00
sql := `SELECT
name
FROM
sqlite_master
WHERE
type='table'
AND
name='jilo_agent_checks';`
2024-10-14 15:23:09 +00:00
row := db.QueryRow(sql)
var name string
err := row.Scan(&name)
2024-10-15 17:48:28 +00:00
return err == nil && name == "jilo_agent_checks"
}
// Get Jilo agents details
func getAgents(db *sql.DB) ([]Agent, error) {
sql := `SELECT
ja.id,
ja.url,
ja.secret_key,
ja.check_period,
jat.endpoint
FROM
jilo_agents ja
JOIN
jilo_agent_types jat ON ja.agent_type_id = jat.id`
rows, err := db.Query(sql)
if err != nil {
return nil, err
}
defer rows.Close()
var agents []Agent
for rows.Next() {
var agent Agent
var endpoint string
var checkPeriod int
// Get the agent details
if err := rows.Scan(&agent.ID, &agent.URL, &agent.Secret, &checkPeriod, &endpoint); err != nil {
return nil, err
}
// Form the whole enpoint
agent.URL += endpoint
agent.CheckPeriod = checkPeriod
agents = append(agents, agent)
}
// We return the endpoints, not all the details
return agents, nil
2024-10-14 15:23:09 +00:00
}
2024-10-15 19:34:08 +00:00
// Format the enpoint URL to use in logs
func getDomainAndPath(fullURL string) string {
parsedURL, err := url.Parse(fullURL)
if err != nil {
// Fallback to the original URL on error
return fullURL
}
return parsedURL.Host + parsedURL.Path
}
// JWT token generation
func generateJWT(secret string) (string, error) {
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
"iat": time.Now().Unix(),
})
tokenString, err := token.SignedString([]byte(secret))
if err != nil {
return "", err
}
return tokenString, nil
}
// Check agent endpoint
2024-10-15 15:05:05 +00:00
func checkEndpoint(agent Agent) (int, int64, string, bool) {
2024-10-15 17:48:28 +00:00
log.Println("Sending HTTP get request to Jilo agent:", agent.URL)
// Create the http request
2024-10-15 17:48:28 +00:00
req, err := http.NewRequest("GET", agent.URL, nil)
if err != nil {
log.Println("Failed to create the HTTP request:", err)
2024-10-15 15:05:05 +00:00
return 0, 0, "", false
}
2024-10-15 15:05:05 +00:00
// Generate the JWT token
if agent.Secret != "" {
token, err := generateJWT(agent.Secret)
if err != nil {
log.Println("Failed to generate JWT token:", err)
return 0, 0, "", false
}
// Set Authorization header
req.Header.Set("Authorization", "Bearer "+token)
}
2024-10-12 22:25:52 +00:00
start := time.Now()
resp, err := http.DefaultClient.Do(req)
2024-10-12 22:25:52 +00:00
if err != nil {
2024-10-13 17:06:36 +00:00
log.Println("Failed to check the endpoint:", err)
2024-10-15 15:05:05 +00:00
return 0, 0, "", false
2024-10-12 22:25:52 +00:00
}
defer resp.Body.Close()
elapsed := time.Since(start).Milliseconds()
// Read the response body
body, err := ioutil.ReadAll(resp.Body)
if err != nil {
log.Println("Failed to read the response body:", err)
2024-10-15 15:05:05 +00:00
return resp.StatusCode, elapsed, "", false
}
log.Printf("Received response: %d, Time taken: %d ms", resp.StatusCode, elapsed)
2024-10-15 15:05:05 +00:00
return resp.StatusCode, elapsed, string(body), true
2024-10-12 22:25:52 +00:00
}
// Insert the checks into the database
2024-10-15 17:48:28 +00:00
func saveData(db *sql.DB, agentID int, statusCode int, responseTime int64, responseContent string) {
sql := `INSERT INTO
jilo_agent_checks
(agent_id, status_code, response_time_ms, response_content)
VALUES
(?, ?, ?, ?)`
_, err := db.Exec(sql, agentID, statusCode, responseTime, responseContent)
2024-10-13 17:06:36 +00:00
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-15 19:51:17 +00:00
// Command-line options
// "--init-db" creates the table
2024-10-14 15:23:09 +00:00
initDB := flag.Bool("init-db", false, "Create database table if not present without prompting")
2024-10-15 19:51:17 +00:00
// Config file
configPath := flag.String("config", "", "Path to the configuration file (use -c or --config)")
flag.StringVar(configPath, "c", "", "Path to the configuration file")
2024-10-15 17:48:28 +00:00
2024-10-14 15:23:09 +00:00
flag.Parse()
2024-10-15 19:51:17 +00:00
// Choosing the config file
finalConfigPath := "./jilo-server.conf" // this is the default we fall to
if *configPath != "" {
if _, err := os.Stat(*configPath); err == nil {
finalConfigPath = *configPath
} else {
log.Printf("Specified file \"%s\" doesn't exist. Falling back to the default \"%s\".", *configPath, finalConfigPath)
}
}
// Config file
2024-10-15 19:51:17 +00:00
log.Printf("Using config file %s", finalConfigPath)
config := readConfig(finalConfigPath)
2024-10-12 22:25:52 +00:00
2024-10-16 06:45:07 +00:00
// Start the health check, if it's enabled in the config file
if config.HealthCheckEnabled {
startHealthCheckServer(config.HealthCheckPort, config.HealthCheckEndpoint)
}
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-15 17:48:28 +00:00
// Prepare the Agents
agents, err := getAgents(db)
if err != nil {
log.Fatal("Failed to fetch the agents:", err)
}
2024-10-12 22:25:52 +00:00
log.Println("Starting endpoint checker...")
// Iterate over the servers and agents
2024-10-15 17:48:28 +00:00
for _, agent := range agents {
if agent.CheckPeriod > 0 {
go func(agent Agent) {
// Ticker for the periodic checks
ticker := time.NewTicker(time.Duration(agent.CheckPeriod) * time.Minute)
defer ticker.Stop()
for {
2024-10-15 17:48:28 +00:00
log.Printf("Checking agent [%d]: %s", agent.ID, agent.URL)
2024-10-15 15:05:05 +00:00
statusCode, responseTime, responseContent, success := checkEndpoint(agent)
if success {
2024-10-15 17:48:28 +00:00
log.Printf("Agent [%d]: Status code: %d, Response time: %d ms", agent.ID, statusCode, responseTime)
saveData(db, agent.ID, statusCode, responseTime, responseContent)
2024-10-15 15:05:05 +00:00
} else {
2024-10-15 19:34:08 +00:00
log.Printf("Check for agent %s (%d) failed, skipping database insert", getDomainAndPath(agent.URL), agent.ID)
2024-10-15 15:05:05 +00:00
}
// Sleep until the next tick
<-ticker.C
}
2024-10-15 17:48:28 +00:00
}(agent)
} else {
2024-10-15 19:34:08 +00:00
log.Printf("Agent %s (%d) has an invalid CheckPeriod (%d), skipping it.", getDomainAndPath(agent.URL), agent.ID, agent.CheckPeriod)
}
2024-10-12 22:25:52 +00:00
}
// Prevent the main from exiting
select {}
2024-10-12 22:25:52 +00:00
}