main.go

package main

import (
	"context"
	"embed"
	"fmt"
	"io/fs"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"bandwidth-monitor/adguard"
	"bandwidth-monitor/collector"
	"bandwidth-monitor/conntrack"
	"bandwidth-monitor/dns"
	"bandwidth-monitor/geoip"
	"bandwidth-monitor/handler"
	"bandwidth-monitor/latency"
	"bandwidth-monitor/liveactivity"
	"bandwidth-monitor/nextdns"
	"bandwidth-monitor/omada"
	"bandwidth-monitor/pihole"
	"bandwidth-monitor/resolver"
	"bandwidth-monitor/speedtest"
	"bandwidth-monitor/talkers"
	"bandwidth-monitor/topology"
	"bandwidth-monitor/unifi"
	"bandwidth-monitor/version"
	"bandwidth-monitor/webassets"
	"bandwidth-monitor/wifi"
)

//go:embed static/*
var staticFiles embed.FS

// env returns the value of the environment variable named by key,
// or fallback if the variable is empty/unset.
func env(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func collectorInterval() (time.Duration, error) {
	const minimum = 100 * time.Millisecond
	raw := env("COLLECTOR_INTERVAL", "1s")
	interval, err := time.ParseDuration(raw)
	if err != nil {
		return 0, fmt.Errorf("invalid duration %q: %w", raw, err)
	}
	if interval < minimum {
		return 0, fmt.Errorf("must be at least %s", minimum)
	}
	return interval, nil
}

func main() {
	listenAddr := env("LISTEN", ":8080")
	listenProto := strings.ToLower(strings.TrimSpace(env("LISTEN_PROTOCOL", "http")))
	tlsCertFile := env("TLS_CERT_FILE", "")
	tlsKeyFile := env("TLS_KEY_FILE", "")
	promiscuous := env("PROMISCUOUS", "true")
	promiscuousBool, _ := strconv.ParseBool(promiscuous)
	debugHTTPLog, _ := strconv.ParseBool(env("DEBUG_HTTP_LOG", "false"))
	statsInterval, err := collectorInterval()
	if err != nil {
		log.Fatalf("COLLECTOR_INTERVAL: %v", err)
	}

	if listenProto != "http" && listenProto != "https" {
		log.Fatalf("LISTEN_PROTOCOL: invalid value %q (expected http or https)", listenProto)
	}
	if listenProto == "https" {
		if tlsCertFile == "" || tlsKeyFile == "" {
			log.Fatal("LISTEN_PROTOCOL=https requires TLS_CERT_FILE and TLS_KEY_FILE")
		}
	}

	// Parse LOCAL_NETS: comma-separated CIDRs for SPAN port direction detection
	// e.g. LOCAL_NETS=192.0.2.0/24,2001:db8::/48
	// If not set, auto-discovers from local interface addresses.
	var localNets []*net.IPNet
	if raw := os.Getenv("LOCAL_NETS"); raw != "" {
		for _, cidr := range strings.Split(raw, ",") {
			cidr = strings.TrimSpace(cidr)
			if cidr == "" {
				continue
			}
			_, ipnet, err := net.ParseCIDR(cidr)
			if err != nil {
				log.Printf("LOCAL_NETS: invalid CIDR %q: %v", cidr, err)
				continue
			}
			localNets = append(localNets, ipnet)
		}
		log.Printf("LOCAL_NETS: %d network(s) from configuration", len(localNets))
	} else {
		// Auto-discover from local interfaces
		ifaces, err := net.Interfaces()
		if err == nil {
			for _, iface := range ifaces {
				if iface.Flags&net.FlagLoopback != 0 {
					continue
				}
				addrs, err := iface.Addrs()
				if err != nil {
					continue
				}
				for _, addr := range addrs {
					ipnet, ok := addr.(*net.IPNet)
					if !ok {
						continue
					}
					// Skip link-local
					if ipnet.IP.IsLinkLocalUnicast() || ipnet.IP.IsLinkLocalMulticast() {
						continue
					}
					localNets = append(localNets, ipnet)
				}
			}
		}
		if len(localNets) > 0 {
			log.Printf("LOCAL_NETS: auto-discovered %d network(s) from interfaces", len(localNets))
			for _, n := range localNets {
				log.Printf("  %s", n.String())
			}
		}
	}

	// GeoIP: prefer City DB (has country + city + coordinates), fall back to Country DB.
	// Check env vars first, then auto-detect from common filenames on disk.
	// Always verify the file exists — stale env vars may point to missing files.
	geoCity := env("GEO_CITY", env("GEO_COUNTRY", ""))
	if geoCity != "" {
		if _, err := os.Stat(geoCity); err != nil {
			log.Printf("GeoIP: configured path %q not found, auto-detecting", geoCity)
			geoCity = ""
		}
	}
	if geoCity == "" {
		// Auto-detect: try City first, fall back to Country
		for _, candidate := range []string{"GeoLite2-City.mmdb", "GeoLite2-Country.mmdb"} {
			if _, err := os.Stat(candidate); err == nil {
				geoCity = candidate
				break
			}
		}
	}
	geoASN := env("GEO_ASN", "GeoLite2-ASN.mmdb")
	adguardURL := env("ADGUARD_URL", "")
	adguardUser := env("ADGUARD_USER", "")
	adguardPass := env("ADGUARD_PASS", "")
	nextdnsProfile := env("NEXTDNS_PROFILE", "")
	nextdnsAPIKey := env("NEXTDNS_API_KEY", "")
	piholeURL := env("PIHOLE_URL", "")
	piholePass := env("PIHOLE_PASSWORD", "")
	unifiURL := env("UNIFI_URL", "")
	unifiUser := env("UNIFI_USER", "")
	unifiPass := env("UNIFI_PASS", "")
	unifiSite := env("UNIFI_SITE", "default")
	omadaURL := env("OMADA_URL", "")
	omadaUser := env("OMADA_USER", "")
	omadaPass := env("OMADA_PASS", "")
	omadaSite := env("OMADA_SITE", "Default")

	geoDB, err := geoip.Open(geoCity, geoASN)
	if err != nil {
		log.Printf("GeoIP: %v (continuing without geo)", err)
		geoDB = nil
	} else if geoDB.Available() {
		log.Println("GeoIP databases loaded")
		defer geoDB.Close()
	} else {
		log.Println("GeoIP: no MMDB files found (continuing without geo)")
	}

	// Parse VPN_STATUS_FILES: comma-separated "iface=path" pairs
	// e.g. VPN_STATUS_FILES=myvpn=/run/myvpn-active,wg0=/run/wg0-active
	vpnStatusFiles := make(map[string]string)
	if raw := os.Getenv("VPN_STATUS_FILES"); raw != "" {
		for _, entry := range strings.Split(raw, ",") {
			parts := strings.SplitN(strings.TrimSpace(entry), "=", 2)
			if len(parts) == 2 {
				vpnStatusFiles[parts[0]] = parts[1]
			}
		}
	}

	// Parse INTERFACES: comma-separated list of interface names to display.
	// If not set, all interfaces are shown.
	var allowedIfaces []string
	if raw := os.Getenv("INTERFACES"); raw != "" {
		for _, name := range strings.Split(raw, ",") {
			name = strings.TrimSpace(name)
			if name != "" {
				allowedIfaces = append(allowedIfaces, name)
			}
		}
		log.Printf("INTERFACES: showing %d interface(s): %s", len(allowedIfaces), strings.Join(allowedIfaces, ", "))
	}

	// Parse WAN_INTERFACE: comma-separated list of interface names to treat as WAN.
	// If not set, WAN is auto-detected (public IP, PPP, or default route).
	var wanIfaces []string
	if raw := os.Getenv("WAN_INTERFACE"); raw != "" {
		for _, name := range strings.Split(raw, ",") {
			name = strings.TrimSpace(name)
			if name != "" {
				wanIfaces = append(wanIfaces, name)
			}
		}
		log.Printf("WAN_INTERFACE: %s", strings.Join(wanIfaces, ", "))
	}

	statsCollector := collector.NewWithInterval(vpnStatusFiles, allowedIfaces, wanIfaces, statsInterval)

	// Shared reverse-DNS resolver — used by talkers, conntrack, and debug.
	dnsResolver := resolver.New()

	// SPAN/mirror port mode: override RX/TX direction on a specific interface
	// using pcap-based packet inspection against LOCAL_NETS.
	spanDevice := env("SPAN_DEVICE", "")
	if spanDevice != "" && len(localNets) > 0 {
		statsCollector.EnableSPAN(spanDevice, promiscuousBool, localNets)
		log.Printf("span: enabled on %s (%d local nets)", spanDevice, len(localNets))
	} else if spanDevice != "" && len(localNets) == 0 {
		log.Printf("span: SPAN_DEVICE=%s set but LOCAL_NETS is empty — disabled", spanDevice)
	}

	go statsCollector.Run()

	talkerTracker := talkers.New(allowedIfaces, promiscuousBool, localNets, geoDB, dnsResolver)
	go talkerTracker.Run()

	// DNS provider: AdGuard Home, NextDNS, or Pi-hole (mutually exclusive; first configured wins)
	var dnsProvider dns.Provider
	if adguardURL != "" {
		ac := adguard.New(adguardURL, adguardUser, adguardPass, 10*time.Second)
		go ac.Run()
		dnsProvider = ac
		log.Printf("DNS integration: AdGuard Home (%s)", adguardURL)
	} else if nextdnsProfile != "" && nextdnsAPIKey != "" {
		nc := nextdns.New(nextdnsProfile, nextdnsAPIKey, 30*time.Second)
		go nc.Run()
		dnsProvider = nc
		log.Printf("DNS integration: NextDNS (profile %s)", nextdnsProfile)
	} else if piholeURL != "" {
		pc := pihole.New(piholeURL, piholePass, 10*time.Second)
		go pc.Run()
		dnsProvider = pc
		log.Printf("DNS integration: Pi-hole (%s)", piholeURL)
	}

	// WiFi provider: UniFi or Omada (mutually exclusive; first configured wins)
	var wifiProvider wifi.Provider
	if unifiURL != "" {
		uc := unifi.New(unifiURL, unifiUser, unifiPass, unifiSite, 15*time.Second)
		go uc.Run()
		wifiProvider = uc
		log.Printf("WiFi integration: UniFi (%s)", unifiURL)
	} else if omadaURL != "" {
		oc := omada.New(omadaURL, omadaUser, omadaPass, omadaSite, 15*time.Second)
		go oc.Run()
		wifiProvider = oc
		log.Printf("WiFi integration: Omada (%s)", omadaURL)
	}

	conntrackTracker := conntrack.New(localNets, geoDB, dnsResolver)
	go conntrackTracker.Run()
	log.Println("Conntrack (NAT) tracking enabled")

	speedtestServer := env("SPEEDTEST_SERVER", "https://speed.ffmuc.net")
	speedTester := speedtest.New(speedtestServer)
	log.Printf("Speed test server: %s", speedtestServer)

	// Latency monitor: continuous ICMP + HTTPS probes.
	latencyMonitoring, err := strconv.ParseBool(env("LATENCY_MONITORING", "true"))
	if err != nil {
		log.Fatalf("LATENCY_MONITORING: invalid value %q (expected true or false)", os.Getenv("LATENCY_MONITORING"))
	}
	var latencyMonitor *latency.Monitor
	if latencyMonitoring {
		// LATENCY_TARGETS: comma-separated hostnames/IPs. Uses defaults when unset.
		var latencyTargets []string
		if raw := os.Getenv("LATENCY_TARGETS"); raw != "" {
			for _, t := range strings.Split(raw, ",") {
				t = strings.TrimSpace(t)
				if t != "" {
					latencyTargets = append(latencyTargets, t)
				}
			}
		}
		latencyMonitor = latency.New(latencyTargets)
		go latencyMonitor.Run()
	} else {
		log.Println("Latency monitoring disabled")
	}

	topoScanner := topology.New(dnsResolver, wifiProvider, localNets, 30*time.Second)
	topoScanner.SetWANInterfacesFunc(func() []string {
		var names []string
		for _, iface := range statsCollector.GetAll() {
			if iface.WAN {
				names = append(names, iface.Name)
			}
		}
		return names
	})
	go topoScanner.Run()
	log.Println("Network topology scanner enabled")

	// Live Activity push (optional): when an APNs auth key is configured, the server pushes iOS
	// Live Activity updates so the Lock Screen view stays live while the app is suspended.
	var liveActivityMgr *liveactivity.Manager
	if keyFile := env("APNS_KEY_FILE", ""); keyFile != "" {
		interval, _ := time.ParseDuration(env("APNS_PUSH_INTERVAL", "10s"))
		apnsEnv := env("APNS_ENV", "production")
		mgr, err := liveactivity.New(liveactivity.Config{
			KeyFile:  keyFile,
			KeyID:    env("APNS_KEY_ID", ""),
			TeamID:   env("APNS_TEAM_ID", ""),
			BundleID: env("APNS_BUNDLE_ID", ""),
			Env:      apnsEnv,
			Interval: interval,
		}, statsCollector)
		if err != nil {
			log.Printf("Live Activity push: disabled (%v)", err)
		} else {
			liveActivityMgr = mgr
			go liveActivityMgr.Run()
			log.Printf("Live Activity push: enabled — key=%s team=%s bundle=%s env=%s interval=%s (server clock: %s)",
				env("APNS_KEY_ID", ""), env("APNS_TEAM_ID", ""), env("APNS_BUNDLE_ID", ""),
				apnsEnv, interval, time.Now().UTC().Format(time.RFC3339))
		}
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/api/interfaces", handler.InterfaceStats(statsCollector))
	mux.HandleFunc("/api/interfaces/history", handler.InterfaceHistory(statsCollector))
	mux.HandleFunc("/api/talkers/bandwidth", handler.TopTalkersBandwidth(talkerTracker))
	mux.HandleFunc("/api/talkers/volume", handler.TopTalkersVolume(talkerTracker))
	mux.HandleFunc("/api/clients/bandwidth", handler.TopClientsBandwidth(talkerTracker))
	mux.HandleFunc("/api/clients/volume", handler.TopClientsVolume(talkerTracker))
	mux.HandleFunc("/api/talkers/country", handler.CountryTalkers(talkerTracker))
	mux.HandleFunc("/api/talkers/asn", handler.ASNTalkers(talkerTracker))
	mux.HandleFunc("/api/dns", handler.DNSSummary(dnsProvider, dnsResolver))
	mux.HandleFunc("/api/wifi", handler.WiFiSummary(wifiProvider))
	mux.HandleFunc("/api/conntrack", handler.ConntrackSummary(conntrackTracker))
	mux.HandleFunc("/api/host", handler.HostDetail(talkerTracker, conntrackTracker, geoDB))
	mux.HandleFunc("/api/host/dns", handler.HostDNSLog(dnsProvider))
	mux.HandleFunc("/api/latency", handler.LatencyStatus(latencyMonitor))
	mux.HandleFunc("/api/speedtest/run", handler.SpeedTestRun(speedTester, statsCollector))
	mux.HandleFunc("/api/speedtest/results", handler.SpeedTestResults(speedTester))
	mux.HandleFunc("/api/speedtest/interfaces", handler.SpeedTestInterfaces(statsCollector))
	mux.HandleFunc("/api/debug/traceroute", handler.DebugTraceroute(dnsResolver, statsCollector))
	mux.HandleFunc("/api/debug/dns", handler.DebugDNS())
	mux.HandleFunc("/api/debug/mtu", handler.DebugMTU(statsCollector))
	mux.HandleFunc("/api/debug/publicip", handler.DebugPublicIP(statsCollector))
	mux.HandleFunc("/api/debug/tcpcheck", handler.DebugTCPCheck(statsCollector))
	mux.HandleFunc("/api/summary", handler.MenuBarSummary(statsCollector, talkerTracker, dnsProvider, wifiProvider, conntrackTracker))
	mux.HandleFunc("/api/topology", handler.TopologySummary(topoScanner))
	mux.HandleFunc("/api/events", handler.SSE(statsCollector, talkerTracker, dnsProvider, wifiProvider, latencyMonitor, topoScanner, dnsResolver, geoDB))
	if liveActivityMgr != nil {
		mux.HandleFunc("/api/liveactivity/register", handler.LiveActivityRegister(liveActivityMgr))
	}
	staticSub, err := fs.Sub(staticFiles, "static")
	if err != nil {
		log.Fatalf("Failed to create sub filesystem: %v", err)
	}

	// Compute content hashes for cache-busting query strings.
	// embed.FS has a zero modtime so http.FileServer omits Last-Modified;
	// Safari caches the response with heuristic expiration and never
	// revalidates — even on Cmd+R.  Injecting ?v=<hash> into the HTML
	// forces the browser to fetch fresh assets after every build.
	indexHTML, err := webassets.BuildIndexHTML(staticSub, version.String())
	if err != nil {
		log.Fatalf("Build embedded index: %v", err)
	}

	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" && r.URL.Path != "/index.html" {
			// Serve other static files with ETag-based caching.
			w.Header().Set("Cache-Control", "no-cache")
			http.FileServer(http.FS(staticSub)).ServeHTTP(w, r)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Header().Set("Cache-Control", "no-store")
		w.Write(indexHTML)
	})

	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

	log.Printf("server: bandwidth-monitor %s starting on %s (%s)", version.String(), listenAddr, strings.ToUpper(listenProto))
	if strings.HasPrefix(listenAddr, ":") {
		log.Printf("server: open %s://localhost%s in your browser", listenProto, listenAddr)
	} else {
		log.Printf("server: open %s://%s in your browser", listenProto, listenAddr)
	}
	if listenProto == "https" {
		log.Printf("server: TLS enabled cert=%s key=%s", tlsCertFile, tlsKeyFile)
	}
	var handler http.Handler = mux
	if debugHTTPLog {
		handler = withRequestLog(handler)
		log.Printf("server: DEBUG_HTTP_LOG enabled — logging every request to stdout")
	}
	handler = withSignature(handler)

	srv := &http.Server{
		Addr:              listenAddr,
		Handler:           handler,
		ReadHeaderTimeout: 10 * time.Second,
		IdleTimeout:       120 * time.Second,
	}

	go func() {
		<-sigCh
		fmt.Println("\nShutting down...")

		// Gracefully shut down the HTTP server (drains active connections).
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		srv.Shutdown(ctx)
	}()

	serveErr := error(nil)
	if listenProto == "https" {
		serveErr = srv.ListenAndServeTLS(tlsCertFile, tlsKeyFile)
	} else {
		serveErr = srv.ListenAndServe()
	}
	if serveErr != nil && serveErr != http.ErrServerClosed {
		log.Fatalf("Server failed: %v", serveErr)
	}

	// Clean up all subsystems — defers (e.g. geoDB.Close) will also run.
	statsCollector.Stop()
	talkerTracker.Stop()
	if dnsProvider != nil {
		dnsProvider.Stop()
	}
	if wifiProvider != nil {
		wifiProvider.Stop()
	}
	conntrackTracker.Stop()
	topoScanner.Stop()
	latencyMonitor.Stop()
	if liveActivityMgr != nil {
		liveActivityMgr.Stop()
	}
	dnsResolver.Stop()
}

// withSignature wraps an http.Handler to inject a X-Bandwidth-Monitor header
// on every response, allowing clients to verify they're talking to the right service.
func withSignature(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Bandwidth-Monitor", version.String())
		h.ServeHTTP(w, r)
	})
}

// responseLogger wraps http.ResponseWriter to capture the status code and
// bytes written, for withRequestLog. It implements http.Flusher (delegating
// to the underlying ResponseWriter) so the SSE handler's type-assertion for
// flush support keeps working when logging is enabled.
type responseLogger struct {
	http.ResponseWriter
	status int
	bytes  int
}

func (rl *responseLogger) WriteHeader(status int) {
	rl.status = status
	rl.ResponseWriter.WriteHeader(status)
}

func (rl *responseLogger) Write(b []byte) (int, error) {
	if rl.status == 0 {
		rl.status = http.StatusOK
	}
	n, err := rl.ResponseWriter.Write(b)
	rl.bytes += n
	return n, err
}

func (rl *responseLogger) Flush() {
	if f, ok := rl.ResponseWriter.(http.Flusher); ok {
		f.Flush()
	}
}

// withRequestLog wraps an http.Handler to log method, path+query, remote
// addr, status, response size, and duration for every request to stdout.
// Opt-in via DEBUG_HTTP_LOG — useful for watching what clients actually
// request and how much data each response actually costs (e.g. confirming a
// browser client is sending ?since=/?iface= and getting a small response
// back, not the full history dump), but noisy for normal use.
func withRequestLog(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		rl := &responseLogger{ResponseWriter: w}
		h.ServeHTTP(rl, r)
		status := rl.status
		if status == 0 {
			status = http.StatusOK
		}
		log.Printf("http: %s %s %s %d %dB %s", r.Method, r.URL.RequestURI(), r.RemoteAddr, status, rl.bytes, time.Since(start))
	})
}