#!/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 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 skipImage, nondestructive, test_main  # noqa
from machinesxmls import USB_HOSTDEV, PCI_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-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("echo \"{0}\" > /tmp/usbhostedxml".format(USB_HOSTDEV))
                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")

        # Test offline attachment of PCI host device
        # A pci device should always be present
        m.execute("virsh destroy subVmTest1")
        b.wait_in_text("#vm-subVmTest1-state", "Shut off")
        m.execute("echo \"{0}\" > /tmp/pcihostedxml".format(PCI_HOSTDEV))
        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")

    @skipImage("exporting devices to qemu:///session not supported yet", "ubuntu-2004")
    def testHostDevAddSessionConnection(self):
        self.testHostDevAdd('session')

    def testHostDevAdd(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("#vm-subVmTest1-hostdev-{0}-product".format(self.vm_dev_id))
                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("input#{0}".format(self.dev_type))
                b.set_checked(".pf-c-table input[name='checkrow{0}']".format(self.dev_id), True)
                self._model = b.text("#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child({0}) td:nth-child(2)".format(self.dev_id + 1))
                self._vendor = b.text("#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child({0}) td:nth-child(3)".format(self.dev_id + 1))
                if self.dev_type == "pci":
                    self._slot = b.text("#vm-subVmTest1-hostdevs-dialog table tbody tr:nth-child({0}) td:nth-child(4) dd".format(self.dev_id + 1))

            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("#vm-subVmTest1-hostdev-{0}-product".format(self.vm_dev_id))
                b.wait_in_text("#vm-subVmTest1-hostdev-{0}-product".format(self.vm_dev_id), self._model)
                b.wait_in_text("#vm-subVmTest1-hostdev-{0}-vendor".format(self.vm_dev_id), 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("cat /tmp/vmdir/nodedevxml | xmllint --xpath 'string(//device/capability/vendor[starts-with(@id, \"{0}\")])' - 2>&1 || true".format(vendor_id))
                            product = m.execute("cat /tmp/vmdir/nodedevxml | xmllint --xpath 'string(//device/capability/product[starts-with(@id, \"{0}\")])' - 2>&1 || true".format(product_id))

                            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]) == domain and int(slot_parts[1]) == bus and int(slot_parts[2]) == slot and int(slot_parts[3]) == func:
                        return

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

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

        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()


if __name__ == '__main__':
    test_main()
