CoreControl/agent/main.go

package main

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

	_ "github.com/lib/pq"
	"gopkg.in/gomail.v2"
)

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, "Test Notification from CoreControl")

		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, message string) error {
	switch providerType {
	case Telegram:
		return sendTelegramNotification(config, test, message)
	case Ntfy:
		return sendNtfyNotification(config, test, message)
	case SMTP:
		return sendSMTPNotification(config, test, message)
	default:
		return fmt.Errorf("unknown provider type: %s", providerType)
	}
}

func sendTelegramNotification(config json.RawMessage, test NotificationTest, message string) 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, message)

	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, message string) error {
	var cfg struct {
		URL   string `json:"url"`
		Token string `json:"token"`
	}

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

	baseURL := strings.TrimSuffix(cfg.URL, "/")

	req, err := http.NewRequest("POST", baseURL, strings.NewReader(message))
	if err != nil {
		return fmt.Errorf("failed to create NTFY request: %w", err)
	}

	if cfg.Token != "" {
		req.Header.Set("Authorization", "Bearer "+cfg.Token)
	}
	req.Header.Set("Content-Type", "text/plain")

	client := &http.Client{Timeout: 5 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("failed to send NTFY request: %w", err)
	}
	defer resp.Body.Close()

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

	log.Printf("Successfully sent NTFY notification to %s", baseURL)
	return nil
}

func sendSMTPNotification(config json.RawMessage, test NotificationTest, message string) error {
	var cfg struct {
		Host     string `json:"host"`
		Port     int    `json:"port"`
		Username string `json:"username"`
		Password string `json:"password"`
		From     string `json:"from"`
		To       string `json:"to"`
		Secure   bool   `json:"secure"`
	}

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

	d := gomail.NewDialer(cfg.Host, cfg.Port, cfg.Username, cfg.Password)
	if cfg.Secure {
		d.SSL = true
	}

	m := gomail.NewMessage()
	m.SetHeader("From", cfg.From)
	m.SetHeader("To", cfg.To)
	m.SetHeader("Subject", "Test Notification from CoreControl")
	m.SetBody("text/plain", message)

	if err := d.DialAndSend(m); err != nil {
		return fmt.Errorf("failed to send email: %w", err)
	}

	log.Printf("Successfully sent SMTP notification to %s", cfg.To)
	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
}