main.go

package main

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

	"github.com/IBM/sarama"
	"github.com/joho/godotenv"
	_ "github.com/lib/pq" // Import the PostgreSQL driver
	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"
)

type CVERecord struct {
	// Define the structure based on the provided JSON schema
	CVEDataType    string `json:"CVE_data_type"`
	CVEDataFormat  string `json:"CVE_data_format"`
	CVEDataVersion string `json:"CVE_data_version"`
	CVEDataMeta    struct {
		ID                string `json:"cveId"`
		AssignerOrgID     string `json:"assignerOrgId"`
		AssignerShortName string `json:"assignerShortName"`
		DatePublished     string `json:"datePublished"`
		DateReserved      string `json:"dateReserved"`
		DateUpdated       string `json:"dateUpdated"`
		State             string `json:"state"`
	} `json:"cveMetadata"`
	Containers struct {
		CNA struct {
			ProviderMetadata struct {
				OrgID     string `json:"orgId"`
				ShortName string `json:"shortName"`
			} `json:"providerMetadata"`
			Title        string `json:"title"`
			Descriptions []struct {
				Lang  string `json:"lang"`
				Value string `json:"value"`
			} `json:"descriptions"`
			Affected []struct {
				Product string `json:"product"`
				Version string `json:"version"`
			} `json:"affected"`
		} `json:"cna"`
	} `json:"containers"`
}

var (
	db          *sql.DB
	kafkaClient sarama.Client
	topic       string
)

func init() {
	godotenv.Load()
	topic = os.Getenv("KAFKA_TOPIC")
	zerolog.TimeFieldFormat = zerolog.TimeFormatUnix
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr, TimeFormat: time.RFC3339})
}

func consumeFromKafka(brokerList []string, topic string, db *sql.DB) error {
	consumer, err := sarama.NewConsumer(brokerList, nil)
	if err != nil {
		return err
	}
	defer consumer.Close()

	partitionConsumer, err := consumer.ConsumePartition(topic, 0, sarama.OffsetNewest)
	if err != nil {
		return err
	}
	defer partitionConsumer.Close()

	for message := range partitionConsumer.Messages() {
		var cve CVERecord
		err := json.Unmarshal(message.Value, &cve)
		if err != nil {
			log.Error().Err(err).Msg("Failed to unmarshal CVE record")
			continue
		}

		err = storeInDatabase(db, cve)
		if err != nil {
			log.Error().Err(err).Msg("Failed to store CVE record in database")
		}
	}

	return nil
}

func storeInDatabase(db *sql.DB, cve CVERecord) error {
	cveJSON, err := json.Marshal(cve)
	if err != nil {
		return fmt.Errorf("failed to marshal CVE: %v", err)
	}

	insertQuery := `
    INSERT INTO cve.cves (cve_id, data)
    VALUES ($1, $2)
    ON CONFLICT (cve_id) DO UPDATE SET data = EXCLUDED.data;
    `
	if _, err := db.Exec(insertQuery, cve.CVEDataMeta.ID, cveJSON); err != nil {
		return fmt.Errorf("failed to insert CVE: %v", err)
	}
	return nil
}

func healthzHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("ok"))
}

func readinessHandler(w http.ResponseWriter, r *http.Request) {
	// Check Kafka connection
	if err := kafkaClient.RefreshMetadata(); err != nil {
		http.Error(w, "kafka connection failed", http.StatusInternalServerError)
		log.Error().Err(err).Msg("Kafka connection failed")
		return
	}

	// Check if Kafka topic exists
	topics, err := kafkaClient.Topics()
	if err != nil || !contains(topics, topic) {
		http.Error(w, "kafka topic check failed", http.StatusInternalServerError)
		log.Error().Err(err).Msg("Kafka topic check failed")
		return
	}

	// Check PostgreSQL connection
	if err := db.Ping(); err != nil {
		http.Error(w, "database connection failed", http.StatusInternalServerError)
		log.Error().Err(err).Msg("Database connection failed")
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("ready"))
}

func contains(slice []string, item string) bool {
	for _, s := range slice {
		if s == item {
			return true
		}
	}
	return false
}

func main() {
	brokerList := []string{os.Getenv("KAFKA_BROKER")}
	var err error

	log.Info().Msg("Connecting to Kafka...")
	kafkaClient, err = sarama.NewClient(brokerList, nil)
	if err != nil {
		log.Fatal().Err(err).Msg("Error connecting to Kafka")
		return
	}
	defer kafkaClient.Close()

	log.Info().Msg("Connecting to PostgreSQL database...")
	connStr := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable", os.Getenv("DB_HOST"), os.Getenv("DB_PORT"), os.Getenv("DB_USER"), os.Getenv("DB_PASSWORD"), os.Getenv("DB_NAME"))
	db, err = sql.Open("postgres", connStr)
	if err != nil {
		log.Fatal().Err(err).Msg("Error connecting to database")
		return
	}
	defer db.Close()

	go func() {
		err = consumeFromKafka(brokerList, topic, db)
		if err != nil {
			log.Error().Err(err).Msg("Failed to consume from Kafka")
		}
	}()

	http.HandleFunc("/healthz", healthzHandler)
	http.HandleFunc("/readiness", readinessHandler)
	log.Fatal().Err(http.ListenAndServe(":8080", nil)).Msg("HTTP server stopped")
}