unifi-adblocker/handlers.go

package main

import (
	"net/http"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

func profileHandler(w http.ResponseWriter, r *http.Request) {
	configMu.RLock()
	mfaEnabled := config.MFASecret != ""
	configMu.RUnlock()

	data := struct {
		MFAEnabled bool
	}{
		MFAEnabled: mfaEnabled,
	}
	renderTemplate(w, r, "profile.html", data)
}

func dashboardHandler(w http.ResponseWriter, r *http.Request) {
	renderTemplate(w, r, "dashboard.html", nil)
}

func urlListsHandler(w http.ResponseWriter, r *http.Request) {
	configMu.RLock()
	urlFileList := config.URLFileList
	defaultURLs := config.DefaultURLs
	configMu.RUnlock()

	if r.Method == "POST" {
		action := r.FormValue("action")

		items, err := readURLListCSV(urlFileList)
		if err != nil {
			renderTemplate(w, r, "urllists.html", struct {
				Items            []URLListItem
				Notification     string
				NotificationType string
			}{
				Items:            items,
				Notification:     "Error reading URL list",
				NotificationType: "error",
			})
			return
		}

		var errorMsg string

		switch action {
		case "add":
			name := sanitizeInput(r.FormValue("name"))
			urlStr := sanitizeInput(r.FormValue("url"))
			group := sanitizeInput(r.FormValue("group"))
			note := sanitizeInput(r.FormValue("note"))

			if name == "" {
				errorMsg = "Name is required"
			} else if !isValidURL(urlStr) {
				errorMsg = "Invalid URL format. Must start with http:// or https://"
			} else {
				items = append(items, URLListItem{
					Name:    name,
					Enabled: true,
					Group:   group,
					URL:     urlStr,
					Note:    note,
				})
			}

		case "remove":
			indexStr := r.FormValue("index")
			index, _ := strconv.Atoi(indexStr)
			if index >= 0 && index < len(items) {
				items = append(items[:index], items[index+1:]...)
			} else {
				errorMsg = "Invalid index"
			}

		case "toggle":
			indexStr := r.FormValue("index")
			index, _ := strconv.Atoi(indexStr)
			if index >= 0 && index < len(items) {
				items[index].Enabled = !items[index].Enabled
			} else {
				errorMsg = "Invalid index"
			}

		case "edit":
			indexStr := r.FormValue("index")
			name := sanitizeInput(r.FormValue("name"))
			urlStr := sanitizeInput(r.FormValue("url"))
			group := sanitizeInput(r.FormValue("group"))
			note := sanitizeInput(r.FormValue("note"))

			index, _ := strconv.Atoi(indexStr)
			if index >= 0 && index < len(items) {
				if name == "" {
					errorMsg = "Name is required"
				} else if !isValidURL(urlStr) {
					errorMsg = "Invalid URL format. Must start with http:// or https://"
				} else {
					items[index].Name = name
					items[index].URL = urlStr
					items[index].Group = group
					items[index].Note = note
				}
			} else {
				errorMsg = "Invalid index"
			}
		}

		if errorMsg != "" {
			renderTemplate(w, r, "urllists.html", struct {
				Items            []URLListItem
				Notification     string
				NotificationType string
			}{
				Items:            items,
				Notification:     errorMsg,
				NotificationType: "error",
			})
			return
		}

		err = writeURLListCSV(urlFileList, items)
		if err != nil {
			renderTemplate(w, r, "urllists.html", struct {
				Items            []URLListItem
				Notification     string
				NotificationType string
			}{
				Items:            items,
				Notification:     "Error saving URL list",
				NotificationType: "error",
			})
			return
		}

		renderTemplate(w, r, "urllists.html", struct {
			Items            []URLListItem
			Notification     string
			NotificationType string
		}{
			Items:            items,
			Notification:     "URL list updated successfully",
			NotificationType: "success",
		})
		return
	}

	items, err := readURLListCSV(urlFileList)
	if err != nil || len(items) == 0 {
		// Use default blocklists
		defaultItems := defaultURLs
		writeURLListCSV(urlFileList, defaultItems)
		items = defaultItems
	}

	renderTemplate(w, r, "urllists.html", struct {
		Items []URLListItem
	}{
		Items: items,
	})
}

func domainsHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		action := r.FormValue("action")
		domain := sanitizeInput(r.FormValue("domain"))

		configMu.RLock()
		blocklistFile := config.BlocklistFile
		configMu.RUnlock()

		domains, err := readLines(blocklistFile)
		if err != nil {
			renderTemplate(w, r, "domains.html", struct {
				Domains          []string
				Query            string
				Notification     string
				NotificationType string
			}{
				Domains:          domains,
				Query:            "",
				Notification:     "Error reading domains",
				NotificationType: "error",
			})
			return
		}

		errorMsg := ""
		switch action {
		case "add":
			if !isValidDomain(domain) {
				errorMsg = "Invalid domain format"
			} else {
				domains = append(domains, domain)
				sort.Strings(domains)
			}
		case "remove":
			for i, d := range domains {
				if d == domain {
					domains = append(domains[:i], domains[i+1:]...)
					break
				}
			}
		}

		if errorMsg != "" {
			renderTemplate(w, r, "domains.html", struct {
				Domains          []string
				Query            string
				Notification     string
				NotificationType string
			}{
				Domains:          domains,
				Query:            "",
				Notification:     errorMsg,
				NotificationType: "error",
			})
			return
		}

		writeLines(blocklistFile, domains)

		renderTemplate(w, r, "domains.html", struct {
			Domains          []string
			Query            string
			Notification     string
			NotificationType string
		}{
			Domains:          domains,
			Query:            "",
			Notification:     "Domain updated successfully",
			NotificationType: "success",
		})
		return
	}

	configMu.RLock()
	blocklistFile := config.BlocklistFile
	configMu.RUnlock()

	query := r.URL.Query().Get("query")
	domains, _ := readLines(blocklistFile)

	var filtered []string
	if query != "" {
		query = strings.ReplaceAll(query, "*", ".*")
		re, err := regexp.Compile("(?i)" + query)
		if err == nil {
			for _, d := range domains {
				if re.MatchString(d) {
					filtered = append(filtered, d)
				}
			}
		}
	} else {
		filtered = domains
	}

	data := struct {
		Domains []string
		Query   string
	}{
		Domains: filtered,
		Query:   query,
	}
	renderTemplate(w, r, "domains.html", data)
}

func applyHandler(w http.ResponseWriter, r *http.Request) {
	userInitiatedMu.Lock()
	userInitiated = true
	userInitiatedMu.Unlock()

	updateBlocklist(1) // Force

	userInitiatedMu.Lock()
	userInitiated = false
	userInitiatedMu.Unlock()

	http.Redirect(w, r, "/", http.StatusSeeOther)
}

func logsHandler(w http.ResponseWriter, r *http.Request) {
	// For simplicity, assume logs are in stdout, or implement file logging
	w.Header().Set("Content-Type", "text/plain")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Logs placeholder"))
}