firewall_containers/network-go/firewall/firewall_test.go

package firewall

import (
	"context"
	"testing"

	"firewall_containers/network-go/config"
	"firewall_containers/network-go/mock"
)

func testConfig() *config.NetworksConfig {
	return &config.NetworksConfig{
		Networks: map[string]config.NetworkConfig{
			"smarthost-loadbalancer": {
				NetworkName: "smarthost-loadbalancer",
				Subnet:      "172.18.103.0/24",
				Gateway:     "172.18.103.1",
			},
			"smarthost_backend-1": {
				NetworkName: "smarthost_backend-1",
				Subnet:      "172.18.104.0/24",
				Gateway:     "172.18.104.1",
			},
		},
		IPs: map[string]config.IPConfig{
			"172.18.103.2": {
				IP:            "172.18.103.2",
				ContainerName: "smarthostloadbalancer",
				Selector:      "smarthostloadbalancer",
				ServiceName:   "smarthost-proxy",
			},
			"172.18.104.2": {
				IP:            "172.18.104.2",
				ContainerName: "smarthostbackend-1",
				Selector:      "smarthostbackend-1",
				ServiceName:   "smarthost-proxy",
			},
		},
		Policies: []config.PolicyConfig{},
	}
}

func TestReconcileAllCreatesNetworks(t *testing.T) {
	cfg := testConfig()
	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.ReconcileAll(ctx, cfg)

	if !docker.EnsureNetworkCalled {
		t.Error("EnsureNetwork was not called")
	}
}

func TestReconcileAllConnectsContainers(t *testing.T) {
	cfg := testConfig()
	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.ReconcileAll(ctx, cfg)

	if !docker.ConnectContainerCalled {
		t.Error("ConnectContainer was not called")
	}
}

func TestReconcileAllEnablesIPForward(t *testing.T) {
	cfg := testConfig()
	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.ReconcileAll(ctx, cfg)

	if !iptables.EnsureIPForwardCalled {
		t.Error("EnsureIPForward was not called")
	}
}

func TestReconcilePoliciesForwardRule(t *testing.T) {
	cfg := testConfig()
	cfg.Policies = []config.PolicyConfig{
		{
			ServiceName:   "smarthost-proxy",
			ContainerName: "smarthostloadbalancer",
			Selector:      "smarthostloadbalancer",
			From:          "publicbackend",
			Port:          80,
			Proto:         "tcp",
		},
	}

	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{BinaryResult: "/usr/sbin/iptables"}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.reconcilePolicies(ctx, cfg)

	// "publicbackend" should be resolved to an IP from the config (if it matches)
	// Since "publicbackend" is not in the IPs map, sourceIP will be empty
	if !iptables.InsertForwardAcceptCalled {
		t.Error("InsertForwardAccept was not called")
	}

	// Should use DOCKER-USER chain (default, even with non-legacy iptables)
	if iptables.InsertForwardAcceptChain != "DOCKER-USER" {
		t.Errorf("expected DOCKER-USER chain, got %s", iptables.InsertForwardAcceptChain)
	}
}

func TestReconcilePoliciesForwardRuleWithLegacy(t *testing.T) {
	cfg := testConfig()
	cfg.Policies = []config.PolicyConfig{
		{
			ServiceName: "smarthost-proxy",
			From:        "publicbackend",
			Port:        80,
			Proto:       "tcp",
		},
	}

	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{BinaryResult: "/usr/sbin/iptables-legacy"}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.reconcilePolicies(ctx, cfg)

	if !iptables.InsertForwardAcceptCalled {
		t.Error("InsertForwardAccept was not called")
	}

	// Should use DOCKER-USER chain when iptables-legacy
	if iptables.InsertForwardAcceptChain != "DOCKER-USER" {
		t.Errorf("expected DOCKER-USER chain, got %s", iptables.InsertForwardAcceptChain)
	}
}

func TestReconcilePoliciesDNATWithInterface(t *testing.T) {
	cfg := testConfig()
	cfg.Policies = []config.PolicyConfig{
		{
			ServiceName: "smarthost-proxy",
			Name:        "wireguardproxy",
			Selector:    "smarthostloadbalancer",
			Iface:       "wg0",
			Nat:         "dnat",
			To:          "smarthostloadbalancer",
			Port:        80,
			Proto:       "tcp",
		},
	}

	docker := &mock.MockDockerClient{
		GetContainerPIDResult: 0, // simulate no PID available
		GetContainerPIDErr:    assertError("container not running"),
	}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.reconcilePolicies(ctx, cfg)

	// When GetContainerPID fails, should fall back to interface-based rule
	if !iptables.InsertPreroutingRuleOnInterfaceCalled {
		t.Error("InsertPreroutingRuleOnInterface was not called (should fall back from nsenter)")
	}

	if len(iptables.InsertPreroutingRuleOnInterfaceArgs) > 0 {
		iface := iptables.InsertPreroutingRuleOnInterfaceArgs[0]
		if iface != "wg0" {
			t.Errorf("expected interface wg0, got %s", iface)
		}
	}
}

func TestReconcilePoliciesDNATWithContainerPID(t *testing.T) {
	cfg := testConfig()
	cfg.Policies = []config.PolicyConfig{
		{
			ServiceName: "smarthost-proxy",
			Name:        "wireguardproxy",
			Selector:    "smarthostloadbalancer",
			Nat:         "dnat",
			To:          "smarthostloadbalancer",
			Port:        80,
			Proto:       "tcp",
		},
	}

	docker := &mock.MockDockerClient{
		GetContainerPIDResult: 1234,
		GetContainerPIDErr:    nil,
	}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.reconcilePolicies(ctx, cfg)

	if !docker.GetContainerPIDCalled {
		t.Error("GetContainerPID was not called")
	}

	if !iptables.InsertPreroutingRuleInContainerCalled {
		t.Error("InsertPreroutingRuleInContainer was not called")
	}

	if iptables.InsertPreroutingRuleInContainerPID != 1234 {
		t.Errorf("expected PID 1234, got %d", iptables.InsertPreroutingRuleInContainerPID)
	}
}

func TestReconcilePoliciesUnresolvedTarget(t *testing.T) {
	cfg := testConfig()
	cfg.Policies = []config.PolicyConfig{
		{
			ServiceName: "test",
			Nat:         "dnat",
			Selector:    "container1",
			To:          "nonexistent-target",
			Port:        80,
		},
	}

	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	ctx := context.Background()
	orch.reconcilePolicies(ctx, cfg)

	// Should not call GetContainerPID when target can't be resolved
	if docker.GetContainerPIDCalled {
		t.Error("GetContainerPID should not be called when target is unresolvable")
	}
}

func TestResolveIPDirectIP(t *testing.T) {
	cfg := testConfig()
	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	// Direct IP should be returned as CIDR
	result := orch.resolveIP("172.18.103.2")
	if result != "172.18.103.2" {
		t.Errorf("expected 172.18.103.2, got %s", result)
	}

	// .0 ending should be converted to /24
	result = orch.resolveIP("172.18.103.0")
	if result != "172.18.103.0/24" {
		t.Errorf("expected 172.18.103.0/24, got %s", result)
	}
}

func TestResolveIPFromConfig(t *testing.T) {
	cfg := testConfig()
	docker := &mock.MockDockerClient{}
	iptables := &mock.MockIPTablesManager{}
	orch := NewOrchestrator(docker, iptables, cfg)

	// Should resolve container name from config
	result := orch.resolveIP("smarthostloadbalancer")
	if result != "172.18.103.2" {
		t.Errorf("expected 172.18.103.2, got %s", result)
	}
}

func TestFindNetworkForIP(t *testing.T) {
	cfg := testConfig()

	tests := []struct {
		ip   string
		want string
	}{
		{"172.18.103.5", "smarthost-loadbalancer"},
		{"172.18.104.5", "smarthost_backend-1"},
		{"10.0.0.1", ""},
		{"", ""},
	}

	for _, tt := range tests {
		got := findNetworkForIP(cfg, tt.ip)
		if got != tt.want {
			t.Errorf("findNetworkForIP(%q) = %q, want %q", tt.ip, got, tt.want)
		}
	}
}

func TestReconcileAllReproducible(t *testing.T) {
	cfg := testConfig()
	cfg.Policies = []config.PolicyConfig{
		{
			ServiceName: "test-svc",
			From:        "publicbackend",
			Port:        80,
			Proto:       "tcp",
		},
	}

	// Run reconciliation twice with separate mocks
	for i := 0; i < 2; i++ {
		docker := &mock.MockDockerClient{}
		iptables := &mock.MockIPTablesManager{}
		orch := NewOrchestrator(docker, iptables, cfg)
		ctx := context.Background()
		orch.ReconcileAll(ctx, cfg)

		if !docker.EnsureNetworkCalled {
			t.Errorf("run %d: EnsureNetwork not called", i)
		}
		if !iptables.InsertForwardAcceptCalled {
			t.Errorf("run %d: InsertForwardAccept not called", i)
		}
	}
}

// assertError is a helper to create a simple error for tests
type simpleErr struct{ msg string }

func (e simpleErr) Error() string { return e.msg }

func assertError(msg string) error {
	return simpleErr{msg: msg}
}