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

# 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 re
import subprocess
import xml.etree.ElementTree as ET

import machineslib
import testlib
from machinesxmls import PCI_HOSTDEV, SCSI_HOST_HOSTDEV, USB_HOSTDEV, USB_HOSTDEV_NONEXISTENT


class HostDevAddDialog(object):
    def __init__(
        self, test_obj, connection_name, dev_type="usb_device", dev_id=0, vm_dev_id=1, remove=True, fail_message=None
    ):
        self.test_obj = test_obj
        self.connection_name = connection_name
        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 = self.test_obj.browser
        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-v5-c-modal-box .pf-v5-c-modal-box__header .pf-v5-c-modal-box__title", "Add host device")
        if self.connection_name != "session":
            b.assert_pixels(".pf-v5-c-modal-box", "vm-hostdevs-add-dialog", skip_layouts=["rtl"])

    def fill(self):
        b = self.test_obj.browser
        b.click(f"input#{self.dev_type}")
        b.set_checked(f".pf-v5-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 = self.test_obj.browser
        b.click(".pf-v5-c-modal-box__footer button:contains(Cancel)")
        b.wait_not_present("#vm-subVmTest1-hostdevs-dialog")

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

    def verify(self):
        b = self.test_obj.browser
        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):
        m = self.test_obj.machine
        self.run_admin(f"virsh -c qemu:///{self.connection_name} dumpxml subVmTest1 > /tmp/vmdir/vmxml2",
                       self.connection_name)
        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)' -").strip()
            product_id = m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/product/@id)' -").strip()

            output = self.run_admin(f"virsh -c qemu:///{self.connection_name} nodedev-list --cap {self.dev_type}",
                                    self.connection_name)
            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:///{self.connection_name} nodedev-dumpxml --device {dev} > /tmp/vmdir/nodedevxml",
                                   self.connection_name)
                    vendor = m.execute(f"xmllint --xpath 'string(//device/capability/vendor[starts-with(@id, \"{vendor_id}\")])' - 2>&1 < /tmp/vmdir/nodedevxml")
                    product = m.execute(f"xmllint --xpath 'string(//device/capability/product[starts-with(@id, \"{product_id}\")])' < /tmp/vmdir/nodedevxml - 2>&1")

                    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)' -"), base=16)
            bus = int(m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/address/@bus)' -"), base=16)
            slot = int(m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/address/@slot)' -"), base=16)
            func = int(m.execute("cat /tmp/vmdir/vmdiff | xmllint --xpath 'string(//hostdev/source/address/@function)' -"), 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 = self.test_obj.browser
        b.click(f"#delete-vm-subVmTest1-hostdev-{self.vm_dev_id}")
        b.wait_in_text(".pf-v5-c-modal-box__body .pf-v5-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-v5-c-modal-box__footer button:contains("Remove")')
        b.wait_not_present("#delete-resource-modal")
        b.wait_not_present(f"#vm-subVmTest1-hostdev-{self.vm_dev_id}-product")


@testlib.nondestructive
class TestMachinesHostDevs(machineslib.VirtualMachinesCase):

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

        self.createVm("subVmTest1")

        self.login_and_go("/machines")
        self.waitPageInit()
        self.waitVmRow("subVmTest1")

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

        self.goToVmPage("subVmTest1")

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

        # Test hot plug 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 #device-1", "1")
                b.wait_in_text("#vm-subVmTest1-hostdev-1-source #bus-1", "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 #device-1", "Unspecified")
        b.wait_in_text("#vm-subVmTest1-hostdev-1-source #bus-1", "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", skip_layouts=["rtl"])
        except subprocess.CalledProcessError:
            pass

        b.wait_in_text("#vm-subVmTest1-hostdev-1-source #slot-1", "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):
        b = self.browser

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

        self.login_and_go("/machines", superuser=False)
        self.waitPageInit()

        self.createVm("subVmTest1", connection="session")

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

        # users can't access host devices
        HostDevAddDialog(
            self,
            "session",
            dev_type="pci",
            fail_message="Host device could not be attached",
        ).execute()

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

        m.execute("mkdir /tmp/vmdir")
        self.addCleanup(m.execute, "rm -rf /tmp/vmdir/")

        # this needs to be dynamic: some TF or custom test machines don't have such devices
        has_usb = m.execute("virsh nodedev-list --cap usb_device").strip() != ""
        has_pci = m.execute("virsh nodedev-list --cap pci").strip() != ""

        self.login_and_go("/machines")
        self.waitPageInit()

        self.createVm("subVmTest1")

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

        # Add USB devices when the VM is running
        b.wait_in_text("#vm-subVmTest1-system-state", "Running")
        if has_usb:
            HostDevAddDialog(self, "system", dev_type="usb_device").execute()

        # Check the error if selecting no devices when the VM is running
        dialog = HostDevAddDialog(self, "system", fail_message="No host device selected")
        dialog.open()
        dialog.add()
        m.execute("while virsh dumpxml subVmTest1 | grep -A 5 hostdev; do sleep 1; done")
        # Check no host devices attached after shutting off the VM
        self.performAction("subVmTest1", "forceOff")
        b.wait_in_text("#vm-subVmTest1-system-state", "Shut off")
        m.execute("while virsh dumpxml subVmTest1 | grep -A 5 hostdev; do sleep 1; done")

        if has_usb:
            HostDevAddDialog(self, "system", dev_type="usb_device").execute()

        if has_pci:
            HostDevAddDialog(self, "system", dev_type="pci",).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")
        self.waitPageInit()
        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-v5-c-modal-box .pf-v5-c-modal-box__header .pf-v5-c-modal-box__title", "Add host device")
        b.click("input#pci")

        b.set_checked(".pf-v5-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-v5-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-v5-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("#slot-1", slot1)

        b.wait_visible("#vm-subVmTest1-hostdev-2-product")
        b.wait_in_text("#slot-2", 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__':
    testlib.test_main()
