wild-central/internal/dnsfilter/manager.go

package dnsfilter

import (
	"bufio"
	"context"
	"crypto/sha256"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/wild-cloud/wild-central/internal/storage"
	"gopkg.in/yaml.v3"
)

// Blocklist represents a subscribed blocklist source.
type Blocklist struct {
	ID          string    `yaml:"id" json:"id"`
	URL         string    `yaml:"url,omitempty" json:"url,omitempty"`
	Name        string    `yaml:"name" json:"name"`
	Enabled     bool      `yaml:"enabled" json:"enabled"`
	EntryCount  int       `yaml:"entryCount,omitempty" json:"entryCount,omitempty"`
	LastUpdated time.Time `yaml:"lastUpdated,omitempty" json:"lastUpdated,omitzero"`
	LastError   string    `yaml:"lastError,omitempty" json:"lastError,omitempty"`
}

// CustomEntry represents a user-managed domain override.
type CustomEntry struct {
	Domain string `yaml:"domain" json:"domain"`
	Type   string `yaml:"type" json:"type"` // "allow" or "block"
}

// FilterData holds all DNS filter configuration persisted in lists.yaml.
type FilterData struct {
	Lists         []Blocklist   `yaml:"lists" json:"lists"`
	CustomEntries []CustomEntry `yaml:"customEntries" json:"customEntries"`
}

// FilterStats holds current filtering statistics.
type FilterStats struct {
	Enabled      bool        `json:"enabled"`
	TotalBlocked int         `json:"totalBlocked"`
	ListCount    int         `json:"listCount"`
	EnabledLists int         `json:"enabledLists"`
	LastUpdated  time.Time   `json:"lastUpdated"`
	QueryStats   *QueryStats `json:"queryStats,omitempty"`
}

// QueryStats holds DNS query statistics from dnsmasq logs.
type QueryStats struct {
	TotalQueries   int     `json:"totalQueries"`
	BlockedQueries int     `json:"blockedQueries"`
	BlockedPercent float64 `json:"blockedPercent"`
	Period         string  `json:"period"` // e.g. "24h"
}

// Manager handles DNS filter list management and compilation.
type Manager struct {
	dataDir       string // e.g. /var/lib/wild-central/dns-filter
	hostsFilePath string // compiled output path (must be readable by dnsmasq)
}

// NewManager creates a new DNS filter manager.
func NewManager(dataDir string) *Manager {
	return &Manager{
		dataDir:       dataDir,
		hostsFilePath: filepath.Join(dataDir, "blocked.conf"), // default, override with SetHostsFilePath
	}
}

// SetHostsFilePath sets the output path for the compiled hosts file.
// Must be readable by the dnsmasq process (e.g. /etc/dnsmasq.d/).
func (m *Manager) SetHostsFilePath(path string) {
	m.hostsFilePath = path
}

func (m *Manager) listsPath() string { return filepath.Join(m.dataDir, "lists.yaml") }
func (m *Manager) cachePath() string { return filepath.Join(m.dataDir, "cache") }
func (m *Manager) lockPath() string  { return filepath.Join(m.dataDir, "lists.yaml.lock") }

// HostsFilePath returns the path to the compiled blocked hosts file.
func (m *Manager) HostsFilePath() string { return m.hostsFilePath }

// EnsureDataDir creates the dns-filter directory structure.
func (m *Manager) EnsureDataDir() error {
	if err := storage.EnsureDir(m.dataDir, 0755); err != nil {
		return err
	}
	// Also ensure the hosts file output directory exists (may differ from dataDir)
	hostsDir := filepath.Dir(m.hostsFilePath)
	if hostsDir != m.dataDir {
		if err := storage.EnsureDir(hostsDir, 0755); err != nil {
			return err
		}
	}
	return storage.EnsureDir(m.cachePath(), 0755)
}

// LoadFilterData reads lists.yaml. Returns empty data if file doesn't exist.
func (m *Manager) LoadFilterData() (*FilterData, error) {
	data, err := os.ReadFile(m.listsPath())
	if err != nil {
		if os.IsNotExist(err) {
			return &FilterData{}, nil
		}
		return nil, fmt.Errorf("reading filter data: %w", err)
	}
	var fd FilterData
	if err := yaml.Unmarshal(data, &fd); err != nil {
		return nil, fmt.Errorf("parsing filter data: %w", err)
	}
	return &fd, nil
}

// SaveFilterData writes lists.yaml atomically with file locking.
func (m *Manager) SaveFilterData(fd *FilterData) error {
	return storage.WithLock(m.lockPath(), func() error {
		data, err := yaml.Marshal(fd)
		if err != nil {
			return fmt.Errorf("marshaling filter data: %w", err)
		}
		return os.WriteFile(m.listsPath(), data, 0644)
	})
}

// listID returns a stable ID for a URL-based list.
func listID(url string) string {
	h := sha256.Sum256([]byte(url))
	return fmt.Sprintf("%x", h[:8])
}

// cacheFile returns the cache file path for a list by ID.
func (m *Manager) cacheFile(id string) string {
	return filepath.Join(m.cachePath(), id+".txt")
}

// AddList adds a URL-based blocklist subscription.
func (m *Manager) AddList(bl Blocklist) error {
	fd, err := m.LoadFilterData()
	if err != nil {
		return err
	}

	if bl.URL != "" {
		bl.ID = listID(bl.URL)
	}
	if bl.ID == "" {
		return fmt.Errorf("list must have a URL or ID")
	}

	// Deduplicate
	for _, existing := range fd.Lists {
		if existing.ID == bl.ID {
			return fmt.Errorf("list already exists: %s", bl.Name)
		}
	}

	bl.Enabled = true
	fd.Lists = append(fd.Lists, bl)
	return m.SaveFilterData(fd)
}

// AddUploadedList saves an uploaded blocklist file and adds it as a list entry.
func (m *Manager) AddUploadedList(name string, content []byte) (string, error) {
	if err := m.EnsureDataDir(); err != nil {
		return "", err
	}

	h := sha256.Sum256(content)
	id := fmt.Sprintf("upload-%x", h[:8])

	// Write content to cache
	if err := os.WriteFile(m.cacheFile(id), content, 0644); err != nil {
		return "", fmt.Errorf("writing uploaded list: %w", err)
	}

	// Count entries
	domains, _ := ParseFile(m.cacheFile(id))

	bl := Blocklist{
		ID:          id,
		Name:        name,
		Enabled:     true,
		EntryCount:  len(domains),
		LastUpdated: time.Now(),
	}

	fd, err := m.LoadFilterData()
	if err != nil {
		return "", err
	}

	// Deduplicate
	for _, existing := range fd.Lists {
		if existing.ID == id {
			return id, fmt.Errorf("identical list already uploaded")
		}
	}

	fd.Lists = append(fd.Lists, bl)
	return id, m.SaveFilterData(fd)
}

// RemoveList removes a list by ID and its cached file.
func (m *Manager) RemoveList(id string) error {
	fd, err := m.LoadFilterData()
	if err != nil {
		return err
	}

	found := false
	filtered := fd.Lists[:0]
	for _, bl := range fd.Lists {
		if bl.ID == id {
			found = true
			_ = os.Remove(m.cacheFile(id))
			continue
		}
		filtered = append(filtered, bl)
	}
	if !found {
		return fmt.Errorf("list not found: %s", id)
	}

	fd.Lists = filtered
	return m.SaveFilterData(fd)
}

// ToggleList enables or disables a list by ID.
func (m *Manager) ToggleList(id string, enabled bool) error {
	fd, err := m.LoadFilterData()
	if err != nil {
		return err
	}

	for i := range fd.Lists {
		if fd.Lists[i].ID == id {
			fd.Lists[i].Enabled = enabled
			return m.SaveFilterData(fd)
		}
	}
	return fmt.Errorf("list not found: %s", id)
}

// SetCustomEntry adds or replaces a custom allow/block entry.
func (m *Manager) SetCustomEntry(entry CustomEntry) error {
	if entry.Type != "allow" && entry.Type != "block" {
		return fmt.Errorf("type must be 'allow' or 'block'")
	}

	fd, err := m.LoadFilterData()
	if err != nil {
		return err
	}

	// Replace if exists
	for i := range fd.CustomEntries {
		if fd.CustomEntries[i].Domain == entry.Domain {
			fd.CustomEntries[i].Type = entry.Type
			return m.SaveFilterData(fd)
		}
	}

	fd.CustomEntries = append(fd.CustomEntries, entry)
	return m.SaveFilterData(fd)
}

// RemoveCustomEntry removes a custom entry by domain.
func (m *Manager) RemoveCustomEntry(domain string) error {
	fd, err := m.LoadFilterData()
	if err != nil {
		return err
	}

	found := false
	filtered := fd.CustomEntries[:0]
	for _, e := range fd.CustomEntries {
		if e.Domain == domain {
			found = true
			continue
		}
		filtered = append(filtered, e)
	}
	if !found {
		return fmt.Errorf("custom entry not found: %s", domain)
	}

	fd.CustomEntries = filtered
	return m.SaveFilterData(fd)
}

// UpdateLists downloads all enabled URL-based lists to cache.
func (m *Manager) UpdateLists(ctx context.Context) error {
	if err := m.EnsureDataDir(); err != nil {
		return err
	}

	fd, err := m.LoadFilterData()
	if err != nil {
		return err
	}

	client := &http.Client{Timeout: 60 * time.Second}
	changed := false

	for i := range fd.Lists {
		bl := &fd.Lists[i]
		if !bl.Enabled || bl.URL == "" {
			continue
		}

		slog.Info("dns-filter: downloading list", "name", bl.Name, "url", bl.URL)

		req, err := http.NewRequestWithContext(ctx, http.MethodGet, bl.URL, nil)
		if err != nil {
			bl.LastError = err.Error()
			changed = true
			continue
		}
		req.Header.Set("User-Agent", "WildCentral/1.0")

		resp, err := client.Do(req)
		if err != nil {
			bl.LastError = fmt.Sprintf("download failed: %v", err)
			changed = true
			slog.Warn("dns-filter: download failed", "name", bl.Name, "error", err)
			continue
		}

		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
			resp.Body.Close()
			bl.LastError = fmt.Sprintf("HTTP %d", resp.StatusCode)
			changed = true
			slog.Warn("dns-filter: download failed", "name", bl.Name, "status", resp.StatusCode)
			continue
		}

		// Limit download to 50MB
		limited := io.LimitReader(resp.Body, 50*1024*1024)
		data, err := io.ReadAll(limited)
		resp.Body.Close()
		if err != nil {
			bl.LastError = fmt.Sprintf("read failed: %v", err)
			changed = true
			continue
		}

		if err := os.WriteFile(m.cacheFile(bl.ID), data, 0644); err != nil {
			bl.LastError = fmt.Sprintf("write failed: %v", err)
			changed = true
			continue
		}

		bl.LastUpdated = time.Now()
		bl.LastError = ""
		changed = true
	}

	if changed {
		_ = m.SaveFilterData(fd)
	}

	return nil
}

// Compile merges all enabled cached lists, applies custom entries,
// and writes the compiled hosts.blocked file. Returns domain count.
func (m *Manager) Compile() (int, error) {
	fd, err := m.LoadFilterData()
	if err != nil {
		return 0, err
	}

	blocked := make(map[string]struct{})

	// Parse all enabled lists
	for i := range fd.Lists {
		bl := &fd.Lists[i]
		if !bl.Enabled {
			continue
		}

		cachePath := m.cacheFile(bl.ID)
		if !storage.FileExists(cachePath) {
			continue
		}

		domains, err := ParseFile(cachePath)
		if err != nil {
			slog.Warn("dns-filter: failed to parse list", "name", bl.Name, "error", err)
			continue
		}

		for d := range domains {
			blocked[d] = struct{}{}
		}
		bl.EntryCount = len(domains)
	}

	// Apply custom entries
	for _, ce := range fd.CustomEntries {
		switch ce.Type {
		case "allow":
			delete(blocked, ce.Domain)
		case "block":
			blocked[ce.Domain] = struct{}{}
		}
	}

	// Write hosts.blocked atomically
	tmpPath := m.HostsFilePath() + ".tmp"
	f, err := os.Create(tmpPath)
	if err != nil {
		return 0, fmt.Errorf("creating hosts file: %w", err)
	}

	// Sort for deterministic output
	sorted := make([]string, 0, len(blocked))
	for d := range blocked {
		sorted = append(sorted, d)
	}
	sort.Strings(sorted)

	w := bufio.NewWriter(f)
	fmt.Fprintln(w, "# Wild Central DNS Filter — auto-generated, do not edit")
	for _, d := range sorted {
		fmt.Fprintf(w, "address=/%s/\n", d)
	}
	if err := w.Flush(); err != nil {
		f.Close()
		os.Remove(tmpPath)
		return 0, fmt.Errorf("writing hosts file: %w", err)
	}
	f.Close()

	if err := os.Rename(tmpPath, m.HostsFilePath()); err != nil {
		return 0, fmt.Errorf("renaming hosts file: %w", err)
	}

	// Update entry counts and save
	_ = m.SaveFilterData(fd)

	slog.Info("dns-filter: compiled blocklist", "domains", len(sorted))
	return len(sorted), nil
}

// GetStats returns current filtering statistics.
func (m *Manager) GetStats() FilterStats {
	fd, err := m.LoadFilterData()
	if err != nil {
		return FilterStats{}
	}

	var enabled int
	var lastUpdated time.Time
	for _, bl := range fd.Lists {
		if bl.Enabled {
			enabled++
		}
		if bl.LastUpdated.After(lastUpdated) {
			lastUpdated = bl.LastUpdated
		}
	}

	// Count lines in hosts.blocked for total
	total := 0
	if f, err := os.Open(m.HostsFilePath()); err == nil {
		scanner := bufio.NewScanner(f)
		for scanner.Scan() {
			total++
		}
		_ = scanner.Err()
		f.Close()
	}

	return FilterStats{
		TotalBlocked: total,
		ListCount:    len(fd.Lists),
		EnabledLists: enabled,
		LastUpdated:  lastUpdated,
	}
}

// GetQueryStats parses dnsmasq journal logs to count total and blocked queries.
// With address=/ directives, blocked queries log as "config <domain> is NXDOMAIN"
// (for address=/domain/) or "config <domain> is 0.0.0.0".
func (m *Manager) GetQueryStats() *QueryStats {
	// Use journalctl to read dnsmasq logs from the last 24 hours.
	// Run via sudo since the wildcloud user may not be in the systemd-journal group.
	cmd := exec.Command("sudo", "journalctl", "-u", "dnsmasq", "--since", "24 hours ago", "--no-pager", "-o", "cat")
	output, err := cmd.Output()
	if err != nil {
		slog.Debug("dns-filter: failed to read dnsmasq journal", "error", err)
		return nil
	}

	var total, blocked int
	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "query[") {
			total++
		}
		// dnsmasq logs address=/ blocks as: "config <domain> is NXDOMAIN"
		// and addn-hosts blocks as: "/path/file <domain> is 0.0.0.0"
		if strings.HasPrefix(line, "config ") && strings.Contains(line, " is NXDOMAIN") {
			blocked++
		}
	}

	var pct float64
	if total > 0 {
		pct = float64(blocked) / float64(total) * 100
	}

	return &QueryStats{
		TotalQueries:   total,
		BlockedQueries: blocked,
		BlockedPercent: pct,
		Period:         "24h",
	}
}

// SuggestedLists returns well-known blocklists for easy adding.
// Prefers dnsmasq-format lists which support wildcard subdomain blocking.
func SuggestedLists() []Blocklist {
	type entry struct {
		url  string
		name string
	}
	lists := []entry{
		{"https://raw.githubusercontent.com/StevenBlack/hosts/master/hosts", "Steven Black - Unified Hosts"},
		{"https://cdn.jsdelivr.net/gh/hagezi/dns-blocklists@latest/dnsmasq/light.txt", "Hagezi - Light"},
		{"https://cdn.jsdelivr.net/gh/hagezi/dns-blocklists@latest/dnsmasq/pro.txt", "Hagezi - Pro"},
		{"https://cdn.jsdelivr.net/gh/hagezi/dns-blocklists@latest/dnsmasq/ultimate.txt", "Hagezi - Ultimate"},
		{"https://big.oisd.nl/dnsmasq2", "OISD - Big"},
		{"https://small.oisd.nl/dnsmasq2", "OISD - Small"},
	}
	result := make([]Blocklist, len(lists))
	for i, l := range lists {
		result[i] = Blocklist{ID: listID(l.url), URL: l.url, Name: l.name}
	}
	return result
}