#!/usr/bin/python3

# This file is part of Cockpit.
#
# Copyright (C) 2021 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import subprocess
import re
import xml.etree.ElementTree as ET

# import Cockpit's machinery for test VMs and its browser test API
TEST_DIR = os.path.dirname(__file__)
sys.path.append(os.path.join(TEST_DIR, "common"))
sys.path.append(os.path.join(os.path.dirname(TEST_DIR), "bots/machine"))

from machineslib import VirtualMachinesCase  # noqa
from testlib import nondestructive, test_main  # noqa
from machinesxmls import USB_HOSTDEV, USB_HOSTDEV_NONEXISTENT, PCI_HOSTDEV, SCSI_HOST_HOSTDEV  # noqa

virt_xml_mock = """#!/usr/bin/python

raise Exception("Mock error message")"""


@nondestructive
class TestMachinesHostDevs(VirtualMachinesCase):

    def testHostDevicesList(self):
        b = self.browser
        m = self.machine

        self.createVm("subVmTest1")

        self.login_and_go("/machines")
        b.wait_in_text("body", "Virtual machines")
        self.waitVmRow("subVmTest1")

        b.wait_in_text("#vm-subVmTest1-system-state", "Running")

        self.goToVmPage("subVmTest1")

        b.wait_in_text("#vm-subVmTest1-hostdevs .pf-c-empty-state__body", "No host devices assigned to this VM")

        # Test hostplug of USB host device
        # A usb device might not always be present
        nodedev_list = m.execute("virsh nodedev-list")
        lines = nodedev_list.partition('\n')
        for line in lines:
            if "usb_usb" in line:
                m.execute(f"echo \"{USB_HOSTDEV}\" > /tmp/usbhostedxml")
                m.execute("virsh attach-device --domain subVmTest1 --file /tmp/usbhostedxml")

                b.wait_in_text("#vm-subVmTest1-hostdev-1-type", "usb")
                b.wait_in_text("#vm-subVmTest1-hostdev-1-vendor", "Linux Foundation")
                b.wait_in_text("#vm-subVmTest1-hostdev-1-product", "1.1 root hub")
                b.wait_in_text("#vm-subVmTest1-hostdev-1-source #1-device", "1")
                b.wait_in_text("#vm-subVmTest1-hostdev-1-source #1-bus", "1")

        m.execute("virsh destroy subVmTest1")
        b.wait_in_text("#vm-subVmTest1-system-state", "Shut off")

        # Test attachment of non-existent host device
        m.execute(f"echo \"{USB_HOSTDEV_NONEXISTENT}\" > /tmp/usbnonexistenthostedxml")
        m.execute("virsh attach-device --domain subVmTest1 --file /tmp/usbnonexistenthostedxml --config")
        b.reload()
        b.enter_page('/machines')

        b.wait_in_text("#vm-subVmTest1-hostdev-1-vendor", "0xffff")
        b.wait_in_text("#vm-subVmTest1-hostdev-1-product", "0xffff")
        b.wait_in_text("#vm-subVmTest1-hostdev-1-source #1-device", "Unspecified")
        b.wait_in_text("#vm-subVmTest1-hostdev-1-source #1-bus", "Unspecified")

        m.execute("virsh detach-device --domain subVmTest1 --file /tmp/usbnonexistenthostedxml --config")

        # Test offline attachment of PCI host device
        # A pci device should always be present
        m.execute(f"echo \"{PCI_HOSTDEV}\" > /tmp/pcihostedxml")
        m.execute("virsh attach-device --domain subVmTest1 --file /tmp/pcihostedxml --persistent")
        b.reload()
        b.enter_page('/machines')

        b.wait_in_text("#vm-subVmTest1-hostdev-1-type", "pci")
        try:
            m.execute("test -d /sys/devices/pci0000\\:00/0000\\:00\\:0f.0/")
            b.wait_in_text("#vm-subVmTest1-hostdev-1-vendor", "Red Hat, Inc")
            b.wait_in_text("#vm-subVmTest1-hostdev-1-product", "Virtio network device")
            b.assert_pixels("#vm-subVmTest1-hostdevs", "vm-details-hostdevs-card")
        except subprocess.CalledProcessError:
            pass

        b.wait_in_text("#vm-subVmTest1-hostdev-1-source #1-slot", "0000:00:0f.0")

        # QEMU version on rhel and centos-8/9-stream doesn't support scsi-host devices yet
        if not m.image.startswith("rhel-8") and not m.image.startswith("rhel-9") and m.image != "centos-8-stream" and m.image != "centos-9-stream":
            # Test the unsupported device type, e.g. scsi_host, doesn't have "Remove" button
            m.execute(f"echo \"{SCSI_HOST_HOSTDEV}\" > /tmp/scsihost_hostdevxml")
            m.execute("virsh attach-device --domain subVmTest1 --file /tmp/scsihost_hostdevxml --persistent")
            b.reload()
            b.enter_page('/machines')

            b.wait_in_text("#vm-subVmTest1-hostdev-2-type", "scsi_host")
            b.wait_not_present("#delete-vm-subVmTest1-hostdev-r2")

    def testHostDevAddSessionConnection(self):
        self.testHostDevAddSingleDevice('session')

    def testHostDevAddSingleDevice(self, connectionName='system'):
        b = self.browser
        m = self.machine

        self.run_admin("mkdir /tmp/vmdir", connectionName)
        self.addCleanup(self.run_admin, "rm -rf /tmp/vmdir/", connectionName)

        self.login_and_go("/machines")
        b.wait_in_text("body", "Virtual machines")

        self.createVm("subVmTest1", running=False, connection=connectionName)

        self.goToVmPage("subVmTest1", connectionName)
        b.wait_visible("#vm-subVmTest1-hostdevs")

        class HostDevAddDialog(object):
            def __init__(
                self, test_obj, dev_type="usb_device", dev_id=0, vm_dev_id=1, remove=True, fail_message=None
            ):
                self.test_obj = test_obj
                self.dev_type = dev_type
                self.dev_id = dev_id
                self.vm_dev_id = vm_dev_id
                self._vendor = None
                self._model = None
                self.fail_message = fail_message
                self.run_admin = test_obj.run_admin
                self.addCleanup = test_obj.addCleanup

            def execute(self):
                self.open()
                self.fill()
                self.add()
                if not self.fail_message:
                    self.verify()
                    self.verify_backend()
                    if self.remove:
                        self.remove()

            def open(self):
                b.wait_not_present(f"#vm-subVmTest1-hostdev-{self.vm_dev_id}-product")
                b.click("button#vm-subVmTest1-hostdevs-add")
                b.wait_in_text(".pf-c-modal-box .pf-c-modal-box__header .pf-c-modal-box__title", "Add host device")
                if connectionName != "session":
                    b.assert_pixels(".pf-c-modal-box", "vm-hostdevs-add-dialog")

            def fill(self):
                b.click(f"input#{self.dev_type}")
                b.set_checked(f".pf-c-table input[name='checkrow{self.dev_id}']", True)
                self._model = b.text(f"#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child({self.dev_id + 1}) td:nth-child(2)")
                self._vendor = b.text(f"#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child({self.dev_id + 1}) td:nth-child(3)")
                if self.dev_type == "pci":
                    self._slot = b.text(f"#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child({self.dev_id + 1}) td:nth-child(4) dd")

            def cancel(self):
                b.click(".pf-c-modal-box__footer button:contains(Cancel)")
                b.wait_not_present("#vm-subVmTest1-hostdevs-dialog")

            def add(self):
                self.run_admin(f"virsh -c qemu:///{connectionName} dumpxml subVmTest1 > /tmp/vmdir/vmxml1", connectionName)
                b.click(".pf-c-modal-box__footer button:contains(Add)")
                if self.fail_message:
                    b.wait_in_text(".pf-c-modal-box__footer .pf-c-alert__title", self.fail_message)
                    b.click(".pf-c-modal-box__footer button:contains(Cancel)")
                b.wait_not_present("#vm-subVmTest1-hostdevs-dialog")

            def verify(self):
                b.wait_visible(f"#vm-subVmTest1-hostdev-{self.vm_dev_id}-product")
                if (self._model != "(Undefined)"):
                    b.wait_in_text(f"#vm-subVmTest1-hostdev-{self.vm_dev_id}-product", self._model)

                b.wait_in_text(f"#vm-subVmTest1-hostdev-{self.vm_dev_id}-vendor", self._vendor)

            def verify_backend(self):
                self.run_admin(f"virsh -c qemu:///{connectionName} dumpxml subVmTest1 > /tmp/vmdir/vmxml2", connectionName)
                m.execute("diff /tmp/vmdir/vmxml1 /tmp/vmdir/vmxml2 | sed -e 's/^>//;1d' > /tmp/vmdir/vmdiff")

                if self.dev_type == "usb_device":
                    vendor_id = m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/vendor/@id)' - 2>&1 || true").strip()
                    product_id = m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/product/@id)' - 2>&1 || true").strip()

                    output = self.run_admin(f"virsh -c qemu:///{connectionName} nodedev-list --cap {self.dev_type}", connectionName)
                    devices = output.splitlines()
                    devices = list(filter(None, devices))
                    for dev in devices:
                        if self.dev_type == "usb_device":
                            self.run_admin(f"virsh -c qemu:///{connectionName} nodedev-dumpxml --device {dev} > /tmp/vmdir/nodedevxml", connectionName)
                            vendor = m.execute(f"cat /tmp/vmdir/nodedevxml | xmllint --xpath 'string(//device/capability/vendor[starts-with(@id, \"{vendor_id}\")])' - 2>&1 || true")
                            product = m.execute(f"cat /tmp/vmdir/nodedevxml | xmllint --xpath 'string(//device/capability/product[starts-with(@id, \"{product_id}\")])' - 2>&1 || true")

                            if vendor.strip() == self._vendor and product.strip() == self._model:
                                return

                elif self.dev_type == "pci":
                    domain = int(m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/address/@domain)' - 2>&1 || true"), base=16)
                    bus = int(m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/address/@bus)' - 2>&1 || true"), base=16)
                    slot = int(m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/address/@slot)' - 2>&1 || true"), base=16)
                    func = int(m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/address/@function)' - 2>&1 || true"), base=16)

                    slot_parts = re.split(r":|\.", self._slot)

                    if int(slot_parts[0], 16) == domain and int(slot_parts[1], 16) == bus and int(slot_parts[2], 16) == slot and int(slot_parts[3], 16) == func:
                        return

                raise Exception("Verification failed. No matching node device was found in VM's xml.")

            def remove(self):
                b.click(f"#delete-vm-subVmTest1-hostdev-{self.vm_dev_id}")
                b.wait_in_text(".pf-c-modal-box__body .pf-c-description-list", "subVmTest1")
                if (self._model != "(Undefined)"):
                    b.wait_in_text("#delete-resource-modal-product", self._model)
                b.wait_in_text("#delete-resource-modal-vendor", self._vendor)

                b.click('.pf-c-modal-box__footer button:contains("Remove")')
                b.wait_not_present(f"#vm-subVmTest1-hostdev-{self.vm_dev_id}-product")

        output = self.run_admin(f"virsh -c qemu:///{connectionName} nodedev-list --cap usb_device", connectionName)
        if output.strip() != "":
            HostDevAddDialog(
                self,
                dev_type="usb_device",
            ).execute()

        output = m.execute(f"virsh -c qemu:///{connectionName} nodedev-list --cap pci")
        if output.strip() != "":
            HostDevAddDialog(
                self,
                dev_type="pci",
            ).execute()

        mock_virt_xml = f"{self.vm_tmpdir}/mock-virt-xml"
        m.execute(f"echo '{virt_xml_mock}' > {mock_virt_xml}; chmod 777 {mock_virt_xml}")
        m.execute(f"mount -o bind {mock_virt_xml} /usr/bin/virt-xml")
        self.addCleanup(m.execute, "umount /usr/bin/virt-xml")
        HostDevAddDialog(
            self,
            dev_type="pci",
            fail_message="Host device could not be attached",
        ).execute()

    def testHostDevAddMultipleDevices(self, connectionName='system'):
        b = self.browser
        m = self.machine

        self.run_admin("mkdir /tmp/vmdir", connectionName)
        self.addCleanup(self.run_admin, "rm -rf /tmp/vmdir/", connectionName)

        self.login_and_go("/machines")
        b.wait_in_text("body", "Virtual machines")
        self.createVm("subVmTest1", running=False, connection=connectionName)
        self.goToVmPage("subVmTest1", connectionName)

        b.wait_visible("#vm-subVmTest1-hostdevs")
        b.wait_not_present("#vm-subVmTest1-hostdev-1-product")

        b.click("button#vm-subVmTest1-hostdevs-add")
        b.wait_in_text(".pf-c-modal-box .pf-c-modal-box__header .pf-c-modal-box__title", "Add host device")
        b.click("input#pci")

        b.set_checked(".pf-c-table input[name='checkrow0']", True)
        slot1 = b.text("#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child(1) td:nth-child(4) dd")

        b.set_checked(".pf-c-table input[name='checkrow1']", True)
        slot2 = b.text("#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child(2) td:nth-child(4) dd")

        # PCI devies will be sorted in the UI by slot
        if slot1 > slot2:
            (slot1, slot2) = (slot2, slot1)

        self.run_admin(f"virsh -c qemu:///{connectionName} dumpxml subVmTest1 > /tmp/vmdir/vmxml1", connectionName)
        b.click(".pf-c-modal-box__footer button:contains(Add)")
        b.wait_not_present("#vm-subVmTest1-hostdevs-dialog")

        b.wait_visible("#vm-subVmTest1-hostdev-1-product")
        b.wait_in_text("#1-slot", slot1)

        b.wait_visible("#vm-subVmTest1-hostdev-2-product")
        b.wait_in_text("#2-slot", slot2)

        self.run_admin(f"virsh -c qemu:///{connectionName} dumpxml subVmTest1 > /tmp/vmdir/vmxml2", connectionName)
        vm_diff = m.execute("diff /tmp/vmdir/vmxml1 /tmp/vmdir/vmxml2 | sed -e 's/^>//;1d'")  # Print difference between XMLs before and after adding host devices
        vm_diff = f'<root>{vm_diff}</root>'  # Diff contains 2 <hostdevice> elements. Add root to ease xml parsing

        root = ET.fromstring(vm_diff)

        hostdev1_address_elem = root[0].find('source').find('address')
        hostdev2_address_elem = root[1].find('source').find('address')

        hostdev1_domain = hostdev1_address_elem.get('domain')[2:]  # Remove '0x' prefix from hex number
        hostdev1_bus = hostdev1_address_elem.get('bus')[2:]
        hostdev1_slot = hostdev1_address_elem.get('slot')[2:]
        hostdev1_function = hostdev1_address_elem.get('function')[2:]
        hostdev2_domain = hostdev2_address_elem.get('domain')[2:]
        hostdev2_bus = hostdev2_address_elem.get('bus')[2:]
        hostdev2_slot = hostdev2_address_elem.get('slot')[2:]
        hostdev2_function = hostdev2_address_elem.get('function')[2:]

        slot_parts1 = re.split(r":|\.", slot1)
        slot_parts2 = re.split(r":|\.", slot2)
        # Cannot guarantee order of host devices in VM's XML, try in different order in case of failure
        try:
            self.assertEqual(slot_parts1[0], hostdev1_domain)
            self.assertEqual(slot_parts1[1], hostdev1_bus)
            self.assertEqual(slot_parts1[2], hostdev1_slot)
            self.assertEqual(slot_parts1[3], hostdev1_function)

            self.assertEqual(slot_parts2[0], hostdev2_domain)
            self.assertEqual(slot_parts2[1], hostdev2_bus)
            self.assertEqual(slot_parts2[2], hostdev2_slot)
            self.assertEqual(slot_parts2[3], hostdev2_function)
        except AssertionError:
            self.assertEqual(slot_parts1[0], hostdev2_domain)
            self.assertEqual(slot_parts1[1], hostdev2_bus)
            self.assertEqual(slot_parts1[2], hostdev2_slot)
            self.assertEqual(slot_parts1[3], hostdev2_function)

            self.assertEqual(slot_parts2[0], hostdev1_domain)
            self.assertEqual(slot_parts2[1], hostdev1_bus)
            self.assertEqual(slot_parts2[2], hostdev1_slot)
            self.assertEqual(slot_parts2[3], hostdev1_function)


if __name__ == '__main__':
    test_main()
