firewall_containers/network-go/firewall/firewall.go

package firewall

import (
	"context"
	"log"
	"net"
	"strconv"
	"time"

	"firewall_containers/network-go/config"
	"firewall_containers/network-go/docker"
	"firewall_containers/network-go/iptables"
	"firewall_containers/network-go/resolver"
)

// Orchestrator reconciles the networks.json configuration into Docker networks
// and iptables firewall rules
type Orchestrator struct {
	dockerClient docker.DockerAPI
	iptablesMgr  iptables.IPTablesAPI
	resolver     *resolver.Resolver
	debug        bool
}

// NewOrchestrator creates a new firewall orchestrator
func NewOrchestrator(dockerClient docker.DockerAPI, iptablesMgr iptables.IPTablesAPI, cfg *config.NetworksConfig) *Orchestrator {
	return &Orchestrator{
		dockerClient: dockerClient,
		iptablesMgr:  iptablesMgr,
		resolver:     resolver.NewResolver(cfg),
	}
}

// ReconcileAll runs the full reconciliation: networks, container connections, and firewall rules
func (o *Orchestrator) ReconcileAll(ctx context.Context, cfg *config.NetworksConfig) {
	log.Println("FIREWALL: starting full reconciliation")

	// Update resolver with latest config
	o.resolver.SetConfig(cfg)

	// Step 0: Enable IP forwarding (may fail in containers with read-only fs)
	if err := o.iptablesMgr.EnsureIPForward(); err != nil {
		log.Printf("FIREWALL: WARNING could not enable ip_forward: %v", err)
	} else {
		log.Println("FIREWALL: IP forwarding enabled")
	}

	// Step 1: Ensure all defined networks exist
	o.reconcileNetworks(ctx, cfg)

	// Step 2: Connect containers to networks with assigned IPs
	o.reconcileIPs(ctx, cfg)

	// Step 3: Reconcile firewall policies
	o.reconcilePolicies(ctx, cfg)

	log.Println("FIREWALL: full reconciliation completed")
}

// reconcileNetworks creates Docker networks if they don't exist
func (o *Orchestrator) reconcileNetworks(ctx context.Context, cfg *config.NetworksConfig) {
	for name, netCfg := range cfg.Networks {
		log.Printf("FIREWALL: ensuring network %s (%s, subnet=%s, gateway=%s)", name, netCfg.NetworkName, netCfg.Subnet, netCfg.Gateway)
		if err := o.dockerClient.EnsureNetwork(ctx, netCfg); err != nil {
			log.Printf("FIREWALL: ERROR ensuring network %s: %v", name, err)
		}
	}
}

// reconcileIPs connects containers to networks with their assigned IPs
func (o *Orchestrator) reconcileIPs(ctx context.Context, cfg *config.NetworksConfig) {
	for ipStr, ipCfg := range cfg.IPs {
		networkName := findNetworkForIP(cfg, ipStr)
		if networkName == "" {
			log.Printf("FIREWALL: WARNING no network found for IP %s (container=%s)", ipStr, ipCfg.ContainerName)
			continue
		}

		// Resolve the actual container name, with fallback to fuzzy matching
		// (old shell script behavior: docker ps | grep $D"-")
		containerName, err := o.dockerClient.FindContainerName(ctx, ipCfg.ContainerName, ipCfg.Selector)
		if err != nil {
			log.Printf("FIREWALL: WARNING container %s (selector=%s) not found: %v, trying connection anyway", ipCfg.ContainerName, ipCfg.Selector, err)
			containerName = ipCfg.ContainerName
		}

		log.Printf("FIREWALL: connecting container %s to network %s with IP %s", containerName, networkName, ipStr)

		waitCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
		if err := o.dockerClient.WaitForContainerRunning(waitCtx, containerName, 10*time.Second); err != nil {
			log.Printf("FIREWALL: WARNING container %s not running yet: %v, connecting anyway", containerName, err)
		}
		cancel()

		if err := o.dockerClient.ConnectContainer(ctx, containerName, networkName, ipStr); err != nil {
			log.Printf("FIREWALL: ERROR connecting container %s to %s: %v", containerName, networkName, err)
		}
	}
}

// reconcilePolicies translates PolicyConfig entries into iptables rules
func (o *Orchestrator) reconcilePolicies(ctx context.Context, cfg *config.NetworksConfig) {
	for i, policy := range cfg.Policies {
		log.Printf("FIREWALL: processing policy[%d]", i)

		proto := policy.Proto
		if proto == "" {
			proto = "tcp"
		}
		port := strconv.Itoa(policy.Port)

		// Build comment for iptables (matches shell script's NAME-COMMENT pattern)
		// Use Name if present, otherwise ServiceName, to avoid trailing dashes
		comment := ""
		if policy.Name != "" {
			comment = policy.Name
		}
		if policy.ServiceName != "" {
			if comment != "" {
				comment += "-" + policy.ServiceName
			} else {
				comment = policy.ServiceName
			}
		}

		// CASE 1: Rule with "from" field — this is a FORWARD ACCEPT rule
		if policy.From != "" {
			o.applyForwardRule(ctx, cfg, policy, proto, port, comment)
			continue
		}

		// CASE 2: Rule with "nat" field — this is a DNAT/MASQUERADE rule
		if policy.Nat != "" {
			o.applyNATRule(ctx, cfg, policy, proto, port, comment)
			continue
		}

		// Unhandled pattern
		log.Printf("FIREWALL: policy[%d] unhandled pattern — service=%s container=%s selector=%s from=%s to=%s port=%d proto=%s nat=%s",
			i, policy.ServiceName, policy.ContainerName, policy.Selector, policy.From, policy.To, policy.Port, policy.Proto, policy.Nat)
	}
}

func (o *Orchestrator) applyForwardRule(ctx context.Context, cfg *config.NetworksConfig, policy config.PolicyConfig, proto, port, comment string) {
	sourceIP := o.resolveIP(policy.From)
	targetIP := ""
	if policy.To != "" {
		targetIP = o.resolveIP(policy.To)
	}

	// Determine the chain: use DOCKER-USER (iptables-legacy) or FORWARD
	chain := "FORWARD"
	if o.iptablesMgr.Binary() == "/usr/sbin/iptables-legacy" {
		chain = "DOCKER-USER"
	}

	// Ensure established/related rule exists at the top
	if err := o.iptablesMgr.EnsureEstablishedRelated(chain); err != nil {
		log.Printf("FIREWALL: ERROR ensuring established/related rule in %s: %v", chain, err)
	}

	// Insert the FORWARD ACCEPT rule
	if err := o.iptablesMgr.InsertForwardAccept(chain, sourceIP, targetIP, proto, "", port, comment); err != nil {
		log.Printf("FIREWALL: ERROR inserting FORWARD ACCEPT rule: %v", err)
	}
}

func (o *Orchestrator) applyNATRule(ctx context.Context, cfg *config.NetworksConfig, policy config.PolicyConfig, proto, port, comment string) {
	to := policy.To

	// Resolve "to" as target IP
	targetIP := o.resolveIP(to)

	if targetIP == "" {
		log.Printf("FIREWALL: WARNING cannot resolve target %s for nat policy", to)
		return
	}

	if policy.Nat == "dnat" {
		// Determine the best container selector from the policy: try Selector, then ContainerName, then Name
		selector := policy.Selector
		if selector == "" {
			selector = policy.ContainerName
		}
		if selector == "" {
			selector = policy.Name
		}

		// Try to insert rules inside the container namespace via nsenter
		usedContainer := false
		if selector != "" {
			pid, err := o.dockerClient.GetContainerPID(ctx, selector)
			if err == nil {
				if err := o.iptablesMgr.InsertPreroutingRuleInContainer(pid, "0.0.0.0/0", proto, port, targetIP, port, comment); err != nil {
					log.Printf("FIREWALL: ERROR inserting container PREROUTING rule: %v", err)
				} else {
					usedContainer = true
				}
			} else {
				log.Printf("FIREWALL: WARNING cannot get PID for container %s: %v, trying host rules", selector, err)
			}
		}

		// Fall back to host-level PREROUTING if container not used
		if !usedContainer && policy.Iface != "" {
			if err := o.iptablesMgr.InsertPreroutingRuleOnInterface(policy.Iface, proto, port, targetIP, port, comment); err != nil {
				log.Printf("FIREWALL: ERROR inserting interface PREROUTING rule: %v", err)
			}
		}
	}
}

// resolveIP resolves a name or IP string to an IP address using networks.json config
func (o *Orchestrator) resolveIP(name string) string {
	// If it's already an IP, return it as CIDR
	if config.IsIP(name) {
		return config.ToCIDR(name)
	}

	// Use the resolver which looks up from networks.json
	ips := o.resolver.Resolve(name)
	if len(ips) > 0 {
		return ips[0]
	}

	return ""
}

// findNetworkForIP finds the network name that contains the given IP in its subnet
func findNetworkForIP(cfg *config.NetworksConfig, ip string) string {
	for _, netCfg := range cfg.Networks {
		subnet, err := netCfg.ParseCIDR()
		if err != nil {
			continue
		}
		parsedIP := net.ParseIP(ip)
		if parsedIP == nil {
			continue
		}
		if subnet.Contains(parsedIP) {
			return netCfg.NetworkName
		}
	}
	return ""
}