#!/usr/bin/python3 -cimport os, sys; os.execv(os.path.dirname(sys.argv[1]) + "/../common/pywrap", sys.argv)

import subprocess
import sys

import netlib
import packagelib
import testlib


class TestWireGuard(packagelib.PackageCase, netlib.NetworkCase):
    provision = {
        "machine1": {"address": "192.168.100.11/24", "address6": "2001:db8:face::1/64"},
        "machine2": {"address": "192.168.100.12/24", "address6": "2001:db8:face::2/64"}
    }

    def testVPN(self):
        m1 = self.machines["machine1"]
        m2 = self.machines["machine2"]
        b = self.browser

        # Disable pre-loading packagekit, dnf needs-restarting (dnf 4) consumes tons of cpu/memory on RHEL-10-1
        self.disable_preload("packagekit", machine=m2)

        self.login_and_go("/network")

        # Peer 1 (client)
        m1_port = 51820
        m1_ip4 = "10.0.0.1"
        m1_ip6 = "2001::1"
        b.click("button:contains('Add VPN')")
        b.wait_visible("#network-wireguard-settings-dialog")
        iface_name = b.val("#network-wireguard-settings-interface-name-input")
        b.wait_visible("#network-wireguard-settings-save:disabled")
        if m1.image.startswith("rhel-8"):
            b.wait_in_text(".pf-v6-c-alert", "wireguard-tools package is not installed")
            b.click("button:contains('Cancel')")
            b.wait_not_present("#network-ip-settings-dialog")
            b.wait_visible("#networking")
            b.wait_not_present(f"#networking-interfaces th:contains('{iface_name}')")
            # Skip the rest of the tests for images without wireguard-tools
            # As without it private/public key, connection over IPv4/IPv6 etc can't be tested
            return

        # Peer 2 (server)
        m2_port = 51820
        m2_ip4 = "10.0.0.2"
        m2_ip6 = "2001::2"
        b2 = self.new_browser(m2)
        m2.start_cockpit()
        if not m2.ostree_image:
            m2.execute(f"firewall-cmd --add-port={m2_port}/udp")
        b2.login_and_go("/network")
        b2.click("button:contains('Add VPN')")
        b2.wait_visible("#network-wireguard-settings-dialog")
        m2_iface_name = b2.val("#network-wireguard-settings-interface-name-input")
        b2.wait_not_val("#network-wireguard-settings-public-key input", "")
        m2_pubkey = b2.val("#network-wireguard-settings-public-key input")
        b2.set_input_text("#network-wireguard-settings-addresses-input", f"{m2_ip4}/24")
        b2.set_input_text("#network-wireguard-settings-listen-port-input", str(m2_port))

        # Validate each field, enter the right value, and then proceed to the next field
        #
        # check private-key
        self.allow_browser_errors("wg: Key is not the correct length or format")
        b.click("#network-wireguard-settings-paste-key")
        b.set_input_text("#network-wireguard-settings-private-key-input", "incorrect key")
        b.set_input_text("#network-wireguard-settings-addresses-input", m1_ip4)
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "key must be 32 bytes base64 encoded")
        b.click("#network-wireguard-settings-generated-key")

        # check public-key
        b.wait_not_val("#network-wireguard-settings-public-key input", "")
        m1_pubkey = b.val("#network-wireguard-settings-public-key input")

        # check listen-port
        b.set_input_text("#network-wireguard-settings-listen-port-input", "66000")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "out of range")
        b.set_input_text("#network-wireguard-settings-listen-port-input", "sometext")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Listen port must be a number")
        b.set_input_text("#network-wireguard-settings-listen-port-input", str(m1_port))

        b.set_input_text("#network-wireguard-settings-addresses-input", "10.0.0.1/24/56")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Addresses are not formatted correctly")
        b.set_input_text("#network-wireguard-settings-addresses-input", "10.0.0")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Invalid IP address '10.0.0'")
        b.set_input_text("#network-wireguard-settings-addresses-input", "ten.one")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Invalid IP address 'ten.one'")
        b.set_input_text("#network-wireguard-settings-addresses-input", "10 1")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Invalid IP address '10'")
        b.set_input_text("#network-wireguard-settings-addresses-input", "1.2.3.4/")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Invalid prefix or netmask")
        b.set_input_text("#network-wireguard-settings-addresses-input", "1.2.3.4  ,  5.6.7.8  1.2.3.4.5")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Invalid IP address '1.2.3.4.5'")
        b.set_input_text("#network-wireguard-settings-addresses-input", f"{m1_ip4}/24 1.2.3.4")

        # peer
        b.click("button:contains('Add peer')")
        b.wait_visible("#network-wireguard-settings-peer-0")
        b.set_input_text("#network-wireguard-settings-publickey-peer-0", m2_pubkey)
        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", "   192.168.100.12  ")  # test that the extra spaces are trimmed
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"  {m2_ip4}  ")  # test that the extra spaces are trimmed
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Peer #1 has invalid endpoint. It must be specified as host:port, e.g. 1.2.3.4:51820, [2001:db8::1]:51820 or example.com:51820")
        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", "192.168.100.12:somestring")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Peer #1 has invalid endpoint port. Port must be a number.")
        b.click("button:contains('Add peer')")
        b.wait_visible("#network-wireguard-settings-peer-1")
        b.set_input_text("#network-wireguard-settings-publickey-peer-1", m2_pubkey)
        b.set_input_text("#network-wireguard-settings-endpoint-peer-1", f"192.168.100.12:{m2_port}")
        b.set_input_text("#network-wireguard-settings-allowedips-peer-1", m2_ip4)
        b.click("button#network-wireguard-settings-btn-close-peer-0")
        b.wait_not_present("#network-wireguard-settings-peer-1")
        b.assert_pixels("#network-wireguard-settings-dialog", "networking-wireguard-add-generated",
                        ignore=["#network-wireguard-settings-private-key-input",
                                "#network-wireguard-settings-public-key",
                                "#network-wireguard-settings-publickey-peer-0"])
        b.click("#network-wireguard-settings-save")
        b.wait_not_present("#network-wireguard-settings-dialog")
        b.wait_in_text(f"#networking-interfaces tr[data-row-id='{iface_name}'] td[data-label='IP address']", f"1.2.3.4/32, {m1_ip4}/24")

        # if some wg properties are not valid, for example, if it was changed by some external tool, don't crash
        # this doesn't work on Ubuntu as the config is done through netplan; but this is just testing
        # handling of internal errors, so it's ok to skip it
        invalid_props_test = not m1.image.startswith("ubuntu")
        if invalid_props_test:
            m1.execute("sed -i '/allowed-ips/d' /etc/NetworkManager/system-connections/con-wg0.nmconnection")
            m1.execute("systemctl restart NetworkManager")
            b.reload()
            b.enter_page("/network")
            b.wait_visible("#networking")

        b.click(f"#networking-interfaces button:contains('{iface_name}')")
        b.wait_visible("#network-interface")
        b.click("#networking-edit-wg")

        if invalid_props_test:
            b.click("#network-wireguard-settings-save")
            b.wait_in_text(".pf-v6-c-alert", "has invalid allowed-ips")

        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", m2_ip4)
        b.click("#network-wireguard-settings-save")
        b.wait_not_present("#network-wireguard-settings-dialog")

        m1.execute(f"until wg show wg0 | grep -q 'allowed ips.*{m2_ip4}/32'; do sleep 1; done")
        m1.execute(f"until ip route | grep -q '10.0.0.0/24 dev wg0 proto kernel scope link src {m1_ip4} metric 50'; do sleep 1; done")

        # endpoint and port is not necessary for a peer if that peer estalishes the connectio first (i.e. the client)
        b2.click("button:contains('Add peer')")
        b2.set_input_text("#network-wireguard-settings-publickey-peer-0", m1_pubkey)
        b2.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"{m1_ip4}/32")
        b2.click("#network-wireguard-settings-save")
        b2.wait_not_present("#network-wireguard-settings-dialog")
        b2.wait_in_text(f"#networking-interfaces tr[data-row-id='{m2_iface_name}'] td[data-label='IP address']", f"{m2_ip4}/24")

        # check connection over ipv4
        try:
            m1.execute(f"ping -c 5 {m2_ip4}")
        except (subprocess.CalledProcessError, testlib.Error):
            print("-------- status on m1 ----------", file=sys.stderr)
            m1.execute("set -x; ip a >&2; ip route >&2; nmcli c >&2; wg >&2")
            print("-------- status on m2 ----------", file=sys.stderr)
            m2.execute("set -x; ip a >&2; ip route >&2; nmcli c >&2; wg >&2")
            raise

        # check connection over ipv6
        b2.click(f"#networking-interfaces button:contains('{m2_iface_name}')")

        b2.click("#networking-edit-wg")
        b2.wait_visible("#network-wireguard-settings-dialog")
        b2.set_input_text("#network-wireguard-settings-addresses-input", f"{m2_ip4}/24, {m2_ip6}/64")
        b2.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"{m1_ip4}/32,{m1_ip6}")
        b2.click("#network-wireguard-settings-save")
        b2.wait_not_present("#network-wireguard-settings-dialog")

        m2.execute(f"until wg show wg0 | grep -q 'allowed ips.*{m1_ip6}/128'; do sleep 1; done")
        m2.execute(f"until ip a show dev {m2_iface_name} | grep -q 'inet6 {m2_ip6}/64 scope global'; do sleep 0.3; done", timeout=10)

        b.click("#networking-edit-wg")
        b.wait_visible("#network-wireguard-settings-dialog")
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"{m2_ip4}/32,{m2_ip6}")

        # add IPv6 to m1 wg iface using the WireGuard modal
        # test IPv6 validation
        b.set_input_text("#network-wireguard-settings-addresses-input", "2001::1/128/64")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Addresses are not formatted correctly")
        b.set_input_text("#network-wireguard-settings-addresses-input", "2001:1/64")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Invalid IP address '2001:1'")
        b.set_input_text("#network-wireguard-settings-addresses-input", "cafe:face:nope::1/64")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Invalid IP address 'cafe:face:nope::1'")
        b.set_input_text("#network-wireguard-settings-addresses-input", "2001:db8:cafe::1/64/64")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Addresses are not formatted correctly")

        # valid IPv6 now
        b.set_input_text("#network-wireguard-settings-addresses-input", f"{m1_ip4}/24, 1.2.3.4/32, {m1_ip6}/64")
        b.click("#network-wireguard-settings-save")
        b.wait_not_present("#network-wireguard-settings-dialog")
        self.wait_for_iface_setting("IPv6", f"Address {m1_ip6}/64")

        m1.execute(f"until wg show wg0 | grep -q 'allowed ips.*{m2_ip6}/128'; do sleep 1; done")
        m1.execute(f"until ip a show dev {iface_name} | grep -q 'inet6 {m1_ip6}/64 scope global'; do sleep 0.3; done", timeout=10)

        # IPv6 is prefilled when opening the WG modal again
        b.click("#networking-edit-wg")
        b.wait_visible("#network-wireguard-settings-dialog")
        b.wait_visible(f"#network-wireguard-settings-addresses-input[value='{m1_ip4}/24, 1.2.3.4/32, {m1_ip6}/64']")
        b.click("#network-wireguard-settings-cancel")
        b.wait_not_present("#network-wireguard-settings-dialog")

        try:
            m1.execute(f"ping -6 -c 5 {m2_ip6}")
        except (subprocess.CalledProcessError, testlib.Error):
            print("-------- status on m1 ----------", file=sys.stderr)
            m1.execute("set -x; ip a >&2; ip -6 route >&2; nmcli c >&2; wg >&2")
            print("-------- status on m2 ----------", file=sys.stderr)
            m2.execute("set -x; ip a >&2; ip -6 route >&2; nmcli c >&2; wg >&2")
            raise

        # IPv6 over IPv6
        b.go("/network")
        b.wait_visible("#networking")
        b2.go("/network")
        b2.wait_visible("#networking")

        b.click("button:contains('Add VPN')")
        b.wait_visible("#network-wireguard-settings-dialog")
        b2.click("button:contains('Add VPN')")
        b2.wait_visible("#network-wireguard-settings-dialog")

        # get public keys
        b.wait_not_val("#network-wireguard-settings-public-key input", "")
        m1_wg1_pubkey = b.val("#network-wireguard-settings-public-key input")
        b2.wait_not_val("#network-wireguard-settings-public-key input", "")
        m2_wg1_pubkey = b2.val("#network-wireguard-settings-public-key input")

        wg1_iface = b.val("#network-wireguard-settings-interface-name-input")
        m1_wg1_port = "51821"
        m1_wg1_ipv6 = "2001:db8:cafe::1"
        m1_endpoint_ipv6 = "2001:db8:face::1"
        m2_wg1_port = "51821"
        m2_wg1_ipv6 = "2001:db8:cafe::2"
        if not m2.ostree_image:
            m2.execute(f"firewall-cmd --add-port={m2_wg1_port}/udp")
        m2_endpoint_ipv6 = "2001:db8:face::2"

        b.set_input_text("#network-wireguard-settings-addresses-input", f"{m1_wg1_ipv6}/64")
        b.set_input_text("#network-wireguard-settings-listen-port-input", m1_wg1_port)
        b.click("button:contains('Add peer')")
        b.wait_visible("#network-wireguard-settings-peer-0")
        b.set_input_text("#network-wireguard-settings-publickey-peer-0", m2_wg1_pubkey)
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", "2001:db8::1/128")

        # Not a valid IPv6 endpoint
        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", "[2001:db8:nope::1]:55852")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Peer #1 has invalid endpoint")

        # No port specified
        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", "2001:db8::1")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Peer #1 has invalid endpoint")

        # Not a valid port number (string)
        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", "[2001:db8::1]:portnum")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Peer #1 has invalid endpoint port")

        # Not a valid port number (high value)
        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", "[2001:db8::1]:80000")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "Peer #1 has invalid endpoint port")

        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", "[2001:db8::1]:50000")

        # Not a valid IPv6
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", "2001:db8:nope::1/128")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "peer #1 has invalid allowed-ips setting")

        # Not a valid netmask
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", "2001:db8::1/129")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "peer #1 has invalid allowed-ips setting")
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", "2001:db8::1/mask")
        b.click("#network-wireguard-settings-save")
        b.wait_in_text(".pf-v6-c-alert", "peer #1 has invalid allowed-ips setting")

        # Set valid configuration
        b.set_input_text("#network-wireguard-settings-endpoint-peer-0", f"[{m2_endpoint_ipv6}]:{m2_wg1_port}")
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"{m2_wg1_ipv6}/128")
        b.click("#network-wireguard-settings-save")
        b.wait_not_present("#network-wireguard-settings-dialog")
        b.wait_in_text(f"#networking-interfaces tr[data-row-id='{wg1_iface}'] td[data-label='IP address']", f"{m1_wg1_ipv6}/64")

        m1.execute(f"until wg show wg1 | grep -q 'allowed ips.*{m2_wg1_ipv6}/128'; do sleep 1; done")
        m1.execute(f"until ip a show dev {wg1_iface} | grep -q 'inet6 {m1_wg1_ipv6}/64 scope global'; do sleep 0.3; done", timeout=10)

        # Configuration on the second machine
        b2.set_input_text("#network-wireguard-settings-addresses-input", f"{m2_wg1_ipv6}/64")
        b2.set_input_text("#network-wireguard-settings-listen-port-input", m2_wg1_port)
        b2.click("button:contains('Add peer')")
        b2.wait_visible("#network-wireguard-settings-peer-0")
        b2.set_input_text("#network-wireguard-settings-publickey-peer-0", m1_wg1_pubkey)
        b2.set_input_text("#network-wireguard-settings-endpoint-peer-0", f"[{m1_endpoint_ipv6}]:{m1_wg1_port}")
        b2.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"{m1_wg1_ipv6}/128")
        b2.click("#network-wireguard-settings-save")
        b2.wait_not_present("#network-wireguard-settings-dialog")
        b2.wait_in_text(f"#networking-interfaces tr[data-row-id='{wg1_iface}'] td[data-label='IP address']", f"{m2_wg1_ipv6}/64")

        m2.execute(f"until wg show wg1 | grep -q 'allowed ips.*{m1_wg1_ipv6}/128'; do sleep 1; done")
        m2.execute(f"until ip a show dev {wg1_iface} | grep -q 'inet6 {m2_wg1_ipv6}/64 scope global'; do sleep 0.3; done", timeout=10)

        try:
            m1.execute(f"ping -6 -c 5 {m2_wg1_ipv6}")
        except (subprocess.CalledProcessError, testlib.Error):
            print("-------- status on m1 ----------", file=sys.stderr)
            m1.execute("set -x; ip a >&2; ip -6 route >&2; nmcli c >&2; wg >&2")
            print("-------- status on m2 ----------", file=sys.stderr)
            m2.execute("set -x; ip a >&2; ip -6 route >&2; nmcli c >&2; wg >&2")
            raise

        # IPv4 over IPv6
        m1_wg1_ipv4 = "10.0.1.1"
        m2_wg1_ipv4 = "10.0.1.2"

        b.click(f"#networking-interfaces button:contains('{wg1_iface}')")
        b.click("#networking-edit-wg")
        b.wait_visible("#network-wireguard-settings-dialog")
        b.set_input_text("#network-wireguard-settings-addresses-input", f"{m1_wg1_ipv6}/64, {m1_wg1_ipv4}/24")
        b.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"{m2_wg1_ipv6}/128,{m2_wg1_ipv4}/32")
        b.click("#network-wireguard-settings-save")
        b.wait_not_present("#network-wireguard-settings-dialog")
        self.wait_for_iface_setting("IPv4", "Address 10.0.1.1/24")

        m1.execute(f"until wg show wg1 | grep -q 'allowed ips.*{m2_wg1_ipv4}/32'; do sleep 1; done")
        m1.execute(f"until ip a show dev {wg1_iface} | grep -q 'inet {m1_wg1_ipv4}/24'; do sleep 0.3; done", timeout=10)

        b2.click(f"#networking-interfaces button:contains('{wg1_iface}')")
        b2.click("#networking-edit-wg")
        b2.wait_visible("#network-wireguard-settings-dialog")
        b2.set_input_text("#network-wireguard-settings-addresses-input", f"{m2_wg1_ipv6}/64, {m2_wg1_ipv4}/24")
        b2.set_input_text("#network-wireguard-settings-allowedips-peer-0", f"{m1_wg1_ipv6}/128,{m1_wg1_ipv4}/32")
        b2.click("#network-wireguard-settings-save")
        b2.wait_not_present("#network-wireguard-settings-dialog")
        b2.wait_in_text("#network-interface-settings div:nth-child(4) .network-interface-settings-text", "Address 10.0.1.2/24")

        m2.execute(f"until wg show wg1 | grep -q 'allowed ips.*{m1_wg1_ipv4}/32'; do sleep 1; done")
        m2.execute(f"until ip a show dev {wg1_iface} | grep -q 'inet {m2_wg1_ipv4}/24'; do sleep 0.3; done", timeout=10)

        try:
            m1.execute(f"ping -4 -c 5 {m2_wg1_ipv4}")
        except (subprocess.CalledProcessError, testlib.Error):
            print("-------- status on m1 ----------", file=sys.stderr)
            m1.execute("set -x; ip a >&2; ip -6 route >&2; nmcli c >&2; wg >&2")
            print("-------- status on m2 ----------", file=sys.stderr)
            m2.execute("set -x; ip a >&2; ip -6 route >&2; nmcli c >&2; wg >&2")
            raise

        b.go("/network")
        # install wireguard-tools from the install dialog
        if not m1.ostree_image:
            self.removePackages(["wireguard-tools"])
            self.createPackage("wireguard-tools", "1", "1")
            self.enableRepo()
            b.click("button:contains('Add VPN')")
            b.wait_visible("#dialog button:contains('Install'):enabled")
            b.click("#dialog button:contains('Install')")
            b.wait_not_present("#dialog")
            b.wait_visible("#network-wireguard-settings-dialog")
            b.wait_visible(".pf-v6-c-alert")
            b.click("button:contains('Cancel')")

        # lastly delete the interface
        b.click(f"#networking-interfaces button:contains('{iface_name}')")
        b.click("#network-interface-delete")
        b.wait_not_present(f"#networking-interfaces th:contains('{iface_name}')")


if __name__ == "__main__":
    testlib.test_main()
