#!/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 os
import time
from datetime import datetime

import machineslib
import testlib


# libvirt-dbus snapshot APIs are available since 1.4.0, see https://github.com/libvirt/libvirt-dbus/commit/642b1b71
def supportsSnapshot(image):
    return not image.startswith("rhel-8")


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

    def testSnapshots(self):
        # Checks if difference of @time1 and @time2 is not greater than @difference (in seconds)
        def checkTimeDiff(time1, time2, difference):
            tmp = time2.split(' ')  # split "today at 13:13:13" into day and time
            if not tmp[2].startswith('00:') and not tmp[0] == "today":
                return False

            diff = datetime.strptime(time1, '%H:%M:%S') - datetime.strptime("".join(tmp[-2:]), '%I:%M%p')
            return -difference < diff.total_seconds() < difference

        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")
        if not supportsSnapshot(m.image):
            b.wait_not_present("#vm-subVmTest1-snapshots")
            return

        b.wait_in_text("#vm-subVmTest1-snapshots .pf-v5-c-empty-state", "No snapshots")

        # Check snapshot for running VM
        m.execute("virsh detach-disk --domain subVmTest1 --target vda --persistent")  # vda is raw disk, which are not supported by internal snapshots
        m.execute("qemu-img create -f qcow2 /var/lib/libvirt/images/foobar.qcow2 1M")
        m.execute("virsh attach-disk --domain subVmTest1 --source /var/lib/libvirt/images/foobar.qcow2 --target vdb --persistent")
        time1 = datetime.now().strftime("%H:%M:%S")
        m.execute("virsh snapshot-create-as --domain subVmTest1 --name snapshotB --description 'Description of snapshotB' --disk-only")

        b.reload()  # snapshots events not available yet: https://gitlab.com/libvirt/libvirt/-/issues/44
        b.enter_page('/machines')
        b.wait_in_text("#vm-subVmTest1-system-state", "Running")

        b.wait_in_text("#vm-subVmTest1-snapshot-0-name", "snapshotB")
        b.wait_in_text("#vm-subVmTest1-snapshot-0-description", "Description of snapshotB")
        b.wait_in_text("#vm-subVmTest1-snapshot-0-type", "no state saved")
        b.wait_in_text("#vm-subVmTest1-snapshot-0-parent", "No parent")
        time2 = b.text("#vm-subVmTest1-snapshot-0-date")
        self.assertTrue(checkTimeDiff(time1, time2, 60))

        # Check snapshot for shutoff VM
        self.performAction("subVmTest1", "forceOff")

        time1 = datetime.now().strftime("%H:%M:%S")
        long_description = "This is a long description" * 10
        m.execute(f"virsh snapshot-create-as --domain subVmTest1 --name snapshotC --description '{long_description}'")

        b.reload()  # snapshots events not available yet: https://gitlab.com/libvirt/libvirt/-/issues/44
        b.enter_page('/machines')

        b.wait_in_text("#vm-subVmTest1-snapshot-0-name", "snapshotC")
        b.wait_in_text("#vm-subVmTest1-snapshot-0-description", long_description)
        b.wait_in_text("#vm-subVmTest1-snapshot-0-type", "shut off")
        b.wait_in_text("#vm-subVmTest1-snapshot-0-parent", "snapshotB")
        time2 = b.text("#vm-subVmTest1-snapshot-0-date")
        self.assertTrue(checkTimeDiff(time1, time2, 60))

        b.assert_pixels(
            "#vm-subVmTest1-snapshots", "vm-snapshost-card",
            ignore=[
                "tr:nth-child(1) .snap-creation-time",
                "tr:nth-child(2) .snap-creation-time",
                "tr:nth-child(2) .tooltip-circle",
            ],
            skip_layouts=["rtl"]
        )

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

        self.createVm("subVmTest1")

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

        if not supportsSnapshot(m.image):
            b.wait_not_present("#vm-subVmTest1-snapshots")
            return
        self.goToVmPage("subVmTest1")
        b.wait_visible("#vm-subVmTest1-snapshots")

        # Shut off domain
        self.performAction("subVmTest1", "forceOff")

        class SnapshotCreateDialog(object):
            def __init__(
                self, test_obj, name=None, description=None, memory_path=None, state="shutoff", snap_num=0,
                vm_name="subVmTest1", expect_external=False, xfail=None, remove=True
            ):
                self.test_obj = test_obj
                self.name = name
                self.description = description
                self.memory_path = memory_path
                self.state = state
                self.snap_num = snap_num
                self.vm_name = vm_name
                self.remove = remove
                self.xfail = xfail
                self.expect_external = expect_external

            def execute(self):
                self.open()
                self.fill()
                self.create()
                if self.xfail is None:
                    self.verify_frontend()
                    self.verify_backend()
                    self.revert()
                    if self.remove:
                        self.cleanup()

            def open(self):
                b.click("#vm-subVmTest1-add-snapshot-button")
                b.wait_in_text(".pf-v5-c-modal-box .pf-v5-c-modal-box__header .pf-v5-c-modal-box__title", "Create snapshot")

            def fill(self):
                if self.name is not None:
                    b.set_input_text("#snapshot-create-dialog-name", self.name)
                if self.description:
                    b.set_input_text("#snapshot-create-dialog-description", self.description)
                if self.memory_path is not None:
                    b.set_input_text("#snapshot-create-dialog-memory-path input", self.memory_path)

            def assert_pixels(self):
                if self.name == 'test_snap_1':
                    b.assert_pixels("#vm-subVmTest1-create-snapshot-modal", "create-snapshot-dialog" + ("" if not self.xfail else "-error"), skip_layouts=["rtl"])

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

            def create(self):
                if not self.xfail:
                    self.assert_pixels()

                if self.xfail is None:
                    b.click(".pf-v5-c-modal-box__footer button:contains(Create)")
                    b.wait_not_present("#vm-subVmTest1-create-snapshot-modal")
                    return

                b.wait_text(".pf-v5-c-modal-box__footer button[aria-disabled=true]", "Create")

                if self.xfail == 'name':
                    self.assert_pixels()
                    b.wait_visible("#snapshot-create-dialog-name[aria-invalid=true]")
                elif self.xfail == 'memory-path':
                    b.wait_visible("#snapshot-create-dialog-memory-path .pf-v5-c-helper-text__item-text")
                else:
                    raise ValueError(f"Unknown xfail: {self.xfail}")

                self.cancel()

            def verify_frontend(self):
                if self.name:
                    b.wait_in_text(f"#vm-subVmTest1-snapshot-{self.snap_num}-name", self.name)
                else:
                    self.name = b.text(f"#vm-subVmTest1-snapshot-{self.snap_num}-name")
                if self.description:
                    b.wait_in_text(f"#vm-subVmTest1-snapshot-{self.snap_num}-description", self.description)
                else:
                    b.wait_in_text(f"#vm-subVmTest1-snapshot-{self.snap_num}-description", "No description")
                if self.state:
                    if self.state == "shutoff":
                        state = "shut off"
                    else:
                        state = self.state
                    b.wait_in_text(f"#vm-subVmTest1-snapshot-{self.snap_num}-type", state)

            def verify_backend(self):
                # Verify libvirt XML
                snap_xml = f"virsh -c qemu:///system snapshot-dumpxml --domain subVmTest1 --snapshotname {self.name}"
                xmllint_element = f"{snap_xml} | xmllint --xpath 'string(//domainsnapshot/{{prop}})' -"

                if (self.name):
                    self.test_obj.assertEqual(self.name, m.execute(xmllint_element.format(prop='name')).strip())
                if (self.description):
                    self.test_obj.assertEqual(self.description, m.execute(xmllint_element.format(prop='description')).strip())
                if (self.state):
                    self.test_obj.assertEqual(self.state, m.execute(xmllint_element.format(prop='state')).strip())

                try:
                    memory_snapshot = m.execute(xmllint_element.format(prop='memory/@snapshot')).strip()
                    ext_int = "external" if self.expect_external else "internal"

                    if self.state == "running":
                        self.test_obj.assertEqual(memory_snapshot, ext_int)
                    else:
                        self.test_obj.assertEqual(memory_snapshot, 'no')

                    disk_snapshot = m.execute(xmllint_element.format(prop='disks/disk/@snapshot')).strip()
                    self.test_obj.assertEqual(disk_snapshot, ext_int)
                    if self.name and self.expect_external:
                        disk_source = m.execute(xmllint_element.format(prop='disks/disk/source/@file')).strip()
                        self.test_obj.assertIn(self.name, disk_source)
                except AssertionError:
                    print("------ snapshot XML -------")
                    print(m.execute(snap_xml))
                    print("------ end snapshot XML -------")
                    raise

            def revert(self):
                b.click(f"#vm-subVmTest1-snapshot-{self.snap_num}-revert")
                b.wait_visible(".pf-v5-c-modal-box")
                b.click('.pf-v5-c-modal-box__footer button:contains("Revert")')
                b.wait_not_present(".pf-v5-c-modal-box")

            def cleanup(self):
                b.click(f"#delete-vm-subVmTest1-snapshot-{self.snap_num}")
                b.wait_in_text(".pf-v5-c-modal-box .pf-v5-c-modal-box__header .pf-v5-c-modal-box__title", "Delete snapshot?")
                b.wait_in_text(".pf-v5-c-modal-box__body .pf-v5-c-description-list", f"{self.name} will be deleted from {self.vm_name}")
                b.click('.pf-v5-c-modal-box__footer button:contains("Delete")')
                b.wait_not_present(f"#vm-subVmTest1-snapshot-{self.snap_num}-name:contains({self.name})")

        # No Snapshots present
        b.wait_visible("#vm-subVmTest1-add-snapshot-button")

        # external memory snapshots introduced in libvirt 9.9.0
        supports_external = not (
            m.image.startswith("rhel-8") or
            m.image in ["debian-stable", "ubuntu-2204", "ubuntu-stable", "fedora-39"])
        # Transitional code while we move ubuntu-stable from 23.10 mantic to 24.04 noble
        if m.image == "ubuntu-stable" and m.execute(". /etc/os-release; echo $VERSION_ID").strip() == "24.04":
            supports_external = True

        # HACK: deleting external snapshots for non-running VMs is broken https://bugs.debian.org/bug=1061725
        # Work around that by temporarily disabling libvirtd's AppArmor profile. AppArmor isn't installed by
        # default in Debian, and we want the rest of the test to run, thus no naughty.
        apparmor_hack = m.image in ["debian-testing"]
        if apparmor_hack:
            # we don't install apparmor-utils, so need to emulate aa-disable
            m.execute("ln -s /etc/apparmor.d/usr.sbin.libvirtd /etc/apparmor.d/disable/usr.sbin.libvirtd")
            m.execute("aa-teardown; systemctl restart apparmor libvirtd")

        # Test snapshot creation with pre-generated values
        SnapshotCreateDialog(
            self,
            expect_external=supports_external,
        ).execute()

        # Test snapshot creation with predefined values
        SnapshotCreateDialog(
            self,
            name="test_snap_1",
            description="Description of test_snap_1",
            state="shutoff",
            expect_external=supports_external,
            remove=False,
        ).execute()

        # Test inline validation for empty name
        SnapshotCreateDialog(
            self,
            name="",
            state="shutoff",
            expect_external=supports_external,
            xfail="name",
        ).execute()

        # Test inline validation for existing snapshot name
        SnapshotCreateDialog(
            self,
            name="test_snap_1",
            state="shutoff",
            expect_external=supports_external,
            xfail="name",
        ).execute()

        # Test create the same description and different name
        SnapshotCreateDialog(
            self,
            name="test_snap_des",
            description="Description of test_snap_1",
            expect_external=supports_external,
            state="shutoff",
        ).execute()

        m.execute("virsh snapshot-delete subVmTest1 test_snap_1")

        # With disk in a pool.
        # This will make a internal snapshot.

        p1 = os.path.join(self.vm_tmpdir, "vm_one")
        m.execute(f"mkdir --mode 777 {p1}")
        m.execute(f"virsh pool-create-as myPoolOne --type dir --target {p1}")
        m.execute("virsh vol-create-as myPoolOne mydisk --capacity 100M --format qcow2")
        testlib.wait(lambda: "mydisk" in m.execute("virsh vol-list myPoolOne"))

        diskXML = """<disk type="volume" device="disk">
          <driver name="qemu"/>
          <source pool="myPoolOne" volume="mydisk"/>
          <target dev="vdb" bus="virtio"/>
        </disk>""".replace("\n", "")

        m.execute(f"echo '{diskXML}' > /tmp/disk.xml; virsh attach-device --config --file /tmp/disk.xml subVmTest1")

        # virsh attach-disk doesn't send an event for offline VM changes
        b.reload()
        b.enter_page('/machines')
        b.wait_visible("#vm-subVmTest1-disks-vdb-source-volume")

        SnapshotCreateDialog(
            self,
            name="withpool",
            state="shutoff",
            expect_external=False,
        ).execute()

        self.deleteDisk("vdb")
        testlib.wait(lambda: "vdb" not in m.execute("virsh domblklist subVmTest1"), delay=1)

        # With CD-ROM with media
        m.execute("touch /var/lib/libvirt/images/mock.iso")
        m.execute("virsh attach-disk --domain subVmTest1 --config --source /var/lib/libvirt/images/mock.iso "
                  "--target sdd --type cdrom")
        # virsh attach-disk doesn't send an event for offline VM changes
        b.reload()
        b.enter_page('/machines')
        b.wait_visible("#vm-subVmTest1-disks-sdd-source-file")
        SnapshotCreateDialog(
            self,
            name="withmedia",
            state="shutoff",
            expect_external=supports_external,
        ).execute()

        # With CD-ROM drive without media
        b.click("#vm-subVmTest1-disks-sdd-eject-button")
        b.click(".pf-v5-c-modal-box__footer button:contains(Eject)")
        b.wait_not_present(".pf-v5-c-modal-box")
        b.wait_visible("#vm-subVmTest1-disks-sdd-insert")
        SnapshotCreateDialog(
            self,
            name="nomedia",
            state="shutoff",
            expect_external=supports_external,
        ).execute()

        # delete CD-DROM again so that the next test can use external snapshot again
        self.deleteDisk("sdd")
        testlib.wait(lambda: "sdd" not in m.execute("virsh domblklist subVmTest1"), delay=1)

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

        if apparmor_hack:
            # should work fine for running VMs, so re-enable AppArmor
            m.execute("rm /etc/apparmor.d/disable/usr.sbin.libvirtd")
            m.execute("aa-teardown; systemctl restart apparmor libvirtd")

        # Test snapshot creation on running VM
        SnapshotCreateDialog(
            self,
            name="test_snap_2",
            description="Description of test_snap_2",
            state="running",
            expect_external=supports_external,
        ).execute()

        if supports_external:
            SnapshotCreateDialog(
                self,
                name="test_snap_3",
                state="running",
                expect_external=True,
                memory_path="",
                xfail="memory-path",
            ).execute()

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

        self.createVm("subVmTest1")

        # Check snapshot for running VM
        m.execute("virsh detach-disk --domain subVmTest1 --target vda --config")  # vda is raw disk, which are not supported by internal snapshots
        m.execute("qemu-img create -f qcow2 /var/lib/libvirt/images/foobar.qcow2 1M")
        m.execute("virsh attach-disk --domain subVmTest1 --source /var/lib/libvirt/images/foobar.qcow2 --target vdb --subdriver qcow2 --config")
        m.execute("virsh snapshot-create-as --domain subVmTest1 --name snapshotA --description 'Description of snapshotA'")
        time.sleep(1)
        m.execute("virsh snapshot-create-as --domain subVmTest1 --name snapshotB --description 'Description of snapshotB'")
        # snapshotB is the current snapshot
        m.execute("virsh snapshot-current --domain subVmTest1 --snapshotname snapshotA")

        self.login_and_go("/machines")
        self.waitPageInit()
        self.waitVmRow("subVmTest1")
        self.goToVmPage("subVmTest1")
        if not supportsSnapshot(m.image):
            b.wait_not_present("#vm-subVmTest1-snapshots")
            return

        b.wait_visible("#vm-subVmTest1-snapshot-1-current")
        b.wait_not_present("#vm-subVmTest1-snapshot-0-current")
        self.assertEqual("yes", m.execute("virsh snapshot-info --domain subVmTest1 --snapshotname snapshotA | grep 'Current:' | awk '{print $2}'").strip())
        self.assertEqual("no", m.execute("virsh snapshot-info --domain subVmTest1 --snapshotname snapshotB | grep 'Current:' | awk '{print $2}'").strip())

        b.click("#vm-subVmTest1-snapshot-0-revert")
        b.wait_in_text(".pf-v5-c-modal-box .pf-v5-c-modal-box__header .pf-v5-c-modal-box__title", "Revert to snapshot snapshotB")
        b.click('.pf-v5-c-modal-box__footer button:contains("Revert")')

        b.wait_not_present("#vm-subVmTest1-snapshot-1-current")
        b.wait_visible("#vm-subVmTest1-snapshot-0-current")
        self.assertEqual("no", m.execute("virsh snapshot-info --domain subVmTest1 --snapshotname snapshotA | grep 'Current:' | awk '{print $2}'").strip())
        self.assertEqual("yes", m.execute("virsh snapshot-info --domain subVmTest1 --snapshotname snapshotB | grep 'Current:' | awk '{print $2}'").strip())

        # Test force revert
        # Create snapshot of shut off VM
        self.performAction("subVmTest1", "forceOff")
        m.execute("virsh snapshot-create-as --domain subVmTest1 --name snapshotC")

        # Create managed save of a running VM
        b.click("#vm-subVmTest1-system-run")
        b.wait_in_text("#vm-subVmTest1-system-state", "Running")
        m.execute("virsh managedsave subVmTest1")

        # When reverting to a snapshot with shut off VM state, but VM also has saved managed state of it's runtime memory,
        # libvirt will require using force options since there is a risk of memory corruption
        b.wait_in_text("#vm-subVmTest1-snapshot-0-name", "snapshotC")
        b.click("#vm-subVmTest1-snapshot-0-revert")
        b.wait_in_text(".pf-v5-c-modal-box .pf-v5-c-modal-box__header .pf-v5-c-modal-box__title", "Revert to snapshot snapshotC")
        b.click('.pf-v5-c-modal-box__footer button:contains("Revert")')
        # Check "Force revert" is present and Regular revert is disabled
        b.wait_in_text(".pf-v5-c-modal-box .pf-v5-c-alert__description", "revert requires force")
        b.wait_visible(".pf-v5-c-modal-box__footer button:contains(Revert):disabled")
        b.wait_visible(".pf-v5-c-modal-box__footer button:contains(Force revert)")

        b.click(".pf-v5-c-modal-box__footer button:contains(Force revert)")
        b.wait_not_present(".pf-v5-c-modal-box")

        # Check VM was reverted to a correct snapshot
        b.wait_not_present("#vm-subVmTest1-snapshot-1-current")
        b.wait_not_present("#vm-subVmTest1-snapshot-2-current")
        b.wait_visible("#vm-subVmTest1-snapshot-0-current")
        self.assertEqual("no", m.execute("virsh snapshot-info --domain subVmTest1 --snapshotname snapshotA | grep 'Current:' | awk '{print $2}'").strip())
        self.assertEqual("no", m.execute("virsh snapshot-info --domain subVmTest1 --snapshotname snapshotB | grep 'Current:' | awk '{print $2}'").strip())
        self.assertEqual("yes", m.execute("virsh snapshot-info --domain subVmTest1 --snapshotname snapshotC | grep 'Current:' | awk '{print $2}'").strip())


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