Restructure and refactor codebase

internal/db/db.go

@@ -0,0 +1,214 @@
+package internal
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log"
+	"os"
+	"os/exec"
+	"sync"
+	"time"
+
+	"github.com/oschwald/maxminddb-golang"
+)
+
+// Database paths
+const (
+	CityDBPath = "./GeoLite2-City.mmdb"
+	ASNDBPath  = "./GeoLite2-ASN.mmdb"
+)
+
+// Common errors
+var (
+	ErrDatabaseNotFound = errors.New("database file not found")
+	ErrDatabaseOpen     = errors.New("failed to open database")
+	ErrDownloadFailed   = errors.New("failed to download database")
+)
+
+// Handles MaxMind GeoIP database operations
+type GeoIPManager struct {
+	cityDB *maxminddb.Reader
+	asnDB  *maxminddb.Reader
+	mu     sync.RWMutex
+}
+
+// Creates and initializes a new GeoIP database
+func NewGeoIPManager() (*GeoIPManager, error) {
+	manager := &GeoIPManager{}
+	if err := manager.Initialize(); err != nil {
+		return nil, fmt.Errorf("initializing GeoIP manager: %w", err)
+	}
+	return manager, nil
+}
+
+// Initialize opens the GeoIP databases, downloading them if necessary
+func (g *GeoIPManager) Initialize() error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	if err := g.openCityDB(); err != nil {
+		return err
+	}
+
+	if err := g.openASNDB(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Opens the city database, downloading it if necessary
+func (g *GeoIPManager) openCityDB() error {
+	var err error
+	g.cityDB, err = maxminddb.Open(CityDBPath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			log.Println("City database not found, attempting to download...")
+			if err := g.downloadDatabases(); err != nil {
+				return fmt.Errorf("%w: %v", ErrDownloadFailed, err)
+			}
+			g.cityDB, err = maxminddb.Open(CityDBPath)
+			if err != nil {
+				return fmt.Errorf("%w (city): %v", ErrDatabaseOpen, err)
+			}
+			return nil
+		}
+		return fmt.Errorf("%w (city): %v", ErrDatabaseOpen, err)
+	}
+	return nil
+}
+
+// Opens the ASN database, downloading it if necessary
+func (g *GeoIPManager) openASNDB() error {
+	var err error
+	g.asnDB, err = maxminddb.Open(ASNDBPath)
+	if err != nil {
+		if os.IsNotExist(err) {
+			log.Println("ASN database not found, attempting to download...")
+			if err := g.downloadDatabases(); err != nil {
+				return fmt.Errorf("%w: %v", ErrDownloadFailed, err)
+			}
+			g.asnDB, err = maxminddb.Open(ASNDBPath)
+			if err != nil {
+				return fmt.Errorf("%w (ASN): %v", ErrDatabaseOpen, err)
+			}
+			return nil
+		}
+		return fmt.Errorf("%w (ASN): %v", ErrDatabaseOpen, err)
+	}
+	return nil
+}
+
+// Downloads both GeoIP databases using geoipupdate
+func (g *GeoIPManager) downloadDatabases() error {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+	defer cancel()
+
+	cmd := exec.CommandContext(ctx, "geoipupdate", "-d", "./")
+	output, err := cmd.CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("failed to download databases: %v. Output: %s. Ensure geoipupdate is installed and configured", err, output)
+	}
+	return nil
+}
+
+// Close properly closes the database readers
+func (g *GeoIPManager) Close() error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	var errs []error
+
+	if g.cityDB != nil {
+		if err := g.cityDB.Close(); err != nil {
+			errs = append(errs, fmt.Errorf("closing city database: %w", err))
+		}
+	}
+
+	if g.asnDB != nil {
+		if err := g.asnDB.Close(); err != nil {
+			errs = append(errs, fmt.Errorf("closing ASN database: %w", err))
+		}
+	}
+
+	if len(errs) > 0 {
+		return fmt.Errorf("errors closing databases: %v", errs)
+	}
+	return nil
+}
+
+// Safely provides read access to the City database
+func (g *GeoIPManager) GetCityDB() *maxminddb.Reader {
+	g.mu.RLock()
+	defer g.mu.RUnlock()
+	return g.cityDB
+}
+
+// Safely provides read access to the ASN database
+func (g *GeoIPManager) GetASNDB() *maxminddb.Reader {
+	g.mu.RLock()
+	defer g.mu.RUnlock()
+	return g.asnDB
+}
+
+// Sets up automatic database updates
+func (g *GeoIPManager) StartUpdater(ctx context.Context, updateInterval time.Duration) {
+	log.Printf("Starting MaxMind GeoIP database updater with interval: %s", updateInterval)
+
+	ticker := time.NewTicker(updateInterval)
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				log.Println("Performing scheduled GeoIP database update")
+				if err := g.UpdateDatabases(); err != nil {
+					log.Printf("Failed to update databases: %v", err)
+				}
+			case <-ctx.Done():
+				ticker.Stop()
+				log.Println("GeoIP database updater stopped")
+				return
+			}
+		}
+	}()
+}
+
+// Downloads fresh copies of the databases and reloads them
+func (g *GeoIPManager) UpdateDatabases() error {
+	// Download new databases
+	if err := g.downloadDatabases(); err != nil {
+		return err
+	}
+
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	// Close existing databases
+	if g.cityDB != nil {
+		if err := g.cityDB.Close(); err != nil {
+			log.Printf("Warning: error closing city database: %v", err)
+		}
+	}
+
+	if g.asnDB != nil {
+		if err := g.asnDB.Close(); err != nil {
+			log.Printf("Warning: error closing ASN database: %v", err)
+		}
+	}
+
+	// Reopen databases
+	var err error
+	g.cityDB, err = maxminddb.Open(CityDBPath)
+	if err != nil {
+		return fmt.Errorf("reopening city database: %w", err)
+	}
+
+	g.asnDB, err = maxminddb.Open(ASNDBPath)
+	if err != nil {
+		return fmt.Errorf("reopening ASN database: %w", err)
+	}
+
+	log.Println("Successfully updated and reloaded GeoIP databases")
+	return nil
+}