CoreControl/agent/main.go

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	_ "github.com/lib/pq"
)

var db *sql.DB

type NotificationType string

const (
	Telegram NotificationType = "TELEGRAM"
	Ntfy     NotificationType = "NTFY"
	SMTP     NotificationType = "SMTP"
)

type NotificationTest struct {
	ID                     int
	NotificationProviderID int
	Sent                   bool
	Success                sql.NullBool
	ProviderType           string
	ProviderConfig         json.RawMessage
}

func main() {
	var err error
	db, err = sql.Open("postgres", os.Getenv("DATABASE_URL"))
	if err != nil {
		log.Fatal("Failed to connect to database:", err)
	}
	defer db.Close()

	if err = db.Ping(); err != nil {
		log.Fatal("Failed to ping database:", err)
	}

	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		processPendingTests()
	}
}

func processPendingTests() {
	log.Println("Checking for pending notification tests...")

	tests, err := getPendingTests()
	if err != nil {
		log.Println("Error fetching tests:", err)
		return
	}

	for _, test := range tests {
		providerType := NotificationType(test.ProviderType)
		err := SendNotification(providerType, test.ProviderConfig, test)

		if updateErr := updateTestStatus(test.ID, err == nil); updateErr != nil {
			log.Printf("Error updating test status: %v", updateErr)
		}

		if err != nil {
			log.Printf("Failed to send test %d: %v", test.ID, err)
			continue
		}

		log.Printf("Successfully sent test %d", test.ID)
	}
}

func SendNotification(providerType NotificationType, config json.RawMessage, test NotificationTest) error {
	switch providerType {
	case Telegram:
		return sendTelegramNotification(config, test)
	case Ntfy:
		return sendNtfyNotification(config, test)
	case SMTP:
		return sendSMTPNotification(config, test)
	default:
		return fmt.Errorf("unknown provider type: %s", providerType)
	}
}

func sendTelegramNotification(config json.RawMessage, test NotificationTest) error {
	var cfg struct {
		Token  string `json:"token"`
		ChatID string `json:"chat_id"`
	}

	if err := json.Unmarshal(config, &cfg); err != nil {
		return fmt.Errorf("invalid Telegram config: %w", err)
	}

	apiUrl := fmt.Sprintf("https://api.telegram.org/bot%s/sendMessage?chat_id=%s&text=%s",
		cfg.Token, cfg.ChatID, "Test Notification")

	resp, err := http.Get(apiUrl)
	if err != nil {
		return fmt.Errorf("failed to send Telegram message: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("telegram API returned error status: %d", resp.StatusCode)
	}

	var result struct {
		OK    bool   `json:"ok"`
		Error string `json:"description,omitempty"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return fmt.Errorf("failed to decode telegram response: %w", err)
	}

	if !result.OK {
		return fmt.Errorf("telegram API error: %s", result.Error)
	}

	log.Printf("Successfully sent Telegram notification to chat %s", cfg.ChatID)
	return nil
}

func sendNtfyNotification(config json.RawMessage, test NotificationTest) error {
	var cfg struct {
		Topic  string `json:"topic"`
		Server string `json:"server"`
	}

	if err := json.Unmarshal(config, &cfg); err != nil {
		return fmt.Errorf("invalid NTFY config: %w", err)
	}

	log.Printf("Sending NTFY test to topic %s on %s", cfg.Topic, cfg.Server)
	return nil
}

func sendSMTPNotification(config json.RawMessage, test NotificationTest) error {
	var cfg struct {
		Server   string `json:"server"`
		Port     int    `json:"port"`
		Username string `json:"username"`
		To       string `json:"to"`
	}

	if err := json.Unmarshal(config, &cfg); err != nil {
		return fmt.Errorf("invalid SMTP config: %w", err)
	}

	log.Printf("Sending SMTP test to %s via %s:%d", cfg.To, cfg.Server, cfg.Port)
	return nil
}

func getPendingTests() ([]NotificationTest, error) {
	query := `
        SELECT t.id, 
               "t"."notificationProviderId",
               t.sent, 
               t.success, 
               p.type, 
               p.config
        FROM notification_tests t
        JOIN notification_providers p 
          ON "t"."notificationProviderId" = p.id
        WHERE t.sent IS NOT TRUE`

	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tests []NotificationTest
	for rows.Next() {
		var t NotificationTest
		var pt string

		if err := rows.Scan(
			&t.ID,
			&t.NotificationProviderID,
			&t.Sent,
			&t.Success,
			&pt,
			&t.ProviderConfig,
		); err != nil {
			return nil, err
		}

		t.ProviderType = pt
		tests = append(tests, t)
	}
	return tests, nil
}

func updateTestStatus(testID int, success bool) error {
	_, err := db.Exec(
		"UPDATE notification_tests SET sent = true, success = $1 WHERE id = $2",
		success, testID,
	)
	return err
}