#!/usr/bin/env python3
"""NAT44 ED output-feature tests"""

import random
import unittest
from scapy.layers.inet import ICMP, Ether, IP, TCP
from scapy.packet import Raw
from scapy.data import IP_PROTOS
from framework import VppTestCase, VppTestRunner
from vpp_papi import VppEnum


def get_nat44_ed_in2out_worker_index(ip, vpp_worker_count):
    if 0 == vpp_worker_count:
        return 0
    numeric = socket.inet_aton(ip)
    numeric = struct.unpack("!L", numeric)[0]
    numeric = socket.htonl(numeric)
    h = numeric + (numeric >> 8) + (numeric >> 16) + (numeric >> 24)
    return 1 + h % vpp_worker_count


class TestNAT44EDOutput(VppTestCase):
    """NAT44 ED output feature Test Case"""

    max_sessions = 1024

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.create_pg_interfaces(range(2))
        cls.interfaces = list(cls.pg_interfaces)

    @classmethod
    def tearDownClass(cls):
        super().tearDownClass()

    def setUp(self):
        super().setUp()
        for i in self.interfaces:
            i.admin_up()
            i.config_ip4()
            i.resolve_arp()
        self.vapi.nat44_ed_plugin_enable_disable(sessions=self.max_sessions, enable=1)

    def tearDown(self):
        if not self.vpp_dead:
            self.logger.debug(self.vapi.cli("show nat44 sessions"))
        super().tearDown()
        if not self.vpp_dead:
            for i in self.pg_interfaces:
                i.unconfig_ip4()
                i.admin_down()
            self.vapi.nat44_ed_plugin_enable_disable(enable=0)

    def test_static_dynamic(self):
        """Create static mapping which matches existing dynamic mapping"""

        config = self.vapi.nat44_show_running_config()
        old_timeouts = config.timeouts
        new_transitory = 2
        self.vapi.nat_set_timeouts(
            udp=old_timeouts.udp,
            tcp_established=old_timeouts.tcp_established,
            icmp=old_timeouts.icmp,
            tcp_transitory=new_transitory,
        )

        local_host = self.pg0.remote_ip4
        remote_host = self.pg1.remote_ip4
        nat_intf = self.pg1
        outside_addr = nat_intf.local_ip4

        self.vapi.nat44_add_del_address_range(
            first_ip_address=outside_addr,
            last_ip_address=outside_addr,
            vrf_id=0xFFFFFFFF,
            is_add=1,
            flags=0,
        )
        self.vapi.nat44_interface_add_del_feature(
            sw_if_index=self.pg0.sw_if_index, is_add=1
        )
        self.vapi.nat44_interface_add_del_feature(
            sw_if_index=self.pg0.sw_if_index,
            flags=VppEnum.vl_api_nat_config_flags_t.NAT_IS_INSIDE,
            is_add=1,
        )
        self.vapi.nat44_ed_add_del_output_interface(
            sw_if_index=self.pg1.sw_if_index, is_add=1
        )

        thread_index = get_nat44_ed_in2out_worker_index(
            local_host, self.vpp_worker_count
        )
        port_per_thread = int((0xFFFF - 1024) / max(1, self.vpp_worker_count))
        local_sport = 1024 + random.randint(1, port_per_thread)
        if self.vpp_worker_count > 0:
            local_sport += port_per_thread * (thread_index - 1)

        remote_dport = 10000

        pg0 = self.pg0
        pg1 = self.pg1

        # first setup a dynamic TCP session

        # SYN packet in->out
        p = (
            Ether(src=pg0.remote_mac, dst=pg0.local_mac)
            / IP(src=local_host, dst=remote_host)
            / TCP(sport=local_sport, dport=remote_dport, flags="S")
        )
        p = self.send_and_expect(pg0, [p], pg1)[0]

        self.assertEqual(p[IP].src, outside_addr)
        self.assertEqual(p[TCP].sport, local_sport)
        outside_port = p[TCP].sport

        # SYN+ACK packet out->in
        p = (
            Ether(src=pg1.remote_mac, dst=pg1.local_mac)
            / IP(src=remote_host, dst=outside_addr)
            / TCP(sport=remote_dport, dport=outside_port, flags="SA")
        )
        self.send_and_expect(pg1, [p], pg0)

        # ACK packet in->out
        p = (
            Ether(src=pg0.remote_mac, dst=pg0.local_mac)
            / IP(src=local_host, dst=remote_host)
            / TCP(sport=local_sport, dport=remote_dport, flags="A")
        )
        self.send_and_expect(pg0, [p], pg1)

        # now we have a session up, create a conflicting static mapping
        self.vapi.nat44_add_del_static_mapping(
            is_add=1,
            local_ip_address=local_host,
            external_ip_address=outside_addr,
            external_sw_if_index=0xFFFFFFFF,
            local_port=local_sport,
            external_port=outside_port,
            protocol=IP_PROTOS.tcp,
            flags=VppEnum.vl_api_nat_config_flags_t.NAT_IS_OUT2IN_ONLY,
        )

        sessions = self.vapi.nat44_user_session_dump(local_host, 0)
        self.assertEqual(1, len(sessions))

        # now send some more data over existing session - it should pass

        # in->out
        p = (
            Ether(src=pg0.remote_mac, dst=pg0.local_mac)
            / IP(src=local_host, dst=remote_host)
            / TCP(sport=local_sport, dport=remote_dport)
            / Raw("zippity zap")
        )
        self.send_and_expect(pg0, [p], pg1)

        # out->in
        p = (
            Ether(src=pg1.remote_mac, dst=pg1.local_mac)
            / IP(src=remote_host, dst=outside_addr)
            / TCP(sport=remote<style>
@media only all and (prefers-color-scheme: dark) {
.highlight .hll { background-color: #49483e }
.highlight .c { color: #75715e } /* Comment */
.highlight .err { color: #960050; background-color: #1e0010 } /* Error */
.highlight .k { color: #66d9ef } /* Keyword */
.highlight .l { color: #ae81ff } /* Literal */
.highlight .n { color: #f8f8f2 } /* Name */
.highlight .o { color: #f92672 } /* Operator */
.highlight .p { color: #f8f8f2 } /* Punctuation */
.highlight .ch { color: #75715e } /* Comment.Hashbang */
.highlight .cm { color: #75715e } /* Comment.Multiline */
.highlight .cp { color: #75715e } /* Comment.Preproc */
.highlight .cpf { color: #75715e } /* Comment.PreprocFile */
.highlight .c1 { color: #75715e } /* Comment.Single */
.highlight .cs { color: #75715e } /* Comment.Special */
.highlight .gd { color: #f92672 } /* Generic.Deleted */
.highlight .ge { font-style: italic } /* Generic.Emph */
.highlight .gi { color: #a6e22e } /* Generic.Inserted */
.highlight .gs { font-weight: bold } /* Generic.Strong */
.highlight .gu { color: #75715e } /* Generic.Subheading */
.highlight .kc { color: #66d9ef } /* Keyword.Constant */
.highlight .kd { color: #66d9ef } /* Keyword.Declaration */
.highlight .kn { color: #f92672 } /* Keyword.Namespace */
.highlight .kp { color: #66d9ef } /* Keyword.Pseudo */
.highlight .kr { color: #66d9ef } /* Keyword.Reserved */
.highlight .kt { color: #66d9ef } /* Keyword.Type */
.highlight .ld { color: #e6db74 } /* Literal.Date */
.highlight .m { color: #ae81ff } /* Literal.Number */
.highlight .s { color: #e6db74 } /* Literal.String */
.highlight .na { color: #a6e22e } /* Name.Attribute */
.highlight .nb { color: #f8f8f2 } /* Name.Builtin */
.highlight .nc { color: #a6e22e } /* Name.Class */
.highlight .no { color: #66d9ef } /* Name.Constant */
.highlight .nd { color: #a6e22e } /* Name.Decorator */
.highlight .ni { color: #f8f8f2 } /* Name.Entity */
.highlight .ne { color: #a6e22e } /* Name.Exception */
.highlight .nf { color: #a6e22e } /* Name.Function */
.highlight .nl { color: #f8f8f2 } /* Name.Label */
.highlight .nn { color: #f8f8f2 } /* Name.Namespace */
.highlight .nx { color: #a6e22e } /* Name.Other */
.highlight .py { color: #f8f8f2 } /* Name.Property */
.highlight .nt { color: #f92672 } /* Name.Tag */
.highlight .nv { color: #f8f8f2 } /* Name.Variable */
.highlight .ow { color: #f92672 } /* Operator.Word */
.highlight .w { color: #f8f8f2 } /* Text.Whitespace */
.highlight .mb { color: #ae81ff } /* Literal.Number.Bin */
.highlight .mf { color: #ae81ff } /* Literal.Number.Float */
.highlight .mh { color: #ae81ff } /* Literal.Number.Hex */
.highlight .mi { color: #ae81ff } /* Literal.Number.Integer */
.highlight .mo { color: #ae81ff } /* Literal.Number.Oct */
.highlight .sa { color: #e6db74 } /* Literal.String.Affix */
.highlight .sb { color: #e6db74 } /* Literal.String.Backtick */
.highlight .sc { color: #e6db74 } /* Literal.String.Char */
.highlight .dl { color: #e6db74 } /* Literal.String.Delimiter */
.highlight .sd { color: #e6db74 } /* Literal.String.Doc */
.highlight .s2 { color: #e6db74 } /* Literal.String.Double */
.highlight .se { color: #ae81ff } /* Literal.String.Escape */
.highlight .sh { color: #e6db74 } /* Literal.String.Heredoc