#!/usr/bin/python3
# Script by Александр Дрожевский from Rebrain
import time
import argparse
import libvirt
from lxml import etree
from lxml.builder import E


def parse_args():
    parser = argparse.ArgumentParser(description='Create snapshot for libvirt machine')
    parser.add_argument('--name', required=True, help='Machine name')
    parser.add_argument('--device', required=True, help='Device path (like vda)')
    return parser.parse_args()

def create_snapshot_node(disk_node):
    device = disk_node.xpath('/domain/devices/disk/target/@dev')[0]
    disk_path = disk_node.xpath('/domain/devices/disk/source/@file')[0]

    now = int(time.time())
    snap_path = disk_path.rsplit('.', 1)[0]
    snap_path = f'{snap_path}-{now}.snap'
    snap_name = snap_path.rsplit('/', 1)[1]

    snap = E.domainsnapshot()
    snap.append(E.name(snap_name))

    disks = E.disks()
    disk = E.disk(name=device)
    disk.append(E.source(file=snap_path))
    disks.append(disk)

    snap.append(disks)
    return snap

def get_disk_by_name(vm, name):
    raw_xml = vm.XMLDesc(0)
    xml = etree.XML(raw_xml)
    disks = xml.xpath(f'/domain/devices/disk/target[@dev="{name}"]/..')
    return disks[0]

def main():
    args = parse_args()
    session = libvirt.open('qemu:///system')
    try:
        vm = session.lookupByName(args.name)
    except libvirt.libvirtError:
        print(f'Machine with name {args.name} not found!')
        session.close()
        return
    disk = get_disk_by_name(vm, args.device)
    snapshot = create_snapshot_node(disk)
    vm.snapshotCreateXML(etree.tostring(snapshot, encoding='unicode'),
            libvirt.VIR_DOMAIN_SNAPSHOT_CREATE_DISK_ONLY | libvirt.VIR_DOMAIN_SNAPSHOT_CREATE_ATOMIC)
    session.close()

if __name__ == '__main__':
    main()