main.go

// main.go
package main

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"sort"
	"strings"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

type IPTablesCollector struct {
	tables  []string
	ipv6    bool
	timeout time.Duration

	descRules          *prometheus.Desc
	descRulesAction    *prometheus.Desc
	descRulesTable     *prometheus.Desc
	descScrapeSuccess  *prometheus.Desc
	descScrapeDuration *prometheus.Desc
}

func NewIPTablesCollector(tables []string, ipv6 bool, timeout time.Duration) *IPTablesCollector {
	return &IPTablesCollector{
		tables:  tables,
		ipv6:    ipv6,
		timeout: timeout,

		descRules: prometheus.NewDesc(
			"iptables_rules",
			"Number of iptables rules (current)",
			[]string{"family", "table", "chain"},
			nil,
		),
		descRulesAction: prometheus.NewDesc(
			"iptables_rules_action",
			"Number of iptables rules by action (current)",
			[]string{"family", "table", "chain", "action"},
			nil,
		),
		descRulesTable: prometheus.NewDesc(
			"iptables_rules_table",
			"Number of iptables rules in table (current)",
			[]string{"family", "table"},
			nil,
		),
		descScrapeSuccess: prometheus.NewDesc(
			"iptables_scrape_success",
			"Whether the last scrape for a family succeeded (1) or failed (0)",
			[]string{"family"},
			nil,
		),
		descScrapeDuration: prometheus.NewDesc(
			"iptables_scrape_duration_seconds",
			"Duration of last scrape for a family",
			[]string{"family"},
			nil,
		),
	}
}

func (c *IPTablesCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.descRules
	ch <- c.descRulesAction
	ch <- c.descRulesTable
	ch <- c.descScrapeSuccess
	ch <- c.descScrapeDuration
}

func (c *IPTablesCollector) Collect(ch chan<- prometheus.Metric) {
	families := []struct {
		family string
		bin    string
	}{
		{family: "ipv4", bin: "iptables-save"},
	}
	if c.ipv6 {
		families = append(families, struct {
			family string
			bin    string
		}{family: "ipv6", bin: "ip6tables-save"})
	}

	for _, fam := range families {
		start := time.Now()
		success := 1.0

		if err := c.collectFamily(ch, fam.family, fam.bin); err != nil {
			success = 0
			log.Printf("collect failed: family=%s bin=%s err=%v", fam.family, fam.bin, err)
		}

		ch <- prometheus.MustNewConstMetric(c.descScrapeSuccess, prometheus.GaugeValue, success, fam.family)
		ch <- prometheus.MustNewConstMetric(c.descScrapeDuration, prometheus.GaugeValue, time.Since(start).Seconds(), fam.family)
	}
}

func (c *IPTablesCollector) collectFamily(ch chan<- prometheus.Metric, family, bin string) error {
	fullPath, err := exec.LookPath(bin)
	if err != nil {
		return fmt.Errorf("binary not found: %s: %w", bin, err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()

	rules, tableTotals, err := scrapeIPTables(ctx, fullPath, c.tables)
	if err != nil {
		return err
	}

	// tableChainRules: map[table][chain]count
	tableChainRules := make(map[string]map[string]int)
	// tableChainActionRules: map[table][chain][action]count
	tableChainActionRules := make(map[string]map[string]map[string]int)

	for k, count := range rules {
		if tableChainRules[k.table] == nil {
			tableChainRules[k.table] = make(map[string]int)
		}
		tableChainRules[k.table][k.chain] += count

		if tableChainActionRules[k.table] == nil {
			tableChainActionRules[k.table] = make(map[string]map[string]int)
		}
		if tableChainActionRules[k.table][k.chain] == nil {
			tableChainActionRules[k.table][k.chain] = make(map[string]int)
		}
		tableChainActionRules[k.table][k.chain][k.action] += count
	}

	// emit metrics (stable order)
	tables := make([]string, 0, len(tableTotals))
	for t := range tableTotals {
		tables = append(tables, t)
	}
	sort.Strings(tables)

	for _, table := range tables {
		ch <- prometheus.MustNewConstMetric(c.descRulesTable, prometheus.GaugeValue, float64(tableTotals[table]), family, table)

		chains := make([]string, 0, len(tableChainRules[table]))
		for cn := range tableChainRules[table] {
			chains = append(chains, cn)
		}
		sort.Strings(chains)

		for _, chain := range chains {
			ch <- prometheus.MustNewConstMetric(c.descRules, prometheus.GaugeValue, float64(tableChainRules[table][chain]), family, table, chain)

			actions := make([]string, 0, len(tableChainActionRules[table][chain]))
			for a := range tableChainActionRules[table][chain] {
				actions = append(actions, a)
			}
			sort.Strings(actions)

			for _, action := range actions {
				ch <- prometheus.MustNewConstMetric(c.descRulesAction, prometheus.GaugeValue, float64(tableChainActionRules[table][chain][action]), family, table, chain, action)
			}
		}
	}

	return nil
}

type ruleKey struct {
	table  string
	chain  string
	action string
}

func scrapeIPTables(ctx context.Context, bin string, allowedTables []string) (map[ruleKey]int, map[string]int, error) {
	cmd := exec.CommandContext(ctx, bin)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, fmt.Errorf("StdoutPipe: %w", err)
	}
	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	if err := cmd.Start(); err != nil {
		return nil, nil, fmt.Errorf("start: %w", err)
	}

	rules := make(map[ruleKey]int)
	tableTotals := make(map[string]int)

	allowed := make(map[string]bool)
	for _, t := range allowedTables {
		allowed[t] = true
	}

	sc := bufio.NewScanner(stdout)
	sc.Buffer(make([]byte, 1024), 1024*1024)

	var currentTable string
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}

		if strings.HasPrefix(line, "*") {
			currentTable = line[1:]
			if allowed[currentTable] && tableTotals[currentTable] == 0 {
				tableTotals[currentTable] = 0 // Ensure table exists in totals
			}
			continue
		}

		if currentTable == "" || !allowed[currentTable] {
			continue
		}

		if strings.HasPrefix(line, ":") {
			// Chain definition: :CHAIN_NAME POLICY [PACKETS:BYTES]
			parts := strings.Fields(line)
			if len(parts) > 0 {
				chain := parts[0][1:]
				key := ruleKey{
					table:  currentTable,
					chain:  chain,
					action: "CONTINUE", // Default for an empty chain
				}
				if _, exists := rules[key]; !exists {
					rules[key] = 0
				}
			}
			continue
		}

		if strings.HasPrefix(line, "-A ") {
			fields := strings.Fields(line)
			if len(fields) >= 2 {
				chain := fields[1]
				tableTotals[currentTable]++

				// Extract action (-j or --jump)
				action := "CONTINUE"
				for i := 0; i < len(fields)-1; i++ {
					if fields[i] == "-j" || fields[i] == "--jump" {
						if i+1 < len(fields) {
							action = fields[i+1]
						}
						break
					}
				}

				key := ruleKey{
					table:  currentTable,
					chain:  chain,
					action: action,
				}
				rules[key]++
			}
		}
	}

	if err := sc.Err(); err != nil {
		_ = cmd.Wait()
		return nil, nil, fmt.Errorf("scan: %w", err)
	}

	if err := cmd.Wait(); err != nil {
		return nil, nil, fmt.Errorf("iptables error: %w: %s", err, strings.TrimSpace(stderr.String()))
	}

	return rules, tableTotals, nil
}

func splitCSV(s string) []string {
	parts := strings.Split(s, ",")
	out := make([]string, 0, len(parts))
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p != "" {
			out = append(out, p)
		}
	}
	return out
}

func getEnv(key, fallback string) string {
	if value, ok := os.LookupEnv(key); ok {
		return value
	}
	return fallback
}

func main() {
	var (
		listen      = flag.String("listen", getEnv("LISTEN_ADDRESS", ":9105"), "listen address")
		metricsPath = flag.String("metrics-path", getEnv("METRICS_PATH", "/metrics"), "metrics path")
		enableIPv6  = flag.Bool("ipv6", getEnv("ENABLE_IPV6", "true") == "true", "collect ip6tables too")
		tablesCSV   = flag.String("tables", getEnv("TABLES", "filter,nat,mangle,raw,security"), "comma-separated tables")
		timeout     = flag.Duration("timeout", func() time.Duration {
			d, err := time.ParseDuration(getEnv("TIMEOUT", "2s"))
			if err != nil {
				return 2 * time.Second
			}
			return d
		}(), "timeout per iptables command")
	)
	flag.Parse()

	collector := NewIPTablesCollector(splitCSV(*tablesCSV), *enableIPv6, *timeout)

	reg := prometheus.NewRegistry()
	reg.MustRegister(collector)
	reg.MustRegister(prometheus.NewGoCollector())
	reg.MustRegister(prometheus.NewProcessCollector(prometheus.ProcessCollectorOpts{}))

	mux := http.NewServeMux()
	mux.Handle(*metricsPath, promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte("ok\n")) })

	srv := &http.Server{
		Addr:              *listen,
		Handler:           mux,
		ReadHeaderTimeout: 5 * time.Second,
	}

	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt, syscall.SIGTERM)

	go func() {
		log.Printf("listening on %s, metrics on %s", *listen, *metricsPath)
		if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			log.Fatalf("listen: %s\n", err)
		}
	}()

	<-done
	log.Print("stopping server...")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := srv.Shutdown(ctx); err != nil {
		log.Fatalf("server shutdown failed: %+v", err)
	}
	log.Print("server stopped")
}