SamWaf/firewall/ipset_linux.go

//go:build linux

package firewall

import (
	"SamWaf/common/wafexec"
	"fmt"
	"os/exec"

	"strings"
)

// 说明：本文件在 Linux 下用 ipset 提供"批量封禁大列表(上万~十万条威胁情报 IP)"的能力，
// 区别于 firewall.go 里逐条 iptables 规则(仅适合少量手工封禁)。
//   - 一个命名 set 装十万 IP/CIDR，配一条 `iptables -m set --match-set <set> src -j DROP` 规则，
//     收包匹配走内核哈希 O(1)，而非 INPUT 链上万条线性规则。
//   - 全量重建用 `ipset restore` 从 stdin 一次灌入 + `swap` 原子替换，单次 fork，避免旧实现的 O(n²)。
// v4/v6 分成两个 set(hash:net 的 family 在 create 时固定)：<set> 走 iptables，<set>_6 走 ip6tables。

const (
	// ipsetMaxElem 单个 set 的最大元素数上限(hash:net 动态增长，此值仅为封顶，防投毒膨胀)
	ipsetMaxElem = "2097152"
	// ipsetProbeName 能力探测用的临时 set 名
	ipsetProbeName = "samwaf_probe"
	// ipsetMaxBaseNameLen set 基础名长度上限：ipset 名总长上限 31，需为 "_6"(v6) 与 "_s"(swap) 预留后缀
	ipsetMaxBaseNameLen = 24
)

// v6SetName 返回某 set 对应的 IPv6 集合名
func v6SetName(setName string) string { return setName + "_6" }

// swapSetName 返回某 set 全量重建时用的临时交换集合名
func swapSetName(setName string) string { return setName + "_s" }

// validateSetName 校验集合名合法(小写字母/数字/下划线、长度受限)，兼作防命令注入兜底
func validateSetName(setName string) error {
	if len(setName) == 0 || len(setName) > ipsetMaxBaseNameLen {
		return fmt.Errorf("ipset 集合名长度需为 1-%d: %q", ipsetMaxBaseNameLen, setName)
	}
	for _, c := range setName {
		if !(c >= 'a' && c <= 'z') && !(c >= '0' && c <= '9') && c != '_' {
			return fmt.Errorf("ipset 集合名仅允许小写字母/数字/下划线: %q", setName)
		}
	}
	return nil
}

// runFirewallCmd 执行一条防火墙相关命令并返回合并输出(仅补 Stdin，不接管 Stdout/Stderr 逻辑)
func (fw *FireWallEngine) runFirewallCmd(name string, args ...string) (string, error) {
	cmd := wafexec.FixStdin(exec.Command(name, args...))
	out, err := cmd.CombinedOutput()
	return string(out), err
}

// hasIP6tables 判断系统是否具备 ip6tables(无则跳过 v6 封禁，仅处理 v4)
func hasIP6tables() bool {
	_, err := exec.LookPath("ip6tables")
	return err == nil
}

// SupportsIPSet 报告当前环境是否可用 ipset 批量封禁(带 30s 缓存，见 available.go)
func (fw *FireWallEngine) SupportsIPSet() bool {
	return cachedSupportsIPSet(fw) == nil
}

// IPSetUpToDate 报告集合内容是否已经就是这份 ip 列表(尽力而为)。
//
// 用 `ipset list -t`(terse，只出头部不出成员)拿 "Number of entries" 与期望条数比对，
// 一次 fork、与集合大小无关，很便宜。用途有二：
//   - 启动重放：进程重启(非机器重启)时内核集合还在，可直接跳过重灌
//   - 落地对账：定时核对系统层是否与快照一致，不一致就用快照覆盖式重建
//
// 只比条数不比内容：逐条比对十万级集合的代价远超收益，而条数不符已经能覆盖
// "集合没了 / 灌了一半 / 灌成了别的快照"这些实际会发生的情况。
// 误判的后果只是多做一次重建(`ipset restore`+`swap` 幂等且原子)，不会出错。
func (fw *FireWallEngine) IPSetUpToDate(setName string, ips []string) bool {
	if !fw.SupportsIPSet() {
		return false
	}
	if err := validateSetName(setName); err != nil {
		return false
	}
	v4, v6 := splitByIPVersion(ips)
	if !fw.setEntryCountEquals(setName, len(v4)) {
		return false
	}
	if hasIP6tables() && !fw.setEntryCountEquals(v6SetName(setName), len(v6)) {
		return false
	}
	return true
}

// setEntryCountEquals 判断某集合当前条目数是否等于 want。集合不存在/解析不出则返回 false。
func (fw *FireWallEngine) setEntryCountEquals(setName string, want int) bool {
	out, err := fw.runFirewallCmd("ipset", "list", "-t", setName)
	if err != nil {
		return false // 集合不存在或 ipset 不可用，都当作"需要重建"
	}
	// 头部形如 "Number of entries: 13157"，字段名在各版本 ipset 中稳定且为英文
	n, ok := parseIpsetEntryCount(out)
	return ok && n == want
}

// supportsIPSet 实际探测：ipset 二进制存在 + 能创建/销毁临时 set(验证内核模块与权限)
func (fw *FireWallEngine) supportsIPSet() error {
	if _, err := exec.LookPath("ipset"); err != nil {
		msg := "当前环境未安装 ipset，无法使用大列表批量封禁，可改用 WAF 应用层黑名单。"
		if isInContainer() {
			msg += containerHint
		}
		return fmt.Errorf("%s", msg)
	}
	if out, err := fw.runFirewallCmd("ipset", "create", ipsetProbeName, "hash:net", "-exist"); err != nil {
		msg := fmt.Sprintf("ipset 不可用(可能缺内核模块或权限): %v, 输出: %s", err, strings.TrimSpace(out))
		if isInContainer() {
			msg += containerHint
		}
		return fmt.Errorf("%s", msg)
	}
	fw.runFirewallCmd("ipset", "destroy", ipsetProbeName)
	return nil
}

// EnsureIPSet 确保 v4/v6 两个集合存在，且 INPUT 上各挂了一条 match-set DROP 引用规则(幂等)
func (fw *FireWallEngine) EnsureIPSet(setName string) error {
	if err := validateSetName(setName); err != nil {
		return err
	}
	if out, err := fw.runFirewallCmd("ipset", "create", setName, "hash:net", "family", "inet", "maxelem", ipsetMaxElem, "-exist"); err != nil {
		return fmt.Errorf("创建 ipset %s 失败: %v, 输出: %s", setName, err, strings.TrimSpace(out))
	}
	fw.ensureDropRule("iptables", setName)
	// v6 尽力而为：无 ip6tables 时跳过
	if hasIP6tables() {
		if _, err := fw.runFirewallCmd("ipset", "create", v6SetName(setName), "hash:net", "family", "inet6", "maxelem", ipsetMaxElem, "-exist"); err == nil {
			fw.ensureDropRule("ip6tables", v6SetName(setName))
		}
	}
	return nil
}

// ensureDropRule 幂等挂载 `<iptablesBin> -I INPUT 1 -m set --match-set <set> src -j DROP`
func (fw *FireWallEngine) ensureDropRule(iptablesBin, setName string) {
	if _, err := fw.runFirewallCmd(iptablesBin, "-C", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP"); err == nil {
		return // 已存在
	}
	fw.runFirewallCmd(iptablesBin, "-I", "INPUT", "1", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
}

// SupportsPortScopedSet 报告本平台能否把集合的封禁范围限制在指定端口上。
// Linux 支持：iptables 的 multiport 匹配可以直接写进那条 match-set 引用规则里。
func (fw *FireWallEngine) SupportsPortScopedSet() bool { return true }

// ApplyIPSetPortScope 调整集合引用规则的作用端口。
//
// ports 为空表示恢复成"封全端口"。用途是让主机防爆破可以只封 SSH/RDP 端口——
// 一旦误封，被挡住的只是远程登录，Web、数据库、业务端口都还通着，
// 排查和自救的余地大得多。
//
// 实现上是"先删掉该集合已有的全部 INPUT 引用规则，再按新范围插一条"，
// 而不是叠加：叠加会在反复改端口后留下一串旧规则，其中任何一条命中都照样 DROP，
// 用户会看到"明明改成只封22了，怎么全端口还是不通"。
func (fw *FireWallEngine) ApplyIPSetPortScope(setName string, tcpPorts []int) error {
	if err := validateSetName(setName); err != nil {
		return err
	}
	if err := fw.EnsureIPSet(setName); err != nil {
		return err
	}
	fw.applyPortScopeOne("iptables", setName, tcpPorts)
	if hasIP6tables() {
		fw.applyPortScopeOne("ip6tables", v6SetName(setName), tcpPorts)
	}
	return nil
}

// applyPortScopeOne 对单个集合(v4 或 v6)重建引用规则
func (fw *FireWallEngine) applyPortScopeOne(iptablesBin, setName string, tcpPorts []int) {
	// 反复 -D 直到删不动为止：同一条规则可能因历史操作存在多份
	base := []string{"-D", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP"}
	for i := 0; i < 16; i++ {
		if _, err := fw.runFirewallCmd(iptablesBin, base...); err != nil {
			break
		}
	}
	portSpec := formatMultiport(tcpPorts)
	if portSpec != "" {
		withPorts := []string{"-D", "INPUT", "-p", "tcp", "-m", "multiport", "--dports", portSpec,
			"-m", "set", "--match-set", setName, "src", "-j", "DROP"}
		for i := 0; i < 16; i++ {
			if _, err := fw.runFirewallCmd(iptablesBin, withPorts...); err != nil {
				break
			}
		}
	}

	if portSpec == "" {
		fw.runFirewallCmd(iptablesBin, "-I", "INPUT", "1", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
		return
	}
	fw.runFirewallCmd(iptablesBin, "-I", "INPUT", "1", "-p", "tcp", "-m", "multiport", "--dports", portSpec,
		"-m", "set", "--match-set", setName, "src", "-j", "DROP")
}

// formatMultiport 把端口列表拼成 multiport 参数。
// multiport 最多支持 15 个端口，超出的截断——SSH/RDP 场景下不可能超。
func formatMultiport(ports []int) string {
	if len(ports) == 0 {
		return ""
	}
	if len(ports) > 15 {
		ports = ports[:15]
	}
	parts := make([]string, 0, len(ports))
	for _, p := range ports {
		if p > 0 && p <= 65535 {
			parts = append(parts, fmt.Sprintf("%d", p))
		}
	}
	return strings.Join(parts, ",")
}

// RestoreIPSet 用给定 IP/CIDR 列表全量原子重建集合(v4/v6 自动分流)。
// 采用 `ipset restore` 单次 fork：create 临时交换集合→逐条 add→swap 原子替换→destroy 临时集合。
func (fw *FireWallEngine) RestoreIPSet(setName string, ips []string) error {
	if err := fw.EnsureIPSet(setName); err != nil {
		return err
	}
	v4, v6 := splitByIPVersion(ips)
	if err := fw.restoreOneSet(setName, "inet", v4); err != nil {
		return err
	}
	if hasIP6tables() {
		if err := fw.restoreOneSet(v6SetName(setName), "inet6", v6); err != nil {
			return err
		}
	}
	return nil
}

// restoreOneSet 通过 ipset restore + swap 原子重建单个集合
func (fw *FireWallEngine) restoreOneSet(setName, family string, ips []string) error {
	swp := swapSetName(setName)
	var b strings.Builder
	fmt.Fprintf(&b, "create %s hash:net family %s maxelem %s -exist\n", swp, family, ipsetMaxElem)
	fmt.Fprintf(&b, "flush %s\n", swp)
	for _, ip := range ips {
		ip = strings.TrimSpace(ip)
		if ip == "" {
			continue
		}
		fmt.Fprintf(&b, "add %s %s\n", swp, ip)
	}
	fmt.Fprintf(&b, "swap %s %s\n", swp, setName)
	fmt.Fprintf(&b, "destroy %s\n", swp)
	return fw.ipsetRestore(b.String())
}

// ipsetRestore 把 restore 脚本从 stdin 灌给 `ipset restore -exist`(忽略重复/不存在，保证幂等)
func (fw *FireWallEngine) ipsetRestore(payload string) error {
	cmd := exec.Command("ipset", "restore", "-exist")
	cmd.Stdin = strings.NewReader(payload)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("ipset restore 失败: %v, 输出: %s", err, strings.TrimSpace(string(out)))
	}
	return nil
}

// SupportsIncrementalIPSet 报告本平台是否支持集合的增量增删。
// Linux 支持：`ipset restore` 可以只灌 add/del 行，一次 fork 搞定，不需要重建整个集合。
// 这对主机防爆破很关键——那边是"来一个封一个"，用 RestoreIPSet 全量重建会把开销放大成 O(n)。
func (fw *FireWallEngine) SupportsIncrementalIPSet() bool { return true }

// AddToIPSet 向集合增量添加 IP/CIDR(v4/v6 自动分流)，-exist 保证重复添加不报错
func (fw *FireWallEngine) AddToIPSet(setName string, ips []string) error {
	return fw.incrementalIPSet(setName, ips, "add")
}

// DelFromIPSet 从集合增量删除 IP/CIDR，-exist 保证删不存在的项不报错
func (fw *FireWallEngine) DelFromIPSet(setName string, ips []string) error {
	return fw.incrementalIPSet(setName, ips, "del")
}

// incrementalIPSet 拼 restore 脚本做批量增量操作
func (fw *FireWallEngine) incrementalIPSet(setName string, ips []string, op string) error {
	if len(ips) == 0 {
		return nil
	}
	if err := fw.EnsureIPSet(setName); err != nil {
		return err
	}
	v4, v6 := splitByIPVersion(ips)
	var b strings.Builder
	for _, ip := range v4 {
		fmt.Fprintf(&b, "%s %s %s\n", op, setName, ip)
	}
	if hasIP6tables() {
		for _, ip := range v6 {
			fmt.Fprintf(&b, "%s %s %s\n", op, v6SetName(setName), ip)
		}
	}
	if b.Len() == 0 {
		return nil
	}
	return fw.ipsetRestore(b.String())
}

// FlushIPSet 清空集合内容(保留集合与引用规则)
func (fw *FireWallEngine) FlushIPSet(setName string) error {
	if err := validateSetName(setName); err != nil {
		return err
	}
	fw.runFirewallCmd("ipset", "flush", setName)
	fw.runFirewallCmd("ipset", "flush", v6SetName(setName))
	return nil
}

// DestroyIPSet 先摘除 INPUT 引用规则，再销毁 v4/v6 集合(set 被引用时无法 destroy)
func (fw *FireWallEngine) DestroyIPSet(setName string) error {
	if err := validateSetName(setName); err != nil {
		return err
	}
	fw.runFirewallCmd("iptables", "-D", "INPUT", "-m", "set", "--match-set", setName, "src", "-j", "DROP")
	if hasIP6tables() {
		fw.runFirewallCmd("ip6tables", "-D", "INPUT", "-m", "set", "--match-set", v6SetName(setName), "src", "-j", "DROP")
	}
	fw.runFirewallCmd("ipset", "destroy", setName)
	fw.runFirewallCmd("ipset", "destroy", v6SetName(setName))
	return nil
}