// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package iptables

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestSetup(t *testing.T) {
	cases := []struct {
		name            string
		cfg             Config
		additionalRules [][]string
		expectedRules   []string
	}{
		{
			"no proxy outbound port provided",
			Config{
				ProxyUserID:      "123",
				ProxyInboundPort: 20000,
				IptablesProvider: &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 15001",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"Consul DNS IP provided",
			Config{
				ProxyUserID:      "123",
				ProxyInboundPort: 20000,
				ConsulDNSIP:      "10.0.34.16",
				IptablesProvider: &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 15001",
				"iptables -t nat -A CONSUL_DNS_REDIRECT -p udp --dport 53 -j DNAT --to-destination 10.0.34.16",
				"iptables -t nat -A CONSUL_DNS_REDIRECT -p tcp --dport 53 -j DNAT --to-destination 10.0.34.16",
				"iptables -t nat -A OUTPUT -p udp --dport 53 -j CONSUL_DNS_REDIRECT",
				"iptables -t nat -A OUTPUT -p tcp --dport 53 -j CONSUL_DNS_REDIRECT",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"Consul DNS port provided",
			Config{
				ProxyUserID:      "123",
				ProxyInboundPort: 20000,
				ConsulDNSPort:    8600,
				IptablesProvider: &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 15001",
				"iptables -t nat -A CONSUL_DNS_REDIRECT -p udp -d 127.0.0.1 --dport 53 -j DNAT --to-destination 127.0.0.1:8600",
				"iptables -t nat -A CONSUL_DNS_REDIRECT -p tcp -d 127.0.0.1 --dport 53 -j DNAT --to-destination 127.0.0.1:8600",
				"iptables -t nat -A OUTPUT -p udp -d 127.0.0.1 --dport 53 -j CONSUL_DNS_REDIRECT",
				"iptables -t nat -A OUTPUT -p tcp -d 127.0.0.1 --dport 53 -j CONSUL_DNS_REDIRECT",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"Consul DNS IP and port provided",
			Config{
				ProxyUserID:      "123",
				ProxyInboundPort: 20000,
				ConsulDNSIP:      "10.0.34.16",
				ConsulDNSPort:    8600,
				IptablesProvider: &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 15001",
				"iptables -t nat -A CONSUL_DNS_REDIRECT -p udp -d 10.0.34.16 --dport 53 -j DNAT --to-destination 10.0.34.16:8600",
				"iptables -t nat -A CONSUL_DNS_REDIRECT -p tcp -d 10.0.34.16 --dport 53 -j DNAT --to-destination 10.0.34.16:8600",
				"iptables -t nat -A OUTPUT -p udp -d 10.0.34.16 --dport 53 -j CONSUL_DNS_REDIRECT",
				"iptables -t nat -A OUTPUT -p tcp -d 10.0.34.16 --dport 53 -j CONSUL_DNS_REDIRECT",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"proxy outbound port is provided",
			Config{
				ProxyUserID:       "123",
				ProxyInboundPort:  20000,
				ProxyOutboundPort: 21000,
				IptablesProvider:  &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 21000",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"exclude inbound ports is set",
			Config{
				ProxyUserID:         "123",
				ProxyInboundPort:    20000,
				ProxyOutboundPort:   21000,
				ExcludeInboundPorts: []string{"22000", "22500"},
				IptablesProvider:    &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 21000",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -I CONSUL_PROXY_INBOUND -p tcp --dport 22000 -j RETURN",
				"iptables -t nat -I CONSUL_PROXY_INBOUND -p tcp --dport 22500 -j RETURN",
			},
		},
		{
			"exclude outbound ports is set",
			Config{
				ProxyUserID:          "123",
				ProxyInboundPort:     20000,
				ProxyOutboundPort:    21000,
				ExcludeOutboundPorts: []string{"22000", "22500"},
				IptablesProvider:     &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 21000",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -p tcp --dport 22000 -j RETURN",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -p tcp --dport 22500 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"exclude outbound CIDRs is set",
			Config{
				ProxyUserID:          "123",
				ProxyInboundPort:     20000,
				ProxyOutboundPort:    21000,
				ExcludeOutboundCIDRs: []string{"1.1.1.1", "2.2.2.2/24"},
				IptablesProvider:     &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 21000",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -d 1.1.1.1 -j RETURN",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -d 2.2.2.2/24 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"exclude UIDs is set",
			Config{
				ProxyUserID:       "123",
				ProxyInboundPort:  20000,
				ProxyOutboundPort: 21000,
				ExcludeUIDs:       []string{"456", "789"},
				IptablesProvider:  &fakeIptablesProvider{},
			},
			nil,
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 21000",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -m owner --uid-owner 456 -j RETURN",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -m owner --uid-owner 789 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
			},
		},
		{
			"additional rules are passed",
			Config{
				ProxyUserID:       "123",
				ProxyInboundPort:  20000,
				ProxyOutboundPort: 21000,
				ExcludeUIDs:       []string{"456", "789"},
				IptablesProvider:  &fakeIptablesProvider{},
			},
			[][]string{
				{"iptables", "-t", "nat", "--policy", "POSTROUTING", "ACCEPT"},
				{"iptables", "-t", "nat", "--policy", "PREROUTING", "ACCEPT"},
			},
			[]string{
				"iptables -t nat -N CONSUL_PROXY_INBOUND",
				"iptables -t nat -N CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat -N CONSUL_PROXY_OUTPUT",
				"iptables -t nat -N CONSUL_PROXY_REDIRECT",
				"iptables -t nat -N CONSUL_DNS_REDIRECT",
				"iptables -t nat -A CONSUL_PROXY_REDIRECT -p tcp -j REDIRECT --to-port 21000",
				"iptables -t nat -A OUTPUT -p tcp -j CONSUL_PROXY_OUTPUT",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -m owner --uid-owner 123 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -d 127.0.0.1/32 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_OUTPUT -j CONSUL_PROXY_REDIRECT",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -m owner --uid-owner 456 -j RETURN",
				"iptables -t nat -I CONSUL_PROXY_OUTPUT -m owner --uid-owner 789 -j RETURN",
				"iptables -t nat -A CONSUL_PROXY_IN_REDIRECT -p tcp -j REDIRECT --to-port 20000",
				"iptables -t nat -A PREROUTING -p tcp -j CONSUL_PROXY_INBOUND",
				"iptables -t nat -A CONSUL_PROXY_INBOUND -p tcp -j CONSUL_PROXY_IN_REDIRECT",
				"iptables -t nat --policy POSTROUTING ACCEPT",
				"iptables -t nat --policy PREROUTING ACCEPT",
			},
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			if c.additionalRules != nil {
				c.cfg.AddAdditionalRulesFn = func(provider Provider) {
					for _, rule := range c.additionalRules {
						provider.AddRule(rule[0], rule[1:]...)
					}
				}
			}

			err := Setup(c.cfg)
			require.NoError(t, err)
			require.Equal(t, c.expectedRules, c.cfg.IptablesProvider.Rules())
		})
	}
}

func TestSetup_errors(t *testing.T) {
	cases := []struct {
		name   string
		cfg    Config
		expErr string
	}{
		{
			"no proxy UID",
			Config{
				IptablesProvider: &iptablesExecutor{},
			},
			"ProxyUserID is required to set up traffic redirection",
		},
		{
			"no proxy inbound port",
			Config{
				ProxyUserID:       "123",
				ProxyOutboundPort: 21000,
				IptablesProvider:  &iptablesExecutor{},
			},
			"ProxyInboundPort is required to set up traffic redirection",
		},
	}

	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			err := Setup(c.cfg)
			require.EqualError(t, err, c.expErr)
		})
	}
}

type fakeIptablesProvider struct {
	rules []string
}

func (f *fakeIptablesProvider) AddRule(name string, args ...string) {
	var rule []string
	rule = append(rule, name)
	rule = append(rule, args...)

	f.rules = append(f.rules, strings.Join(rule, " "))
}

func (f *fakeIptablesProvider) ApplyRules() error {
	return nil
}

func (f *fakeIptablesProvider) Rules() []string {
	return f.rules
}