Improvements for haven-notify

haven-notify/main.go

@@ -1,13 +1,16 @@
-
 package main
 
 import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"html/template"
+	"io/ioutil"
 	"log"
 	"net/http"
 	"os"
+	"strings"
+	"time"
 )
 
 // Notification payload
@@ -18,6 +21,9 @@ type Notification struct {
 
 func main() {
 	http.HandleFunc("/notify", notifyHandler)
+	http.HandleFunc("/ready", readinessHandler)
+	http.HandleFunc("/live", livenessHandler)
+	http.HandleFunc("/template/notify/", templateNotifyHandler)
 	log.Println("Starting server on :8080...")
 	log.Fatal(http.ListenAndServe(":8080", nil))
 }
@@ -53,6 +59,18 @@ func notifyHandler(w http.ResponseWriter, r *http.Request) {
 	w.Write([]byte("Notification sent"))
 }
 
+// Readiness handler
+func readinessHandler(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("Ready"))
+}
+
+// Liveness handler
+func livenessHandler(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("Alive"))
+}
+
 func sendDiscordNotification(title, message string) error {
 	webhookURL := os.Getenv("WEBHOOK_URL")
 	if webhookURL == "" {
@@ -90,3 +108,107 @@ func sendDiscordNotification(title, message string) error {
 	log.Printf("Discord notification sent successfully: Title='%s'", title)
 	return nil
 }
+
+func templateNotifyHandler(w http.ResponseWriter, r *http.Request) {
+	log.Printf("Incoming %s request from %s to %s", r.Method, r.RemoteAddr, r.URL.Path)
+	if r.Method != http.MethodPost {
+		log.Printf("Method not allowed: %s", r.Method)
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		w.Write([]byte("Method not allowed"))
+		return
+	}
+
+	templateName := r.URL.Path[len("/template/notify/"):] // Extract template name
+	if templateName == "" {
+		log.Printf("Template name not provided")
+		w.WriteHeader(http.StatusBadRequest)
+		w.Write([]byte("Template name not provided"))
+		return
+	}
+
+	templatePath := "template/" + templateName + ".tmpl"
+	templateData, err := ioutil.ReadFile(templatePath)
+	if err != nil {
+		log.Printf("Failed to read template: %v", err)
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte("Failed to read template"))
+		return
+	}
+
+	tmpl, err := template.New(templateName).Funcs(template.FuncMap{
+		"formatSize": func(size float64) string {
+			if size > 1024 {
+				return fmt.Sprintf("%.2f GiB", size/1024)
+			}
+			return fmt.Sprintf("%.2f MiB", size)
+		},
+		"upper": strings.ToUpper,
+		"lower": strings.ToLower,
+		"title": strings.Title,
+		"now": func() string {
+			return fmt.Sprintf("%d", time.Now().Unix())
+		},
+		"formatTime": func(timestamp string) string {
+			if timestamp == "" {
+				return time.Now().Format("2006-01-02T15:04:05Z")
+			}
+			return timestamp
+		},
+	}).Parse(string(templateData))
+	if err != nil {
+		log.Printf("Failed to parse template: %v", err)
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte("Failed to parse template"))
+		return
+	}
+
+	var rawPayload map[string]interface{}
+	if err := json.NewDecoder(r.Body).Decode(&rawPayload); err != nil {
+		log.Printf("Invalid payload: %v", err)
+		w.WriteHeader(http.StatusBadRequest)
+		w.Write([]byte("Invalid payload"))
+		return
+	}
+
+	// Normalize keys to lowercase for case-insensitive parsing
+	payload := make(map[string]interface{})
+	for key, value := range rawPayload {
+		payload[strings.ToLower(key)] = value
+	}
+
+	var filledTemplate bytes.Buffer
+	if err := tmpl.Execute(&filledTemplate, payload); err != nil {
+		log.Printf("Failed to execute template: %v", err)
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte("Failed to execute template"))
+		return
+	}
+
+	webhookURL := os.Getenv("WEBHOOK_URL")
+	if webhookURL == "" {
+		log.Printf("WEBHOOK_URL environment variable not set")
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte("WEBHOOK_URL environment variable not set"))
+		return
+	}
+
+	resp, err := http.Post(webhookURL, "application/json", &filledTemplate)
+	if err != nil {
+		log.Printf("Error posting to Discord webhook: %v", err)
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte("Failed to send notification"))
+		return
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		log.Printf("Discord webhook returned status: %s", resp.Status)
+		w.WriteHeader(http.StatusInternalServerError)
+		w.Write([]byte("Failed to send notification"))
+		return
+	}
+
+	log.Printf("Notification sent successfully using template '%s'", templateName)
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("Notification sent"))
+}