unifi-adblocker/blocklist.go

package main

import (
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
)

func updateBlocklist(force int) {
	configMu.RLock()
	delayDays := config.UpdateDelay
	configMu.RUnlock()

	if checkIfUpdateNeeded(delayDays, force) {
		createDefaultURLList()
		fetchAndMergeSources()
		applyRemovalRules()
		finalizeBlocklist()
	} else {
		handleSkippedUpdate(delayDays)
	}
}

func checkIfUpdateNeeded(delayDays, force int) bool {
	configMu.RLock()
	lastUpdateFile := config.LastUpdateFolder + "blocklist_Default_lastupdate.txt"
	configMu.RUnlock()

	data, err := os.ReadFile(lastUpdateFile)
	if err != nil {
		return true
	}
	lastTS, _ := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64)
	nowTS := time.Now().Unix()

	delaySeconds := int64(delayDays * 86400)

	if force == 0 && nowTS-lastTS < delaySeconds {
		return false
	}
	return true
}

func handleSkippedUpdate(delayDays int) {
	configMu.RLock()
	blocklistFile := ""
	if len(config.Filters) > 0 {
		blocklistFile = config.DomainListFolder + config.Filters[0].BlocklistFile
	}
	tmpFile := config.TmpFolder + "blocklist_Default.tmp"
	configMu.RUnlock()

	log.Printf("%s: Full Update skipped — last update was less than %d days ago. Use force to override.", time.Now(), delayDays)

	blockData, _ := os.ReadFile(blocklistFile)
	tmpData, _ := os.ReadFile(tmpFile)

	if string(blockData) == string(tmpData) {
		log.Printf("%s: Custom Blocklist still in use, doing nothing.", time.Now())
		return
	}

	time.Sleep(10 * time.Second)
	log.Printf("%s: Copying existing blocklist to %s", time.Now(), blocklistFile)

	merged := mergeUnique(string(blockData), string(tmpData))
	os.WriteFile(blocklistFile, []byte(merged), 0644)

	log.Printf("%s: Restarting CoreDNS to apply new blocklist file.", time.Now())
	exec.Command("pkill", "coredns").Run()
}

func createDefaultURLList() {
	configMu.RLock()
	var blocklistURLs []URLListItem
	if len(config.Filters) > 0 {
		blocklistURLs = config.Filters[0].BlocklistURLs
	}
	configMu.RUnlock()

	urlFileList := "blocklist_Default_urls.csv"
	_, err := os.Stat(urlFileList)
	if os.IsNotExist(err) {
		log.Printf("%s: Using Default url list: %s", time.Now(), urlFileList)
		writeURLListCSV(urlFileList, blocklistURLs)
	}
}

func fetchAndMergeSources() {
	configMu.RLock()
	tmpFile := config.TmpFolder + "blocklist_Default.tmp"
	blocklistFile := ""
	if len(config.Filters) > 0 {
		blocklistFile = config.DomainListFolder + config.Filters[0].BlocklistFile
	}
	urlFileList := "blocklist_Default_urls.csv"
	configMu.RUnlock()

	os.Create(tmpFile)

	blockData, _ := os.ReadFile(blocklistFile)
	var builder strings.Builder
	builder.Write(blockData)

	items, _ := readURLListCSV(urlFileList)
	for _, item := range items {
		if !item.Enabled {
			continue
		}
		resp, err := http.Get(item.URL)
		if err != nil {
			log.Printf("%s: Failed to fetch %s: %v", time.Now(), item.Name, err)
			continue
		}
		body, _ := io.ReadAll(resp.Body)
		resp.Body.Close()

		lines := strings.Split(string(body), "\n")
		for _, line := range lines {
			line = strings.TrimSpace(line)
			if line == "" || strings.HasPrefix(line, "!") || strings.HasPrefix(line, "#") {
				continue
			}
			line = strings.TrimPrefix(line, "||")
			line = strings.TrimSuffix(line, "^")
			builder.WriteString(line + "\n")
		}
	}

	lines := strings.Split(builder.String(), "\n")
	sort.Strings(lines)
	unique := []string{}
	seen := make(map[string]bool)
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line == "" || seen[line] {
			continue
		}
		if !isValidDomain(line) {
			continue
		}
		unique = append(unique, line)
		seen[line] = true
	}

	os.WriteFile(tmpFile, []byte(strings.Join(unique, "\n")), 0644)
	log.Printf("%s: Combined List cleanup.", time.Now())
}

func applyRemovalRules() {
	configMu.RLock()
	removeFile := ""
	if len(config.Filters) > 0 {
		removeFile = config.DomainListFolder + config.Filters[0].WhitelistFile
	}
	tmpFile := config.TmpFolder + "blocklist_Default.tmp"
	configMu.RUnlock()

	removeLines, _ := readLines(removeFile)
	var patterns []*regexp.Regexp
	for _, rule := range removeLines {
		rule = strings.TrimSpace(rule)
		if rule == "" {
			continue
		}
		rule = regexp.QuoteMeta(rule)
		rule = strings.ReplaceAll(rule, "\\*", ".*")
		re, _ := regexp.Compile("^" + rule + "$")
		patterns = append(patterns, re)
	}

	data, _ := os.ReadFile(tmpFile)
	lines := strings.Split(string(data), "\n")
	var filtered []string
	for _, line := range lines {
		drop := false
		for _, re := range patterns {
			if re.MatchString(line) {
				drop = true
				break
			}
		}
		if !drop {
			filtered = append(filtered, line)
		}
	}

	os.WriteFile(tmpFile+".filtered", []byte(strings.Join(filtered, "\n")), 0644)
}

func finalizeBlocklist() {
	configMu.RLock()
	tmpFile := config.TmpFolder + "blocklist_Default.tmp"
	blocklistFile := ""
	if len(config.Filters) > 0 {
		blocklistFile = config.DomainListFolder + config.Filters[0].BlocklistFile
	}
	lastUpdateFile := config.LastUpdateFolder + "blocklist_Default_lastupdate.txt"
	configMu.RUnlock()

	os.Rename(tmpFile+".filtered", tmpFile)
	copyFile(tmpFile, blocklistFile)

	now := time.Now().Unix()
	os.WriteFile(lastUpdateFile, []byte(strconv.FormatInt(now, 10)), 0644)

	log.Printf("%s: Restarting CoreDNS to apply new blocklist file", time.Now())
	exec.Command("pkill", "coredns").Run()

	stat, _ := os.Stat(blocklistFile)
	size := stat.Size()
	lines, _ := readLines(blocklistFile)
	log.Printf("%s: Blocklist created at: %s with size %d and %d lines", time.Now(), blocklistFile, size, len(lines))
}