package main

import (
	"fmt"
	"os/exec"
)

func initNFTables() error {
	cmds := [][]string{
		{"nft", "add", "table", "inet", "wgrplane"},
		{"nft", "add", "chain", "inet", "wgrplane", "forward", "{", "type", "filter", "hook", "forward", "priority", "0;", "}"},
		{"nft", "add", "set", "inet", "wgrplane", "wg_access", "{", "type", "ipv4_addr", ".", "ipv4_addr;", "}"},
		{"nft", "add", "set", "inet", "wgrplane", "wg_internet", "{", "type", "ipv4_addr;", "}"},
	}
	for _, args := range cmds {
		if err := exec.Command(args[0], args[1:]...).Run(); err != nil {
			return fmt.Errorf("nftables init failed: %v", err)
		}
	}
	return nil
}

func AddAccessRule(peerIP, targetCIDR string) error {
	element := fmt.Sprintf("{ %s . %s }", peerIP, targetCIDR)
	return exec.Command("nft", "add", "element", "inet", "wgrplane", "wg_access", element).Run()
}

func RemoveAccessRule(peerIP, targetCIDR string) error {
	element := fmt.Sprintf("{ %s . %s }", peerIP, targetCIDR)
	return exec.Command("nft", "delete", "element", "inet", "wgrplane", "wg_access", element).Run()
}

func SetInternetAccess(peerIP string, enabled bool) error {
	if enabled {
		element := fmt.Sprintf("{ %s }", peerIP)
		return exec.Command("nft", "add", "element", "inet", "wgrplane", "wg_internet", element).Run()
	}
	element := fmt.Sprintf("{ %s }", peerIP)
	return exec.Command("nft", "delete", "element", "inet", "wgrplane", "wg_internet", element).Run()
}